Scheduler: Move setting of env variables into base class

aiida/engine/processes/calcjobs/calcjob.py

@@ -840,7 +840,6 @@ def presubmit(self, folder: Folder) -> CalcInfo:
         job_tmpl = JobTemplate()
         job_tmpl.submit_as_hold = False
         job_tmpl.rerunnable = self.options.get('rerunnable', False)
-        job_tmpl.job_environment = {}
         # 'email', 'email_on_started', 'email_on_terminated',
         job_tmpl.job_name = f'aiida-{self.node.pk}'
         job_tmpl.sched_output_path = self.options.scheduler_stdout

aiida/schedulers/datastructures.py

@@ -261,7 +261,7 @@ class JobTemplate(DefaultFieldsAttributeDict):  # pylint: disable=too-many-insta
         before the execution of the code.
       * ``environment_variables_double_quotes``: if set to True, use double quotes
         instead of single quotes to escape the environment variables specified
-        in ``environment_variables``.
+        in ``job_environment``.
       * ``working_directory``: the working directory for this job. During
         submission, the transport will first do a 'chdir' to this directory,
         and then possibly set a scheduler parameter, if this is supported

aiida/schedulers/plugins/direct.py

@@ -160,9 +160,6 @@ def _get_submit_script_header(self, job_tmpl):
         if job_tmpl.job_resource and job_tmpl.job_resource.num_cores_per_mpiproc:
             lines.append(f'export OMP_NUM_THREADS={job_tmpl.job_resource.num_cores_per_mpiproc}')
 
