From 8bc0149b3569c8ba284eca0619b215bab9b7bd2a Mon Sep 17 00:00:00 2001 From: wakonig_k Date: Wed, 28 Oct 2026 17:16:30 +0100 Subject: [PATCH] fix: dispatch terminal statuses and log callback failures --- slurm_manager/job_tracker.py | 115 +++++++++++++++++++++------------ slurm_manager/slurm_manager.py | 26 ++++---- test/test_job_tracker.py | 73 +++++++++++++++++++++ 3 files changed, 160 insertions(+), 54 deletions(-) diff --git a/slurm_manager/job_tracker.py b/slurm_manager/job_tracker.py index 6a58857..53fd3dc 100644 --- a/slurm_manager/job_tracker.py +++ b/slurm_manager/job_tracker.py @@ -189,22 +189,29 @@ class JobTracker: same cleanup path as an explicit error. """ info = message.status.split(":", 1) - status = info[0] - try: # Catch irregular formatted status messages to avoid crashing the callback loop. - status = JobStatus(status) + try: + status = JobStatus(info[0]) except ValueError: - logger.warning("Received unknown job status '%s' for job %s", status, future.name) - future._update_status( - "error", - f"Received unknown job status '{status}' from message: {message.model_dump()}", - ) - self.remove_subscriptions_for_job(future._job_id, skip_event_callback=True) - return - exc_info = info[1] if len(info) > 1 else None - future._update_status(status, exc_info) - # Clean up subscriptions for finished job - if status in [JobStatus.ERROR, JobStatus.FINISHED, JobStatus.CANCELLED]: - self.remove_subscriptions_for_job(future._job_id, skip_event_callback=True) + logger.warning("Received unknown job status '%s' for job %s", info[0], future.name) + detail = f"Received unknown job status '{info[0]}' from message: {message.model_dump()}" + message = StatusMessage(status=f"error:{detail}", metadata=message.metadata.copy()) + status = JobStatus.ERROR + info = [status.value, detail] + + with self._lock: + if future.done(): + return + topic = TopicInfo(future._job_id).event + callbacks = list(self._subscriptions_by_topic.get(topic, {}).values()) + exc_info = info[1] if len(info) > 1 else None + future._update_status(status, exc_info) + # Queue notifications before cleanup or concurrent shutdown can discard them. + # User code runs on the executor, never on this thread under the state lock. + for subscription in callbacks: + if subscription.callback is not None: + subscription.callback(message) + if status in (JobStatus.ERROR, JobStatus.FINISHED, JobStatus.CANCELLED): + self.remove_subscriptions_for_job(future._job_id, skip_event_callback=True) def get_missing_heartbeats(self) -> list[SubscriptionInfo]: """Select stale heartbeat subscriptions whose scheduler check is due. @@ -325,26 +332,24 @@ class JobTracker: "slurm_job_id": slurm_job_id, } ) - executor.submit(callback, message, **(callback_kwargs or {})) + executor.submit(self._run_callback, callback, message, callback_kwargs or {}) - # 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, - ), - ) - + dispatch = partial( + wrapped_callback, + executor=self._executor, + callback_kwargs=callback_kwargs, + job_id=job_id, + slurm_job_id=future.slurm_job_id, + ) + # Registration may block. Let lifecycle updates proceed, then validate ownership + # under the state lock and undo the Redis subscription if the job finished. + redis_callback = self._ignore_status_message if key == "status" else dispatch + cb_id = self._sub_client.subscribe(topic_str, redis_callback) + with self._lock: 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") @@ -353,7 +358,7 @@ class JobTracker: job_id=job_id, topic=topic_str, slurm_job_id=future.slurm_job_id, - callback=wrapped_callback, + callback=dispatch, last_heartbeat=None, callback_id=cb_id, heartbeat_received=False, @@ -421,17 +426,8 @@ class JobTracker: This only removes local monitoring; it does not cancel the Slurm allocation. """ with self._lock: - topics = self._active_subscriptions.pop(job_id, set()) - for topic in topics: - callbacks = self._subscriptions_by_topic.pop(topic, {}) - for callback_id in callbacks: - self._topic_by_subscription_id.pop(callback_id, None) - self._sub_client.unsubscribe_topic(topic) future = self._future_registry.get(job_id) - if future is None: - return - # This callback will trigger a cleanup. - if not skip_event_callback: + if future is not None and not future.done() and not skip_event_callback: self.apply_status( StatusMessage( status=f"error:Removal requested by user; timestamp {datetime.now().isoformat()}", @@ -443,7 +439,14 @@ class JobTracker: ), future, ) - self._remove_job_future(future._job_id) + return + topics = self._active_subscriptions.pop(job_id, set()) + for topic in topics: + callbacks = self._subscriptions_by_topic.pop(topic, {}) + for callback_id in callbacks: + self._topic_by_subscription_id.pop(callback_id, None) + self._sub_client.unsubscribe_topic(topic) + self._remove_job_future(job_id) ################################################ # Public Lifecycle # @@ -470,6 +473,34 @@ class JobTracker: # Internal Subscriptions # ################################################ + @staticmethod + def _run_callback( + callback: Callable[[SlurmMessage], None], + message: SlurmMessage, + callback_kwargs: dict[str, Any], + ) -> None: + """Execute a user callback and log asynchronous failures. + + Args: + callback (Callable[[SlurmMessage], None]): User callback to execute. + message (SlurmMessage): Parsed job message delivered to the callback. + callback_kwargs (dict[str, Any]): Additional callback keyword arguments. + """ + try: + callback(message, **callback_kwargs) + except Exception: + logger.exception( + "Error in %s callback for job %s", message.msg_type, message.metadata.get("job_id") + ) + + @staticmethod + def _ignore_status_message(message: SlurmMessage) -> None: + """Leave status dispatch to the internal lifecycle subscription. + + Args: + message (SlurmMessage): Redis status already handled by apply_status. + """ + def _subscribe_to_internal_topics(self, future: JobFuture, topic_info: TopicInfo) -> None: """Attach the manager status and heartbeat callbacks before release. diff --git a/slurm_manager/slurm_manager.py b/slurm_manager/slurm_manager.py index 5a08cd5..87fc523 100644 --- a/slurm_manager/slurm_manager.py +++ b/slurm_manager/slurm_manager.py @@ -188,18 +188,20 @@ class SlurmManager: job_future = self._job_tracker.get_future(job_id) self.rest_client.cancel(slurm_job_id) - if job_future: - self._job_tracker.apply_status( - StatusMessage( - status="cancelled", - metadata={ - "msg_received": datetime.now().isoformat(), - "job_id": job_future._job_id, - "slurm_job_id": job_future.slurm_job_id, - }, - ), - future=job_future, - ) + if job_future is None: + return + + self._job_tracker.apply_status( + StatusMessage( + status="cancelled", + metadata={ + "msg_received": datetime.now().isoformat(), + "job_id": job_future._job_id, + "slurm_job_id": job_future.slurm_job_id, + }, + ), + future=job_future, + ) ################################################ # Public Event Subscriptions # diff --git a/test/test_job_tracker.py b/test/test_job_tracker.py index 8a1d33a..10192ba 100644 --- a/test/test_job_tracker.py +++ b/test/test_job_tracker.py @@ -66,3 +66,76 @@ def test_scheduler_check_skips_a_job_removed_after_snapshot(tracker): assert tracker.begin_scheduler_check("first") is None assert tracker.begin_scheduler_check("unknown") is None + + +@pytest.mark.parametrize("status", ["finished", "cancelled", "error:workload failed", "unknown"]) +def test_synthesized_terminal_status_notifies_listeners_once(tracker, status): + future = JobFuture("first", manager=None, slurm_job_id="123") + tracker.track(future) + received = [] + tracker.add_subscription("first", "status", received.append) + topic = TopicInfo("first").event + queued_redis_callbacks = list(tracker._sub_client._subscriptions_by_topic[topic].values()) + message = StatusMessage(status=status) + + tracker.apply_status(message, future) + # Simulate a Redis message whose callbacks were captured before local cleanup. + for subscription in queued_redis_callbacks: + subscription.callback(message) + tracker.shutdown(timeout=3) + + assert len(received) == 1 + assert received[0].metadata["job_id"] == "first" + assert received[0].metadata["slurm_job_id"] == "123" + assert received[0].status == status or received[0].status.startswith("error:") + assert future.done() + + +def test_redis_status_is_dispatched_once_and_shutdown_drains_it(tracker): + future = JobFuture("first", manager=None, slurm_job_id="123") + tracker.track(future) + received = [] + tracker.add_subscription("first", "status", received.append) + tracker.wait_until_ready("first") + tracker._sub_client._redis.publish( + TopicInfo("first").event, '{"msg_type":"status","status":"finished"}' + ) + future.wait(timeout=1) + tracker.shutdown(timeout=3) + + assert [message.status for message in received] == ["finished"] + + +def test_removing_tracking_reports_failure_to_status_listeners(tracker): + future = JobFuture("first", manager=None, slurm_job_id="123") + tracker.track(future) + received = [] + tracker.add_subscription("first", "status", received.append) + + tracker.remove_subscriptions_for_job("first") + tracker.shutdown(timeout=3) + + assert len(received) == 1 + assert received[0].status.startswith("error:Removal requested by user") + assert tracker.get_future("first") is None + assert tracker._sub_client.registered_topics == set() + + +def test_callback_errors_are_logged_without_stopping_other_callbacks(tracker, caplog): + import threading + + future = JobFuture("first", manager=None, slurm_job_id="123") + tracker.track(future) + received = threading.Event() + + def fail(message): + raise RuntimeError("callback failed") + + tracker.add_subscription("first", "status", fail) + tracker.add_subscription("first", "status", lambda message: received.set()) + tracker.apply_status(StatusMessage(status="finished"), future) + tracker.shutdown(timeout=3) + + assert received.is_set() + assert "Error in status callback for job first" in caplog.text + assert "RuntimeError: callback failed" in caplog.text