diff --git a/.nvim.lua b/.nvim.lua new file mode 100644 index 000000000..ec0ca6948 --- /dev/null +++ b/.nvim.lua @@ -0,0 +1,3 @@ +vim.lsp.config("pyright", { + root_markers = { ".git" }, +}) diff --git a/docs/inference_runtime_inputs_implementation.md b/docs/inference_runtime_inputs_implementation.md index 75460cced..edbc7399d 100644 --- a/docs/inference_runtime_inputs_implementation.md +++ b/docs/inference_runtime_inputs_implementation.md @@ -11,7 +11,10 @@ what is intentionally still outside this layer. Implementation lives in `flashdreams.runtime`: -- `flashdreams/flashdreams/runtime/inputs.py` — the input types and schemas +- `flashdreams/flashdreams/runtime/inputs.py` — user/canonical input types + and schemas +- `flashdreams/flashdreams/runtime/inference_session.py` — model-ready + `InferenceInput`, `InferenceInputSchema`, and session lifecycle - `flashdreams/flashdreams/runtime/canonical.py` — raw device to canonical modality conversion - `flashdreams/flashdreams/runtime/mapping.py` — canonical to encoded mapping @@ -266,7 +269,7 @@ Named here only so the boundary is explicit; these are not gaps in the input layer: - **`FrameStream`**, which the architecture diagrams place between - `InferenceSession` and `Output Target`. The code writes `StepResult` straight + `InferenceSession` and `Output Target`. The code writes `InferenceOutput` straight to `OutputTarget.write()`. Output shape is T5. - **Declared output modalities**, so an output target or quality-eval can state what it requires and be matched the way inputs now are. T5/T8. diff --git a/docs/source/_static/diagrams/flashdreams-runtime-data-flow.png b/docs/source/_static/diagrams/flashdreams-runtime-data-flow.png new file mode 100644 index 000000000..d94ac7357 Binary files /dev/null and b/docs/source/_static/diagrams/flashdreams-runtime-data-flow.png differ diff --git a/docs/source/_static/diagrams/flashdreams-runtime.png b/docs/source/_static/diagrams/flashdreams-runtime.png new file mode 100644 index 000000000..7178cc6af Binary files /dev/null and b/docs/source/_static/diagrams/flashdreams-runtime.png differ diff --git a/docs/source/developer_guides/flashdreams_runtime.rst b/docs/source/developer_guides/flashdreams_runtime.rst new file mode 100644 index 000000000..85159ec20 --- /dev/null +++ b/docs/source/developer_guides/flashdreams_runtime.rst @@ -0,0 +1,307 @@ +.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +.. SPDX-License-Identifier: Apache-2.0 +.. +.. Licensed under the Apache License, Version 2.0 (the "License"); +.. you may not use this file except in compliance with the License. +.. You may obtain a copy of the License at +.. +.. http://www.apache.org/licenses/LICENSE-2.0 +.. +.. Unless required by applicable law or agreed to in writing, software +.. distributed under the License is distributed on an "AS IS" BASIS, +.. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +.. See the License for the specific language governing permissions and +.. limitations under the License. + +FlashDreams runtime architecture +================================ + +The FlashDreams runtime connects interactive controls to an autoregressive +inference pipeline and routes the generated frames to an output destination. +This page defines the target architecture and the boundaries between its +components. It is a design contract rather than a reference for an existing +Python API. + +Architecture overview +--------------------- + +.. figure:: /_static/diagrams/flashdreams-runtime.png + :alt: FlashDreams runtime architecture and its input and output data flow. + :align: center + :width: 100% + :class: zoomable + + The application owns the runtime components. Dashed arrows show data flow; + solid lines with diamonds show ownership. Solid boxes are classes and + dashed boxes are data. + +.. image:: /_static/diagrams/flashdreams-runtime-data-flow.png + :alt: FlashDreams user and global conditioning data flow into an inference session. + :align: center + :width: 70% + :class: zoomable + +Application layer +----------------- + +``Application`` is the composition and lifecycle boundary for an interactive +FlashDreams runtime. It owns ``InputSystem``, ``InputMapping``, +``OutputTarget``, and the main ``InferenceSession`` that runs the inference +pipeline. These components are passed to the application through dependency +injection. The application connects them, drives the runtime loop, and shuts +them down in a defined order. + +Keeping orchestration in the application gives each child component a narrow +responsibility. Device handling stays out of model execution, model-specific +input conversion stays out of device handling, and presentation stays out of +the inference session. + +InputSystem +~~~~~~~~~~~ + +``InputSystem`` accepts an ordered list of timestamped raw input events +supplied by the application or an upstream input framework. It converts that +list into an ordered stream of timestamped, canonicalized user-input events. +Raw events can include keyboard key-down and key-up events, digital wheel +readings, controller joystick readings, or Meta Quest hand-tracking readings. + +`Unity's Input System +`_ +provides similar concepts for devices and actions. FlashDreams has a narrower +boundary: it does not poll devices or configure key bindings. Device polling, +event collection, and binding configuration are handled upstream. + +.. admonition:: Preserving events during slow inference + :class: note + + A call to ``InferenceSession.step()`` can take approximately 100--1000 ms + because it runs a latent-diffusion step. The upstream input source must + preserve every raw input change that occurs while inference is running + rather than retain only the most recent device state. On the next runtime + iteration, the application passes the complete ordered raw-event list to + ``InputSystem``, which converts every event without collapsing intermediate + states. + + For example, assume positive *x* means right and positive *y* means forward. + The raw-event list contains a W key-down at 5 ms, a D key-down at 10 ms, a W + key-up at 50 ms, an S key-down at 60 ms, a D key-up at 70 ms, and an S + key-up at 80 ms. ``InputSystem`` produces this canonical event stream: + + .. code-block:: text + + [ + ( 5 ms, vec2(0, 1)), # W pressed + (10 ms, vec2(1, 1)), # D pressed; W remains pressed + (50 ms, vec2(1, 0)), # W released + (60 ms, vec2(1, -1)), # S pressed; D remains pressed + (70 ms, vec2(0, -1)), # D released + (80 ms, vec2(0, 0)), # S released + ] + + The timestamps are the original raw-event times, not the time at which + ``InputSystem`` processes the list. + +This layer handles raw-to-canonical conversion concerns such as device-value +normalization, dead zones, axis conventions, and event ordering. Its output +describes the user's intent in a stable, device-independent form. It does not +create model embeddings or know how a particular inference pipeline represents +conditioning. + +InputMapping +~~~~~~~~~~~~ + +``InputMapping`` consumes the ordered list of timestamped, canonicalized +user-input events returned by ``InputSystem`` and produces the model-ready, +per-step inference conditioning expected by an ``InferenceSession``. Depending +on the model, this conversion can include embedding control values, rendering +a control representation, changing layouts, or assembling tensors. + +.. admonition:: Example + :class: note + + ``integrations/omnidreams`` represents canonicalized driving input as a + floating-point steering-wheel angle and a floating-point paddle/brake value. + Its ``InputMapping`` runs the vehicle-dynamics simulation, renders the + resulting HD map with the Ludus renderer, and produces the rendered RGB HD + map as the per-step user-input condition. The resulting tensor has shape + ``(3, H, W)``. + +The ``(3, H, W)`` output is specific to OmniDreams, not a universal +``InputMapping`` contract. Another ``InferenceSession`` might expect an image +embedding as its condition. In that case, ``InputMapping`` can use an image +encoder to encode the frame and return the resulting embedding instead. + +This is the boundary between application-level control semantics and +model-specific conditioning. Replacing a keyboard with a controller should +usually affect the ``InputSystem``; replacing the model or its control encoder +should usually affect the ``InputMapping``. + +OutputTarget +~~~~~~~~~~~~ + +``OutputTarget`` consumes the ``FrameStream`` produced by the inference +session. A target can present frames in a native window, send them to a video +encoder, publish them through a WebRTC host, or adapt them for another output +system. + +The output target owns presentation and transport concerns. It must not be +responsible for interpreting user controls or advancing model inference. +Buffering and backpressure policies belong at this output boundary so that a +slow consumer does not silently redefine inference behavior. + +InferenceSession +~~~~~~~~~~~~~~~~ + +``InferenceSession`` is the execution boundary for the main inference +pipeline. It accepts inference input, maintains the state required across +autoregressive steps, runs the pipeline, and exposes generated output as a +``FrameStream``. + +The session receives model-ready data only. It does not poll devices, +canonicalize user intent, or present generated frames. After accepting global +inference conditioning, it retains the active global condition across later +steps until the application supplies an update or the session ends. + +Input data flow +--------------- + +The input path deliberately separates physical device readings, semantic +controls, and model-ready conditioning. This separation allows devices and +models to evolve independently. + +User conditioning +~~~~~~~~~~~~~~~~~ + +Raw user input +^^^^^^^^^^^^^^ + +**Raw user input** is a reading or event in the vocabulary of a physical input +source. Examples include: + +* WASD key presses and releases; +* digital wheel input; +* controller joystick readings; and +* Meta Quest hand-tracking readings. + +Raw events can depend on a particular device, driver, and upstream input +framework. They are collected and timestamped before entering FlashDreams. The +application passes the ordered raw-event list to ``InputSystem``; raw events +must not be passed directly to ``InferenceSession``. + +Canonicalized user input +^^^^^^^^^^^^^^^^^^^^^^^^ + +**Canonicalized user input** expresses user intent in the vocabulary of the +interaction, independent of the device that produced it. Typical structures +include: + +* a 2D character-movement vector and a 2D camera-movement vector for + character-control games; +* a floating-point wheel value and a floating-point paddle value for driving; + and +* hand-tracking positions, correction vectors, or another agreed semantic hand + control for Cosmos-style interaction models. + +The exact structure depends on the interaction type and can evolve as its +semantics become clearer. ``InputSystem`` emits these values as an ordered +stream of timestamped, canonicalized events. The important invariant is that +equivalent intent from different devices has the same canonical +representation. + +Per-step inference conditioning +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Per-step inference conditioning** is the model-ready encoding of +canonicalized user input for one inference step. ``InputMapping`` produces it by +performing whatever embedding, rendering, or tensor conversion the selected +model requires. + +This term distinguishes the changing user control for one step from global +conditioning, which normally remains stable across many steps. + +Inference input +^^^^^^^^^^^^^^^ + +**Inference input** is the complete input delivered to ``InferenceSession``. +It is the runtime boundary object, not another name for a raw or canonicalized +control. It can carry: + +* the per-step inference conditioning; and +* optional global inference conditioning. + +The first inference step generally carries both. On later steps, the global +condition remains active inside the session, so the application normally sends +only new per-step inference conditioning. Omitting global conditioning means +"continue using the active global condition"; it must not mean "clear the +global condition." + +Global conditioning +~~~~~~~~~~~~~~~~~~~ + +Global conditioning establishes the scene-level context for generation and can +contain model-specific data. Two of the most common examples are: + +* a **global conditioning frame**, sometimes called an initial frame by a + model; and +* a **global conditioning prompt**, containing the text description for the + run. + +The runtime uses *global conditioning frame* instead of *initial frame* +because the condition is not inherently limited to initialization. A future +runtime can replace it while a session is already running. + +Raw and canonicalized global conditioning +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Raw global conditioning** is the application-facing representation. For +example, the prompt is text and the conditioning frame is image data. + +**Canonicalized global conditioning** is the model-ready representation sent +through inference input as **global inference conditioning**. A text prompt is +typically converted into embedded tokens. A conditioning frame is +model-dependent: one model might convert it into CLIP embeddings, while +another might retain a frame or spatial representation such as an HD-map +condition. + +For that reason, *canonicalized* is preferred over *embedded* for the combined +global condition. It does not incorrectly imply that every part of the global +condition must become an embedding. + +Updating global conditioning during a run +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Although mid-run global-conditioning updates are not implemented yet, the +runtime contracts must leave room for them. An application could, for example, +submit a new conditioning frame and prompt to make an OmniDreams driving scene +transition suddenly to rainy weather. + +A later inference input can therefore carry new global inference conditioning +alongside its per-step conditioning. The session then treats it as a change to +the active global context. When no update is present, the session reuses the +previous context. The exact effects on model history, caches, and transition +behavior are model-specific and must be defined by the corresponding pipeline; +the application-level contract must not assume that global conditioning is +initialization-only. + +End-to-end runtime loop +----------------------- + +At a conceptual level, one runtime iteration follows these steps: + +#. ``Application`` passes ``InputSystem`` the ordered list of timestamped + raw input events accumulated upstream since the previous runtime iteration. +#. ``InputSystem`` converts that list into an ordered stream of timestamped, + canonicalized user-input events. +#. ``InputMapping`` consumes this canonical event stream and produces per-step + inference conditioning. +#. ``Application`` packages that data as inference input, adding canonicalized + global conditioning on the first step or whenever it changes. +#. ``InferenceSession`` advances the pipeline and emits generated frames + through ``FrameStream``. +#. ``OutputTarget`` consumes the stream for display, encoding, transport, or + another presentation path. + +These boundaries are the central architectural constraint: an upstream input +source collects timestamped raw events, ``InputSystem`` produces canonical +events, ``InputMapping`` produces model-ready conditioning, +``InferenceSession`` runs the model, and ``OutputTarget`` delivers the result. diff --git a/docs/source/developer_guides/index.rst b/docs/source/developer_guides/index.rst index 7cc2bac85..7b7b50573 100644 --- a/docs/source/developer_guides/index.rst +++ b/docs/source/developer_guides/index.rst @@ -60,6 +60,7 @@ generated clip, see :doc:`/quickstart/index`. :hidden: :maxdepth: 1 + flashdreams_runtime inference_pipeline_overview config_system new_integration diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index 5f89196ba..01a2f9c76 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -17,13 +17,17 @@ ScriptedModality, ) from flashdreams.runtime.config import ExecutionBackend, InferenceConfig, Precision +from flashdreams.runtime.inference_session import ( + InferenceInput, + InferenceInputSchema, + InferenceOutput, + InferenceSessionConfig, +) from flashdreams.runtime.inputs import ( INPUT_PHASES, CanonicalInputs, CanonicalInputSchema, CanonicalModality, - InferenceInput, - InferenceInputSchema, InputField, InputPhase, TimeWindow, @@ -77,8 +81,10 @@ "InferenceConfig", "InferenceInput", "InferenceInputSchema", + "InferenceOutput", "InferenceRuntime", "InferenceSession", + "InferenceSessionConfig", "InMemoryMetricsRecorder", "INPUT_PHASES", "InputCanonicalizer", diff --git a/flashdreams/flashdreams/runtime/inference_session.py b/flashdreams/flashdreams/runtime/inference_session.py new file mode 100644 index 000000000..50a07acc2 --- /dev/null +++ b/flashdreams/flashdreams/runtime/inference_session.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Inference session lifecycle, model-input envelope, and schema.""" + +from abc import ABC, abstractmethod +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +from flashdreams.infra.pipeline import ( + StreamInferencePipeline, + StreamInferencePipelineCache, + StreamInferencePipelineConfig, +) +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.inputs import InputField, TimeWindow, check_payload +from flashdreams.runtime.types import StepRequest + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InferenceInput: + """Global and per-step conditioning for one inference call.""" + + __hash__ = None + + global_conditioning: Mapping[str, Any] = field(default_factory=dict) + per_step_conditioning: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__( + self, "global_conditioning", freeze_mapping(self.global_conditioning) + ) + object.__setattr__( + self, + "per_step_conditioning", + freeze_mapping(self.per_step_conditioning), + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InferenceOutput: + """Generated output and metadata for one inference step.""" + + __hash__ = None + + step_index: int + """Zero-based index of the completed inference step.""" + + output: Any = None + """Generated payload for the step.""" + + frame_count: int | None = None + """Number of generated frames when the output is frame-based.""" + + output_window: TimeWindow | None = None + """Session time window represented by the generated output.""" + + metadata: Mapping[str, Any] = field(default_factory=dict) + """Output metadata supplied by the session or model adapter.""" + + metrics: Mapping[str, float | int] = field(default_factory=dict) + """Per-step numeric measurements.""" + + def __post_init__(self) -> None: + if self.step_index < 0: + raise ValueError("InferenceOutput.step_index must be >= 0.") + if self.frame_count is not None and self.frame_count < 0: + raise ValueError("InferenceOutput.frame_count must be >= 0.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + object.__setattr__(self, "metrics", freeze_mapping(self.metrics)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InferenceInputSchema: + """Required and optional fields in each inference input payload.""" + + global_fields: tuple[InputField, ...] = () + """Model inputs required before starting the initial generation/session.""" + + per_step_fields: tuple[InputField, ...] = () + """Per-step model inputs required after the session starts.""" + + def check_global_payload(self, inputs: InferenceInput) -> None: + """Check that required global fields are present in ``inputs``.""" + check_payload(self.global_fields, inputs.global_conditioning) + + def check_per_step_payload(self, inputs: InferenceInput) -> None: + """Check that required per-step fields are present in ``inputs``.""" + check_payload(self.per_step_fields, inputs.per_step_conditioning) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InferenceSessionConfig: + """Configuration for constructing an inference session.""" + + __hash__ = None + + pipeline: StreamInferencePipelineConfig + """Pipeline configuration to instantiate.""" + + +class InferenceSession(ABC): + """Stateful inference pipeline session.""" + + _pipeline_cache: StreamInferencePipelineCache[Any, Any, Any] | None + """Pipeline cache for the active rollout; ``None`` before its first step.""" + + _step_index: int + """Zero-based index assigned to the next generated output.""" + + def __init__(self, config: InferenceSessionConfig) -> None: + """Initialize the inference pipeline. + + Args: + config: Session configuration. + """ + self.config = config + # Initialize the inference pipeline from the provided configuration. + self.pipeline: StreamInferencePipeline = self.config.pipeline.setup() + self._pipeline_cache = None + self._step_index = 0 + + def __del__(self) -> None: + """Release session resources.""" + if hasattr(self, "pipeline"): + del self.pipeline + + def next_step_request(self) -> StepRequest: + """Return input requirements for the next pipeline step.""" + return StepRequest(step_index=self._step_index) + + @abstractmethod + def reset(self) -> None: + """Reset the pipeline and discard the active rollout state.""" + pipeline_reset = getattr(self.pipeline, "reset", None) + if callable(pipeline_reset): + pipeline_reset() + self._pipeline_cache = None + self._step_index = 0 + + @abstractmethod + def step(self, inference_input: InferenceInput) -> InferenceOutput: + """Run one inference step. + + Args: + inference_input: Model-ready inputs for the step. + + Returns: + Generated output for the step. + + Raises: + ValueError: Global conditioning is supplied after the rollout starts. + """ + request = self.next_step_request() + input_schema = request.inference_input_schema + if self._pipeline_cache is None: + if input_schema is not None: + input_schema.check_global_payload(inference_input) + self._pipeline_cache = self.pipeline.initialize_cache( + **inference_input.global_conditioning + ) + elif inference_input.global_conditioning: + raise ValueError( + "InferenceInput.global_conditioning can only be supplied on the " + "first step after reset()." + ) + + if input_schema is not None: + input_schema.check_per_step_payload(inference_input) + pipeline_input = inference_input.per_step_conditioning or None + output = self.pipeline.generate( + autoregressive_index=request.step_index, + cache=self._pipeline_cache, + input=pipeline_input, + ) + self.pipeline.finalize( + autoregressive_index=request.step_index, + cache=self._pipeline_cache, + ) + inference_output = InferenceOutput( + step_index=request.step_index, + output=output, + output_window=request.user_input_window, + metadata=request.metadata, + metrics={}, # No metrics for now, inject to pipeline later. + ) + self._step_index = request.step_index + 1 + return inference_output diff --git a/flashdreams/flashdreams/runtime/inputs.py b/flashdreams/flashdreams/runtime/inputs.py index 9174b6a84..4f7cec75a 100644 --- a/flashdreams/flashdreams/runtime/inputs.py +++ b/flashdreams/flashdreams/runtime/inputs.py @@ -1,7 +1,7 @@ # 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.""" +"""User and canonical input envelopes and schemas for the runtime API.""" from __future__ import annotations @@ -189,7 +189,11 @@ def validate_event(self, event: "UserInputEvent") -> None: def missing_snapshot(self, inputs: "UserInputs") -> tuple[str, ...]: """Return required snapshot fields absent from ``inputs``.""" - return _missing_required(self.snapshot_fields, inputs.snapshot) + return tuple( + input_field.name + for input_field in self.snapshot_fields + if input_field.required and input_field.name not in inputs.snapshot + ) def require_snapshot(self, inputs: "UserInputs") -> None: """Raise if required snapshot fields are absent.""" @@ -354,7 +358,8 @@ class CanonicalModality: Modalities describe live user control only. Global conditioning such as a prompt or conditioning frame is application-owned and reaches - :class:`InferenceInput` directly, without passing through this layer. + :class:`flashdreams.runtime.inference_session.InferenceInput` directly, + without passing through this layer. """ name: str @@ -409,8 +414,8 @@ class CanonicalInputs: Values are level-triggered and normally present every step: a key held down emits no events but still means full throttle. Global conditioning does not - appear here; it is application-owned and reaches :class:`InferenceInput` - directly. + appear here; it is application-owned and reaches + :class:`flashdreams.runtime.inference_session.InferenceInput` directly. """ __hash__ = None @@ -466,3 +471,5 @@ def _missing_required( for input_field in fields if input_field.required and input_field.name not in payload ) + if missing: + raise ValueError(f"Missing required input(s): {missing}") diff --git a/flashdreams/flashdreams/runtime/interfaces.py b/flashdreams/flashdreams/runtime/interfaces.py index 5c5054dc4..385e69d62 100644 --- a/flashdreams/flashdreams/runtime/interfaces.py +++ b/flashdreams/flashdreams/runtime/interfaces.py @@ -8,13 +8,14 @@ from typing import Protocol, runtime_checkable from flashdreams.runtime.config import InferenceConfig -from flashdreams.runtime.inputs import ( - CanonicalInputSchema, +from flashdreams.runtime.inference_session import ( InferenceInput, InferenceInputSchema, + InferenceOutput, ) +from flashdreams.runtime.inputs import CanonicalInputSchema from flashdreams.runtime.mapping import InputMapping -from flashdreams.runtime.types import StepRequest, StepResult +from flashdreams.runtime.types import StepRequest @runtime_checkable @@ -25,7 +26,7 @@ def next_step_request(self) -> StepRequest | None: """Return the next step's runtime request, or ``None`` when complete.""" ... - def step(self, inputs: InferenceInput) -> StepResult: + def step(self, inputs: InferenceInput) -> InferenceOutput: """Run one sequential inference step.""" ... diff --git a/flashdreams/flashdreams/runtime/mapping.py b/flashdreams/flashdreams/runtime/mapping.py index 6dfb5cc45..e6e2b9b3e 100644 --- a/flashdreams/flashdreams/runtime/mapping.py +++ b/flashdreams/flashdreams/runtime/mapping.py @@ -10,13 +10,12 @@ from typing import Any, Protocol, runtime_checkable from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.inference_session import InferenceInput, InferenceInputSchema from flashdreams.runtime.inputs import ( INPUT_PHASES, CanonicalInputs, CanonicalInputSchema, CanonicalModality, - InferenceInput, - InferenceInputSchema, InputField, InputPhase, ) @@ -271,7 +270,15 @@ def _build_compatibility( unavailable.append(mapping_schema) usable = combine_mapping_schemas(feedable, name=reported_schema.name) - required = inference_input_schema.required_fields() + declared_fields: tuple[tuple[InputPhase, InputField], ...] = tuple( + (phase, input_field) + for phase, fields in ( + ("global", inference_input_schema.global_fields), + ("step", inference_input_schema.per_step_fields), + ) + for input_field in fields + ) + required = tuple(declared for declared in declared_fields if declared[1].required) missing_required = tuple( (phase, input_field) for phase, input_field in required @@ -284,7 +291,8 @@ def _build_compatibility( ) available_optional = tuple( (phase, input_field) - for phase, input_field in inference_input_schema.optional_fields() + for phase, input_field in declared_fields + if not input_field.required if usable.can_produce(phase, input_field) ) @@ -364,10 +372,14 @@ def undeclared_inference_inputs( Mapping tests can use this to keep the declared compatibility surface honest. """ + payloads: tuple[tuple[InputPhase, Mapping[str, Any]], ...] = ( + ("global", inputs.global_conditioning), + ("step", inputs.per_step_conditioning), + ) return tuple( (phase, key) - for phase in INPUT_PHASES - for key in inputs.for_phase(phase) + for phase, payload in payloads + for key in payload if not any( declared.name == key for declared in mapping_schema.produces_for(phase) ) diff --git a/flashdreams/flashdreams/runtime/output.py b/flashdreams/flashdreams/runtime/output.py index aac341ee1..675b04d60 100644 --- a/flashdreams/flashdreams/runtime/output.py +++ b/flashdreams/flashdreams/runtime/output.py @@ -10,7 +10,7 @@ from typing import Any, Protocol, runtime_checkable from flashdreams.runtime._utils import freeze_mapping -from flashdreams.runtime.types import StepResult +from flashdreams.runtime.inference_session import InferenceOutput @dataclass(frozen=True, kw_only=True, slots=True) @@ -39,7 +39,7 @@ def open(self) -> None: """Prepare the target for a new run.""" ... - def write(self, result: StepResult) -> None: + def write(self, result: InferenceOutput) -> None: """Consume one generated step result.""" ... @@ -54,7 +54,7 @@ class NullOutputTarget: store_results: bool = False output_count: int = field(default=0, init=False) - results: list[StepResult] = field(default_factory=list, init=False) + results: list[InferenceOutput] = field(default_factory=list, init=False) _opened: bool = field(default=False, init=False, repr=False) @property @@ -66,7 +66,7 @@ def open(self) -> None: self.output_count = 0 self.results.clear() - def write(self, result: StepResult) -> None: + def write(self, result: InferenceOutput) -> None: if not self._opened: raise RuntimeError("Cannot write to a closed output target.") self.output_count += 1 diff --git a/flashdreams/flashdreams/runtime/types.py b/flashdreams/flashdreams/runtime/types.py index 51d3846db..ecf045166 100644 --- a/flashdreams/flashdreams/runtime/types.py +++ b/flashdreams/flashdreams/runtime/types.py @@ -7,10 +7,13 @@ from collections.abc import Mapping from dataclasses import dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any from flashdreams.runtime._utils import freeze_mapping -from flashdreams.runtime.inputs import InferenceInputSchema, TimeWindow +from flashdreams.runtime.inputs import TimeWindow + +if TYPE_CHECKING: + from flashdreams.runtime.inference_session import InferenceInputSchema @dataclass(frozen=True, kw_only=True, slots=True) diff --git a/flashdreams/tests/test_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py index 42f75d688..9f8091512 100644 --- a/flashdreams/tests/test_inference_runtime_api.py +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -26,6 +26,9 @@ UserInputs, UserInputSchema, ) +from flashdreams.runtime.inference_session import ( + InferenceSession as PipelineInferenceSession, +) pytestmark = pytest.mark.ci_cpu @@ -67,8 +70,8 @@ def test_inference_config_rejects_empty_model_id() -> None: ), (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: InferenceOutput(step_index=-1), "step_index"), + (lambda: InferenceOutput(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"), @@ -91,17 +94,49 @@ def test_schema_validates_global_conditioning_and_step_payloads() -> None: InputField(name="prompt"), InputField(name="global_conditioning_frame"), ), - step_fields=(InputField(name="camera_poses"),), + per_step_fields=(InputField(name="camera_poses"),), ) inputs = InferenceInput( - global_conditioning={"prompt": "drive", "global_conditioning_frame": object()} + global_conditioning={ + "prompt": "drive", + "global_conditioning_frame": object(), + }, + per_step_conditioning={"camera_poses": object()}, ) schema.require_global_conditioning(inputs) assert schema.missing_step(inputs) == ("camera_poses",) + with pytest.raises(ValueError, match="prompt"): + schema.check_global_payload(InferenceInput()) with pytest.raises(ValueError, match="camera_poses"): - schema.require_step(inputs) + schema.check_per_step_payload(InferenceInput()) + + +def test_inference_input_and_schema_only_declare_two_fields() -> None: + assert tuple(InferenceInput.__dataclass_fields__) == ( + "global_conditioning", + "per_step_conditioning", + ) + assert tuple(InferenceInputSchema.__dataclass_fields__) == ( + "global_fields", + "per_step_fields", + ) + + +def test_inference_output_matches_step_result_fields() -> None: + assert tuple(field.name for field in fields(InferenceOutput)) == tuple( + field.name for field in fields(StepResult) + ) + + +def test_inference_session_uses_step_requests_instead_of_a_schema_property() -> None: + assert "inference_input_schema" not in PipelineInferenceSession.__dict__ + assert callable(PipelineInferenceSession.next_step_request) + + +def test_pipeline_inference_session_requires_reset_and_step_implementations() -> None: + assert PipelineInferenceSession.__abstractmethods__ == frozenset({"reset", "step"}) def test_user_inputs_filter_timestamped_event_windows() -> None: @@ -164,7 +199,8 @@ def test_user_input_schema_validates_required_snapshot_fields() -> None: def test_identity_input_mapping_leaves_inference_input_unchanged() -> None: mapping = IdentityInputMapping() inference_input = InferenceInput( - global_conditioning={"prompt": "fixed"}, step={"hdmap": object()} + global_conditioning={"prompt": "fixed"}, + per_step_conditioning={"hdmap": object()}, ) request = StepRequest(step_index=0) @@ -187,7 +223,7 @@ def test_identity_input_mapping_leaves_inference_input_unchanged() -> None: def test_null_output_target_counts_and_optionally_stores_results() -> None: target = NullOutputTarget(store_results=True) - result = StepResult(step_index=0, output=b"frame") + result = InferenceOutput(step_index=0, output=b"frame") assert target.closed with pytest.raises(RuntimeError, match="closed output target"): @@ -203,22 +239,22 @@ def test_null_output_target_counts_and_optionally_stores_results() -> None: assert target.output_count == 1 assert target.results == [result] with pytest.raises(RuntimeError, match="closed output target"): - target.write(StepResult(step_index=1)) + target.write(InferenceOutput(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.write(InferenceOutput(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")) + target.write(InferenceOutput(step_index=0, output=b"second")) assert target.output_count == 1 - assert target.results == [StepResult(step_index=0, output=b"second")] + assert target.results == [InferenceOutput(step_index=0, output=b"second")] def test_in_memory_metrics_recorder_uses_seconds_for_timing() -> None: diff --git a/flashdreams/tests/test_inference_session.py b/flashdreams/tests/test_inference_session.py new file mode 100644 index 000000000..261b53843 --- /dev/null +++ b/flashdreams/tests/test_inference_session.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for inference-session pipeline orchestration.""" + +from typing import Any, cast + +import pytest +import torch + +from flashdreams.infra.pipeline import ( + StreamInferencePipeline, + StreamInferencePipelineConfig, +) +from flashdreams.runtime.inference_session import ( + InferenceInput, + InferenceInputSchema, + InferenceOutput, + InferenceSession, + InferenceSessionConfig, +) +from flashdreams.runtime.inputs import TimeWindow +from flashdreams.runtime.types import StepRequest + +pytestmark = pytest.mark.ci_cpu + + +class FakeStreamInferencePipeline(StreamInferencePipeline[Any, Any, Any]): + """Record session orchestration without constructing model components.""" + + def __init__(self, output: Any) -> None: + torch.nn.Module.__init__(self) + self.output = output + self.cache = object() + self.reset_calls = 0 + self.initialize_cache_calls: list[dict[str, Any]] = [] + self.generate_calls: list[dict[str, Any]] = [] + self.finalize_calls: list[dict[str, Any]] = [] + + def reset(self) -> None: + self.reset_calls += 1 + + def initialize_cache(self, **global_conditioning: Any) -> object: + self.initialize_cache_calls.append(global_conditioning) + self.cache = object() + return self.cache + + def generate( + self, autoregressive_index: int, cache: object, input: Any = None + ) -> Any: + self.generate_calls.append( + { + "autoregressive_index": autoregressive_index, + "cache": cache, + "input": input, + } + ) + return self.output + + def finalize(self, autoregressive_index: int, cache: object) -> dict[str, float]: + self.finalize_calls.append( + {"autoregressive_index": autoregressive_index, "cache": cache} + ) + return {"total_ms": 4.0} + + +class _FakePipelineConfig: + def __init__(self, pipeline: FakeStreamInferencePipeline) -> None: + self.pipeline = pipeline + self.setup_calls = 0 + + def setup(self) -> FakeStreamInferencePipeline: + self.setup_calls += 1 + return self.pipeline + + +class _ConcreteInferenceSession(InferenceSession): + def next_step_request(self) -> StepRequest: + return StepRequest( + step_index=self._step_index, + inference_input_schema=InferenceInputSchema(), + user_input_window=TimeWindow( + start_s=float(self._step_index), + end_s=float(self._step_index + 1), + ), + metadata={"request": "fake"}, + ) + + def reset(self) -> None: + super().reset() + + def step(self, inference_input: InferenceInput) -> InferenceOutput: + return super().step(inference_input) + + +def _create_session( + pipeline: FakeStreamInferencePipeline, +) -> tuple[_ConcreteInferenceSession, _FakePipelineConfig]: + pipeline_config = _FakePipelineConfig(pipeline) + config = InferenceSessionConfig( + pipeline=cast(StreamInferencePipelineConfig, pipeline_config) + ) + return _ConcreteInferenceSession(config), pipeline_config + + +def test_constructor_initializes_the_configured_pipeline() -> None: + pipeline = FakeStreamInferencePipeline(output=object()) + + session, pipeline_config = _create_session(pipeline) + + assert session.pipeline is pipeline + assert pipeline_config.setup_calls == 1 + + +def test_reset_resets_the_pipeline_and_rollout_state() -> None: + pipeline = FakeStreamInferencePipeline(output=object()) + session, _ = _create_session(pipeline) + session.step( + InferenceInput( + global_conditioning={"prompt": "first"}, + per_step_conditioning={"control": 1}, + ) + ) + + session.reset() + result = session.step( + InferenceInput( + global_conditioning={"prompt": "second"}, + per_step_conditioning={"control": 2}, + ) + ) + + assert pipeline.reset_calls == 1 + assert pipeline.initialize_cache_calls == [ + {"prompt": "first"}, + {"prompt": "second"}, + ] + assert result.step_index == 0 + + +def test_step_converts_inference_input_and_wraps_pipeline_output() -> None: + generated = object() + pipeline = FakeStreamInferencePipeline(output=generated) + session, _ = _create_session(pipeline) + inference_input = InferenceInput( + global_conditioning={"prompt": "drive"}, + per_step_conditioning={"steering": 0.25}, + ) + + result = session.step(inference_input) + + assert pipeline.initialize_cache_calls == [{"prompt": "drive"}] + assert pipeline.generate_calls == [ + { + "autoregressive_index": 0, + "cache": pipeline.cache, + "input": {"steering": 0.25}, + } + ] + assert pipeline.finalize_calls == [ + {"autoregressive_index": 0, "cache": pipeline.cache} + ] + assert result == InferenceOutput( + step_index=0, + output=generated, + output_window=TimeWindow(start_s=0.0, end_s=1.0), + metadata={"request": "fake"}, + metrics={"total_ms": 4.0}, + ) + + +def test_step_reuses_the_pipeline_cache_and_advances_the_index() -> None: + pipeline = FakeStreamInferencePipeline(output=object()) + session, _ = _create_session(pipeline) + session.step(InferenceInput(global_conditioning={"prompt": "drive"})) + + result = session.step(InferenceInput(per_step_conditioning={"steering": -0.5})) + + assert pipeline.initialize_cache_calls == [{"prompt": "drive"}] + assert pipeline.generate_calls[-1] == { + "autoregressive_index": 1, + "cache": pipeline.cache, + "input": {"steering": -0.5}, + } + assert result.step_index == 1 + assert session.next_step_request().step_index == 2 diff --git a/flashdreams/tests/test_runtime_canonical.py b/flashdreams/tests/test_runtime_canonical.py index cfe2e3646..838fe505d 100644 --- a/flashdreams/tests/test_runtime_canonical.py +++ b/flashdreams/tests/test_runtime_canonical.py @@ -64,7 +64,7 @@ consumes=(DRIVER_COMMAND,), produces_step=(InputField(name="steering"),), ) -STEERING_MODEL = InferenceInputSchema(step_fields=(InputField(name="steering"),)) +STEERING_MODEL = InferenceInputSchema(per_step_fields=(InputField(name="steering"),)) WINDOW = TimeWindow(start_s=0.0, end_s=1.0) NEXT_WINDOW = TimeWindow(start_s=1.0, end_s=2.0) diff --git a/flashdreams/tests/test_runtime_input_mapping.py b/flashdreams/tests/test_runtime_input_mapping.py index b774371f4..3a7e3bb97 100644 --- a/flashdreams/tests/test_runtime_input_mapping.py +++ b/flashdreams/tests/test_runtime_input_mapping.py @@ -226,7 +226,8 @@ def test_invalid_phase_is_rejected() -> None: def test_inference_input_expose_payload_per_phase() -> None: inputs = InferenceInput( - global_conditioning={"prompt": "drive"}, step={"steering": 0.25} + global_conditioning={"prompt": "drive"}, + per_step_conditioning={"steering": 0.25}, ) assert inputs.for_phase("global_conditioning")["prompt"] == "drive" @@ -452,7 +453,7 @@ def test_combine_rejects_non_schema_entries() -> None: def test_undeclared_inference_input_catches_schema_drift() -> None: produced = InferenceInput( - global_conditioning={"prompt": "drive"}, step={"steering": 0.0} + global_conditioning={"prompt": "drive"}, per_step_conditioning={"steering": 0.0} ) undeclared = undeclared_inference_inputs(produced, PROMPT_MAPPING) @@ -463,7 +464,7 @@ def test_undeclared_inference_input_catches_schema_drift() -> None: def test_declared_outputs_report_no_drift() -> None: combined = combine_mapping_schemas((PROMPT_MAPPING, STEERING_MAPPING)) produced = InferenceInput( - global_conditioning={"prompt": "drive"}, step={"steering": 0.0} + global_conditioning={"prompt": "drive"}, per_step_conditioning={"steering": 0.0} ) assert undeclared_inference_inputs(produced, combined) == () @@ -476,7 +477,7 @@ def test_identity_mapping_needs_no_declared_surface() -> None: """Fixed-input runs stay possible without any schema declaration.""" mapping = IdentityInputMapping() fixed = InferenceInput( - global_conditioning={"prompt": "fixed"}, step={"steering": 0.0} + global_conditioning={"prompt": "fixed"}, per_step_conditioning={"steering": 0.0} ) mapped = mapping.map_step_inputs( @@ -485,7 +486,7 @@ def test_identity_mapping_needs_no_declared_surface() -> None: request=StepRequest(step_index=0), ) - assert mapped.step["steering"] == 0.0 + assert mapped.per_step_conditioning["steering"] == 0.0 def test_empty_mapping_set_cannot_satisfy_a_required_field() -> None: