Merge pull request #31 from amq92/amq92-patch-scancel-squeue

Minor updates

simple_slurm/scancel.py

@@ -15,7 +15,7 @@ class SlurmScancelWrapper:
     def __init__(self, staledelta=timedelta(minutes=30)):
         self.stale_delta = staledelta
 
-    def cancel_job(self, job_id):
+    def cancel_job(self, job_id: int):
         """Sends a straightforward scancel to a job"""
         job_id = str(job_id)
         result = subprocess.run(
@@ -27,9 +27,9 @@ def cancel_job(self, job_id):
         if result.returncode != 0:
             raise RuntimeError(f"Error cancelling job: {result.stderr.strip()}")
 
-    def signal_job(self, job_id):
-        """First time it is sent to a job, tries send a sigtem to the job id
-        If sent again to the same job, attempts a sigkill instead
+    def signal_job(self, job_id: int):
+        """First time it is sent to a job, tries send a SIGTERM to the job id
+        If sent again to the same job, attempts a SIGKILL instead
         if that fails as well, involkes scancel without term arguments
         """
         job_id = str(job_id)
@@ -42,9 +42,9 @@ def signal_job(self, job_id):
             self.sigmkills[job_id] = datetime.now()
             logger.warning(f"Failed to SIGTERM {job_id}. Sending SIGKILL")
         # Just straight up kills the node via slurm
-        elif job_id in self.sigmtems:
+        elif job_id in self.sigmkills:
             signal = ""
-            logger.warning(f"Failed to SIGKILL {job_id}. Terminating!")
+            logger.warning(f"Failed to SIGKILL {job_id}. Terminating with scancel")
         result = subprocess.run(
             ["scancel", signal, job_id],
             stdout=subprocess.PIPE,

simple_slurm/squeue.py

@@ -25,11 +25,11 @@ def update_squeue(self):
         )
 
         if result.returncode != 0:
-            raise RuntimeError(f"Error running squeue: {result.stderr.strip()}")
+            raise RuntimeError(f"Error running squeue: {result.stderr}")
 
-        self.jobs = self._parse_output(result.stdout.strip())
+        self.jobs = self._parse_output(result.stdout)
 
-    def _is_valid_csv_format(self, format_str):
+    def _is_valid_csv_format(self, format_str: str):
         """validates that the output is a valid csv"""
         try:
             sniffer = csv.Sniffer()
@@ -40,11 +40,11 @@ def _is_valid_csv_format(self, format_str):
         except csv.Error:
             return False
 
-    def _parse_output(self, output):
+    def _parse_output(self, output: str):
         """converts the stdout into a python dictionary
         each key is a jobid as integer
         """
-        csv_file = StringIO(output)
+        csv_file = StringIO(output.strip())
         reader = csv.DictReader(
             csv_file, delimiter=",", quotechar='"', skipinitialspace=True
         )
@@ -58,7 +58,7 @@ def display_jobs(self):
         for job in self.jobs.values():
             print(job)
 
-    def get_filtered_jobs(self, name_seek):
+    def get_filtered_jobs(self, name_seek: str):
         """Filters jobs by name"""
         matching_jobs = {}
         for job_id, job in self.jobs.items():

test/test_core.py

@@ -312,7 +312,7 @@ def test_19_sbatch_execution_with_job_file(self):
         self.assertIn("Hello!", contents)
         self.assertIn(f"Submitted batch job {job_id}", stdout)
 
-    def test_19_add_cmd_single(self):
+    def test_20_add_cmd_single(self):
         slurm = Slurm(
             "-a",
             "3-11",
@@ -342,7 +342,7 @@ def test_19_add_cmd_single(self):
         )
         self.assertEqual(self.script + "\n" + self.commands, str(slurm))
 
-    def test_19_add_cmd_multiple(self):
+    def test_21_add_cmd_multiple(self):
         slurm = Slurm(
             "-a",
             "3-11",
@@ -366,7 +366,7 @@ 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):
+    def test_22_parsable_sbatch_execution(self):
         with io.StringIO() as buffer:
             with contextlib.redirect_stdout(buffer):
                 slurm = Slurm(contiguous=True, parsable=True)