diff --git a/src/executorlib/standalone/command_pysqa.py b/src/executorlib/standalone/command_pysqa.py index 77525a54..6267d488 100644 --- a/src/executorlib/standalone/command_pysqa.py +++ b/src/executorlib/standalone/command_pysqa.py @@ -1,9 +1,67 @@ import contextlib +import os import subprocess +from time import monotonic from typing import Optional, Union from pysqa import QueueAdapter +# Minimum time between two queries of the queuing system for the status of a task whose output +# file has not appeared yet. Detecting a dead job (timeout, OOM, node failure, scancel, ...) relies +# 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 + + +def pysqa_job_output_validation( + task_key: str, + file_name: str, + queue_id: int, + status_check_dict: Optional[dict], + pysqa_config_directory: Optional[str] = None, + backend: Optional[str] = None, + job_status_check_interval: float = _JOB_STATUS_CHECK_INTERVAL, +) -> 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. + + 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 + (status None), while flux's "flux jobs -a" keeps listing inactive jobs and instead reports + pysqa's terminal-failure status "error" (the same status pysqa_terminate already treats as + not alive). + + Args: + task_key (str): The key of the task. + file_name (str): Path of the expected output HDF5 file. + queue_id (int, optional): The queuing system ID of the task. + pysqa_config_directory (str, optional): path to the pysqa config directory. + backend (str, optional): name of the backend used to spawn tasks ["slurm", "flux"]. + status_check_dict (dict): Dictionary tracking when each task's job status was last queried. + job_status_check_interval (float): Minimum time interval between job status checks for the same task. + + Returns: + bool: True if the job is no longer known to the queuing system, or is reported as having + 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 + ) + if now - last_checked < job_status_check_interval: + return False + if status_check_dict is not None: + status_check_dict[task_key] = now + status = pysqa_get_status_of_job( + queue_id=queue_id, + config_directory=pysqa_config_directory, + backend=backend, + ) + return (status is None or status == "error") and not os.path.exists(file_name) + def pysqa_terminate( queue_id: int, @@ -29,6 +87,31 @@ def pysqa_terminate( qa.delete_job(process_id=queue_id) +def pysqa_get_status_of_job( + queue_id: int, + config_directory: Optional[str] = None, + backend: Optional[str] = None, +) -> Optional[str]: + """ + Query the status of a job from the queuing system. + + Args: + queue_id (int): Queuing system ID of the job to check. + config_directory (str, optional): path to the config directory. + backend (str, optional): name of the backend used to spawn tasks ["slurm", "flux"]. + + Returns: + str: status of the job as reported by the queuing system, None if the queuing system no + longer knows about the job (e.g. it timed out, was killed or already finished). + """ + qa = QueueAdapter( + directory=config_directory, + queue_type=backend, + execute_command=pysqa_execute_command, + ) + return qa.get_status_of_job(process_id=queue_id) + + def pysqa_execute_command( commands: str, working_directory: Optional[str] = None, diff --git a/src/executorlib/task_scheduler/file/shared.py b/src/executorlib/task_scheduler/file/shared.py index 8d346b71..d6deb74b 100644 --- a/src/executorlib/task_scheduler/file/shared.py +++ b/src/executorlib/task_scheduler/file/shared.py @@ -59,6 +59,7 @@ def execute_tasks_h5( execute_function: Callable, executor_kwargs: dict, terminate_function: Optional[Callable] = None, + validate_function: Optional[Callable] = None, pysqa_config_directory: Optional[str] = None, backend: Optional[str] = None, disable_dependencies: bool = False, @@ -76,6 +77,7 @@ def execute_tasks_h5( - cwd (str/None): current working directory where the parallel python task is executed execute_function (Callable): The function to execute the tasks. terminate_function (Callable): The function to terminate the tasks. + validate_function (Callable): The function to validate the tasks. pysqa_config_directory (str, optional): path to the pysqa config directory (only for pysqa based backend). backend (str, optional): name of the backend used to spawn tasks. disable_dependencies (boolean): Disable resolving future objects during the submission. @@ -92,6 +94,7 @@ def execute_tasks_h5( cache_dir_dict: dict = {} file_name_dict: dict = {} duplicate_dict: dict = {} + status_check_dict: dict = {} while True: task_dict = None with contextlib.suppress(queue.Empty): @@ -104,7 +107,9 @@ def execute_tasks_h5( process_dict=process_dict, duplicate_dict=duplicate_dict, cache_dir_dict=cache_dir_dict, + status_check_dict=status_check_dict, terminate_function=terminate_function, + validate_function=validate_function, pysqa_config_directory=pysqa_config_directory, backend=backend, refresh_rate=refresh_rate, @@ -188,7 +193,9 @@ def execute_tasks_h5( cache_dir_dict=cache_dir_dict, process_dict=process_dict, duplicate_dict=duplicate_dict, + status_check_dict=status_check_dict, terminate_function=terminate_function, + validate_function=validate_function, pysqa_config_directory=pysqa_config_directory, backend=backend, refresh_rate=refresh_rate, @@ -199,15 +206,30 @@ def _check_task_output( task_key: str, future_obj: Future, cache_directory: str, + queue_id: Optional[int] = None, + pysqa_config_directory: Optional[str] = None, + backend: Optional[str] = None, + validate_function: Optional[Callable] = None, + status_check_dict: Optional[dict] = None, duplicate_dict: Optional[dict] = None, ) -> Future: """ Check the output of a task and set the result of the future object if available. + If the output file is missing and the task is running on a queuing system backend, this also + detects jobs which died without producing output (e.g. walltime TIMEOUT, OOM, NODE_FAIL or an + external scancel) by periodically querying the job status via pysqa and fails the future + instead of leaving it pending forever. + Args: task_key (str): The key of the task. future_obj (Future): The future object associated with the task. cache_directory (str): The directory where the HDF5 files are stored. + queue_id (int, optional): The queuing system ID of the task, if submitted via pysqa. + pysqa_config_directory (str, optional): path to the pysqa config directory. + backend (str, optional): name of the backend used to spawn tasks ["slurm", "flux"]. + status_check_dict (dict): Dictionary tracking when each task's job status was last queried, + used to throttle calls to the queuing system. duplicate_dict (dict): The dictionary mapping task keys to their associated duplicate future objects. Returns: Future: The updated future object. @@ -215,8 +237,32 @@ def _check_task_output( """ file_name = os.path.join(cache_directory, task_key + "_o.h5") if not os.path.exists(file_name): - return future_obj - exec_flag, no_error_flag, result = get_output(file_name=file_name) + if ( + backend is None + or queue_id is None + or validate_function is None + or not validate_function( + task_key=task_key, + file_name=file_name, + queue_id=queue_id, + pysqa_config_directory=pysqa_config_directory, + backend=backend, + status_check_dict=status_check_dict, + ) + ): + return future_obj + exec_flag, no_error_flag, result = ( + True, + False, + RuntimeError( + f"executorlib: queue job {queue_id} for task {task_key} terminated without " + "producing output (timeout, out-of-memory, node failure or cancellation)." + ), + ) + else: + exec_flag, no_error_flag, result = get_output(file_name=file_name) + if status_check_dict is not None: + status_check_dict.pop(task_key, None) _update_future( future_obj=future_obj, exec_flag=exec_flag, @@ -321,7 +367,9 @@ def _refresh_memory_dict( cache_dir_dict: dict, process_dict: dict, duplicate_dict: Optional[dict] = None, + status_check_dict: Optional[dict] = None, terminate_function: Optional[Callable] = None, + validate_function: Optional[Callable] = None, pysqa_config_directory: Optional[str] = None, backend: Optional[str] = None, refresh_rate: float = 0.01, @@ -334,7 +382,10 @@ def _refresh_memory_dict( cache_dir_dict (dict): dictionary with task keys and cache directories process_dict (dict): dictionary with task keys and process reference. duplicate_dict (dict): dictionary with task keys and duplicate future objects. + status_check_dict (dict): dictionary with task keys and the last time their queuing system + job status was queried, used to throttle detection of jobs that died without output. terminate_function (callable): The function to terminate the tasks. + validate_function (callable): The function to validate the tasks. pysqa_config_directory (str): path to the pysqa config directory (only for pysqa based backend). backend (str): name of the backend used to spawn tasks. refresh_rate (float): The rate at which to refresh the result. Defaults to 0.01. @@ -351,11 +402,19 @@ def _refresh_memory_dict( pysqa_config_directory=pysqa_config_directory, backend=backend, ) + if status_check_dict is not None: + for key in cancelled_lst: + status_check_dict.pop(key, None) memory_updated_dict = { key: _check_task_output( task_key=key, future_obj=value, cache_directory=cache_dir_dict[key], + queue_id=process_dict.get(key), + pysqa_config_directory=pysqa_config_directory, + backend=backend, + validate_function=validate_function, + status_check_dict=status_check_dict, duplicate_dict=duplicate_dict, ) for key, value in memory_dict.items() @@ -443,7 +502,9 @@ def _shutdown_executor( process_dict: dict, cache_dir_dict: dict, duplicate_dict: Optional[dict] = None, + status_check_dict: Optional[dict] = None, terminate_function: Optional[Callable] = None, + validate_function: Optional[Callable] = None, pysqa_config_directory: Optional[str] = None, backend: Optional[str] = None, refresh_rate: float = 0.01, @@ -466,6 +527,9 @@ def _shutdown_executor( process_dict (dict): Mapping of task keys to process handles or queue IDs. duplicate_dict (dict): Mapping of task keys to lists of duplicate Future objects. cache_dir_dict (dict): Mapping of task keys to the cache directory for each task. + status_check_dict (dict): Mapping of task keys to the last time their queuing system job + status was queried, used to throttle detection of jobs that died without output. + validate_function (Callable, optional): Function used to validate the tasks. terminate_function (Callable, optional): Function used to terminate running processes. pysqa_config_directory (str, optional): Path to the pysqa config directory. backend (str, optional): Name of the backend ("slurm", "flux", or None for subprocess). @@ -478,7 +542,9 @@ def _shutdown_executor( cache_dir_dict=cache_dir_dict, process_dict=process_dict, duplicate_dict=duplicate_dict, + status_check_dict=status_check_dict, terminate_function=terminate_function, + validate_function=validate_function, pysqa_config_directory=pysqa_config_directory, backend=backend, refresh_rate=refresh_rate, @@ -493,7 +559,9 @@ def _shutdown_executor( cache_dir_dict=cache_dir_dict, process_dict=process_dict, duplicate_dict=duplicate_dict, + status_check_dict=status_check_dict, terminate_function=terminate_function, + validate_function=validate_function, pysqa_config_directory=pysqa_config_directory, backend=backend, refresh_rate=refresh_rate, @@ -512,7 +580,9 @@ def _shutdown_executor( cache_dir_dict=cache_dir_dict, process_dict=process_dict, duplicate_dict=duplicate_dict, + status_check_dict=status_check_dict, terminate_function=terminate_function, + validate_function=validate_function, pysqa_config_directory=pysqa_config_directory, backend=backend, refresh_rate=refresh_rate, diff --git a/src/executorlib/task_scheduler/file/task_scheduler.py b/src/executorlib/task_scheduler/file/task_scheduler.py index ef989ba9..95fbfce6 100644 --- a/src/executorlib/task_scheduler/file/task_scheduler.py +++ b/src/executorlib/task_scheduler/file/task_scheduler.py @@ -17,12 +17,16 @@ ) try: - from executorlib.standalone.command_pysqa import pysqa_terminate + from executorlib.standalone.command_pysqa import ( + pysqa_job_output_validation, + pysqa_terminate, + ) from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa except ImportError: # If pysqa is not available fall back to executing tasks in a subprocess execute_with_pysqa = subprocess_execute # type: ignore pysqa_terminate = None # type: ignore + pysqa_job_output_validation = None # type: ignore class FileTaskScheduler(TaskSchedulerBase): @@ -74,6 +78,7 @@ def __init__( "future_queue": self._future_queue, "execute_function": execute_function, "terminate_function": terminate_function, + "validate_function": pysqa_job_output_validation, "pysqa_config_directory": pysqa_config_directory, "backend": backend, "disable_dependencies": disable_dependencies, diff --git a/tests/unit/executor/test_flux_cluster.py b/tests/unit/executor/test_flux_cluster.py index 21b23b5e..df86a3b0 100644 --- a/tests/unit/executor/test_flux_cluster.py +++ b/tests/unit/executor/test_flux_cluster.py @@ -3,13 +3,15 @@ import unittest import shutil from time import sleep +from unittest.mock import patch -from executorlib import FluxClusterExecutor +from executorlib import FluxClusterExecutor, get_cache_data from executorlib.standalone.serialize import cloudpickle_register from executorlib.standalone.command import get_cache_execute_command try: import flux.job + from pysqa import QueueAdapter from executorlib import terminate_tasks_in_cache, terminate_task_in_cache from executorlib.standalone.hdf import dump from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa @@ -233,6 +235,50 @@ def test_executor_no_cwd(self): self.assertEqual(len(os.listdir("executorlib_cache")), 2) self.assertTrue(fs1.done()) + def test_executor_future_fails_when_job_dies_without_output(self): + # Regression test for https://github.com/pyiron/executorlib/issues/1037 : a queuing + # system job which dies before ever writing its output file (walltime TIMEOUT, OOM, + # NODE_FAIL, or an external `flux cancel`/scancel) must fail the submitting future + # instead of leaving it pending forever. This mirrors the SLURM reproducer from the + # issue, but runs it against a live flux instance so the fix is exercised end-to-end + # rather than through a mocked pysqa status query. + with patch( + "executorlib.standalone.command_pysqa._JOB_STATUS_CHECK_INTERVAL", 1.0 + ): + with FluxClusterExecutor( + resource_dict={"cores": 1, "cwd": "executorlib_cache"}, + block_allocation=False, + cache_directory="executorlib_cache", + pmi_mode=pmi, + ) as exe: + cloudpickle_register(ind=1) + future = exe.submit(long_running_function, 1) + + queue_id = None + for _ in range(200): + for entry in get_cache_data(cache_directory="executorlib_cache"): + if entry.get("queue_id") is not None: + queue_id = entry["queue_id"] + if queue_id is not None: + break + sleep(0.1) + self.assertIsNotNone( + queue_id, msg="task was never submitted to the flux queue" + ) + + # Kill the job the same way a scheduler would on a walltime TIMEOUT, OOM, or + # NODE_FAIL, or how an operator would with an external `flux cancel` - executorlib + # never sees this happen directly and has to notice it via a status query. + QueueAdapter(queue_type="flux").delete_job(process_id=queue_id) + + error = future.exception(timeout=30) + self.assertIsInstance( + error, + RuntimeError, + msg="the future must fail once its job is gone instead of hanging forever", + ) + self.assertIn("terminated without producing output", str(error)) + def test_pysqa_interface(self): queue_id = execute_with_pysqa( command=get_cache_execute_command( diff --git a/tests/unit/task_scheduler/file/test_backend.py b/tests/unit/task_scheduler/file/test_backend.py index bf189f7e..07fec7c1 100644 --- a/tests/unit/task_scheduler/file/test_backend.py +++ b/tests/unit/task_scheduler/file/test_backend.py @@ -1,7 +1,9 @@ from concurrent.futures import Future import os import shutil +import sys import unittest +from unittest.mock import patch from executorlib.standalone.select import FutureSelector @@ -16,6 +18,14 @@ except ImportError: skip_h5io_test = True +try: + import pysqa # noqa: F401 + from executorlib.standalone.command_pysqa import pysqa_job_output_validation + + skip_pysqa_test = False +except ImportError: + skip_pysqa_test = True + def my_funct(a, b): return a + b @@ -219,5 +229,123 @@ def test_execute_function_error(self): with self.assertRaises(ValueError): future_file_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_dead_job_without_output(self): + # Reproduces https://github.com/pyiron/executorlib/issues/1037 : a queuing system job + # which dies without ever writing its output file (e.g. walltime TIMEOUT, OOM, NODE_FAIL + # or an external scancel) must fail the future instead of leaving it pending forever. + cache_directory = os.path.abspath("executorlib_cache") + os.makedirs(cache_directory, exist_ok=True) + task_key, data_dict = 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, + ) as status_mock: + _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_mock.assert_called_once() + 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_dead_job_reported_as_error_status(self): + # On some backends (e.g. flux, whose "flux jobs -a" keeps listing inactive jobs) a + # cancelled/failed job is never removed from the queue listing - unlike slurm's squeue, + # which drops the job entirely (status None) once it is gone. Instead pysqa reports it + # with its terminal-failure status "error" (see command_pysqa.pysqa_terminate, which + # already treats "error" as not alive). That status must also fail the future - see + # https://github.com/pyiron/executorlib/issues/1037. + cache_directory = os.path.abspath("executorlib_cache") + os.makedirs(cache_directory, exist_ok=True) + task_key, data_dict = 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="error", + ) as status_mock: + _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_mock.assert_called_once() + 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") + os.makedirs(cache_directory, exist_ok=True) + task_key, data_dict = 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="running", + ) as status_mock: + _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_mock.assert_called_once() + self.assertFalse(future_obj.done()) + + @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_status_check_is_throttled(self): + cache_directory = os.path.abspath("executorlib_cache") + os.makedirs(cache_directory, exist_ok=True) + task_key, data_dict = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) + status_check_dict = {} + with patch( + "executorlib.standalone.command_pysqa.pysqa_get_status_of_job", + return_value="running", + ) as status_mock: + for _ in range(3): + _check_task_output( + task_key=task_key, + future_obj=Future(), + cache_directory=cache_directory, + queue_id=123, + backend="slurm", + validate_function=pysqa_job_output_validation, + status_check_dict=status_check_dict, + ) + status_mock.assert_called_once() + + @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_no_backend_never_queries_status(self): + # subprocess-backed tasks (backend=None) must never trigger a queuing system status check. + cache_directory = os.path.abspath("executorlib_cache") + os.makedirs(cache_directory, exist_ok=True) + task_key, data_dict = 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", + ) as status_mock: + _check_task_output( + task_key=task_key, + future_obj=future_obj, + cache_directory=cache_directory, + queue_id=123, + backend=None, + ) + status_mock.assert_not_called() + self.assertFalse(future_obj.done()) + def tearDown(self): shutil.rmtree("executorlib_cache", ignore_errors=True)