fix: route clustered jobs and serialize callback registration

slurm_manager/job_future.py

@@ -57,12 +57,13 @@ class JobFuture:
     ):
         self._job_id = job_id
         self._manager = manager
-        self._slurm_job_id = slurm_job_id
+        self._slurm_job_id, _, cluster = slurm_job_id.strip().partition(";")
+        self._slurm_cluster = cluster or None
         self._registry = registry if registry is not None else {}
         self._future = Future()
         if name is None:
             name = "Job"
-        self.name = f"{name}-{slurm_job_id}"
+        self.name = f"{name}-{self._slurm_job_id}"
         self._status: str = JobStatus.PENDING.value
         self._shutdown_event = threading.Event()
         self._lock = threading.RLock()
@@ -74,9 +75,14 @@ class JobFuture:
 
     @property
     def slurm_job_id(self) -> str:
-        """The SLURM ID of the job, as returned by sbatch."""
+        """The numeric SLURM job ID, without an optional cluster suffix."""
         return self._slurm_job_id
 
+    @property
+    def slurm_cluster(self) -> str | None:
+        """The cluster returned by sbatch, or None for the default cluster."""
+        return self._slurm_cluster
+
     @property
     def status(self) -> str | None:
         """The current status of the job."""
@@ -84,7 +90,9 @@ class JobFuture:
 
     def cancel(self):
         """Cancel the job."""
-        self._manager.cancel_job(slurm_job_id=self.slurm_job_id, job_id=self._job_id)
+        self._manager.cancel_job(
+            slurm_job_id=self.slurm_job_id, job_id=self._job_id, slurm_cluster=self.slurm_cluster
+        )
 
     def exception(self) -> Exception | None:
         """If the job has finished with an error, return the exception. Otherwise, return None."""
