diff --git a/aiida/schedulers/plugins/direct.py b/aiida/schedulers/plugins/direct.py index 3c40a047f1..983ad9bd08 100644 --- a/aiida/schedulers/plugins/direct.py +++ b/aiida/schedulers/plugins/direct.py @@ -150,32 +150,6 @@ def _get_submit_script_header(self, job_tmpl): if job_tmpl.custom_scheduler_commands: lines.append(job_tmpl.custom_scheduler_commands) - env_lines = [] - - if job_tmpl.job_resource and job_tmpl.job_resource.num_cores_per_mpiproc: - # since this was introduced after the environment injection below, - # it is intentionally put before it to avoid breaking current users script by overruling - # any explicit OMP_NUM_THREADS they may have set in their job_environment - env_lines.append(f'export OMP_NUM_THREADS={job_tmpl.job_resource.num_cores_per_mpiproc}') - - # Job environment variables are to be set on one single line. - # This is a tough job due to the escaping of commas, etc. - # moreover, I am having issues making it work. - # Therefore, I assume that this is bash and export variables by - # and. - if job_tmpl.job_environment: - if not isinstance(job_tmpl.job_environment, dict): - raise ValueError('If you provide job_environment, it must be a dictionary') - for key, value in job_tmpl.job_environment.items(): - env_lines.append(f'export {key.strip()}={escape_for_bash(value)}') - - if env_lines: - lines.append(empty_line) - lines.append('# ENVIRONMENT VARIABLES BEGIN ###') - lines += env_lines - lines.append('# ENVIRONMENT VARIABLES END ###') - lines.append(empty_line) - if job_tmpl.rerunnable: self.logger.warning( "The 'rerunnable' option is set to 'True', but has no effect when using the direct scheduler." @@ -199,6 +173,21 @@ 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. diff --git a/aiida/schedulers/plugins/lsf.py b/aiida/schedulers/plugins/lsf.py index cee5aac892..35a8025314 100644 --- a/aiida/schedulers/plugins/lsf.py +++ b/aiida/schedulers/plugins/lsf.py @@ -311,8 +311,6 @@ def _get_submit_script_header(self, job_tmpl): import re import string - empty_line = '' - lines = [] if job_tmpl.submit_as_hold: lines.append('#BSUB -H') @@ -434,23 +432,6 @@ def _get_submit_script_header(self, job_tmpl): if job_tmpl.custom_scheduler_commands: lines.append(job_tmpl.custom_scheduler_commands) - # Job environment variables are to be set on one single line. - # This is a tough job due to the escaping of commas, etc. - # moreover, I am having issues making it work. - # Therefore, I assume that this is bash and export variables by - # hand. - if job_tmpl.job_environment: - lines.append(empty_line) - lines.append('# ENVIRONMENT VARIABLES BEGIN ###') - if not isinstance(job_tmpl.job_environment, dict): - raise ValueError('If you provide job_environment, it must be a dictionary') - for key, value in job_tmpl.job_environment.items(): - lines.append(f'export {key.strip()}={escape_for_bash(value)}') - lines.append('# ENVIRONMENT VARIABLES END ###') - lines.append(empty_line) - - lines.append(empty_line) - # 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...) diff --git a/aiida/schedulers/plugins/pbsbaseclasses.py b/aiida/schedulers/plugins/pbsbaseclasses.py index 0d1e2ae9db..e760c9e3d8 100644 --- a/aiida/schedulers/plugins/pbsbaseclasses.py +++ b/aiida/schedulers/plugins/pbsbaseclasses.py @@ -297,22 +297,6 @@ def _get_submit_script_header(self, job_tmpl): if job_tmpl.custom_scheduler_commands: lines.append(job_tmpl.custom_scheduler_commands) - # Job environment variables are to be set on one single line. - # This is a tough job due to the escaping of commas, etc. - # moreover, I am having issues making it work. - # Therefore, I assume that this is bash and export variables by - # and. - - if job_tmpl.job_environment: - lines.append(empty_line) - lines.append('# ENVIRONMENT VARIABLES BEGIN ###') - if not isinstance(job_tmpl.job_environment, dict): - raise ValueError('If you provide job_environment, it must be a dictionary') - for key, value in job_tmpl.job_environment.items(): - lines.append(f'export {key.strip()}={escape_for_bash(value)}') - lines.append('# ENVIRONMENT VARIABLES END ###') - lines.append(empty_line) - # Required to change directory to the working directory, that is # the one from which the job was submitted lines.append('cd "$PBS_O_WORKDIR"') diff --git a/aiida/schedulers/plugins/sge.py b/aiida/schedulers/plugins/sge.py index 839fc7d73a..5f4663d993 100644 --- a/aiida/schedulers/plugins/sge.py +++ b/aiida/schedulers/plugins/sge.py @@ -150,8 +150,6 @@ def _get_submit_script_header(self, job_tmpl): import re import string - empty_line = '' - lines = [] # SGE provides flags for wd and cwd @@ -267,22 +265,6 @@ def _get_submit_script_header(self, job_tmpl): if job_tmpl.custom_scheduler_commands: lines.append(job_tmpl.custom_scheduler_commands) - # TAKEN FROM PBSPRO: - # Job environment variables are to be set on one single line. - # This is a tough job due to the escaping of commas, etc. - # moreover, I am having issues making it work. - # Therefore, I assume that this is bash and export variables by - # and. - if job_tmpl.job_environment: - lines.append(empty_line) - lines.append('# ENVIRONMENT VARIABLES BEGIN ###') - if not isinstance(job_tmpl.job_environment, dict): - raise ValueError('If you provide job_environment, it must be a dictionary') - for key, value in job_tmpl.job_environment.items(): - lines.append(f'export {key.strip()}={escape_for_bash(value)}') - lines.append('# ENVIRONMENT VARIABLES END ###') - lines.append(empty_line) - return '\n'.join(lines) def _get_submit_command(self, submit_script): diff --git a/aiida/schedulers/plugins/slurm.py b/aiida/schedulers/plugins/slurm.py index 839acc692c..07f31a119a 100644 --- a/aiida/schedulers/plugins/slurm.py +++ b/aiida/schedulers/plugins/slurm.py @@ -13,7 +13,6 @@ """ import re -from aiida.common.escaping import escape_for_bash from aiida.common.lang import type_check from aiida.schedulers import Scheduler, SchedulerError from aiida.schedulers.datastructures import JobInfo, JobState, NodeNumberJobResource @@ -263,8 +262,6 @@ def _get_submit_script_header(self, job_tmpl): # pylint: disable=too-many-statements,too-many-branches import string - empty_line = '' - lines = [] if job_tmpl.submit_as_hold: lines.append('#SBATCH -H') @@ -398,24 +395,6 @@ def _get_submit_script_header(self, job_tmpl): if job_tmpl.custom_scheduler_commands: lines.append(job_tmpl.custom_scheduler_commands) - # Job environment variables are to be set on one single line. - # This is a tough job due to the escaping of commas, etc. - # moreover, I am having issues making it work. - # Therefore, I assume that this is bash and export variables by - # and. - - if job_tmpl.job_environment: - lines.append(empty_line) - lines.append('# ENVIRONMENT VARIABLES BEGIN ###') - if not isinstance(job_tmpl.job_environment, dict): - raise ValueError('If you provide job_environment, it must be a dictionary') - for key, value in job_tmpl.job_environment.items(): - lines.append(f'export {key.strip()}={escape_for_bash(value)}') - lines.append('# ENVIRONMENT VARIABLES END ###') - lines.append(empty_line) - - lines.append(empty_line) - return '\n'.join(lines) def _get_submit_command(self, submit_script): diff --git a/aiida/schedulers/scheduler.py b/aiida/schedulers/scheduler.py index a152aaddd3..f87ccfca09 100644 --- a/aiida/schedulers/scheduler.py +++ b/aiida/schedulers/scheduler.py @@ -152,6 +152,13 @@ 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) @@ -170,6 +177,21 @@ def get_submit_script(self, job_tmpl): return '\n'.join(script_lines) + def _get_submit_script_environment_variables(self, template): # pylint: disable=no-self-use + """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. + """ + 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()] + return '\n'.join(lines) + @abc.abstractmethod def _get_submit_script_header(self, job_tmpl): """Return the submit script header, using the parameters from the job template. diff --git a/tests/schedulers/test_sge.py b/tests/schedulers/test_sge.py index 5c53fef1a0..d92e2b2cab 100644 --- a/tests/schedulers/test_sge.py +++ b/tests/schedulers/test_sge.py @@ -12,6 +12,7 @@ import logging import unittest +from aiida.common.datastructures import CodeRunMode from aiida.schedulers.datastructures import JobState from aiida.schedulers.plugins.sge import SgeScheduler from aiida.schedulers.scheduler import SchedulerError, SchedulerParsingError @@ -310,6 +311,8 @@ def test_submit_script(self): sge = SgeScheduler() job_tmpl = JobTemplate() + job_tmpl.codes_info = [] + job_tmpl.codes_run_mode = CodeRunMode.SERIAL job_tmpl.job_resource = sge.create_job_resource(parallel_env='mpi8', tot_num_mpiprocs=16) job_tmpl.working_directory = '/home/users/dorigm7s/test' job_tmpl.submit_as_hold = None @@ -325,14 +328,12 @@ def test_submit_script(self): job_tmpl.max_wallclock_seconds = '3600' # "23:59:59" job_tmpl.job_environment = {'HOME': '/home/users/dorigm7s/', 'WIENROOT': '$HOME:/WIEN2k'} - submit_script_text = sge._get_submit_script_header(job_tmpl) + submit_script_text = sge.get_submit_script(job_tmpl) self.assertTrue('#$ -wd /home/users/dorigm7s/test' in submit_script_text) self.assertTrue('#$ -N BestJobEver' in submit_script_text) self.assertTrue('#$ -q FavQ.q' in submit_script_text) self.assertTrue('#$ -l h_rt=01:00:00' in submit_script_text) - # self.assertTrue( 'export HOME=/home/users/dorigm7s/' - # in submit_script_text ) self.assertTrue('# ENVIRONMENT VARIABLES BEGIN ###' in submit_script_text) self.assertTrue("export HOME='/home/users/dorigm7s/'" in submit_script_text) self.assertTrue("export WIENROOT='$HOME:/WIEN2k'" in submit_script_text) @@ -345,15 +346,17 @@ def test_submit_script_rerunnable(self): # pylint: disable=no-self-use sge = SgeScheduler() job_tmpl = JobTemplate() + job_tmpl.codes_info = [] + job_tmpl.codes_run_mode = CodeRunMode.SERIAL job_tmpl.job_resource = sge.create_job_resource(parallel_env='mpi8', tot_num_mpiprocs=16) job_tmpl.rerunnable = True - submit_script_text = sge._get_submit_script_header(job_tmpl) + submit_script_text = sge.get_submit_script(job_tmpl) assert '#$ -r yes' in submit_script_text assert '#$ -r no' not in submit_script_text job_tmpl.rerunnable = False - submit_script_text = sge._get_submit_script_header(job_tmpl) + submit_script_text = sge.get_submit_script(job_tmpl) assert '#$ -r yes' not in submit_script_text assert '#$ -r no' in submit_script_text