Make slurm worker throw an error on bad job ID

lanl · May 30, 2024 · 2cb37d6 · 2cb37d6
1 parent 3a77653
commit 2cb37d6
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/beeflow/common/worker/slurm_worker.py b/beeflow/common/worker/slurm_worker.py
@@ -171,7 +171,6 @@ def query_task(self, job_id):
                 try:
                     job_state = data['jobs'][0]['job_state']
                 except (KeyError, IndexError) as exc:
-                    job_state = 'ZOMBIE'
                     raise WorkerError(f'Failed to query job {job_id}') from exc
             else:
                 # If slurmrestd does not find job make last attempt with sacct command

diff --git a/beeflow/common/worker/utils.py b/beeflow/common/worker/utils.py
@@ -1,6 +1,7 @@
 """Worker utility functions."""
 
 import subprocess
+from beeflow.common.worker.worker import WorkerError
 
 
 def get_state_sacct(job_id):
@@ -15,7 +16,7 @@ def get_state_sacct(job_id):
         state_idx = header.index('State')
         return info[state_idx]
     except (subprocess.CalledProcessError, ValueError, KeyError):
-        return 'UNKNOWN'
+        raise WorkerError(f'sacct query failed for job {job_id}')
 
 
 def parse_key_val(pair):

diff --git a/beeflow/tests/test_slurm_worker.py b/beeflow/tests/test_slurm_worker.py
@@ -7,6 +7,7 @@
 import pytest
 from beeflow.common.config_driver import BeeConfig as bc
 from beeflow.common.worker_interface import WorkerInterface
+from beeflow.common.worker.worker import WorkerError
 from beeflow.common.worker.slurm_worker import SlurmWorker
 from beeflow.common.wf_data import Task
 
@@ -95,8 +96,8 @@ def test_bad_task(slurm_worker):
 
 def test_query_bad_job_id(slurm_worker):
     """Test querying a bad job ID."""
-    job_state = slurm_worker.query_task(888)
-    assert job_state == 'UNKNOWN'
+    with pytest.raises(WorkerError):
+        job_state = slurm_worker.query_task(888)
 
 
 def test_cancel_good_job(slurm_worker):