diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md index 6a4eea9dc..2f0ba19f8 100644 --- a/docs/inference_runtime_api_design.md +++ b/docs/inference_runtime_api_design.md @@ -59,25 +59,25 @@ Initial scope: ## Task Tracker -| ID | Workstream | Can run in parallel? | Depends on | Done when | -| --- | --- | --- | --- | --- | -| T0 | Create experimental branch and contribution rules. | No, this starts the work. | None. | Branch exists, PR target is agreed, and main merge criteria are written down. | -| T1 | Minimal API envelope and naming. | Partly. | T0. | `InferenceConfig`, `UserInputs`, `ModelInputs`, runtime/session, output target, and mapping boundaries are defined well enough for demos to use. | -| T2 | Event-based `UserInputs`. | Yes, after T1 direction is agreed. | T1. | User inputs are primarily timestamped events; replay traces and derived snapshots are supported where needed. | -| T3 | `ModelInputs`, schemas, and mapping boundary. | Yes, after T1 direction is agreed. | T1. | Models can declare required initial/per-step inputs, and mappings can convert user events into model inputs. | -| T4 | `ModelRunner`, `InferenceRuntime`, and `InferenceSession` skeleton. | Partly. | T1. | A minimal standard loop can initialize a runtime, run at least one sequential session, and close cleanly. | -| T5 | Output mode selection. | Yes, after the result/output shape is agreed. | T1, T4. | A run can choose output behavior such as MP4, JPEG/MJPEG stream, WebRTC, benchmark artifact, or headless/null without changing model code. | -| T6 | LingBot migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | LingBot runs through the new API path with its event inputs mapped into model inputs. | -| T7 | OmniDreams migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | OmniDreams runs through the new API path with its model-specific inputs and mapping preserved. | -| T8 | Benchmark/smoke verification for LingBot and OmniDreams. | Preparation can run early; final gate is late. | T5, T6, T7. | Existing or updated benchmark tooling can run both migrated demos and produce enough evidence that they still work. | -| T9 | Metrics and profiling normalization for the branch. | Yes, but final integration is late. | T4, T5, T8. | Basic canonical metrics are emitted for migrated demos; deeper metrics can remain follow-up work. | -| T10 | CLI compatibility and migration cleanup. | Yes, after demo migrations start. | T6, T7. | Required demo commands are restored or replaced, temporary hacks are removed, and user-facing docs/notes match the branch behavior. | -| T11 | Stabilize and merge experimental branch to `main`. | No, final integration step. | T6-T10. | LingBot and OmniDreams pass agreed smoke/benchmark checks, review feedback is addressed, and the branch can merge as one API transition. | +| ID | Status | Workstream | Can run in parallel? | Depends on | Done when | +| --- | --- | --- | --- | --- | --- | +| T0 | Complete | Create experimental branch and contribution rules. | No, this starts the work. | None. | Branch exists, PR target is agreed, and main merge criteria are written down. | +| T1 | Complete | Minimal API envelope and naming. | Partly. | T0. | `InferenceConfig`, `UserInputs`, `ModelInputs`, runtime/session, output target, and mapping boundaries are defined well enough for demos to use. | +| T2 | Planned | Event-based `UserInputs`. | Yes, after T1 direction is agreed. | T1. | User inputs are primarily timestamped events; replay traces and derived snapshots are supported where needed. | +| T3 | Planned | `ModelInputs`, schemas, and mapping boundary. | Yes, after T1 direction is agreed. | T1. | Models can declare required initial/per-step inputs, and mappings can convert user events into model inputs. | +| T4 | Planned | `ModelRunner`, `InferenceRuntime`, and `InferenceSession` skeleton. | Partly. | T1. | A minimal standard loop can initialize a runtime, run at least one sequential session, and close cleanly. | +| T5 | Planned | Output mode selection. | Yes, after the result/output shape is agreed. | T1, T4. | A run can choose output behavior such as MP4, JPEG/MJPEG stream, WebRTC, benchmark artifact, or headless/null without changing model code. | +| T6 | Planned | LingBot migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | LingBot runs through the new API path with its event inputs mapped into model inputs. | +| T7 | Planned | OmniDreams migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | OmniDreams runs through the new API path with its model-specific inputs and mapping preserved. | +| T8 | Planned | Benchmark/smoke verification for LingBot and OmniDreams. | Preparation can run early; final gate is late. | T5, T6, T7. | Existing or updated benchmark tooling can run both migrated demos and produce enough evidence that they still work. | +| T9 | Planned | Metrics and profiling normalization for the branch. | Yes, but final integration is late. | T4, T5, T8. | Basic canonical metrics are emitted for migrated demos; deeper metrics can remain follow-up work. | +| T10 | Planned | CLI compatibility and migration cleanup. | Yes, after demo migrations start. | T6, T7. | Required demo commands are restored or replaced, temporary hacks are removed, and user-facing docs/notes match the branch behavior. | +| T11 | Planned | Stabilize and merge experimental branch to `main`. | No, final integration step. | T6-T10. | LingBot and OmniDreams pass agreed smoke/benchmark checks, review feedback is addressed, and the branch can merge as one API transition. | Suggested parallel split: -- one person owns T1/T4, because the API envelope and standard loop are the - critical path; +- one person owns T4 and keeps it aligned with the completed T1 envelope, + because the standard loop is now the critical path; - one person owns T2/T3, because event inputs, schemas, and mapping need to stay coherent; - one person owns T5/T8/T9, because outputs, benchmarks, and metrics are tightly @@ -182,7 +182,7 @@ local model implementation, a Dynamo-like backend, or a hosted service. | Model/preset registry | Lists what can run: model/preset slugs, scenarios, capabilities, resource hints, and supported output modes. | Must remain cheap to query and must not load checkpoints. | | App / integration / benchmark / transport | Owns the user-facing mode: CLI, native integration, WebRTC, benchmark, hosted request, or replay. | Supplies run setup, user inputs, model inputs, and output target selection. | | User input library | Normalizes live or replayed controls into FlashDreams-supported user input events/windows. | Shared primitives for keyboard, reset, prompt/image updates, traces, and future scalar controls. | -| Input mapping | Converts user/app inputs plus initial model inputs into the model-specific inputs needed by the session. | Owned by model/application code; may be a no-op for simple runs. | +| Input mapping | Converts user/app inputs plus initial model inputs into the model-specific inputs needed by the session. | A model adapter may provide a default mapping; runtimes, applications, benchmarks, and replay tools may override it without changing the model step. | | ModelRunner / standard loop | Orchestrates one run from setup through runtime initialization, stepping, output, metrics, and teardown. | Shared orchestration layer used by CLIs, benchmarks, MP4 runs, and simple realtime flows. | | InferenceRuntime | Owns heavyweight lifecycle: distributed init, model construction, checkpoint loading, compile/capture, warmup, hosted-service connection, and teardown. | Long-lived reusable runtime created from `InferenceConfig`; lets FlashDreams load/warm once and create sessions sequentially unless the backend supports concurrency. | | InferenceSession | Owns one rollout or stream: initial inputs, cache state, current step, reset behavior, step requirements, and step execution. | Per-rollout interface consumed by the standard loop; keeps state isolated across prompts, browser clients, replay scenarios, or benchmark repeats. | @@ -281,8 +281,11 @@ native app, replay trace, synthetic benchmark driver, or no-op source. User inputs should primarily be represented as timestamped events. This gives live apps, replay traces, and benchmarks the same basic shape, and lets -FlashDreams resample or window those events when a model session asks for the -next chunk of inputs. +FlashDreams route, drain, or window those events when a model session asks for +the next chunk of inputs. Resampling and interpolation should remain +input-specific mapping or helper behavior, because controls such as rotations, +poses, or controller state may need semantics that generic runtime code cannot +infer safely. Initial supported user input types should stay close to what FlashDreams already uses: @@ -359,8 +362,8 @@ shapes or architecture details, but it usually does not fully define: - preprocessing, encoder, decoder, mask, prompt, or cache rules. Therefore, a FlashDreams-supported model should have an adapter or integration -layer that declares its model input requirements and prepares inputs for the -underlying model implementation. +layer that declares its model input requirements, declares any user inputs it can +map by default, and prepares inputs for the underlying model implementation. Users running an existing FlashDreams-supported model should not need to write that adapter. Developers bringing a new world model to FlashDreams should expect @@ -407,21 +410,25 @@ unless the checkpoint already matches a supported generic adapter. ## Input Mapping Input mapping is required whenever `UserInputs` need to become per-step -`ModelInputs`. The exact implementation does not need to be a required top-level -object. It could be: - -- a method on the model adapter; -- a method on an app/runtime adapter; -- a separate mapper object; -- a default no-op or identity mapping for simple T2V/I2V/fixed-input runs. +`ModelInputs`. In the T1 envelope this boundary is represented by a separate +`InputMapping` protocol. A model adapter may provide the default mapper because +it knows how its supported user controls affect model-facing inputs. Applications, +benchmarks, replay tools, or hosted runtimes may replace that mapper when they +need a different wire surface or aggregation policy. There are two separate moments to keep clear: - before runtime initialization, FlashDreams should select the mapping and check obvious compatibility between the app event source and the model; -- during the standard loop, the runner uses the mapping to build initial or - per-step `ModelInputs` from the relevant event window, often after the session - reports what it needs next. +- during the standard loop, the runtime or runner queues and timestamps user + events, then uses the selected mapping to build initial or per-step + `ModelInputs` from the relevant event window, often after the session reports + what it needs next. + +This keeps the Reactor-style contract intact: the model-side integration can +declare user inputs, declare model inputs, and provide a default mapping, while +the runtime owns transport, event validation, timestamping, input queue/window +selection, output delivery, and optional overrides. Examples: @@ -465,6 +472,11 @@ the app or output layer to guess. This matters because AR step 0 can differ from steady-state steps, and encoder/decoder temporal compression can produce different input and output frame windows. +Input and output timing should share a session timeline even when raw capture +rates and presentation rates differ. A session can request a user-input window +for mapping, then return an output window or equivalent metadata so an output +target can present the generated chunk at the intended cadence. + ## Output Targets Output handling should be separate from model execution. The model session @@ -598,20 +610,40 @@ back to `main` safely. - Public/internal boundaries must remain clean. Internal adapters, slugs, and scenarios should not leak into the public repo. -## Decisions To Make Before Implementation +## Decisions Made In T1 + +Task T1 settles the initial package and naming envelope without committing to a +registry, standard loop, concrete output modes, or model migrations: + +- The experimental API lives under `flashdreams.runtime`. +- The model-specific integration boundary is named `ModelAdapter`. +- Heavyweight lifecycle is split into `InferenceRuntime` and + `InferenceSession`. +- Step data carriers are named `StepRequest` and `StepResult`; a session returns + `None` from `next_step_request()` when the rollout is complete. +- User-facing inputs use `UserInputs`; model-facing inputs use `ModelInputs`. + Both remain lightweight payload envelopes with shallow read-only mappings. +- `UserInputSchema` and `ModelInputSchema` stay intentionally small: they + declare supported event types and required named fields for early validation, + not a full type system. +- Input mapping is represented by a separate `InputMapping` protocol. Model + adapters may provide a default mapping; runtimes and applications may override + it while preserving the `UserInputs` to `ModelInputs` boundary. Simple + fixed-input runs can use `IdentityInputMapping`. +- Output handling is represented by `OutputTarget`; `NullOutputTarget` is the + initial headless implementation. +- Metrics collection is represented by `MetricsRecorder`; timing samples use + seconds as the canonical unit. +- The minimum v0 user input shape is timestamped `UserInputEvent` records plus + optional snapshot data. Concrete event-type catalogs are left to T2 and demo + migrations. + +## Remaining Decisions -- What should the top-level package/API be called? -- Should the main registered object be called an adapter, integration, runtime - factory, or something else? - What direct-Python API should let users pass an external adapter without registering it? - What package registration mechanism should third-party and internal adapters use for CLI discovery and benchmarks? -- How lightweight should `UserInputSchema` and `ModelInputSchema` be? -- Where should input mapping live: model adapter, app adapter, separate object, - or a mix? -- What should the output abstraction be called? -- What is the minimum v0 set of supported user input events? - What is the first public model to migrate? - What metrics are required for every benchmark run? - What metadata must be discoverable without loading checkpoints? diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py new file mode 100644 index 000000000..03e6202b0 --- /dev/null +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Experimental inference runtime API envelope. + +This package defines the small v0 boundary above ``flashdreams.infra``. It is +intentionally additive while integrations migrate onto it. +""" + +from flashdreams.runtime.config import ExecutionBackend, InferenceConfig, Precision +from flashdreams.runtime.inputs import ( + InputField, + ModelInputs, + ModelInputSchema, + TimeWindow, + UserInputEvent, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.interfaces import ( + InferenceRuntime, + InferenceSession, + ModelAdapter, +) +from flashdreams.runtime.mapping import IdentityInputMapping, InputMapping +from flashdreams.runtime.metrics import ( + InMemoryMetricsRecorder, + MetricsRecorder, + NullMetricsRecorder, + RuntimeMetricSample, +) +from flashdreams.runtime.output import NullOutputTarget, OutputArtifact, OutputTarget +from flashdreams.runtime.types import StepRequest, StepResult + +__all__ = [ + "ExecutionBackend", + "IdentityInputMapping", + "InferenceConfig", + "InferenceRuntime", + "InferenceSession", + "InMemoryMetricsRecorder", + "InputField", + "InputMapping", + "MetricsRecorder", + "ModelAdapter", + "ModelInputs", + "ModelInputSchema", + "NullMetricsRecorder", + "NullOutputTarget", + "OutputArtifact", + "OutputTarget", + "Precision", + "RuntimeMetricSample", + "StepRequest", + "StepResult", + "TimeWindow", + "UserInputEvent", + "UserInputs", + "UserInputSchema", +] diff --git a/flashdreams/flashdreams/runtime/_utils.py b/flashdreams/flashdreams/runtime/_utils.py new file mode 100644 index 000000000..d8016c6b7 --- /dev/null +++ b/flashdreams/flashdreams/runtime/_utils.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small helpers shared by the experimental runtime API.""" + +from __future__ import annotations + +from collections.abc import Mapping +from types import MappingProxyType +from typing import TypeVar + +ValueT = TypeVar("ValueT") + + +def freeze_mapping(value: Mapping[str, ValueT]) -> Mapping[str, ValueT]: + """Return a read-only shallow copy of ``value``.""" + return MappingProxyType(dict(value)) diff --git a/flashdreams/flashdreams/runtime/config.py b/flashdreams/flashdreams/runtime/config.py new file mode 100644 index 000000000..4b8752f13 --- /dev/null +++ b/flashdreams/flashdreams/runtime/config.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runtime-facing configuration envelope.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + +from flashdreams.runtime._utils import freeze_mapping + +ExecutionBackend = Literal["local", "local-distributed", "external", "hosted"] +"""Where and how inference compute is run.""" + +Precision = Literal["auto", "fp32", "fp16", "bf16"] +"""Coarse runtime precision choices.""" + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InferenceConfig: + """Runtime settings that affect model execution. + + Prompts, user controls, browser settings, output paths, and benchmark + directories intentionally live outside this object. The typed optimization + fields cover common cross-backend knobs; open-ended adapter-specific choices + can use :attr:`runtime_options`. + """ + + __hash__ = None + + model_id: str + """Stable identity for the model adapter or runtime integration.""" + + preset_id: str | None = None + """Optional preset identity under :attr:`model_id`.""" + + checkpoint: str | Path | None = None + """Optional checkpoint or model-asset selector understood by the adapter.""" + + backend: ExecutionBackend = "local" + """Execution placement and backend family for inference compute.""" + + device: str | None = None + """Optional device selector such as ``cuda`` or ``cuda:0``; ``None`` leaves placement to the adapter/backend.""" + + precision: Precision = "auto" + """Preferred compute precision.""" + + compile: bool | None = None + """Optional - Whether model compilation is requested or disabled. `None` means left to the adapter to decide.""" + + cuda_graph: bool | None = None + """Optional - Whether CUDA graph capture is requested or disabled. `None` means left to the adapter to decide.""" + + attention_backend: str | None = None + """Optional attention implementation selector; ``None`` leaves the choice to the adapter.""" + + cache_policy: str | None = None + """Optional cache policy selector; ``None`` leaves the choice to the adapter.""" + + runtime_options: Mapping[str, Any] = field(default_factory=dict) + """Adapter/backend-specific runtime options.""" + + resource_hints: Mapping[str, Any] = field(default_factory=dict) + """Resource hints for launchers, schedulers, or hosted backends.""" + + def __post_init__(self) -> None: + if not self.model_id.strip(): + raise ValueError("InferenceConfig.model_id must be non-empty.") + object.__setattr__( + self, "runtime_options", freeze_mapping(self.runtime_options) + ) + object.__setattr__(self, "resource_hints", freeze_mapping(self.resource_hints)) diff --git a/flashdreams/flashdreams/runtime/inputs.py b/flashdreams/flashdreams/runtime/inputs.py new file mode 100644 index 000000000..e14b35722 --- /dev/null +++ b/flashdreams/flashdreams/runtime/inputs.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""User- and model-input envelopes for the experimental runtime API.""" + +from __future__ import annotations + +import math +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field +from typing import Any + +from flashdreams.runtime._utils import freeze_mapping + + +@dataclass(frozen=True, kw_only=True, slots=True) +class TimeWindow: + """Half-open time window in seconds since session start.""" + + start_s: float + end_s: float + + def __post_init__(self) -> None: + if not math.isfinite(self.start_s) or not math.isfinite(self.end_s): + raise ValueError("TimeWindow bounds must be finite seconds.") + if self.start_s < 0 or self.end_s < 0: + raise ValueError("TimeWindow bounds must be non-negative.") + if self.end_s < self.start_s: + raise ValueError("TimeWindow.end_s must be >= start_s.") + + def contains(self, timestamp_s: float) -> bool: + """Return whether ``timestamp_s`` falls within this half-open window.""" + return self.start_s <= timestamp_s < self.end_s + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InputField: + """Lightweight schema field for user snapshots or model inputs.""" + + name: str + required: bool = True + semantic_type: str | None = None + description: str = "" + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("InputField.name must be non-empty.") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class UserInputSchema: + """Minimal metadata for user events a source or mapping can provide.""" + + event_types: frozenset[str] = field(default_factory=frozenset) + snapshot_fields: tuple[InputField, ...] = () + description: str = "" + + def supports_event_types(self, event_types: Iterable[str]) -> bool: + """Return whether every requested event type is declared supported.""" + requested = frozenset(event_types) + if not requested: + return True + return requested.issubset(self.event_types) + + def missing_snapshot(self, inputs: "UserInputs") -> tuple[str, ...]: + """Return required snapshot fields absent from ``inputs``.""" + return _missing_required(self.snapshot_fields, inputs.snapshot) + + def require_snapshot(self, inputs: "UserInputs") -> None: + """Raise if required snapshot fields are absent.""" + missing = self.missing_snapshot(inputs) + if missing: + raise ValueError(f"Missing required user snapshot field(s): {missing}") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ModelInputSchema: + """Minimal metadata for model-facing initial and per-step inputs.""" + + initial_fields: tuple[InputField, ...] = () + """Model inputs required before starting the initial generation/session.""" + + step_fields: tuple[InputField, ...] = () + """Per-step model inputs required after the session starts.""" + + description: str = "" + + def missing_initial(self, inputs: "ModelInputs") -> tuple[str, ...]: + """Return required initial fields absent from ``inputs``.""" + return _missing_required(self.initial_fields, inputs.initial) + + def missing_step(self, inputs: "ModelInputs") -> tuple[str, ...]: + """Return required per-step fields absent from ``inputs``.""" + return _missing_required(self.step_fields, inputs.step) + + def require_initial(self, inputs: "ModelInputs") -> None: + """Raise if required initial fields are absent.""" + missing = self.missing_initial(inputs) + if missing: + raise ValueError(f"Missing required initial model input(s): {missing}") + + def require_step(self, inputs: "ModelInputs") -> None: + """Raise if required per-step fields are absent.""" + missing = self.missing_step(inputs) + if missing: + raise ValueError(f"Missing required step model input(s): {missing}") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class UserInputEvent: + """User-facing input event timestamped in seconds since session start. + + Live runtimes, transports, replay loaders, or benchmark drivers stamp events + before queuing them for input mapping. Payload schema is intentionally minimal + in T1; concrete event catalogs belong to follow-up input-mapping work. + """ + + __hash__ = None + + timestamp_s: float + event_type: str + payload: Mapping[str, Any] = field(default_factory=dict) + source: str | None = None + source_event_id: str | None = None + + def __post_init__(self) -> None: + if not math.isfinite(self.timestamp_s) or self.timestamp_s < 0: + raise ValueError("UserInputEvent.timestamp_s must be finite and >= 0.") + if not self.event_type.strip(): + raise ValueError("UserInputEvent.event_type must be non-empty.") + object.__setattr__(self, "payload", freeze_mapping(self.payload)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class UserInputs: + """Transport-neutral user input batch or window. + + Events must be in non-decreasing timestamp order. Runtimes can pass the full + input history, a drained queue batch, or a session-requested time window to an + ``InputMapping``. + """ + + __hash__ = None + + events: tuple[UserInputEvent, ...] = () + snapshot: Mapping[str, Any] = field(default_factory=dict) + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + previous_timestamp_s = -math.inf + for event in self.events: + if event.timestamp_s < previous_timestamp_s: + raise ValueError( + "UserInputs.events must be sorted by non-decreasing timestamp_s." + ) + previous_timestamp_s = event.timestamp_s + object.__setattr__(self, "snapshot", freeze_mapping(self.snapshot)) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + def window(self, time_window: TimeWindow) -> "UserInputs": + """Return inputs with events filtered to ``time_window``.""" + return UserInputs( + events=tuple( + event + for event in self.events + if time_window.contains(event.timestamp_s) + ), + snapshot=self.snapshot, + metadata=self.metadata, + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ModelInputs: + """Model-facing payloads split by initial and per-step use.""" + + __hash__ = None + + initial: Mapping[str, Any] = field(default_factory=dict) + step: Mapping[str, Any] = field(default_factory=dict) + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "initial", freeze_mapping(self.initial)) + object.__setattr__(self, "step", freeze_mapping(self.step)) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + def with_step(self, step: Mapping[str, Any]) -> "ModelInputs": + """Return a copy with replaced per-step payload.""" + return ModelInputs(initial=self.initial, step=step, metadata=self.metadata) + + +def _missing_required( + fields: tuple[InputField, ...], payload: Mapping[str, Any] +) -> tuple[str, ...]: + return tuple( + input_field.name + for input_field in fields + if input_field.required and input_field.name not in payload + ) diff --git a/flashdreams/flashdreams/runtime/interfaces.py b/flashdreams/flashdreams/runtime/interfaces.py new file mode 100644 index 000000000..9b6a064fd --- /dev/null +++ b/flashdreams/flashdreams/runtime/interfaces.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Protocols for model adapters, reusable runtimes, and sessions.""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +from flashdreams.runtime.config import InferenceConfig +from flashdreams.runtime.inputs import ( + ModelInputs, + ModelInputSchema, + UserInputSchema, +) +from flashdreams.runtime.mapping import InputMapping +from flashdreams.runtime.types import StepRequest, StepResult + + +@runtime_checkable +class InferenceSession(Protocol): + """One rollout or stream with isolated model/cache state.""" + + def next_step_request(self) -> StepRequest | None: + """Describe the next step's inputs, or return ``None`` when complete.""" + ... + + def step(self, inputs: ModelInputs) -> StepResult: + """Run one sequential inference step.""" + ... + + def reset(self, inputs: ModelInputs | None = None) -> None: + """Reset this session's rollout state when the backend supports it.""" + ... + + def close(self) -> None: + """Release per-session resources.""" + ... + + +@runtime_checkable +class InferenceRuntime(Protocol): + """Heavyweight reusable runtime created from :class:`InferenceConfig`.""" + + def start_session(self, inputs: ModelInputs) -> InferenceSession: + """Create an isolated session from initial model inputs.""" + ... + + def close(self) -> None: + """Release model/backend resources.""" + ... + + +# Do not mark ModelAdapter runtime-checkable: properties make issubclass() +# unreliable, and isinstance() would only verify attribute presence. +class ModelAdapter(Protocol): + """Model-specific boundary that declares defaults and creates runtimes. + + Adapters declare model-facing input requirements, optional user-input + capabilities, and an optional default mapping between the two. Runtime, + application, or benchmark code may override that mapping while preserving the + same ``UserInputs`` to ``ModelInputs`` boundary. + """ + + @property + def model_id(self) -> str: + """Stable identity for the model adapter or runtime integration.""" + ... + + @property + def model_input_schema(self) -> ModelInputSchema: + """Model-facing initial and per-step input requirements.""" + ... + + @property + def user_input_schema(self) -> UserInputSchema | None: + """User inputs supported by the adapter's default mapping, if any.""" + ... + + def default_input_mapping(self) -> InputMapping | None: + """Return the model-provided default user-to-model mapping, if any.""" + ... + + def validate_config(self, config: InferenceConfig) -> None: + """Fail early for unsupported runtime settings.""" + ... + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + """Initialize and return the heavyweight runtime.""" + ... diff --git a/flashdreams/flashdreams/runtime/mapping.py b/flashdreams/flashdreams/runtime/mapping.py new file mode 100644 index 000000000..756351081 --- /dev/null +++ b/flashdreams/flashdreams/runtime/mapping.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Input mapping boundary from user input windows to model inputs.""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +from flashdreams.runtime.inputs import ( + ModelInputs, + ModelInputSchema, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.types import StepRequest + + +@runtime_checkable +class InputMapping(Protocol): + """Convert user-facing inputs into model-facing inputs. + + A mapping may be supplied by the model adapter as a default or by an + application/runtime override. Step mappings usually receive a timestamped + event window selected by the runner for the current model step or chunk. + """ + + def validate( + self, + *, + user_schema: UserInputSchema | None = None, + model_schema: ModelInputSchema | None = None, + ) -> None: + """Fail early for obvious app, event-source, and model mismatches.""" + ... + + def map_initial_inputs( + self, + *, + user_inputs: UserInputs, + model_inputs: ModelInputs, + ) -> ModelInputs: + """Build initial model inputs before a session starts.""" + ... + + def map_step_inputs( + self, + *, + user_inputs: UserInputs, + model_inputs: ModelInputs, + request: StepRequest, + ) -> ModelInputs: + """Build model inputs for one session step from the current input window.""" + ... + + +class IdentityInputMapping: + """No-op mapper for fixed model-input or simple generation flows.""" + + def validate( + self, + *, + user_schema: UserInputSchema | None = None, + model_schema: ModelInputSchema | None = None, + ) -> None: + del user_schema, model_schema + + def map_initial_inputs( + self, + *, + user_inputs: UserInputs, + model_inputs: ModelInputs, + ) -> ModelInputs: + del user_inputs + return model_inputs + + def map_step_inputs( + self, + *, + user_inputs: UserInputs, + model_inputs: ModelInputs, + request: StepRequest, + ) -> ModelInputs: + del user_inputs, request + return model_inputs diff --git a/flashdreams/flashdreams/runtime/metrics.py b/flashdreams/flashdreams/runtime/metrics.py new file mode 100644 index 000000000..4286204f6 --- /dev/null +++ b/flashdreams/flashdreams/runtime/metrics.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runtime metrics boundary for inference sessions.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +from flashdreams.runtime._utils import freeze_mapping + + +@dataclass(frozen=True, kw_only=True, slots=True) +class RuntimeMetricSample: + """One runtime metric sample. + + Timing samples should use seconds as their canonical unit. + """ + + __hash__ = None + + name: str + value: float | int + unit: str = "s" + step_index: int | None = None + category: str = "runtime" + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("RuntimeMetricSample.name must be non-empty.") + if isinstance(self.value, bool) or not isinstance(self.value, (int, float)): + raise TypeError("RuntimeMetricSample.value must be numeric.") + if not math.isfinite(float(self.value)): + raise ValueError("RuntimeMetricSample.value must be finite.") + if self.step_index is not None and self.step_index < 0: + raise ValueError("RuntimeMetricSample.step_index must be >= 0.") + if not self.unit.strip(): + raise ValueError("RuntimeMetricSample.unit must be non-empty.") + if self.category == "timing" and self.unit != "s": + raise ValueError("Timing metric samples must use unit='s'.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@runtime_checkable +class MetricsRecorder(Protocol): + """Collector for runtime metrics.""" + + def record(self, sample: RuntimeMetricSample) -> None: + """Record one metric sample.""" + ... + + def record_timing( + self, + name: str, + duration_s: float, + *, + step_index: int | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> None: + """Record one timing sample in seconds.""" + ... + + def close(self) -> None: + """Finalize metric collection.""" + ... + + +@dataclass(slots=True) +class InMemoryMetricsRecorder: + """Simple metrics recorder useful for tests, smoke runs, and adapters.""" + + samples: list[RuntimeMetricSample] = field(default_factory=list) + closed: bool = False + + def record(self, sample: RuntimeMetricSample) -> None: + if self.closed: + raise RuntimeError("Cannot record metrics after close().") + self.samples.append(sample) + + def record_timing( + self, + name: str, + duration_s: float, + *, + step_index: int | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> None: + self.record( + RuntimeMetricSample( + name=name, + value=duration_s, + unit="s", + step_index=step_index, + category="timing", + metadata={} if metadata is None else metadata, + ) + ) + + def close(self) -> None: + self.closed = True + + +class NullMetricsRecorder: + """Metrics recorder that intentionally drops all samples.""" + + def record(self, sample: RuntimeMetricSample) -> None: + del sample + + def record_timing( + self, + name: str, + duration_s: float, + *, + step_index: int | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> None: + del name, duration_s, step_index, metadata + + def close(self) -> None: + return None diff --git a/flashdreams/flashdreams/runtime/output.py b/flashdreams/flashdreams/runtime/output.py new file mode 100644 index 000000000..aac341ee1 --- /dev/null +++ b/flashdreams/flashdreams/runtime/output.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Output target boundary for generated inference results.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.types import StepResult + + +@dataclass(frozen=True, kw_only=True, slots=True) +class OutputArtifact: + """Artifact produced by an output target.""" + + __hash__ = None + + kind: str + uri: str + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.kind.strip(): + raise ValueError("OutputArtifact.kind must be non-empty.") + if not self.uri.strip(): + raise ValueError("OutputArtifact.uri must be non-empty.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@runtime_checkable +class OutputTarget(Protocol): + """Consumes generated session outputs for presentation or persistence.""" + + def open(self) -> None: + """Prepare the target for a new run.""" + ... + + def write(self, result: StepResult) -> None: + """Consume one generated step result.""" + ... + + def close(self) -> Sequence[OutputArtifact]: + """Finalize and return any produced artifacts.""" + ... + + +@dataclass(slots=True) +class NullOutputTarget: + """Output target for headless runs and throughput measurements.""" + + store_results: bool = False + output_count: int = field(default=0, init=False) + results: list[StepResult] = field(default_factory=list, init=False) + _opened: bool = field(default=False, init=False, repr=False) + + @property + def closed(self) -> bool: + return not self._opened + + def open(self) -> None: + self._opened = True + self.output_count = 0 + self.results.clear() + + def write(self, result: StepResult) -> None: + if not self._opened: + raise RuntimeError("Cannot write to a closed output target.") + self.output_count += 1 + if self.store_results: + self.results.append(result) + + def close(self) -> Sequence[OutputArtifact]: + self._opened = False + return () diff --git a/flashdreams/flashdreams/runtime/types.py b/flashdreams/flashdreams/runtime/types.py new file mode 100644 index 000000000..52bf82166 --- /dev/null +++ b/flashdreams/flashdreams/runtime/types.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Plain data carriers shared by runtime protocols and adapters.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.inputs import ModelInputSchema, TimeWindow + + +@dataclass(frozen=True, kw_only=True, slots=True) +class StepRequest: + """Model-session request for the next step's inputs. + + ``user_input_window`` lets a runner drain or slice timestamped user events for + the current step before invoking the selected ``InputMapping``. + """ + + __hash__ = None + + step_index: int + model_input_schema: ModelInputSchema | None = None + user_input_window: TimeWindow | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.step_index < 0: + raise ValueError("StepRequest.step_index must be >= 0.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class StepResult: + """Generated output and metadata for one inference step.""" + + __hash__ = None + + step_index: int + output: Any = None + frame_count: int | None = None + output_window: TimeWindow | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + metrics: Mapping[str, float | int] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.step_index < 0: + raise ValueError("StepResult.step_index must be >= 0.") + if self.frame_count is not None and self.frame_count < 0: + raise ValueError("StepResult.frame_count must be >= 0.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + object.__setattr__(self, "metrics", freeze_mapping(self.metrics)) diff --git a/flashdreams/tests/test_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py new file mode 100644 index 000000000..1474383a0 --- /dev/null +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -0,0 +1,492 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import fields +from typing import Any, cast + +import pytest + +from flashdreams.runtime import ( + IdentityInputMapping, + InferenceConfig, + InferenceRuntime, + InferenceSession, + InMemoryMetricsRecorder, + InputField, + InputMapping, + MetricsRecorder, + ModelAdapter, + ModelInputs, + ModelInputSchema, + NullOutputTarget, + OutputArtifact, + OutputTarget, + RuntimeMetricSample, + StepRequest, + StepResult, + TimeWindow, + UserInputEvent, + UserInputs, + UserInputSchema, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_inference_config_keeps_runtime_settings_separate() -> None: + denied_app_fields = {"prompt", "output_dir", "browser_settings"} + config = InferenceConfig( + model_id="lingbot-world", + preset_id="fast-taehv", + backend="local", + precision="bf16", + compile=False, + runtime_options={"chunk_size": 3}, + ) + + assert config.model_id == "lingbot-world" + assert config.preset_id == "fast-taehv" + assert config.runtime_options["chunk_size"] == 3 + assert denied_app_fields.isdisjoint(field.name for field in fields(InferenceConfig)) + with pytest.raises(TypeError): + cast(Any, config.runtime_options)["chunk_size"] = 4 + + +def test_inference_config_rejects_empty_model_id() -> None: + with pytest.raises(ValueError, match="model_id"): + InferenceConfig(model_id=" ") + + +@pytest.mark.parametrize( + ("factory", "match"), + [ + (lambda: InputField(name=" "), "InputField.name"), + (lambda: TimeWindow(start_s=1.0, end_s=0.0), "end_s"), + (lambda: TimeWindow(start_s=-1.0, end_s=0.0), "non-negative"), + (lambda: TimeWindow(start_s=0.0, end_s=float("nan")), "finite"), + ( + lambda: UserInputEvent(timestamp_s=-1.0, event_type="keydown"), + "timestamp_s", + ), + (lambda: UserInputEvent(timestamp_s=0.0, event_type=" "), "event_type"), + (lambda: StepRequest(step_index=-1), "step_index"), + (lambda: StepResult(step_index=-1), "step_index"), + (lambda: StepResult(step_index=0, frame_count=-1), "frame_count"), + (lambda: RuntimeMetricSample(name=" ", value=1.0), "name"), + (lambda: RuntimeMetricSample(name="sample", value=float("nan")), "finite"), + (lambda: OutputArtifact(kind=" ", uri="artifact://demo"), "kind"), + (lambda: OutputArtifact(kind="mp4", uri=" "), "uri"), + ], +) +def test_runtime_envelopes_reject_invalid_values(factory: object, match: str) -> None: + with pytest.raises(ValueError, match=match): + cast(Any, factory)() + + +def test_runtime_metric_sample_rejects_bool_values() -> None: + with pytest.raises(TypeError, match="numeric"): + RuntimeMetricSample(name="sample", value=True) + + +def test_model_input_schema_validates_initial_and_step_payloads() -> None: + schema = ModelInputSchema( + initial_fields=( + InputField(name="prompt"), + InputField(name="first_frame"), + ), + step_fields=(InputField(name="camera_poses"),), + ) + inputs = ModelInputs(initial={"prompt": "drive", "first_frame": object()}) + + schema.require_initial(inputs) + assert schema.missing_step(inputs) == ("camera_poses",) + + with pytest.raises(ValueError, match="camera_poses"): + schema.require_step(inputs) + + +def test_user_inputs_filter_timestamped_event_windows() -> None: + inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.1, + event_type="keyboard.keydown", + payload={"key": "w"}, + ), + UserInputEvent( + timestamp_s=0.4, + event_type="keyboard.keyup", + payload={"key": "w"}, + ), + UserInputEvent(timestamp_s=0.8, event_type="reset"), + ) + ) + + windowed = inputs.window(TimeWindow(start_s=0.25, end_s=0.75)) + + assert [event.event_type for event in windowed.events] == ["keyboard.keyup"] + + +def test_user_inputs_require_sorted_events() -> None: + with pytest.raises(ValueError, match="non-decreasing"): + UserInputs( + events=( + UserInputEvent(timestamp_s=1.0, event_type="late"), + UserInputEvent(timestamp_s=0.5, event_type="early"), + ) + ) + + +def test_user_input_schema_declares_event_capabilities() -> None: + schema = UserInputSchema( + event_types=frozenset({"keyboard.keydown", "keyboard.keyup", "reset"}) + ) + + assert schema.supports_event_types(["keyboard.keydown", "reset"]) + assert not schema.supports_event_types(["prompt.update"]) + + +def test_user_input_schema_validates_required_snapshot_fields() -> None: + schema = UserInputSchema( + snapshot_fields=( + InputField(name="pressed_keys"), + InputField(name="prompt", required=False), + ) + ) + inputs = UserInputs(snapshot={"pressed_keys": frozenset({"w"})}) + + schema.require_snapshot(inputs) + assert schema.missing_snapshot(UserInputs()) == ("pressed_keys",) + + with pytest.raises(ValueError, match="pressed_keys"): + schema.require_snapshot(UserInputs()) + + +def test_identity_input_mapping_leaves_model_inputs_unchanged() -> None: + mapping = IdentityInputMapping() + model_inputs = ModelInputs(initial={"prompt": "fixed"}, step={"hdmap": object()}) + request = StepRequest(step_index=0) + + assert ( + mapping.map_initial_inputs( + user_inputs=UserInputs(), + model_inputs=model_inputs, + ) + is model_inputs + ) + assert ( + mapping.map_step_inputs( + user_inputs=UserInputs(), + model_inputs=model_inputs, + request=request, + ) + is model_inputs + ) + + +def test_null_output_target_counts_and_optionally_stores_results() -> None: + target = NullOutputTarget(store_results=True) + result = StepResult(step_index=0, output=b"frame") + + assert target.closed + with pytest.raises(RuntimeError, match="closed output target"): + target.write(result) + + target.open() + assert not target.closed + target.write(result) + artifacts = target.close() + + assert target.closed + assert artifacts == () + assert target.output_count == 1 + assert target.results == [result] + with pytest.raises(RuntimeError, match="closed output target"): + target.write(StepResult(step_index=1)) + + +def test_null_output_target_open_resets_per_run_state() -> None: + target = NullOutputTarget(store_results=True) + + target.open() + target.write(StepResult(step_index=0, output=b"first")) + target.close() + target.open() + + assert target.output_count == 0 + assert target.results == [] + target.write(StepResult(step_index=0, output=b"second")) + assert target.output_count == 1 + assert target.results == [StepResult(step_index=0, output=b"second")] + + +def test_in_memory_metrics_recorder_uses_seconds_for_timing() -> None: + recorder = InMemoryMetricsRecorder() + + recorder.record_timing("model_step", 0.125, step_index=2) + + assert len(recorder.samples) == 1 + sample = recorder.samples[0] + assert sample.name == "model_step" + assert sample.value == pytest.approx(0.125) + assert sample.unit == "s" + assert sample.category == "timing" + assert sample.step_index == 2 + + +def test_timing_metric_samples_must_use_seconds() -> None: + with pytest.raises(ValueError, match="unit='s'"): + RuntimeMetricSample( + name="model_step", + value=12.5, + unit="ms", + category="timing", + ) + + +def test_runtime_api_components_compose_for_sequential_session() -> None: + adapter = _FakeAdapter() + config = InferenceConfig(model_id="fake-model") + user_inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.25, + event_type="keyboard.keydown", + payload={"key": "w"}, + ), + ) + ) + model_inputs = ModelInputs(initial={"prompt": "drive forward"}) + output = NullOutputTarget(store_results=True) + metrics = InMemoryMetricsRecorder() + + adapter.validate_config(config) + mapping = adapter.default_input_mapping() + assert mapping is not None + _drive_two_step_session( + adapter=adapter, + config=config, + mapping=mapping, + user_inputs=user_inputs, + model_inputs=model_inputs, + output=output, + metrics=metrics, + ) + + assert output.output_count == 2 + assert [result.output for result in output.results] == ["chunk-0", "chunk-1"] + assert [result.frame_count for result in output.results] == [3, 3] + assert output.results[0].output_window == TimeWindow(start_s=0.0, end_s=0.5) + assert [sample.step_index for sample in metrics.samples] == [0, 1] + assert metrics.closed + + +def test_reference_loop_validates_mapping_before_runtime_creation() -> None: + mapping = _OrderCheckingMapping() + adapter = _OrderCheckingAdapter(mapping=mapping) + + _drive_two_step_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=mapping, + user_inputs=UserInputs(), + model_inputs=ModelInputs(initial={"prompt": "drive forward"}), + output=NullOutputTarget(), + metrics=InMemoryMetricsRecorder(), + ) + + assert mapping.validated + assert adapter.created_runtime_after_validate + + +def test_reference_loop_closes_runtime_when_session_start_fails() -> None: + adapter = _FailingStartAdapter() + output = NullOutputTarget() + metrics = InMemoryMetricsRecorder() + + with pytest.raises(RuntimeError, match="start failed"): + _drive_two_step_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=IdentityInputMapping(), + user_inputs=UserInputs(), + model_inputs=ModelInputs(initial={"prompt": "drive forward"}), + output=output, + metrics=metrics, + ) + + assert adapter.runtime is not None + assert adapter.runtime.closed + assert output.closed + assert metrics.closed + + +def _drive_two_step_session( + *, + adapter: ModelAdapter, + config: InferenceConfig, + mapping: InputMapping, + user_inputs: UserInputs, + model_inputs: ModelInputs, + output: OutputTarget, + metrics: MetricsRecorder, +) -> None: + mapping.validate( + user_schema=adapter.user_input_schema, + model_schema=adapter.model_input_schema, + ) + initial_inputs = mapping.map_initial_inputs( + user_inputs=user_inputs, + model_inputs=model_inputs, + ) + runtime = adapter.create_runtime(config) + session: InferenceSession | None = None + output_opened = False + try: + session = runtime.start_session(initial_inputs) + output.open() + output_opened = True + while (request := session.next_step_request()) is not None: + step_inputs = mapping.map_step_inputs( + user_inputs=( + user_inputs.window(request.user_input_window) + if request.user_input_window is not None + else user_inputs + ), + model_inputs=ModelInputs( + initial=initial_inputs.initial, + step={"chunk_index": request.step_index}, + ), + request=request, + ) + result = session.step(step_inputs) + output.write(result) + metrics.record_timing( + "model_step", + float(result.metrics["model_step_s"]), + step_index=result.step_index, + ) + finally: + if output_opened: + output.close() + if session is not None: + session.close() + runtime.close() + metrics.close() + + +class _FakeAdapter: + model_id = "fake-model" + model_input_schema = ModelInputSchema( + initial_fields=(InputField(name="prompt"),), + step_fields=(InputField(name="chunk_index"),), + ) + user_input_schema = UserInputSchema(event_types=frozenset({"keyboard.keydown"})) + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError(f"Unsupported model_id={config.model_id!r}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + return _FakeRuntime(model_input_schema=self.model_input_schema) + + +class _FakeRuntime: + def __init__(self, *, model_input_schema: ModelInputSchema) -> None: + self._model_input_schema = model_input_schema + self.closed = False + + def start_session(self, inputs: ModelInputs) -> InferenceSession: + self._model_input_schema.require_initial(inputs) + return _FakeSession(model_input_schema=self._model_input_schema) + + def close(self) -> None: + self.closed = True + + +class _FailingRuntime(_FakeRuntime): + def start_session(self, inputs: ModelInputs) -> InferenceSession: + del inputs + raise RuntimeError("start failed") + + +class _FakeSession: + def __init__(self, *, model_input_schema: ModelInputSchema) -> None: + self._model_input_schema = model_input_schema + self.step_index = 0 + self.closed = False + + def next_step_request(self) -> StepRequest | None: + if self.step_index >= 2: + return None + return StepRequest( + step_index=self.step_index, + model_input_schema=self._model_input_schema, + user_input_window=TimeWindow( + start_s=0.5 * self.step_index, + end_s=0.5 * (self.step_index + 1), + ), + ) + + def step(self, inputs: ModelInputs) -> StepResult: + self._model_input_schema.require_step(inputs) + result = StepResult( + step_index=self.step_index, + output=f"chunk-{self.step_index}", + frame_count=3, + output_window=TimeWindow( + start_s=0.5 * self.step_index, + end_s=0.5 * (self.step_index + 1), + ), + metrics={"model_step_s": 0.01}, + ) + self.step_index += 1 + return result + + def reset(self, inputs: ModelInputs | None = None) -> None: + del inputs + self.step_index = 0 + + def close(self) -> None: + self.closed = True + + +class _OrderCheckingMapping(IdentityInputMapping): + def __init__(self) -> None: + self.validated = False + + def validate( + self, + *, + user_schema: UserInputSchema | None = None, + model_schema: ModelInputSchema | None = None, + ) -> None: + super().validate(user_schema=user_schema, model_schema=model_schema) + self.validated = True + + +class _OrderCheckingAdapter(_FakeAdapter): + def __init__(self, *, mapping: _OrderCheckingMapping) -> None: + self._mapping = mapping + self.created_runtime_after_validate = False + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + self.created_runtime_after_validate = self._mapping.validated + return _FakeRuntime(model_input_schema=self.model_input_schema) + + +class _FailingStartAdapter(_FakeAdapter): + def __init__(self) -> None: + self.runtime: _FailingRuntime | None = None + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + self.runtime = _FailingRuntime(model_input_schema=self.model_input_schema) + return self.runtime