@@ -146,15 +154,22 @@ class JobFuture:
         Args:
             callback_id: The unique ID of the callback to unregister, as returned by listen()
         """
+        removed = False
         with self._lock:
             for event, callbacks in self._callback_registry.items():
                 for cb in callbacks:
                     if cb.id == callback_id:
                         callbacks.remove(cb)
-                        self._manager.remove_subscription_by_id(callback_id=callback_id)
-                        if not callbacks:  # if no more callbacks for event, remove the event
+                        if not callbacks:
                             self._callback_registry.pop(event, None)
-                        return
+                        removed = True
+                        break
+                if removed:
+                    break
+        # Manager cleanup may acquire this future's lock while holding its own.
+        # Release the future lock before acquiring the manager lock.
+        if removed:
+            self._manager.remove_subscription_by_id(callback_id=callback_id)
 
     ###################
     ## Internal Methods

slurm_manager/slurm_manager.py

@@ -28,6 +28,11 @@ logger = getLogger(__name__)
 _STOP = object()
 
 
+def _cluster_args(cluster: str | None) -> list[str]:
+    """Route Slurm CLI operations to the cluster that accepted the job."""
+    return ["--clusters", cluster] if cluster else []
+
+
 @dataclass(frozen=True, slots=True)
 class TopicInfo:
     """Helper for topic specification within Redis bec-slurm pub/sub interface."""
@@ -139,7 +144,14 @@ class SlurmManager:
         sub_info.last_scheduler_check = time.time()
         try:
             queued = self._run_subprocess(
-                ["squeue", "--noheader", "--jobs", future.slurm_job_id, "--format=%T"]
+                [
+                    "squeue",
+                    *_cluster_args(future.slurm_cluster),
+                    "--noheader",
+                    "--jobs",
+                    future.slurm_job_id,
+                    "--format=%T",
+                ]
             )
             if queued.returncode == 0 and queued.stdout.strip():
                 # Pending, configuring, and running jobs may legitimately have no heartbeat.
@@ -147,6 +159,7 @@ class SlurmManager:
             accounting = self._run_subprocess(
                 [
                     "sacct",
+                    *_cluster_args(future.slurm_cluster),
                     "--noheader",
                     "--parsable2",
                     "--allocations",
@@ -220,19 +233,21 @@ class SlurmManager:
         job_future = JobFuture(
             job_id=job_id, manager=self, registry=self._future_registry, slurm_job_id=slurm_job_id
         )
+        slurm_job_id = job_future.slurm_job_id
+        cluster_args = _cluster_args(job_future.slurm_cluster)
         topic_info = TopicInfo(job_id=job_id)
         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])
+            released = self._run_subprocess(["scontrol", *cluster_args, "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])
+                cancelled = self._run_subprocess(["scancel", *cluster_args, 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):
@@ -243,14 +258,18 @@ class SlurmManager:
             raise
         return job_future
 
-    def cancel_job(self, slurm_job_id: str, job_id: str) -> None:
-        """Cancel a job via scancel."""
-        result = self._run_subprocess(["scancel", slurm_job_id])
+    def cancel_job(
+        self, slurm_job_id: str, job_id: str, *, slurm_cluster: str | None = None
+    ) -> None:
+        """Cancel a job on its recorded cluster via scancel."""
+        job_future = self._future_registry.get(job_id)
+        slurm_job_id, _, cluster = slurm_job_id.partition(";")
+        cluster = slurm_cluster or cluster or (job_future.slurm_cluster if job_future else None)
+        result = self._run_subprocess(["scancel", *_cluster_args(cluster), slurm_job_id])
         if result.returncode != 0:
             stderr = (result.stderr or "").strip()
             raise RuntimeError(f"scancel failed: {stderr}")
 
-        job_future = self._future_registry.get(job_id)
         if job_future:
             self._event_callback(
                 StatusMessage(
@@ -425,35 +444,39 @@ class SlurmManager:
             )
             executor.submit(callback, message, **(callback_kwargs or {}))
 
-        future = self._future_registry.get(job_id)
-        if future is None:
-            logger.warning(
-                "Trying to add subscription for job_id %s which does not exist in future registry.",
-                job_id,
+        # Completion removes subscriptions under this same lock. Keep ownership
+        # registration atomic with cleanup, and undo Redis registration if the
+        # future became terminal while subscribe() was in progress.
+        with self._lock:
+            future = self._future_registry.get(job_id)
+            if future is None or future.done():
+                raise RuntimeError(f"Job {job_id} is no longer active")
+            cb_id = self.sub_client.subscribe(
+                topic_str,
+                partial(
+                    wrapped_callback,
+                    executor=self.executor,
+                    callback_kwargs=callback_kwargs,
+                    job_id=job_id,
+                    slurm_job_id=future.slurm_job_id,
+                ),
             )
-            return ""
-        cb_id = self.sub_client.subscribe(
-            topic_str,
-            partial(
-                wrapped_callback,
-                executor=self.executor,
-                callback_kwargs=callback_kwargs,
-                job_id=job_id,
-                slurm_job_id=future.slurm_job_id,
-            ),
-        )
 
-        sub_info = SubscriptionInfo(
-            job_id=job_id,
-            topic=topic_str,
-            slurm_job_id=future.slurm_job_id,
-            callback=wrapped_callback,
-            last_heartbeat=None,
-            callback_id=cb_id,
-            heartbeat_received=False,
-        )
-        self._register_subscription_for_job(sub_info=sub_info)
-        return cb_id
+            if self._future_registry.get(job_id) is not future or future.done():
+                self.sub_client.unsubscribe(cb_id)
+                raise RuntimeError(f"Job {job_id} finished during subscription registration")
+
+            sub_info = SubscriptionInfo(
+                job_id=job_id,
+                topic=topic_str,
+                slurm_job_id=future.slurm_job_id,
+                callback=wrapped_callback,
+                last_heartbeat=None,
+                callback_id=cb_id,
+                heartbeat_received=False,
+            )
+            self._register_subscription_for_job(sub_info=sub_info)
+            return cb_id
 
     def _remove_job_future(self, job_id: str) -> None:
         """Remove a job future from the registry."""

test/test_review_regressions.py

@@ -11,7 +11,7 @@ from unittest import mock
 import fakeredis
 import pytest
 
-from slurm_manager.event_messages import HeartBeatMessage
+from slurm_manager.event_messages import HeartBeatMessage, StatusMessage
 from slurm_manager.job_future import JobFuture
 from slurm_manager.redis_sub_client import RedisSubClient
 from slurm_manager.slurm_manager import SlurmManager, TopicInfo
@@ -220,3 +220,162 @@ def test_removing_last_log_callback_unsubscribes_topic(manager):
     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
+
+
+def cluster_submission(manager):
+    original = manager._run_subprocess.side_effect
+
+    def slurm(command):
+        if command[0] == "sbatch":
+            return subprocess.CompletedProcess(command, 0, "123;cluster_a\n", "")
+        return original(command)
+
+    manager._run_subprocess.side_effect = slurm
+
+
+def test_cluster_qualified_job_is_released_and_cancelled_on_its_cluster(manager):
+    cluster_submission(manager)
+    future = submit(manager)
+    assert future.slurm_job_id == "123"
+    assert future.slurm_cluster == "cluster_a"
+    assert mock.call(["scontrol", "--clusters", "cluster_a", "release", "123"]) in (
+        manager._run_subprocess.call_args_list
+    )
+    future.cancel()
+    assert manager._run_subprocess.call_args == mock.call(
+        ["scancel", "--clusters", "cluster_a", "123"]
+    )
+
+
+def test_cluster_qualified_setup_failure_cancels_on_its_cluster(manager):
+    cluster_submission(manager)
+    with (
+        mock.patch.object(
+            manager.sub_client, "wait_until_subscribed", side_effect=TimeoutError("no ack")
+        ),
+        pytest.raises(TimeoutError),
+    ):
+        submit(manager)
+    assert manager._run_subprocess.call_args == mock.call(
+        ["scancel", "--clusters", "cluster_a", "123"]
+    )
+    assert not manager._future_registry
+    assert not manager.sub_client.registered_topics
+
+
+def test_cluster_qualified_job_monitoring_queries_its_cluster(manager):
+    cluster_submission(manager)
+    future = submit(manager)
+    manager._run_subprocess.reset_mock()
+    manager._run_subprocess.side_effect = [
+        subprocess.CompletedProcess([], 0, "", ""),
+        subprocess.CompletedProcess([], 0, "COMPLETED|0:0\n", ""),
+    ]
+    manager._check_slurm_job(heartbeat_info(manager, future))
+    commands = [call.args[0] for call in manager._run_subprocess.call_args_list]
+    assert [command[0] for command in commands] == ["squeue", "sacct"]
+    for command in commands:
+        assert command[1:3] == ["--clusters", "cluster_a"]
+        assert command[command.index("--jobs") + 1] == "123"
+    assert future.success()
+
+
+def test_completion_during_subscription_registration_leaves_no_orphan(manager):
+    future = submit(manager)
+    subscribed = threading.Event()
+    resume = threading.Event()
+    cleaned = threading.Event()
+    errors = []
+    original_subscribe = manager.sub_client.subscribe
+    original_event = manager._event_callback
+
+    def subscribe(topic, callback):
+        callback_id = original_subscribe(topic, callback)
+        subscribed.set()
+        assert resume.wait(2)
+        return callback_id
+
+    def event(message, future):
+        original_event(message, future)
+        cleaned.set()
+
+    def listen():
+        try:
+            future.listen("log", lambda _: None)
+        except RuntimeError as exc:
+            errors.append(exc)
+
+    with (
+        mock.patch.object(manager.sub_client, "subscribe", side_effect=subscribe),
+        mock.patch.object(manager, "_event_callback", side_effect=event),
+    ):
+        thread = threading.Thread(target=listen)
+        thread.start()
+        try:
+            assert subscribed.wait(1)
+            manager.sub_client._redis.publish(
+                TopicInfo(future._job_id).event, '{"msg_type":"status","status":"finished"}'
+            )
+            future.wait(timeout=1)
+        finally:
+            resume.set()
+            thread.join(timeout=2)
+        assert not thread.is_alive()
+        assert cleaned.wait(1)
+    assert len(errors) == 1
+    assert not manager._future_registry
+    assert not manager._active_subscriptions
+    assert not manager._topic_by_subscription_id
+    assert not manager.sub_client.registered_topics
+    assert not manager.sub_client._topic_by_subscription_id
+
+
+def test_listen_after_completion_does_not_register_callback(manager):
+    future = submit(manager)
+    manager._event_callback(StatusMessage(status="finished"), future)
+    with pytest.raises(RuntimeError, match="no longer active"):
+        future.listen("log", lambda _: None)
+    assert not manager.sub_client.registered_topics
+
+
+def test_unlisten_allows_concurrent_job_cleanup(manager):
+    future = submit(manager)
+    callback_id = future.listen("log", lambda _: None)
+    cleaned = threading.Event()
+    errors = []
+    completed_during_removal = []
+    original_remove = manager.remove_subscription_by_id
+
+    def cleanup():
+        try:
+            manager.remove_subscriptions_for_job(future._job_id)
+        except Exception as exc:
+            errors.append(exc)
+        finally:
+            cleaned.set()
+
+    thread = threading.Thread(target=cleanup)
+
+    def remove(callback_id):
+        # Cleanup takes the manager lock and then updates the future. It must
+        # finish even while unlisten is paused at the manager call boundary.
+        thread.start()
+        completed = cleaned.wait(1)
+        completed_during_removal.append(completed)
+        if completed:
+            original_remove(callback_id)
+        # On regression, return without taking the manager lock so unlisten
+        # releases its lock and the test can join cleanup instead of hanging.
+
+    with mock.patch.object(manager, "remove_subscription_by_id", side_effect=remove):
+        try:
+            future.unlisten(callback_id)
+        finally:
+            thread.join(timeout=2)
+    assert not thread.is_alive()
+    assert not errors
+    assert completed_during_removal == [True]
+    assert future.done()
+    assert not future._callback_registry
+    assert not manager._future_registry
+    assert not manager.sub_client.registered_topics