Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 83 additions & 0 deletions src/executorlib/standalone/command_pysqa.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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,
Expand Down
74 changes: 72 additions & 2 deletions src/executorlib/task_scheduler/file/shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.
Expand All @@ -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):
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -199,24 +206,63 @@ 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.

"""
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,
Expand Down Expand Up @@ -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,
Expand All @@ -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.
Expand All @@ -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()
Expand Down Expand Up @@ -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,
Expand All @@ -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).
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
7 changes: 6 additions & 1 deletion src/executorlib/task_scheduler/file/task_scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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,
Expand Down
48 changes: 47 additions & 1 deletion tests/unit/executor/test_flux_cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Comment on lines +248 to +264

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Use a test-specific cache directory.

get_cache_data() scans all entries in executorlib_cache. If another entry has a queue_id, this test can cancel that unrelated job and leave future running. Create a unique cache directory for this test and use it for both FluxClusterExecutor and get_cache_data().

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/executor/test_flux_cluster.py` around lines 248 - 264, Use a
unique test-specific cache directory in the FluxClusterExecutor test, and pass
that same directory to both the executor’s cache_directory parameter and
get_cache_data(). Ensure queue_id discovery only scans entries created by this
test, avoiding interference from unrelated jobs.

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(
Expand Down
Loading
Loading