diff --git a/aiida/schedulers/plugins/direct.py b/aiida/schedulers/plugins/direct.py index 3c40a047f1..81e6c10cd7 100644 --- a/aiida/schedulers/plugins/direct.py +++ b/aiida/schedulers/plugins/direct.py @@ -150,31 +150,11 @@ 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. + lines.append(f'export OMP_NUM_THREADS={job_tmpl.job_resource.num_cores_per_mpiproc}') + 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) + lines.append(self._get_submit_script_environment_variables(job_tmpl)) if job_tmpl.rerunnable: self.logger.warning( diff --git a/aiida/schedulers/plugins/lsf.py b/aiida/schedulers/plugins/lsf.py index cee5aac892..df2dd06d87 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,22 +432,8 @@ 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) + 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 diff --git a/aiida/schedulers/plugins/pbsbaseclasses.py b/aiida/schedulers/plugins/pbsbaseclasses.py index 0d1e2ae9db..941064e99b 100644 --- a/aiida/schedulers/plugins/pbsbaseclasses.py +++ b/aiida/schedulers/plugins/pbsbaseclasses.py @@ -297,21 +297,8 @@ 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(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 diff --git a/aiida/schedulers/plugins/sge.py b/aiida/schedulers/plugins/sge.py index 839fc7d73a..1002016842 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,21 +265,8 @@ 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) + lines.append(self._get_submit_script_environment_variables(job_tmpl)) return '\n'.join(lines) diff --git a/aiida/schedulers/plugins/slurm.py b/aiida/schedulers/plugins/slurm.py index 839acc692c..685439d652 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,23 +395,8 @@ 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) + lines.append(self._get_submit_script_environment_variables(job_tmpl)) return '\n'.join(lines) diff --git a/aiida/schedulers/scheduler.py b/aiida/schedulers/scheduler.py index a152aaddd3..867964870a 100644 --- a/aiida/schedulers/scheduler.py +++ b/aiida/schedulers/scheduler.py @@ -170,6 +170,24 @@ 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 not isinstance(template.job_environment, dict): + raise ValueError('If you provide job_environment, it must be a dictionary') + + 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 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