diff --git a/flashdreams/flashdreams/core/attention/kvcache.py b/flashdreams/flashdreams/core/attention/kvcache.py index 6b6ae639..5673a32a 100644 --- a/flashdreams/flashdreams/core/attention/kvcache.py +++ b/flashdreams/flashdreams/core/attention/kvcache.py @@ -36,14 +36,13 @@ class BlockKVCache: non-overlapping: each update adds one chunk of ``chunk_size`` tokens at the next logical position in the full sequence. - Note: Currently only supports ``total_size`` (``sink_size + window_size``) divisible by ``chunk_size``. - Phases: - Filling: cache not yet full; tokens are written contiguously; ``cached_k()`` / ``cached_v()`` return only the valid prefix. - - Steady-state: cache full; each new chunk triggers a left-roll of the - local window and overwrites the rightmost positions; - ``cached_k()`` / ``cached_v()`` return the full buffer. + - Steady-state: if adding a chunk would exceed the fixed cache size, the + local window rolls left by the overflow amount and the new chunk + overwrites the rightmost positions; ``cached_k()`` / ``cached_v()`` + return the full buffer. The argument ``chunk_idx`` (0, 1, 2, ...) is the index of the new chunk in the full sequence (not an index into the cache). If ``chunk_idx`` is greater than @@ -150,10 +149,6 @@ def __post_init__(self) -> None: f"k_shape[seq_dim] ({self.k_shape[self.seq_dim]}) must equal sink_size + window_size ({expected_length})" ) - assert (self.window_size + self.sink_size) % self.chunk_size == 0, ( - f"window_size + sink_size ({self.window_size + self.sink_size}) must be divisible by chunk_size ({self.chunk_size})" - ) - self._k = torch.empty(self.k_shape, device=self.device, dtype=self.dtype) self._v = torch.empty(self.v_shape, device=self.device, dtype=self.dtype) @@ -163,17 +158,19 @@ def _seq_slice(self, start: int | None, end: int | None) -> tuple[slice | int, . idx[self.seq_dim] = slice(start, end) return tuple(idx) - def _roll_local_window_left(self) -> None: - """Shift the local window left by chunk_size tokens (steady-state only).""" + def _roll_local_window_left(self, shift_size: int) -> None: + """Shift valid local-window tokens left by ``shift_size`` tokens.""" total_size = self._k.shape[self.seq_dim] - assert total_size == self._n_cached, ( - f"Expected full cache: {total_size=} != {self._n_cached=}" + assert 0 < shift_size <= self.chunk_size, ( + f"shift_size ({shift_size}) must be in (0, {self.chunk_size}]" ) - tokens_to_keep = self.window_size - self.chunk_size + valid_end = min(self._n_cached, total_size) + valid_local_size = max(0, valid_end - self.sink_size) + tokens_to_keep = max(0, valid_local_size - shift_size) if tokens_to_keep > 0: - src_start = self.sink_size + self.chunk_size - src_end = total_size + src_start = self.sink_size + shift_size + src_end = src_start + tokens_to_keep dst_start = self.sink_size dst_end = self.sink_size + tokens_to_keep @@ -181,6 +178,10 @@ def _roll_local_window_left(self) -> None: src_slice = self._seq_slice(src_start, src_end) self._k[dst_slice] = self._k[src_slice].clone() self._v[dst_slice] = self._v[src_slice].clone() + # before_update() only rolls when the next contiguous chunk would overflow + # this fixed buffer; update() must immediately write that chunk into the + # newly freed right edge. + self._n_cached = total_size def _current_chunk_overlaps_sink(self) -> bool: assert self._curr_chunk_idx is not None, ( @@ -286,8 +287,11 @@ def before_update(self, chunk_idx: int) -> None: "Expected the new chunk_idx to be +1 from the previous chunk_idx, " f"got {chunk_idx} != {self._prev_chunk_idx} + 1" ) - if self.is_steady_state(): - self._roll_local_window_left() + total_size = self._k.shape[self.seq_dim] + if not self._current_chunk_overlaps_sink(): + overflow = self._n_cached + self.chunk_size - total_size + if overflow > 0: + self._roll_local_window_left(overflow) def update(self, k: Tensor, v: Tensor) -> None: """ @@ -332,7 +336,8 @@ def after_update(self, chunk_idx: int) -> None: if self.is_steady_state(): pass else: - self._n_cached += self.chunk_size + total_size = self._k.shape[self.seq_dim] + self._n_cached = min(self._n_cached + self.chunk_size, total_size) self._prev_chunk_idx += 1 elif self._curr_chunk_idx == self._prev_chunk_idx: pass diff --git a/flashdreams/flashdreams/serving/realtime/input.py b/flashdreams/flashdreams/serving/realtime/input.py index 1b488bb6..6baf92dd 100644 --- a/flashdreams/flashdreams/serving/realtime/input.py +++ b/flashdreams/flashdreams/serving/realtime/input.py @@ -135,21 +135,21 @@ class KeyboardResampler: def __init__( self, *, - fps: int, + fps: float, start_v: float = 0.0, supported_keys: frozenset[str] = DEFAULT_SUPPORTED_KEYS, ) -> None: if fps <= 0: raise ValueError("fps must be > 0") - self._fps = fps - self._dt = 1.0 / fps + self._fps = float(fps) + self._dt = 1.0 / self._fps self._supported_keys = supported_keys self.next_chunk_start_v = start_v self._event_log: deque[tuple[float, dict[str, str]]] = deque() self._carried_state = KeyboardState(supported_keys=supported_keys) @property - def fps(self) -> int: + def fps(self) -> float: return self._fps @property diff --git a/flashdreams/flashdreams/serving/realtime/media.py b/flashdreams/flashdreams/serving/realtime/media.py index 266e7ac5..9b93eaca 100644 --- a/flashdreams/flashdreams/serving/realtime/media.py +++ b/flashdreams/flashdreams/serving/realtime/media.py @@ -141,6 +141,18 @@ def tensor_chunk_to_rgb_frames( "minus_one_one" if video_chunk.is_floating_point() else "uint8" ) if video_chunk.ndim == 4: + if video_chunk.shape[-1] == 3 and video_chunk.shape[1] != 3: + return rgb_array_to_uint8_frames( + video_chunk, + layout="thwc", + value_range=value_range, + sync_device=sync_device, + ) + if video_chunk.shape[1] != 3: + raise ValueError( + "Expected video chunk [T, C, H, W] or [T, H, W, C] with 3 RGB " + f"channels, got {tuple(video_chunk.shape)}" + ) return rgb_array_to_uint8_frames( video_chunk, layout="tchw", @@ -155,7 +167,8 @@ def tensor_chunk_to_rgb_frames( sync_device=sync_device, ) raise ValueError( - "Expected video chunk [T, C, H, W] or [1, 1, T, 3, H, W], " + "Expected video chunk [T, C, H, W], [T, H, W, C], " + "or [1, 1, T, 3, H, W], " f"got {tuple(video_chunk.shape)}" ) diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index d0b7d896..6ad20743 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -53,11 +53,27 @@ # How often the liveness watchdog wakes to re-check the elapsed-since-last-message. _CLIENT_LIVENESS_CHECK_INTERVAL_S = 1.0 +_DEFAULT_PERF_LOG_INTERVAL_CHUNKS = 5 _RuntimeT = TypeVar("_RuntimeT", bound=WebRTCSessionRuntime) _RuntimeConfigT = TypeVar("_RuntimeConfigT", bound=WebRTCRuntimeConfig) +def _stat_float(stats: dict[str, float], name: str, default: float = 0.0) -> float: + value = stats.get(name) + if value is None: + return default + return float(value) + + +def _stat_ms(stats: dict[str, float], name: str, default_ms: float = 0.0) -> float: + return _stat_float(stats, name, default_ms / 1e3) * 1e3 + + +def _stat_int(stats: dict[str, float], name: str) -> int: + return int(round(_stat_float(stats, name))) + + class WebRTCControlSignal(IntEnum): """Rank-orchestration signals shared by the single-session runtimes.""" @@ -123,6 +139,7 @@ class BaseWebRTCSessionManager(Generic[_RuntimeT, _RuntimeConfigT]): _runtime_error_types: tuple[type[Exception], ...] = (RuntimeError,) _close_session_on_generation_error: bool = False _resampler_supported_keys: AbstractSet[str] | None = None + _perf_log_interval_chunks: int = _DEFAULT_PERF_LOG_INTERVAL_CHUNKS def __init__( self, @@ -160,14 +177,72 @@ async def _reset_runtime_for_session(self, session_input: Any) -> None: await self._runtime.reset_for_new_session() def _make_resampler(self, *, start_v: float) -> KeyboardResampler: + return self._make_resampler_at_fps(start_v=start_v, fps=self.fps) + + def _make_resampler_at_fps( + self, *, start_v: float, fps: float + ) -> KeyboardResampler: if self._resampler_supported_keys is None: - return KeyboardResampler(fps=self.fps, start_v=start_v) + return KeyboardResampler(fps=fps, start_v=start_v) return KeyboardResampler( - fps=self.fps, + fps=fps, start_v=start_v, supported_keys=frozenset(self._resampler_supported_keys), ) + @staticmethod + def _positive_int_runtime_value(value: Any, *, label: str) -> int: + try: + parsed = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{label} must be an integer.") from exc + if parsed <= 0: + raise ValueError(f"{label} must be > 0.") + return parsed + + @staticmethod + def _positive_float_runtime_value(value: Any, *, label: str) -> float: + try: + parsed = float(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{label} must be numeric.") from exc + if parsed <= 0.0: + raise ValueError(f"{label} must be > 0.") + return parsed + + def _runtime_input_fps(self, runtime: Any) -> float: + method = getattr(runtime, "peek_input_fps", None) + if callable(method): + return self._positive_float_runtime_value( + method(), + label="peek_input_fps", + ) + return float(self.fps) + + def _runtime_next_input_num_frames(self, runtime: Any) -> int: + method = getattr(runtime, "peek_next_input_num_frames", None) + if callable(method): + return self._positive_int_runtime_value( + method(), + label="peek_next_input_num_frames", + ) + return self._positive_int_runtime_value( + runtime.peek_next_chunk_num_frames(), + label="peek_next_chunk_num_frames", + ) + + def _runtime_steady_output_num_frames(self, runtime: Any) -> int: + method = getattr(runtime, "peek_steady_output_num_frames", None) + if callable(method): + return self._positive_int_runtime_value( + method(), + label="peek_steady_output_num_frames", + ) + return self._positive_int_runtime_value( + runtime.peek_steady_chunk_num_frames(), + label="peek_steady_chunk_num_frames", + ) + def _register_extra_peer_handlers(self, peer_connection: Any) -> None: """Register optional extra peer-connection event handlers.""" @@ -289,13 +364,16 @@ async def _create_answer_with_runtime_ready_locked( # is throttled to the consumer's drain rate. AR step 0 emits fewer # frames than steady state; sizing to it would force a per-chunk # stall, so we size to the steady-state count. - num_frames = self._runtime.peek_steady_chunk_num_frames() + num_frames = self._runtime_steady_output_num_frames(self._runtime) video_track = BufferedVideoTrack(fps=self.fps, maxsize=num_frames) peer_connection.addTrack(video_track) # Start the resampler's virtual clock at 0; the real anchor is set # in the ``on_datachannel`` handler so chunk 0's window starts when # input can actually arrive. - resampler = self._make_resampler(start_v=0.0) + resampler = self._make_resampler_at_fps( + start_v=0.0, + fps=self._runtime_input_fps(self._runtime), + ) loop = asyncio.get_running_loop() managed_session = ManagedWebRTCSession( runtime=self._runtime, @@ -561,15 +639,19 @@ async def _generation_worker( "First action received; starting generation at start_v={:.3f}", resampler.next_chunk_start_v, ) + perf_log_interval = max(0, int(self._perf_log_interval_chunks)) + perf_window_start = loop.time() + perf_window_chunks = 0 + perf_window_frames = 0 try: while not managed_session.closed: try: - num_frames = runtime.peek_next_chunk_num_frames() + input_num_frames = self._runtime_next_input_num_frames(runtime) except self._runtime_error_types: logger.exception("Runtime not ready; stopping generation worker.") return # Trigger when wallclock reaches the chunk's window end. - chunk_duration = num_frames * resampler.dt + chunk_duration = input_num_frames * resampler.dt trigger_wall = resampler.next_chunk_start_v + chunk_duration delay = trigger_wall - loop.time() if delay > 0: @@ -588,7 +670,7 @@ async def _generation_worker( resampler.next_chunk_start_v = now - chunk_duration t_before_gen = loop.time() - segments, frame_times = resampler.sample_chunk(num_frames) + segments, frame_times = resampler.sample_chunk(input_num_frames) chunk_end_v = resampler.next_chunk_start_v consumed_action_arrivals: list[float] = [] while ( @@ -624,11 +706,65 @@ async def _generation_worker( if consumed_action_arrivals else None ) + perf_window_chunks += 1 + perf_window_frames += result.num_frames + if result.chunk_index == 0 or ( + perf_log_interval > 0 + and result.chunk_index % perf_log_interval == 0 + ): + interval_s = max(t_after_enqueue - perf_window_start, 1.0e-6) + interval_fps = perf_window_frames / interval_s + gen_fps = result.num_frames / max( + t_after_gen - t_before_gen, 1.0e-6 + ) + stats = result.stats or {} + logger.info( + "WebRTC perf chunk={} interval_chunks={} frames={} " + "gen_fps={:.1f} interval_fps={:.1f} playback_fps={} " + "gen_ms={:.0f} enqueue_ms={:.0f} model_ms={:.0f} " + "denoise_ms={:.0f} decode_ms={:.0f} pixel_post_ms={:.0f} " + "copy_ms={:.0f} cache_ms={:.0f} " + "cache_wait_ms={:.0f} cache_submit_ms={:.0f} " + "queue_depth={} lag_ms={:.0f} control_latency_ms={} " + "compile_active={} compile_start_step={} cuda_graph={} " + "cache_frames={} cache_tokens={}", + result.chunk_index, + perf_window_chunks, + perf_window_frames, + gen_fps, + interval_fps, + video_track.fps, + gen_ms, + enqueue_ms, + _stat_ms(stats, "model_step_s", gen_ms), + _stat_ms(stats, "denoise_s"), + _stat_ms(stats, "decode_s"), + _stat_ms(stats, "pixel_post_s"), + _stat_ms(stats, "gpu_to_cpu_copy_s"), + _stat_ms(stats, "cache_seed_prune_s"), + _stat_ms(stats, "cache_update_wait_s"), + _stat_ms(stats, "cache_update_submit_s"), + video_track.qsize(), + lag_ms, + "-" + if control_latency_ms is None + else f"{control_latency_ms:.0f}", + _stat_int(stats, "compile_denoise_active"), + _stat_int(stats, "compile_denoise_start_step"), + _stat_int(stats, "cuda_graph_captured"), + _stat_int(stats, "cache_frames"), + _stat_int(stats, "cache_tokens"), + ) + perf_window_start = t_after_enqueue + perf_window_chunks = 0 + perf_window_frames = 0 logger.debug( - "Chunk done chunk={} num_frames={} segments={} enqueued={} " + "Chunk done chunk={} input_frames={} output_frames={} " + "segments={} enqueued={} " "gen_ms={:.1f} enqueue_ms={:.1f} play_ms={:.1f} queue_depth={} " "lag_ms={:.1f}", result.chunk_index, + input_num_frames, result.num_frames, len(segments), enqueued, diff --git a/flashdreams/flashdreams/serving/webrtc/runtime.py b/flashdreams/flashdreams/serving/webrtc/runtime.py index d9cf0d03..55c1b6a5 100644 --- a/flashdreams/flashdreams/serving/webrtc/runtime.py +++ b/flashdreams/flashdreams/serving/webrtc/runtime.py @@ -19,6 +19,7 @@ class WebRTCStepResult: """One generated chunk handed back by a WebRTC model runtime.""" chunk_index: int + # Number of output video frames in ``video_chunk``. num_frames: int video_chunk: torch.Tensor stats: dict[str, float] | None @@ -39,6 +40,15 @@ class WebRTCGenerationRuntime(Protocol): Integrations keep their model-specific state, checkpoints, conditioning, and cache logic inside their concrete runtime. The shared manager only needs this lifecycle and chunk-generation surface. + + By default, ``peek_next_chunk_num_frames`` and + ``peek_steady_chunk_num_frames`` are used for both input sampling and + output queue sizing. Runtimes whose model input clock differs from their + output video clock may also implement these optional methods: + + - ``peek_input_fps() -> float`` for the control/input sampling clock. + - ``peek_next_input_num_frames() -> int`` for the length of ``frame_times``. + - ``peek_steady_output_num_frames() -> int`` for video queue sizing. """ async def initialize(self) -> None: ... diff --git a/flashdreams/tests/test_kvcache.py b/flashdreams/tests/test_kvcache.py index 6e0ee738..90042e33 100644 --- a/flashdreams/tests/test_kvcache.py +++ b/flashdreams/tests/test_kvcache.py @@ -91,20 +91,16 @@ def update(self, chunk_idx: int, k: torch.Tensor, v: torch.Tensor) -> None: self._cache_v[:, -length:] = v return - if self._cache_k.shape[1] == self.total_size: - tail_k = self._cache_k[:, self.sink_size + self.chunk_size :] - tail_v = self._cache_v[:, self.sink_size + self.chunk_size :] - window_k = torch.cat([tail_k, k], dim=1)[:, -self.window_size :] - window_v = torch.cat([tail_v, v], dim=1)[:, -self.window_size :] - self._cache_k = torch.cat( - [self._cache_k[:, : self.sink_size], window_k], dim=1 - ) - self._cache_v = torch.cat( - [self._cache_v[:, : self.sink_size], window_v], dim=1 - ) - else: - self._cache_k = torch.cat([self._cache_k, k], dim=1) - self._cache_v = torch.cat([self._cache_v, v], dim=1) + sink_k = self._cache_k[:, : self.sink_size] + sink_v = self._cache_v[:, : self.sink_size] + window_k = torch.cat([self._cache_k[:, self.sink_size :], k], dim=1)[ + :, -self.window_size : + ] + window_v = torch.cat([self._cache_v[:, self.sink_size :], v], dim=1)[ + :, -self.window_size : + ] + self._cache_k = torch.cat([sink_k, window_k], dim=1) + self._cache_v = torch.cat([sink_v, window_v], dim=1) self._prev_chunk_idx += 1 def cached_k(self) -> torch.Tensor: @@ -176,7 +172,9 @@ def dtype() -> torch.dtype: @pytest.mark.ci_cpu -@pytest.mark.parametrize("sink_size,window_size", [(0, 8), (0, 24), (3, 5), (3, 21)]) +@pytest.mark.parametrize( + "sink_size,window_size", [(0, 8), (0, 24), (3, 5), (3, 21), (1, 16)] +) def test_block_kvcache_matches_baseline( device: torch.device, dtype: torch.dtype, @@ -350,7 +348,9 @@ def test_kvcache_relative_rope_cp_freqs_match_cache_chunks() -> None: not torch.cuda.is_available(), reason="CUDA required for cudagraph test" ) @pytest.mark.ci_gpu -@pytest.mark.parametrize("sink_size,window_size", [(0, 8), (0, 24), (3, 5), (3, 21)]) +@pytest.mark.parametrize( + "sink_size,window_size", [(0, 8), (0, 24), (3, 5), (3, 21), (1, 16)] +) def test_block_kvcache_cudagraph_matches_baseline( dtype: torch.dtype, sink_size: int, diff --git a/flashdreams/tests/test_realtime_serving.py b/flashdreams/tests/test_realtime_serving.py index 5c89d6c3..33b21476 100644 --- a/flashdreams/tests/test_realtime_serving.py +++ b/flashdreams/tests/test_realtime_serving.py @@ -144,6 +144,29 @@ def test_realtime_media_supports_omnidreams_uint8_layout() -> None: assert frames[1][0, 0, 0] == 255 +def test_realtime_media_supports_thwc_uint8_chunks() -> None: + chunk = torch.zeros((2, 4, 5, 3), dtype=torch.uint8) + chunk[1, 0, 0] = torch.tensor([10, 20, 30], dtype=torch.uint8) + + frames = tensor_chunk_to_rgb_frames(chunk) + + assert len(frames) == 2 + assert frames[0].shape == (4, 5, 3) + assert frames[0].dtype == np.uint8 + assert frames[1][0, 0].tolist() == [10, 20, 30] + + +def test_realtime_media_prefers_tchw_when_width_is_three() -> None: + chunk = torch.zeros((2, 3, 4, 3), dtype=torch.uint8) + chunk[1, :, 0, 0] = torch.tensor([10, 20, 30], dtype=torch.uint8) + + frames = tensor_chunk_to_rgb_frames(chunk) + + assert len(frames) == 2 + assert frames[1].shape == (4, 3, 3) + assert frames[1][0, 0].tolist() == [10, 20, 30] + + @pytest.mark.parametrize( "chunk", [ diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index 69bb05ce..ebcafdbb 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -11,6 +11,7 @@ import pytest import torch +from flashdreams.serving.webrtc import manager as manager_module from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, @@ -75,6 +76,11 @@ def sample_chunk( return [(0.0, 0.0, frozenset({"w"}))], [0.0] +class _CountingVideoTrack(_FakeVideoTrack): + async def enqueue_chunk(self, chunk: Any) -> int: + return int(chunk.shape[0]) + + class _BaseTestManager(BaseWebRTCSessionManager): def _model_name(self) -> str: return "fake-model" @@ -90,6 +96,45 @@ def _make_manager( ) +def test_runtime_frame_timing_hooks_default_to_legacy_methods() -> None: + class _LegacyRuntime: + def peek_next_chunk_num_frames(self) -> int: + return 2 + + def peek_steady_chunk_num_frames(self) -> int: + return 3 + + runtime = _LegacyRuntime() + manager = _make_manager(_BaseTestManager, runtime) + + assert manager._runtime_input_fps(runtime) == pytest.approx(30.0) + assert manager._runtime_next_input_num_frames(runtime) == 2 + assert manager._runtime_steady_output_num_frames(runtime) == 3 + + +def test_runtime_frame_timing_hooks_can_split_input_and_output() -> None: + class _SplitRuntime: + def peek_input_fps(self) -> float: + return 6.0 + + def peek_next_input_num_frames(self) -> int: + return 4 + + def peek_steady_output_num_frames(self) -> int: + return 16 + + runtime = _SplitRuntime() + manager = _make_manager(_BaseTestManager, runtime) + resampler = manager._make_resampler_at_fps( + start_v=0.0, + fps=manager._runtime_input_fps(runtime), + ) + + assert resampler.dt == pytest.approx(1.0 / 6.0) + assert manager._runtime_next_input_num_frames(runtime) == 4 + assert manager._runtime_steady_output_num_frames(runtime) == 16 + + def _managed_session( runtime: Any, ) -> tuple[ManagedWebRTCSession, _FakeVideoTrack, _FakePeerConnection, _FakeChannel]: @@ -231,6 +276,139 @@ def _chunk_done_extra(self) -> dict[str, Any]: assert payload["resolution"] == {"width": 8, "height": 4} +@pytest.mark.asyncio +async def test_generation_worker_uses_split_input_and_output_frame_counts() -> None: + class _SplitResampler: + dt = 0.0 + next_chunk_start_v = 0.0 + + def __init__(self) -> None: + self.sampled_num_frames: list[int] = [] + + def sample_chunk( + self, num_frames: int + ) -> tuple[list[tuple[float, float, frozenset[str]]], list[float]]: + self.sampled_num_frames.append(num_frames) + return ( + [(0.0, 0.0, frozenset({"w"}))], + [float(index) for index in range(num_frames)], + ) + + class _SplitRuntime: + def __init__(self) -> None: + self.managed_session: ManagedWebRTCSession | None = None + self.frame_times: list[float] | None = None + + def peek_input_fps(self) -> float: + return 6.0 + + def peek_next_input_num_frames(self) -> int: + return 2 + + async def generate_chunk( + self, *, segments: Any, frame_times: list[float] + ) -> WebRTCStepResult: + del segments + self.frame_times = frame_times + if self.managed_session is not None: + self.managed_session.closed = True + return WebRTCStepResult( + chunk_index=0, + num_frames=5, + video_chunk=torch.zeros((5, 1, 1, 3, 2, 2), dtype=torch.uint8), + stats=None, + ) + + runtime = _SplitRuntime() + manager = _make_manager(_BaseTestManager, runtime) + managed, _video_track, _peer, channel = _managed_session(runtime) + resampler = _SplitResampler() + managed.video_track = _CountingVideoTrack() # ty:ignore[invalid-assignment] + managed.resampler = resampler # ty:ignore[invalid-assignment] + runtime.managed_session = managed + manager._active_session = managed + + task = asyncio.create_task(manager._generation_worker(managed_session=managed)) + managed.generation_task = task + await asyncio.wait_for(task, timeout=5.0) + + assert resampler.sampled_num_frames == [2] + assert runtime.frame_times == [0.0, 1.0] + chunk_done = [ + json.loads(message) + for message in channel.messages + if json.loads(message).get("type") == "chunk_done" + ] + assert len(chunk_done) == 1 + assert chunk_done[0]["num_frames"] == 5 + assert chunk_done[0]["enqueued_frames"] == 5 + assert chunk_done[0]["play_ms"] == pytest.approx(5 * 1000 / 30, abs=0.05) + + +@pytest.mark.asyncio +async def test_generation_worker_logs_periodic_perf_stats( + monkeypatch: pytest.MonkeyPatch, +) -> None: + perf_logs: list[tuple[str, tuple[Any, ...]]] = [] + + def _record_info(message: str, *args: Any, **_kwargs: Any) -> None: + if message.startswith("WebRTC perf"): + perf_logs.append((message, args)) + + class _StatsRuntime: + def __init__(self) -> None: + self.managed_session: ManagedWebRTCSession | None = None + self.chunk_index = 0 + + def peek_next_chunk_num_frames(self) -> int: + return 1 + + async def generate_chunk( + self, *, segments: Any, frame_times: Any + ) -> WebRTCStepResult: + del segments, frame_times + chunk_index = self.chunk_index + self.chunk_index += 1 + if chunk_index >= 2 and self.managed_session is not None: + self.managed_session.closed = True + return WebRTCStepResult( + chunk_index=chunk_index, + num_frames=4, + video_chunk=torch.zeros((4, 1, 1, 3, 2, 2), dtype=torch.uint8), + stats={ + "model_step_s": 0.02, + "denoise_s": 0.01, + "decode_s": 0.004, + "pixel_post_s": 0.003, + "gpu_to_cpu_copy_s": 0.002, + "compile_denoise_active": 1.0, + "compile_denoise_start_step": 3.0, + "cache_frames": 13.0, + "cache_tokens": 512.0, + }, + ) + + class _FrequentLogManager(_BaseTestManager): + _perf_log_interval_chunks = 2 + + monkeypatch.setattr(manager_module.logger, "info", _record_info) + runtime = _StatsRuntime() + manager = _make_manager(_FrequentLogManager, runtime) + managed, _video_track, _peer, _channel = _managed_session(runtime) + runtime.managed_session = managed + manager._active_session = managed + + task = asyncio.create_task(manager._generation_worker(managed_session=managed)) + managed.generation_task = task + await asyncio.wait_for(task, timeout=5.0) + + assert [args[0] for _message, args in perf_logs] == [0, 2] + assert "compile_active" in perf_logs[0][0] + assert "pixel_post_ms" in perf_logs[0][0] + assert "copy_ms" in perf_logs[0][0] + assert perf_logs[0][1][-2:] == (13, 512) + + @pytest.mark.asyncio async def test_create_answer_raises_busy_with_subclass_message() -> None: class _BusyManager(_BaseTestManager):