add parsable flag

simple_slurm/arguments.txt

@@ -70,6 +70,7 @@ open_mode
 output,o
 overcommit,O
 oversubscribe,s
+parsable
 partition,p
 power
 prefer

simple_slurm/core.py

@@ -134,6 +134,9 @@ def srun(self, run_cmd: str, srun_cmd: str = 'srun') -> int:
 
         result = subprocess.run(cmd, shell=True, check=True)
         return result.returncode
+    @property
+    def is_parsable(self) -> bool:
+        return getattr(self.namespace, 'parsable', None) is not None
 
     def sbatch(self, *run_cmd: str, convert: bool = True,
                verbose: bool = True, sbatch_cmd: str = 'sbatch',
@@ -177,12 +180,24 @@ def sbatch(self, *run_cmd: str, convert: bool = True,
                 'EOF',
             ))
         result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE)
-        success_msg = 'Submitted batch job'
-        stdout = result.stdout.decode('utf-8')
-        assert success_msg in stdout, result.stderr
+        # init for clarity
+        job_id = None
+        stdout = ''
+        if self.is_parsable:
+            # gather the first line from stdout
+            stdout = result.stdout.decode().strip()
+            # parsable will be of format job_id[:cluster]
+            # ref: https://slurm.schedmd.com/sbatch.html#OPT_parsable
+            job_id = int(stdout.split(':')[0])
+            assert result.returncode == 0, result.stderr
+        else:
+            success_msg = 'Submitted batch job'
+            stdout = result.stdout.decode()
+            assert success_msg in stdout, result.stderr
+            job_id = int(stdout.split(' ')[3])
+        assert job_id is not None, 'this should never happen, assert for linter'
         if verbose:
             print(stdout)
-        job_id = int(stdout.split(' ')[3])
         return job_id
 
 

test/test_core.py

@@ -211,7 +211,7 @@ def test_15_sbatch_execution(self):
         with open(out_file, 'r') as fid:
             contents = fid.read()
         os.remove(out_file)
-
+        self.assertFalse(slurm.is_parsable)
         self.assertIsInstance(job_id, int)
         self.assertIn('Hello!', contents)
         self.assertIn(f'Submitted batch job {job_id}', stdout)
@@ -326,6 +326,28 @@ def test_19_add_cmd_multiple(self):
         slurm.add_cmd('echo "done"')
         self.assertEqual(self.script + '\n' + self.commands, str(slurm))
 
+    def test_20_parsable_sbatch_execution(self):
+        with io.StringIO() as buffer:
+            with contextlib.redirect_stdout(buffer):
+                slurm = Slurm(contiguous=True, parsable=True)
+                if shutil.which('sbatch') is not None:
+                    job_id = slurm.sbatch('echo Hello!')
+                else:
+                    with patch('subprocess.run', subprocess_sbatch_parsable):
+                        job_id = slurm.sbatch('echo Hello!')
+                stdout = buffer.getvalue()
+
+        out_file = f'slurm-{job_id}.out'
+        while True:  # wait for job to finalize
+            if os.path.isfile(out_file):
+                break
+        with open(out_file, 'r') as fid:
+            contents = fid.read()
+        os.remove(out_file)
+        self.assertTrue(slurm.is_parsable)
+        self.assertIsInstance(job_id, int)
+        self.assertIn('Hello!', contents)
+        self.assertEqual(f'{job_id}\n', stdout)
 
 
 def subprocess_srun(*args, **kwargs):
@@ -342,6 +364,15 @@ def subprocess_sbatch(*args, **kwargs):
     return subprocess.CompletedProcess(*args, returncode=1,
                                        stdout=stdout.encode('utf-8'))
 
+def subprocess_sbatch_parsable(*args, **kwargs):
+    job_id = 1234
+    out_file = f'slurm-{job_id}.out'
+    with open(out_file, 'w') as fid:
+        fid.write('Hello!!!\n')
+    stdout = str(job_id)
+    return subprocess.CompletedProcess(*args, returncode=0,
+                                       stdout=stdout.encode('utf-8'))
+
 
 if __name__ == '__main__':
     unittest.main()