Scheduler: Move setting of env variables into base class

The `Scheduler` base class implements the concrete method
`_get_submit_script_environment_variables` which formats the lines for
the submission script that set the environment variables that were
defined in the `metadata.options.environment_variables` input.

However, it was left up to the plugins to actually call this method in
the `_get_submit_script_header` and append the lines to the created
submission script header. This would leave it open to plugins forgetting
to add this call and essentially have the environment variables
specified by the caller of the `CalcJob` be ignored.

Here we move the call of `_get_submit_script_environment_variables` to
the `get_submit_script` of the base class, guaranteeing that it will
always be called. This still leaves the possibility to the plugin to
override the method and customize the behavior.

A check is added just before adding the new call, to check whether the
environment variables have already been added indicating the plugin is
also still calling the method. In this case a deprecation warning is
emitted suggesting that the call can be removed from the plugin.

With this change we do assume that setting the environment variables are
idempotent as it is possible that the lines will be included twice until
all external plugins have been updated, but this should be a relatively
safe assumption.

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