diff --git a/pysqa/ext/modular.py b/pysqa/ext/modular.py index 274a217f..d1e703a6 100644 --- a/pysqa/ext/modular.py +++ b/pysqa/ext/modular.py @@ -25,6 +25,7 @@ def submit_job( run_time_max=None, dependency_list=None, command=None, + **kwargs ): """ @@ -49,6 +50,7 @@ def submit_job( memory_max=memory_max, run_time_max=run_time_max, command=command, + **kwargs ) cluster_module = self._queue_to_cluster_dict[queue] commands = self._switch_cluster_command( diff --git a/pysqa/utils/basic.py b/pysqa/utils/basic.py index 4ef77abd..98b3c48d 100644 --- a/pysqa/utils/basic.py +++ b/pysqa/utils/basic.py @@ -123,6 +123,7 @@ def submit_job( run_time_max=None, dependency_list=None, command=None, + **kwargs ): """ @@ -151,6 +152,7 @@ def submit_job( memory_max=memory_max, run_time_max=run_time_max, command=command, + **kwargs ) out = self._execute_command( commands=self._list_command_to_be_executed( @@ -319,6 +321,7 @@ def _write_queue_script( memory_max=None, run_time_max=None, command=None, + **kwargs ): """ @@ -344,6 +347,7 @@ def _write_queue_script( memory_max=memory_max, run_time_max=run_time_max, command=command, + **kwargs ) if not os.path.exists(working_directory): os.makedirs(working_directory) @@ -361,6 +365,7 @@ def _job_submission_template( memory_max=None, run_time_max=None, command=None, + **kwargs ): """ @@ -397,6 +402,7 @@ def _job_submission_template( memory_max=memory_max, run_time_max=run_time_max, command=command, + **kwargs ) @staticmethod diff --git a/tests/config/slurm/queue.yaml b/tests/config/slurm/queue.yaml index 0da45996..417eb2db 100644 --- a/tests/config/slurm/queue.yaml +++ b/tests/config/slurm/queue.yaml @@ -2,3 +2,4 @@ queue_type: SLURM queue_primary: slurm queues: slurm: {cores_max: 100, cores_min: 10, run_time_max: 259200, script: slurm.sh} + slurm_extra: {cores_max: 100, cores_min: 10, run_time_max: 259200, script: slurm_extra.sh} diff --git a/tests/config/slurm/slurm_extra.sh b/tests/config/slurm/slurm_extra.sh new file mode 100644 index 00000000..f7f177d2 --- /dev/null +++ b/tests/config/slurm/slurm_extra.sh @@ -0,0 +1,21 @@ +#!/bin/bash +#SBATCH --output=time.out +#SBATCH --job-name={{job_name}} +#SBATCH --chdir={{working_directory}} +#SBATCH --get-user-env=L +#SBATCH --partition=slurm +{%- if account %} +#SBATCH --account={{account}} +{%- endif %} +{%- if exclude %} +#SBATCH --exclude={{exclude}} +{%- endif %} +{%- if run_time_max %} +#SBATCH --time={{ [1, run_time_max // 60]|max }} +{%- endif %} +{%- if memory_max %} +#SBATCH --mem={{memory_max}}G +{%- endif %} +#SBATCH --cpus-per-task={{cores}} + +{{command}} diff --git a/tests/test_queueadapter.py b/tests/test_queueadapter.py index 3fbeed13..fcb99c13 100644 --- a/tests/test_queueadapter.py +++ b/tests/test_queueadapter.py @@ -452,3 +452,31 @@ def test_write_queue(self): echo \"hello\"""" self.assertEqual(content, output) os.remove("run_queue.sh") + + def test_write_queue_extra_keywords(self): + self.slurm._adapter._write_queue_script( + queue="slurm_extra", + job_name=None, + working_directory=None, + cores=None, + memory_max=None, + run_time_max=None, + command="echo \"hello\"", + account="123456" + ) + with open("run_queue.sh", "r") as f: + content = f.read() + output = """\ +#!/bin/bash +#SBATCH --output=time.out +#SBATCH --job-name=None +#SBATCH --chdir=. +#SBATCH --get-user-env=L +#SBATCH --partition=slurm +#SBATCH --account=123456 +#SBATCH --time=4320 +#SBATCH --cpus-per-task=10 + +echo \"hello\"""" + self.assertEqual(content, output) + os.remove("run_queue.sh") \ No newline at end of file