Let real generations preempt the background compile prewarm

The prewarm registered its cancel event in _active_generate_cancel, but a
begin_generate arriving mid-warmup overwrote that slot with its own event
and then queued its worker behind the full warmup on _generate_lock. From
that point unload and cancel_generate signalled the wrong event, so the
warmup could no longer be aborted and the first real request waited out
the 9-54s the prewarm exists to hide.

Track the prewarm's event in a dedicated _prewarm_cancel slot (cleared
identity-checked alongside _active_generate_cancel) and signal it from
begin_generate before registering the real job's event, and from direct
generate() calls that skip begin_generate. The warmup then aborts at its
next step boundary and the real job takes the lock, while unload/cancel
keep working against whichever run is actually active.

studio/backend/core/inference/video.py

@@ -490,6 +490,11 @@ class VideoBackend:
         # The post-load background compile prewarm thread (None until a compiled
         # load spawns one); kept for tests/diagnostics, never joined on the hot path.
         self._prewarm_thread: Optional[threading.Thread] = None
+        # The prewarm's cancel event, set (in addition to _active_generate_cancel)
+        # only while the prewarm runs. Real generations signal it on entry so they
+        # preempt the warmup at its next step boundary instead of queueing behind
+        # it; unlike _active_generate_cancel it can never point at a real job.
+        self._prewarm_cancel: Optional[threading.Event] = None
 
     # ── validation ───────────────────────────────────────────────────────────
 
@@ -1872,6 +1877,7 @@ class VideoBackend:
                 if self._generate_job_active or self._gen.get("active"):
                     return  # a real request beat us; it absorbs the warmup itself
                 self._active_generate_cancel = cancel
+                self._prewarm_cancel = cancel
             started = time.monotonic()
             try:
                 pipe = state.pipe
@@ -1947,6 +1953,8 @@ class VideoBackend:
                 with self._lock:
                     if self._active_generate_cancel is cancel:
                         self._active_generate_cancel = None
+                    if self._prewarm_cancel is cancel:
+                        self._prewarm_cancel = None
 
     # ── generation ───────────────────────────────────────────────────────────
 
@@ -2002,6 +2010,13 @@ class VideoBackend:
                 raise RuntimeError(VIDEO_NOT_LOADED_MSG)
             if self._generate_job_active:
                 raise RuntimeError(VIDEO_GENERATION_BUSY_MSG)
+            # A background compile prewarm may hold _generate_lock. Signal its
+            # dedicated cancel handle BEFORE registering ours so the real job
+            # preempts the warmup at its next step boundary instead of queueing
+            # behind the full warmup (which also left unload/cancel pointing at
+            # the wrong event once _active_generate_cancel was overwritten below).
+            if self._prewarm_cancel is not None:
+                self._prewarm_cancel.set()
             self._generate_job_active = True
             # Register the cancel event BEFORE the worker starts so a cancel (or an
             # unload) that lands in the spawn window still stops the run instead of
@@ -2147,6 +2162,12 @@ class VideoBackend:
         # begin_generate passes the event it already registered (so a cancel in the
         # spawn window is honoured); a direct call makes its own.
         cancel = cancel_event if cancel_event is not None else threading.Event()
+        if cancel_event is None:
+            # Direct callers skip begin_generate's preemption, so signal a
+            # running compile prewarm here too rather than queueing behind it.
+            with self._lock:
+                if self._prewarm_cancel is not None:
+                    self._prewarm_cancel.set()
         with self._generate_lock:
             with self._lock:
                 state = self._state

studio/backend/tests/test_video_backend.py

@@ -2023,3 +2023,57 @@ def test_compile_prewarm_yields_to_generations_and_stale_tokens(fake_runtime):
     backend._compile_prewarm(backend._load_token)  # a request beat the warmup
     assert backend._state.pipe.last_kwargs is None
     backend._generate_job_active = False
+
+
+def test_begin_generate_preempts_running_prewarm(fake_runtime, monkeypatch):
+    # A real generation arriving while the prewarm holds _generate_lock must
+    # signal the prewarm's dedicated cancel handle (so the warmup aborts at its
+    # next step boundary instead of running to completion in front of the user
+    # job), and the prewarm must clear that handle when it exits.
+    import inspect
+    import threading as _threading
+    import time
+
+    backend = VideoBackend()
+    backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
+
+    prewarm_entered = _threading.Event()
+    release_prewarm = _threading.Event()
+    pipe = backend._state.pipe
+    real_call = pipe.__class__.__call__
+
+    def _blocking_call(self, **kwargs):
+        if kwargs.get("prompt") == "warmup":
+            prewarm_entered.set()
+            release_prewarm.wait(timeout = 5)
+        return real_call(self, **kwargs)
+
+    # Keep the real signature visible: _compile_prewarm picks its cancel plumbing
+    # by inspecting pipe.__call__ for callback_on_step_end.
+    _blocking_call.__signature__ = inspect.signature(real_call)
+    monkeypatch.setattr(pipe.__class__, "__call__", _blocking_call)
+
+    prewarm = _threading.Thread(
+        target = backend._compile_prewarm, args = (backend._load_token,), daemon = True,
+    )
+    prewarm.start()
+    assert prewarm_entered.wait(timeout = 5), "prewarm never reached the pipe"
+    prewarm_cancel = backend._prewarm_cancel
+    assert prewarm_cancel is not None and not prewarm_cancel.is_set()
+
+    # The user job lands mid-warmup: it must fire the prewarm's cancel handle
+    # and own the active-cancel slot for the run that follows.
+    backend.begin_generate(prompt = "real request")
+    assert prewarm_cancel.is_set()
+    assert backend._active_generate_cancel is not prewarm_cancel
+
+    release_prewarm.set()
+    prewarm.join(timeout = 5)
+    assert not prewarm.is_alive()
+    assert backend._prewarm_cancel is None
+
+    deadline = time.monotonic() + 5
+    while backend._generate_job_active and time.monotonic() < deadline:
+        time.sleep(0.02)
+    assert backend._generate_job_active is False
+    backend.unload()