Wait for shared-filesystem output of completed Slurm jobs

src/executorlib/standalone/command_pysqa.py

@@ -10,6 +10,15 @@
 # on this status query, but it must not be issued on every poll of the (much faster) refresh_rate
 # loop, as that would flood the queuing system commands (e.g. squeue/sacct) with requests.
 _JOB_STATUS_CHECK_INTERVAL = 30.0
+_JOB_OUTPUT_GRACE_PERIOD = 60.0
+_SLURM_RUNNING_STATES = {
+    "COMPLETING",
+    "CONFIGURING",
+    "PENDING",
+    "RESIZING",
+    "RUNNING",
+    "SUSPENDED",
+}
 
 
 def pysqa_job_output_validation(
@@ -23,9 +32,10 @@ def pysqa_job_output_validation(
 ) -> bool:
     """
     Check whether the queuing system job backing a task has died without ever writing its output
-    file. Only applies to queuing system backends (pysqa) and is throttled to at most once every
-    ``job_status_check_interval`` seconds per task, to avoid flooding the queuing system with
-    status queries on every poll of the (much faster) refresh_rate loop.
+    file. Successful Slurm and Flux jobs are given a grace period for their output to appear on a
+    shared filesystem. Status checks are throttled to at most once every
+    ``job_status_check_interval`` seconds per task, to avoid flooding the queuing system on every
+    poll of the refresh loop.
 
     A dead job is recognized in two ways, since queuing systems differ in whether they drop
     terminated jobs from their listing: slurm's squeue removes a job as soon as it is gone
@@ -47,21 +57,87 @@ def pysqa_job_output_validation(
             errored out, and still has no output.
     """
     now = monotonic()
-    last_checked = (
-        status_check_dict.get(task_key, 0.0) if status_check_dict is not None else 0.0
+    check_state = (
+        status_check_dict.get(task_key, {}) if status_check_dict is not None else {}
     )
+    if isinstance(check_state, dict):
+        last_checked = check_state.get("last_checked", 0.0)
+        completed_at = check_state.get("completed_at")
+    else:
+        last_checked = check_state
+        completed_at = None
+    if completed_at is not None:
+        return now - completed_at >= _JOB_OUTPUT_GRACE_PERIOD
     if now - last_checked < job_status_check_interval:
         return False
     if status_check_dict is not None:
-        status_check_dict[task_key] = now
+        status_check_dict[task_key] = {"last_checked": now}
     status = pysqa_get_status_of_job(
         queue_id=queue_id,
         config_directory=pysqa_config_directory,
         backend=backend,
     )
+    if backend == "slurm" and status is None:
+        return _validate_slurm_job_output(
+            task_key=task_key,
+            file_name=file_name,
+            queue_id=queue_id,
+            now=now,
+            status_check_dict=status_check_dict,
+        )
+    if backend == "flux" and status == "finished":
+        return _mark_job_completed(task_key, now, status_check_dict)
     return (status is None or status == "error") and not os.path.exists(file_name)
 
 
+def _mark_job_completed(
+    task_key: str,
+    now: float,
+    status_check_dict: dict | None,
+) -> bool:
+    if status_check_dict is None:
+        return True
+    status_check_dict[task_key] = {"last_checked": now, "completed_at": now}
+    return False
+
+
+def _validate_slurm_job_output(
+    task_key: str,
+    file_name: str,
+    queue_id: int,
+    now: float,
+    status_check_dict: dict | None,
+) -> bool:
+    state = pysqa_get_slurm_job_state(queue_id=queue_id)
+    if state == "COMPLETED":
+        return _mark_job_completed(task_key, now, status_check_dict)
+    return state not in _SLURM_RUNNING_STATES and not os.path.exists(file_name)
+
+
+def pysqa_get_slurm_job_state(queue_id: int) -> str | None:
+    """
+    Query Slurm accounting for a job that is no longer in the active queue.
+
+    Returns the first allocation state reported by sacct, or None when no state is available.
+    """
+    output = pysqa_execute_command(
+        [
+            "sacct",
+            "--noheader",
+            "--allocations",
+            "--jobs",
+            str(queue_id),
+            "--format=State",
+        ]
+    )
+    lines = output.splitlines() if isinstance(output, str) else output
+    for line in lines:
+        state = line.partition("|")[0].strip().upper()
+        if state:
+            return state
+    return None
+
+
 def pysqa_terminate(
     queue_id: int,
     config_directory: str | None = None,
@@ -112,7 +188,7 @@ def pysqa_get_status_of_job(
 
 
 def pysqa_execute_command(
-    commands: str,
+    commands: str | list[str],
     working_directory: str | None = None,
     split_output: bool = True,
     shell: bool = False,

tests/unit/task_scheduler/file/test_backend.py

@@ -246,7 +246,10 @@ def test_check_task_output_dead_job_without_output(self):
         with patch(
             "executorlib.standalone.command_pysqa.pysqa_get_status_of_job",
             return_value=None,
-        ) as status_mock:
+        ) as status_mock, patch(
+            "executorlib.standalone.command_pysqa.pysqa_get_slurm_job_state",
+            return_value=None,
+        ):
             _check_task_output(
                 task_key=task_key,
                 future_obj=future_obj,
@@ -289,6 +292,227 @@ def test_check_task_output_dead_job_reported_as_error_status(self):
         with self.assertRaises(RuntimeError):
             future_obj.result()
 
+    @unittest.skipIf(sys.platform == "win32" or skip_pysqa_test, "pysqa module patching not supported on Windows or when pysqa is not installed")
+    def test_check_task_output_completed_slurm_job_waits_for_output(self):
+        cache_directory = os.path.abspath("executorlib_cache")
+        os.makedirs(cache_directory, exist_ok=True)
+        task_key, _, _ = serialize_funct(
+            fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}
+        )
+        future_obj = Future()
+        status_check_dict = {}
+        with (
+            patch(
+                "executorlib.standalone.command_pysqa.pysqa_get_status_of_job",
+                return_value=None,
+            ),
+            patch(
+                "executorlib.standalone.command_pysqa.pysqa_get_slurm_job_state",
+                return_value="COMPLETED",
+            ),
+            patch(
+                "executorlib.standalone.command_pysqa.monotonic",
+                side_effect=[100.0, 101.0],
+            ),
+        ):
+            _check_task_output(
+                task_key=task_key,
+                future_obj=future_obj,
+                cache_directory=cache_directory,
+                queue_id=123,
+                backend="slurm",
+                validate_function=pysqa_job_output_validation,
+                status_check_dict=status_check_dict,
+            )
+            self.assertFalse(future_obj.done())
+            dump(
+                file_name=os.path.join(cache_directory, task_key + "_o.h5"),
+                data_dict={"output": 3},
+            )
+            _check_task_output(
+                task_key=task_key,
+                future_obj=future_obj,
+                cache_directory=cache_directory,
+                queue_id=123,
+                backend="slurm",
+                validate_function=pysqa_job_output_validation,
+                status_check_dict=status_check_dict,
+            )
+        self.assertTrue(future_obj.done())
+        self.assertEqual(future_obj.result(), 3)
+
+    @unittest.skipIf(sys.platform == "win32" or skip_pysqa_test, "pysqa module patching not supported on Windows or when pysqa is not installed")
+    def test_check_task_output_completed_flux_job_waits_for_output(self):
+        cache_directory = os.path.abspath("executorlib_cache")
+        os.makedirs(cache_directory, exist_ok=True)
+        task_key, _, _ = serialize_funct(
+            fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}
+        )
+        future_obj = Future()
+        status_check_dict = {}
+        with (
+            patch(
+                "executorlib.standalone.command_pysqa.pysqa_get_status_of_job",
+                return_value="finished",
+            ),
+            patch(
+                "executorlib.standalone.command_pysqa.monotonic",
+                side_effect=[100.0],
+            ),
+        ):
+            _check_task_output(
+                task_key=task_key,
+                future_obj=future_obj,
+                cache_directory=cache_directory,
+                queue_id=123,
+                backend="flux",
+                validate_function=pysqa_job_output_validation,
+                status_check_dict=status_check_dict,
+            )
+            self.assertFalse(future_obj.done())
+            dump(
+                file_name=os.path.join(cache_directory, task_key + "_o.h5"),
+                data_dict={"output": 3},
+            )
+            _check_task_output(
+                task_key=task_key,
+                future_obj=future_obj,
+                cache_directory=cache_directory,
+                queue_id=123,
+                backend="flux",
+                validate_function=pysqa_job_output_validation,
+                status_check_dict=status_check_dict,
+            )
+        self.assertTrue(future_obj.done())
+        self.assertEqual(future_obj.result(), 3)
+
+    @unittest.skipIf(sys.platform == "win32" or skip_pysqa_test, "pysqa module patching not supported on Windows or when pysqa is not installed")
+    def test_check_task_output_failed_slurm_job_does_not_wait(self):
+        cache_directory = os.path.abspath("executorlib_cache")
+        os.makedirs(cache_directory, exist_ok=True)
+        task_key, _, _ = serialize_funct(
+            fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}
+        )
+        future_obj = Future()
+        with (
+            patch(
+                "executorlib.standalone.command_pysqa.pysqa_get_status_of_job",
+                return_value=None,
+            ),
+            patch(
+                "executorlib.standalone.command_pysqa.pysqa_get_slurm_job_state",
+                return_value="TIMEOUT",
+            ),
+        ):
+            _check_task_output(
+                task_key=task_key,
+                future_obj=future_obj,
+                cache_directory=cache_directory,
+                queue_id=123,
+                backend="slurm",
+                validate_function=pysqa_job_output_validation,
+                status_check_dict={},
+            )
+        self.assertTrue(future_obj.done())
+        with self.assertRaises(RuntimeError):
+            future_obj.result()
+
+    @unittest.skipIf(sys.platform == "win32" or skip_pysqa_test, "pysqa module patching not supported on Windows or when pysqa is not installed")
+    def test_check_task_output_completed_slurm_job_wait_is_bounded(self):
+        cache_directory = os.path.abspath("executorlib_cache")
+        os.makedirs(cache_directory, exist_ok=True)
+        task_key, _, _ = serialize_funct(
+            fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}
+        )
+        future_obj = Future()
+        status_check_dict = {}
+        with (
+            patch(
+                "executorlib.standalone.command_pysqa.pysqa_get_status_of_job",
+                return_value=None,
+            ),
+            patch(
+                "executorlib.standalone.command_pysqa.pysqa_get_slurm_job_state",
+                return_value="COMPLETED",
+            ),
+            patch(
+                "executorlib.standalone.command_pysqa.monotonic",
+                side_effect=[100.0, 160.0],
+            ),
+            patch(
+                "executorlib.standalone.command_pysqa._JOB_OUTPUT_GRACE_PERIOD",
+                60.0,
+            ),
+        ):
+            _check_task_output(
+                task_key=task_key,
+                future_obj=future_obj,
+                cache_directory=cache_directory,
+                queue_id=123,
+                backend="slurm",
+                validate_function=pysqa_job_output_validation,
+                status_check_dict=status_check_dict,
+            )
+            self.assertFalse(future_obj.done())
+            _check_task_output(
+                task_key=task_key,
+                future_obj=future_obj,
+                cache_directory=cache_directory,
+                queue_id=123,
+                backend="slurm",
+                validate_function=pysqa_job_output_validation,
+                status_check_dict=status_check_dict,
+            )
+        self.assertTrue(future_obj.done())
+        with self.assertRaises(RuntimeError):
+            future_obj.result()
+
+    @unittest.skipIf(sys.platform == "win32" or skip_pysqa_test, "pysqa module patching not supported on Windows or when pysqa is not installed")
+    def test_check_task_output_completed_flux_job_wait_is_bounded(self):
+        cache_directory = os.path.abspath("executorlib_cache")
+        os.makedirs(cache_directory, exist_ok=True)
+        task_key, _, _ = serialize_funct(
+            fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}
+        )
+        future_obj = Future()
+        status_check_dict = {}
+        with (
+            patch(
+                "executorlib.standalone.command_pysqa.pysqa_get_status_of_job",
+                return_value="finished",
+            ),
+            patch(
+                "executorlib.standalone.command_pysqa.monotonic",
+                side_effect=[100.0, 160.0],
+            ),
+            patch(
+                "executorlib.standalone.command_pysqa._JOB_OUTPUT_GRACE_PERIOD",
+                60.0,
+            ),
+        ):
+            _check_task_output(
+                task_key=task_key,
+                future_obj=future_obj,
+                cache_directory=cache_directory,
+                queue_id=123,
+                backend="flux",
+                validate_function=pysqa_job_output_validation,
+                status_check_dict=status_check_dict,
+            )
+            self.assertFalse(future_obj.done())
+            _check_task_output(
+                task_key=task_key,
+                future_obj=future_obj,
+                cache_directory=cache_directory,
+                queue_id=123,
+                backend="flux",
+                validate_function=pysqa_job_output_validation,
+                status_check_dict=status_check_dict,
+            )
+        self.assertTrue(future_obj.done())
+        with self.assertRaises(RuntimeError):
+            future_obj.result()
+
     @unittest.skipIf(sys.platform == "win32" or skip_pysqa_test, "pysqa module patching not supported on Windows or when pysqa is not installed")
     def test_check_task_output_job_still_running(self):
         cache_directory = os.path.abspath("executorlib_cache")