Move pysqa_job_output_validation() function

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,

src/executorlib/task_scheduler/file/shared.py

@@ -2,20 +2,14 @@
 import os
 import queue
 from concurrent.futures import Future
-from time import monotonic, sleep
+from time import sleep
 from typing import Any, Callable, Optional
 
 from executorlib.standalone.command import get_cache_execute_command
 from executorlib.standalone.hdf import get_cache_files, get_output, get_queue_id
 from executorlib.standalone.serialize import serialize_funct
 from executorlib.task_scheduler.file.spawner_subprocess import subprocess_terminate
 
-# 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
-
 
 class FutureItem:
     def __init__(self, file_name: str, selector: Optional[int | str] = None):
@@ -65,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,
@@ -82,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.
@@ -113,6 +109,7 @@ def execute_tasks_h5(
                 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,
@@ -198,6 +195,7 @@ def execute_tasks_h5(
                 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,
@@ -211,6 +209,7 @@ def _check_task_output(
     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:
@@ -238,13 +237,18 @@ def _check_task_output(
     """
     file_name = os.path.join(cache_directory, task_key + "_o.h5")
     if not os.path.exists(file_name):
-        if not _job_died_without_output(
-            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,
+        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 = (
@@ -277,63 +281,6 @@ def _check_task_output(
     return future_obj
 
 
-def _job_died_without_output(
-    task_key: str,
-    file_name: str,
-    queue_id: Optional[int],
-    pysqa_config_directory: Optional[str],
-    backend: Optional[str],
-    status_check_dict: Optional[dict],
-) -> 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.
-
-    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.
-    """
-    if backend is None or queue_id is None:
-        return False
-    try:
-        # Imported lazily so subprocess-only (non-pysqa) task submissions - including every
-        # cache_serial.py backend subprocess spawned for local execution - never pay the cost of
-        # importing pysqa.
-        from executorlib.standalone.command_pysqa import pysqa_get_status_of_job
-    except ImportError:
-        return False
-    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 _update_future(
     future_obj: Future, exec_flag: bool, no_error_flag: bool, result: Any
 ) -> None:
@@ -422,6 +369,7 @@ def _refresh_memory_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,
@@ -437,6 +385,7 @@ def _refresh_memory_dict(
         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.
@@ -464,6 +413,7 @@ def _refresh_memory_dict(
             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,
         )
@@ -554,6 +504,7 @@ def _shutdown_executor(
     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,
@@ -578,6 +529,7 @@ def _shutdown_executor(
         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).
@@ -592,6 +544,7 @@ def _shutdown_executor(
                 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,
@@ -608,6 +561,7 @@ def _shutdown_executor(
                 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,
@@ -628,6 +582,7 @@ def _shutdown_executor(
             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,

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,

tests/unit/executor/test_flux_cluster.py

@@ -243,7 +243,7 @@ def test_executor_future_fails_when_job_dies_without_output(self):
         # 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.task_scheduler.file.shared._JOB_STATUS_CHECK_INTERVAL", 1.0
+            "executorlib.standalone.command_pysqa._JOB_STATUS_CHECK_INTERVAL", 1.0
         ):
             with FluxClusterExecutor(
                 resource_dict={"cores": 1, "cwd": "executorlib_cache"},

tests/unit/task_scheduler/file/test_backend.py

@@ -20,6 +20,7 @@
 
 try:
     import pysqa  # noqa: F401
+    from executorlib.standalone.command_pysqa import pysqa_job_output_validation
 
     skip_pysqa_test = False
 except ImportError:
@@ -247,6 +248,7 @@ def test_check_task_output_dead_job_without_output(self):
                 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())
@@ -275,6 +277,7 @@ def test_check_task_output_dead_job_reported_as_error_status(self):
                 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())
@@ -297,6 +300,7 @@ def test_check_task_output_job_still_running(self):
                 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())
@@ -318,6 +322,7 @@ def test_check_task_output_status_check_is_throttled(self):
                     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()