From afcf5fbe65c5293ed60a129cc5e4eba2eb04d492 Mon Sep 17 00:00:00 2001 From: wakonig_k Date: Thu, 22 Oct 2026 14:40:47 +0200 Subject: [PATCH] fix: address Slurm job monitoring and subscription review findings --- .github/workflows/formatter.yml | 73 +++------- README.md | 109 +++++++-------- pyproject.toml | 36 ++--- slurm_manager/job_future.py | 24 ++-- slurm_manager/redis_sub_client.py | 35 ++++- slurm_manager/slurm_manager.py | 119 +++++++++++----- test/test_redis_pubsub_client.py | 12 +- test/test_review_regressions.py | 222 ++++++++++++++++++++++++++++++ 8 files changed, 440 insertions(+), 190 deletions(-) create mode 100644 test/test_review_regressions.py diff --git a/.github/workflows/formatter.yml b/.github/workflows/formatter.yml index e977335..1ff8b06 100644 --- a/.github/workflows/formatter.yml +++ b/.github/workflows/formatter.yml @@ -1,60 +1,31 @@ -name: Formatter and Pylint jobs +name: Ruff and unit tests on: [workflow_call] jobs: - Formatter: + Ruff: runs-on: ubuntu-latest - steps: - - name: Check out repository code - uses: actions/checkout@v4 - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: "3.13" - - - name: Run black and isort - run: | - pip install black isort - pip install -e .[dev] - isort --check --diff --line-length=100 --profile=black --multi-line=3 --trailing-comma ./ - black --check --diff --color --line-length=100 --skip-magic-trailing-comma ./ - Pylint: - runs-on: ubuntu-latest - steps: - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 + - uses: actions/setup-python@v5 with: python-version: "3.13" - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install pylint pylint-exit anybadge - - - name: Run Pylint - run: | - mkdir -p ./pylint - set +e - pylint ./slurm_manager --output-format=text > ./pylint/pylint.log - pylint-exit $? - set -e - - - name: Extract Pylint Score - id: score - run: | - SCORE=$(sed -n 's/^Your code has been rated at \([-0-9.]*\)\/.*/\1/p' ./pylint/pylint.log) - echo "score=$SCORE" >> $GITHUB_OUTPUT - - - name: Create Badge - run: | - anybadge --label=Pylint --file=./pylint/pylint.svg --value="${{ steps.score.outputs.score }}" 2=red 4=orange 8=yellow 10=green - - - name: Upload Artifacts - uses: actions/upload-artifact@v4 + run: python -m pip install -e '.[dev]' + - name: Check imports and lint + run: ruff check . + - name: Check formatting + run: ruff format --check . + UnitTests: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.14"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 with: - name: pylint-artifacts - path: | - # ./pylint/pylint.log # not sure why this isn't working - ./pylint/pylint.svg + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: python -m pip install -e '.[dev]' + - name: Run unit tests + run: python -m pytest --random-order -q test diff --git a/README.md b/README.md index 05bdf04..47a172f 100644 --- a/README.md +++ b/README.md @@ -1,69 +1,54 @@ # slurm_manager -The slurm manager is a tool for managing slurm jobs. It provides a simple interface for submitting, monitoring, and canceling slurm jobs. It also provides a way to view the status of all slurm jobs in a cluster. +Submit Slurm jobs from Python and receive live status, logs, and heartbeats through Redis. - -## SlurmManager Manager (Python) -* submits jobs using the slurm client -* sets up the configuration for slurm jobs -* provides an overview to query the status of all slurm jobs that are currently running or have recently finished - -## SlurmManager Job Future (Python) -* wrapper around a slurm submission command -* provides an interface to query the status of a slurm job and to cancel a slurm job using a future-like interface -* provides an interface to listen to the stream of stdout and stderr of a slurm job through redis subscriptions - -## SlurmManager Job Wrapper (Bash script + ?) -* Wraps around the user's command to be run as a slurm job -* streams stdout and stderr of the job to redis -* emits events to redis when the job starts and finishes and when it fails -* Looking forward, this may also start a local server that listens to commands from the user to emit custom events to redis (bec-slurm) to stream out data. +## Setup -## Open Points -* We want to allow that user scripts can be written in a way that they emit custom events to redis (bec-slurm). -* Proper procedure to cancel/kill a slurm job, and clean up with proper event emission to redis. Alternatively, client side will implement also a monitoring of the slurm job status, and if the job vanishes from the slurm job list it will be considered as finished respectively failed if no finished event was emitted. - -### Example Code - -```python - -manager = SlurmManager(redis_host="localhost", redis_port=6379) -manager.set_default_job_config( - time="00:30:00", - partition="short", - nodes=1, - ntasks_per_node=1, - cpus_per_task=4, - mem="16G", - environments={"default": "/path/to/env/script.sh", "env2": "/path/to/env2/script.sh"}, - default_environment="default", -) -job = manager.submit_job( - job_name="my_job", - command="python my_script.py", -) -job_status = job.status() -print(f"Job status: {job_status}") -try: - job.wait(timeout=20) -except TimeoutError: - print("Job did not finish within the timeout period.") - job.cancel() - -job = manager.submit_job( - job_name="my_job", - command="python my_script.py", - environment="env2", -) - -# Listen to job events -job.listen(event_type="status", callback=lambda event: print(f"Received event: {event}")) -job.listen(event_type="stdout", callback=lambda event: print(f"Received stdout: {event}")) -job.listen(event_type="stderr", callback=lambda event: print(f"Received stderr: {event}")) - - -# Listen to all job events -manager.listen(event_type="status", callback=lambda event: print(f"Job started: {event}")) +```bash +pip install -e . ``` +The submitter needs Python and the Slurm client commands `sbatch`, `scontrol`, +`scancel`, `squeue`, and `sacct`. Compute nodes use the Bash wrapper and Redis worker; +they do not need this Python package or the submitter's Python environment. +Compute nodes need `redis-cli`, `jq`, and `stdbuf`, with the site's Redis module +available. The wrapper, Redis worker, workload, and activation script must be +accessible at the paths passed to Slurm. The Redis worker currently publishes to +`bec-slurm.psi.ch`. + +## Submit a job + +```python +from slurm_manager.slurm_manager import SlurmManager + +manager = SlurmManager(redis_host="bec-slurm.psi.ch") +try: + job = manager.submit_job( + script_path="/shared/run_task.sh", env_path="/shared/activate_workload.sh" + ) + job.listen("log", lambda message: print(message.log)) + job.wait(timeout=60) +finally: + manager.shutdown() +``` + +The activation script selects the workload's environment on the compute node. +Jobs are submitted on hold and released after Redis acknowledges the monitoring +subscriptions. Missing heartbeats trigger Slurm queries; queued jobs remain pending. + +`job.wait()` raises on failure or timeout. A timeout does not cancel the job; +use `job.cancel()` for that. `job.listen()` accepts `status`, `log`, and `heartbeat`; +remove a callback with `job.unlisten(callback_id)`. Manager shutdown disconnects the +client without cancelling remote jobs. + +## Development + +```bash +pip install -e '.[dev]' +ruff check . +ruff format --check . +python -m pytest --random-order -q test +``` + +Unit tests run on Python 3.11 and 3.14. Other CI checks use Python 3.13. diff --git a/pyproject.toml b/pyproject.toml index e1e2820..4f498cc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,17 +15,15 @@ classifiers = [ dependencies = [ "hiredis~=3.0", - "pydantic~=2.8, <2.12.0", + "pydantic~=2.12.0", "redis~=6.2", ] [project.optional-dependencies] dev = [ - "black~=26.0", - "isort~=5.13, >=5.13.2", "coverage~=7.0", "pytest~=8.0", - "pylint~=3.0", + "ruff~=0.17.1", "pytest-random-order~=1.1", "fakeredis~=2.23, >=2.23.2", ] @@ -40,27 +38,21 @@ changelog = "https://github.com/bec-project/slurm_manager/blob/main/CHANGELOG.md [tool.hatch.build.targets.wheel] include = ["*"] -[tool.isort] -profile = "black" -line_length = 100 -multi_line_output = 3 -include_trailing_comma = true - -[tool.black] +[tool.ruff] line-length = 100 +target-version = "py311" + +[tool.ruff.format] skip-magic-trailing-comma = true -[tool.pylint.basic] -# Good variable names regexes, separated by a comma. If names match any regex, -# they will always be accepted -good-names-rgxs = [ - ".*scanID.*", - ".*RID.*", - ".*pointID.*", - ".*ID.*", - ".*_2D.*", - ".*_1D.*", -] +[tool.ruff.lint] +select = ["E4", "E7", "E9", "F", "I", "PL"] + +[tool.ruff.lint.isort] +split-on-trailing-comma = false + +[tool.ruff.lint.per-file-ignores] +"test/**" = ["PL"] [tool.semantic_release] build_command = "python -m build" diff --git a/slurm_manager/job_future.py b/slurm_manager/job_future.py index 15818a5..95604a4 100644 --- a/slurm_manager/job_future.py +++ b/slurm_manager/job_future.py @@ -9,7 +9,6 @@ from __future__ import annotations import logging import threading -import time from concurrent.futures import Future from dataclasses import dataclass from enum import StrEnum @@ -109,15 +108,7 @@ class JobFuture: Args: timeout: Maximum time to wait for the job to finish, in seconds. If None, wait indefinitely. """ - self._raise_if_failed() - - start_time = time.time() - poll_interval = 0.05 # seconds - while not self.done(): - if timeout is not None and (time.time() - start_time) > timeout: - raise TimeoutError(f"{self.name} did not complete within {timeout} seconds.") - if self._shutdown_event.wait(timeout=poll_interval): - break + self._future.result(timeout=timeout) def listen( self, @@ -197,8 +188,11 @@ class JobFuture: raise self._future.exception() def _update_status(self, status: str, exc_info: str | None = None) -> None: - self._status = status - if status == JobStatus.FINISHED: - self._set_finished() - elif status in (JobStatus.CANCELLED, JobStatus.ERROR): - self._set_exception(status, exc_info) + with self._lock: + if self._future.done(): + return + self._status = status + if status == JobStatus.FINISHED: + self._set_finished() + elif status in (JobStatus.CANCELLED, JobStatus.ERROR): + self._set_exception(status, exc_info) diff --git a/slurm_manager/redis_sub_client.py b/slurm_manager/redis_sub_client.py index 6db8d06..0ba56a4 100644 --- a/slurm_manager/redis_sub_client.py +++ b/slurm_manager/redis_sub_client.py @@ -8,10 +8,11 @@ from __future__ import annotations import logging import threading +import time +import traceback from dataclasses import dataclass from typing import Any, Callable from uuid import uuid4 -import traceback from redis.client import Redis from redis.exceptions import RedisError @@ -59,12 +60,13 @@ class RedisSubClient: host, port = bootstrap.split(":", maxsplit=1) self._redis = redis_cls(host=host, port=int(port), **redis_kwargs) - self._pubsub = self._redis.pubsub(ignore_subscribe_messages=True) + self._pubsub = self._redis.pubsub(ignore_subscribe_messages=False) self._lock = threading.RLock() self._subscriptions_by_topic: dict[str, dict[str, Subscription]] = {} self._topic_by_subscription_id: dict[str, str] = {} + self._subscription_ready: dict[str, threading.Event] = {} self._stop_listener = threading.Event() self._listener_thread: threading.Thread | None = None @@ -82,8 +84,9 @@ class RedisSubClient: with self._lock: callbacks = self._subscriptions_by_topic.get(topic) if callbacks is None: - self._subscriptions_by_topic[topic] = {sub_id: subscription} + self._subscription_ready[topic] = threading.Event() self._pubsub.subscribe(topic) + self._subscriptions_by_topic[topic] = {sub_id: subscription} else: callbacks[sub_id] = subscription self._topic_by_subscription_id[sub_id] = topic @@ -112,6 +115,7 @@ class RedisSubClient: # No more callbacks for this topic, we can unsubscribe from Redis and clean up self._subscriptions_by_topic.pop(topic, None) + self._subscription_ready.pop(topic, None) self._pubsub.unsubscribe(topic) return True @@ -121,11 +125,21 @@ class RedisSubClient: callbacks = self._subscriptions_by_topic.pop(topic, None) if not callbacks: # No callbacks for this topic, nothing happens return 0 + self._subscription_ready.pop(topic, None) for sub_id in callbacks: # Clean up subscription id mapping self._topic_by_subscription_id.pop(sub_id, None) self._pubsub.unsubscribe(topic) # Unsubscribe from Redis return len(callbacks) + def wait_until_subscribed(self, topics: list[str], timeout: float = 5.0) -> None: + """Wait for Redis to acknowledge each subscription before releasing a job.""" + deadline = time.monotonic() + timeout + with self._lock: + ready = [(topic, self._subscription_ready[topic]) for topic in topics] + for topic, event in ready: + if not event.wait(max(0, deadline - time.monotonic())): + raise TimeoutError(f"Redis did not acknowledge subscription to {topic}") + def shutdown(self, timeout: float | None = 1.0) -> None: """Stops worker threads and closes Redis resources.""" self._stop_listener.set() @@ -138,6 +152,7 @@ class RedisSubClient: with self._lock: self._subscriptions_by_topic.clear() self._topic_by_subscription_id.clear() + self._subscription_ready.clear() self._pubsub.close() self._redis.close() @@ -155,7 +170,9 @@ class RedisSubClient: """Listener thread loop that polls for messages and dispatches them to callbacks.""" while not self._stop_listener.is_set(): try: - raw_message = self._pubsub.get_message(timeout=0.2) + with self._lock: + raw_message = self._pubsub.get_message(timeout=0.05) + self._stop_listener.wait(0.001) except RedisError: logger.error("Error while polling redis pub/sub") self._stop_listener.wait(0.5) @@ -163,6 +180,16 @@ class RedisSubClient: if raw_message is None: continue + if raw_message["type"] == "subscribe": + channel = raw_message["channel"] + topic = channel.decode() if isinstance(channel, bytes) else channel + with self._lock: + event = self._subscription_ready.get(topic) + if event is not None: + event.set() + continue + if raw_message["type"] == "unsubscribe": + continue try: topic, message = self._normalize_message(raw_message) diff --git a/slurm_manager/slurm_manager.py b/slurm_manager/slurm_manager.py index 7bd4b83..1d3746d 100644 --- a/slurm_manager/slurm_manager.py +++ b/slurm_manager/slurm_manager.py @@ -67,6 +67,7 @@ class SubscriptionInfo(BaseModel): callback_id: str | None = None last_heartbeat: float | None = None heartbeat_received: bool = False + last_scheduler_check: float | None = None HEARTBEAT_TIMEOUT = 10 # seconds @@ -120,37 +121,68 @@ class SlurmManager: infos = self._subscriptions_by_topic.get(topic, {}) for sub_info in infos.values(): if sub_info and sub_info.last_heartbeat: - if time.time() - sub_info.last_heartbeat > HEARTBEAT_TIMEOUT: + now = time.time() + if now - sub_info.last_heartbeat > HEARTBEAT_TIMEOUT and ( + sub_info.last_scheduler_check is None + or now - sub_info.last_scheduler_check > HEARTBEAT_TIMEOUT + ): missing_heartbeats.append(sub_info) for sub_info in missing_heartbeats: self._check_slurm_job(sub_info) def _check_slurm_job(self, sub_info: SubscriptionInfo) -> None: - """Check if a SLURM job is active""" - # TODO check if job is Pending through SLURM Rest API.. - # Currently we just trigger a cleanup if heartbeat is missing. + """Resolve missing heartbeats using Slurm, never assuming silence means failure.""" future = self._future_registry.get(sub_info.job_id) - if not future: + if future is None or future.done(): return - cancel_msg = f"error:No heartbeat received within {HEARTBEAT_TIMEOUT}s for topic {sub_info.topic} and slurm JobID {sub_info.slurm_job_id}." - if sub_info.heartbeat_received is True: - cancel_msg += ( - " Last heartbeat from job at " - f"{datetime.fromtimestamp(sub_info.last_heartbeat).isoformat()}." + sub_info.last_scheduler_check = time.time() + try: + queued = self._run_subprocess( + ["squeue", "--noheader", "--jobs", future.slurm_job_id, "--format=%T"] ) - - self._event_callback( - StatusMessage( - status=cancel_msg, - metadata={ - "msg_received": datetime.now().isoformat(), - "job_id": future._job_id, - "slurm_job_id": future.slurm_job_id, - }, - ), - future=future, - ) + if queued.returncode == 0 and queued.stdout.strip(): + # Pending, configuring, and running jobs may legitimately have no heartbeat. + return + accounting = self._run_subprocess( + [ + "sacct", + "--noheader", + "--parsable2", + "--allocations", + "--jobs", + future.slurm_job_id, + "--format=State,ExitCode", + ] + ) + except (OSError, subprocess.SubprocessError): + logger.exception("Could not query Slurm job %s", future.slurm_job_id) + return + if accounting.returncode != 0 or not accounting.stdout.strip(): + # Accounting may be disabled, delayed, or temporarily unavailable. + return + fields = accounting.stdout.strip().splitlines()[0].split("|") + state = fields[0].split()[0].rstrip("+") + exit_code = fields[1].strip() if len(fields) > 1 else "" + if state == "COMPLETED" and exit_code == "0:0": + status = JobStatus.FINISHED.value + elif state == "CANCELLED": + status = JobStatus.CANCELLED.value + elif state in { + "COMPLETED", + "FAILED", + "TIMEOUT", + "OUT_OF_MEMORY", + "NODE_FAIL", + "BOOT_FAIL", + "DEADLINE", + "PREEMPTED", + "REVOKED", + }: + status = f"error:Slurm state {state}, exit code {exit_code or 'unknown'}" + else: + return + self._event_callback(StatusMessage(status=status), future) def submit_job(self, script_path: str, env_path: str) -> JobFuture: """Create a temporary script from command and submit via sbatch wrapper.""" @@ -167,6 +199,7 @@ class SlurmManager: sbatch_cmd = [ "sbatch", "--parsable", + "--hold", os.path.abspath(self._wrapper_path), env_path, script_path, @@ -188,8 +221,26 @@ class SlurmManager: job_id=job_id, manager=self, registry=self._future_registry, slurm_job_id=slurm_job_id ) topic_info = TopicInfo(job_id=job_id) - self._future_registry[job_id] = job_future - self._subscribe_to_internal_topics(future=job_future, topic_info=topic_info) + with self._lock: + self._future_registry[job_id] = job_future + try: + self._subscribe_to_internal_topics(future=job_future, topic_info=topic_info) + self.sub_client.wait_until_subscribed([topic_info.heartbeat, topic_info.event]) + released = self._run_subprocess(["scontrol", "release", slurm_job_id]) + if released.returncode != 0: + raise RuntimeError(f"scontrol release failed: {(released.stderr or '').strip()}") + except Exception as exc: + # A setup failure must not leave a held allocation behind. + try: + cancelled = self._run_subprocess(["scancel", slurm_job_id]) + if cancelled.returncode != 0: + logger.error("Could not cancel held job %s: %s", slurm_job_id, cancelled.stderr) + except (OSError, subprocess.SubprocessError): + logger.exception("Could not cancel held job %s", slurm_job_id) + self._event_callback( + StatusMessage(status=f"error:Submission setup failed: {exc}"), job_future + ) + raise return job_future def cancel_job(self, slurm_job_id: str, job_id: str) -> None: @@ -255,14 +306,15 @@ class SlurmManager: if topic is not None: callbacks = self._subscriptions_by_topic.get(topic, {}) sub_info = callbacks.pop(cb_id, None) - if sub_info: - job_id = sub_info.job_id - self._active_subscriptions[job_id].discard(topic) + self.sub_client.unsubscribe(cb_id) - # If no more callbacks for the topic, remove the topic subscription - if len(callbacks) == 0: + # Other callbacks, including heartbeat monitoring, still own this topic. + if not callbacks: self._subscriptions_by_topic.pop(topic, None) - self.sub_client.unsubscribe_topic(topic) + if sub_info: + topics = self._active_subscriptions.get(sub_info.job_id) + if topics is not None: + topics.discard(topic) def remove_subscription_by_topic(self, topic: str) -> None: """Remove a subscription by its topic.""" @@ -318,6 +370,11 @@ class SlurmManager: def shutdown(self, timeout: float | None = 1.0) -> None: """Shuts down the manager, including the heartbeat monitoring thread and Redis subscriptions.""" self._stop_heartbeat_thread.set() + self._heartbeat_monitoring_thread.join(timeout) + if self._heartbeat_monitoring_thread.is_alive(): + raise TimeoutError( + "Heartbeat monitor did not stop; retry after its Slurm query finishes" + ) with self._lock: self._subscriptions_by_topic.clear() self._topic_by_subscription_id.clear() @@ -404,7 +461,7 @@ class SlurmManager: self._future_registry.pop(job_id, None) def _run_subprocess(self, cmd: list[str]) -> subprocess.CompletedProcess: - return subprocess.run(cmd, capture_output=True, text=True, check=False) + return subprocess.run(cmd, capture_output=True, text=True, check=False, timeout=30) def _subscribe_to_internal_topics(self, future: JobFuture, topic_info: TopicInfo) -> None: self._subscribe_to_heartbeat(topic_info, future) diff --git a/test/test_redis_pubsub_client.py b/test/test_redis_pubsub_client.py index f6440a6..3566683 100644 --- a/test/test_redis_pubsub_client.py +++ b/test/test_redis_pubsub_client.py @@ -16,7 +16,9 @@ def client(): """FakeRedis-based RedisServer fixture for testing the RedisSubClient.""" server = fakeredis.FakeServer() redis_cls = partial(fakeredis.FakeRedis, server=server, decode_responses=True) - return RedisSubClient(bootstrap="localhost:6379", redis_cls=redis_cls) + client = RedisSubClient(bootstrap="localhost:6379", redis_cls=redis_cls) + yield client + client.shutdown() def test_redis_pubsub_listener(client): @@ -55,7 +57,7 @@ def test_redis_pubsub_unsubscribe_topic(client): client.subscribe("a", partial(_callback, text="cb2")) assert client.unsubscribe_topic("a") == 2 - client._redis.publish("job/1/status", '{"msg_type":"status", "status": "running"}') + client._redis.publish("a", '{"msg_type":"status", "status": "running"}') assert not received_event.wait(timeout=1) assert not got, "Callbacks should not be called after unsubscribing from topic" client.shutdown() @@ -69,9 +71,9 @@ def test_redis_pubsub_shutdown(client): ): client._ensure_threads_started() client.shutdown() - assert ( - not client._listener_thread.is_alive() - ), "Listener thread should be stopped after shutdown" + assert not client._listener_thread.is_alive(), ( + "Listener thread should be stopped after shutdown" + ) assert mock_close.called, "PubSub close method should be called on shutdown" assert mock_redis_close.called, "Redis close method should be called on shutdown" diff --git a/test/test_review_regressions.py b/test/test_review_regressions.py new file mode 100644 index 0000000..50585d3 --- /dev/null +++ b/test/test_review_regressions.py @@ -0,0 +1,222 @@ +"""Regressions for job submission, waiting, and subscription ownership.""" + +from __future__ import annotations + +import subprocess +import threading +import time +from functools import partial +from unittest import mock + +import fakeredis +import pytest + +from slurm_manager.event_messages import HeartBeatMessage +from slurm_manager.job_future import JobFuture +from slurm_manager.redis_sub_client import RedisSubClient +from slurm_manager.slurm_manager import SlurmManager, TopicInfo + + +@pytest.fixture +def manager(): + client = RedisSubClient(redis_cls=partial(fakeredis.FakeRedis, decode_responses=True)) + with mock.patch("slurm_manager.slurm_manager.RedisSubClient", return_value=client): + manager = SlurmManager() + + def slurm(command): + output = ( + "123\n" if command[0] == "sbatch" else "PENDING\n" if command[0] == "squeue" else "" + ) + return subprocess.CompletedProcess(command, 0, output, "") + + manager._run_subprocess = mock.Mock(side_effect=slurm) + try: + yield manager + finally: + manager.shutdown(timeout=3) + assert not manager._heartbeat_monitoring_thread.is_alive() + listener = client._listener_thread + assert listener is None or not listener.is_alive() + + +def submit(manager): + # The workload and activation environment belong to the remote node. They + # deliberately do not exist in the submitter's Python environment/filesystem. + return manager.submit_job("/remote/node/workload.sh", "/remote/node/activate.sh") + + +def heartbeat_info(manager, future): + return next(iter(manager._subscriptions_by_topic[TopicInfo(future._job_id).heartbeat].values())) + + +def test_wait_propagates_failure_arriving_during_wait(): + future = JobFuture("job", manager=None, slurm_job_id="123") + timer = threading.Timer(0.05, lambda: future._update_status("error", "exit code 7")) + timer.start() + try: + with pytest.raises(Exception, match="exit code 7"): + future.wait(timeout=1) + finally: + timer.join() + + +def test_wait_timeout_does_not_complete_future(): + future = JobFuture("job", manager=None, slurm_job_id="123") + with pytest.raises(TimeoutError): + future.wait(timeout=0) + assert not future.done() + future._update_status("finished") + future.wait(timeout=0) + assert future.success() + + +def test_terminal_state_ignores_late_status_updates(): + future = JobFuture("job", manager=None, slurm_job_id="123") + future._update_status("finished") + future._update_status("running") + assert future.status == "finished" + assert future.success() + + +def test_job_is_held_until_redis_acknowledges_subscriptions(manager): + commands = [] + + def slurm(command): + commands.append(command) + if command[0] == "sbatch": + assert "--hold" in command + assert "/remote/node/activate.sh" in command + assert "/remote/node/workload.sh" in command + assert command[-1].endswith("redis_worker.sh") + return subprocess.CompletedProcess(command, 0, "123\n", "") + assert command == ["scontrol", "release", "123"] + future = next(iter(manager._future_registry.values())) + topics = TopicInfo(future._job_id) + assert manager.sub_client._subscription_ready[topics.event].is_set() + assert manager.sub_client._subscription_ready[topics.heartbeat].is_set() + # Simulate a remote job completing before scontrol returns. + manager.sub_client._redis.publish(topics.event, '{"msg_type":"status","status":"finished"}') + future.wait(timeout=1) + return subprocess.CompletedProcess(command, 0, "", "") + + manager._run_subprocess.side_effect = slurm + future = submit(manager) + assert future.success() + assert [command[0] for command in commands] == ["sbatch", "scontrol"] + assert not manager.sub_client.registered_topics + + +@pytest.mark.parametrize("failure", ["subscription", "acknowledgment", "release"]) +def test_setup_failure_cancels_held_job_and_cleans_subscriptions(manager, failure): + original = manager._run_subprocess.side_effect + + def slurm(command): + if command[0] == "scontrol": + return subprocess.CompletedProcess(command, 1, "", "release rejected") + return original(command) + + manager._run_subprocess.side_effect = slurm + if failure == "subscription": + patch = mock.patch.object( + manager, "_subscribe_to_event", side_effect=RuntimeError("offline") + ) + elif failure == "acknowledgment": + patch = mock.patch.object( + manager.sub_client, "wait_until_subscribed", side_effect=TimeoutError("no ack") + ) + else: + patch = mock.patch.object( + manager.sub_client, + "wait_until_subscribed", + wraps=manager.sub_client.wait_until_subscribed, + ) + with patch, pytest.raises((RuntimeError, TimeoutError)): + submit(manager) + assert mock.call(["scancel", "123"]) in manager._run_subprocess.call_args_list + assert not manager._future_registry + assert not manager._active_subscriptions + assert not manager.sub_client.registered_topics + + +@pytest.mark.parametrize("state", ["PENDING", "CONFIGURING", "RUNNING"]) +def test_missing_heartbeat_does_not_fail_active_slurm_job(manager, state): + future = submit(manager) + info = heartbeat_info(manager, future) + info.last_heartbeat = time.time() - 100 + manager._run_subprocess.return_value = subprocess.CompletedProcess([], 0, f"{state}\n", "") + manager._run_subprocess.side_effect = None + manager._check_slurm_job(info) + assert not future.done() + assert future._job_id in manager._future_registry + assert TopicInfo(future._job_id).heartbeat in manager._active_subscriptions[future._job_id] + assert manager._run_subprocess.call_args.args[0][0] == "squeue" + + +@pytest.mark.parametrize("accounting", ["", "PENDING|0:0\n"]) +def test_unknown_or_delayed_accounting_is_not_failure(manager, accounting): + future = submit(manager) + manager._run_subprocess.side_effect = [ + subprocess.CompletedProcess([], 0, "", ""), + subprocess.CompletedProcess([], 0, accounting, ""), + ] + manager._check_slurm_job(heartbeat_info(manager, future)) + assert not future.done() + + +@pytest.mark.parametrize( + "accounting,status,success", + [ + ("COMPLETED|0:0\n", "finished", True), + ("FAILED|7:0\n", "error", False), + ("CANCELLED by 1000|0:15\n", "cancelled", False), + ], +) +def test_terminal_accounting_resolves_missed_remote_status(manager, accounting, status, success): + future = submit(manager) + manager._run_subprocess.side_effect = [ + subprocess.CompletedProcess([], 0, "", ""), + subprocess.CompletedProcess([], 0, accounting, ""), + ] + manager._check_slurm_job(heartbeat_info(manager, future)) + assert future.done() + assert future.status == status + assert future.success() == success + assert not manager._future_registry + assert not manager.sub_client.registered_topics + + +def test_scheduler_error_keeps_job_unresolved(manager): + future = submit(manager) + manager._run_subprocess.side_effect = FileNotFoundError("squeue unavailable") + manager._check_slurm_job(heartbeat_info(manager, future)) + assert not future.done() + + +def test_unlisten_removes_only_selected_callback_and_preserves_monitoring(manager): + future = submit(manager) + removed = threading.Event() + received = threading.Event() + callback_id = future.listen("heartbeat", lambda _: removed.set()) + future.listen("heartbeat", lambda _: received.set()) + future.unlisten(callback_id) + topic = TopicInfo(future._job_id).heartbeat + assert callback_id not in manager.sub_client._topic_by_subscription_id + assert topic in manager._active_subscriptions[future._job_id] + manager.sub_client._redis.publish( + topic, HeartBeatMessage(timestamp=time.time()).model_dump_json() + ) + assert received.wait(1) + assert not removed.is_set() + assert heartbeat_info(manager, future).heartbeat_received + manager.remove_subscriptions_for_job(future._job_id) + assert not manager.sub_client.registered_topics + + +def test_removing_last_log_callback_unsubscribes_topic(manager): + future = submit(manager) + callback_id = future.listen("log", lambda _: None) + future.unlisten(callback_id) + topic = TopicInfo(future._job_id).log + assert topic not in manager._active_subscriptions[future._job_id] + assert topic not in manager.sub_client.registered_topics + assert TopicInfo(future._job_id).heartbeat in manager.sub_client.registered_topics