From 85008e424c4a201ea73ff259fb590886942ca37d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 1 Jul 2026 23:57:10 +0000
Subject: [PATCH] Diffusion training service: join the old pump outside the
 lock

start() joined a finished job's pump thread while holding the service lock,
but the pump's final state writes need that same lock, so the join always
burned its full timeout and a stale pump could then overwrite the new job's
state. Join outside the lock (with a re-check after), and fence _apply_event
and the exit handler by process identity so a superseded pump can never touch
the current job's state. Adds regression tests for both.
---
 .../training/diffusion_training_service.py    | 28 +++++++++++++++----
 .../backend/tests/test_diffusion_training.py  | 28 +++++++++++++++++++
 2 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py
index bdefcfe25b..091f7b881c 100644
--- a/studio/backend/core/training/diffusion_training_service.py
+++ b/studio/backend/core/training/diffusion_training_service.py
@@ -89,11 +89,21 @@ class DiffusionTrainingService:
 
         _config_from_dict(config).normalized()
 
+        # Join a finished job's pump OUTSIDE the lock: its final state writes take this
+        # lock (via _apply_event / the exit handler), so joining under it would stall
+        # the start for the whole timeout and then let the stale pump overwrite the new
+        # job's state once the lock was released.
         with self._lock:
             if self._proc is not None and self._proc.is_alive():
                 raise RuntimeError("A diffusion training job is already running.")
-            if self._pump is not None and self._pump.is_alive():
-                self._pump.join(timeout = 5.0)
+            pump = self._pump
+        if pump is not None and pump.is_alive():
+            pump.join(timeout = 5.0)
+
+        with self._lock:
+            # Re-check: another start() may have won the race while we joined.
+            if self._proc is not None and self._proc.is_alive():
+                raise RuntimeError("A diffusion training job is already running.")
 
             job_id = uuid.uuid4().hex
             event_queue = self._ctx.Queue()
@@ -162,11 +172,13 @@ class DiffusionTrainingService:
                     drained = False
                     while True:
                         try:
-                            self._apply_event(event_queue.get_nowait())
+                            self._apply_event(event_queue.get_nowait(), proc = proc)
                             drained = True
                         except Exception:  # noqa: BLE001
                             break
                     with self._lock:
+                        if self._proc is not proc:
+                            return  # superseded by a newer job; don't touch its state
                         if self._state.get("status") not in ("completed", "stopped", "error"):
                             self._state.update(
                                 active = False,
@@ -177,15 +189,19 @@ class DiffusionTrainingService:
                     _ = drained
                     return
                 continue
-            self._apply_event(ev)
+            self._apply_event(ev, proc = proc)
             if ev.get("type") in _TERMINAL:
                 return
 
-    def _apply_event(self, ev: dict[str, Any]) -> None:
+    def _apply_event(self, ev: dict[str, Any], proc: Any = None) -> None:
         """Fold one trainer event into the status snapshot. Pure state update -- unit
-        tested by feeding events directly."""
+        tested by feeding events directly. ``proc`` (when given) fences a stale pump:
+        an event from a superseded job's process must not touch the current job's
+        state."""
         etype = ev.get("type")
         with self._lock:
+            if proc is not None and self._proc is not proc:
+                return
             s = self._state
             s["updated_at"] = time.time()
             if etype == "model_load_started":
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index ffc231673c..32d239ecde 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -289,3 +289,31 @@ def test_route_status_and_stop(client):
     # After stopping, a stop with nothing running reports idle.
     st2 = client.post("/api/train/diffusion/stop")
     assert st2.json()["status"] == "idle"
+
+
+def test_service_restart_after_completion():
+    # A finished job's pump is joined OUTSIDE the lock (it needs the lock for its
+    # final state writes), so a second start neither stalls nor deadlocks.
+    svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
+    svc.start(dict(_CFG))
+    _wait_status(svc, "completed")
+    t0 = time.time()
+    job2 = svc.start(dict(_CFG))
+    assert job2
+    assert time.time() - t0 < 4.0  # no 5s join-under-lock stall
+    st = _wait_status(svc, "completed")
+    assert st["status"] == "completed"
+
+
+def test_stale_pump_events_cannot_corrupt_new_job():
+    # An event carrying a superseded job's proc identity must be dropped, so a
+    # straggler pump can never overwrite the state of a newly started job.
+    svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
+    svc.start(dict(_CFG))
+    _wait_status(svc, "completed")
+    current = svc._proc
+    svc._apply_event({"type": "error", "message": "stale boom"}, proc = object())
+    assert svc.status()["message"] != "stale boom"
+    # The current job's events still apply.
+    svc._apply_event({"type": "progress", "step": 9}, proc = current)
+    assert svc.status()["step"] == 9