Merge pull request #36 from amq92/srun-add_cmd

srun add cmd

Author: amq92

simple_slurm/__about__.py

@@ -1 +1 @@
-__version__ = "0.3.4"
+__version__ = "0.3.5"

simple_slurm/core.py

@@ -129,17 +129,30 @@ def script(self, shell: str = "/bin/sh", convert: bool = True):
         script = "\n".join((arguments, commands)).strip() + "\n"
         return script
 
-    def srun(self, run_cmd: str, srun_cmd: str = "srun") -> int:
+    def srun(self, *run_cmd: str, connector: str = ";", srun_cmd: str = "srun") -> int:
         """Run the srun command with all the (previously) set arguments and
-        the provided command in 'run_cmd'.
+        the provided commands in 'run_cmd' alongside with the previously set
+        commands using 'add_cmd'.
+
+        Note that 'run_cmd' can accept multiple arguments. Thus, any of the
+        other arguments must be given as key-value pairs :
+            > slurm.srun('echo "Hello"')
+            > slurm.srun('echo', '"Hello"')
+            > slurm.srun('echo', '"Hello"', verbose=False)
+
+        The 'connector' parameter defines how multiple commands are connected:
+            - A ; B   : Run A and then B, regardless of success of A.
+            - A && B  : Run B if and only if A succeeded.
+            - A || B  : Run B if and only if A failed.
         """
         args = (
             f"--{self._valid_key(k)}" + f"={v}" if len(v) else ""
             for k, v in vars(self.namespace).items()
             if v is not None
         )
-        cmd = " ".join((srun_cmd, *args, run_cmd))
-
+        self.add_cmd(*run_cmd)
+        commands = f" {connector} ".join(self.run_cmds)
+        cmd = " ".join((srun_cmd, *args, commands))
         result = subprocess.run(cmd, shell=True, check=True)
         return result.returncode
 
@@ -165,6 +178,7 @@ def sbatch(
             > slurm.sbatch('python main.py')
             > slurm.sbatch('python', 'main.py')
             > slurm.sbatch('python', 'main.py', verbose=False)
+
         This function employs the 'here document' syntax, which requires that
         bash variables be scaped. This behavior is default, set 'convert'
         to False to disable it.
@@ -176,7 +190,7 @@ def sbatch(
             > EOF
 
         For such reason if any bash variable is employed by the 'run_command',
-        the '$' should be scaped into '\$'. This behavior is default, set
+        the '$' should be scaped into '\\$'. This behavior is default, set
         'convert' to False to disable it.
 
         If the argument 'job_file' is used, the script will be written to the