diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 3924950793d..14817974b94 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -18,6 +18,7 @@ from __future__ import annotations +from contextlib import contextmanager import functools import inspect import json @@ -1094,12 +1095,21 @@ def __init__(self) -> None: self._load_token = 0 # Set by unload() to abort an in-flight download. Replaced, never cleared, so a cancelled worker stays cancelled. self._cancel_event = threading.Event() + # Cancellation state has its own tiny lock so Stop remains responsive while a replacement + # load holds _lock for model construction. + self._generation_cancel_lock = threading.Lock() # Cancel Event of the in-flight generation; per-generation so a cancel can't be lost or leak. self._active_generate_cancel: Optional[threading.Event] = None - # Unloads / superseding loads waiting on _generate_lock to free this pipeline. A generation queued behind the active - # one holds no cancel event yet, so without this fence it could win the lock after an eject and denoise anyway. A - # count, not a flag, so concurrent teardowns each own their own release. + # Requests waiting behind teardown are cancellable too, but cannot share the active + # slot: multiple HTTP callers may queue while another generation is still denoising. + self._queued_generate_cancels: set[threading.Event] = set() + # Unloads / superseding loads waiting on _generate_lock to free this pipeline. A queued + # generation is not the ACTIVE one teardown should cancel, so without this fence it could + # win the lock after an eject and denoise anyway. A count lets concurrent teardowns reserve. self._teardown_waiters = 0 + # Wakes a generation that yielded the generation lock to a pending teardown. It shares + # _lock so checking the count and sleeping cannot miss a completed teardown. + self._teardown_drained = threading.Condition(self._lock) # Written by the callback, read lock-free by generate_progress(). self._gen: Optional[_GenState] = None # img2img/inpaint pipes built via from_pipe (shared modules, no extra VRAM); cleared on unload. @@ -1122,6 +1132,76 @@ def _pick_device_and_dtype(self, ordinal: Optional[int] = None) -> tuple[str, An # The INDEXED string, so _resolve_device_target can rebuild a selection an override would erase. return target.torch_device, target.dtype + def _release_teardown_locked(self) -> None: + """Release one teardown reservation and wake generations when the last one leaves. + + Call only while holding ``_lock``. A count is necessary because an unload and a + superseding load can both be queued behind the active generation. + """ + assert self._teardown_waiters > 0, "teardown reservation released without an owner" + self._teardown_waiters -= 1 + if self._teardown_waiters == 0: + self._teardown_drained.notify_all() + + @contextmanager + def _generation_slot(self, cancel: threading.Event): + """Hold the generation lock, yielding to teardown and remaining cancellable. + + Lock acquisition is not FIFO. If a generation wins the lock after a load or unload + has raised its fence, it must let that teardown run before reading ``_state``. Once + the final fence drops, the teardown still owns ``_generate_lock`` until its model + transition has settled, so the retried acquisition observes the new truthful state. + + The zero-fence check and active-cancel registration share one ``_lock`` section. A + teardown starting after that check therefore either sees this event and cancels it, + or reserved before the check and makes this request yield. + """ + admitted = False + with self._generation_cancel_lock: + self._queued_generate_cancels.add(cancel) + try: + while True: + # A replacement holds this lock throughout construction. Timed acquisition keeps + # a queued HTTP request responsive to Stop without weakening that barrier. + while not self._generate_lock.acquire(timeout = 0.1): + if cancel.is_set(): + raise RuntimeError(DIFFUSION_CANCELLED_MSG) + with self._lock: + if not self._teardown_waiters: + # Lock order is state -> cancellation everywhere teardown touches both. + # Registration is therefore atomic with the zero-fence observation. + with self._generation_cancel_lock: + cancelled = cancel.is_set() + if not cancelled: + self._queued_generate_cancels.discard(cancel) + self._active_generate_cancel = cancel + admitted = True + else: + cancelled = cancel.is_set() + if admitted: + break + self._generate_lock.release() + if cancelled: + raise RuntimeError(DIFFUSION_CANCELLED_MSG) + with self._teardown_drained: + while self._teardown_waiters and not cancel.is_set(): + # Cancellation uses its independent lock and cannot notify this condition + # while a load owns _lock, so wake periodically only while actually queued. + self._teardown_drained.wait(timeout = 0.1) + if cancel.is_set(): + raise RuntimeError(DIFFUSION_CANCELLED_MSG) + try: + yield + finally: + with self._generation_cancel_lock: + if self._active_generate_cancel is cancel: + self._active_generate_cancel = None + self._generate_lock.release() + finally: + if not admitted: + with self._generation_cancel_lock: + self._queued_generate_cancels.discard(cancel) + # Memory requests whose offload policy is decided by the REQUEST rather than by the measured # footprint. `fast` and `auto` are measurements and so cannot be judged network-free. def assert_precision_available( @@ -3087,8 +3167,9 @@ def load_pipeline( # Bail before signalling if this load was superseded, else a stale worker aborts a live one. if _load_token is not None and _load_token != self._load_token: raise RuntimeError("Diffusion load was cancelled.") - if self._active_generate_cancel is not None: - self._active_generate_cancel.set() + with self._generation_cancel_lock: + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() # Same fence unload() takes: a queued generation must not run on the pipeline this load is about to free. self._teardown_waiters += 1 with self._generate_lock: @@ -3102,7 +3183,7 @@ def load_pipeline( self._unload_locked() finally: # Released here, not at the end of the load: the old pipe is gone and the rest of the load holds _generate_lock. - self._teardown_waiters -= 1 + self._release_teardown_locked() # Single-file kinds resolve a checkpoint path; the pipeline kind has none. single_file_path = ( @@ -5281,16 +5362,13 @@ def generate( # Per-generation cancel Event that unload()/a superseding load set (under _lock) to abort just this denoise. cancel = threading.Event() - with self._generate_lock: + with self._generation_slot(cancel): with self._lock: - # A teardown is waiting for this lock and Python locks are not FIFO, so refuse rather than start a denoise on a pipeline that is already being torn down. - if self._teardown_waiters: - raise RuntimeError(DIFFUSION_CANCELLED_MSG) state = self._state if state is None: raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) - # Register under _lock so unload()/a load can signal THIS generation. - self._active_generate_cancel = cancel + if cancel.is_set(): + raise RuntimeError(DIFFUSION_CANCELLED_MSG) # Publish an active (step 0) state before the slow pre-denoise setup so a reload mount probe does not read idle. self._gen = _GenState(total_steps = steps) try: @@ -5715,11 +5793,11 @@ def _on_step(pipe, step_index, timestep, callback_kwargs): # artifact per shape, so that save is not instant) and the page still shows Stop # for as long as progress reads active, so a Stop landing there was answered # cancelled = true and then contradicted by the image the route persisted. - # Check and deregister under _lock, which is the lock cancel_generate takes, so the + # Check and deregister under the cancellation lock, which cancel_generate takes, so the # two cannot interleave: a cancel that saw this event registered ran strictly # before the check, and one that arrives after finds nothing to set and answers # false. The finally below repeats the clear for every other exit. - with self._lock: + with self._generation_cancel_lock: if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) if self._active_generate_cancel is cancel: @@ -5750,18 +5828,23 @@ def _on_step(pipe, step_index, timestep, callback_kwargs): } finally: # Deregister so a later unload/load can't poke a finished generation (if still ours). - with self._lock: + with self._generation_cancel_lock: if self._active_generate_cancel is cancel: self._active_generate_cancel = None + with self._lock: # Sole clear of the published progress state, on every exit, so a crashed generation never leaves the UI stuck. self._gen = None def generate_progress(self) -> dict[str, Any]: - """Live per-step progress for an in-flight generation (lock-free read).""" + """Live per-step progress for an in-flight or teardown-queued generation.""" gen = self._gen if gen is None or gen.total_steps <= 0: + with self._generation_cancel_lock: + pending = bool( + self._queued_generate_cancels or self._active_generate_cancel is not None + ) return { - "active": False, + "active": pending, "step": 0, "total_steps": 0, "fraction": 0.0, @@ -5786,11 +5869,14 @@ def cancel_generate(self) -> bool: Best effort by construction: the sampler stops at the NEXT step callback, so a cancel during the VAE decode or the encode that precedes step 0 lands when that finishes. Same contract as the video backend.""" - with self._lock: - cancel = self._active_generate_cancel - if cancel is None: + with self._generation_cancel_lock: + cancels = set(self._queued_generate_cancels) + if self._active_generate_cancel is not None: + cancels.add(self._active_generate_cancel) + if not cancels: return False - cancel.set() + for cancel in cancels: + cancel.set() return True def unload(self) -> dict[str, Any]: @@ -5799,9 +5885,11 @@ def unload(self) -> dict[str, Any]: # rebinds this attribute, so an unlocked read could set an event the current load no longer watches. self._cancel_event.set() # Abort an in-flight denoise via ITS cancel event. - if self._active_generate_cancel is not None: - self._active_generate_cancel.set() - # Fence queued generations too: they hold no cancel event yet, so the signal above cannot reach them. + with self._generation_cancel_lock: + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() + # Fence queued generations too: they are intentionally not cancelled by model + # lifecycle changes, so they must wait and observe the post-teardown state. self._teardown_waiters += 1 # Cancel any in-flight load (its worker checks this token) and drop the marker. self._load_token += 1 @@ -5815,7 +5903,7 @@ def unload(self) -> dict[str, Any]: finally: # Released in a finally, exactly like begin_load: _unload_locked ends in clear_gpu_cache(), which raises on a # sticky CUDA fault, and an un-drained fence would refuse every later generation for the life of the process. - self._teardown_waiters -= 1 + self._release_teardown_locked() return self.status() def _unload_locked(self) -> None: diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 5bc86449586..bda4fb44111 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -31,6 +31,7 @@ from __future__ import annotations import contextlib +import functools import inspect import os import tempfile @@ -542,6 +543,11 @@ class _VideoLoadState: class _VideoLoadingState: repo_id: str base_repo: str + # The model the in-flight load will commit, so a generation queued behind a replacement + # is validated synchronously against the family it will actually run on. + family: Optional[VideoFamily] = None + engine: Optional[str] = None + h3_task: Optional[str] = None expected_bytes: Optional[int] = None error: Optional[str] = None # Companion repos this load is ALSO pulling from, beyond repo_id / base_repo (the native H3 @@ -1005,6 +1011,29 @@ def _probe_target(request_shape: dict[str, Any]) -> Any: return types.SimpleNamespace(device = request_shape.get("device"), dtype = dtype) +@dataclass +class _TeardownReservation: + active: bool = False + + +def _drain_teardown_after_load(load): + """Keep a standard-load teardown reservation until the replacement settles.""" + + @functools.wraps(load) + def wrapped(self, *args, **kwargs): + reservation = _TeardownReservation() + kwargs["_teardown_reservation"] = reservation + try: + return load(self, *args, **kwargs) + finally: + with self._lock: + if reservation.active: + reservation.active = False + self._release_teardown_locked() + + return wrapped + + class VideoBackend: """One loaded video pipeline; loads swap it atomically (same model as images).""" @@ -1015,11 +1044,25 @@ def __init__(self) -> None: self._loading: Optional[_VideoLoadingState] = None self._load_token = 0 self._cancel_event = threading.Event() + # Cancellation state has its own tiny lock so a user Stop stays responsive while a + # load holds the state lock. Everything touching _queued_generate_cancels or + # _active_generate_cancel takes this; code holding _lock may nest it, never the reverse. + self._generation_cancel_lock = threading.Lock() self._active_generate_cancel: Optional[threading.Event] = None + # Cancel events of jobs queued behind a teardown but not yet admitted to the generation + # slot. Only the ADMITTED event lives in _active_generate_cancel, so a load/unload cancels + # the denoising job without killing one that is merely waiting for a replacement, while a + # user Stop (cancel_generate) signals both. + self._queued_generate_cancels: set[threading.Event] = set() # How many unloads / superseding loads are waiting on _generate_lock to free this pipeline. A generation queued behind # the active one holds no cancel event yet, so without this fence it could win the lock after an eject and denoise a # whole new clip against a pipeline being freed. A count, so concurrent teardowns each own their own release. self._teardown_waiters = 0 + # Wakes a generation that yielded the generation lock to a pending teardown. A plain Event, + # NOT a Condition on _lock: Event.wait() returns without reacquiring the state lock, so a + # load holding _lock cannot stall a queued job's cancellation (the wait loop re-checks the + # fence under _lock and the cancel Event lock-free each wake). + self._teardown_drained = threading.Event() # Generation progress, written by the step callback / phase transitions. self._gen: dict[str, Any] = {"active": False} # True from begin_generate() until its worker records a terminal state, so a second call is refused while it runs. @@ -1038,6 +1081,90 @@ def _device_target(self, ordinal: Optional[int] = None) -> DiffusionDeviceTarget apply_diffusion_device_ordinal(target) return target + def _release_teardown_locked(self) -> None: + """Release one teardown reservation and wake generations when the last one leaves.""" + assert self._teardown_waiters > 0, "teardown reservation released without an owner" + self._teardown_waiters -= 1 + if self._teardown_waiters == 0: + self._teardown_drained.set() + + def _cancel_active_generation_locked(self) -> bool: + """Cancel the ADMITTED generation. + + Queued jobs are deliberately left alone: a model replacement must not cancel a job + that is merely waiting for teardown, it should survive and run against the + replacement. cancel_generate() signals those separately. The cancel state lives + under its own lock, so this stays cheap for the load/unload callers that hold _lock.""" + with self._generation_cancel_lock: + cancel = self._active_generate_cancel + if cancel is None: + return False + cancel.set() + return True + + @contextlib.contextmanager + def _generation_slot(self, cancel_event: Optional[threading.Event] = None): + """Hold the generation lock, yielding to queued teardown or cancellation. + + A queued job registers its cancel event separately so a load/unload cannot cancel + it. Admission promotes the event to the active slot in the SAME ``_lock`` section + that observes the zero-fence state: a teardown starting after admission either sees + this event and cancels it, or reserved before the check and made this request yield. + While queued, cancellation is observed WITHOUT the state lock: the wake is a plain + Event whose wait returns without reacquiring ``_lock``, so a load holding ``_lock`` + cannot stall a queued Stop. + """ + if cancel_event is not None: + with self._generation_cancel_lock: + self._queued_generate_cancels.add(cancel_event) + admitted = False + try: + while True: + # A replacement holds this lock while its final memory placement and state + # commit run. Polling lets a queued background job observe cancellation instead + # of remaining active until that potentially lengthy placement finishes. + while not self._generate_lock.acquire(timeout = 0.1): + if cancel_event is not None and cancel_event.is_set(): + raise RuntimeError(VIDEO_CANCELLED_MSG) + with self._lock: + cancelled = cancel_event is not None and cancel_event.is_set() + if not cancelled and not self._teardown_waiters: + if cancel_event is not None: + with self._generation_cancel_lock: + self._queued_generate_cancels.discard(cancel_event) + self._active_generate_cancel = cancel_event + admitted = True + if admitted: + break + self._generate_lock.release() + if cancelled: + raise RuntimeError(VIDEO_CANCELLED_MSG) + while True: + # Lock-free: reading the cancel Event needs no state lock, so cancellation + # stays responsive while a load holds _lock for its construction. + if cancel_event is not None and cancel_event.is_set(): + raise RuntimeError(VIDEO_CANCELLED_MSG) + with self._lock: + if self._teardown_waiters == 0: + break + # Event.wait returns without reacquiring _lock (unlike a Condition bound + # to it); the 100 ms timeout bounds how late a cancel is observed. + self._teardown_drained.wait(timeout = 0.1) + self._teardown_drained.clear() + try: + yield + finally: + if cancel_event is not None: + with self._generation_cancel_lock: + if self._active_generate_cancel is cancel_event: + self._active_generate_cancel = None + self._queued_generate_cancels.discard(cancel_event) + self._generate_lock.release() + finally: + if not admitted and cancel_event is not None: + with self._generation_cancel_lock: + self._queued_generate_cancels.discard(cancel_event) + def _state_device_target(self, state: _VideoLoadState) -> DiffusionDeviceTarget: """The resident pipeline's target, pinned onto the calling thread. Every worker that touches the loaded pipeline goes through this: the weights are on ``state.gpu_ordinal`` @@ -1347,6 +1474,13 @@ def begin_load( self._loading = _VideoLoadingState( repo_id = repo_id, base_repo = fam.base_repo, + family = fam, + engine = "sd_cpp" if h3_native else "diffusers", + # Mirrors _run_load_h3_native's derivation; for diffusers loads the workflow + # param is the task (None for non-H3 families). + h3_task = ( + h3_transformer_task(gguf_filename) if h3_native and gguf_filename else h3_task + ), asset_repos = claimed_assets, ) @@ -1881,8 +2015,7 @@ def _run_load_h3_native( with self._lock: if token is not None and token != self._load_token: raise RuntimeError("Video load was cancelled or superseded.") - if self._active_generate_cancel is not None: - self._active_generate_cancel.set() + self._cancel_active_generation_locked() self._teardown_waiters += 1 with self._generate_lock: with self._lock: @@ -1921,7 +2054,7 @@ def _run_load_h3_native( ), ) finally: - self._teardown_waiters -= 1 + self._release_teardown_locked() if native_device == "cpu": # /video/load acquired the VIDEO GPU claim off the resolved device target, but no # accelerator binary was available and this runtime committed to the CPU build, so it @@ -3139,6 +3272,7 @@ def loaded_repo_ids(self) -> tuple[str, ...]: # ── the load itself ────────────────────────────────────────────────────── + @_drain_teardown_after_load def load_pipeline( self, repo_id: str, @@ -3162,6 +3296,7 @@ def load_pipeline( _base_local_dir: Optional[str] = None, _te_prequant_skipped: tuple[str, ...] = (), _h3_auto_denoiser_planned: Optional[str] = None, + _teardown_reservation: Optional[_TeardownReservation] = None, ) -> dict[str, Any]: fam = self.validate_load_request( repo_id, @@ -3206,22 +3341,19 @@ def load_pipeline( if _load_token is not None and _load_token != self._load_token: raise RuntimeError("Video load was cancelled or superseded.") # Signal a generation from the PREVIOUS model (the token check above bailed a superseded worker). - if self._active_generate_cancel is not None: - self._active_generate_cancel.set() + self._cancel_active_generation_locked() # Same fence unload() takes, raised BEFORE the barrier: a queued generation holds no cancel event, so the signal # above cannot reach it and it would slip through the moment the barrier released _generate_lock. + assert _teardown_reservation is not None + _teardown_reservation.active = True self._teardown_waiters += 1 # Barrier: wait for the signalled generation to exit before teardown, or two models coexist in VRAM. with self._generate_lock: with self._lock: - try: - # The barrier wait can outlive this load (superseded by a newer load/unload); recheck before touching shared state. - if _load_token is not None and _load_token != self._load_token: - raise RuntimeError("Video load was cancelled or superseded.") - self._teardown_state_locked() - finally: - # Released here, not at the end of the load: the old pipe is gone (or this load bailed), and a raising teardown must not leave the fence up for the life of the process. - self._teardown_waiters -= 1 + # The barrier wait can outlive this load (superseded by a newer load/unload); recheck before touching shared state. + if _load_token is not None and _load_token != self._load_token: + raise RuntimeError("Video load was cancelled or superseded.") + self._teardown_state_locked() target = self._device_target(gpu_ordinal) device = target.device @@ -4704,47 +4836,74 @@ def begin_generate( sentinels the route maps to 409. """ cancel = threading.Event() - # Snapshot load state before decoding outside the backend lock. + # Snapshot load state before decoding outside the backend lock. A replacement holds a + # teardown reservation while it builds with _state cleared, so "nothing loaded" means + # either truly unloaded (refuse now) or a replacement in flight (queue the job; its + # worker validates against the replacement state once it is available). with self._lock: - if self._state is None: - raise RuntimeError(VIDEO_NOT_LOADED_MSG) if self._generate_job_active: raise RuntimeError(VIDEO_GENERATION_BUSY_MSG) - family = self._state.family - task = self._state.h3_task - engine = self._state.engine - # Validate conditioning before creating the asynchronous job so request errors return 400. - _, _, canvas_w, canvas_h, _ = self._resolve_keyframes( - family, task, first_frame, last_frame, width, height - ) - self._resolve_references( - family, - task, - engine, - reference_images, - reference_videos, - reference_audios, - reference_image_size, - canvas_w, - canvas_h, - ) - self._resolve_flow_shifts(family, engine, flow_shift, audio_flow_shift) - with self._lock: - if self._state is None: + if self._state is None and self._teardown_waiters == 0: raise RuntimeError(VIDEO_NOT_LOADED_MSG) + state = self._state + # A standard replacement tears the old pipeline down before committing the new one, + # so validate against the model the IN-FLIGHT load will commit: malformed image data, + # unsupported references or invalid flow shifts must still 400/422 here rather than + # fail the queued job asynchronously. The worker re-checks against the committed state. + if state is not None: + family, task, engine = state.family, state.h3_task, state.engine + elif self._loading is not None: + family, task, engine = ( + self._loading.family, + self._loading.h3_task, + self._loading.engine, + ) + else: + family = task = engine = None + if family is not None: + # Validate conditioning before creating the asynchronous job so request errors return 400. + _, _, canvas_w, canvas_h, _ = self._resolve_keyframes( + family, task, first_frame, last_frame, width, height + ) + self._resolve_references( + family, + task, + engine, + reference_images, + reference_videos, + reference_audios, + reference_image_size, + canvas_w, + canvas_h, + ) + self._resolve_flow_shifts(family, engine, flow_shift, audio_flow_shift) + with self._lock: if self._generate_job_active: raise RuntimeError(VIDEO_GENERATION_BUSY_MSG) + if self._state is None and self._teardown_waiters == 0: + raise RuntimeError(VIDEO_NOT_LOADED_MSG) # Under the SAME lock that reserves the state this job will run against. A load # commits its new state here too, so judging the shape from a separate earlier read # could accept a size for the family being replaced and then denoise it with the new # one, or reject a size the new family supports. getattr, so a state carrying no - # family degrades to the old snapping rather than raising. + # family degrades to the old snapping rather than raising. During a replacement the + # resident state is gone, so the shape is judged against the incoming family. fam = getattr(self._state, "family", None) + if fam is None and self._loading is not None: + fam = self._loading.family if fam is not None: validate_video_request_shape(fam, width = width, height = height, num_frames = num_frames) + validated_state = self._state + # A job accepted while a replacement was building had no committed state to judge, + # so its worker repeats the shape check against the state that actually denoises. + defer_shape_validation = validated_state is None self._generate_job_active = True - # Register BEFORE the worker starts so a cancel/unload in the spawn window still stops the run. - self._active_generate_cancel = cancel + # Registered BEFORE the worker starts so a cancel in the spawn window still stops the + # run. QUEUED, not active: a load/unload cancels only the admitted generation, so a + # job waiting behind a replacement survives it. The generation slot promotes this + # event to _active_generate_cancel when the job is admitted. + with self._generation_cancel_lock: + self._queued_generate_cancels.add(cancel) self._gen = { "active": True, "phase": "queued", @@ -4774,6 +4933,8 @@ def begin_generate( flow_shift = flow_shift, audio_flow_shift = audio_flow_shift, cancel_event = cancel, + _validated_state = validated_state, + _defer_shape_validation = defer_shape_validation, ), daemon = True, ).start() @@ -4864,9 +5025,13 @@ def _finish_generate_job( moment a new begin_generate() can start is after the outcome is visible.""" with self._lock: self._generate_job_active = False - if cancel_event is not None and self._active_generate_cancel is cancel_event: - # Covers a worker that failed before reaching generate()'s finally; identity-guarded so a direct generate() keeps its handle. - self._active_generate_cancel = None + if cancel_event is not None: + with self._generation_cancel_lock: + # Covers a worker that failed before reaching generate()'s finally; identity-guarded so a direct generate() keeps its handle. + if self._active_generate_cancel is cancel_event: + self._active_generate_cancel = None + # A job cancelled or failed while still queued never reached the slot's admission, so drop its queued registration here. + self._queued_generate_cancels.discard(cancel_event) if error is not None: self._gen = { "active": False, @@ -4908,18 +5073,23 @@ def generate( flow_shift: Optional[float] = None, audio_flow_shift: Optional[float] = None, cancel_event: Optional[threading.Event] = None, + _validated_state: Optional[_VideoLoadState] = None, + # True when begin_generate accepted the job without a committed state to judge (a + # replacement was building), so the shape must be validated against the state that + # actually denoises. Direct callers leave it False: their shape is judged elsewhere. + _defer_shape_validation: bool = False, ) -> dict[str, Any]: # begin_generate passes its already-registered event; a direct call makes its own. cancel = cancel_event if cancel_event is not None else threading.Event() - with self._generate_lock: + with self._generation_slot(cancel): with self._lock: - # A teardown is waiting for this lock and Python locks are not FIFO, so refuse rather than denoise against a pipeline that is already being torn down. - if self._teardown_waiters: - raise RuntimeError(VIDEO_CANCELLED_MSG) state = self._state if state is None: raise RuntimeError(VIDEO_NOT_LOADED_MSG) - self._active_generate_cancel = cancel + # The slot registered this event on admission; a teardown that signalled it + # while we waited must still refuse instead of denoising a cancelled run. + if cancel.is_set(): + raise RuntimeError(VIDEO_CANCELLED_MSG) # Bound below, once the request is resolved. None means the failure beat the # resolution, and there is nothing truthful to report. request_shape: Optional[dict[str, Any]] = None @@ -4930,6 +5100,18 @@ def generate( # the pipeline sits on the selected one. self._state_device_target(state) fam = state.family + # begin_generate validates synchronously against the state visible when the job + # is accepted. Admission may then wait through a replacement, so repeat the shape + # check against the exact state that will denoise instead of silently snapping a + # request that only the previous family supported. A job accepted while a + # replacement was building had no committed state to judge at acceptance, so it + # validates against the admitted state here. + if _defer_shape_validation or ( + _validated_state is not None and state is not _validated_state + ): + validate_video_request_shape( + fam, width = width, height = height, num_frames = num_frames + ) first_pil, last_pil, width, height, conditioning = self._resolve_keyframes( fam, state.h3_task, first_frame, last_frame, width, height ) @@ -5307,7 +5489,7 @@ def _on_scheduler_step(done: int) -> None: _log_failed_generation(request_shape, exc) raise finally: - with self._lock: + with self._generation_cancel_lock: if self._active_generate_cancel is cancel: self._active_generate_cancel = None @@ -5757,13 +5939,24 @@ def forget_terminal_video(self, video_id: Optional[str] = None) -> bool: return True def cancel_generate(self) -> bool: - """Signal the in-flight generation to stop at its next step callback.""" - with self._lock: - cancel = self._active_generate_cancel - if cancel is None: - return False - cancel.set() - return True + """Signal the in-flight generation to stop at its next step callback. + + Cancels the ADMITTED generation and any job queued behind a teardown: a user Stop + should abort a waiting job, unlike a load/unload which only cancels the admitted + generation so queued jobs survive a replacement. Runs under the independent + cancellation lock, so Stop stays responsive while a load holds the state lock.""" + with self._generation_cancel_lock: + active = self._active_generate_cancel is not None + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() + queued = list(self._queued_generate_cancels) + for cancel in queued: + cancel.set() + if queued: + # Wakes queued workers so they observe their cancel promptly; the wait is a + # lock-free Event, so this needs no state lock. + self._teardown_drained.set() + return active or bool(queued) # ── teardown + status ──────────────────────────────────────────────────── @@ -5786,8 +5979,7 @@ def unload(self) -> dict[str, Any]: self._load_token += 1 self._cancel_event.set() self._loading = None - if self._active_generate_cancel is not None: - self._active_generate_cancel.set() + self._cancel_active_generation_locked() # Fence generations queued behind the active one too: they hold no cancel event, so the signal cannot reach them. self._teardown_waiters += 1 # Barrier: wait for the signalled generation to exit before freeing the pipeline, else we report the VRAM free while @@ -5798,7 +5990,7 @@ def unload(self) -> dict[str, Any]: self._teardown_state_locked() finally: # Released in a finally: _teardown_state_locked ends in clear_gpu_cache(), which raises on a sticky CUDA fault, and an un-drained fence refuses every later generation. - self._teardown_waiters -= 1 + self._release_teardown_locked() logger.info("video.unloaded") return self.status() diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 557537547f6..a2d1d622092 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -14,8 +14,10 @@ import re import sys import threading +import time import types from pathlib import Path +from typing import Optional import pytest @@ -32,6 +34,7 @@ import core.inference.diffusion_eager_patches # noqa: E402,F401 import core.inference.diffusion_arch_patches # noqa: E402,F401 from core.inference.diffusion_families import ( + DIFFUSION_CANCELLED_MSG, _GATED_MIRROR_PAIRS, _MIRROR_PAIRS, _UNGATED_MIRROR_PAIRS, @@ -7140,7 +7143,8 @@ def test_download_plan_still_plans_an_unrecognised_gguf_given_an_explicit_base(m def test_unload_fences_queued_generations_while_it_waits(fake_runtime, tmp_path): - # A queued generation holds no cancel event, so unload's signal cannot reach it, and Python locks are not FIFO, so it could get in ahead and denoise after the eject. + # A queued generation is not the active denoise unload signals, and Python locks are not FIFO, + # so without the fence it could get in ahead and denoise after the eject. (tmp_path / "model.gguf").write_bytes(b"weights") backend = DiffusionBackend() backend.load_pipeline( @@ -7199,8 +7203,60 @@ def _sticky(*_args, **_kwargs): assert backend.generate(prompt = "after", steps = 2)["images"] -def test_generation_refuses_while_a_teardown_is_waiting(fake_runtime, tmp_path): - # The fence's effect: with a teardown waiting on _generate_lock, a generation that wins the lock refuses instead of denoising on a pipeline being freed. +class _RecordingCondition(threading.Condition): + def __init__( + self, + lock, + waiting: threading.Event, + waiting_again: Optional[threading.Event] = None, + ): + super().__init__(lock) + self._waiting = waiting + self._waiting_again = waiting_again + self._wait_count = 0 + + def wait(self, timeout = None): + self._wait_count += 1 + if self._wait_count == 1: + self._waiting.set() + elif self._waiting_again is not None: + self._waiting_again.set() + return super().wait(timeout) + + +class _AdmissionHookLock: + """Lock wrapper that pauses generation after atomic admission releases state.""" + + def __init__(self, backend, on_admitted): + self._lock = threading.Lock() + self._backend = backend + self._on_admitted = on_admitted + self._fired = False + + def acquire(self, *args, **kwargs): + return self._lock.acquire(*args, **kwargs) + + def release(self): + self._lock.release() + if ( + not self._fired + and threading.current_thread().name == "generation-under-test" + and self._backend._active_generate_cancel is not None + ): + self._fired = True + self._on_admitted() + + def __enter__(self): + self.acquire() + return self + + def __exit__(self, *_args): + self.release() + + +def test_generation_waits_for_all_pending_teardowns(fake_runtime, tmp_path, monkeypatch): + # A generation that wins the lock race must yield it to the queued teardown, not + # misreport a cancellation. Two reservations prove it does not resume early. (tmp_path / "model.gguf").write_bytes(b"weights") backend = DiffusionBackend() backend.load_pipeline( @@ -7211,14 +7267,264 @@ def test_generation_refuses_while_a_teardown_is_waiting(fake_runtime, tmp_path): ) assert backend.generate(prompt = "before", steps = 2)["images"] - backend._teardown_waiters = 1 - with pytest.raises(RuntimeError, match = "cancelled"): - backend.generate(prompt = "during", steps = 2) - # Still loaded: the refusal is about the pending teardown, not a missing model. - assert backend._state is not None + waiting = threading.Event() + waiting_again = threading.Event() + denoise_entered = threading.Event() + pipe_type = type(backend._state.pipe) + real_call = pipe_type.__call__ - backend._teardown_waiters = 0 - assert backend.generate(prompt = "after", steps = 2)["images"] + def record_denoise(self, *args, **kwargs): + denoise_entered.set() + return real_call(self, *args, **kwargs) + + monkeypatch.setattr(pipe_type, "__call__", record_denoise) + backend._teardown_drained = _RecordingCondition(backend._lock, waiting, waiting_again) + with backend._lock: + backend._teardown_waiters = 2 + + outcome: dict = {} + worker = threading.Thread( + target = lambda: outcome.setdefault("result", backend.generate(prompt = "during", steps = 2)), + daemon = True, + ) + worker.start() + assert waiting.wait(5), "generation did not yield to the pending teardown" + + with backend._lock: + backend._release_teardown_locked() + # Exercise the predicate against a spurious wake-up. It must go back to sleep + # while the second reservation is still live. + backend._teardown_drained.notify_all() + assert waiting_again.wait(5), "generation did not re-wait for the final teardown" + assert not denoise_entered.is_set(), "generation denoised before every teardown drained" + + with backend._lock: + backend._release_teardown_locked() + worker.join(5) + assert not worker.is_alive(), "generation did not resume after the teardown drained" + assert denoise_entered.is_set() + assert outcome["result"]["images"] + + +def test_cancel_wakes_generation_waiting_for_replacement(fake_runtime, tmp_path, monkeypatch): + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + ) + + replacement_build_started = threading.Event() + allow_replacement_commit = threading.Event() + real_from_single_file = _FakeTransformer.from_single_file + + def blocking_from_single_file(cls, path, **kwargs): + replacement_build_started.set() + assert allow_replacement_commit.wait(5), "replacement load was not released" + return real_from_single_file(path, **kwargs) + + monkeypatch.setattr( + _FakeTransformer, "from_single_file", classmethod(blocking_from_single_file) + ) + + load_outcome: dict = {} + + def replace_model(): + try: + load_outcome["result"] = backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + ) + except BaseException as exc: # noqa: BLE001 - surface worker failures in the test thread + load_outcome["error"] = exc + + loader = threading.Thread(target = replace_model, daemon = True) + loader.start() + assert replacement_build_started.wait(5), load_outcome + + outcome: dict = {} + + def generate(): + try: + backend.generate(prompt = "cancel while queued", steps = 2) + except RuntimeError as exc: + outcome["error"] = str(exc) + + worker = threading.Thread(target = generate, daemon = True) + worker.start() + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + with backend._generation_cancel_lock: + if backend._queued_generate_cancels: + break + time.sleep(0.01) + else: + allow_replacement_commit.set() + pytest.fail("generation was not published while waiting for replacement") + assert backend.generate_progress()["active"] is True + + assert backend.cancel_generate() is True + worker.join(5) + assert not worker.is_alive(), "cancelled generation waited for replacement to finish" + assert outcome["error"] == DIFFUSION_CANCELLED_MSG + assert not backend._queued_generate_cancels + assert backend.generate_progress()["active"] is False + assert loader.is_alive(), "replacement unexpectedly finished before the queued cancel" + + allow_replacement_commit.set() + loader.join(5) + assert not loader.is_alive(), "replacement load did not finish" + assert "error" not in load_outcome, load_outcome + + +def test_cancel_stops_every_generation_queued_behind_teardown(fake_runtime, tmp_path): + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + ) + with backend._lock: + backend._teardown_waiters = 1 + + errors: list[str] = [] + + def generate(prompt): + try: + backend.generate(prompt = prompt, steps = 2) + except RuntimeError as exc: + errors.append(str(exc)) + + workers = [ + threading.Thread(target = generate, args = (f"queued-{index}",), daemon = True) + for index in range(2) + ] + for worker in workers: + worker.start() + + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + with backend._generation_cancel_lock: + if len(backend._queued_generate_cancels) == 2: + break + time.sleep(0.01) + else: + pytest.fail("both generations did not register for queued cancellation") + + try: + assert backend.cancel_generate() is True + for worker in workers: + worker.join(5) + assert not worker.is_alive() + assert errors == [DIFFUSION_CANCELLED_MSG, DIFFUSION_CANCELLED_MSG] + assert not backend._queued_generate_cancels + finally: + with backend._lock: + if backend._teardown_waiters: + backend._release_teardown_locked() + + +def test_admission_registers_cancel_before_teardown_can_reserve(fake_runtime, tmp_path): + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + ) + + start_teardown = threading.Event() + teardown_reserved = threading.Event() + saw_active_cancel: list[bool] = [] + + def after_admission(): + start_teardown.set() + assert teardown_reserved.wait(5), "teardown did not reserve after admission" + + backend._lock = _AdmissionHookLock(backend, after_admission) + backend._teardown_drained = threading.Condition(backend._lock) + + def teardown(): + assert start_teardown.wait(5), "generation never reached admission" + with backend._lock: + with backend._generation_cancel_lock: + cancel = backend._active_generate_cancel + saw_active_cancel.append(cancel is not None) + if cancel is not None: + cancel.set() + backend._teardown_waiters += 1 + teardown_reserved.set() + with backend._generate_lock: + with backend._lock: + try: + backend._unload_locked() + finally: + backend._release_teardown_locked() + + teardown_worker = threading.Thread(target = teardown, daemon = True) + teardown_worker.start() + + outcome: dict = {} + + def generate(): + try: + backend.generate(prompt = "atomic admission", steps = 2) + except RuntimeError as exc: + outcome["error"] = str(exc) + + generation_worker = threading.Thread(target = generate, name = "generation-under-test", daemon = True) + generation_worker.start() + generation_worker.join(5) + teardown_worker.join(5) + + assert not generation_worker.is_alive(), "generation ignored teardown cancellation" + assert not teardown_worker.is_alive(), "teardown remained blocked behind generation" + assert saw_active_cancel == [True] + assert outcome["error"] == DIFFUSION_CANCELLED_MSG + assert backend._state is None + assert backend._teardown_waiters == 0 + + +def test_generation_reports_not_loaded_after_waiting_for_unload(fake_runtime, tmp_path): + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + ) + + waiting = threading.Event() + backend._teardown_drained = _RecordingCondition(backend._lock, waiting) + with backend._lock: + backend._teardown_waiters = 1 + + outcome: dict = {} + + def generate(): + try: + backend.generate(prompt = "during", steps = 2) + except RuntimeError as exc: + outcome["error"] = str(exc) + + worker = threading.Thread(target = generate, daemon = True) + worker.start() + assert waiting.wait(5), "generation did not wait for unload" + + with backend._lock: + backend._unload_locked() + backend._release_teardown_locked() + worker.join(5) + assert not worker.is_alive(), "generation remained blocked after unload" + assert outcome["error"] == "No diffusion model is loaded." def test_a_superseding_load_fences_queued_generations_too(fake_runtime, tmp_path): diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index c32f1059d18..dcb71314163 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -9,12 +9,14 @@ import builtins import contextlib import dataclasses +import functools import sys import threading import time import types from dataclasses import replace from pathlib import Path +from typing import Optional import pytest @@ -4604,10 +4606,13 @@ def test_each_video_load_gets_its_own_cancel_event(monkeypatch): backend = VideoBackend.__new__(VideoBackend) backend._lock = threading.RLock() backend._generate_lock = threading.RLock() + backend._generation_cancel_lock = threading.Lock() + backend._teardown_drained = threading.Event() backend._cancel_event = threading.Event() backend._load_token = 0 backend._loading = None backend._active_generate_cancel = None + backend._queued_generate_cancels = set() backend._state = None started: list[threading.Event] = [] @@ -4664,6 +4669,30 @@ def __exit__(self, *_exc) -> None: self.release() +class _ObservedLock: + """Lock wrapper that reports an acquisition attempt made outside its owner thread.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._owner = threading.current_thread() + self.waiting = threading.Event() + + def acquire(self, *args, **kwargs): + if threading.current_thread() is not self._owner: + self.waiting.set() + return self._lock.acquire(*args, **kwargs) + + def release(self) -> None: + self._lock.release() + + def __enter__(self): + self.acquire() + return self + + def __exit__(self, *_exc) -> None: + self.release() + + def _run_teardown_race(backend, teardown): """Park a generation behind ``teardown``'s _generate_lock barrier, admit it the instant that barrier releases the lock, and report what it did.""" @@ -4714,34 +4743,347 @@ def test_unload_fences_a_generation_queued_behind_its_barrier(fake_runtime, tmp_ assert backend._teardown_waiters == 0 # the fence drained -def test_superseding_load_fences_a_generation_queued_behind_its_barrier(fake_runtime, tmp_path): - # The load path takes the same barrier before tearing the old model down, so it has the same hole. +class _RecordingEvent(threading.Event): + """Event that reports each wait() so a test can observe a worker parking on it.""" + + def __init__( + self, + waiting: threading.Event, + waiting_again: Optional[threading.Event] = None, + ): + super().__init__() + self._waiting = waiting + self._waiting_again = waiting_again + self._wait_count = 0 + + def wait(self, timeout = None): + self._wait_count += 1 + if self._wait_count == 1: + self._waiting.set() + elif self._waiting_again is not None: + self._waiting_again.set() + return super().wait(timeout) + + +def test_generation_waits_for_all_pending_teardowns(fake_runtime, tmp_path, monkeypatch): + # A generation that wins the lock race must yield it to the queued teardown, not + # misreport a cancellation. Two reservations prove it does not resume early. backend = VideoBackend() _load_gguf(backend, tmp_path) + assert backend.generate(prompt = "before", steps = 2)["mp4_bytes"] == b"MP4" - queued = _run_teardown_race(backend, lambda: _load_gguf(backend, tmp_path)) + waiting = threading.Event() + waiting_again = threading.Event() + denoise_entered = threading.Event() + pipe_type = type(backend._state.pipe) + real_call = pipe_type.__call__ + + @functools.wraps(real_call) + def record_denoise(self, *args, **kwargs): + denoise_entered.set() + return real_call(self, *args, **kwargs) + + monkeypatch.setattr(pipe_type, "__call__", record_denoise) + backend._teardown_drained = _RecordingEvent(waiting, waiting_again) + with backend._lock: + backend._teardown_waiters = 2 + + outcome: dict = {} + worker = threading.Thread( + target = lambda: outcome.setdefault("result", backend.generate(prompt = "during", steps = 2)), + daemon = True, + ) + worker.start() + assert waiting.wait(5), "generation did not yield to the pending teardown" + + with backend._lock: + backend._release_teardown_locked() + backend._teardown_drained.set() + assert waiting_again.wait(5), "generation did not re-wait for the final teardown" + assert not denoise_entered.is_set(), "generation denoised before every teardown drained" + + with backend._lock: + backend._release_teardown_locked() + worker.join(5) + assert not worker.is_alive(), "generation did not resume after the teardown drained" + assert denoise_entered.is_set() + assert outcome["result"]["mp4_bytes"] == b"MP4" + + +def test_background_generation_waits_for_replacement_and_completes( + fake_runtime, tmp_path, monkeypatch +): + backend = VideoBackend() + _load_gguf(backend, tmp_path) + old_pipe = backend._state.pipe + + replacement_build_started = threading.Event() + allow_replacement_commit = threading.Event() + real_from_single_file = _FakeTransformer.from_single_file + + def blocking_from_single_file(cls, path, **kwargs): + replacement_build_started.set() + assert allow_replacement_commit.wait(5), "replacement load was not released" + return real_from_single_file(path, **kwargs) + + monkeypatch.setattr( + _FakeTransformer, "from_single_file", classmethod(blocking_from_single_file) + ) + + generated_with: list[object] = [] + real_call = _FakePipe.__call__ + + @functools.wraps(real_call) + def record_pipe(self, *args, **kwargs): + generated_with.append(self) + return real_call(self, *args, **kwargs) + + monkeypatch.setattr(_FakePipe, "__call__", record_pipe) + monkeypatch.setattr( + "core.inference.video_gallery.save", + lambda *_args, **_kwargs: {"id": "replacement-race"}, + ) + + load_outcome: dict = {} + + def replace_model(): + try: + load_outcome["result"] = _load_gguf(backend, tmp_path) + except BaseException as exc: # noqa: BLE001 - surface worker failures in the test thread + load_outcome["error"] = exc + + # Hold the barrier until both contenders are queued. The replacement reserves teardown + # first; begin_generate can still validate the old committed state, and its worker must + # yield even if Python admits it to the generation lock before the loader. + backend._generate_lock.acquire() + loader = threading.Thread(target = replace_model, daemon = True) + loader.start() + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + with backend._lock: + if backend._teardown_waiters == 1: + break + time.sleep(0.01) + else: + backend._generate_lock.release() + pytest.fail("replacement did not reserve teardown") + + backend.begin_generate(prompt = "during replacement", steps = 2) + backend._generate_lock.release() + + assert replacement_build_started.wait(5), load_outcome + assert backend.generate_progress()["active"] is True + allow_replacement_commit.set() + loader.join(5) + assert not loader.is_alive(), "replacement load did not finish" + assert "error" not in load_outcome, load_outcome + + deadline = time.monotonic() + 5 + while backend.generate_progress()["active"] and time.monotonic() < deadline: + time.sleep(0.01) + progress = backend.generate_progress() + assert progress["phase"] == "completed", progress + assert generated_with == [backend._state.pipe] + assert generated_with[0] is not old_pipe - assert ( - "out" not in queued - ), "a generation queued behind the load barrier ran against a pipeline being torn down" - assert queued.get("error") in (VIDEO_NOT_LOADED_MSG, VIDEO_CANCELLED_MSG), queued - assert backend._teardown_waiters == 0 # the fence drained +def test_queued_background_generation_revalidates_the_replacement_family( + fake_runtime, tmp_path, monkeypatch +): + import core.inference.video as video_mod -def test_generation_refuses_while_a_teardown_is_waiting(fake_runtime, tmp_path): - # The fence's effect: with a teardown waiting on _generate_lock, a generation that wins the lock refuses instead of denoising. backend = VideoBackend() _load_gguf(backend, tmp_path) - assert backend.generate(prompt = "before", steps = 2)["mp4_bytes"] == b"MP4" + old_family = backend._state.family + replacement_family = dataclasses.replace(old_family, name = "replacement-family") - backend._teardown_waiters = 1 - with pytest.raises(RuntimeError, match = "cancelled"): - backend.generate(prompt = "during", steps = 2) - # Still loaded: the refusal is about the pending teardown, not a missing model. - assert backend._state is not None + validated: list[str] = [] - backend._teardown_waiters = 0 - assert backend.generate(prompt = "after", steps = 2)["mp4_bytes"] == b"MP4" + def validate_shape(family, **_kwargs): + validated.append(family.name) + if family is replacement_family: + raise ValueError("replacement family rejects this shape") + + monkeypatch.setattr(video_mod, "validate_video_request_shape", validate_shape) + + waiting = threading.Event() + backend._teardown_drained = _RecordingEvent(waiting) + with backend._lock: + backend._teardown_waiters = 1 + + backend.begin_generate(prompt = "queued for another family", width = 768, height = 512) + assert waiting.wait(5), "background generation did not wait for replacement" + + with backend._lock: + backend._state = dataclasses.replace(backend._state, family = replacement_family) + backend._release_teardown_locked() + + deadline = time.monotonic() + 5 + while backend.generate_progress()["active"] and time.monotonic() < deadline: + time.sleep(0.01) + progress = backend.generate_progress() + assert progress["phase"] == "failed", progress + assert progress["error"] == "replacement family rejects this shape" + assert validated == [old_family.name, replacement_family.name] + + +def test_cancel_wakes_a_background_generation_waiting_for_teardown(fake_runtime, tmp_path): + backend = VideoBackend() + _load_gguf(backend, tmp_path) + + waiting = threading.Event() + backend._teardown_drained = _RecordingEvent(waiting) + with backend._lock: + backend._teardown_waiters = 1 + + try: + backend.begin_generate(prompt = "cancel while queued", steps = 2) + assert waiting.wait(5), "background generation did not wait for teardown" + assert backend.cancel_generate() is True + + deadline = time.monotonic() + 5 + while backend.generate_progress()["active"] and time.monotonic() < deadline: + time.sleep(0.01) + progress = backend.generate_progress() + assert progress["phase"] == "failed", progress + assert progress["error"] == VIDEO_CANCELLED_MSG + assert backend._teardown_waiters == 1, "test teardown drained before cancellation exited" + finally: + with backend._lock: + if backend._teardown_waiters: + backend._release_teardown_locked() + + +def test_cancel_interrupts_background_generation_waiting_for_generation_lock( + fake_runtime, tmp_path +): + backend = VideoBackend() + _load_gguf(backend, tmp_path) + + # Stand in for a replacement's final placement, which holds _generate_lock after + # teardown has drained. The queued worker must reach a terminal state while this + # lock remains held rather than waiting for placement to finish. + generate_lock = _ObservedLock() + backend._generate_lock = generate_lock + generate_lock.acquire() + try: + backend.begin_generate(prompt = "cancel during placement", steps = 2) + assert generate_lock.waiting.wait(5), "background generation did not wait for placement" + assert backend.cancel_generate() is True + + deadline = time.monotonic() + 5 + while backend.generate_progress()["active"] and time.monotonic() < deadline: + time.sleep(0.01) + progress = backend.generate_progress() + assert progress["phase"] == "failed", progress + assert progress["error"] == VIDEO_CANCELLED_MSG + finally: + generate_lock.release() + + +def test_generation_reports_not_loaded_after_waiting_for_unload(fake_runtime, tmp_path): + backend = VideoBackend() + _load_gguf(backend, tmp_path) + + waiting = threading.Event() + backend._teardown_drained = _RecordingEvent(waiting) + with backend._lock: + backend._teardown_waiters = 1 + + outcome: dict = {} + + def generate(): + try: + backend.generate(prompt = "during", steps = 2) + except RuntimeError as exc: + outcome["error"] = str(exc) + + worker = threading.Thread(target = generate, daemon = True) + worker.start() + assert waiting.wait(5), "generation did not wait for unload" + + with backend._lock: + backend._teardown_state_locked() + backend._release_teardown_locked() + worker.join(5) + assert not worker.is_alive(), "generation remained blocked after unload" + assert outcome["error"] == VIDEO_NOT_LOADED_MSG + + +def test_cancel_generate_does_not_need_the_state_lock(fake_runtime, tmp_path): + # A queued Stop must not block on _lock: cancel_generate signals under the independent + # cancellation lock and wakes the queued worker via a lock-free Event, so it works while + # a load holds the state lock for its (multi-minute) construction. + backend = VideoBackend() + _load_gguf(backend, tmp_path) + + with backend._lock: + backend._teardown_waiters = 1 + backend.begin_generate(prompt = "cancel while a load holds _lock", steps = 2) + + deadline = time.monotonic() + 5 + while not backend._queued_generate_cancels and time.monotonic() < deadline: + time.sleep(0.01) + assert backend._queued_generate_cancels, "generation never queued behind the teardown" + + try: + # The stand-in for the load's construction: _lock is HELD here, yet the Stop still + # lands and wakes the worker (neither the signal nor the wake needs the state lock). + with backend._lock: + assert backend.cancel_generate() is True + + deadline = time.monotonic() + 5 + while backend.generate_progress()["active"] and time.monotonic() < deadline: + time.sleep(0.01) + progress = backend.generate_progress() + assert progress["phase"] == "failed", progress + assert progress["error"] == VIDEO_CANCELLED_MSG + finally: + with backend._lock: + if backend._teardown_waiters: + backend._release_teardown_locked() + + +def test_generation_queued_behind_replacement_is_validated_against_the_incoming_family( + fake_runtime, tmp_path, monkeypatch +): + # Queuing behind a replacement must not bypass the input contract: with the old pipeline + # torn down but the new one not committed, the request is validated synchronously against + # the IN-FLIGHT load's family, so malformed conditioning 400s here instead of failing the + # queued job asynchronously through polling. + import core.inference.video as video_mod + + backend = VideoBackend() + _load_gguf(backend, tmp_path) + replacement_family = dataclasses.replace(backend._state.family, name = "replacement-family") + + judged: list[str] = [] + + def record_flow_shifts(fam, engine, flow_shift, audio_flow_shift): + judged.append(fam.name) + raise ValueError("replacement family rejects this flow shift") + + monkeypatch.setattr(VideoBackend, "_resolve_flow_shifts", staticmethod(record_flow_shifts)) + + with backend._lock: + backend._state = None + backend._teardown_waiters = 1 + backend._loading = video_mod._VideoLoadingState( + repo_id = "unsloth/replacement", + base_repo = "unsloth/replacement", + family = replacement_family, + engine = "diffusers", + ) + + try: + with pytest.raises(ValueError, match = "replacement family rejects this flow shift"): + backend.begin_generate(prompt = "queued", flow_shift = 1.5) + assert judged == [replacement_family.name] + assert not backend._generate_job_active, "a refused request must not reserve the job slot" + finally: + with backend._lock: + if backend._teardown_waiters: + backend._release_teardown_locked() def test_a_raising_teardown_still_drains_the_fence(fake_runtime, tmp_path, monkeypatch):