Add option to add aditional arguments to SLURM submission scripts

pyiron · Feb 29, 2024 · 0c7b519 · 0c7b519
1 parent 56f0158
commit 0c7b519
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/pympipool/shared/interface.py b/pympipool/shared/interface.py
@@ -76,6 +76,7 @@ def __init__(
         threads_per_core=1,
         gpus_per_core=0,
         oversubscribe=False,
+        command_line_argument_lst=[],
     ):
         super().__init__(
             cwd=cwd,
@@ -84,6 +85,7 @@ def __init__(
         )
         self._threads_per_core = threads_per_core
         self._gpus_per_core = gpus_per_core
+        self._command_line_argument_lst = command_line_argument_lst
 
     def generate_command(self, command_lst):
         command_prepend_lst = generate_slurm_command(
@@ -92,6 +94,7 @@ def generate_command(self, command_lst):
             threads_per_core=self._threads_per_core,
             gpus_per_core=self._gpus_per_core,
             oversubscribe=self._oversubscribe,
+            command_line_argument_lst=self._command_line_argument_lst,
         )
         return super().generate_command(
             command_lst=command_prepend_lst + command_lst,
@@ -109,7 +112,7 @@ def generate_mpiexec_command(cores, oversubscribe=False):
 
 
 def generate_slurm_command(
-    cores, cwd, threads_per_core=1, gpus_per_core=0, oversubscribe=False
+    cores, cwd, threads_per_core=1, gpus_per_core=0, oversubscribe=False, command_line_argument_lst=[],
 ):
     command_prepend_lst = [SLURM_COMMAND, "-n", str(cores)]
     if cwd is not None:
@@ -120,4 +123,6 @@ def generate_slurm_command(
         command_prepend_lst += ["--gpus-per-task=" + str(gpus_per_core)]
     if oversubscribe:
         command_prepend_lst += ["--oversubscribe"]
+    if len(command_line_argument_lst) > 0:
+        command_prepend_lst += command_line_argument_lst
     return command_prepend_lst
diff --git a/pympipool/slurm/executor.py b/pympipool/slurm/executor.py
@@ -60,6 +60,7 @@ def __init__(
         init_function=None,
         cwd=None,
         hostname_localhost=False,
+        command_line_argument_lst=[],
     ):
         super().__init__()
         self._set_process(
@@ -78,6 +79,7 @@ def __init__(
                         "gpus_per_core": int(gpus_per_worker / cores_per_worker),
                         "cwd": cwd,
                         "oversubscribe": oversubscribe,
+                        "command_line_argument_lst": command_line_argument_lst,
                     },
                 )
                 for _ in range(max_workers)