From 47c202eee13dc4d6294e5f2db43e9c764e4f639d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 08:10:03 +0000 Subject: [PATCH] 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 | 21 +++++++++ studio/backend/tests/test_video_backend.py | 54 ++++++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 35dcce802c..b939b4455c 100644 --- a/studio/backend/core/inference/video.py +++ b/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 diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 18b1d5fcca..722d475625 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/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()