diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2895bb895..e11384925 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -356,6 +356,12 @@ The nested `integrations_v2/omnidreams/impl/ludus-renderer` workspace package is named `ludus-renderer` and is installed as part of Omnidreams workflows that need it. +## Benchmarking + +Trigger benchmarking via the demo argument (before `--`): `--stats-path .json`. + +`MetricsOutputSink::close` is currently where the schema of our "standardized" benchmark artifact is defined. + ## Licensing of contributions By submitting a pull request to this repository, you agree that your diff --git a/apps/t2v/t2v/testing.py b/apps/t2v/t2v/testing.py deleted file mode 100644 index d59f7fccc..000000000 --- a/apps/t2v/t2v/testing.py +++ /dev/null @@ -1,445 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Test support for text-to-video integrations, shipped for them to import. - -Nothing here runs in production. It is the shared check an integration's tests -call, and the stand-in model they run it against, named as ``numpy.testing`` and -``torch.testing`` are. -""" - -import os -import shutil -from collections.abc import Sequence -from dataclasses import dataclass, field, replace -from pathlib import Path -from typing import Any -from unittest.mock import patch - -import numpy as np -import numpy.typing as npt -import torch - -from flashdreams.api_v2.client_window import IClientWindow -from flashdreams.api_v2.output_sink import OutputSink -from flashdreams.runtime_v2.blit_model_output_to_screen_loop import ( - BlitModelOutputToScreenLoop, -) -from flashdreams.runtime_v2.mp4_output_sink import Mp4OutputSink -from flashdreams.runtime_v2.session_desc import ( - BackpressureMode, - PresentationMode, - SessionDesc, -) -from flashdreams.runtime_v2.session_runner import run_session -from flashdreams.runtime_v2.step_result import StepResult -from flashdreams.runtime_v2.user_input_events import UserInputEvents -from flashdreams.runtime_v2.video_encoder import result_to_rgb24_frames -from t2v.application import T2VApplication -from t2v.ui import T2VImGuiUILoop - -_REAL_CLIP_LUMINANCE = (16.0, 240.0) -"""Mean pixel value a real clip has to land inside, from ``0`` to ``255``. - -Loose on purpose: what it catches is a run that came back blank. -""" - -_REAL_CLIP_FRAME_DIFFERENCE = 0.5 -"""How much consecutive frames of a real clip have to differ, on that scale.""" - - -@dataclass(frozen=True, kw_only=True, slots=True) -class ExpectedFrameStats: - """What a caller expects a run to have generated. - - Every field is optional, and one left out is not checked. A model that - samples cannot be expected to produce a particular picture, but it can be - expected to produce a picture at all. - """ - - frame_count: int | None = None - """Frames the whole run should generate.""" - - mean_luminance: tuple[float, float] | None = None - """Range the mean pixel value should land in, from ``0`` to ``255``.""" - - min_frame_difference: float | None = None - """Smallest mean change from one frame to the next, on that scale.""" - - -@dataclass(frozen=True, kw_only=True, slots=True) -class T2VCheckResult: - """What a run generated, and how it compared to what was expected.""" - - failures: tuple[str, ...] - """Expectations the run did not meet, in the order they were checked.""" - - frames_per_step: tuple[int, ...] - """Frames each step generated. A model whose first chunk is a different - size to the rest shows up here.""" - - mean_luminance: float - """Mean pixel value over every frame, from ``0`` to ``255``.""" - - frame_difference: float - """Mean change from one frame to the next, over the whole run.""" - - metrics: tuple[dict[str, float | int], ...] = field(default_factory=tuple) - """Whatever each step reported, such as generation timings.""" - - mp4_path: Path | None = None - """File written, when one was asked for.""" - - @property - def passed(self) -> bool: - """Whether the run met every expectation it was given.""" - return not self.failures - - @property - def frame_count(self) -> int: - """Frames the whole run generated.""" - return sum(self.frames_per_step) - - -def check_t2v_model_impl( - application: T2VApplication, - session_desc: SessionDesc | None = None, - *, - steps: int, - expected: ExpectedFrameStats, - commandline_args: Sequence[str] = (), - mp4_path: str | Path | None = None, -) -> T2VCheckResult: - """Run a text-to-video application for ``steps`` steps and inspect the video. - - The coverage an integration gets from one call: the application loads, - resolves a session, generates, and what it generated is a video rather than - a run that merely finished. It is initialized and closed here, and frames - are read the way a sink reads them. Blocking backpressure and new-frame-only - presentation keep the inspected frames identical to model output. - - Args: - application: Uninitialized application to run. - session_desc: Session to ask for, or ``None`` to take the application's - own. A stand-in generating some other size says so here. - steps: Steps to generate. Enough to reach steady state, since a model - whose first chunk differs is only interesting from the second. - expected: What the generated video should look like. - commandline_args: Application arguments, such as a prompt. - mp4_path: File to write as well, for a person to watch. - - Returns: - What the run generated, and which expectations it missed. - - Raises: - Whatever the run raises. A model that fails to load or generate is a - failure of the integration rather than of an expectation. - """ - if steps <= 0: - raise ValueError(f"steps must be > 0, got {steps}.") - - inspector = _FrameInspector(Mp4OutputSink(mp4_path) if mp4_path else None) - application.init(commandline_args) - try: - if session_desc is None: - session_desc = application.session_desc() - session_desc = replace( - session_desc, - backpressure_mode=BackpressureMode.BLOCK, - presentation_mode=PresentationMode.ON_DEMAND, - ) - with patch.multiple( - T2VImGuiUILoop, - _initialize_loop_state=BlitModelOutputToScreenLoop._initialize_loop_state, - step=BlitModelOutputToScreenLoop.step, - is_finished=BlitModelOutputToScreenLoop.is_finished, - reset=BlitModelOutputToScreenLoop.reset, - ): - run_session( - application.create_session(session_desc), - _DiscardingClientWindow(), - metrics_output_sink=inspector, - steps=steps, - ) - finally: - application.close() - - return _compare(inspector, expected, Path(mp4_path) if mp4_path else None) - - -def real_model_run_skip_reason(run_env: str) -> str | None: - """Return why the real model cannot be run here, or ``None`` if it can. - - Such a run needs a GPU and downloads tens of gigabytes of checkpoint, so it - is asked for rather than automatic. It carries ``ci_gpu`` and skips unless - ``run_env`` is set; the ``manual`` marker describes it better and cannot be - used, since ``pytest-manual-marker`` xfails those at setup. - """ - if not os.environ.get(run_env): - return f"set {run_env}=1 to download the checkpoints and generate a clip" - if not torch.cuda.is_available(): - return "the model needs a GPU" - if shutil.which("ffmpeg") is None: - return "writing an MP4 needs ffmpeg on PATH" - return None - - -def check_real_model_generates_a_clip( - application: T2VApplication, - *, - prompt: str, - steps: int, - frame_count: int, - mp4_path: str | Path, -) -> T2VCheckResult: - """Generate a clip with a real checkpoint and check that it is a video. - - Every integration's real-model run is this one, and only the numbers differ. - No session is described, since the clip worth watching is the one the model - was trained for. Where it landed is printed, so a run made with ``-s`` says - where to look. - - Args: - application: Uninitialized application over the real model. - prompt: Text to generate from, usually the integration's own default. - steps: Blocks to generate. - frame_count: Frames those blocks should decode to. - mp4_path: File to write. - """ - result = check_t2v_model_impl( - application, - steps=steps, - # Compilation costs minutes and buys back milliseconds a block, which is - # the wrong trade for a handful of blocks. - commandline_args=["--prompt", prompt, "--no-compile"], - expected=ExpectedFrameStats( - frame_count=frame_count, - mean_luminance=_REAL_CLIP_LUMINANCE, - min_frame_difference=_REAL_CLIP_FRAME_DIFFERENCE, - ), - mp4_path=mp4_path, - ) - print(f"\nwrote {mp4_path}\n{result}") - return result - - -class FakeT2VPipeline: - """A model's worth of behaviour, without a model. - - Generates frames of the shape and range a real text-to-video pipeline does, - so an integration's tests can cover the seam a checkpoint plugs into on a - CPU. Every call is recorded, so a test can assert the rollout was driven in - order. - """ - - def __init__( - self, - *, - width: int = 128, - height: int = 64, - compression_ratio: int = 8, - first_block_frames: int = 9, - block_frames: int = 12, - fail_at: int | None = None, - ) -> None: - """ - Args: - width: Frame width to generate. Not square by default, so a - transposed frame cannot pass unnoticed. - height: Frame height to generate. - compression_ratio: Pixels one latent covers in each direction. - first_block_frames: Frames the first block decodes, which a causal - decoder usually has fewer of than the rest. - block_frames: Frames every block after the first decodes. - fail_at: Step to fail generating at, for covering a run that gave - up part way through. - """ - self.decoder = _FakeDecoder(compression_ratio) - self.width = width - self.height = height - self.first_block_frames = first_block_frames - self.block_frames = block_frames - self.device: str | None = None - self.eval_count = 0 - self.caches: list[dict[str, Any]] = [] - self.generated: list[int] = [] - self.finalized: list[int] = [] - self.closed = False - self._fail_at = fail_at - self._frames_generated = 0 - - def to(self, device: str) -> "FakeT2VPipeline": - self.device = device - return self - - def eval(self) -> "FakeT2VPipeline": - self.eval_count += 1 - return self - - def initialize_cache(self, **kwargs: Any) -> object: - self.caches.append(kwargs) - self._frames_generated = 0 - return object() - - def generate(self, *, autoregressive_index: int, cache: object) -> torch.Tensor: - del cache - self.generated.append(autoregressive_index) - if autoregressive_index == self._fail_at: - raise RuntimeError("generate failed") - count = ( - self.first_block_frames if autoregressive_index == 0 else self.block_frames - ) - frames = torch.stack( - [self._frame(self._frames_generated + index) for index in range(count)] - ) - self._frames_generated += count - return frames - - def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, float]: - del cache - self.finalized.append(autoregressive_index) - return {"total_ms": 1.5} - - def close(self) -> None: - self.closed = True - - def _frame(self, frame_index: int) -> torch.Tensor: - """Return a grey frame whose shade moves with time. - - Mid grey rather than black or white, and moving rather than still, so - the checks made of a real video are meaningful here too. - """ - shade = -0.5 + (frame_index % 8) / 8.0 - return torch.full((3, self.height, self.width), shade, dtype=torch.float32) - - -class FakeT2VPipelineConfig: - """A pipeline config that builds a stand-in rather than loading a model.""" - - def __init__(self, pipeline: FakeT2VPipeline | None = None) -> None: - """ - Args: - pipeline: Stand-in to build. A default one is made when none is - given, for a test that only cares that something was built. - """ - self.pipeline = pipeline if pipeline is not None else FakeT2VPipeline() - self.setup_count = 0 - - def setup(self) -> FakeT2VPipeline: - self.setup_count += 1 - return self.pipeline - - -class _FakeDecoder: - """The one thing a session asks a decoder for.""" - - def __init__(self, spatial_compression_ratio: int) -> None: - self.spatial_compression_ratio = spatial_compression_ratio - - -class _FrameInspector(OutputSink): - """Measure what a run generates, and pass it on to a file when asked to.""" - - def __init__(self, mp4: Mp4OutputSink | None) -> None: - """ - Args: - mp4: File sink to write as well, or ``None`` to only measure. - """ - self._mp4 = mp4 - self._session_desc: SessionDesc | None = None - self._last_frame: npt.NDArray[np.uint8] | None = None - self.frames_per_step: list[int] = [] - self.metrics: list[dict[str, float | int]] = [] - self.luminance_sum = 0.0 - self.difference_sum = 0.0 - self.difference_count = 0 - - def open(self, session_desc: SessionDesc) -> None: - self._session_desc = session_desc - if self._mp4 is not None: - self._mp4.open(session_desc) - - def write(self, result: StepResult) -> None: - if self._session_desc is None: - raise RuntimeError("open() must run before write().") - frames = result_to_rgb24_frames(result, self._session_desc) - self.frames_per_step.append(len(frames)) - self.metrics.append(dict(result.metrics)) - self.luminance_sum += float(frames.mean()) * len(frames) - # The previous step's last frame leads this one, so the change across a - # step boundary counts like any other. - sequence = frames - if self._last_frame is not None: - sequence = np.concatenate([self._last_frame[np.newaxis], frames]) - if len(sequence) > 1: - change = np.abs(np.diff(sequence.astype(np.int16), axis=0)) - self.difference_sum += float(change.mean()) * (len(sequence) - 1) - self.difference_count += len(sequence) - 1 - self._last_frame = frames[-1] - if self._mp4 is not None: - self._mp4.write(result) - - def close(self) -> None: - if self._mp4 is not None: - self._mp4.close() - - -class _DiscardingClientWindow(IClientWindow): - """Drive a checked run without retaining its composed UI output.""" - - def get_user_input_events(self) -> UserInputEvents: - return UserInputEvents([]) - - def open(self, session_desc: SessionDesc) -> None: - del session_desc - - def write(self, result: StepResult) -> None: - del result - - def close(self) -> None: - return - - -def _compare( - inspector: _FrameInspector, - expected: ExpectedFrameStats, - mp4_path: Path | None, -) -> T2VCheckResult: - """Measure what was generated and collect the expectations it missed.""" - frame_count = sum(inspector.frames_per_step) - luminance = inspector.luminance_sum / frame_count if frame_count else 0.0 - difference = ( - inspector.difference_sum / inspector.difference_count - if inspector.difference_count - else 0.0 - ) - - failures: list[str] = [] - if expected.frame_count is not None and frame_count != expected.frame_count: - failures.append( - f"Expected {expected.frame_count} frames, generated {frame_count} " - f"as {inspector.frames_per_step}." - ) - if expected.mean_luminance is not None: - low, high = expected.mean_luminance - if not low <= luminance <= high: - failures.append( - f"Mean luminance {luminance:.1f} is outside [{low}, {high}]." - ) - if ( - expected.min_frame_difference is not None - and difference < expected.min_frame_difference - ): - failures.append( - f"Frames change by {difference:.2f} on average, less than the " - f"{expected.min_frame_difference} expected of a video." - ) - - return T2VCheckResult( - failures=tuple(failures), - frames_per_step=tuple(inspector.frames_per_step), - mean_luminance=luminance, - frame_difference=difference, - metrics=tuple(inspector.metrics), - mp4_path=mp4_path, - ) diff --git a/flashdreams/flashdreams/runtime_v2/README.md b/flashdreams/flashdreams/runtime_v2/README.md index 4c239a3eb..184cfdf1e 100644 --- a/flashdreams/flashdreams/runtime_v2/README.md +++ b/flashdreams/flashdreams/runtime_v2/README.md @@ -107,7 +107,7 @@ loops to stop and performs their normal cleanup. An in-flight model step must return before the process can finish cleaning up. Synchronous application or session initialization likewise cannot be interrupted mid-call. -`--stats-path` adds a `MetricsOutputSink`. It receives the **model** loop's +`--stats-path .json` adds a `MetricsOutputSink`. It receives the **model** loop's results as they are published, not the UI loop's output, so a benchmark measures what the model generated while the window still sees one composited frame per tick. See [`configs/v2_model_benchmarks.json`](../../../configs/v2_model_benchmarks.json) diff --git a/flashdreams/flashdreams/runtime_v2/application_runner.py b/flashdreams/flashdreams/runtime_v2/application_runner.py index 4e518d62d..23d79981c 100644 --- a/flashdreams/flashdreams/runtime_v2/application_runner.py +++ b/flashdreams/flashdreams/runtime_v2/application_runner.py @@ -12,6 +12,7 @@ from flashdreams.api_v2.application import IApplication from flashdreams.api_v2.client_window import IClientWindow from flashdreams.api_v2.output_sink import OutputSink +from flashdreams.runtime_v2.metrics_output_sink import MetricsOutputSink from flashdreams.runtime_v2.session_desc import SessionDesc from flashdreams.runtime_v2.session_runner import run_session @@ -27,7 +28,7 @@ def __init__( application: IApplication, client_window: IClientWindow, *, - metrics_output_sink: OutputSink | None = None, + metrics_output_sink: MetricsOutputSink | None = None, ) -> None: """ Args: diff --git a/flashdreams/flashdreams/runtime_v2/metrics_output_sink.py b/flashdreams/flashdreams/runtime_v2/metrics_output_sink.py index a1433f2fa..0dc0b589b 100644 --- a/flashdreams/flashdreams/runtime_v2/metrics_output_sink.py +++ b/flashdreams/flashdreams/runtime_v2/metrics_output_sink.py @@ -8,6 +8,8 @@ from pathlib import Path from typing import Any +import numpy as np + from flashdreams.api_v2.output_sink import OutputSink from flashdreams.runtime_v2.session_desc import SessionDesc from flashdreams.runtime_v2.step_result import StepResult @@ -38,6 +40,7 @@ def __init__(self, path: str | Path) -> None: self._session_desc: SessionDesc | None = None self._steps: list[dict[str, Any]] = [] self._samples: list[dict[str, Any]] = [] + self._model_step_info: list[tuple[float, int]] = [] self._written = False def open(self, session_desc: SessionDesc) -> None: @@ -45,8 +48,26 @@ def open(self, session_desc: SessionDesc) -> None: self._session_desc = session_desc self._steps = [] self._samples = [] + self._model_step_info = [] self._written = False + def record_model_step_elapsed_s( + self, step_elapsed_s: float, step_results: list[StepResult] + ) -> None: + """Record wall time and the largest channel frame count for one step. + + Visual update cadence follows the largest channel in the step, so FPS + is that frame count divided by elapsed seconds. + """ + assert step_results and len(step_results) > 0 + + # Ensure all results are ready before we measure the chunk size. + # Since we are measuring "elapsed ms average", the cost of presenting + # all 8 at once is being measured via this wait. + for result in step_results: + result.wait_until_ready() + self._model_step_info.append((step_elapsed_s, step_results[0].frame_count)) + def write(self, result: StepResult) -> None: """Record one model result's frame count and metrics. @@ -83,6 +104,10 @@ def close(self) -> None: }, "steps": self._steps, "samples": self._samples, + "model_step_info": self._model_step_info, + # We provide a p90 FPS metric for all demos since we need to standardize measuring within a percentile of samples. + # Otherwise, reporting will include outliers from model start-up, etc... + "model_step_p90_fps": _model_step_p90_fps(self._model_step_info), } self._path.parent.mkdir(parents=True, exist_ok=True) self._path.write_text( @@ -114,6 +139,18 @@ def _samples_from(result: StepResult) -> list[dict[str, Any]]: return samples +def _model_step_p90_fps(step_info: list[tuple[float, int]]) -> float | None: + """Return p90 of largest-channel frames per elapsed second.""" + fps_samples = [ + chunk_size / elapsed_s + for elapsed_s, chunk_size in step_info + if elapsed_s > 0 and chunk_size > 0 + ] + if not fps_samples: + return None + return float(np.percentile(fps_samples, 90)) + + def _normalized_sample( name: str, value: float | int ) -> tuple[str, float | int, str, str]: diff --git a/flashdreams/flashdreams/runtime_v2/session_runner.py b/flashdreams/flashdreams/runtime_v2/session_runner.py index 7d9ed3ec8..6aaa7578a 100644 --- a/flashdreams/flashdreams/runtime_v2/session_runner.py +++ b/flashdreams/flashdreams/runtime_v2/session_runner.py @@ -15,6 +15,7 @@ from flashdreams.api_v2.output_sink import OutputSink from flashdreams.api_v2.session import ISession from flashdreams.runtime_v2.event_buffer import EventBuffer +from flashdreams.runtime_v2.metrics_output_sink import MetricsOutputSink from flashdreams.runtime_v2.session_desc import PresentationMode, SessionDesc from flashdreams.runtime_v2.step_result import StepResult @@ -47,7 +48,7 @@ def run_session( session: ISession, window: IClientWindow, *, - metrics_output_sink: OutputSink | None = None, + metrics_output_sink: MetricsOutputSink | None = None, steps: int | None = None, timeout_seconds: float | None = None, ) -> SessionDesc | None: @@ -168,6 +169,7 @@ def publish_model_results( step_elapsed_s=step_elapsed_s, ) if metrics_output_sink is not None: + metrics_output_sink.record_model_step_elapsed_s(step_elapsed_s, results) for result in results: metrics_output_sink.write(result) diff --git a/flashdreams/flashdreams/runtime_v2/step_result.py b/flashdreams/flashdreams/runtime_v2/step_result.py index 1e2ab6474..80ee8a131 100644 --- a/flashdreams/flashdreams/runtime_v2/step_result.py +++ b/flashdreams/flashdreams/runtime_v2/step_result.py @@ -60,6 +60,10 @@ def __post_init__(self, output: Tensor) -> None: event.record(torch.cuda.current_stream(output.device)) object.__setattr__(self, "_output_ready_event", event) + def wait_until_ready(self) -> None: + """Wait until the generated output is available to consumers.""" + _ = self.read_output(sync_with_event=True) + def read_output(self, *, sync_with_event: bool = True) -> Tensor: """Return the output, optionally ordered before the current CUDA stream. diff --git a/flashdreams/test_v2/test_application_runner.py b/flashdreams/test_v2/test_application_runner.py index 85d53c90a..f5bc41386 100644 --- a/flashdreams/test_v2/test_application_runner.py +++ b/flashdreams/test_v2/test_application_runner.py @@ -3,9 +3,11 @@ """CPU tests for the v2 application runner.""" +import json import logging from collections.abc import Sequence from dataclasses import replace +from pathlib import Path import pytest import torch @@ -16,6 +18,7 @@ from flashdreams.api_v2.loop import IModelLoop, IUILoop from flashdreams.api_v2.session import ISession from flashdreams.runtime_v2.application_runner import ApplicationRunner +from flashdreams.runtime_v2.metrics_output_sink import MetricsOutputSink from flashdreams.runtime_v2.session_desc import PresentationMode, SessionDesc from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import CloseUserInputEvent @@ -89,6 +92,7 @@ def step(self, step_index: int, events: UserInputEvents) -> StepResult: output=torch.full((1, 3, 1, 2, 2), step_index), frame_count=1, output_layout=VideoTensorLayout.bcthw, + metrics={"total_ms": 10.0}, ) def close(self) -> None: @@ -203,25 +207,6 @@ def open(self, session_desc: SessionDesc) -> None: self._sessions_opened += 1 -class _MetricsSink: - """Record model results delivered independently of the client window.""" - - def __init__(self, calls: list[str]) -> None: - self._calls = calls - self.results: list[StepResult] = [] - - def open(self, session_desc: SessionDesc) -> None: - del session_desc - self._calls.append("metrics.open") - - def write(self, result: StepResult) -> None: - self.results.append(result) - self._calls.append(f"metrics.write({result.step_index})") - - def close(self) -> None: - self._calls.append("metrics.close") - - def _session_desc() -> SessionDesc: return SessionDesc( output_layout=VideoTensorLayout.bcthw, @@ -303,37 +288,34 @@ def test_application_timeout_must_be_positive_and_finite(timeout: float) -> None assert calls == [] -def test_application_runner_keeps_metrics_output_separate_from_the_window() -> None: +def test_application_runner_keeps_metrics_output_separate_from_the_window( + tmp_path: Path, +) -> None: calls: list[str] = [] window = _ClosingAfterWritesWindow(calls, 2) - metrics = _MetricsSink(calls) + stats_path = tmp_path / "stats.json" ApplicationRunner( _Application(calls, session_length=2), window, - metrics_output_sink=metrics, + metrics_output_sink=MetricsOutputSink(stats_path), ).run(_session_desc()) assert [ result.read_output()[0, 0, 0, 0, 0].item() for result in window.results ] == [0, 1] - assert [result.step_index for result in metrics.results] == [0, 1] - assert calls.index("metrics.open") < calls.index("metrics.write(0)") - assert calls.index("metrics.write(1)") < calls.index("metrics.close") + payload = json.loads(stats_path.read_text(encoding="utf-8")) + assert [step["step_index"] for step in payload["steps"]] == [0, 1] + assert [sample["name"] for sample in payload["samples"]] == ["total_s", "total_s"] def test_application_runner_replaces_a_session_before_closing_the_window() -> None: calls: list[str] = [] application = _Application(calls, replace_first_session=True) window = _SecondSessionClosingWindow(calls) - metrics = _MetricsSink(calls) session_desc = _session_desc() - ApplicationRunner( - application, - window, - metrics_output_sink=metrics, - ).run(session_desc) + ApplicationRunner(application, window).run(session_desc) create_indexes = [ index @@ -348,8 +330,6 @@ def test_application_runner_replaces_a_session_before_closing_the_window() -> No assert close_indexes[0] < create_indexes[1] assert calls.count("window.open") == 2 assert calls.count("window.close") == 1 - assert calls.count("metrics.open") == 2 - assert calls.count("metrics.close") == 2 assert application.requested_session_descs == [session_desc, session_desc] assert application.requested_session_descs[1] is session_desc assert calls.count("application.init([])") == 1 @@ -378,27 +358,6 @@ def create_session(self, session_desc: SessionDesc) -> ISession: assert calls[-1] == "application.close" -def test_replacement_stops_if_per_session_metrics_fail_to_close() -> None: - calls: list[str] = [] - - class FailingMetricsSink(_MetricsSink): - def close(self) -> None: - super().close() - raise RuntimeError("metrics close failed") - - with pytest.raises(RuntimeError, match="metrics close failed"): - ApplicationRunner( - _Application(calls, replace_first_session=True), - _SecondSessionClosingWindow(calls), - metrics_output_sink=FailingMetricsSink(calls), - ).run(_session_desc()) - - assert calls.count("application.create_session") == 1 - assert calls.count("session.close") == 1 - assert calls.count("window.close") == 1 - assert calls[-1] == "application.close" - - def test_setup_failure_closes_every_runner_owned_resource() -> None: calls: list[str] = [] bad_trace_desc = replace( @@ -407,16 +366,11 @@ def test_setup_failure_closes_every_runner_owned_resource() -> None: ) with pytest.raises(TypeError, match="trace_chunk_lifecycle_path"): - ApplicationRunner( - _Application(calls), - _Window(calls), - metrics_output_sink=_MetricsSink(calls), - ).run(bad_trace_desc) + ApplicationRunner(_Application(calls), _Window(calls)).run(bad_trace_desc) assert calls == [ "application.init([])", "application.create_session", - "metrics.close", "window.close", "session.close", "application.close", diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py index 60ae42492..df3014169 100644 --- a/flashdreams/test_v2/test_session_runner.py +++ b/flashdreams/test_v2/test_session_runner.py @@ -3,6 +3,7 @@ """CPU test for the v2 session loop, independent of any application.""" +import json import logging import queue import threading @@ -10,6 +11,7 @@ from collections.abc import Iterator from contextlib import contextmanager from dataclasses import replace +from pathlib import Path import pytest import torch @@ -29,6 +31,7 @@ BlitModelOutputToScreenLoop, ) from flashdreams.runtime_v2.event_buffer import EventBuffer +from flashdreams.runtime_v2.metrics_output_sink import MetricsOutputSink from flashdreams.runtime_v2.presentation_manager import ( _PRESENTATION_DRAIN_MARGIN, PresentationManager, @@ -352,6 +355,7 @@ def step(self, step_index: int, events: UserInputEvents) -> StepResult: output=torch.full((1, 3, 1, 1, 1), step_index, dtype=torch.float32), frame_count=1, output_layout=self._session_desc.output_layout, + metrics={"total_ms": 10.0}, ) def run_ui(self, step_index: int, events: UserInputEvents) -> StepResult | None: @@ -365,7 +369,6 @@ def run_ui(self, step_index: int, events: UserInputEvents) -> StepResult | None: output=frame.unsqueeze(0).unsqueeze(2), frame_count=1, output_layout=self.session_desc.output_layout, - metrics={"ui_ms": 0.25}, ) def is_finished(self) -> bool: @@ -575,6 +578,37 @@ def test_run_session_presents_every_step_in_order() -> None: assert steps == ["session.step(0)", "session.step(1)", "session.step(2)"] +def test_run_session_records_model_metrics_separately_from_the_window( + tmp_path: Path, +) -> None: + log = CallLog() + session = FakeSession(_session_desc(), log) + window = RecordingClientWindow(log) + stats_path = tmp_path / "stats.json" + + run_session( + session, + window, + metrics_output_sink=MetricsOutputSink(stats_path), + steps=3, + ) + + assert [ + result.read_output()[0, 0, 0, 0, 0].item() for result in window.results + ] == [0, 1, 2] + payload = json.loads(stats_path.read_text(encoding="utf-8")) + assert [step["step_index"] for step in payload["steps"]] == [0, 1, 2] + assert [sample["name"] for sample in payload["samples"]] == [ + "total_s", + "total_s", + "total_s", + ] + assert all(result.metrics == {} for result in window.results) + assert len(payload["model_step_info"]) == 3 + assert all(chunk_size == 1 for _, chunk_size in payload["model_step_info"]) + assert isinstance(payload["model_step_p90_fps"], float) + + def test_run_session_opens_before_writing_and_closes_after() -> None: log = CallLog() session = FakeSession(_session_desc(), log) @@ -924,52 +958,6 @@ def test_composite_clamps_alpha_before_interpolation() -> None: assert torch.equal(composited[:, :, 1], overlay[:3, :, 1]) -def test_default_ui_presents_each_frame_from_a_model_chunk() -> None: - log = CallLog() - - class MultiFrameSession(FakeSession): - def step(self, step_index: int, events: UserInputEvents) -> StepResult: - del events - self._log.record(f"session.step({step_index})") - return StepResult( - step_index=step_index, - output=torch.arange(36, dtype=torch.float32).reshape(1, 3, 12, 1, 1), - frame_count=12, - output_layout=self.session_desc.output_layout, - metrics={"total_ms": 1.5}, - ) - - class RecordingMetricsSink: - def __init__(self) -> None: - self.results: list[StepResult] = [] - - def open(self, session_desc: SessionDesc) -> None: - del session_desc - - def write(self, result: StepResult) -> None: - self.results.append(result) - - def close(self) -> None: - return - - window = RecordingClientWindow(log) - metrics = RecordingMetricsSink() - run_session( - MultiFrameSession(_session_desc(), log), - window, - metrics_output_sink=metrics, - steps=1, - ) - - assert [result.frame_count for result in window.results] == [1] * 12 - assert [ - result.read_output()[0, 0, 0, 0, 0].item() for result in window.results - ] == list(range(12)) - assert [result.metrics for result in window.results] == [{"ui_ms": 0.25}] * 12 - assert len(metrics.results) == 1 - assert metrics.results[0].metrics == {"total_ms": 1.5} - - def test_default_ui_does_not_redraw_an_unchanged_model_frame() -> None: log = CallLog()