Alternate implementation

Have `_get_submit_script_environment_variables` include the starting and
end markers, and do not call it in `Scheduler.get_submit_script`, but
instead let the subclasses call it in `_get_submit_script_header`
themselves. This prevents existing external plugins printing the
environment variables twice, and still the logic is provided in a single
place.

This still require external plugins to update to use the new method to
print the environment variables for it to pick up the changes if
`aiida-core` decides to improve the implementation.

aiida/schedulers/plugins/direct.py

@@ -150,6 +150,12 @@ 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_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."
@@ -173,21 +179,6 @@ def _get_submit_script_header(self, job_tmpl):
 
         return '\n'.join(lines)
 
-    def _get_submit_script_environment_variables(self, template):
-        """Return the part of the submit script header that defines environment variables.
-
-        :parameter template: a `aiida.schedulers.datastrutures.JobTemplate` instance.
-        :return: string containing environment variable declarations.
-        """
-        result = super()._get_submit_script_environment_variables(template)
-
-        if template.job_resource and template.job_resource.num_cores_per_mpiproc:
-            # This should be prepended to the environment variables from the template, such that it does not overrule
-            # any explicit OMP_NUM_THREADS that may have been defined in the ``template.job_environment``.
-            result = f'export OMP_NUM_THREADS={template.job_resource.num_cores_per_mpiproc}\n{result}'
-
-        return result
-
     def _get_submit_command(self, submit_script):
         """
         Return the string to execute to submit a given script.

aiida/schedulers/plugins/lsf.py

@@ -432,6 +432,9 @@ 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,6 +297,9 @@ 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,6 +265,9 @@ 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

@@ -395,6 +395,9 @@ 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

@@ -152,13 +152,6 @@ def get_submit_script(self, job_tmpl):
         script_lines.append(self._get_submit_script_header(job_tmpl))
         script_lines.append(empty_line)
 
-        environment_variables = self._get_submit_script_environment_variables(job_tmpl)
-        if environment_variables:
-            script_lines.append('# ENVIRONMENT VARIABLES BEGIN ###')
-            script_lines.append(environment_variables)
-            script_lines.append('# ENVIRONMENT VARIABLES END ###')
-            script_lines.append(empty_line)
-
         if job_tmpl.prepend_text:
             script_lines.append(job_tmpl.prepend_text)
             script_lines.append(empty_line)
@@ -183,13 +176,16 @@ def _get_submit_script_environment_variables(self, template):  # pylint: disable
         :parameter template: a `aiida.schedulers.datastrutures.JobTemplate` instance.
         :return: string containing environment variable declarations.
         """
-        if template.job_environment is None:
-            return ''
-
         if not isinstance(template.job_environment, dict):
             raise ValueError('If you provide job_environment, it must be a dictionary')
 
-        lines = [f'export {key.strip()}={escape_for_bash(value)}' for key, value in template.job_environment.items()]
+        lines = ['# ENVIRONMENT VARIABLES BEGIN ###']
+
+        for key, value in template.job_environment.items():
+            lines.append(f'export {key.strip()}={escape_for_bash(value)}')
+
+        lines.append('# ENVIRONMENT VARIABLES END ###')
+
         return '\n'.join(lines)
 
     @abc.abstractmethod