diff --git a/slurm_manager/job_future.py b/slurm_manager/job_future.py index 95604a4..a75a845 100644 --- a/slurm_manager/job_future.py +++ b/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 diff --git a/slurm_manager/slurm_manager.py b/slurm_manager/slurm_manager.py index 1d3746d..da8b0b0 100644 --- a/slurm_manager/slurm_manager.py +++ b/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.""" diff --git a/test/test_review_regressions.py b/test/test_review_regressions.py index 50585d3..855261a 100644 --- a/test/test_review_regressions.py +++ b/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