diff --git a/invokeai/app/util/video_thumbnails.py b/invokeai/app/util/video_thumbnails.py index 01416dcc311..678e0847373 100644 --- a/invokeai/app/util/video_thumbnails.py +++ b/invokeai/app/util/video_thumbnails.py @@ -224,8 +224,15 @@ def _iter_video_frames_unbounded( video_path: Path, timeout: float = VIDEO_DECODE_TIMEOUT_SECONDS, is_canceled: Optional[Callable[[], bool]] = None, + first_frame_timeout: Optional[float] = None, ) -> Iterator[np.ndarray]: - """Streams decoded frames from an isolated worker with bounded memory and wait time.""" + """Streams decoded frames from an isolated worker with bounded memory and wait time. + + ``timeout`` bounds decoder *inactivity*: it is restarted after every frame, so a long + video is not killed for being long. ``first_frame_timeout`` overrides that budget for + the first frame only, letting a caller that already spent part of the budget waiting + for capacity charge that wait against the same deadline instead of granting a fresh one. + """ proc = _spawn_worker( "stream", str(video_path), @@ -286,7 +293,7 @@ def read_frames() -> None: stderr_reader = threading.Thread(target=drain_stderr, name="video-stderr-reader", daemon=True) reader.start() stderr_reader.start() - deadline = time.monotonic() + timeout + deadline = time.monotonic() + (timeout if first_frame_timeout is None else first_frame_timeout) try: while True: if is_canceled is not None and is_canceled(): @@ -349,7 +356,18 @@ def iter_video_frames( if slot.acquire(timeout=min(0.1, remaining)): acquired.append(slot) break - yield from _iter_video_frames_unbounded(video_path, timeout, is_canceled) + # Charge the capacity wait against the same deadline as the first frame, the way + # _run_worker does. Handing the decoder a fresh full timeout here would let a + # caller that waited just under `timeout` for a slot block for nearly 2 * timeout + # before failing — twice the bound the callers (upload probing, node decodes) + # believe they are enforcing. Later frames still get a full `timeout` each: after + # the first frame the budget is an inactivity bound, not a queueing one. + yield from _iter_video_frames_unbounded( + video_path, + timeout, + is_canceled, + first_frame_timeout=max(0.0, capacity_deadline - time.monotonic()), + ) finally: for slot in reversed(acquired): slot.release() diff --git a/tests/app/util/test_video_thumbnails.py b/tests/app/util/test_video_thumbnails.py index 494a36fd4ac..a8fb7272770 100644 --- a/tests/app/util/test_video_thumbnails.py +++ b/tests/app/util/test_video_thumbnails.py @@ -279,6 +279,56 @@ def consume() -> None: assert len(errors) == 1 assert isinstance(errors[0], TimeoutError) + def test_capacity_wait_and_first_frame_share_one_deadline(self, hanging_worker, tmp_path: Path) -> None: + """Waiting for capacity and waiting for the first frame draw on ONE budget. + + Each used to get a full ``timeout``: a caller that waited nearly the whole budget + for a stream slot then got a fresh full budget on the hung decoder, so the call + could take ~2x the bound it advertised. Later frames still get a full timeout + each — after the first frame the budget is an inactivity bound, not a queueing one. + """ + target = tmp_path / "malicious.mp4" + target.write_bytes(b"pretend this hangs the decoder") + timeout = 2.0 + # Occupy the single stream slot for most (but not all) of the budget, so the + # decode still starts and the two waits are distinguishable in the total. + hold = 1.6 + errors: list[BaseException] = [] + finished = Event() + assert video_thumbnails._VIDEO_STREAM_SLOTS.acquire(timeout=1) + released = False + + def consume() -> None: + try: + next(iter_video_frames(target, timeout=timeout)) + except BaseException as error: + errors.append(error) + finally: + finished.set() + + started = time.monotonic() + thread = threading.Thread(target=consume) + thread.start() + try: + time.sleep(hold) + video_thumbnails._VIDEO_STREAM_SLOTS.release() + released = True + assert finished.wait(timeout=10), "decoder never gave up" + finally: + if not released: + video_thumbnails._VIDEO_STREAM_SLOTS.release() + thread.join(timeout=10) + elapsed = time.monotonic() - started + + assert len(errors) == 1 + # "Timed out decoding" (not "Timed out waiting to decode") — the slot was acquired + # and the worker was spawned, so this is the first-frame wait expiring. + assert isinstance(errors[0], TimeoutError) + assert "Timed out decoding" in str(errors[0]) + # Midway between the fixed behavior (~timeout) and the old one (hold + timeout), + # leaving room for worker-spawn and scheduling jitter on CI. + assert elapsed < timeout + hold / 2, f"first frame got a fresh timeout after the capacity wait ({elapsed:.2f}s)" + def test_probe_timeout_includes_waiting_for_decoder_capacity(self, tmp_path: Path) -> None: target = tmp_path / "never-opened.mp4" errors: list[BaseException] = []