From 784d8b5bdbc5f7e64bb0b7693a8d4dbf4141d6c9 Mon Sep 17 00:00:00 2001 From: jan-janssen <3854739+jan-janssen@users.noreply.github.com> Date: Thu, 5 Nov 2026 13:47:23 +0000 Subject: [PATCH 1/3] Wait for shared-filesystem output of completed Slurm jobs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/executorlib/standalone/command_pysqa.py | 77 +++++++++- .../unit/task_scheduler/file/test_backend.py | 135 +++++++++++++++++- 2 files changed, 205 insertions(+), 7 deletions(-) diff --git a/src/executorlib/standalone/command_pysqa.py b/src/executorlib/standalone/command_pysqa.py index b6846f29..b4ec7ddd 100644 --- a/src/executorlib/standalone/command_pysqa.py +++ b/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 +_SLURM_OUTPUT_GRACE_PERIOD = 60.0 +_SLURM_RUNNING_STATES = { + "COMPLETING", + "CONFIGURING", + "PENDING", + "RESIZING", + "RUNNING", + "SUSPENDED", +} def pysqa_job_output_validation( @@ -23,9 +32,9 @@ 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. For Slurm, a successful job is given a grace period for its 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 +56,77 @@ 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 >= _SLURM_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, + ) return (status is None or status == "error") and not os.path.exists(file_name) +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": + if status_check_dict is None: + return True + status_check_dict[task_key] = {"last_checked": now, "completed_at": now} + return False + 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, diff --git a/tests/unit/task_scheduler/file/test_backend.py b/tests/unit/task_scheduler/file/test_backend.py index 35a62f7b..153fc424 100644 --- a/tests/unit/task_scheduler/file/test_backend.py +++ b/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,136 @@ 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_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._SLURM_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_job_still_running(self): cache_directory = os.path.abspath("executorlib_cache") From f4fe1ec8937e621a8cc5c495c0c59aa529ef1ec9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 5 Nov 2026 14:14:54 +0000 Subject: [PATCH 2/3] Fix pysqa command argument typing Co-authored-by: jan-janssen <3854739+jan-janssen@users.noreply.github.com> --- src/executorlib/standalone/command_pysqa.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/executorlib/standalone/command_pysqa.py b/src/executorlib/standalone/command_pysqa.py index b4ec7ddd..b6a1c8a9 100644 --- a/src/executorlib/standalone/command_pysqa.py +++ b/src/executorlib/standalone/command_pysqa.py @@ -177,7 +177,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, From 35167ab1b7c77f981818defe469f8857ed88fa15 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 5 Nov 2026 14:54:13 +0000 Subject: [PATCH 3/3] Extend output grace period to Flux Co-authored-by: jan-janssen <3854739+jan-janssen@users.noreply.github.com> --- src/executorlib/standalone/command_pysqa.py | 29 ++++-- .../unit/task_scheduler/file/test_backend.py | 93 ++++++++++++++++++- 2 files changed, 112 insertions(+), 10 deletions(-) diff --git a/src/executorlib/standalone/command_pysqa.py b/src/executorlib/standalone/command_pysqa.py index b6a1c8a9..12dd607b 100644 --- a/src/executorlib/standalone/command_pysqa.py +++ b/src/executorlib/standalone/command_pysqa.py @@ -10,7 +10,7 @@ # 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 -_SLURM_OUTPUT_GRACE_PERIOD = 60.0 +_JOB_OUTPUT_GRACE_PERIOD = 60.0 _SLURM_RUNNING_STATES = { "COMPLETING", "CONFIGURING", @@ -32,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. For Slurm, a successful job is given a grace period for its 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. + 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 @@ -66,7 +67,7 @@ def pysqa_job_output_validation( last_checked = check_state completed_at = None if completed_at is not None: - return now - completed_at >= _SLURM_OUTPUT_GRACE_PERIOD + 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: @@ -84,9 +85,22 @@ def pysqa_job_output_validation( 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, @@ -96,10 +110,7 @@ def _validate_slurm_job_output( ) -> bool: state = pysqa_get_slurm_job_state(queue_id=queue_id) if state == "COMPLETED": - if status_check_dict is None: - return True - status_check_dict[task_key] = {"last_checked": now, "completed_at": now} - return False + return _mark_job_completed(task_key, now, status_check_dict) return state not in _SLURM_RUNNING_STATES and not os.path.exists(file_name) diff --git a/tests/unit/task_scheduler/file/test_backend.py b/tests/unit/task_scheduler/file/test_backend.py index 153fc424..e38721f6 100644 --- a/tests/unit/task_scheduler/file/test_backend.py +++ b/tests/unit/task_scheduler/file/test_backend.py @@ -341,6 +341,51 @@ def test_check_task_output_completed_slurm_job_waits_for_output(self): 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") @@ -395,7 +440,7 @@ def test_check_task_output_completed_slurm_job_wait_is_bounded(self): side_effect=[100.0, 160.0], ), patch( - "executorlib.standalone.command_pysqa._SLURM_OUTPUT_GRACE_PERIOD", + "executorlib.standalone.command_pysqa._JOB_OUTPUT_GRACE_PERIOD", 60.0, ), ): @@ -422,6 +467,52 @@ def test_check_task_output_completed_slurm_job_wait_is_bounded(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_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")