fix: address Slurm job monitoring and subscription review findings

.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

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.

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"

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)

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)

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)

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"
 

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