-        if job_tmpl.job_environment:
-            lines.append(self._get_submit_script_environment_variables(job_tmpl))
-
         if job_tmpl.rerunnable:
             self.logger.warning(
                 "The 'rerunnable' option is set to 'True', but has no effect when using the direct scheduler."

aiida/schedulers/plugins/lsf.py

@@ -439,9 +439,6 @@ def _get_submit_script_header(self, job_tmpl):
         if job_tmpl.custom_scheduler_commands:
             lines.append(job_tmpl.custom_scheduler_commands)
 
-        if job_tmpl.job_environment:
-            lines.append(self._get_submit_script_environment_variables(job_tmpl))
-
         # The following seems to be the only way to copy the input files
         # to the node where the computation are actually launched (the
         # -f option of bsub that does not always work...)

aiida/schedulers/plugins/pbsbaseclasses.py

@@ -297,9 +297,6 @@ def _get_submit_script_header(self, job_tmpl):
         if job_tmpl.custom_scheduler_commands:
             lines.append(job_tmpl.custom_scheduler_commands)
 
-        if job_tmpl.job_environment:
-            lines.append(self._get_submit_script_environment_variables(job_tmpl))
-
         # Required to change directory to the working directory, that is
         # the one from which the job was submitted
         lines.append('cd "$PBS_O_WORKDIR"')

aiida/schedulers/plugins/sge.py

@@ -265,9 +265,6 @@ def _get_submit_script_header(self, job_tmpl):
         if job_tmpl.custom_scheduler_commands:
             lines.append(job_tmpl.custom_scheduler_commands)
 
-        if job_tmpl.job_environment:
-            lines.append(self._get_submit_script_environment_variables(job_tmpl))
-
         return '\n'.join(lines)
 
     def _get_submit_command(self, submit_script):

aiida/schedulers/plugins/slurm.py

@@ -393,9 +393,6 @@ def _get_submit_script_header(self, job_tmpl):
         if job_tmpl.custom_scheduler_commands:
             lines.append(job_tmpl.custom_scheduler_commands)
 
-        if job_tmpl.job_environment:
-            lines.append(self._get_submit_script_environment_variables(job_tmpl))
-
         return '\n'.join(lines)
 
     def _get_submit_command(self, submit_script):

aiida/schedulers/scheduler.py

@@ -13,7 +13,7 @@
 import abc
 import typing as t
 
-from aiida.common import exceptions, log
+from aiida.common import exceptions, log, warnings
 from aiida.common.datastructures import CodeRunMode
 from aiida.common.escaping import escape_for_bash
 from aiida.common.lang import classproperty
@@ -155,9 +155,22 @@ def get_submit_script(self, job_tmpl: JobTemplate) -> str:
             script_lines.append('#!/bin/bash')
         else:
             raise ValueError(f'Invalid shebang set: {job_tmpl.shebang}')
-        script_lines.append(self._get_submit_script_header(job_tmpl))
+
+        script_header = self._get_submit_script_header(job_tmpl)
+        script_lines.append(script_header)
         script_lines.append(empty_line)
 
+        if '# ENVIRONMENT VARIABLES BEGIN ###' in script_header:
+            warnings.warn_deprecation(
+                f'Environment variables added by `{self.__class__.__name__}._get_submit_script_environment_variables`, '
+                'however, this is no longer necessary and automatically done by the base `Scheduler` class.',
+                version=3
+            )
+
+        if job_tmpl.job_environment:
+            script_lines.append(self._get_submit_script_environment_variables(job_tmpl))
+            script_lines.append(empty_line)
+
         if job_tmpl.prepend_text:
             script_lines.append(job_tmpl.prepend_text)
             script_lines.append(empty_line)

tests/schedulers/test_all.py

@@ -0,0 +1,66 @@
+# -*- coding: utf-8 -*-
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+# pylint: disable=redefined-outer-name
+"""Tests run for all plugins in :mod:`aiida.schedulers.plugins`."""
+from __future__ import annotations
+
+import pytest
+
+from aiida.common.datastructures import CodeRunMode
+from aiida.plugins import SchedulerFactory, entry_point
+from aiida.schedulers import Scheduler
+from aiida.schedulers.datastructures import JobTemplate, JobTemplateCodeInfo, NodeNumberJobResource
+
+
+def get_scheduler_entry_point_names() -> list[str]:
+    """Return the list of entry point names for the scheduler plugins registered by ``aiida-core``.
+
+    :return: List of entry point names.
+    """
+    entry_points_names = entry_point.get_entry_point_names('aiida.schedulers')
+    return [name for name in entry_points_names if name.startswith('core.')]
+
+
+def get_scheduler_job_resource(scheduler):
+    """Create a ``JobResource`` instance for the given scheduler with default resources."""
+    if issubclass(scheduler.job_resource_class, NodeNumberJobResource):
+        resources = {'num_machines': 1, 'num_mpiprocs_per_machine': 1}
+    else:
+        resources = {'parallel_env': 'env', 'tot_num_mpiprocs': 1}
+
+    return scheduler.create_job_resource(**resources)
+
+
+@pytest.fixture(scope='function', params=get_scheduler_entry_point_names())
+def scheduler(request) -> Scheduler:
+    """Fixture that parametrizes over all the ``Scheduler`` implementations registered by ``aiida-core``."""
+    return SchedulerFactory(request.param)()
+
+
+@pytest.mark.parametrize('environment_variables_double_quotes', (True, False))
+def test_job_environment(scheduler, environment_variables_double_quotes):
+    """Test that ``JobTemplate.job_environment`` make it into the submission script.
+
+    Also tests that the ``JobTemplate.environment_variables_double_quotes`` is respected.
+    """
+    job_template = JobTemplate()
+    job_template.codes_info = [JobTemplateCodeInfo()]
+    job_template.codes_run_mode = CodeRunMode.SERIAL
+    job_template.job_resource = get_scheduler_job_resource(scheduler)
+    job_template.environment_variables_double_quotes = environment_variables_double_quotes
+    job_template.job_environment = {'SOME_STRING': 'value', 'SOME_INTEGER': 1}
+    script = scheduler.get_submit_script(job_template)
+
+    if environment_variables_double_quotes:
+        assert 'export SOME_STRING="value"' in script
+        assert 'export SOME_INTEGER="1"' in script
+    else:
+        assert "export SOME_STRING='value'" in script
+        assert "export SOME_INTEGER='1'" in script