fix: dispatch terminal statuses and log callback failures

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.
 

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                  #

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