Add extra keyword arguments

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(

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

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}

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}}

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")