From 27e2d8949b3ad715fc71be5bd7ccf525bf238c03 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Tue, 4 Aug 2026 10:35:19 -0700 Subject: [PATCH 1/7] WIP implementation of T2, T3, partial T4 --- docs/inference_runtime_api_design.md | 36 +- ...inference_runtime_inputs_implementation.md | 385 +++++++ ...ence_runtime_supported_inputs_inventory.md | 333 ++++++ flashdreams/flashdreams/inference/__init__.py | 50 + flashdreams/flashdreams/inference/inputs.py | 717 +++++++++++++ flashdreams/tests/test_inference_inputs.py | 954 ++++++++++++++++++ 6 files changed, 2470 insertions(+), 5 deletions(-) create mode 100644 docs/inference_runtime_inputs_implementation.md create mode 100644 docs/inference_runtime_supported_inputs_inventory.md create mode 100644 flashdreams/flashdreams/inference/__init__.py create mode 100644 flashdreams/flashdreams/inference/inputs.py create mode 100644 flashdreams/tests/test_inference_inputs.py diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md index 6a4eea9dc..5376a60e2 100644 --- a/docs/inference_runtime_api_design.md +++ b/docs/inference_runtime_api_design.md @@ -30,6 +30,12 @@ integration-specific runner code: - metrics/profiling: timings, memory, traces, NVTX ranges, and benchmark outputs. +Current T2/T3 implementation notes are in +`docs/inference_runtime_inputs_implementation.md`. + +The supported-model input inventory used to revisit T2/T3 is in +`docs/inference_runtime_supported_inputs_inventory.md`. + The API should standardize the envelope and lifecycle. It should not pretend that all world models have the same inputs, that all models use the same optimization stack, or that a raw checkpoint can fully describe how to run the @@ -322,6 +328,17 @@ Model input payloads should use semantic names, not only modality names. For example, a first frame and an HD map frame should be distinct inputs even if both are image-like values. +Model input metadata may also include a lightweight lifecycle label, such as +runtime config, cache initialization, rollout binding, per-step input, or +session update. This should remain query metadata, not model-specific tensor +validation. + +Model input names, payload kinds, lifecycle labels, and schema metadata should +be open-ended. Supported integrations such as SANA-WM, LingBot, Omnidreams, and +future external adapters may need different semantic fields. Adding a new model +should usually mean adding adapter-owned schema declarations and mappings, not +changing a central FlashDreams enum. + For interactive runs, most `ModelInputs` will be initial values plus per-step inputs produced by input mapping. For MP4 generation and benchmarking, the API should also support fixed per-step model inputs so runs can be deterministic. @@ -342,6 +359,13 @@ The purpose is to fail early before expensive model initialization, produce clearer errors, make fixed scenarios easier to validate, and avoid ambiguous dict payloads where keys only describe modality. +Schema objects may carry open-ended metadata for query-time hints such as +coordinate frame, units, rough shape summary, accepted file suffixes, schema +URI, model family, or source/transport details. Metadata should help humans and +adapter selection code, but compatibility should still be based on the declared +event capabilities, semantic model fields, payload representation hints, and +lifecycle labels. + For simple CLI text-to-video or image-to-video runs, `UserInputSchema` can be trivial or omitted because there may be no live controls. `ModelInputSchema` is more important because each supported model still needs to declare the @@ -406,9 +430,10 @@ 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: +Input mapping is required whenever `UserInputs` need to become `ModelInputs`. +The selected mapping may be one mapper or a composed set of mapper schemas. 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; @@ -417,8 +442,9 @@ object. It could be: 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; +- before runtime initialization, FlashDreams should select the mapping or mapper + set 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. diff --git a/docs/inference_runtime_inputs_implementation.md b/docs/inference_runtime_inputs_implementation.md new file mode 100644 index 000000000..e6cda3727 --- /dev/null +++ b/docs/inference_runtime_inputs_implementation.md @@ -0,0 +1,385 @@ + + +# Inference Runtime Inputs API Implementation Notes + +This note documents the current T2/T3 implementation of the inference runtime +input contracts. It is meant as an evaluator's guide: what exists, how the +pieces fit together, what the compatibility query answers, and what is +intentionally still outside this layer. + +Implementation lives in `flashdreams.inference`: + +- `flashdreams/flashdreams/inference/__init__.py` +- `flashdreams/flashdreams/inference/inputs.py` +- `flashdreams/tests/test_inference_inputs.py` + +The current supported-model input inventory that informed this revision is in +`docs/inference_runtime_supported_inputs_inventory.md`. + +## What This Implements + +The implementation covers the T2/T3 contract from +`docs/inference_runtime_api_design.md`: + +- `UserInputs` are represented as timestamped events. +- `UserInputSchema` describes what an app, transport, replay trace, or + benchmark source can provide. +- `ModelInputs` are semantic model-facing payloads split into initial and + per-step inputs. +- `ModelInputSchema` describes what a model or session requires or optionally + accepts. +- `InputMapper` is the conversion contract between user events and model inputs. +- Schema objects support open-ended `metadata` maps for lightweight hints that + future model adapters can expose without changing the core API. +- `check_mapping_compatibility()` answers whether a selected source can drive a + selected model through a selected mapper before expensive runtime setup. +- `check_mapping_set_compatibility()` answers the same question for a composed + set of mapper schemas, such as prompt plus first-frame plus live-control + mappings. + +This does not migrate LingBot, OmniDreams, WebRTC, CLI runners, or the standard +runtime/session loop. Those are T4+ and migration tasks. + +## User Inputs + +`UserInputEvent` is the primary user-input representation. Every user-facing +input is modeled as an event in session time, including static startup values. + +Examples: + +- `prompt_set` at `timestamp_s=0.0` +- `initial_frame_set` at `timestamp_s=0.0` +- `scene_selected` at `timestamp_s=0.0` +- `key_down` / `key_up` during a session +- `controller_axis` during a session +- `camera_pose` events in a replay trace + +`UserInputTrace` stores events in deterministic timestamp order and can slice a +`UserInputWindow` for one runtime step. Events with equal timestamps preserve +their original order. + +```python +from flashdreams.inference import UserInputEvent, UserInputTrace + +trace = UserInputTrace.from_events( + [ + UserInputEvent( + timestamp_s=0.0, + kind="prompt_set", + payload={"prompt": "drive forward"}, + ), + UserInputEvent( + timestamp_s=0.5, + kind="key_down", + payload={"key": "w"}, + ), + ] +) + +step_window = trace.window(start_s=0.0, end_s=1.0) +``` + +Snapshots are intentionally not primary inputs. A mapper or runtime can derive +snapshots from a trace/window when a model wants snapshot-style controls. + +## User Input Schema + +`UserInputSchema` is lightweight metadata for the source side. It answers: + +- Which event kinds can this source provide? +- Which payload fields are present on those events? +- Is this source live, replayed, fixed, or otherwise identified? + +It does not assign model semantics. For example, `key_down` does not mean +steering or camera motion until a mapper says so. + +```python +from flashdreams.inference import UserInputCapability, UserInputSchema + +browser_schema = UserInputSchema( + name="browser", + source_kind="live", + metadata={"transport": "webrtc"}, + capabilities=( + UserInputCapability( + event_kind="prompt_set", + payload_kind="text", + payload_fields=frozenset({"prompt"}), + ), + UserInputCapability( + event_kind="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_kind="key_up", + payload_fields=frozenset({"key"}), + ), + ), +) +``` + +The schema can validate that a source event has an event kind and payload fields +the source claims to provide. It does not validate tensor shape, image format, +camera coordinate system, or model-specific units. + +Capabilities and schemas also carry optional `metadata`. This is for query-time +hints such as source type, UI widget, file suffixes, units, coordinate frame, or +schema URI. Metadata is intentionally not part of compatibility matching. + +## Model Inputs + +`ModelInputs` is the model-facing payload container. It separates values needed +to start or reset a rollout from values needed for one generated step/chunk. + +```python +from flashdreams.inference import ModelInputs + +inputs = ModelInputs( + initial={"prompt": "drive forward", "first_frame": first_frame}, + step={"steering": 0.25}, +) +``` + +`ModelInputSchema` declares the semantic fields a model or session expects. It +uses names like `prompt`, `first_frame`, `steering`, `camera_trajectory`, or +`hdmap_frames`, rather than generic modality-only keys. + +```python +from flashdreams.inference import ModelInputField, ModelInputSchema + +model_schema = ModelInputSchema( + name="driving-model", + fields=( + ModelInputField( + name="prompt", + phase="initial", + required=True, + payload_kind="text", + update_policy="step_boundary", + lifecycle="cache_init", + ), + ModelInputField( + name="steering", + phase="step", + required=True, + lifecycle="step_input", + ), + ModelInputField( + name="first_frame", + phase="initial", + required=False, + lifecycle="cache_init", + ), + ModelInputField( + name="camera_trajectory_c2w", + phase="initial", + required=False, + payload_kind="c2w_sequence", + lifecycle="rollout_binding", + metadata={"shape": "[F,4,4]", "coordinates": "opencv_c2w"}, + ), + ), + metadata={"model_family": "example-driving-model"}, +) +``` + +`update_policy` is deliberately plain metadata. It lets a model advertise facts +such as "prompt updates can happen at step boundaries" without making this +schema layer responsible for implementing or deeply validating that behavior. + +`lifecycle` is also plain metadata. It lets a model distinguish initial values +used at different adapter moments, such as `runtime_config`, `cache_init`, +`rollout_binding`, `step_input`, or `session_update`. Compatibility requires +lifecycle agreement only when both the model field and mapper output specify a +lifecycle; otherwise simple schemas remain permissive. + +Model input names, payload kinds, lifecycle labels, and metadata are +open-ended. They are not a FlashDreams-wide enum. A SANA-WM-like adapter can +declare fields such as `camera_trajectory_c2w`, `camera_intrinsics_vec4`, +`stage1_sampling`, or `streaming_chunking`; another model can declare different +semantic names. The adapter and mapper own the deep interpretation. + +## Input Mappers + +`InputMapper` is a protocol with three responsibilities: + +- expose an `InputMapperSchema`; +- build initial `ModelInputs` from a `UserInputTrace`; +- build per-step `ModelInputs` from a `UserInputWindow`. + +`InputMapperSchema` declares the mapper's compatibility surface: + +- `consumes`: user event capabilities required by the mapper; +- `produces`: model input fields the mapper can produce. + +Example: keyboard events can be mapped into steering for one model, camera +trajectory for another model, or ignored entirely. That meaning is owned by the +mapper, not by `UserInputEvent`. + +```python +from flashdreams.inference import ( + InputMapperSchema, + ModelInputField, + UserInputCapability, +) + +keyboard_to_steering_schema = InputMapperSchema( + name="keyboard-to-steering", + consumes=( + UserInputCapability( + event_kind="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_kind="key_up", + payload_fields=frozenset({"key"}), + ), + ), + produces=( + ModelInputField(name="steering", phase="step"), + ), +) +``` + +For fixed runs that already have model-facing inputs, `StaticInputMapper` +provides the no-live-controls path. It consumes no user events and returns the +configured initial/per-step `ModelInputs`. + +```python +from flashdreams.inference import ModelInputs, StaticInputMapper + +static_mapper = StaticInputMapper.from_inputs( + inputs=ModelInputs( + initial={"prompt": "fixed prompt"}, + step={"camera_trajectory": camera_poses}, + ), + name="fixed-scenario", +) +``` + +Mapper schemas can also be combined for compatibility checks. This matches the +current supported model inventory: a run may select one mapping for prompt +events, another for first-frame events, and another for keyboard or controller +events. + +```python +from flashdreams.inference import check_mapping_set_compatibility + +compatibility = check_mapping_set_compatibility( + source_schema=browser_schema, + model_schema=lingbot_schema, + mapper_schemas=( + prompt_mapper.schema, + first_frame_mapper.schema, + keyboard_to_camera_mapper.schema, + ), +) +``` + +## Compatibility Query + +`check_mapping_compatibility()` is the central query for T2/T3. Given a source +schema, model schema, and mapper schema, it reports: + +- whether the source can drive the model through this mapper; +- source capabilities the mapper needs but the source lacks; +- required model fields the mapper cannot produce; +- required model fields that are satisfied; +- optional model fields that can be enabled. + +```python +from flashdreams.inference import check_mapping_compatibility + +compatibility = check_mapping_compatibility( + source_schema=browser_schema, + model_schema=model_schema, + mapper_schema=keyboard_to_steering_schema, +) + +if not compatibility.can_drive: + compatibility.raise_if_incompatible() +``` + +This is an early compatibility check. It is intended to fail before expensive +runtime initialization when the mismatch is obvious. It is not a guarantee that +the model run will succeed. + +Use `check_mapping_set_compatibility()` when the selected mapping is composed +from multiple mapper schemas. It returns the same `MappingCompatibility` report +after combining mapper inputs and outputs. + +## What This Does Not Validate + +The schemas intentionally avoid becoming a rich type system. The following +remain the responsibility of the model adapter, runtime, session, or mapper +implementation: + +- tensor shape and dtype; +- image decode details; +- camera coordinate systems; +- pose and timestamp units; +- prompt-embedding swap mechanics; +- reset semantics; +- whether a specific model can actually apply an update policy at runtime; +- deep validation of scene, HD map, or actor-state data. + +The schema layer should be enough to answer "can this source plausibly drive +this model through this mapper?" It should not replace model-owned validation. + +## Current Tests + +The focused CPU tests are in `flashdreams/tests/test_inference_inputs.py`. +They cover: + +- deterministic event ordering and trace windowing; +- startup values represented as events; +- source capability declaration and basic event validation; +- required and optional model input declarations; +- prompt update metadata on a model field; +- lifecycle metadata on model fields and mapper outputs; +- open-ended schema metadata that stays queryable but does not constrain + compatibility; +- missing required model inputs; +- missing source capabilities; +- optional model inputs becoming available only when mapper support exists; +- composed mapper-set compatibility; +- SANA-WM-like model-specific inputs without importing the SANA integration; +- fake prompt, initial-frame, keyboard-to-steering, and camera-trajectory + mappers; +- `StaticInputMapper` for fixed model-input scenarios. + +Targeted validation command: + +```bash +.venv/bin/pytest flashdreams/tests/test_inference_inputs.py -q +.venv/bin/ty check flashdreams/flashdreams/inference flashdreams/tests/test_inference_inputs.py +.venv/bin/python -m py_compile \ + flashdreams/flashdreams/inference/__init__.py \ + flashdreams/flashdreams/inference/inputs.py \ + flashdreams/tests/test_inference_inputs.py +``` + +At the time this note was written, these targeted checks passed. + +## Evaluation Checklist + +Use this checklist to judge whether the implementation matches the T2/T3 design: + +- Are all `UserInputs` represented as timestamped events? +- Can a source declare what event capabilities it provides? +- Can a model declare required and optional initial/per-step inputs? +- Are model input names semantic rather than modality-only? +- Can new models add model-specific field names and metadata without changing + core dataclasses? +- Can a mapper declare what it consumes and what it produces? +- Can multiple mapper schemas be checked as one selected mapping surface? +- Does compatibility checking report both missing source capabilities and + missing required model inputs? +- Are optional model inputs reported separately from required inputs? +- Can model fields distinguish cache initialization, rollout binding, per-step + inputs, and active-session update support without deep tensor validation? +- Is deep model/tensor validation kept out of the lightweight schema layer? +- Is fixed input/model-input replay possible without live user events? diff --git a/docs/inference_runtime_supported_inputs_inventory.md b/docs/inference_runtime_supported_inputs_inventory.md new file mode 100644 index 000000000..42f1f450e --- /dev/null +++ b/docs/inference_runtime_supported_inputs_inventory.md @@ -0,0 +1,333 @@ + + +# Supported Model Input Inventory For T2/T3 + +This note inventories the inputs used by the currently supported FlashDreams +runners and interactive runtimes, plus the SANA-WM input surface on `main`, then +records the T2/T3 API implications. It is intentionally about input contracts, +not tensor shape validation or model quality. + +## Inventory + +WAN 2.1 T2V, Self-Forcing WAN 2.1 T2V, Causal-Forcing T2V, +FastVideo Causal WAN 2.2 T2V, and Cosmos Predict2 T2V: + +- Source/app inputs: prompt text or prompt text file, pixel height/width, and + fps or block count depending on runner. +- Model-facing initial inputs: prompt text plus latent/output height and width + derived from run config. +- Model-facing step/update inputs: no live controls; AR loop steps with fixed + session state. + +WAN 2.1 I2V, Causal-Forcing I2V, and Cosmos Predict2 I2V: + +- Source/app inputs: prompt text or prompt file, first-frame image path or URL, + and pixel height/width. +- Model-facing initial inputs: prompt text and decoded first-frame tensor. +- Model-facing step/update inputs: no live controls. + +FlashVSR: + +- Source/app inputs: input video path or URL, chunk size, crop region, sparse + ratio, and optional output FPS. +- Model-facing initial inputs: no explicit prompt at runner time; the prompt + tensor is configured in the pipeline. Input video dimensions affect + per-video runtime/pipeline setup. +- Model-facing step/update inputs: video chunks passed to + `pipeline.generate(input=clip)`. + +LingBot CLI: + +- Source/app inputs: prompt or prompt path, first-frame image path, pose path, + intrinsics path, total blocks, dimensions, and fps. +- Model-facing initial inputs: prompt text and first-frame tensor. +- Model-facing step/update inputs: `CamCtrlInput` with intrinsics, camera poses, + and world scale. + +LingBot WebRTC: + +- Source/app inputs: session prompt, uploaded/remote/default first-frame image, + keyboard events, reset requests, text-event catalog, and trigger events. +- Model-facing initial inputs: prompt text, first-frame tensor, base text + embeddings, precomputed text-event embeddings, base intrinsics, and world + scale. +- Model-facing step/update inputs: keyboard event windows become pose segments + and camera trajectories. Text-event triggers can replace rollout text + embeddings when the model supports it. + +HY-WorldPlay WAN I2V: + +- Source/app inputs: prompt or prompt path, first-frame image path or example + image, pose string or pose JSON, memory-selection settings, dimensions, fps, + and seed. +- Model-facing initial inputs: prompt text and first-frame tensor for cache + initialization. +- Model-facing step/update inputs: pose data is bound for the rollout as action + labels, view matrices, intrinsics, and memory-selection state before AR steps. + +Omnidreams CLI: + +- Source/app inputs: shared prompt or per-camera prompts, HDMap video paths, + first-frame image/video paths, camera names, example-data UUID, and optional + embedding save/load paths. +- Model-facing initial inputs: prompt list, first-frame tensor, view names; or + precomputed text/image/negative-text embeddings. +- Model-facing step/update inputs: HDMap video chunks passed per AR step. + +Omnidreams WebRTC: + +- Source/app inputs: scene directory or scene UUID, scene variant, camera name, + prompt/first-frame assets resolved from the scene, keyboard events, reset + requests, and optional postprocess preset. +- Model-facing initial inputs: scene data, renderer, first-frame tensor, prompt, + camera calibration/extrinsics, initial ego pose, and initial timestamp. +- Model-facing step/update inputs: keyboard event windows become ego poses, + camera poses per view, and frame timestamps. The wrapper renders HDMap + conditioning internally for each step. + +Omnidreams interactive drive: + +- Source/app inputs: scene bundle, keyboard events or wheel/controller samples, + view-mode/reset/scene-exit controls, and vehicle/chunk config. +- Model-facing initial inputs: scene bundle, selected camera, prompt, initial + RGB frame, initial rig pose, and initial timestamp. +- Model-facing step/update inputs: `DriverCommand` samples become trajectory + chunks, rendered frames, and world-model conditioning. + +Template recipe: + +- Source/app inputs: synthetic runner config: batch size, height, width, context + tokens, AR steps, and seed. +- Model-facing initial inputs: synthetic transformer context, optional negative + context, height, and width. +- Model-facing step/update inputs: optional synthetic control tensor. + +WAN 2.2 TI2V pipeline config: + +- Source/app inputs: downstream runners use this rather than a standalone runner + in this tree. +- Model-facing initial inputs: prompt text and first-frame image for TI2V-style + cache initialization. +- Model-facing step/update inputs: downstream runners decide controls; + HY-WorldPlay currently binds action/camera state around it. + +SANA-WM bidirectional and streaming on `main`: + +- Source/app inputs: first-frame image path, prompt or prompt path, optional + negative prompt, camera trajectory path or action DSL, optional intrinsics + path or derived intrinsics, frame count, fps, Stage-1 sampling knobs, seed, + precision/refiner options, and streaming chunk/block settings. +- Model-facing initial inputs: decoder context such as prompt, fps, + `save_stage1`, refiner seed, sink size, and streaming refiner window/block + parameters. +- Model-facing step/update inputs: bidirectional passes one + `SanaWMI2VConditioningRequest` into the single generation step. Streaming + passes one `SanaWMStreamingI2VConditioningRequest` repeatedly; the + conditioning encoder caches rollout-wide prompt, first-frame, camera, latent + shape, and chunk-boundary state, then slices per AR chunk. +- Model-facing semantic fields include prompt, negative prompt, first frame, + camera-to-world trajectory, intrinsics vec4 sequence, frame count, fps, + sampling parameters, seed, and streaming chunking parameters. + +## API Implications + +The inventory changes the T2/T3 shape in four concrete ways. + +First, a selected mapping is often a composition. A LingBot-like run needs prompt +mapping, first-frame mapping, and keyboard-to-camera mapping. Omnidreams may add +scene selection, camera selection, and HDMap mapping. The implementation should +support checking a set of mapper schemas as one compatibility surface, while +still allowing a single mapper object when that is simpler. + +Second, `ModelInputSchema` needs a lightweight lifecycle tag in addition to the +`initial` versus `step` phase. The phase answers when the value is needed at the +standard-loop level. The lifecycle tag distinguishes where the model adapter +uses it, such as: + +- `runtime_config`: values that affect setup before model/runtime construction, + such as FlashVSR input-video dimensions; +- `cache_init`: values passed when initializing or resetting a rollout cache, + such as prompts, first frames, view names, and precomputed embeddings; +- `rollout_binding`: values bound after cache initialization but before AR + steps, such as HY-WorldPlay action labels, camera tensors, and memory state; +- `step_input`: values consumed for one generated chunk, such as HDMap frames, + camera trajectories, driver commands, video chunks, and timestamps; +- `session_update`: values that can update an active session when supported, + such as LingBot text-event embedding swaps. + +The lifecycle tag is metadata, not a new deep type system. If both a model field +and mapper output specify lifecycle, compatibility should require them to agree. +If either side omits it, matching stays permissive for simple schemas. + +Third, `payload_kind` should be treated as a representation hint rather than a +universal semantic type. For example, `prompt` may arrive as inline text or a +path but become prompt text or text embeddings; `first_frame` may arrive as a +path, URL, bytes, or decoded tensor; camera motion may arrive as keys, pose JSON, +Numpy arrays, or integrated tensors. The semantic input name is still the main +contract. + +Fourth, schema objects need open-ended metadata for future adapters. This lets a +SANA-WM-like adapter advertise that `camera_trajectory_c2w` uses an +`[F,4,4]` OpenCV camera-to-world sequence, or lets another model advertise a +schema URI, units, coordinate frame, accepted file suffixes, cardinality hints, +or update notes. Metadata should remain query information and should not become +the compatibility type system. + +Fifth, `UserInputSchema` should describe raw source capabilities, while +mapper schemas describe derived model-facing semantics. A browser may provide +`key_down`, `key_up`, `prompt_set`, and `initial_frame_set` events. Whether that +can drive `steering`, `camera_trajectory`, `driver_command`, or text embedding +updates depends on the selected mapper and model schema. + +## Revised T2/T3 Plan + +The implementation plan after this inventory is: + +1. Keep `UserInputEvent`, `UserInputTrace`, and `UserInputWindow` as the primary + event-based input API. Static startup values remain timestamp-zero events. +2. Keep `UserInputSchema` lightweight and source-facing. It declares event kinds, + payload representation hints, and payload fields that a source can provide. +3. Keep `ModelInputs` split into `initial` and `step` payload maps. This still + matches the standard loop's session-start and per-step moments. +4. Extend `ModelInputField` with optional lifecycle metadata so models can + distinguish runtime config, cache initialization, rollout binding, per-step + inputs, and supported active-session updates. +5. Keep `InputMapperSchema` as the mapping boundary, but add mapper-set + compatibility helpers for composed mappings. +6. Keep input names, payload kinds, lifecycle labels, and metadata open-ended. + Adding a new model should usually mean adding adapter-owned schema + declarations and mappers, not changing the core input dataclasses. +7. Leave deep validation to model adapters, sessions, and mappers. The schema + layer should catch obvious source/mapping/model mismatches before expensive + runtime initialization, not validate every tensor and coordinate convention. + +## Extensibility Contract + +The inventory above is not a vocabulary freeze. The core API should not contain +a closed enum of allowed model input names. New adapters can introduce semantic +field names that match the model boundary they own. + +Use these conventions when adding future model schemas: + +- Prefer semantic names over modality names, such as `camera_trajectory_c2w` + instead of `array`, or `hdmap_frames` instead of `image`. +- Use `payload_kind` for a coarse representation hint, such as `path`, + `decoded_tensor`, `c2w_sequence`, `intrinsics_vec4_sequence`, `driver_command`, + or `embedding`. +- Use `lifecycle` to say where the adapter consumes the value, such as + `runtime_config`, `cache_init`, `rollout_binding`, `step_input`, or + `session_update`. +- Use `metadata` for query hints: units, coordinate frame, shape summary, + accepted suffixes, schema URI, model family, value ranges, or cardinality. +- Keep deep validation in the adapter/mapper. The lightweight schemas answer + whether the selected source and mapping can plausibly drive the model before + expensive initialization. + +## Representative Schema Sketches + +These sketches are not migration work for T4+, but they show that the current +T2/T3 primitives can describe the supported input surfaces. + +```python +lingbot_model = ModelInputSchema( + name="lingbot-world", + fields=( + ModelInputField("prompt", "initial", lifecycle="cache_init"), + ModelInputField("first_frame", "initial", lifecycle="cache_init"), + ModelInputField("camera_trajectory", "step", lifecycle="step_input"), + ModelInputField( + "text_embeddings", + "step", + required=False, + update_policy="step_boundary", + lifecycle="session_update", + ), + ), +) +``` + +```python +omnidreams_model = ModelInputSchema( + name="omnidreams", + fields=( + ModelInputField("prompts", "initial", lifecycle="cache_init"), + ModelInputField("first_frames", "initial", lifecycle="cache_init"), + ModelInputField("view_names", "initial", lifecycle="cache_init"), + ModelInputField("hdmap_frames", "step", lifecycle="step_input"), + ModelInputField( + "text_embeddings", + "initial", + required=False, + lifecycle="cache_init", + ), + ModelInputField( + "image_embeddings", + "initial", + required=False, + lifecycle="cache_init", + ), + ), +) +``` + +```python +hy_worldplay_model = ModelInputSchema( + name="hy-worldplay", + fields=( + ModelInputField("prompt", "initial", lifecycle="cache_init"), + ModelInputField("first_frame", "initial", lifecycle="cache_init"), + ModelInputField("action_labels", "initial", lifecycle="rollout_binding"), + ModelInputField("camera_viewmats", "initial", lifecycle="rollout_binding"), + ModelInputField("camera_intrinsics", "initial", lifecycle="rollout_binding"), + ModelInputField("memory_config", "initial", lifecycle="rollout_binding"), + ), +) +``` + +```python +sana_wm_model = ModelInputSchema( + name="sana-wm", + fields=( + ModelInputField("prompt", "initial", lifecycle="cache_init"), + ModelInputField( + "negative_prompt", + "initial", + required=False, + lifecycle="cache_init", + ), + ModelInputField("first_frame", "initial", lifecycle="cache_init"), + ModelInputField( + "camera_trajectory_c2w", + "initial", + payload_kind="c2w_sequence", + lifecycle="rollout_binding", + metadata={"shape": "[F,4,4]", "coordinates": "opencv_c2w"}, + ), + ModelInputField( + "camera_intrinsics_vec4", + "initial", + required=False, + payload_kind="intrinsics_vec4_sequence", + lifecycle="rollout_binding", + metadata={"shape": "[F,4]"}, + ), + ModelInputField( + "stage1_sampling", + "initial", + lifecycle="rollout_binding", + metadata={"fields": ("steps", "cfg_scale", "flow_shift", "seed")}, + ), + ModelInputField( + "streaming_chunking", + "initial", + required=False, + lifecycle="rollout_binding", + metadata={"fields": ("num_frame_per_block", "cached_blocks")}, + ), + ), + metadata={"model_family": "sana-wm"}, +) +``` diff --git a/flashdreams/flashdreams/inference/__init__.py b/flashdreams/flashdreams/inference/__init__.py new file mode 100644 index 000000000..e22df109b --- /dev/null +++ b/flashdreams/flashdreams/inference/__init__.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Inference runtime input contracts. + +This package contains the lightweight ``UserInputs`` / ``ModelInputs`` +contracts used by the experimental runtime API. The objects here describe +capabilities and mapping boundaries; model-specific tensor validation stays in +the model adapter or session implementation. +""" + +from flashdreams.inference.inputs import ( + InputMapper, + InputMapperSchema, + InputPhase, + MappingCompatibility, + ModelInputField, + ModelInputSchema, + ModelInputs, + StaticInputMapper, + UserInputCapability, + UserInputEvent, + UserInputSchema, + UserInputTrace, + UserInputWindow, + check_mapping_compatibility, + check_mapping_set_compatibility, + combine_mapper_schemas, + missing_required_inputs, +) + +__all__ = [ + "InputMapper", + "InputMapperSchema", + "InputPhase", + "MappingCompatibility", + "ModelInputField", + "ModelInputSchema", + "ModelInputs", + "StaticInputMapper", + "UserInputCapability", + "UserInputEvent", + "UserInputSchema", + "UserInputTrace", + "UserInputWindow", + "check_mapping_compatibility", + "check_mapping_set_compatibility", + "combine_mapper_schemas", + "missing_required_inputs", +] diff --git a/flashdreams/flashdreams/inference/inputs.py b/flashdreams/flashdreams/inference/inputs.py new file mode 100644 index 000000000..5216f455e --- /dev/null +++ b/flashdreams/flashdreams/inference/inputs.py @@ -0,0 +1,717 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""User-input, model-input, and mapping contracts for inference runtimes.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from math import isfinite +from types import MappingProxyType +from typing import Any, Literal, Protocol, cast + +InputPhase = Literal["initial", "step"] + + +def _normalized_token(value: str, *, field_name: str) -> str: + normalized = value.strip() + if not normalized: + raise ValueError(f"{field_name} must be a non-empty string.") + return normalized + + +def _normalized_optional_token(value: str | None, *, field_name: str) -> str | None: + if value is None: + return None + return _normalized_token(value, field_name=field_name) + + +def _normalized_payload_fields(values: Iterable[str]) -> frozenset[str]: + return frozenset( + _normalized_token(value, field_name="payload field") for value in values + ) + + +def _immutable_metadata_map(values: Mapping[str, Any]) -> Mapping[str, Any]: + normalized: dict[str, Any] = {} + for key, value in values.items(): + normalized[_normalized_token(str(key), field_name="metadata key")] = value + return MappingProxyType(normalized) + + +def _validate_phase(value: str) -> InputPhase: + if value not in {"initial", "step"}: + raise ValueError(f"phase must be 'initial' or 'step', got {value!r}.") + return cast(InputPhase, value) + + +def _field_key(field: "ModelInputField") -> tuple[InputPhase, str]: + return field.phase, field.name + + +def _model_field_matches( + produced: "ModelInputField", + required: "ModelInputField", +) -> bool: + if _field_key(produced) != _field_key(required): + return False + payload_kind_ok = ( + produced.payload_kind is None + or required.payload_kind is None + or produced.payload_kind == required.payload_kind + ) + lifecycle_ok = ( + produced.lifecycle is None + or required.lifecycle is None + or produced.lifecycle == required.lifecycle + ) + return payload_kind_ok and lifecycle_ok + + +@dataclass(frozen=True, slots=True) +class UserInputEvent: + """One user-facing input event in session time. + + Static session setup values, such as prompts or initial frames, are still + represented as events. Snapshots are derived views over events rather than + primary inputs. + """ + + timestamp_s: float + kind: str + payload: Mapping[str, Any] = field(default_factory=dict) + session_id: str | None = None + source: str | None = None + + def __post_init__(self) -> None: + timestamp_s = float(self.timestamp_s) + if not isfinite(timestamp_s) or timestamp_s < 0: + raise ValueError("timestamp_s must be finite and >= 0.") + object.__setattr__(self, "timestamp_s", timestamp_s) + object.__setattr__( + self, + "kind", + _normalized_token(self.kind, field_name="event kind"), + ) + object.__setattr__(self, "payload", MappingProxyType(dict(self.payload))) + object.__setattr__( + self, + "session_id", + _normalized_optional_token(self.session_id, field_name="session_id"), + ) + object.__setattr__( + self, + "source", + _normalized_optional_token(self.source, field_name="source"), + ) + + +@dataclass(frozen=True, slots=True) +class UserInputWindow: + """A deterministic time window over user input events.""" + + start_s: float + end_s: float + events: tuple[UserInputEvent, ...] = () + + def __post_init__(self) -> None: + start_s = float(self.start_s) + end_s = float(self.end_s) + if not isfinite(start_s) or not isfinite(end_s): + raise ValueError("window bounds must be finite.") + if end_s < start_s: + raise ValueError("end_s must be >= start_s.") + object.__setattr__(self, "start_s", start_s) + object.__setattr__(self, "end_s", end_s) + object.__setattr__(self, "events", _sorted_events(self.events)) + + def events_of_kind(self, kind: str) -> tuple[UserInputEvent, ...]: + """Return all events of ``kind`` inside this window.""" + kind = _normalized_token(kind, field_name="event kind") + return tuple(event for event in self.events if event.kind == kind) + + def latest(self, kind: str) -> UserInputEvent | None: + """Return the latest event of ``kind`` inside this window, if any.""" + events = self.events_of_kind(kind) + return events[-1] if events else None + + +@dataclass(frozen=True, slots=True) +class UserInputTrace: + """Ordered replayable user-input event trace.""" + + events: tuple[UserInputEvent, ...] = () + + def __post_init__(self) -> None: + object.__setattr__(self, "events", _sorted_events(self.events)) + + @classmethod + def from_events(cls, events: Sequence[UserInputEvent]) -> "UserInputTrace": + """Build a trace from any event sequence.""" + return cls(events=tuple(events)) + + def window( + self, + *, + start_s: float, + end_s: float, + include_start: bool = True, + include_end: bool = False, + ) -> UserInputWindow: + """Slice the trace into a deterministic event window.""" + start_s = float(start_s) + end_s = float(end_s) + if end_s < start_s: + raise ValueError("end_s must be >= start_s.") + + def _starts_in(event: UserInputEvent) -> bool: + if include_start: + return event.timestamp_s >= start_s + return event.timestamp_s > start_s + + def _ends_in(event: UserInputEvent) -> bool: + if include_end: + return event.timestamp_s <= end_s + return event.timestamp_s < end_s + + return UserInputWindow( + start_s=start_s, + end_s=end_s, + events=tuple( + event for event in self.events if _starts_in(event) and _ends_in(event) + ), + ) + + def events_of_kind(self, kind: str) -> tuple[UserInputEvent, ...]: + """Return all events of ``kind`` in this trace.""" + kind = _normalized_token(kind, field_name="event kind") + return tuple(event for event in self.events if event.kind == kind) + + def latest(self, kind: str) -> UserInputEvent | None: + """Return the latest event of ``kind`` in this trace, if any.""" + events = self.events_of_kind(kind) + return events[-1] if events else None + + +def _sorted_events(events: Sequence[UserInputEvent]) -> tuple[UserInputEvent, ...]: + indexed_events = tuple(enumerate(events)) + for _, event in indexed_events: + if not isinstance(event, UserInputEvent): + raise TypeError(f"Expected UserInputEvent, got {type(event).__name__}.") + return tuple( + event + for _, event in sorted( + indexed_events, + key=lambda item: (item[1].timestamp_s, item[0]), + ) + ) + + +@dataclass(frozen=True, slots=True) +class UserInputCapability: + """Lightweight metadata for a user event a source or mapper can provide.""" + + event_kind: str + payload_kind: str | None = None + payload_fields: frozenset[str] = field(default_factory=frozenset) + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + description: str = "" + + def __post_init__(self) -> None: + object.__setattr__( + self, + "event_kind", + _normalized_token(self.event_kind, field_name="event kind"), + ) + object.__setattr__( + self, + "payload_kind", + _normalized_optional_token( + self.payload_kind, + field_name="payload_kind", + ), + ) + object.__setattr__( + self, + "payload_fields", + _normalized_payload_fields(self.payload_fields), + ) + object.__setattr__(self, "metadata", _immutable_metadata_map(self.metadata)) + + def is_satisfied_by(self, provider: "UserInputCapability") -> bool: + """Return whether ``provider`` can satisfy this consumed capability.""" + if self.event_kind != provider.event_kind: + return False + payload_kind_ok = ( + self.payload_kind is None + or provider.payload_kind is None + or self.payload_kind == provider.payload_kind + ) + return payload_kind_ok and self.payload_fields.issubset( + provider.payload_fields + ) + + +@dataclass(frozen=True, slots=True) +class UserInputSchema: + """Metadata describing what a source can provide.""" + + capabilities: tuple[UserInputCapability, ...] = () + name: str = "user-input-source" + source_kind: str | None = None + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "name", + _normalized_token(self.name, field_name="schema name"), + ) + object.__setattr__( + self, + "source_kind", + _normalized_optional_token(self.source_kind, field_name="source_kind"), + ) + object.__setattr__(self, "metadata", _immutable_metadata_map(self.metadata)) + object.__setattr__(self, "capabilities", tuple(self.capabilities)) + for capability in self.capabilities: + if not isinstance(capability, UserInputCapability): + raise TypeError( + "capabilities must contain UserInputCapability objects." + ) + + def supports(self, capability: UserInputCapability) -> bool: + """Return whether this source can satisfy ``capability``.""" + return any( + capability.is_satisfied_by(provider) for provider in self.capabilities + ) + + def validate_event(self, event: UserInputEvent) -> None: + """Validate an event against this source schema.""" + matching = [ + capability + for capability in self.capabilities + if capability.event_kind == event.kind + ] + if not matching: + raise ValueError( + f"User input source {self.name!r} does not provide event " + f"kind {event.kind!r}." + ) + payload_keys = set(event.payload) + if not any( + capability.payload_fields.issubset(payload_keys) + for capability in matching + ): + expected = sorted( + { + field_name + for capability in matching + for field_name in capability.payload_fields + } + ) + raise ValueError( + f"Event {event.kind!r} payload is missing required fields " + f"for source {self.name!r}: {expected}." + ) + + +@dataclass(frozen=True, slots=True) +class ModelInputField: + """Lightweight metadata for one semantic model-facing input field.""" + + name: str + phase: InputPhase + required: bool = True + payload_kind: str | None = None + update_policy: str | None = None + lifecycle: str | None = None + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + description: str = "" + + def __post_init__(self) -> None: + object.__setattr__( + self, + "name", + _normalized_token(self.name, field_name="model input name"), + ) + object.__setattr__(self, "phase", _validate_phase(self.phase)) + object.__setattr__( + self, + "payload_kind", + _normalized_optional_token( + self.payload_kind, + field_name="payload_kind", + ), + ) + object.__setattr__( + self, + "update_policy", + _normalized_optional_token( + self.update_policy, + field_name="update_policy", + ), + ) + object.__setattr__( + self, + "lifecycle", + _normalized_optional_token( + self.lifecycle, + field_name="lifecycle", + ), + ) + object.__setattr__(self, "metadata", _immutable_metadata_map(self.metadata)) + + +@dataclass(frozen=True, slots=True) +class ModelInputSchema: + """Metadata describing what a model or session expects.""" + + fields: tuple[ModelInputField, ...] = () + name: str = "model-input-consumer" + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "name", + _normalized_token(self.name, field_name="schema name"), + ) + object.__setattr__(self, "metadata", _immutable_metadata_map(self.metadata)) + object.__setattr__(self, "fields", tuple(self.fields)) + seen: set[tuple[InputPhase, str]] = set() + for field_def in self.fields: + if not isinstance(field_def, ModelInputField): + raise TypeError("fields must contain ModelInputField objects.") + key = _field_key(field_def) + if key in seen: + phase, name = key + raise ValueError( + f"Duplicate model input field {name!r} for phase {phase!r}." + ) + seen.add(key) + + def required_fields( + self, + *, + phase: InputPhase | None = None, + ) -> tuple[ModelInputField, ...]: + """Return required fields, optionally filtered by phase.""" + return tuple( + field_def + for field_def in self.fields + if field_def.required and (phase is None or field_def.phase == phase) + ) + + def optional_fields( + self, + *, + phase: InputPhase | None = None, + ) -> tuple[ModelInputField, ...]: + """Return optional fields, optionally filtered by phase.""" + return tuple( + field_def + for field_def in self.fields + if not field_def.required and (phase is None or field_def.phase == phase) + ) + + def field(self, *, name: str, phase: InputPhase) -> ModelInputField | None: + """Return one field definition, if present.""" + name = _normalized_token(name, field_name="model input name") + phase = _validate_phase(phase) + for field_def in self.fields: + if field_def.name == name and field_def.phase == phase: + return field_def + return None + + +@dataclass(frozen=True, slots=True) +class ModelInputs: + """Semantic model-facing input payloads split by runtime phase.""" + + initial: Mapping[str, Any] = field(default_factory=dict) + step: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "initial", _immutable_payload_map(self.initial)) + object.__setattr__(self, "step", _immutable_payload_map(self.step)) + + @classmethod + def initial_only(cls, values: Mapping[str, Any]) -> "ModelInputs": + """Create model inputs containing only initial values.""" + return cls(initial=values) + + @classmethod + def step_only(cls, values: Mapping[str, Any]) -> "ModelInputs": + """Create model inputs containing only per-step values.""" + return cls(step=values) + + def for_phase(self, phase: InputPhase) -> Mapping[str, Any]: + """Return the payload mapping for ``phase``.""" + phase = _validate_phase(phase) + return self.initial if phase == "initial" else self.step + + +def _immutable_payload_map(values: Mapping[str, Any]) -> Mapping[str, Any]: + normalized: dict[str, Any] = {} + for key, value in values.items(): + normalized[_normalized_token(str(key), field_name="model input key")] = value + return MappingProxyType(normalized) + + +def missing_required_inputs( + inputs: ModelInputs, + schema: ModelInputSchema, + *, + phase: InputPhase | None = None, +) -> tuple[ModelInputField, ...]: + """Return required model fields absent from ``inputs``.""" + return tuple( + field_def + for field_def in schema.required_fields(phase=phase) + if field_def.name not in inputs.for_phase(field_def.phase) + ) + + +@dataclass(frozen=True, slots=True) +class InputMapperSchema: + """Metadata for a mapper that converts user events into model inputs.""" + + consumes: tuple[UserInputCapability, ...] = () + produces: tuple[ModelInputField, ...] = () + name: str = "input-mapper" + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "name", + _normalized_token(self.name, field_name="mapper name"), + ) + object.__setattr__(self, "metadata", _immutable_metadata_map(self.metadata)) + object.__setattr__(self, "consumes", tuple(self.consumes)) + object.__setattr__(self, "produces", tuple(self.produces)) + for capability in self.consumes: + if not isinstance(capability, UserInputCapability): + raise TypeError("consumes must contain UserInputCapability objects.") + for field_def in self.produces: + if not isinstance(field_def, ModelInputField): + raise TypeError("produces must contain ModelInputField objects.") + + def can_produce(self, model_field: ModelInputField) -> bool: + """Return whether this mapper can produce ``model_field``.""" + return any( + _model_field_matches(produced, model_field) + for produced in self.produces + ) + + +class InputMapper(Protocol): + """Contract for user-event to model-input conversion.""" + + @property + def schema(self) -> InputMapperSchema: + """Return mapper metadata used for compatibility checks.""" + ... + + def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: + """Build session-start model inputs from a user-input trace.""" + ... + + def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: + """Build per-step model inputs from one user-input window.""" + ... + + +@dataclass(frozen=True, slots=True) +class StaticInputMapper: + """Mapper for fixed or already prepared model inputs. + + This is the no-op mapping path for runs that do not need live user + controls, such as prompt-only CLI runs or model-input replay. It consumes + no user events and returns configured model-facing payloads. + """ + + schema: InputMapperSchema + inputs: ModelInputs = field(default_factory=ModelInputs) + + @classmethod + def from_inputs( + cls, + *, + inputs: ModelInputs, + name: str = "static-input-mapper", + ) -> "StaticInputMapper": + """Create a static mapper whose produced fields come from ``inputs``.""" + produces = tuple( + ModelInputField(name=field_name, phase=phase, required=False) + for phase, values in ( + ("initial", inputs.initial), + ("step", inputs.step), + ) + for field_name in values + ) + return cls( + schema=InputMapperSchema(name=name, produces=produces), + inputs=inputs, + ) + + def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: + """Return configured initial model inputs.""" + del trace + return ModelInputs(initial=self.inputs.initial) + + def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: + """Return configured per-step model inputs.""" + del window + return ModelInputs(step=self.inputs.step) + + +@dataclass(frozen=True, slots=True) +class MappingCompatibility: + """Compatibility report for one source, model schema, and mapper.""" + + source_schema: UserInputSchema + model_schema: ModelInputSchema + mapper_schema: InputMapperSchema + missing_source_capabilities: tuple[UserInputCapability, ...] + missing_required_model_fields: tuple[ModelInputField, ...] + satisfied_required_model_fields: tuple[ModelInputField, ...] + available_optional_model_fields: tuple[ModelInputField, ...] + + @property + def can_drive(self) -> bool: + """Return whether this source can drive this model through the mapper.""" + return ( + not self.missing_source_capabilities + and not self.missing_required_model_fields + ) + + def raise_if_incompatible(self) -> None: + """Raise a compact error when this mapping cannot drive the model.""" + if self.can_drive: + return + problems: list[str] = [] + if self.missing_source_capabilities: + missing = ", ".join( + capability.event_kind + for capability in self.missing_source_capabilities + ) + problems.append(f"missing source capabilities: {missing}") + if self.missing_required_model_fields: + missing = ", ".join( + f"{field_def.phase}:{field_def.name}" + for field_def in self.missing_required_model_fields + ) + problems.append(f"missing required model inputs: {missing}") + raise ValueError( + f"Input mapper {self.mapper_schema.name!r} cannot drive model " + f"{self.model_schema.name!r} from source {self.source_schema.name!r}: " + + "; ".join(problems) + ) + + +def check_mapping_compatibility( + *, + source_schema: UserInputSchema, + model_schema: ModelInputSchema, + mapper_schema: InputMapperSchema, +) -> MappingCompatibility: + """Check whether a user-input source can drive a model through a mapper.""" + missing_source_capabilities = tuple( + capability + for capability in mapper_schema.consumes + if not source_schema.supports(capability) + ) + required_fields = model_schema.required_fields() + missing_required_model_fields = tuple( + field_def + for field_def in required_fields + if not mapper_schema.can_produce(field_def) + ) + satisfied_required_model_fields = tuple( + field_def + for field_def in required_fields + if mapper_schema.can_produce(field_def) + ) + available_optional_model_fields = tuple( + field_def + for field_def in model_schema.optional_fields() + if mapper_schema.can_produce(field_def) + ) + return MappingCompatibility( + source_schema=source_schema, + model_schema=model_schema, + mapper_schema=mapper_schema, + missing_source_capabilities=missing_source_capabilities, + missing_required_model_fields=missing_required_model_fields, + satisfied_required_model_fields=satisfied_required_model_fields, + available_optional_model_fields=available_optional_model_fields, + ) + + +def combine_mapper_schemas( + mapper_schemas: Sequence[InputMapperSchema], + *, + name: str = "input-mapper-set", +) -> InputMapperSchema: + """Combine independently declared mappers into one compatibility surface.""" + consumes: list[UserInputCapability] = [] + produces: list[ModelInputField] = [] + seen_consumes: set[UserInputCapability] = set() + seen_produces: set[ModelInputField] = set() + + for mapper_schema in mapper_schemas: + if not isinstance(mapper_schema, InputMapperSchema): + raise TypeError("mapper_schemas must contain InputMapperSchema objects.") + for capability in mapper_schema.consumes: + if capability not in seen_consumes: + consumes.append(capability) + seen_consumes.add(capability) + for field_def in mapper_schema.produces: + if field_def not in seen_produces: + produces.append(field_def) + seen_produces.add(field_def) + + return InputMapperSchema( + name=name, + consumes=tuple(consumes), + produces=tuple(produces), + ) + + +def check_mapping_set_compatibility( + *, + source_schema: UserInputSchema, + model_schema: ModelInputSchema, + mapper_schemas: Sequence[InputMapperSchema], + name: str = "input-mapper-set", +) -> MappingCompatibility: + """Check compatibility for a composed set of input mappers.""" + mapper_schema = combine_mapper_schemas(mapper_schemas, name=name) + return check_mapping_compatibility( + source_schema=source_schema, + model_schema=model_schema, + mapper_schema=mapper_schema, + ) diff --git a/flashdreams/tests/test_inference_inputs.py b/flashdreams/tests/test_inference_inputs.py new file mode 100644 index 000000000..e07f30da0 --- /dev/null +++ b/flashdreams/tests/test_inference_inputs.py @@ -0,0 +1,954 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +from flashdreams.inference import ( + InputMapperSchema, + ModelInputField, + ModelInputSchema, + ModelInputs, + StaticInputMapper, + UserInputCapability, + UserInputEvent, + UserInputSchema, + UserInputTrace, + UserInputWindow, + check_mapping_compatibility, + check_mapping_set_compatibility, + combine_mapper_schemas, + missing_required_inputs, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_user_input_trace_orders_events_and_slices_windows() -> None: + key_down = UserInputEvent( + timestamp_s=1.0, + kind="key_down", + payload={"key": "w"}, + ) + key_up = UserInputEvent( + timestamp_s=1.0, + kind="key_up", + payload={"key": "w"}, + ) + prompt = UserInputEvent( + timestamp_s=0.0, + kind="prompt_set", + payload={"prompt": "drive forward"}, + ) + + trace = UserInputTrace.from_events([key_down, key_up, prompt]) + + assert [event.kind for event in trace.events] == [ + "prompt_set", + "key_down", + "key_up", + ] + window = trace.window(start_s=0.5, end_s=1.0, include_end=True) + assert window.events == (key_down, key_up) + assert window.latest("key_down") is key_down + + +def test_static_startup_values_are_user_input_events() -> None: + trace = UserInputTrace.from_events( + [ + UserInputEvent( + timestamp_s=0.0, + kind="initial_frame_set", + payload={"image": b"encoded-image"}, + ), + UserInputEvent( + timestamp_s=0.0, + kind="scene_selected", + payload={"scene_id": "scene-a"}, + ), + ] + ) + + initial_frame = trace.latest("initial_frame_set") + scene = trace.latest("scene_selected") + assert initial_frame is not None + assert scene is not None + assert initial_frame.payload["image"] == b"encoded-image" + assert scene.payload["scene_id"] == "scene-a" + + +def test_user_input_schema_declares_and_validates_source_capabilities() -> None: + schema = UserInputSchema( + name="browser", + source_kind="live", + metadata={"transport": "webrtc"}, + capabilities=( + UserInputCapability( + event_kind="prompt_set", + payload_kind="text", + payload_fields=frozenset({"prompt"}), + metadata={"source_widget": "prompt-box"}, + ), + UserInputCapability( + event_kind="key_down", + payload_fields=frozenset({"key"}), + ), + ), + ) + + assert schema.supports( + UserInputCapability( + event_kind="prompt_set", + payload_kind="text", + payload_fields=frozenset({"prompt"}), + ) + ) + assert schema.metadata["transport"] == "webrtc" + assert schema.capabilities[0].metadata["source_widget"] == "prompt-box" + schema.validate_event( + UserInputEvent( + timestamp_s=0.0, + kind="prompt_set", + payload={"prompt": "a prompt"}, + ) + ) + with pytest.raises(ValueError, match="missing required fields"): + schema.validate_event( + UserInputEvent( + timestamp_s=0.0, + kind="prompt_set", + payload={"text": "wrong key"}, + ) + ) + + +def test_model_input_schema_declares_required_optional_and_update_metadata() -> None: + schema = ModelInputSchema( + name="steering-model", + metadata={"model_family": "example"}, + fields=( + ModelInputField( + name="prompt", + phase="initial", + required=True, + payload_kind="text", + update_policy="step_boundary", + lifecycle="cache_init", + metadata={"max_tokens": 256}, + ), + ModelInputField( + name="steering", + phase="step", + required=True, + lifecycle="step_input", + ), + ModelInputField(name="first_frame", phase="initial", required=False), + ), + ) + + assert [field.name for field in schema.required_fields()] == [ + "prompt", + "steering", + ] + assert [field.name for field in schema.optional_fields()] == ["first_frame"] + prompt_field = schema.field(name="prompt", phase="initial") + assert prompt_field is not None + assert prompt_field.update_policy == "step_boundary" + assert prompt_field.lifecycle == "cache_init" + assert prompt_field.metadata["max_tokens"] == 256 + assert schema.metadata["model_family"] == "example" + + +def test_model_inputs_report_missing_required_fields() -> None: + schema = ModelInputSchema( + fields=( + ModelInputField(name="prompt", phase="initial"), + ModelInputField(name="steering", phase="step"), + ) + ) + inputs = ModelInputs(initial={"prompt": "drive"}, step={}) + + missing = missing_required_inputs(inputs, schema) + + assert [(field.phase, field.name) for field in missing] == [ + ("step", "steering") + ] + + +def test_mapping_compatibility_reports_satisfied_missing_and_optional_fields() -> None: + source = UserInputSchema( + name="keyboard-app", + capabilities=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + UserInputCapability( + event_kind="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_kind="key_up", + payload_fields=frozenset({"key"}), + ), + ), + ) + model = ModelInputSchema( + name="drive-model", + fields=( + ModelInputField(name="prompt", phase="initial"), + ModelInputField(name="steering", phase="step"), + ModelInputField(name="first_frame", phase="initial", required=False), + ), + ) + mapper = InputMapperSchema( + name="keyboard-drive", + consumes=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + UserInputCapability( + event_kind="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_kind="key_up", + payload_fields=frozenset({"key"}), + ), + ), + produces=( + ModelInputField(name="prompt", phase="initial"), + ModelInputField(name="steering", phase="step"), + ), + ) + + compatibility = check_mapping_compatibility( + source_schema=source, + model_schema=model, + mapper_schema=mapper, + ) + + assert compatibility.can_drive + assert [field.name for field in compatibility.satisfied_required_model_fields] == [ + "prompt", + "steering", + ] + assert compatibility.available_optional_model_fields == () + + +def test_mapping_compatibility_reports_available_optional_model_input() -> None: + source = UserInputSchema( + capabilities=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + UserInputCapability( + event_kind="initial_frame_set", + payload_fields=frozenset({"image"}), + ), + ) + ) + model = ModelInputSchema( + fields=( + ModelInputField(name="prompt", phase="initial"), + ModelInputField(name="first_frame", phase="initial", required=False), + ) + ) + mapper = InputMapperSchema( + consumes=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + UserInputCapability( + event_kind="initial_frame_set", + payload_fields=frozenset({"image"}), + ), + ), + produces=( + ModelInputField(name="prompt", phase="initial"), + ModelInputField(name="first_frame", phase="initial", required=False), + ), + ) + + compatibility = check_mapping_compatibility( + source_schema=source, + model_schema=model, + mapper_schema=mapper, + ) + + assert compatibility.can_drive + assert [ + (field.phase, field.name) + for field in compatibility.available_optional_model_fields + ] == [("initial", "first_frame")] + + +def test_mapping_compatibility_reports_missing_required_model_input() -> None: + source = UserInputSchema( + capabilities=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + ) + ) + model = ModelInputSchema( + fields=( + ModelInputField(name="prompt", phase="initial"), + ModelInputField(name="camera_trajectory", phase="step"), + ) + ) + mapper = InputMapperSchema( + consumes=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + ), + produces=(ModelInputField(name="prompt", phase="initial"),), + ) + + compatibility = check_mapping_compatibility( + source_schema=source, + model_schema=model, + mapper_schema=mapper, + ) + + assert not compatibility.can_drive + assert [ + (field.phase, field.name) + for field in compatibility.missing_required_model_fields + ] == [("step", "camera_trajectory")] + with pytest.raises(ValueError, match="missing required model inputs"): + compatibility.raise_if_incompatible() + + +def test_mapping_compatibility_reports_missing_source_capability() -> None: + source = UserInputSchema( + capabilities=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + ) + ) + model = ModelInputSchema( + fields=(ModelInputField(name="steering", phase="step"),) + ) + mapper = InputMapperSchema( + consumes=( + UserInputCapability( + event_kind="controller_axis", + payload_fields=frozenset({"axis", "value"}), + ), + ), + produces=(ModelInputField(name="steering", phase="step"),), + ) + + compatibility = check_mapping_compatibility( + source_schema=source, + model_schema=model, + mapper_schema=mapper, + ) + + assert not compatibility.can_drive + assert [ + capability.event_kind + for capability in compatibility.missing_source_capabilities + ] == ["controller_axis"] + + +def test_mapping_set_compatibility_supports_composed_model_inputs() -> None: + source = UserInputSchema( + name="browser-with-controls", + capabilities=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + UserInputCapability( + event_kind="initial_frame_set", + payload_fields=frozenset({"image"}), + ), + UserInputCapability( + event_kind="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_kind="key_up", + payload_fields=frozenset({"key"}), + ), + ), + ) + model = ModelInputSchema( + name="lingbot-like", + fields=( + ModelInputField( + name="prompt", + phase="initial", + lifecycle="cache_init", + ), + ModelInputField( + name="first_frame", + phase="initial", + lifecycle="cache_init", + ), + ModelInputField( + name="camera_trajectory", + phase="step", + lifecycle="step_input", + ), + ), + ) + prompt_mapper = InputMapperSchema( + name="prompt", + consumes=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + ), + produces=( + ModelInputField( + name="prompt", + phase="initial", + lifecycle="cache_init", + ), + ), + ) + frame_mapper = InputMapperSchema( + name="first-frame", + consumes=( + UserInputCapability( + event_kind="initial_frame_set", + payload_fields=frozenset({"image"}), + ), + ), + produces=( + ModelInputField( + name="first_frame", + phase="initial", + lifecycle="cache_init", + ), + ), + ) + camera_mapper = InputMapperSchema( + name="keyboard-to-camera", + consumes=( + UserInputCapability( + event_kind="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_kind="key_up", + payload_fields=frozenset({"key"}), + ), + ), + produces=( + ModelInputField( + name="camera_trajectory", + phase="step", + lifecycle="step_input", + ), + ), + ) + + compatibility = check_mapping_set_compatibility( + source_schema=source, + model_schema=model, + mapper_schemas=(prompt_mapper, frame_mapper, camera_mapper), + name="browser-lingbot", + ) + + assert compatibility.can_drive + assert compatibility.mapper_schema.name == "browser-lingbot" + assert [field.name for field in compatibility.satisfied_required_model_fields] == [ + "prompt", + "first_frame", + "camera_trajectory", + ] + assert [capability.event_kind for capability in compatibility.mapper_schema.consumes] == [ + "prompt_set", + "initial_frame_set", + "key_down", + "key_up", + ] + + +def test_mapper_schema_combination_deduplicates_shared_capabilities() -> None: + shared_prompt = UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ) + prompt_field = ModelInputField(name="prompt", phase="initial") + + combined = combine_mapper_schemas( + ( + InputMapperSchema( + name="prompt-a", + consumes=(shared_prompt,), + produces=(prompt_field,), + ), + InputMapperSchema( + name="prompt-b", + consumes=(shared_prompt,), + produces=(prompt_field,), + ), + ) + ) + + assert combined.consumes == (shared_prompt,) + assert combined.produces == (prompt_field,) + + +def test_lifecycle_mismatch_does_not_satisfy_model_field() -> None: + source = UserInputSchema(name="fixed") + model = ModelInputSchema( + fields=( + ModelInputField( + name="video_dimensions", + phase="initial", + lifecycle="runtime_config", + ), + ) + ) + mapper = InputMapperSchema( + produces=( + ModelInputField( + name="video_dimensions", + phase="initial", + lifecycle="cache_init", + ), + ) + ) + + compatibility = check_mapping_compatibility( + source_schema=source, + model_schema=model, + mapper_schema=mapper, + ) + + assert not compatibility.can_drive + assert [ + (field.phase, field.name, field.lifecycle) + for field in compatibility.missing_required_model_fields + ] == [("initial", "video_dimensions", "runtime_config")] + + +def test_metadata_is_queryable_but_not_part_of_compatibility_matching() -> None: + source = UserInputSchema( + capabilities=( + UserInputCapability( + event_kind="pose_path_set", + payload_kind="path", + payload_fields=frozenset({"path"}), + metadata={"file_format": "npy"}, + ), + ), + metadata={"owner": "future-integration"}, + ) + model = ModelInputSchema( + fields=( + ModelInputField( + name="camera_trajectory", + phase="initial", + payload_kind="c2w_sequence", + lifecycle="rollout_binding", + metadata={"coordinates": "opencv_c2w"}, + ), + ), + metadata={"model_family": "future-world-model"}, + ) + mapper = InputMapperSchema( + consumes=( + UserInputCapability( + event_kind="pose_path_set", + payload_kind="path", + payload_fields=frozenset({"path"}), + metadata={"accepted_suffixes": (".npy",)}, + ), + ), + produces=( + ModelInputField( + name="camera_trajectory", + phase="initial", + payload_kind="c2w_sequence", + lifecycle="rollout_binding", + metadata={"shape": "[F,4,4]"}, + ), + ), + metadata={"mapper_family": "camera-path-loader"}, + ) + + compatibility = check_mapping_compatibility( + source_schema=source, + model_schema=model, + mapper_schema=mapper, + ) + + assert compatibility.can_drive + assert source.metadata["owner"] == "future-integration" + assert model.fields[0].metadata["coordinates"] == "opencv_c2w" + assert mapper.produces[0].metadata["shape"] == "[F,4,4]" + + +def test_sana_wm_like_schema_uses_open_ended_model_inputs() -> None: + source = UserInputSchema( + name="sana-wm-cli-like", + capabilities=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + UserInputCapability( + event_kind="initial_frame_set", + payload_fields=frozenset({"image"}), + ), + UserInputCapability( + event_kind="camera_action_set", + payload_fields=frozenset({"action"}), + metadata={"dsl": "-"}, + ), + UserInputCapability( + event_kind="intrinsics_set", + payload_fields=frozenset({"intrinsics"}), + metadata={"optional": True}, + ), + UserInputCapability( + event_kind="rollout_parameter_set", + payload_fields=frozenset({"name", "value"}), + ), + ), + ) + model = ModelInputSchema( + name="sana-wm-like", + fields=( + ModelInputField(name="prompt", phase="initial", lifecycle="cache_init"), + ModelInputField( + name="negative_prompt", + phase="initial", + required=False, + lifecycle="cache_init", + ), + ModelInputField( + name="first_frame", + phase="initial", + lifecycle="cache_init", + ), + ModelInputField( + name="camera_trajectory_c2w", + phase="initial", + payload_kind="c2w_sequence", + lifecycle="rollout_binding", + metadata={"shape": "[F,4,4]"}, + ), + ModelInputField( + name="camera_intrinsics_vec4", + phase="initial", + required=False, + payload_kind="intrinsics_vec4_sequence", + lifecycle="rollout_binding", + metadata={"shape": "[F,4]"}, + ), + ModelInputField( + name="stage1_sampling", + phase="initial", + lifecycle="rollout_binding", + metadata={"fields": ("steps", "cfg_scale", "flow_shift", "seed")}, + ), + ModelInputField( + name="streaming_chunking", + phase="initial", + required=False, + lifecycle="rollout_binding", + metadata={"fields": ("num_frame_per_block", "cached_blocks")}, + ), + ), + metadata={"model_family": "sana-wm"}, + ) + prompt_mapper = InputMapperSchema( + produces=(ModelInputField(name="prompt", phase="initial"),), + consumes=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + ), + ) + frame_mapper = InputMapperSchema( + produces=(ModelInputField(name="first_frame", phase="initial"),), + consumes=( + UserInputCapability( + event_kind="initial_frame_set", + payload_fields=frozenset({"image"}), + ), + ), + ) + camera_mapper = InputMapperSchema( + produces=( + ModelInputField( + name="camera_trajectory_c2w", + phase="initial", + payload_kind="c2w_sequence", + lifecycle="rollout_binding", + ), + ModelInputField( + name="camera_intrinsics_vec4", + phase="initial", + payload_kind="intrinsics_vec4_sequence", + lifecycle="rollout_binding", + ), + ), + consumes=( + UserInputCapability( + event_kind="camera_action_set", + payload_fields=frozenset({"action"}), + ), + UserInputCapability( + event_kind="intrinsics_set", + payload_fields=frozenset({"intrinsics"}), + ), + ), + ) + sampling_mapper = InputMapperSchema( + produces=( + ModelInputField( + name="stage1_sampling", + phase="initial", + lifecycle="rollout_binding", + ), + ), + consumes=( + UserInputCapability( + event_kind="rollout_parameter_set", + payload_fields=frozenset({"name", "value"}), + ), + ), + ) + + compatibility = check_mapping_set_compatibility( + source_schema=source, + model_schema=model, + mapper_schemas=(prompt_mapper, frame_mapper, camera_mapper, sampling_mapper), + ) + + assert compatibility.can_drive + assert [field.name for field in compatibility.satisfied_required_model_fields] == [ + "prompt", + "first_frame", + "camera_trajectory_c2w", + "stage1_sampling", + ] + assert [field.name for field in compatibility.available_optional_model_fields] == [ + "camera_intrinsics_vec4", + ] + + +@dataclass(frozen=True) +class PromptMapper: + schema: InputMapperSchema = InputMapperSchema( + name="prompt", + consumes=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + ), + produces=(ModelInputField(name="prompt", phase="initial"),), + ) + + def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: + prompt = trace.latest("prompt_set") + return ModelInputs.initial_only( + {"prompt": prompt.payload["prompt"]} if prompt is not None else {} + ) + + def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: + del window + return ModelInputs() + + +@dataclass(frozen=True) +class InitialFrameMapper: + schema: InputMapperSchema = InputMapperSchema( + name="initial-frame", + consumes=( + UserInputCapability( + event_kind="initial_frame_set", + payload_fields=frozenset({"image"}), + ), + ), + produces=( + ModelInputField(name="first_frame", phase="initial", required=False), + ), + ) + + def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: + event = trace.latest("initial_frame_set") + return ModelInputs.initial_only( + {"first_frame": event.payload["image"]} if event is not None else {} + ) + + def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: + del window + return ModelInputs() + + +@dataclass(frozen=True) +class KeyboardSteeringMapper: + schema: InputMapperSchema = InputMapperSchema( + name="keyboard-steering", + consumes=( + UserInputCapability( + event_kind="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_kind="key_up", + payload_fields=frozenset({"key"}), + ), + ), + produces=(ModelInputField(name="steering", phase="step"),), + ) + + def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: + del trace + return ModelInputs() + + def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: + steering = 0.0 + for event in window.events: + key = event.payload.get("key") + if event.kind == "key_down" and key == "a": + steering = 1.0 + elif event.kind == "key_down" and key == "d": + steering = -1.0 + elif event.kind == "key_up" and key in {"a", "d"}: + steering = 0.0 + return ModelInputs.step_only({"steering": steering}) + + +@dataclass(frozen=True) +class CameraTrajectoryMapper: + schema: InputMapperSchema = InputMapperSchema( + name="camera-trajectory", + consumes=( + UserInputCapability( + event_kind="camera_pose", + payload_fields=frozenset({"pose"}), + ), + ), + produces=(ModelInputField(name="camera_trajectory", phase="step"),), + ) + + def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: + del trace + return ModelInputs() + + def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: + return ModelInputs.step_only( + { + "camera_trajectory": tuple( + event.payload["pose"] + for event in window.events_of_kind("camera_pose") + ) + } + ) + + +def test_fake_prompt_and_initial_frame_mappers_build_initial_model_inputs() -> None: + trace = UserInputTrace.from_events( + [ + UserInputEvent( + timestamp_s=0.0, + kind="prompt_set", + payload={"prompt": "a road"}, + ), + UserInputEvent( + timestamp_s=0.0, + kind="initial_frame_set", + payload={"image": b"frame"}, + ), + ] + ) + + assert PromptMapper().build_initial_inputs(trace).initial == { + "prompt": "a road" + } + assert InitialFrameMapper().build_initial_inputs(trace).initial == { + "first_frame": b"frame" + } + + +def test_static_input_mapper_supports_fixed_model_inputs() -> None: + inputs = ModelInputs( + initial={"prompt": "fixed prompt"}, + step={"camera_trajectory": ("pose-a", "pose-b")}, + ) + mapper = StaticInputMapper.from_inputs(inputs=inputs, name="fixed") + source = UserInputSchema(name="no-live-controls") + model = ModelInputSchema( + fields=( + ModelInputField(name="prompt", phase="initial"), + ModelInputField(name="camera_trajectory", phase="step"), + ) + ) + + compatibility = check_mapping_compatibility( + source_schema=source, + model_schema=model, + mapper_schema=mapper.schema, + ) + + assert compatibility.can_drive + assert mapper.build_initial_inputs(UserInputTrace()).initial == { + "prompt": "fixed prompt" + } + assert mapper.build_step_inputs(UserInputWindow(start_s=0.0, end_s=1.0)).step == { + "camera_trajectory": ("pose-a", "pose-b") + } + + +def test_fake_keyboard_mapper_builds_step_model_inputs() -> None: + trace = UserInputTrace.from_events( + [ + UserInputEvent(timestamp_s=0.1, kind="key_down", payload={"key": "a"}), + UserInputEvent(timestamp_s=0.2, kind="key_up", payload={"key": "a"}), + UserInputEvent(timestamp_s=0.3, kind="key_down", payload={"key": "d"}), + ] + ) + + inputs = KeyboardSteeringMapper().build_step_inputs( + trace.window(start_s=0.0, end_s=0.4) + ) + + assert inputs.step == {"steering": -1.0} + + +def test_fake_camera_trajectory_mapper_builds_step_model_inputs() -> None: + trace = UserInputTrace.from_events( + [ + UserInputEvent( + timestamp_s=0.1, + kind="camera_pose", + payload={"pose": "pose-a"}, + ), + UserInputEvent( + timestamp_s=0.2, + kind="camera_pose", + payload={"pose": "pose-b"}, + ), + ] + ) + + inputs = CameraTrajectoryMapper().build_step_inputs( + trace.window(start_s=0.0, end_s=0.3) + ) + + assert inputs.step == {"camera_trajectory": ("pose-a", "pose-b")} From d58ace2223ee983e916af22991fc703296dcf9da Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Tue, 4 Aug 2026 11:20:34 -0700 Subject: [PATCH 2/7] Fix issues found by Claude --- ...inference_runtime_inputs_implementation.md | 37 ++- flashdreams/flashdreams/inference/__init__.py | 2 + flashdreams/flashdreams/inference/inputs.py | 230 ++++++++++++++---- flashdreams/tests/test_inference_inputs.py | 222 ++++++++++++++++- 4 files changed, 433 insertions(+), 58 deletions(-) diff --git a/docs/inference_runtime_inputs_implementation.md b/docs/inference_runtime_inputs_implementation.md index e6cda3727..465319b8c 100644 --- a/docs/inference_runtime_inputs_implementation.md +++ b/docs/inference_runtime_inputs_implementation.md @@ -59,7 +59,9 @@ Examples: `UserInputTrace` stores events in deterministic timestamp order and can slice a `UserInputWindow` for one runtime step. Events with equal timestamps preserve -their original order. +their original order. A `UserInputWindow` enforces its own invariant: every +event it holds must fall inside `[start_s, end_s]`, so a directly constructed +window cannot silently disagree with its bounds. ```python from flashdreams.inference import UserInputEvent, UserInputTrace @@ -128,6 +130,8 @@ camera coordinate system, or model-specific units. Capabilities and schemas also carry optional `metadata`. This is for query-time hints such as source type, UI widget, file suffixes, units, coordinate frame, or schema URI. Metadata is intentionally not part of compatibility matching. +Because metadata is an open-ended pass-through, its keys are validated (they +must be non-empty strings) but never rewritten, so a key round-trips unchanged. ## Model Inputs @@ -260,6 +264,11 @@ static_mapper = StaticInputMapper.from_inputs( ) ``` +A mapper schema is a hand-written declaration, so it can drift from what +`build_initial_inputs` / `build_step_inputs` actually return. +`undeclared_model_inputs()` reports payload keys a mapper produced but did not +declare, which keeps mapper tests honest about the compatibility surface. + Mapper schemas can also be combined for compatibility checks. This matches the current supported model inventory: a run may select one mapping for prompt events, another for first-frame events, and another for keyboard or controller @@ -288,7 +297,8 @@ schema, model schema, and mapper schema, it reports: - source capabilities the mapper needs but the source lacks; - required model fields the mapper cannot produce; - required model fields that are satisfied; -- optional model fields that can be enabled. +- optional model fields that can be enabled; +- mappers dropped because the source cannot feed them. ```python from flashdreams.inference import check_mapping_compatibility @@ -309,7 +319,22 @@ the model run will succeed. Use `check_mapping_set_compatibility()` when the selected mapping is composed from multiple mapper schemas. It returns the same `MappingCompatibility` report -after combining mapper inputs and outputs. +for the composed mapping. + +Compatibility is evaluated per mapper rather than over a flattened bag of +capabilities, so each mapper keeps its own `consumes`/`produces` link. A mapper +the source cannot feed is dropped and reported in `unavailable_mapper_schemas`; +it costs only the model inputs that mapper produced. This means: + +- a dropped mapper that produced only optional fields degrades the run instead + of vetoing it, and those fields are correctly absent from + `available_optional_model_fields`; +- a dropped mapper that was the only producer of a required field still blocks, + and its unmet source capabilities are reported in + `missing_source_capabilities`; +- `satisfied_required_model_fields` and `available_optional_model_fields` name + only fields this source can really produce, not every field some mapper + declared. ## What This Does Not Validate @@ -346,6 +371,12 @@ They cover: - missing source capabilities; - optional model inputs becoming available only when mapper support exists; - composed mapper-set compatibility; +- graceful degradation when an optional mapper cannot be fed, and blocking when + a required one cannot; +- metadata merging when duplicate declarations collapse during mapper-set + combination; +- event hashability, window bound enforcement, and metadata key round-tripping; +- mappers producing only model inputs their schema declares; - SANA-WM-like model-specific inputs without importing the SANA integration; - fake prompt, initial-frame, keyboard-to-steering, and camera-trajectory mappers; diff --git a/flashdreams/flashdreams/inference/__init__.py b/flashdreams/flashdreams/inference/__init__.py index e22df109b..327f9189a 100644 --- a/flashdreams/flashdreams/inference/__init__.py +++ b/flashdreams/flashdreams/inference/__init__.py @@ -27,6 +27,7 @@ check_mapping_set_compatibility, combine_mapper_schemas, missing_required_inputs, + undeclared_model_inputs, ) __all__ = [ @@ -47,4 +48,5 @@ "check_mapping_set_compatibility", "combine_mapper_schemas", "missing_required_inputs", + "undeclared_model_inputs", ] diff --git a/flashdreams/flashdreams/inference/inputs.py b/flashdreams/flashdreams/inference/inputs.py index 5216f455e..f97ec51b2 100644 --- a/flashdreams/flashdreams/inference/inputs.py +++ b/flashdreams/flashdreams/inference/inputs.py @@ -6,10 +6,10 @@ from __future__ import annotations from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from math import isfinite from types import MappingProxyType -from typing import Any, Literal, Protocol, cast +from typing import Any, Literal, Protocol, cast, runtime_checkable InputPhase = Literal["initial", "step"] @@ -34,12 +34,31 @@ def _normalized_payload_fields(values: Iterable[str]) -> frozenset[str]: def _immutable_metadata_map(values: Mapping[str, Any]) -> Mapping[str, Any]: + """Copy metadata into a read-only map without rewriting its keys. + + Metadata is an open-ended pass-through for adapter hints, so keys are + validated but never normalized: a key must round-trip unchanged. + """ normalized: dict[str, Any] = {} for key, value in values.items(): - normalized[_normalized_token(str(key), field_name="metadata key")] = value + if not isinstance(key, str): + raise TypeError(f"metadata keys must be strings, got {type(key).__name__}.") + if not key: + raise ValueError("metadata keys must be non-empty strings.") + normalized[key] = value return MappingProxyType(normalized) +def _merged_metadata( + first: Mapping[str, Any], + second: Mapping[str, Any], +) -> Mapping[str, Any]: + """Union two metadata maps, keeping ``first`` on conflicting keys.""" + merged = dict(second) + merged.update(first) + return merged + + def _validate_phase(value: str) -> InputPhase: if value not in {"initial", "step"}: raise ValueError(f"phase must be 'initial' or 'step', got {value!r}.") @@ -106,6 +125,12 @@ def __post_init__(self) -> None: _normalized_optional_token(self.source, field_name="source"), ) + def __hash__(self) -> int: + # ``payload`` is an arbitrary mapping and therefore unhashable, so the + # generated hash would raise. Hash the identity fields instead; equality + # still compares payloads, and equal events still hash equal. + return hash((self.timestamp_s, self.kind, self.session_id, self.source)) + @dataclass(frozen=True, slots=True) class UserInputWindow: @@ -124,7 +149,14 @@ def __post_init__(self) -> None: raise ValueError("end_s must be >= start_s.") object.__setattr__(self, "start_s", start_s) object.__setattr__(self, "end_s", end_s) - object.__setattr__(self, "events", _sorted_events(self.events)) + events = _sorted_events(self.events) + for event in events: + if not start_s <= event.timestamp_s <= end_s: + raise ValueError( + f"Event {event.kind!r} at t={event.timestamp_s} is outside " + f"window [{start_s}, {end_s}]." + ) + object.__setattr__(self, "events", events) def events_of_kind(self, kind: str) -> tuple[UserInputEvent, ...]: """Return all events of ``kind`` inside this window.""" @@ -252,9 +284,7 @@ def is_satisfied_by(self, provider: "UserInputCapability") -> bool: or provider.payload_kind is None or self.payload_kind == provider.payload_kind ) - return payload_kind_ok and self.payload_fields.issubset( - provider.payload_fields - ) + return payload_kind_ok and self.payload_fields.issubset(provider.payload_fields) @dataclass(frozen=True, slots=True) @@ -309,8 +339,7 @@ def validate_event(self, event: UserInputEvent) -> None: ) payload_keys = set(event.payload) if not any( - capability.payload_fields.issubset(payload_keys) - for capability in matching + capability.payload_fields.issubset(payload_keys) for capability in matching ): expected = sorted( { @@ -432,7 +461,7 @@ def optional_fields( if not field_def.required and (phase is None or field_def.phase == phase) ) - def field(self, *, name: str, phase: InputPhase) -> ModelInputField | None: + def field_for(self, *, name: str, phase: InputPhase) -> ModelInputField | None: """Return one field definition, if present.""" name = _normalized_token(name, field_name="model input name") phase = _validate_phase(phase) @@ -490,6 +519,28 @@ def missing_required_inputs( ) +def undeclared_model_inputs( + inputs: ModelInputs, + mapper_schema: "InputMapperSchema", +) -> tuple[tuple[InputPhase, str], ...]: + """Return payload keys a mapper produced but did not declare in its schema. + + Mapper schemas are hand-written, so they can drift from what + ``build_initial_inputs`` / ``build_step_inputs`` actually return. Mapper + tests can use this to keep the declared compatibility surface honest. + """ + declared = { + (field_def.phase, field_def.name) for field_def in mapper_schema.produces + } + phases: tuple[InputPhase, ...] = ("initial", "step") + return tuple( + (phase, key) + for phase in phases + for key in inputs.for_phase(phase) + if (phase, key) not in declared + ) + + @dataclass(frozen=True, slots=True) class InputMapperSchema: """Metadata for a mapper that converts user events into model inputs.""" @@ -522,11 +573,11 @@ def __post_init__(self) -> None: def can_produce(self, model_field: ModelInputField) -> bool: """Return whether this mapper can produce ``model_field``.""" return any( - _model_field_matches(produced, model_field) - for produced in self.produces + _model_field_matches(produced, model_field) for produced in self.produces ) +@runtime_checkable class InputMapper(Protocol): """Contract for user-event to model-input conversion.""" @@ -590,7 +641,14 @@ def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: @dataclass(frozen=True, slots=True) class MappingCompatibility: - """Compatibility report for one source, model schema, and mapper.""" + """Compatibility report for one source, model schema, and mapping. + + ``mapper_schema`` is the full requested mapping surface. Mappers whose + consumed capabilities the source cannot provide are reported separately in + ``unavailable_mapper_schemas`` and are excluded from the satisfied/available + model-field reports, so those lists only name fields that can really be + produced by this source. + """ source_schema: UserInputSchema model_schema: ModelInputSchema @@ -599,15 +657,25 @@ class MappingCompatibility: missing_required_model_fields: tuple[ModelInputField, ...] satisfied_required_model_fields: tuple[ModelInputField, ...] available_optional_model_fields: tuple[ModelInputField, ...] + unavailable_mapper_schemas: tuple[InputMapperSchema, ...] = () @property def can_drive(self) -> bool: - """Return whether this source can drive this model through the mapper.""" + """Return whether this source can drive this model through the mapping. + + Mappers that the source cannot feed do not block the run unless they + were the only way to produce a required model input. + """ return ( not self.missing_source_capabilities and not self.missing_required_model_fields ) + @property + def unavailable_mapper_names(self) -> tuple[str, ...]: + """Return names of mappers dropped because the source cannot feed them.""" + return tuple(schema.name for schema in self.unavailable_mapper_schemas) + def raise_if_incompatible(self) -> None: """Raise a compact error when this mapping cannot drive the model.""" if self.can_drive: @@ -615,8 +683,7 @@ def raise_if_incompatible(self) -> None: problems: list[str] = [] if self.missing_source_capabilities: missing = ", ".join( - capability.event_kind - for capability in self.missing_source_capabilities + capability.event_kind for capability in self.missing_source_capabilities ) problems.append(f"missing source capabilities: {missing}") if self.missing_required_model_fields: @@ -625,6 +692,10 @@ def raise_if_incompatible(self) -> None: for field_def in self.missing_required_model_fields ) problems.append(f"missing required model inputs: {missing}") + if self.unavailable_mapper_schemas: + problems.append( + "unavailable mappers: " + ", ".join(self.unavailable_mapper_names) + ) raise ValueError( f"Input mapper {self.mapper_schema.name!r} cannot drive model " f"{self.model_schema.name!r} from source {self.source_schema.name!r}: " @@ -632,42 +703,87 @@ def raise_if_incompatible(self) -> None: ) -def check_mapping_compatibility( +def _mapper_is_feedable( + source_schema: UserInputSchema, + mapper_schema: InputMapperSchema, +) -> bool: + return all( + source_schema.supports(capability) for capability in mapper_schema.consumes + ) + + +def _build_compatibility( *, source_schema: UserInputSchema, model_schema: ModelInputSchema, - mapper_schema: InputMapperSchema, + mapper_schemas: Sequence[InputMapperSchema], + reported_schema: InputMapperSchema, ) -> MappingCompatibility: - """Check whether a user-input source can drive a model through a mapper.""" - missing_source_capabilities = tuple( - capability - for capability in mapper_schema.consumes - if not source_schema.supports(capability) - ) + feedable: list[InputMapperSchema] = [] + unavailable: list[InputMapperSchema] = [] + for mapper_schema in mapper_schemas: + if _mapper_is_feedable(source_schema, mapper_schema): + feedable.append(mapper_schema) + else: + unavailable.append(mapper_schema) + + usable = combine_mapper_schemas(feedable, name=reported_schema.name) required_fields = model_schema.required_fields() missing_required_model_fields = tuple( - field_def - for field_def in required_fields - if not mapper_schema.can_produce(field_def) + field_def for field_def in required_fields if not usable.can_produce(field_def) ) satisfied_required_model_fields = tuple( - field_def - for field_def in required_fields - if mapper_schema.can_produce(field_def) + field_def for field_def in required_fields if usable.can_produce(field_def) ) available_optional_model_fields = tuple( field_def for field_def in model_schema.optional_fields() - if mapper_schema.can_produce(field_def) + if usable.can_produce(field_def) ) + + # Only capabilities that block a required model input make the mapping + # unusable. A dropped mapper that fed nothing but optional fields degrades + # instead of vetoing the run. + missing_source_capabilities: list[UserInputCapability] = [] + seen_missing: set[UserInputCapability] = set() + for mapper_schema in unavailable: + if not any( + mapper_schema.can_produce(field_def) + for field_def in missing_required_model_fields + ): + continue + for capability in mapper_schema.consumes: + if source_schema.supports(capability) or capability in seen_missing: + continue + seen_missing.add(capability) + missing_source_capabilities.append(capability) + return MappingCompatibility( source_schema=source_schema, model_schema=model_schema, - mapper_schema=mapper_schema, - missing_source_capabilities=missing_source_capabilities, + mapper_schema=reported_schema, + missing_source_capabilities=tuple(missing_source_capabilities), missing_required_model_fields=missing_required_model_fields, satisfied_required_model_fields=satisfied_required_model_fields, available_optional_model_fields=available_optional_model_fields, + unavailable_mapper_schemas=tuple(unavailable), + ) + + +def check_mapping_compatibility( + *, + source_schema: UserInputSchema, + model_schema: ModelInputSchema, + mapper_schema: InputMapperSchema, +) -> MappingCompatibility: + """Check whether a user-input source can drive a model through a mapper.""" + if not isinstance(mapper_schema, InputMapperSchema): + raise TypeError("mapper_schema must be an InputMapperSchema object.") + return _build_compatibility( + source_schema=source_schema, + model_schema=model_schema, + mapper_schemas=(mapper_schema,), + reported_schema=mapper_schema, ) @@ -676,23 +792,44 @@ def combine_mapper_schemas( *, name: str = "input-mapper-set", ) -> InputMapperSchema: - """Combine independently declared mappers into one compatibility surface.""" + """Combine independently declared mappers into one compatibility surface. + + Duplicate entries are collapsed. Because ``metadata`` is excluded from + equality, the metadata of collapsed duplicates is merged rather than + dropped, with the first declaration winning on conflicting keys. + """ consumes: list[UserInputCapability] = [] produces: list[ModelInputField] = [] - seen_consumes: set[UserInputCapability] = set() - seen_produces: set[ModelInputField] = set() + consumes_index: dict[UserInputCapability, int] = {} + produces_index: dict[ModelInputField, int] = {} for mapper_schema in mapper_schemas: if not isinstance(mapper_schema, InputMapperSchema): raise TypeError("mapper_schemas must contain InputMapperSchema objects.") for capability in mapper_schema.consumes: - if capability not in seen_consumes: + index = consumes_index.get(capability) + if index is None: + consumes_index[capability] = len(consumes) consumes.append(capability) - seen_consumes.add(capability) + elif capability.metadata: + consumes[index] = replace( + consumes[index], + metadata=_merged_metadata( + consumes[index].metadata, capability.metadata + ), + ) for field_def in mapper_schema.produces: - if field_def not in seen_produces: + index = produces_index.get(field_def) + if index is None: + produces_index[field_def] = len(produces) produces.append(field_def) - seen_produces.add(field_def) + elif field_def.metadata: + produces[index] = replace( + produces[index], + metadata=_merged_metadata( + produces[index].metadata, field_def.metadata + ), + ) return InputMapperSchema( name=name, @@ -708,10 +845,15 @@ def check_mapping_set_compatibility( mapper_schemas: Sequence[InputMapperSchema], name: str = "input-mapper-set", ) -> MappingCompatibility: - """Check compatibility for a composed set of input mappers.""" - mapper_schema = combine_mapper_schemas(mapper_schemas, name=name) - return check_mapping_compatibility( + """Check compatibility for a composed set of input mappers. + + Each mapper keeps its own ``consumes``/``produces`` link, so a mapper the + source cannot feed only costs the model inputs that mapper produced. + """ + mapper_schemas = tuple(mapper_schemas) + return _build_compatibility( source_schema=source_schema, model_schema=model_schema, - mapper_schema=mapper_schema, + mapper_schemas=mapper_schemas, + reported_schema=combine_mapper_schemas(mapper_schemas, name=name), ) diff --git a/flashdreams/tests/test_inference_inputs.py b/flashdreams/tests/test_inference_inputs.py index e07f30da0..b3d090a7a 100644 --- a/flashdreams/tests/test_inference_inputs.py +++ b/flashdreams/tests/test_inference_inputs.py @@ -4,10 +4,12 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Any import pytest from flashdreams.inference import ( + InputMapper, InputMapperSchema, ModelInputField, ModelInputSchema, @@ -22,6 +24,7 @@ check_mapping_set_compatibility, combine_mapper_schemas, missing_required_inputs, + undeclared_model_inputs, ) pytestmark = pytest.mark.ci_cpu @@ -154,7 +157,7 @@ def test_model_input_schema_declares_required_optional_and_update_metadata() -> "steering", ] assert [field.name for field in schema.optional_fields()] == ["first_frame"] - prompt_field = schema.field(name="prompt", phase="initial") + prompt_field = schema.field_for(name="prompt", phase="initial") assert prompt_field is not None assert prompt_field.update_policy == "step_boundary" assert prompt_field.lifecycle == "cache_init" @@ -173,9 +176,7 @@ def test_model_inputs_report_missing_required_fields() -> None: missing = missing_required_inputs(inputs, schema) - assert [(field.phase, field.name) for field in missing] == [ - ("step", "steering") - ] + assert [(field.phase, field.name) for field in missing] == [("step", "steering")] def test_mapping_compatibility_reports_satisfied_missing_and_optional_fields() -> None: @@ -338,9 +339,7 @@ def test_mapping_compatibility_reports_missing_source_capability() -> None: ), ) ) - model = ModelInputSchema( - fields=(ModelInputField(name="steering", phase="step"),) - ) + model = ModelInputSchema(fields=(ModelInputField(name="steering", phase="step"),)) mapper = InputMapperSchema( consumes=( UserInputCapability( @@ -473,7 +472,9 @@ def test_mapping_set_compatibility_supports_composed_model_inputs() -> None: "first_frame", "camera_trajectory", ] - assert [capability.event_kind for capability in compatibility.mapper_schema.consumes] == [ + assert [ + capability.event_kind for capability in compatibility.mapper_schema.consumes + ] == [ "prompt_set", "initial_frame_set", "key_down", @@ -481,6 +482,89 @@ def test_mapping_set_compatibility_supports_composed_model_inputs() -> None: ] +def test_unfeedable_optional_mapper_degrades_instead_of_blocking_the_run() -> None: + source = UserInputSchema( + name="prompt-only", + capabilities=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + ), + ) + model = ModelInputSchema( + fields=( + ModelInputField(name="prompt", phase="initial"), + ModelInputField(name="steering", phase="step", required=False), + ) + ) + prompt_mapper = InputMapperSchema( + name="prompt", + consumes=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + ), + ), + produces=(ModelInputField(name="prompt", phase="initial"),), + ) + gamepad_mapper = InputMapperSchema( + name="gamepad-steering", + consumes=( + UserInputCapability( + event_kind="controller_axis", + payload_fields=frozenset({"axis", "value"}), + ), + ), + produces=(ModelInputField(name="steering", phase="step", required=False),), + ) + + compatibility = check_mapping_set_compatibility( + source_schema=source, + model_schema=model, + mapper_schemas=(prompt_mapper, gamepad_mapper), + ) + + assert compatibility.can_drive + assert compatibility.unavailable_mapper_names == ("gamepad-steering",) + assert compatibility.missing_source_capabilities == () + assert [field.name for field in compatibility.satisfied_required_model_fields] == [ + "prompt" + ] + # The optional field is not reported as available: nothing can produce it. + assert compatibility.available_optional_model_fields == () + + +def test_unfeedable_required_mapper_still_blocks_the_run() -> None: + source = UserInputSchema(name="prompt-only") + model = ModelInputSchema(fields=(ModelInputField(name="steering", phase="step"),)) + gamepad_mapper = InputMapperSchema( + name="gamepad-steering", + consumes=( + UserInputCapability( + event_kind="controller_axis", + payload_fields=frozenset({"axis", "value"}), + ), + ), + produces=(ModelInputField(name="steering", phase="step"),), + ) + + compatibility = check_mapping_set_compatibility( + source_schema=source, + model_schema=model, + mapper_schemas=(gamepad_mapper,), + ) + + assert not compatibility.can_drive + assert [ + capability.event_kind + for capability in compatibility.missing_source_capabilities + ] == ["controller_axis"] + assert compatibility.unavailable_mapper_names == ("gamepad-steering",) + with pytest.raises(ValueError, match="unavailable mappers: gamepad-steering"): + compatibility.raise_if_incompatible() + + def test_mapper_schema_combination_deduplicates_shared_capabilities() -> None: shared_prompt = UserInputCapability( event_kind="prompt_set", @@ -507,6 +591,80 @@ def test_mapper_schema_combination_deduplicates_shared_capabilities() -> None: assert combined.produces == (prompt_field,) +def test_mapper_schema_combination_merges_metadata_of_collapsed_duplicates() -> None: + combined = combine_mapper_schemas( + ( + InputMapperSchema( + name="prompt-a", + consumes=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + metadata={"widget": "prompt-box", "owner": "app-a"}, + ), + ), + ), + InputMapperSchema( + name="prompt-b", + consumes=( + UserInputCapability( + event_kind="prompt_set", + payload_fields=frozenset({"prompt"}), + metadata={"max_length": 512, "owner": "app-b"}, + ), + ), + ), + ) + ) + + assert len(combined.consumes) == 1 + assert dict(combined.consumes[0].metadata) == { + "widget": "prompt-box", + "max_length": 512, + "owner": "app-a", + } + + +def test_user_input_event_is_hashable_despite_mapping_payload() -> None: + event = UserInputEvent( + timestamp_s=0.5, + kind="key_down", + payload={"key": "w", "modifiers": ["shift"]}, + ) + same = UserInputEvent( + timestamp_s=0.5, + kind="key_down", + payload={"key": "w", "modifiers": ["shift"]}, + ) + + assert {event, same} == {event} + assert hash(event) == hash(same) + + +def test_user_input_window_rejects_events_outside_its_bounds() -> None: + with pytest.raises(ValueError, match="outside window"): + UserInputWindow( + start_s=0.0, + end_s=1.0, + events=(UserInputEvent(timestamp_s=2.0, kind="key_down"),), + ) + + +def test_metadata_keys_round_trip_unchanged() -> None: + capability = UserInputCapability( + event_kind="prompt_set", + metadata={" spaced key ": "kept", "schema.uri": "urn:example"}, + ) + + assert dict(capability.metadata) == { + " spaced key ": "kept", + "schema.uri": "urn:example", + } + non_string_keys: dict[Any, Any] = {1: "coerced"} + with pytest.raises(TypeError, match="metadata keys must be strings"): + UserInputCapability(event_kind="prompt_set", metadata=non_string_keys) + + def test_lifecycle_mismatch_does_not_satisfy_model_field() -> None: source = UserInputSchema(name="fixed") model = ModelInputSchema( @@ -878,9 +1036,7 @@ def test_fake_prompt_and_initial_frame_mappers_build_initial_model_inputs() -> N ] ) - assert PromptMapper().build_initial_inputs(trace).initial == { - "prompt": "a road" - } + assert PromptMapper().build_initial_inputs(trace).initial == {"prompt": "a road"} assert InitialFrameMapper().build_initial_inputs(trace).initial == { "first_frame": b"frame" } @@ -931,6 +1087,50 @@ def test_fake_keyboard_mapper_builds_step_model_inputs() -> None: assert inputs.step == {"steering": -1.0} +@pytest.mark.parametrize( + "mapper", + [ + PromptMapper(), + InitialFrameMapper(), + KeyboardSteeringMapper(), + CameraTrajectoryMapper(), + StaticInputMapper.from_inputs( + inputs=ModelInputs(initial={"prompt": "fixed"}, step={"steering": 0.0}), + ), + ], +) +def test_mappers_only_produce_model_inputs_they_declare(mapper: InputMapper) -> None: + trace = UserInputTrace.from_events( + [ + UserInputEvent( + timestamp_s=0.0, + kind="prompt_set", + payload={"prompt": "a road"}, + ), + UserInputEvent( + timestamp_s=0.0, + kind="initial_frame_set", + payload={"image": b"frame"}, + ), + UserInputEvent(timestamp_s=0.1, kind="key_down", payload={"key": "a"}), + UserInputEvent( + timestamp_s=0.2, + kind="camera_pose", + payload={"pose": "pose-a"}, + ), + ] + ) + window = trace.window(start_s=0.0, end_s=1.0) + + assert isinstance(mapper, InputMapper) + assert ( + undeclared_model_inputs(mapper.build_initial_inputs(trace), mapper.schema) == () + ) + assert ( + undeclared_model_inputs(mapper.build_step_inputs(window), mapper.schema) == () + ) + + def test_fake_camera_trajectory_mapper_builds_step_model_inputs() -> None: trace = UserInputTrace.from_events( [ From 4f391164355a8da24b1e20b9a0f6de3392556eb0 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Tue, 4 Aug 2026 14:27:25 -0700 Subject: [PATCH 3/7] Rewrite based on discussion, port after merge --- flashdreams/flashdreams/inference/__init__.py | 52 - flashdreams/flashdreams/inference/inputs.py | 859 ------------ flashdreams/flashdreams/runtime/__init__.py | 61 +- flashdreams/flashdreams/runtime/canonical.py | 374 ++++++ flashdreams/flashdreams/runtime/inputs.py | 388 +++++- flashdreams/flashdreams/runtime/interfaces.py | 30 +- flashdreams/flashdreams/runtime/mapping.py | 356 ++++- flashdreams/flashdreams/runtime/types.py | 4 +- flashdreams/tests/test_inference_inputs.py | 1154 ----------------- .../tests/test_inference_runtime_api.py | 152 ++- flashdreams/tests/test_runtime_canonical.py | 456 +++++++ .../tests/test_runtime_input_mapping.py | 558 ++++++++ 12 files changed, 2254 insertions(+), 2190 deletions(-) delete mode 100644 flashdreams/flashdreams/inference/__init__.py delete mode 100644 flashdreams/flashdreams/inference/inputs.py create mode 100644 flashdreams/flashdreams/runtime/canonical.py delete mode 100644 flashdreams/tests/test_inference_inputs.py create mode 100644 flashdreams/tests/test_runtime_canonical.py create mode 100644 flashdreams/tests/test_runtime_input_mapping.py diff --git a/flashdreams/flashdreams/inference/__init__.py b/flashdreams/flashdreams/inference/__init__.py deleted file mode 100644 index 327f9189a..000000000 --- a/flashdreams/flashdreams/inference/__init__.py +++ /dev/null @@ -1,52 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Inference runtime input contracts. - -This package contains the lightweight ``UserInputs`` / ``ModelInputs`` -contracts used by the experimental runtime API. The objects here describe -capabilities and mapping boundaries; model-specific tensor validation stays in -the model adapter or session implementation. -""" - -from flashdreams.inference.inputs import ( - InputMapper, - InputMapperSchema, - InputPhase, - MappingCompatibility, - ModelInputField, - ModelInputSchema, - ModelInputs, - StaticInputMapper, - UserInputCapability, - UserInputEvent, - UserInputSchema, - UserInputTrace, - UserInputWindow, - check_mapping_compatibility, - check_mapping_set_compatibility, - combine_mapper_schemas, - missing_required_inputs, - undeclared_model_inputs, -) - -__all__ = [ - "InputMapper", - "InputMapperSchema", - "InputPhase", - "MappingCompatibility", - "ModelInputField", - "ModelInputSchema", - "ModelInputs", - "StaticInputMapper", - "UserInputCapability", - "UserInputEvent", - "UserInputSchema", - "UserInputTrace", - "UserInputWindow", - "check_mapping_compatibility", - "check_mapping_set_compatibility", - "combine_mapper_schemas", - "missing_required_inputs", - "undeclared_model_inputs", -] diff --git a/flashdreams/flashdreams/inference/inputs.py b/flashdreams/flashdreams/inference/inputs.py deleted file mode 100644 index f97ec51b2..000000000 --- a/flashdreams/flashdreams/inference/inputs.py +++ /dev/null @@ -1,859 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""User-input, model-input, and mapping contracts for inference runtimes.""" - -from __future__ import annotations - -from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, field, replace -from math import isfinite -from types import MappingProxyType -from typing import Any, Literal, Protocol, cast, runtime_checkable - -InputPhase = Literal["initial", "step"] - - -def _normalized_token(value: str, *, field_name: str) -> str: - normalized = value.strip() - if not normalized: - raise ValueError(f"{field_name} must be a non-empty string.") - return normalized - - -def _normalized_optional_token(value: str | None, *, field_name: str) -> str | None: - if value is None: - return None - return _normalized_token(value, field_name=field_name) - - -def _normalized_payload_fields(values: Iterable[str]) -> frozenset[str]: - return frozenset( - _normalized_token(value, field_name="payload field") for value in values - ) - - -def _immutable_metadata_map(values: Mapping[str, Any]) -> Mapping[str, Any]: - """Copy metadata into a read-only map without rewriting its keys. - - Metadata is an open-ended pass-through for adapter hints, so keys are - validated but never normalized: a key must round-trip unchanged. - """ - normalized: dict[str, Any] = {} - for key, value in values.items(): - if not isinstance(key, str): - raise TypeError(f"metadata keys must be strings, got {type(key).__name__}.") - if not key: - raise ValueError("metadata keys must be non-empty strings.") - normalized[key] = value - return MappingProxyType(normalized) - - -def _merged_metadata( - first: Mapping[str, Any], - second: Mapping[str, Any], -) -> Mapping[str, Any]: - """Union two metadata maps, keeping ``first`` on conflicting keys.""" - merged = dict(second) - merged.update(first) - return merged - - -def _validate_phase(value: str) -> InputPhase: - if value not in {"initial", "step"}: - raise ValueError(f"phase must be 'initial' or 'step', got {value!r}.") - return cast(InputPhase, value) - - -def _field_key(field: "ModelInputField") -> tuple[InputPhase, str]: - return field.phase, field.name - - -def _model_field_matches( - produced: "ModelInputField", - required: "ModelInputField", -) -> bool: - if _field_key(produced) != _field_key(required): - return False - payload_kind_ok = ( - produced.payload_kind is None - or required.payload_kind is None - or produced.payload_kind == required.payload_kind - ) - lifecycle_ok = ( - produced.lifecycle is None - or required.lifecycle is None - or produced.lifecycle == required.lifecycle - ) - return payload_kind_ok and lifecycle_ok - - -@dataclass(frozen=True, slots=True) -class UserInputEvent: - """One user-facing input event in session time. - - Static session setup values, such as prompts or initial frames, are still - represented as events. Snapshots are derived views over events rather than - primary inputs. - """ - - timestamp_s: float - kind: str - payload: Mapping[str, Any] = field(default_factory=dict) - session_id: str | None = None - source: str | None = None - - def __post_init__(self) -> None: - timestamp_s = float(self.timestamp_s) - if not isfinite(timestamp_s) or timestamp_s < 0: - raise ValueError("timestamp_s must be finite and >= 0.") - object.__setattr__(self, "timestamp_s", timestamp_s) - object.__setattr__( - self, - "kind", - _normalized_token(self.kind, field_name="event kind"), - ) - object.__setattr__(self, "payload", MappingProxyType(dict(self.payload))) - object.__setattr__( - self, - "session_id", - _normalized_optional_token(self.session_id, field_name="session_id"), - ) - object.__setattr__( - self, - "source", - _normalized_optional_token(self.source, field_name="source"), - ) - - def __hash__(self) -> int: - # ``payload`` is an arbitrary mapping and therefore unhashable, so the - # generated hash would raise. Hash the identity fields instead; equality - # still compares payloads, and equal events still hash equal. - return hash((self.timestamp_s, self.kind, self.session_id, self.source)) - - -@dataclass(frozen=True, slots=True) -class UserInputWindow: - """A deterministic time window over user input events.""" - - start_s: float - end_s: float - events: tuple[UserInputEvent, ...] = () - - def __post_init__(self) -> None: - start_s = float(self.start_s) - end_s = float(self.end_s) - if not isfinite(start_s) or not isfinite(end_s): - raise ValueError("window bounds must be finite.") - if end_s < start_s: - raise ValueError("end_s must be >= start_s.") - object.__setattr__(self, "start_s", start_s) - object.__setattr__(self, "end_s", end_s) - events = _sorted_events(self.events) - for event in events: - if not start_s <= event.timestamp_s <= end_s: - raise ValueError( - f"Event {event.kind!r} at t={event.timestamp_s} is outside " - f"window [{start_s}, {end_s}]." - ) - object.__setattr__(self, "events", events) - - def events_of_kind(self, kind: str) -> tuple[UserInputEvent, ...]: - """Return all events of ``kind`` inside this window.""" - kind = _normalized_token(kind, field_name="event kind") - return tuple(event for event in self.events if event.kind == kind) - - def latest(self, kind: str) -> UserInputEvent | None: - """Return the latest event of ``kind`` inside this window, if any.""" - events = self.events_of_kind(kind) - return events[-1] if events else None - - -@dataclass(frozen=True, slots=True) -class UserInputTrace: - """Ordered replayable user-input event trace.""" - - events: tuple[UserInputEvent, ...] = () - - def __post_init__(self) -> None: - object.__setattr__(self, "events", _sorted_events(self.events)) - - @classmethod - def from_events(cls, events: Sequence[UserInputEvent]) -> "UserInputTrace": - """Build a trace from any event sequence.""" - return cls(events=tuple(events)) - - def window( - self, - *, - start_s: float, - end_s: float, - include_start: bool = True, - include_end: bool = False, - ) -> UserInputWindow: - """Slice the trace into a deterministic event window.""" - start_s = float(start_s) - end_s = float(end_s) - if end_s < start_s: - raise ValueError("end_s must be >= start_s.") - - def _starts_in(event: UserInputEvent) -> bool: - if include_start: - return event.timestamp_s >= start_s - return event.timestamp_s > start_s - - def _ends_in(event: UserInputEvent) -> bool: - if include_end: - return event.timestamp_s <= end_s - return event.timestamp_s < end_s - - return UserInputWindow( - start_s=start_s, - end_s=end_s, - events=tuple( - event for event in self.events if _starts_in(event) and _ends_in(event) - ), - ) - - def events_of_kind(self, kind: str) -> tuple[UserInputEvent, ...]: - """Return all events of ``kind`` in this trace.""" - kind = _normalized_token(kind, field_name="event kind") - return tuple(event for event in self.events if event.kind == kind) - - def latest(self, kind: str) -> UserInputEvent | None: - """Return the latest event of ``kind`` in this trace, if any.""" - events = self.events_of_kind(kind) - return events[-1] if events else None - - -def _sorted_events(events: Sequence[UserInputEvent]) -> tuple[UserInputEvent, ...]: - indexed_events = tuple(enumerate(events)) - for _, event in indexed_events: - if not isinstance(event, UserInputEvent): - raise TypeError(f"Expected UserInputEvent, got {type(event).__name__}.") - return tuple( - event - for _, event in sorted( - indexed_events, - key=lambda item: (item[1].timestamp_s, item[0]), - ) - ) - - -@dataclass(frozen=True, slots=True) -class UserInputCapability: - """Lightweight metadata for a user event a source or mapper can provide.""" - - event_kind: str - payload_kind: str | None = None - payload_fields: frozenset[str] = field(default_factory=frozenset) - metadata: Mapping[str, Any] = field( - default_factory=dict, - compare=False, - hash=False, - ) - description: str = "" - - def __post_init__(self) -> None: - object.__setattr__( - self, - "event_kind", - _normalized_token(self.event_kind, field_name="event kind"), - ) - object.__setattr__( - self, - "payload_kind", - _normalized_optional_token( - self.payload_kind, - field_name="payload_kind", - ), - ) - object.__setattr__( - self, - "payload_fields", - _normalized_payload_fields(self.payload_fields), - ) - object.__setattr__(self, "metadata", _immutable_metadata_map(self.metadata)) - - def is_satisfied_by(self, provider: "UserInputCapability") -> bool: - """Return whether ``provider`` can satisfy this consumed capability.""" - if self.event_kind != provider.event_kind: - return False - payload_kind_ok = ( - self.payload_kind is None - or provider.payload_kind is None - or self.payload_kind == provider.payload_kind - ) - return payload_kind_ok and self.payload_fields.issubset(provider.payload_fields) - - -@dataclass(frozen=True, slots=True) -class UserInputSchema: - """Metadata describing what a source can provide.""" - - capabilities: tuple[UserInputCapability, ...] = () - name: str = "user-input-source" - source_kind: str | None = None - metadata: Mapping[str, Any] = field( - default_factory=dict, - compare=False, - hash=False, - ) - - def __post_init__(self) -> None: - object.__setattr__( - self, - "name", - _normalized_token(self.name, field_name="schema name"), - ) - object.__setattr__( - self, - "source_kind", - _normalized_optional_token(self.source_kind, field_name="source_kind"), - ) - object.__setattr__(self, "metadata", _immutable_metadata_map(self.metadata)) - object.__setattr__(self, "capabilities", tuple(self.capabilities)) - for capability in self.capabilities: - if not isinstance(capability, UserInputCapability): - raise TypeError( - "capabilities must contain UserInputCapability objects." - ) - - def supports(self, capability: UserInputCapability) -> bool: - """Return whether this source can satisfy ``capability``.""" - return any( - capability.is_satisfied_by(provider) for provider in self.capabilities - ) - - def validate_event(self, event: UserInputEvent) -> None: - """Validate an event against this source schema.""" - matching = [ - capability - for capability in self.capabilities - if capability.event_kind == event.kind - ] - if not matching: - raise ValueError( - f"User input source {self.name!r} does not provide event " - f"kind {event.kind!r}." - ) - payload_keys = set(event.payload) - if not any( - capability.payload_fields.issubset(payload_keys) for capability in matching - ): - expected = sorted( - { - field_name - for capability in matching - for field_name in capability.payload_fields - } - ) - raise ValueError( - f"Event {event.kind!r} payload is missing required fields " - f"for source {self.name!r}: {expected}." - ) - - -@dataclass(frozen=True, slots=True) -class ModelInputField: - """Lightweight metadata for one semantic model-facing input field.""" - - name: str - phase: InputPhase - required: bool = True - payload_kind: str | None = None - update_policy: str | None = None - lifecycle: str | None = None - metadata: Mapping[str, Any] = field( - default_factory=dict, - compare=False, - hash=False, - ) - description: str = "" - - def __post_init__(self) -> None: - object.__setattr__( - self, - "name", - _normalized_token(self.name, field_name="model input name"), - ) - object.__setattr__(self, "phase", _validate_phase(self.phase)) - object.__setattr__( - self, - "payload_kind", - _normalized_optional_token( - self.payload_kind, - field_name="payload_kind", - ), - ) - object.__setattr__( - self, - "update_policy", - _normalized_optional_token( - self.update_policy, - field_name="update_policy", - ), - ) - object.__setattr__( - self, - "lifecycle", - _normalized_optional_token( - self.lifecycle, - field_name="lifecycle", - ), - ) - object.__setattr__(self, "metadata", _immutable_metadata_map(self.metadata)) - - -@dataclass(frozen=True, slots=True) -class ModelInputSchema: - """Metadata describing what a model or session expects.""" - - fields: tuple[ModelInputField, ...] = () - name: str = "model-input-consumer" - metadata: Mapping[str, Any] = field( - default_factory=dict, - compare=False, - hash=False, - ) - - def __post_init__(self) -> None: - object.__setattr__( - self, - "name", - _normalized_token(self.name, field_name="schema name"), - ) - object.__setattr__(self, "metadata", _immutable_metadata_map(self.metadata)) - object.__setattr__(self, "fields", tuple(self.fields)) - seen: set[tuple[InputPhase, str]] = set() - for field_def in self.fields: - if not isinstance(field_def, ModelInputField): - raise TypeError("fields must contain ModelInputField objects.") - key = _field_key(field_def) - if key in seen: - phase, name = key - raise ValueError( - f"Duplicate model input field {name!r} for phase {phase!r}." - ) - seen.add(key) - - def required_fields( - self, - *, - phase: InputPhase | None = None, - ) -> tuple[ModelInputField, ...]: - """Return required fields, optionally filtered by phase.""" - return tuple( - field_def - for field_def in self.fields - if field_def.required and (phase is None or field_def.phase == phase) - ) - - def optional_fields( - self, - *, - phase: InputPhase | None = None, - ) -> tuple[ModelInputField, ...]: - """Return optional fields, optionally filtered by phase.""" - return tuple( - field_def - for field_def in self.fields - if not field_def.required and (phase is None or field_def.phase == phase) - ) - - def field_for(self, *, name: str, phase: InputPhase) -> ModelInputField | None: - """Return one field definition, if present.""" - name = _normalized_token(name, field_name="model input name") - phase = _validate_phase(phase) - for field_def in self.fields: - if field_def.name == name and field_def.phase == phase: - return field_def - return None - - -@dataclass(frozen=True, slots=True) -class ModelInputs: - """Semantic model-facing input payloads split by runtime phase.""" - - initial: Mapping[str, Any] = field(default_factory=dict) - step: Mapping[str, Any] = field(default_factory=dict) - - def __post_init__(self) -> None: - object.__setattr__(self, "initial", _immutable_payload_map(self.initial)) - object.__setattr__(self, "step", _immutable_payload_map(self.step)) - - @classmethod - def initial_only(cls, values: Mapping[str, Any]) -> "ModelInputs": - """Create model inputs containing only initial values.""" - return cls(initial=values) - - @classmethod - def step_only(cls, values: Mapping[str, Any]) -> "ModelInputs": - """Create model inputs containing only per-step values.""" - return cls(step=values) - - def for_phase(self, phase: InputPhase) -> Mapping[str, Any]: - """Return the payload mapping for ``phase``.""" - phase = _validate_phase(phase) - return self.initial if phase == "initial" else self.step - - -def _immutable_payload_map(values: Mapping[str, Any]) -> Mapping[str, Any]: - normalized: dict[str, Any] = {} - for key, value in values.items(): - normalized[_normalized_token(str(key), field_name="model input key")] = value - return MappingProxyType(normalized) - - -def missing_required_inputs( - inputs: ModelInputs, - schema: ModelInputSchema, - *, - phase: InputPhase | None = None, -) -> tuple[ModelInputField, ...]: - """Return required model fields absent from ``inputs``.""" - return tuple( - field_def - for field_def in schema.required_fields(phase=phase) - if field_def.name not in inputs.for_phase(field_def.phase) - ) - - -def undeclared_model_inputs( - inputs: ModelInputs, - mapper_schema: "InputMapperSchema", -) -> tuple[tuple[InputPhase, str], ...]: - """Return payload keys a mapper produced but did not declare in its schema. - - Mapper schemas are hand-written, so they can drift from what - ``build_initial_inputs`` / ``build_step_inputs`` actually return. Mapper - tests can use this to keep the declared compatibility surface honest. - """ - declared = { - (field_def.phase, field_def.name) for field_def in mapper_schema.produces - } - phases: tuple[InputPhase, ...] = ("initial", "step") - return tuple( - (phase, key) - for phase in phases - for key in inputs.for_phase(phase) - if (phase, key) not in declared - ) - - -@dataclass(frozen=True, slots=True) -class InputMapperSchema: - """Metadata for a mapper that converts user events into model inputs.""" - - consumes: tuple[UserInputCapability, ...] = () - produces: tuple[ModelInputField, ...] = () - name: str = "input-mapper" - metadata: Mapping[str, Any] = field( - default_factory=dict, - compare=False, - hash=False, - ) - - def __post_init__(self) -> None: - object.__setattr__( - self, - "name", - _normalized_token(self.name, field_name="mapper name"), - ) - object.__setattr__(self, "metadata", _immutable_metadata_map(self.metadata)) - object.__setattr__(self, "consumes", tuple(self.consumes)) - object.__setattr__(self, "produces", tuple(self.produces)) - for capability in self.consumes: - if not isinstance(capability, UserInputCapability): - raise TypeError("consumes must contain UserInputCapability objects.") - for field_def in self.produces: - if not isinstance(field_def, ModelInputField): - raise TypeError("produces must contain ModelInputField objects.") - - def can_produce(self, model_field: ModelInputField) -> bool: - """Return whether this mapper can produce ``model_field``.""" - return any( - _model_field_matches(produced, model_field) for produced in self.produces - ) - - -@runtime_checkable -class InputMapper(Protocol): - """Contract for user-event to model-input conversion.""" - - @property - def schema(self) -> InputMapperSchema: - """Return mapper metadata used for compatibility checks.""" - ... - - def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: - """Build session-start model inputs from a user-input trace.""" - ... - - def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: - """Build per-step model inputs from one user-input window.""" - ... - - -@dataclass(frozen=True, slots=True) -class StaticInputMapper: - """Mapper for fixed or already prepared model inputs. - - This is the no-op mapping path for runs that do not need live user - controls, such as prompt-only CLI runs or model-input replay. It consumes - no user events and returns configured model-facing payloads. - """ - - schema: InputMapperSchema - inputs: ModelInputs = field(default_factory=ModelInputs) - - @classmethod - def from_inputs( - cls, - *, - inputs: ModelInputs, - name: str = "static-input-mapper", - ) -> "StaticInputMapper": - """Create a static mapper whose produced fields come from ``inputs``.""" - produces = tuple( - ModelInputField(name=field_name, phase=phase, required=False) - for phase, values in ( - ("initial", inputs.initial), - ("step", inputs.step), - ) - for field_name in values - ) - return cls( - schema=InputMapperSchema(name=name, produces=produces), - inputs=inputs, - ) - - def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: - """Return configured initial model inputs.""" - del trace - return ModelInputs(initial=self.inputs.initial) - - def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: - """Return configured per-step model inputs.""" - del window - return ModelInputs(step=self.inputs.step) - - -@dataclass(frozen=True, slots=True) -class MappingCompatibility: - """Compatibility report for one source, model schema, and mapping. - - ``mapper_schema`` is the full requested mapping surface. Mappers whose - consumed capabilities the source cannot provide are reported separately in - ``unavailable_mapper_schemas`` and are excluded from the satisfied/available - model-field reports, so those lists only name fields that can really be - produced by this source. - """ - - source_schema: UserInputSchema - model_schema: ModelInputSchema - mapper_schema: InputMapperSchema - missing_source_capabilities: tuple[UserInputCapability, ...] - missing_required_model_fields: tuple[ModelInputField, ...] - satisfied_required_model_fields: tuple[ModelInputField, ...] - available_optional_model_fields: tuple[ModelInputField, ...] - unavailable_mapper_schemas: tuple[InputMapperSchema, ...] = () - - @property - def can_drive(self) -> bool: - """Return whether this source can drive this model through the mapping. - - Mappers that the source cannot feed do not block the run unless they - were the only way to produce a required model input. - """ - return ( - not self.missing_source_capabilities - and not self.missing_required_model_fields - ) - - @property - def unavailable_mapper_names(self) -> tuple[str, ...]: - """Return names of mappers dropped because the source cannot feed them.""" - return tuple(schema.name for schema in self.unavailable_mapper_schemas) - - def raise_if_incompatible(self) -> None: - """Raise a compact error when this mapping cannot drive the model.""" - if self.can_drive: - return - problems: list[str] = [] - if self.missing_source_capabilities: - missing = ", ".join( - capability.event_kind for capability in self.missing_source_capabilities - ) - problems.append(f"missing source capabilities: {missing}") - if self.missing_required_model_fields: - missing = ", ".join( - f"{field_def.phase}:{field_def.name}" - for field_def in self.missing_required_model_fields - ) - problems.append(f"missing required model inputs: {missing}") - if self.unavailable_mapper_schemas: - problems.append( - "unavailable mappers: " + ", ".join(self.unavailable_mapper_names) - ) - raise ValueError( - f"Input mapper {self.mapper_schema.name!r} cannot drive model " - f"{self.model_schema.name!r} from source {self.source_schema.name!r}: " - + "; ".join(problems) - ) - - -def _mapper_is_feedable( - source_schema: UserInputSchema, - mapper_schema: InputMapperSchema, -) -> bool: - return all( - source_schema.supports(capability) for capability in mapper_schema.consumes - ) - - -def _build_compatibility( - *, - source_schema: UserInputSchema, - model_schema: ModelInputSchema, - mapper_schemas: Sequence[InputMapperSchema], - reported_schema: InputMapperSchema, -) -> MappingCompatibility: - feedable: list[InputMapperSchema] = [] - unavailable: list[InputMapperSchema] = [] - for mapper_schema in mapper_schemas: - if _mapper_is_feedable(source_schema, mapper_schema): - feedable.append(mapper_schema) - else: - unavailable.append(mapper_schema) - - usable = combine_mapper_schemas(feedable, name=reported_schema.name) - required_fields = model_schema.required_fields() - missing_required_model_fields = tuple( - field_def for field_def in required_fields if not usable.can_produce(field_def) - ) - satisfied_required_model_fields = tuple( - field_def for field_def in required_fields if usable.can_produce(field_def) - ) - available_optional_model_fields = tuple( - field_def - for field_def in model_schema.optional_fields() - if usable.can_produce(field_def) - ) - - # Only capabilities that block a required model input make the mapping - # unusable. A dropped mapper that fed nothing but optional fields degrades - # instead of vetoing the run. - missing_source_capabilities: list[UserInputCapability] = [] - seen_missing: set[UserInputCapability] = set() - for mapper_schema in unavailable: - if not any( - mapper_schema.can_produce(field_def) - for field_def in missing_required_model_fields - ): - continue - for capability in mapper_schema.consumes: - if source_schema.supports(capability) or capability in seen_missing: - continue - seen_missing.add(capability) - missing_source_capabilities.append(capability) - - return MappingCompatibility( - source_schema=source_schema, - model_schema=model_schema, - mapper_schema=reported_schema, - missing_source_capabilities=tuple(missing_source_capabilities), - missing_required_model_fields=missing_required_model_fields, - satisfied_required_model_fields=satisfied_required_model_fields, - available_optional_model_fields=available_optional_model_fields, - unavailable_mapper_schemas=tuple(unavailable), - ) - - -def check_mapping_compatibility( - *, - source_schema: UserInputSchema, - model_schema: ModelInputSchema, - mapper_schema: InputMapperSchema, -) -> MappingCompatibility: - """Check whether a user-input source can drive a model through a mapper.""" - if not isinstance(mapper_schema, InputMapperSchema): - raise TypeError("mapper_schema must be an InputMapperSchema object.") - return _build_compatibility( - source_schema=source_schema, - model_schema=model_schema, - mapper_schemas=(mapper_schema,), - reported_schema=mapper_schema, - ) - - -def combine_mapper_schemas( - mapper_schemas: Sequence[InputMapperSchema], - *, - name: str = "input-mapper-set", -) -> InputMapperSchema: - """Combine independently declared mappers into one compatibility surface. - - Duplicate entries are collapsed. Because ``metadata`` is excluded from - equality, the metadata of collapsed duplicates is merged rather than - dropped, with the first declaration winning on conflicting keys. - """ - consumes: list[UserInputCapability] = [] - produces: list[ModelInputField] = [] - consumes_index: dict[UserInputCapability, int] = {} - produces_index: dict[ModelInputField, int] = {} - - for mapper_schema in mapper_schemas: - if not isinstance(mapper_schema, InputMapperSchema): - raise TypeError("mapper_schemas must contain InputMapperSchema objects.") - for capability in mapper_schema.consumes: - index = consumes_index.get(capability) - if index is None: - consumes_index[capability] = len(consumes) - consumes.append(capability) - elif capability.metadata: - consumes[index] = replace( - consumes[index], - metadata=_merged_metadata( - consumes[index].metadata, capability.metadata - ), - ) - for field_def in mapper_schema.produces: - index = produces_index.get(field_def) - if index is None: - produces_index[field_def] = len(produces) - produces.append(field_def) - elif field_def.metadata: - produces[index] = replace( - produces[index], - metadata=_merged_metadata( - produces[index].metadata, field_def.metadata - ), - ) - - return InputMapperSchema( - name=name, - consumes=tuple(consumes), - produces=tuple(produces), - ) - - -def check_mapping_set_compatibility( - *, - source_schema: UserInputSchema, - model_schema: ModelInputSchema, - mapper_schemas: Sequence[InputMapperSchema], - name: str = "input-mapper-set", -) -> MappingCompatibility: - """Check compatibility for a composed set of input mappers. - - Each mapper keeps its own ``consumes``/``produces`` link, so a mapper the - source cannot feed only costs the model inputs that mapper produced. - """ - mapper_schemas = tuple(mapper_schemas) - return _build_compatibility( - source_schema=source_schema, - model_schema=model_schema, - mapper_schemas=mapper_schemas, - reported_schema=combine_mapper_schemas(mapper_schemas, name=name), - ) diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index 03e6202b0..f1f6d0f48 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -7,22 +7,50 @@ intentionally additive while integrations migrate onto it. """ +from flashdreams.runtime.canonical import ( + CONDITIONING_FRAME, + CONDITIONING_PROMPT, + DRIVER_COMMAND, + DeviceConverter, + DeviceConverterSchema, + InputCanonicalizer, + KeyboardToDriverCommand, + LatestEventToModality, +) from flashdreams.runtime.config import ExecutionBackend, InferenceConfig, Precision from flashdreams.runtime.inputs import ( + INPUT_PHASES, + SESSION_START_ONLY, + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + InferenceInput, + InferenceInputSchema, InputField, - ModelInputs, - ModelInputSchema, + InputPhase, TimeWindow, + UserInputCapability, UserInputEvent, UserInputs, UserInputSchema, + validate_phase, ) from flashdreams.runtime.interfaces import ( InferenceRuntime, InferenceSession, ModelAdapter, ) -from flashdreams.runtime.mapping import IdentityInputMapping, InputMapping +from flashdreams.runtime.mapping import ( + DeclaresMappingSchema, + IdentityInputMapping, + InputMapping, + InputMappingSchema, + MappingCompatibility, + check_mapping_compatibility, + check_mapping_set_compatibility, + combine_mapping_schemas, + undeclared_inference_input, +) from flashdreams.runtime.metrics import ( InMemoryMetricsRecorder, MetricsRecorder, @@ -33,28 +61,51 @@ from flashdreams.runtime.types import StepRequest, StepResult __all__ = [ + "CanonicalInputs", + "CanonicalInputSchema", + "CanonicalModality", + "check_mapping_compatibility", + "check_mapping_set_compatibility", + "combine_mapping_schemas", + "CONDITIONING_FRAME", + "CONDITIONING_PROMPT", + "DeclaresMappingSchema", + "DeviceConverter", + "DeviceConverterSchema", + "DRIVER_COMMAND", "ExecutionBackend", "IdentityInputMapping", "InferenceConfig", + "InferenceInput", + "InferenceInputSchema", "InferenceRuntime", "InferenceSession", "InMemoryMetricsRecorder", + "INPUT_PHASES", + "InputCanonicalizer", "InputField", "InputMapping", + "InputMappingSchema", + "InputPhase", + "KeyboardToDriverCommand", + "LatestEventToModality", + "MappingCompatibility", "MetricsRecorder", "ModelAdapter", - "ModelInputs", - "ModelInputSchema", "NullMetricsRecorder", "NullOutputTarget", "OutputArtifact", "OutputTarget", "Precision", "RuntimeMetricSample", + "SESSION_START_ONLY", "StepRequest", "StepResult", "TimeWindow", + "undeclared_inference_input", + "UserInputCapability", "UserInputEvent", "UserInputs", "UserInputSchema", + "validate_phase", ] diff --git a/flashdreams/flashdreams/runtime/canonical.py b/flashdreams/flashdreams/runtime/canonical.py new file mode 100644 index 000000000..5bd21fc32 --- /dev/null +++ b/flashdreams/flashdreams/runtime/canonical.py @@ -0,0 +1,374 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Raw device input to canonical modality conversion. + +This is the ``raw input -> canonicalized input`` leg. Applications consume +:class:`~flashdreams.runtime.inputs.CanonicalInputs`; they never read raw device +events. Adding a keyboard, gamepad, or force-feedback wheel is therefore a +:meth:`InputCanonicalizer.register` call that touches no application, mapping, +or model code. + +Converters are stateful, because HID input is edge-triggered while per-step +conditioning is level-triggered: a key held across a step emits no events yet +still means full throttle. Feed windows in session order and call +:meth:`InputCanonicalizer.reset` at a rollout boundary; replaying the same +window sequence then reproduces the same canonical inputs. + +Global-conditioning converters behave differently on purpose. They emit only +when a value actually changed in the window, so a downstream non-empty global +slot means "update this", not "re-apply it every step". +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.inputs import ( + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + TimeWindow, + UserInputCapability, + UserInputs, + UserInputSchema, +) +from flashdreams.serving.realtime.input import ( + DRIVING_SUPPORTED_KEYS, + KeyboardState, + normalize_key, +) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class DeviceConverterSchema: + """Metadata for one device-to-canonical-modality converter.""" + + name: str + produces: CanonicalModality + consumes: tuple[UserInputCapability, ...] = () + device_kind: str | None = None + priority: int = 0 + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("DeviceConverterSchema.name must be non-empty.") + if not isinstance(self.produces, CanonicalModality): + raise TypeError("produces must be a CanonicalModality object.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@runtime_checkable +class DeviceConverter(Protocol): + """Contract for turning one device's raw events into a canonical modality.""" + + @property + def schema(self) -> DeviceConverterSchema: + """Return converter metadata used for source selection.""" + ... + + def reset(self) -> None: + """Drop accumulated device state at a session or rollout boundary.""" + ... + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + """Return the modality value for ``window``, or ``None`` if inactive. + + ``user_inputs`` is already filtered to ``window``. Returning ``None`` + lets a present-but-idle device yield to a lower-priority one, and lets a + global-conditioning converter stay silent when nothing changed. + """ + ... + + +DRIVER_COMMAND = CanonicalModality( + name="driver_command", + phase="step", + payload_fields=frozenset({"throttle", "brake", "steer", "stop", "reverse"}), + description=( + "Normalized driving intent. throttle/brake are in [0, 1], steer is in " + "[-1, 1] with positive meaning left." + ), +) + +CONDITIONING_PROMPT = CanonicalModality( + name="conditioning_prompt", + phase="global", + payload_fields=frozenset({"prompt"}), + description="Global conditioning prompt; may change mid-rollout.", +) + +CONDITIONING_FRAME = CanonicalModality( + name="conditioning_frame", + phase="global", + payload_fields=frozenset({"image"}), + description="Global conditioning frame; may change mid-rollout.", +) + + +class KeyboardToDriverCommand: + """Convert keyboard edges into :data:`DRIVER_COMMAND` level state. + + Mirrors the mapping the Omnidreams interactive-drive keyboard backend + already uses, so a keyboard reaches a model through the shared layer with + the same semantics it has today. + """ + + def __init__( + self, + *, + name: str = "keyboard-to-driver-command", + supported_keys: frozenset[str] = DRIVING_SUPPORTED_KEYS, + priority: int = 0, + ) -> None: + self._supported_keys = supported_keys + self._state = KeyboardState(supported_keys=supported_keys) + self._schema = DeviceConverterSchema( + name=name, + produces=DRIVER_COMMAND, + device_kind="keyboard", + priority=priority, + consumes=( + UserInputCapability( + event_type="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_type="key_up", + payload_fields=frozenset({"key"}), + ), + ), + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + self._state = KeyboardState(supported_keys=self._supported_keys) + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + del window + for event in user_inputs.events: + if event.event_type not in {"key_down", "key_up"}: + continue + key = event.payload.get("key") + if not isinstance(key, str): + continue + self._state.apply_event( + event="keydown" if event.event_type == "key_down" else "keyup", + key=key, + ) + + pressed = {normalize_key(key) for key in self._state.snapshot()} + steer = 0.0 + if {"a", "left"} & pressed: + steer += 1.0 + if {"d", "right"} & pressed: + steer -= 1.0 + return DRIVER_COMMAND.value( + { + "throttle": 1.0 if {"w", "up"} & pressed else 0.0, + "brake": 1.0 if {"s", "down"} & pressed else 0.0, + "steer": steer, + "stop": "space" in pressed, + "reverse": False, + } + ) + + +class LatestEventToModality: + """Emit a global-conditioning modality when its source event fires. + + Global conditioning is transient by design: this returns ``None`` in windows + where nothing changed, so downstream code only sees an update request when + the user actually changed something. + """ + + def __init__( + self, + *, + modality: CanonicalModality, + event_type: str, + payload_fields: frozenset[str] | None = None, + name: str | None = None, + device_kind: str | None = None, + priority: int = 0, + ) -> None: + if modality.phase != "global": + raise ValueError( + "LatestEventToModality is for global conditioning; " + f"{modality.name!r} declares phase {modality.phase!r}." + ) + self._modality = modality + self._event_type = event_type + fields = modality.payload_fields if payload_fields is None else payload_fields + self._schema = DeviceConverterSchema( + name=name or f"{event_type}-to-{modality.name}", + produces=modality, + device_kind=device_kind, + priority=priority, + consumes=( + UserInputCapability(event_type=event_type, payload_fields=fields), + ), + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + return None + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + del window + latest = None + for event in user_inputs.events: + if event.event_type == self._event_type: + latest = event + if latest is None: + return None + return self._modality.value( + { + name: latest.payload[name] + for name in self._modality.payload_fields + if name in latest.payload + } + ) + + +class InputCanonicalizer: + """Registry of device converters plus the raw-to-canonical rewrite. + + Registration is the whole extension point: a new device is a converter + registered against an existing modality, and a new modality is a converter + registered with a new :class:`CanonicalModality`. + """ + + def __init__(self, converters: Iterable[DeviceConverter] = ()) -> None: + self._converters: list[DeviceConverter] = [] + for converter in converters: + self.register(converter) + + def register(self, converter: DeviceConverter) -> None: + """Register one device converter.""" + if not isinstance(converter, DeviceConverter): + raise TypeError("converter must implement the DeviceConverter protocol.") + name = converter.schema.name + if any(existing.schema.name == name for existing in self._converters): + raise ValueError( + f"A device converter named {name!r} is already registered." + ) + self._converters.append(converter) + + @property + def converters(self) -> tuple[DeviceConverter, ...]: + """Return every registered converter.""" + return tuple(self._converters) + + def reset(self) -> None: + """Reset every registered converter's device state.""" + for converter in self._converters: + converter.reset() + + def converters_for( + self, + source_schema: UserInputSchema, + ) -> tuple[DeviceConverter, ...]: + """Return converters this source can feed, highest priority first.""" + feedable = [ + converter + for converter in self._converters + if all( + source_schema.supports(capability) + for capability in converter.schema.consumes + ) + ] + # Sort is stable, so equal-priority converters keep registration order. + return tuple(sorted(feedable, key=lambda each: -each.schema.priority)) + + def unavailable_converters( + self, + source_schema: UserInputSchema, + ) -> tuple[DeviceConverter, ...]: + """Return converters this source cannot feed, for diagnostics.""" + feedable = {id(converter) for converter in self.converters_for(source_schema)} + return tuple( + converter for converter in self._converters if id(converter) not in feedable + ) + + def canonical_schema( + self, + source_schema: UserInputSchema, + ) -> CanonicalInputSchema: + """Return the canonical modalities this raw source can supply. + + This is the boundary an application declares against. A mapping that + consumes ``driver_command`` then matches a keyboard source, a wheel + source, or any device registered later. + """ + modalities: list[CanonicalModality] = [] + for converter in self.converters_for(source_schema): + modality = converter.schema.produces + if modality not in modalities: + modalities.append(modality) + return CanonicalInputSchema( + modalities=tuple(modalities), + description=source_schema.description, + ) + + def canonicalize( + self, + user_inputs: UserInputs, + *, + window: TimeWindow, + source_schema: UserInputSchema, + ) -> CanonicalInputs: + """Convert one raw window into canonical inputs. + + Every feedable converter sees the window so its device state stays + current even while another device has precedence; that way unplugging + the higher-priority device does not resume from stale state. Among + converters producing the same modality, the highest-priority one that + returned a value wins. + """ + windowed = user_inputs.window(window) + by_phase: dict[str, dict[str, Any]] = {"global": {}, "step": {}} + sources: dict[str, str] = {} + for converter in self.converters_for(source_schema): + value = converter.convert(windowed, window) + modality = converter.schema.produces + slot = by_phase[modality.phase] + if value is not None and modality.name not in slot: + slot[modality.name] = value + if converter.schema.device_kind is not None: + sources[modality.name] = converter.schema.device_kind + + metadata: dict[str, Any] = {} + if sources: + metadata["canonical_sources"] = freeze_mapping(sources) + return CanonicalInputs( + global_conditioning=by_phase["global"], + per_step=by_phase["step"], + metadata=metadata, + ) diff --git a/flashdreams/flashdreams/runtime/inputs.py b/flashdreams/flashdreams/runtime/inputs.py index e14b35722..cacab635f 100644 --- a/flashdreams/flashdreams/runtime/inputs.py +++ b/flashdreams/flashdreams/runtime/inputs.py @@ -8,10 +8,29 @@ import math from collections.abc import Iterable, Mapping from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal, cast from flashdreams.runtime._utils import freeze_mapping +InputPhase = Literal["global", "step"] + +INPUT_PHASES: tuple[InputPhase, ...] = ("global", "step") + +SESSION_START_ONLY = "session_start" +"""``InputField.update_policy`` value meaning "supply at session start only". + +``update_policy`` is otherwise an open, adapter-owned vocabulary. This is the +one reserved token, because the runtime needs to distinguish a conditioning +value that can be swapped mid-rollout from one that cannot. +""" + + +def validate_phase(value: str) -> InputPhase: + """Return ``value`` as a validated :data:`InputPhase`.""" + if value not in INPUT_PHASES: + raise ValueError(f"phase must be 'global' or 'step', got {value!r}.") + return cast(InputPhase, value) + @dataclass(frozen=True, kw_only=True, slots=True) class TimeWindow: @@ -35,16 +54,70 @@ def contains(self, timestamp_s: float) -> bool: @dataclass(frozen=True, kw_only=True, slots=True) class InputField: - """Lightweight schema field for user snapshots or model inputs.""" + """Lightweight schema field for user snapshots or model inputs. + + ``update_policy`` and ``lifecycle`` are plain query metadata. They let a + model advertise facts such as "prompt updates land at step boundaries" or + "this value is consumed at cache init" without making this layer + responsible for implementing or deeply validating that behavior. + """ name: str required: bool = True semantic_type: str | None = None + update_policy: str | None = None + lifecycle: str | None = None + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) description: str = "" def __post_init__(self) -> None: if not self.name.strip(): raise ValueError("InputField.name must be non-empty.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class UserInputCapability: + """One user event a source or mapping can provide, at payload granularity. + + ``UserInputSchema.event_types`` declares only that an event type exists. A + capability additionally pins the payload fields carried by that event, so a + mapping can state that it needs ``key_down`` events that actually carry a + ``key``. + """ + + event_type: str + semantic_type: str | None = None + payload_fields: frozenset[str] = field(default_factory=frozenset) + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + description: str = "" + + def __post_init__(self) -> None: + if not self.event_type.strip(): + raise ValueError("UserInputCapability.event_type must be non-empty.") + for payload_field in self.payload_fields: + if not payload_field.strip(): + raise ValueError("payload field names must be non-empty.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + def is_satisfied_by(self, provider: "UserInputCapability") -> bool: + """Return whether ``provider`` can satisfy this consumed capability.""" + if self.event_type != provider.event_type: + return False + semantic_ok = ( + self.semantic_type is None + or provider.semantic_type is None + or self.semantic_type == provider.semantic_type + ) + return semantic_ok and self.payload_fields.issubset(provider.payload_fields) @dataclass(frozen=True, kw_only=True, slots=True) @@ -53,6 +126,7 @@ class UserInputSchema: event_types: frozenset[str] = field(default_factory=frozenset) snapshot_fields: tuple[InputField, ...] = () + capabilities: tuple[UserInputCapability, ...] = () description: str = "" def supports_event_types(self, event_types: Iterable[str]) -> bool: @@ -60,7 +134,63 @@ def supports_event_types(self, event_types: Iterable[str]) -> bool: requested = frozenset(event_types) if not requested: return True - return requested.issubset(self.event_types) + return requested.issubset(self.declared_event_types()) + + def declared_event_types(self) -> frozenset[str]: + """Return event types from ``event_types`` and from ``capabilities``.""" + return self.event_types | frozenset( + capability.event_type for capability in self.capabilities + ) + + def declared_capabilities(self) -> tuple[UserInputCapability, ...]: + """Return capabilities, widened with bare ``event_types`` entries. + + A plain ``event_types`` entry carries no payload promise, so it is + modeled as a capability with no payload fields. Coarse schemas written + before capabilities existed therefore still satisfy any consumer that + does not require specific payload fields. + """ + declared = list(self.capabilities) + covered = {capability.event_type for capability in declared} + declared.extend( + UserInputCapability(event_type=event_type) + for event_type in sorted(self.event_types - covered) + ) + return tuple(declared) + + def supports(self, capability: UserInputCapability) -> bool: + """Return whether this source can satisfy ``capability``.""" + return any( + capability.is_satisfied_by(provider) + for provider in self.declared_capabilities() + ) + + def validate_event(self, event: "UserInputEvent") -> None: + """Validate one event against the event types this source declares.""" + matching = [ + capability + for capability in self.declared_capabilities() + if capability.event_type == event.event_type + ] + if not matching: + raise ValueError( + f"User input source does not provide event type {event.event_type!r}." + ) + payload_keys = set(event.payload) + if not any( + capability.payload_fields.issubset(payload_keys) for capability in matching + ): + expected = sorted( + { + payload_field + for capability in matching + for payload_field in capability.payload_fields + } + ) + raise ValueError( + f"Event {event.event_type!r} payload is missing required " + f"fields: {expected}." + ) def missing_snapshot(self, inputs: "UserInputs") -> tuple[str, ...]: """Return required snapshot fields absent from ``inputs``.""" @@ -74,10 +204,10 @@ def require_snapshot(self, inputs: "UserInputs") -> None: @dataclass(frozen=True, kw_only=True, slots=True) -class ModelInputSchema: +class InferenceInputSchema: """Minimal metadata for model-facing initial and per-step inputs.""" - initial_fields: tuple[InputField, ...] = () + global_fields: tuple[InputField, ...] = () """Model inputs required before starting the initial generation/session.""" step_fields: tuple[InputField, ...] = () @@ -85,21 +215,81 @@ class ModelInputSchema: description: str = "" - def missing_initial(self, inputs: "ModelInputs") -> tuple[str, ...]: + def fields_for(self, phase: InputPhase) -> tuple[InputField, ...]: + """Return every declared field for ``phase``.""" + return ( + self.global_fields + if validate_phase(phase) == "global" + else self.step_fields + ) + + def required_fields( + self, + phase: InputPhase | None = None, + ) -> tuple[tuple[InputPhase, InputField], ...]: + """Return required fields as ``(phase, field)``, optionally filtered.""" + return self._select(phase, required=True) + + def optional_fields( + self, + phase: InputPhase | None = None, + ) -> tuple[tuple[InputPhase, InputField], ...]: + """Return optional fields as ``(phase, field)``, optionally filtered.""" + return self._select(phase, required=False) + + def field_for(self, *, name: str, phase: InputPhase) -> InputField | None: + """Return one declared field, if present.""" + for input_field in self.fields_for(phase): + if input_field.name == name: + return input_field + return None + + def _select( + self, + phase: InputPhase | None, + *, + required: bool, + ) -> tuple[tuple[InputPhase, InputField], ...]: + phases = INPUT_PHASES if phase is None else (validate_phase(phase),) + return tuple( + (each_phase, input_field) + for each_phase in phases + for input_field in self.fields_for(each_phase) + if input_field.required is required + ) + + def unsupported_global_updates(self, inputs: "InferenceInput") -> tuple[str, ...]: + """Return requested conditioning updates this model cannot apply. + + A field whose ``update_policy`` is :data:`SESSION_START_ONLY` can be + supplied when the session starts but not changed mid-rollout. Any other + policy, including ``None``, is treated as permissive here; the adapter + still owns whether the swap actually succeeds. + """ + return tuple( + name + for name in inputs.global_conditioning + if (declared := self.field_for(name=name, phase="global")) is not None + and declared.update_policy == SESSION_START_ONLY + ) + + def missing_global(self, inputs: "InferenceInput") -> tuple[str, ...]: """Return required initial fields absent from ``inputs``.""" - return _missing_required(self.initial_fields, inputs.initial) + return _missing_required(self.global_fields, inputs.global_conditioning) - def missing_step(self, inputs: "ModelInputs") -> tuple[str, ...]: + def missing_step(self, inputs: "InferenceInput") -> 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: + def require_global(self, inputs: "InferenceInput") -> None: """Raise if required initial fields are absent.""" - missing = self.missing_initial(inputs) + missing = self.missing_global(inputs) if missing: - raise ValueError(f"Missing required initial model input(s): {missing}") + raise ValueError( + f"Missing required global conditioning input(s): {missing}" + ) - def require_step(self, inputs: "ModelInputs") -> None: + def require_step(self, inputs: "InferenceInput") -> None: """Raise if required per-step fields are absent.""" missing = self.missing_step(inputs) if missing: @@ -171,23 +361,183 @@ def window(self, time_window: TimeWindow) -> "UserInputs": @dataclass(frozen=True, kw_only=True, slots=True) -class ModelInputs: - """Model-facing payloads split by initial and per-step use.""" +class CanonicalModality: + """A device-independent user input an application consumes. + + This is the middle layer of ``raw input -> canonicalized input -> encoded + inference input``. Applications and benchmarks declare and consume + modalities; they never read raw device events, so adding a new device is a + converter registration rather than an application change. + + ``phase`` says which conditioning slot the modality feeds: ``"global"`` for + one-shot rollout conditioning such as a prompt or conditioning frame, and + ``"step"`` for per-step conditioning such as steering. + """ + + name: str + phase: InputPhase = "step" + payload_fields: frozenset[str] = field(default_factory=frozenset) + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + description: str = "" + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("CanonicalModality.name must be non-empty.") + object.__setattr__(self, "phase", validate_phase(self.phase)) + for payload_field in self.payload_fields: + if not payload_field.strip(): + raise ValueError("payload field names must be non-empty.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + def is_satisfied_by(self, provider: "CanonicalModality") -> bool: + """Return whether ``provider`` can satisfy this consumed modality.""" + return ( + self.name == provider.name + and self.phase == provider.phase + and self.payload_fields.issubset(provider.payload_fields) + ) + + def value(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Return ``payload`` frozen, checking it covers this modality.""" + missing = sorted(self.payload_fields - set(payload)) + if missing: + raise ValueError( + f"Canonical modality {self.name!r} requires payload fields " + f"{missing}, which the converter did not produce." + ) + return freeze_mapping(payload) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class CanonicalInputSchema: + """Canonical modalities an application can be fed by a given source.""" + + modalities: tuple[CanonicalModality, ...] = () + description: str = "" + + def supports(self, modality: CanonicalModality) -> bool: + """Return whether this source can supply ``modality``.""" + return any(modality.is_satisfied_by(provided) for provided in self.modalities) + + def modalities_for(self, phase: InputPhase) -> tuple[CanonicalModality, ...]: + """Return supplied modalities feeding ``phase``.""" + validate_phase(phase) + return tuple( + modality for modality in self.modalities if modality.phase == phase + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class CanonicalInputs: + """Canonicalized user input for one step, split by conditioning slot. + + ``per_step`` is level-triggered and normally present every step: a key held + down emits no events but still means full throttle. ``global_conditioning`` + is transient and populated only when a value actually changed in this + window, which is what makes a non-empty global slot downstream mean "update + this" rather than "re-apply every step". + """ + + __hash__ = None + + global_conditioning: Mapping[str, Any] = field(default_factory=dict) + per_step: Mapping[str, Any] = field(default_factory=dict) + metadata: 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", freeze_mapping(self.per_step)) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + @property + def has_global_change(self) -> bool: + """Return whether a global conditioning value changed in this window.""" + return bool(self.global_conditioning) + + def for_phase(self, phase: InputPhase) -> Mapping[str, Any]: + """Return the canonical payload for ``phase``.""" + return ( + self.global_conditioning + if validate_phase(phase) == "global" + else self.per_step + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InferenceInput: + """Encoded inputs for one :class:`InferenceSession` call. + + Two conditioning slots: + + - ``global_conditioning``: values that condition the whole rollout, such as + the conditioning frame or prompt. Normally supplied when the session + starts. + - ``step``: values needed to generate the next chunk or frame. + + A non-empty ``global_conditioning`` on a mid-rollout input is an *update + request*, not a reset. The session should apply it when the model supports + that; resetting rollout state is a separate, explicit + :meth:`InferenceSession.reset` call. Whether a given value can be updated + mid-rollout is declared per field by ``InputField.update_policy``; see + :meth:`InferenceInputSchema.unsupported_global_updates`. + """ __hash__ = None - initial: Mapping[str, Any] = field(default_factory=dict) + global_conditioning: 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, "global_conditioning", freeze_mapping(self.global_conditioning) + ) 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) + @property + def requests_global_update(self) -> bool: + """Return whether this input asks the session to update conditioning.""" + return bool(self.global_conditioning) + + def with_step(self, step: Mapping[str, Any]) -> "InferenceInput": + """Return a copy with replaced per-step payload. + + The global slot is carried through unchanged, so a mid-rollout input + built this way keeps whatever update request it already had. Use + :meth:`without_global_update` for the common steady-state case. + """ + return InferenceInput( + global_conditioning=self.global_conditioning, + step=step, + metadata=self.metadata, + ) + + def with_global_update( + self, global_conditioning: Mapping[str, Any] + ) -> "InferenceInput": + """Return a copy requesting a mid-rollout conditioning update.""" + return InferenceInput( + global_conditioning=global_conditioning, + step=self.step, + metadata=self.metadata, + ) + + def without_global_update(self) -> "InferenceInput": + """Return a copy that requests no conditioning update.""" + return InferenceInput(step=self.step, metadata=self.metadata) + + def for_phase(self, phase: InputPhase) -> Mapping[str, Any]: + """Return the payload mapping for ``phase``.""" + return ( + self.global_conditioning if validate_phase(phase) == "global" else self.step + ) def _missing_required( diff --git a/flashdreams/flashdreams/runtime/interfaces.py b/flashdreams/flashdreams/runtime/interfaces.py index 9b6a064fd..852a77f1c 100644 --- a/flashdreams/flashdreams/runtime/interfaces.py +++ b/flashdreams/flashdreams/runtime/interfaces.py @@ -9,9 +9,9 @@ from flashdreams.runtime.config import InferenceConfig from flashdreams.runtime.inputs import ( - ModelInputs, - ModelInputSchema, - UserInputSchema, + CanonicalInputSchema, + InferenceInput, + InferenceInputSchema, ) from flashdreams.runtime.mapping import InputMapping from flashdreams.runtime.types import StepRequest, StepResult @@ -25,11 +25,11 @@ def next_step_request(self) -> StepRequest | None: """Describe the next step's inputs, or return ``None`` when complete.""" ... - def step(self, inputs: ModelInputs) -> StepResult: + def step(self, inputs: InferenceInput) -> StepResult: """Run one sequential inference step.""" ... - def reset(self, inputs: ModelInputs | None = None) -> None: + def reset(self, inputs: InferenceInput | None = None) -> None: """Reset this session's rollout state when the backend supports it.""" ... @@ -42,8 +42,8 @@ def close(self) -> None: 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 start_session(self, inputs: InferenceInput) -> InferenceSession: + """Create an isolated session from global conditioning inputs.""" ... def close(self) -> None: @@ -56,10 +56,10 @@ def close(self) -> None: 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. + Adapters declare model-facing input requirements, the canonical modalities + their default mapping consumes, and an optional default mapping between the + two. Runtime, application, or benchmark code may override that mapping while + preserving the same ``CanonicalInputs`` to ``InferenceInput`` boundary. """ @property @@ -68,17 +68,17 @@ def model_id(self) -> str: ... @property - def model_input_schema(self) -> ModelInputSchema: + def inference_input_schema(self) -> InferenceInputSchema: """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 canonical_input_schema(self) -> CanonicalInputSchema | None: + """Canonical modalities the adapter's default mapping consumes.""" ... def default_input_mapping(self) -> InputMapping | None: - """Return the model-provided default user-to-model mapping, if any.""" + """Return the model-provided default canonical-to-model mapping.""" ... def validate_config(self, config: InferenceConfig) -> None: diff --git a/flashdreams/flashdreams/runtime/mapping.py b/flashdreams/flashdreams/runtime/mapping.py index 756351081..1bb41d0c2 100644 --- a/flashdreams/flashdreams/runtime/mapping.py +++ b/flashdreams/flashdreams/runtime/mapping.py @@ -1,17 +1,24 @@ # 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.""" +"""Input mapping boundary from canonical inputs to encoded inference inputs.""" from __future__ import annotations -from typing import Protocol, runtime_checkable +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field, replace +from typing import Any, Protocol, runtime_checkable +from flashdreams.runtime._utils import freeze_mapping from flashdreams.runtime.inputs import ( - ModelInputs, - ModelInputSchema, - UserInputs, - UserInputSchema, + INPUT_PHASES, + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + InferenceInput, + InferenceInputSchema, + InputField, + InputPhase, ) from flashdreams.runtime.types import StepRequest @@ -28,28 +35,28 @@ class InputMapping(Protocol): def validate( self, *, - user_schema: UserInputSchema | None = None, - model_schema: ModelInputSchema | None = None, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, ) -> None: """Fail early for obvious app, event-source, and model mismatches.""" ... - def map_initial_inputs( + def map_global_inputs( self, *, - user_inputs: UserInputs, - model_inputs: ModelInputs, - ) -> ModelInputs: - """Build initial model inputs before a session starts.""" + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + """Build global conditioning inputs before a session starts.""" ... def map_step_inputs( self, *, - user_inputs: UserInputs, - model_inputs: ModelInputs, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, request: StepRequest, - ) -> ModelInputs: + ) -> InferenceInput: """Build model inputs for one session step from the current input window.""" ... @@ -60,26 +67,315 @@ class IdentityInputMapping: def validate( self, *, - user_schema: UserInputSchema | None = None, - model_schema: ModelInputSchema | None = None, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, ) -> None: - del user_schema, model_schema + del canonical_schema, inference_input_schema - def map_initial_inputs( + def map_global_inputs( self, *, - user_inputs: UserInputs, - model_inputs: ModelInputs, - ) -> ModelInputs: - del user_inputs - return model_inputs + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + del canonical_inputs + return inference_input def map_step_inputs( self, *, - user_inputs: UserInputs, - model_inputs: ModelInputs, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, request: StepRequest, - ) -> ModelInputs: - del user_inputs, request - return model_inputs + ) -> InferenceInput: + del canonical_inputs, request + return inference_input + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InputMappingSchema: + """Declarative compatibility surface for one mapping. + + ``InputMapping.validate`` fails a run late and opaquely: it raises, but it + cannot answer which optional model inputs a source would enable, or which + missing user capability is responsible for an unreachable model input. This + schema makes those questions answerable before runtime initialization. + """ + + name: str = "input-mapping" + consumes: tuple[CanonicalModality, ...] = () + produces_global: tuple[InputField, ...] = () + produces_step: tuple[InputField, ...] = () + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("InputMappingSchema.name must be non-empty.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + def produces_for(self, phase: InputPhase) -> tuple[InputField, ...]: + """Return the fields this mapping produces for ``phase``.""" + return self.produces_global if phase == "global" else self.produces_step + + def can_produce(self, phase: InputPhase, required: InputField) -> bool: + """Return whether this mapping can produce ``required`` in ``phase``.""" + return any( + _field_matches(produced, required) for produced in self.produces_for(phase) + ) + + +def _field_matches(produced: InputField, required: InputField) -> bool: + if produced.name != required.name: + return False + semantic_ok = ( + produced.semantic_type is None + or required.semantic_type is None + or produced.semantic_type == required.semantic_type + ) + lifecycle_ok = ( + produced.lifecycle is None + or required.lifecycle is None + or produced.lifecycle == required.lifecycle + ) + return semantic_ok and lifecycle_ok + + +@dataclass(frozen=True, kw_only=True, slots=True) +class MappingCompatibility: + """Compatibility report for one source, model schema, and mapping set. + + Mappings whose consumed capabilities the source cannot provide are reported + in ``unavailable_mapping_schemas`` and excluded from the satisfied/available + reports, so those lists only name model inputs that can really be produced. + """ + + __hash__ = None + + canonical_schema: CanonicalInputSchema + model_schema: InferenceInputSchema + mapping_schema: InputMappingSchema + missing_modalities: tuple[CanonicalModality, ...] = () + missing_required_model_fields: tuple[tuple[InputPhase, InputField], ...] = () + satisfied_required_model_fields: tuple[tuple[InputPhase, InputField], ...] = () + available_optional_model_fields: tuple[tuple[InputPhase, InputField], ...] = () + unavailable_mapping_schemas: tuple[InputMappingSchema, ...] = () + + @property + def can_drive(self) -> bool: + """Return whether this source can drive this model through the mapping. + + A mapping the source cannot feed does not block the run unless it was + the only way to produce a required model input. + """ + return not (self.missing_required_model_fields or self.missing_modalities) + + @property + def unavailable_mapping_names(self) -> tuple[str, ...]: + """Return names of mappings dropped because the source cannot feed them.""" + return tuple(schema.name for schema in self.unavailable_mapping_schemas) + + def raise_if_incompatible(self) -> None: + """Raise a compact error when this mapping cannot drive the model.""" + if self.can_drive: + return + problems: list[str] = [] + if self.missing_modalities: + missing = ", ".join(modality.name for modality in self.missing_modalities) + problems.append(f"missing canonical modalities: {missing}") + if self.missing_required_model_fields: + missing = ", ".join( + f"{phase}:{input_field.name}" + for phase, input_field in self.missing_required_model_fields + ) + problems.append(f"missing required model inputs: {missing}") + if self.unavailable_mapping_schemas: + problems.append( + "unavailable mappings: " + ", ".join(self.unavailable_mapping_names) + ) + raise ValueError( + f"Input mapping {self.mapping_schema.name!r} cannot drive this model " + f"from the selected source: " + "; ".join(problems) + ) + + +def _source_can_feed( + canonical_schema: CanonicalInputSchema, + mapping_schema: InputMappingSchema, +) -> bool: + return all( + canonical_schema.supports(modality) for modality in mapping_schema.consumes + ) + + +def combine_mapping_schemas( + mapping_schemas: Sequence[InputMappingSchema], + *, + name: str = "input-mapping-set", +) -> InputMappingSchema: + """Combine independently declared mappings into one compatibility surface. + + Duplicates are collapsed. Because ``metadata`` is excluded from equality, + the metadata of collapsed duplicates is merged rather than dropped, with the + first declaration winning on conflicting keys. + """ + consumes: list[CanonicalModality] = [] + produces: dict[InputPhase, list[InputField]] = {"global": [], "step": []} + + def _merge(target: list[Any], value: Any) -> None: + for index, existing in enumerate(target): + if existing == value: + if value.metadata: + target[index] = replace( + existing, + metadata={**dict(value.metadata), **dict(existing.metadata)}, + ) + return + target.append(value) + + for mapping_schema in mapping_schemas: + if not isinstance(mapping_schema, InputMappingSchema): + raise TypeError("mapping_schemas must contain InputMappingSchema objects.") + for modality in mapping_schema.consumes: + _merge(consumes, modality) + for phase in INPUT_PHASES: + for input_field in mapping_schema.produces_for(phase): + _merge(produces[phase], input_field) + + return InputMappingSchema( + name=name, + consumes=tuple(consumes), + produces_global=tuple(produces["global"]), + produces_step=tuple(produces["step"]), + ) + + +def _build_compatibility( + *, + canonical_schema: CanonicalInputSchema, + model_schema: InferenceInputSchema, + mapping_schemas: Sequence[InputMappingSchema], + reported_schema: InputMappingSchema, +) -> MappingCompatibility: + feedable: list[InputMappingSchema] = [] + unavailable: list[InputMappingSchema] = [] + for mapping_schema in mapping_schemas: + if _source_can_feed(canonical_schema, mapping_schema): + feedable.append(mapping_schema) + else: + unavailable.append(mapping_schema) + + usable = combine_mapping_schemas(feedable, name=reported_schema.name) + required = model_schema.required_fields() + missing_required = tuple( + (phase, input_field) + for phase, input_field in required + if not usable.can_produce(phase, input_field) + ) + satisfied_required = tuple( + (phase, input_field) + for phase, input_field in required + if usable.can_produce(phase, input_field) + ) + available_optional = tuple( + (phase, input_field) + for phase, input_field in model_schema.optional_fields() + if usable.can_produce(phase, input_field) + ) + + # Only capabilities that block a required model input make the mapping + # unusable. A dropped mapping that fed nothing but optional fields degrades + # the run instead of vetoing it. + missing_modalities: list[CanonicalModality] = [] + for mapping_schema in unavailable: + if not any( + mapping_schema.can_produce(phase, input_field) + for phase, input_field in missing_required + ): + continue + for modality in mapping_schema.consumes: + if canonical_schema.supports(modality) or modality in missing_modalities: + continue + missing_modalities.append(modality) + + return MappingCompatibility( + canonical_schema=canonical_schema, + model_schema=model_schema, + mapping_schema=reported_schema, + missing_modalities=tuple(missing_modalities), + missing_required_model_fields=missing_required, + satisfied_required_model_fields=satisfied_required, + available_optional_model_fields=available_optional, + unavailable_mapping_schemas=tuple(unavailable), + ) + + +def check_mapping_compatibility( + *, + canonical_schema: CanonicalInputSchema, + model_schema: InferenceInputSchema, + mapping_schema: InputMappingSchema, +) -> MappingCompatibility: + """Check whether a user-input source can drive a model through a mapping.""" + if not isinstance(mapping_schema, InputMappingSchema): + raise TypeError("mapping_schema must be an InputMappingSchema object.") + return _build_compatibility( + canonical_schema=canonical_schema, + model_schema=model_schema, + mapping_schemas=(mapping_schema,), + reported_schema=mapping_schema, + ) + + +def check_mapping_set_compatibility( + *, + canonical_schema: CanonicalInputSchema, + model_schema: InferenceInputSchema, + mapping_schemas: Sequence[InputMappingSchema], + name: str = "input-mapping-set", +) -> MappingCompatibility: + """Check compatibility for a composed set of mappings. + + Each mapping keeps its own consumes/produces link, so a mapping the source + cannot feed only costs the model inputs that mapping produced. + """ + mapping_schemas = tuple(mapping_schemas) + return _build_compatibility( + canonical_schema=canonical_schema, + model_schema=model_schema, + mapping_schemas=mapping_schemas, + reported_schema=combine_mapping_schemas(mapping_schemas, name=name), + ) + + +def undeclared_inference_input( + inputs: InferenceInput, + mapping_schema: InputMappingSchema, +) -> tuple[tuple[InputPhase, str], ...]: + """Return payload keys a mapping produced but did not declare. + + Mapping schemas are hand-written, so they drift from what + ``map_global_inputs``/``map_step_inputs`` actually return. Mapping tests + can use this to keep the declared compatibility surface honest. + """ + return tuple( + (phase, key) + for phase in INPUT_PHASES + for key in inputs.for_phase(phase) + if not any( + declared.name == key for declared in mapping_schema.produces_for(phase) + ) + ) + + +@runtime_checkable +class DeclaresMappingSchema(Protocol): + """Optional refinement of :class:`InputMapping` that declares its surface.""" + + @property + def mapping_schema(self) -> InputMappingSchema: + """Return the declarative compatibility surface for this mapping.""" + ... diff --git a/flashdreams/flashdreams/runtime/types.py b/flashdreams/flashdreams/runtime/types.py index 52bf82166..467753026 100644 --- a/flashdreams/flashdreams/runtime/types.py +++ b/flashdreams/flashdreams/runtime/types.py @@ -10,7 +10,7 @@ from typing import Any from flashdreams.runtime._utils import freeze_mapping -from flashdreams.runtime.inputs import ModelInputSchema, TimeWindow +from flashdreams.runtime.inputs import InferenceInputSchema, TimeWindow @dataclass(frozen=True, kw_only=True, slots=True) @@ -24,7 +24,7 @@ class StepRequest: __hash__ = None step_index: int - model_input_schema: ModelInputSchema | None = None + inference_input_schema: InferenceInputSchema | None = None user_input_window: TimeWindow | None = None metadata: Mapping[str, Any] = field(default_factory=dict) diff --git a/flashdreams/tests/test_inference_inputs.py b/flashdreams/tests/test_inference_inputs.py deleted file mode 100644 index b3d090a7a..000000000 --- a/flashdreams/tests/test_inference_inputs.py +++ /dev/null @@ -1,1154 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - -import pytest - -from flashdreams.inference import ( - InputMapper, - InputMapperSchema, - ModelInputField, - ModelInputSchema, - ModelInputs, - StaticInputMapper, - UserInputCapability, - UserInputEvent, - UserInputSchema, - UserInputTrace, - UserInputWindow, - check_mapping_compatibility, - check_mapping_set_compatibility, - combine_mapper_schemas, - missing_required_inputs, - undeclared_model_inputs, -) - -pytestmark = pytest.mark.ci_cpu - - -def test_user_input_trace_orders_events_and_slices_windows() -> None: - key_down = UserInputEvent( - timestamp_s=1.0, - kind="key_down", - payload={"key": "w"}, - ) - key_up = UserInputEvent( - timestamp_s=1.0, - kind="key_up", - payload={"key": "w"}, - ) - prompt = UserInputEvent( - timestamp_s=0.0, - kind="prompt_set", - payload={"prompt": "drive forward"}, - ) - - trace = UserInputTrace.from_events([key_down, key_up, prompt]) - - assert [event.kind for event in trace.events] == [ - "prompt_set", - "key_down", - "key_up", - ] - window = trace.window(start_s=0.5, end_s=1.0, include_end=True) - assert window.events == (key_down, key_up) - assert window.latest("key_down") is key_down - - -def test_static_startup_values_are_user_input_events() -> None: - trace = UserInputTrace.from_events( - [ - UserInputEvent( - timestamp_s=0.0, - kind="initial_frame_set", - payload={"image": b"encoded-image"}, - ), - UserInputEvent( - timestamp_s=0.0, - kind="scene_selected", - payload={"scene_id": "scene-a"}, - ), - ] - ) - - initial_frame = trace.latest("initial_frame_set") - scene = trace.latest("scene_selected") - assert initial_frame is not None - assert scene is not None - assert initial_frame.payload["image"] == b"encoded-image" - assert scene.payload["scene_id"] == "scene-a" - - -def test_user_input_schema_declares_and_validates_source_capabilities() -> None: - schema = UserInputSchema( - name="browser", - source_kind="live", - metadata={"transport": "webrtc"}, - capabilities=( - UserInputCapability( - event_kind="prompt_set", - payload_kind="text", - payload_fields=frozenset({"prompt"}), - metadata={"source_widget": "prompt-box"}, - ), - UserInputCapability( - event_kind="key_down", - payload_fields=frozenset({"key"}), - ), - ), - ) - - assert schema.supports( - UserInputCapability( - event_kind="prompt_set", - payload_kind="text", - payload_fields=frozenset({"prompt"}), - ) - ) - assert schema.metadata["transport"] == "webrtc" - assert schema.capabilities[0].metadata["source_widget"] == "prompt-box" - schema.validate_event( - UserInputEvent( - timestamp_s=0.0, - kind="prompt_set", - payload={"prompt": "a prompt"}, - ) - ) - with pytest.raises(ValueError, match="missing required fields"): - schema.validate_event( - UserInputEvent( - timestamp_s=0.0, - kind="prompt_set", - payload={"text": "wrong key"}, - ) - ) - - -def test_model_input_schema_declares_required_optional_and_update_metadata() -> None: - schema = ModelInputSchema( - name="steering-model", - metadata={"model_family": "example"}, - fields=( - ModelInputField( - name="prompt", - phase="initial", - required=True, - payload_kind="text", - update_policy="step_boundary", - lifecycle="cache_init", - metadata={"max_tokens": 256}, - ), - ModelInputField( - name="steering", - phase="step", - required=True, - lifecycle="step_input", - ), - ModelInputField(name="first_frame", phase="initial", required=False), - ), - ) - - assert [field.name for field in schema.required_fields()] == [ - "prompt", - "steering", - ] - assert [field.name for field in schema.optional_fields()] == ["first_frame"] - prompt_field = schema.field_for(name="prompt", phase="initial") - assert prompt_field is not None - assert prompt_field.update_policy == "step_boundary" - assert prompt_field.lifecycle == "cache_init" - assert prompt_field.metadata["max_tokens"] == 256 - assert schema.metadata["model_family"] == "example" - - -def test_model_inputs_report_missing_required_fields() -> None: - schema = ModelInputSchema( - fields=( - ModelInputField(name="prompt", phase="initial"), - ModelInputField(name="steering", phase="step"), - ) - ) - inputs = ModelInputs(initial={"prompt": "drive"}, step={}) - - missing = missing_required_inputs(inputs, schema) - - assert [(field.phase, field.name) for field in missing] == [("step", "steering")] - - -def test_mapping_compatibility_reports_satisfied_missing_and_optional_fields() -> None: - source = UserInputSchema( - name="keyboard-app", - capabilities=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - UserInputCapability( - event_kind="key_down", - payload_fields=frozenset({"key"}), - ), - UserInputCapability( - event_kind="key_up", - payload_fields=frozenset({"key"}), - ), - ), - ) - model = ModelInputSchema( - name="drive-model", - fields=( - ModelInputField(name="prompt", phase="initial"), - ModelInputField(name="steering", phase="step"), - ModelInputField(name="first_frame", phase="initial", required=False), - ), - ) - mapper = InputMapperSchema( - name="keyboard-drive", - consumes=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - UserInputCapability( - event_kind="key_down", - payload_fields=frozenset({"key"}), - ), - UserInputCapability( - event_kind="key_up", - payload_fields=frozenset({"key"}), - ), - ), - produces=( - ModelInputField(name="prompt", phase="initial"), - ModelInputField(name="steering", phase="step"), - ), - ) - - compatibility = check_mapping_compatibility( - source_schema=source, - model_schema=model, - mapper_schema=mapper, - ) - - assert compatibility.can_drive - assert [field.name for field in compatibility.satisfied_required_model_fields] == [ - "prompt", - "steering", - ] - assert compatibility.available_optional_model_fields == () - - -def test_mapping_compatibility_reports_available_optional_model_input() -> None: - source = UserInputSchema( - capabilities=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - UserInputCapability( - event_kind="initial_frame_set", - payload_fields=frozenset({"image"}), - ), - ) - ) - model = ModelInputSchema( - fields=( - ModelInputField(name="prompt", phase="initial"), - ModelInputField(name="first_frame", phase="initial", required=False), - ) - ) - mapper = InputMapperSchema( - consumes=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - UserInputCapability( - event_kind="initial_frame_set", - payload_fields=frozenset({"image"}), - ), - ), - produces=( - ModelInputField(name="prompt", phase="initial"), - ModelInputField(name="first_frame", phase="initial", required=False), - ), - ) - - compatibility = check_mapping_compatibility( - source_schema=source, - model_schema=model, - mapper_schema=mapper, - ) - - assert compatibility.can_drive - assert [ - (field.phase, field.name) - for field in compatibility.available_optional_model_fields - ] == [("initial", "first_frame")] - - -def test_mapping_compatibility_reports_missing_required_model_input() -> None: - source = UserInputSchema( - capabilities=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - ) - ) - model = ModelInputSchema( - fields=( - ModelInputField(name="prompt", phase="initial"), - ModelInputField(name="camera_trajectory", phase="step"), - ) - ) - mapper = InputMapperSchema( - consumes=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - ), - produces=(ModelInputField(name="prompt", phase="initial"),), - ) - - compatibility = check_mapping_compatibility( - source_schema=source, - model_schema=model, - mapper_schema=mapper, - ) - - assert not compatibility.can_drive - assert [ - (field.phase, field.name) - for field in compatibility.missing_required_model_fields - ] == [("step", "camera_trajectory")] - with pytest.raises(ValueError, match="missing required model inputs"): - compatibility.raise_if_incompatible() - - -def test_mapping_compatibility_reports_missing_source_capability() -> None: - source = UserInputSchema( - capabilities=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - ) - ) - model = ModelInputSchema(fields=(ModelInputField(name="steering", phase="step"),)) - mapper = InputMapperSchema( - consumes=( - UserInputCapability( - event_kind="controller_axis", - payload_fields=frozenset({"axis", "value"}), - ), - ), - produces=(ModelInputField(name="steering", phase="step"),), - ) - - compatibility = check_mapping_compatibility( - source_schema=source, - model_schema=model, - mapper_schema=mapper, - ) - - assert not compatibility.can_drive - assert [ - capability.event_kind - for capability in compatibility.missing_source_capabilities - ] == ["controller_axis"] - - -def test_mapping_set_compatibility_supports_composed_model_inputs() -> None: - source = UserInputSchema( - name="browser-with-controls", - capabilities=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - UserInputCapability( - event_kind="initial_frame_set", - payload_fields=frozenset({"image"}), - ), - UserInputCapability( - event_kind="key_down", - payload_fields=frozenset({"key"}), - ), - UserInputCapability( - event_kind="key_up", - payload_fields=frozenset({"key"}), - ), - ), - ) - model = ModelInputSchema( - name="lingbot-like", - fields=( - ModelInputField( - name="prompt", - phase="initial", - lifecycle="cache_init", - ), - ModelInputField( - name="first_frame", - phase="initial", - lifecycle="cache_init", - ), - ModelInputField( - name="camera_trajectory", - phase="step", - lifecycle="step_input", - ), - ), - ) - prompt_mapper = InputMapperSchema( - name="prompt", - consumes=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - ), - produces=( - ModelInputField( - name="prompt", - phase="initial", - lifecycle="cache_init", - ), - ), - ) - frame_mapper = InputMapperSchema( - name="first-frame", - consumes=( - UserInputCapability( - event_kind="initial_frame_set", - payload_fields=frozenset({"image"}), - ), - ), - produces=( - ModelInputField( - name="first_frame", - phase="initial", - lifecycle="cache_init", - ), - ), - ) - camera_mapper = InputMapperSchema( - name="keyboard-to-camera", - consumes=( - UserInputCapability( - event_kind="key_down", - payload_fields=frozenset({"key"}), - ), - UserInputCapability( - event_kind="key_up", - payload_fields=frozenset({"key"}), - ), - ), - produces=( - ModelInputField( - name="camera_trajectory", - phase="step", - lifecycle="step_input", - ), - ), - ) - - compatibility = check_mapping_set_compatibility( - source_schema=source, - model_schema=model, - mapper_schemas=(prompt_mapper, frame_mapper, camera_mapper), - name="browser-lingbot", - ) - - assert compatibility.can_drive - assert compatibility.mapper_schema.name == "browser-lingbot" - assert [field.name for field in compatibility.satisfied_required_model_fields] == [ - "prompt", - "first_frame", - "camera_trajectory", - ] - assert [ - capability.event_kind for capability in compatibility.mapper_schema.consumes - ] == [ - "prompt_set", - "initial_frame_set", - "key_down", - "key_up", - ] - - -def test_unfeedable_optional_mapper_degrades_instead_of_blocking_the_run() -> None: - source = UserInputSchema( - name="prompt-only", - capabilities=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - ), - ) - model = ModelInputSchema( - fields=( - ModelInputField(name="prompt", phase="initial"), - ModelInputField(name="steering", phase="step", required=False), - ) - ) - prompt_mapper = InputMapperSchema( - name="prompt", - consumes=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - ), - produces=(ModelInputField(name="prompt", phase="initial"),), - ) - gamepad_mapper = InputMapperSchema( - name="gamepad-steering", - consumes=( - UserInputCapability( - event_kind="controller_axis", - payload_fields=frozenset({"axis", "value"}), - ), - ), - produces=(ModelInputField(name="steering", phase="step", required=False),), - ) - - compatibility = check_mapping_set_compatibility( - source_schema=source, - model_schema=model, - mapper_schemas=(prompt_mapper, gamepad_mapper), - ) - - assert compatibility.can_drive - assert compatibility.unavailable_mapper_names == ("gamepad-steering",) - assert compatibility.missing_source_capabilities == () - assert [field.name for field in compatibility.satisfied_required_model_fields] == [ - "prompt" - ] - # The optional field is not reported as available: nothing can produce it. - assert compatibility.available_optional_model_fields == () - - -def test_unfeedable_required_mapper_still_blocks_the_run() -> None: - source = UserInputSchema(name="prompt-only") - model = ModelInputSchema(fields=(ModelInputField(name="steering", phase="step"),)) - gamepad_mapper = InputMapperSchema( - name="gamepad-steering", - consumes=( - UserInputCapability( - event_kind="controller_axis", - payload_fields=frozenset({"axis", "value"}), - ), - ), - produces=(ModelInputField(name="steering", phase="step"),), - ) - - compatibility = check_mapping_set_compatibility( - source_schema=source, - model_schema=model, - mapper_schemas=(gamepad_mapper,), - ) - - assert not compatibility.can_drive - assert [ - capability.event_kind - for capability in compatibility.missing_source_capabilities - ] == ["controller_axis"] - assert compatibility.unavailable_mapper_names == ("gamepad-steering",) - with pytest.raises(ValueError, match="unavailable mappers: gamepad-steering"): - compatibility.raise_if_incompatible() - - -def test_mapper_schema_combination_deduplicates_shared_capabilities() -> None: - shared_prompt = UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ) - prompt_field = ModelInputField(name="prompt", phase="initial") - - combined = combine_mapper_schemas( - ( - InputMapperSchema( - name="prompt-a", - consumes=(shared_prompt,), - produces=(prompt_field,), - ), - InputMapperSchema( - name="prompt-b", - consumes=(shared_prompt,), - produces=(prompt_field,), - ), - ) - ) - - assert combined.consumes == (shared_prompt,) - assert combined.produces == (prompt_field,) - - -def test_mapper_schema_combination_merges_metadata_of_collapsed_duplicates() -> None: - combined = combine_mapper_schemas( - ( - InputMapperSchema( - name="prompt-a", - consumes=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - metadata={"widget": "prompt-box", "owner": "app-a"}, - ), - ), - ), - InputMapperSchema( - name="prompt-b", - consumes=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - metadata={"max_length": 512, "owner": "app-b"}, - ), - ), - ), - ) - ) - - assert len(combined.consumes) == 1 - assert dict(combined.consumes[0].metadata) == { - "widget": "prompt-box", - "max_length": 512, - "owner": "app-a", - } - - -def test_user_input_event_is_hashable_despite_mapping_payload() -> None: - event = UserInputEvent( - timestamp_s=0.5, - kind="key_down", - payload={"key": "w", "modifiers": ["shift"]}, - ) - same = UserInputEvent( - timestamp_s=0.5, - kind="key_down", - payload={"key": "w", "modifiers": ["shift"]}, - ) - - assert {event, same} == {event} - assert hash(event) == hash(same) - - -def test_user_input_window_rejects_events_outside_its_bounds() -> None: - with pytest.raises(ValueError, match="outside window"): - UserInputWindow( - start_s=0.0, - end_s=1.0, - events=(UserInputEvent(timestamp_s=2.0, kind="key_down"),), - ) - - -def test_metadata_keys_round_trip_unchanged() -> None: - capability = UserInputCapability( - event_kind="prompt_set", - metadata={" spaced key ": "kept", "schema.uri": "urn:example"}, - ) - - assert dict(capability.metadata) == { - " spaced key ": "kept", - "schema.uri": "urn:example", - } - non_string_keys: dict[Any, Any] = {1: "coerced"} - with pytest.raises(TypeError, match="metadata keys must be strings"): - UserInputCapability(event_kind="prompt_set", metadata=non_string_keys) - - -def test_lifecycle_mismatch_does_not_satisfy_model_field() -> None: - source = UserInputSchema(name="fixed") - model = ModelInputSchema( - fields=( - ModelInputField( - name="video_dimensions", - phase="initial", - lifecycle="runtime_config", - ), - ) - ) - mapper = InputMapperSchema( - produces=( - ModelInputField( - name="video_dimensions", - phase="initial", - lifecycle="cache_init", - ), - ) - ) - - compatibility = check_mapping_compatibility( - source_schema=source, - model_schema=model, - mapper_schema=mapper, - ) - - assert not compatibility.can_drive - assert [ - (field.phase, field.name, field.lifecycle) - for field in compatibility.missing_required_model_fields - ] == [("initial", "video_dimensions", "runtime_config")] - - -def test_metadata_is_queryable_but_not_part_of_compatibility_matching() -> None: - source = UserInputSchema( - capabilities=( - UserInputCapability( - event_kind="pose_path_set", - payload_kind="path", - payload_fields=frozenset({"path"}), - metadata={"file_format": "npy"}, - ), - ), - metadata={"owner": "future-integration"}, - ) - model = ModelInputSchema( - fields=( - ModelInputField( - name="camera_trajectory", - phase="initial", - payload_kind="c2w_sequence", - lifecycle="rollout_binding", - metadata={"coordinates": "opencv_c2w"}, - ), - ), - metadata={"model_family": "future-world-model"}, - ) - mapper = InputMapperSchema( - consumes=( - UserInputCapability( - event_kind="pose_path_set", - payload_kind="path", - payload_fields=frozenset({"path"}), - metadata={"accepted_suffixes": (".npy",)}, - ), - ), - produces=( - ModelInputField( - name="camera_trajectory", - phase="initial", - payload_kind="c2w_sequence", - lifecycle="rollout_binding", - metadata={"shape": "[F,4,4]"}, - ), - ), - metadata={"mapper_family": "camera-path-loader"}, - ) - - compatibility = check_mapping_compatibility( - source_schema=source, - model_schema=model, - mapper_schema=mapper, - ) - - assert compatibility.can_drive - assert source.metadata["owner"] == "future-integration" - assert model.fields[0].metadata["coordinates"] == "opencv_c2w" - assert mapper.produces[0].metadata["shape"] == "[F,4,4]" - - -def test_sana_wm_like_schema_uses_open_ended_model_inputs() -> None: - source = UserInputSchema( - name="sana-wm-cli-like", - capabilities=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - UserInputCapability( - event_kind="initial_frame_set", - payload_fields=frozenset({"image"}), - ), - UserInputCapability( - event_kind="camera_action_set", - payload_fields=frozenset({"action"}), - metadata={"dsl": "-"}, - ), - UserInputCapability( - event_kind="intrinsics_set", - payload_fields=frozenset({"intrinsics"}), - metadata={"optional": True}, - ), - UserInputCapability( - event_kind="rollout_parameter_set", - payload_fields=frozenset({"name", "value"}), - ), - ), - ) - model = ModelInputSchema( - name="sana-wm-like", - fields=( - ModelInputField(name="prompt", phase="initial", lifecycle="cache_init"), - ModelInputField( - name="negative_prompt", - phase="initial", - required=False, - lifecycle="cache_init", - ), - ModelInputField( - name="first_frame", - phase="initial", - lifecycle="cache_init", - ), - ModelInputField( - name="camera_trajectory_c2w", - phase="initial", - payload_kind="c2w_sequence", - lifecycle="rollout_binding", - metadata={"shape": "[F,4,4]"}, - ), - ModelInputField( - name="camera_intrinsics_vec4", - phase="initial", - required=False, - payload_kind="intrinsics_vec4_sequence", - lifecycle="rollout_binding", - metadata={"shape": "[F,4]"}, - ), - ModelInputField( - name="stage1_sampling", - phase="initial", - lifecycle="rollout_binding", - metadata={"fields": ("steps", "cfg_scale", "flow_shift", "seed")}, - ), - ModelInputField( - name="streaming_chunking", - phase="initial", - required=False, - lifecycle="rollout_binding", - metadata={"fields": ("num_frame_per_block", "cached_blocks")}, - ), - ), - metadata={"model_family": "sana-wm"}, - ) - prompt_mapper = InputMapperSchema( - produces=(ModelInputField(name="prompt", phase="initial"),), - consumes=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - ), - ) - frame_mapper = InputMapperSchema( - produces=(ModelInputField(name="first_frame", phase="initial"),), - consumes=( - UserInputCapability( - event_kind="initial_frame_set", - payload_fields=frozenset({"image"}), - ), - ), - ) - camera_mapper = InputMapperSchema( - produces=( - ModelInputField( - name="camera_trajectory_c2w", - phase="initial", - payload_kind="c2w_sequence", - lifecycle="rollout_binding", - ), - ModelInputField( - name="camera_intrinsics_vec4", - phase="initial", - payload_kind="intrinsics_vec4_sequence", - lifecycle="rollout_binding", - ), - ), - consumes=( - UserInputCapability( - event_kind="camera_action_set", - payload_fields=frozenset({"action"}), - ), - UserInputCapability( - event_kind="intrinsics_set", - payload_fields=frozenset({"intrinsics"}), - ), - ), - ) - sampling_mapper = InputMapperSchema( - produces=( - ModelInputField( - name="stage1_sampling", - phase="initial", - lifecycle="rollout_binding", - ), - ), - consumes=( - UserInputCapability( - event_kind="rollout_parameter_set", - payload_fields=frozenset({"name", "value"}), - ), - ), - ) - - compatibility = check_mapping_set_compatibility( - source_schema=source, - model_schema=model, - mapper_schemas=(prompt_mapper, frame_mapper, camera_mapper, sampling_mapper), - ) - - assert compatibility.can_drive - assert [field.name for field in compatibility.satisfied_required_model_fields] == [ - "prompt", - "first_frame", - "camera_trajectory_c2w", - "stage1_sampling", - ] - assert [field.name for field in compatibility.available_optional_model_fields] == [ - "camera_intrinsics_vec4", - ] - - -@dataclass(frozen=True) -class PromptMapper: - schema: InputMapperSchema = InputMapperSchema( - name="prompt", - consumes=( - UserInputCapability( - event_kind="prompt_set", - payload_fields=frozenset({"prompt"}), - ), - ), - produces=(ModelInputField(name="prompt", phase="initial"),), - ) - - def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: - prompt = trace.latest("prompt_set") - return ModelInputs.initial_only( - {"prompt": prompt.payload["prompt"]} if prompt is not None else {} - ) - - def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: - del window - return ModelInputs() - - -@dataclass(frozen=True) -class InitialFrameMapper: - schema: InputMapperSchema = InputMapperSchema( - name="initial-frame", - consumes=( - UserInputCapability( - event_kind="initial_frame_set", - payload_fields=frozenset({"image"}), - ), - ), - produces=( - ModelInputField(name="first_frame", phase="initial", required=False), - ), - ) - - def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: - event = trace.latest("initial_frame_set") - return ModelInputs.initial_only( - {"first_frame": event.payload["image"]} if event is not None else {} - ) - - def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: - del window - return ModelInputs() - - -@dataclass(frozen=True) -class KeyboardSteeringMapper: - schema: InputMapperSchema = InputMapperSchema( - name="keyboard-steering", - consumes=( - UserInputCapability( - event_kind="key_down", - payload_fields=frozenset({"key"}), - ), - UserInputCapability( - event_kind="key_up", - payload_fields=frozenset({"key"}), - ), - ), - produces=(ModelInputField(name="steering", phase="step"),), - ) - - def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: - del trace - return ModelInputs() - - def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: - steering = 0.0 - for event in window.events: - key = event.payload.get("key") - if event.kind == "key_down" and key == "a": - steering = 1.0 - elif event.kind == "key_down" and key == "d": - steering = -1.0 - elif event.kind == "key_up" and key in {"a", "d"}: - steering = 0.0 - return ModelInputs.step_only({"steering": steering}) - - -@dataclass(frozen=True) -class CameraTrajectoryMapper: - schema: InputMapperSchema = InputMapperSchema( - name="camera-trajectory", - consumes=( - UserInputCapability( - event_kind="camera_pose", - payload_fields=frozenset({"pose"}), - ), - ), - produces=(ModelInputField(name="camera_trajectory", phase="step"),), - ) - - def build_initial_inputs(self, trace: UserInputTrace) -> ModelInputs: - del trace - return ModelInputs() - - def build_step_inputs(self, window: UserInputWindow) -> ModelInputs: - return ModelInputs.step_only( - { - "camera_trajectory": tuple( - event.payload["pose"] - for event in window.events_of_kind("camera_pose") - ) - } - ) - - -def test_fake_prompt_and_initial_frame_mappers_build_initial_model_inputs() -> None: - trace = UserInputTrace.from_events( - [ - UserInputEvent( - timestamp_s=0.0, - kind="prompt_set", - payload={"prompt": "a road"}, - ), - UserInputEvent( - timestamp_s=0.0, - kind="initial_frame_set", - payload={"image": b"frame"}, - ), - ] - ) - - assert PromptMapper().build_initial_inputs(trace).initial == {"prompt": "a road"} - assert InitialFrameMapper().build_initial_inputs(trace).initial == { - "first_frame": b"frame" - } - - -def test_static_input_mapper_supports_fixed_model_inputs() -> None: - inputs = ModelInputs( - initial={"prompt": "fixed prompt"}, - step={"camera_trajectory": ("pose-a", "pose-b")}, - ) - mapper = StaticInputMapper.from_inputs(inputs=inputs, name="fixed") - source = UserInputSchema(name="no-live-controls") - model = ModelInputSchema( - fields=( - ModelInputField(name="prompt", phase="initial"), - ModelInputField(name="camera_trajectory", phase="step"), - ) - ) - - compatibility = check_mapping_compatibility( - source_schema=source, - model_schema=model, - mapper_schema=mapper.schema, - ) - - assert compatibility.can_drive - assert mapper.build_initial_inputs(UserInputTrace()).initial == { - "prompt": "fixed prompt" - } - assert mapper.build_step_inputs(UserInputWindow(start_s=0.0, end_s=1.0)).step == { - "camera_trajectory": ("pose-a", "pose-b") - } - - -def test_fake_keyboard_mapper_builds_step_model_inputs() -> None: - trace = UserInputTrace.from_events( - [ - UserInputEvent(timestamp_s=0.1, kind="key_down", payload={"key": "a"}), - UserInputEvent(timestamp_s=0.2, kind="key_up", payload={"key": "a"}), - UserInputEvent(timestamp_s=0.3, kind="key_down", payload={"key": "d"}), - ] - ) - - inputs = KeyboardSteeringMapper().build_step_inputs( - trace.window(start_s=0.0, end_s=0.4) - ) - - assert inputs.step == {"steering": -1.0} - - -@pytest.mark.parametrize( - "mapper", - [ - PromptMapper(), - InitialFrameMapper(), - KeyboardSteeringMapper(), - CameraTrajectoryMapper(), - StaticInputMapper.from_inputs( - inputs=ModelInputs(initial={"prompt": "fixed"}, step={"steering": 0.0}), - ), - ], -) -def test_mappers_only_produce_model_inputs_they_declare(mapper: InputMapper) -> None: - trace = UserInputTrace.from_events( - [ - UserInputEvent( - timestamp_s=0.0, - kind="prompt_set", - payload={"prompt": "a road"}, - ), - UserInputEvent( - timestamp_s=0.0, - kind="initial_frame_set", - payload={"image": b"frame"}, - ), - UserInputEvent(timestamp_s=0.1, kind="key_down", payload={"key": "a"}), - UserInputEvent( - timestamp_s=0.2, - kind="camera_pose", - payload={"pose": "pose-a"}, - ), - ] - ) - window = trace.window(start_s=0.0, end_s=1.0) - - assert isinstance(mapper, InputMapper) - assert ( - undeclared_model_inputs(mapper.build_initial_inputs(trace), mapper.schema) == () - ) - assert ( - undeclared_model_inputs(mapper.build_step_inputs(window), mapper.schema) == () - ) - - -def test_fake_camera_trajectory_mapper_builds_step_model_inputs() -> None: - trace = UserInputTrace.from_events( - [ - UserInputEvent( - timestamp_s=0.1, - kind="camera_pose", - payload={"pose": "pose-a"}, - ), - UserInputEvent( - timestamp_s=0.2, - kind="camera_pose", - payload={"pose": "pose-b"}, - ), - ] - ) - - inputs = CameraTrajectoryMapper().build_step_inputs( - trace.window(start_s=0.0, end_s=0.3) - ) - - assert inputs.step == {"camera_trajectory": ("pose-a", "pose-b")} diff --git a/flashdreams/tests/test_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py index 1474383a0..32fe7bdc2 100644 --- a/flashdreams/tests/test_inference_runtime_api.py +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -9,17 +9,20 @@ import pytest from flashdreams.runtime import ( + CanonicalInputs, + CanonicalInputSchema, IdentityInputMapping, InferenceConfig, + InferenceInput, + InferenceInputSchema, InferenceRuntime, InferenceSession, InMemoryMetricsRecorder, + InputCanonicalizer, InputField, InputMapping, MetricsRecorder, ModelAdapter, - ModelInputs, - ModelInputSchema, NullOutputTarget, OutputArtifact, OutputTarget, @@ -27,6 +30,7 @@ StepRequest, StepResult, TimeWindow, + UserInputCapability, UserInputEvent, UserInputs, UserInputSchema, @@ -35,6 +39,18 @@ pytestmark = pytest.mark.ci_cpu +_SESSION_HORIZON_S = 3600.0 + +_KEYBOARD_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="keyboard.keydown", payload_fields=frozenset({"key"}) + ), + ) +) +_KEYBOARD_CANONICALIZER = InputCanonicalizer() + + def test_inference_config_keeps_runtime_settings_separate() -> None: denied_app_fields = {"prompt", "output_dir", "browser_settings"} config = InferenceConfig( @@ -90,17 +106,19 @@ def test_runtime_metric_sample_rejects_bool_values() -> None: RuntimeMetricSample(name="sample", value=True) -def test_model_input_schema_validates_initial_and_step_payloads() -> None: - schema = ModelInputSchema( - initial_fields=( +def test_inference_input_schema_validates_initial_and_step_payloads() -> None: + schema = InferenceInputSchema( + global_fields=( InputField(name="prompt"), - InputField(name="first_frame"), + InputField(name="global_conditioning_frame"), ), step_fields=(InputField(name="camera_poses"),), ) - inputs = ModelInputs(initial={"prompt": "drive", "first_frame": object()}) + inputs = InferenceInput( + global_conditioning={"prompt": "drive", "global_conditioning_frame": object()} + ) - schema.require_initial(inputs) + schema.require_global(inputs) assert schema.missing_step(inputs) == ("camera_poses",) with pytest.raises(ValueError, match="camera_poses"): @@ -164,25 +182,27 @@ def test_user_input_schema_validates_required_snapshot_fields() -> None: schema.require_snapshot(UserInputs()) -def test_identity_input_mapping_leaves_model_inputs_unchanged() -> None: +def test_identity_input_mapping_leaves_inference_input_unchanged() -> None: mapping = IdentityInputMapping() - model_inputs = ModelInputs(initial={"prompt": "fixed"}, step={"hdmap": object()}) + inference_input = InferenceInput( + global_conditioning={"prompt": "fixed"}, step={"hdmap": object()} + ) request = StepRequest(step_index=0) assert ( - mapping.map_initial_inputs( - user_inputs=UserInputs(), - model_inputs=model_inputs, + mapping.map_global_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=inference_input, ) - is model_inputs + is inference_input ) assert ( mapping.map_step_inputs( - user_inputs=UserInputs(), - model_inputs=model_inputs, + canonical_inputs=CanonicalInputs(), + inference_input=inference_input, request=request, ) - is model_inputs + is inference_input ) @@ -258,7 +278,7 @@ def test_runtime_api_components_compose_for_sequential_session() -> None: ), ) ) - model_inputs = ModelInputs(initial={"prompt": "drive forward"}) + inference_input = InferenceInput(global_conditioning={"prompt": "drive forward"}) output = NullOutputTarget(store_results=True) metrics = InMemoryMetricsRecorder() @@ -269,8 +289,10 @@ def test_runtime_api_components_compose_for_sequential_session() -> None: adapter=adapter, config=config, mapping=mapping, + canonicalizer=_KEYBOARD_CANONICALIZER, + source_schema=_KEYBOARD_SOURCE, user_inputs=user_inputs, - model_inputs=model_inputs, + inference_input=inference_input, output=output, metrics=metrics, ) @@ -291,8 +313,10 @@ def test_reference_loop_validates_mapping_before_runtime_creation() -> None: adapter=adapter, config=InferenceConfig(model_id="fake-model"), mapping=mapping, + canonicalizer=_KEYBOARD_CANONICALIZER, + source_schema=_KEYBOARD_SOURCE, user_inputs=UserInputs(), - model_inputs=ModelInputs(initial={"prompt": "drive forward"}), + inference_input=InferenceInput(global_conditioning={"prompt": "drive forward"}), output=NullOutputTarget(), metrics=InMemoryMetricsRecorder(), ) @@ -311,8 +335,12 @@ def test_reference_loop_closes_runtime_when_session_start_fails() -> None: adapter=adapter, config=InferenceConfig(model_id="fake-model"), mapping=IdentityInputMapping(), + canonicalizer=_KEYBOARD_CANONICALIZER, + source_schema=_KEYBOARD_SOURCE, user_inputs=UserInputs(), - model_inputs=ModelInputs(initial={"prompt": "drive forward"}), + inference_input=InferenceInput( + global_conditioning={"prompt": "drive forward"} + ), output=output, metrics=metrics, ) @@ -328,18 +356,24 @@ def _drive_two_step_session( adapter: ModelAdapter, config: InferenceConfig, mapping: InputMapping, + canonicalizer: InputCanonicalizer, + source_schema: UserInputSchema, user_inputs: UserInputs, - model_inputs: ModelInputs, + inference_input: InferenceInput, output: OutputTarget, metrics: MetricsRecorder, ) -> None: mapping.validate( - user_schema=adapter.user_input_schema, - model_schema=adapter.model_input_schema, + canonical_schema=adapter.canonical_input_schema, + inference_input_schema=adapter.inference_input_schema, ) - initial_inputs = mapping.map_initial_inputs( - user_inputs=user_inputs, - model_inputs=model_inputs, + initial_inputs = mapping.map_global_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=TimeWindow(start_s=0.0, end_s=_SESSION_HORIZON_S), + source_schema=source_schema, + ), + inference_input=inference_input, ) runtime = adapter.create_runtime(config) session: InferenceSession | None = None @@ -350,13 +384,14 @@ def _drive_two_step_session( 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 + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window + or TimeWindow(start_s=0.0, end_s=_SESSION_HORIZON_S), + source_schema=source_schema, ), - model_inputs=ModelInputs( - initial=initial_inputs.initial, + inference_input=InferenceInput( + global_conditioning=initial_inputs.global_conditioning, step={"chunk_index": request.step_index}, ), request=request, @@ -379,11 +414,11 @@ def _drive_two_step_session( class _FakeAdapter: model_id = "fake-model" - model_input_schema = ModelInputSchema( - initial_fields=(InputField(name="prompt"),), + inference_input_schema = InferenceInputSchema( + global_fields=(InputField(name="prompt"),), step_fields=(InputField(name="chunk_index"),), ) - user_input_schema = UserInputSchema(event_types=frozenset({"keyboard.keydown"})) + canonical_input_schema = CanonicalInputSchema() def default_input_mapping(self) -> InputMapping: return IdentityInputMapping() @@ -394,31 +429,31 @@ def validate_config(self, config: InferenceConfig) -> None: def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: self.validate_config(config) - return _FakeRuntime(model_input_schema=self.model_input_schema) + return _FakeRuntime(inference_input_schema=self.inference_input_schema) class _FakeRuntime: - def __init__(self, *, model_input_schema: ModelInputSchema) -> None: - self._model_input_schema = model_input_schema + def __init__(self, *, inference_input_schema: InferenceInputSchema) -> None: + self._inference_input_schema = inference_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 start_session(self, inputs: InferenceInput) -> InferenceSession: + self._inference_input_schema.require_global(inputs) + return _FakeSession(inference_input_schema=self._inference_input_schema) def close(self) -> None: self.closed = True class _FailingRuntime(_FakeRuntime): - def start_session(self, inputs: ModelInputs) -> InferenceSession: + def start_session(self, inputs: InferenceInput) -> InferenceSession: del inputs raise RuntimeError("start failed") class _FakeSession: - def __init__(self, *, model_input_schema: ModelInputSchema) -> None: - self._model_input_schema = model_input_schema + def __init__(self, *, inference_input_schema: InferenceInputSchema) -> None: + self._inference_input_schema = inference_input_schema self.step_index = 0 self.closed = False @@ -427,15 +462,15 @@ def next_step_request(self) -> StepRequest | None: return None return StepRequest( step_index=self.step_index, - model_input_schema=self._model_input_schema, + inference_input_schema=self._inference_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) + def step(self, inputs: InferenceInput) -> StepResult: + self._inference_input_schema.require_step(inputs) result = StepResult( step_index=self.step_index, output=f"chunk-{self.step_index}", @@ -449,7 +484,7 @@ def step(self, inputs: ModelInputs) -> StepResult: self.step_index += 1 return result - def reset(self, inputs: ModelInputs | None = None) -> None: + def reset(self, inputs: InferenceInput | None = None) -> None: del inputs self.step_index = 0 @@ -464,14 +499,19 @@ def __init__(self) -> None: def validate( self, *, - user_schema: UserInputSchema | None = None, - model_schema: ModelInputSchema | None = None, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, ) -> None: - super().validate(user_schema=user_schema, model_schema=model_schema) + super().validate( + canonical_schema=canonical_schema, + inference_input_schema=inference_input_schema, + ) self.validated = True class _OrderCheckingAdapter(_FakeAdapter): + canonical_input_schema = CanonicalInputSchema() + def __init__(self, *, mapping: _OrderCheckingMapping) -> None: self._mapping = mapping self.created_runtime_after_validate = False @@ -479,14 +519,18 @@ def __init__(self, *, mapping: _OrderCheckingMapping) -> None: 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) + return _FakeRuntime(inference_input_schema=self.inference_input_schema) class _FailingStartAdapter(_FakeAdapter): + canonical_input_schema = CanonicalInputSchema() + 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) + self.runtime = _FailingRuntime( + inference_input_schema=self.inference_input_schema + ) return self.runtime diff --git a/flashdreams/tests/test_runtime_canonical.py b/flashdreams/tests/test_runtime_canonical.py new file mode 100644 index 000000000..80637e15c --- /dev/null +++ b/flashdreams/tests/test_runtime_canonical.py @@ -0,0 +1,456 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the raw-input to canonical-modality layer. + +These cover the middle leg of ``raw input -> canonicalized input -> encoded +inference input``: applications consume canonical modalities, never raw device +events, so adding a device is a registration rather than an application change. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import pytest + +from flashdreams.runtime import ( + CONDITIONING_PROMPT, + DRIVER_COMMAND, + CanonicalInputs, + CanonicalModality, + DeviceConverterSchema, + InferenceInputSchema, + InputCanonicalizer, + InputField, + InputMappingSchema, + KeyboardToDriverCommand, + LatestEventToModality, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, + check_mapping_compatibility, +) + +pytestmark = pytest.mark.ci_cpu + +KEYBOARD_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability(event_type="key_down", payload_fields=frozenset({"key"})), + UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})), + ) +) +WHEEL_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="wheel_axis", payload_fields=frozenset({"axis", "value"}) + ), + ) +) +PROMPT_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="prompt_set", payload_fields=frozenset({"prompt"}) + ), + ) +) + +# Written once against the canonical modality. It names no key and no axis. +STEERING_MAPPING = InputMappingSchema( + name="driver-command-to-steering", + consumes=(DRIVER_COMMAND,), + produces_step=(InputField(name="steering"),), +) +STEERING_MODEL = InferenceInputSchema(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) + + +class WheelToDriverCommand: + """Minimal wheel converter standing in for a real evdev profile.""" + + def __init__(self, *, priority: int = 10) -> None: + self._steer = 0.0 + self._seen = False + self._schema = DeviceConverterSchema( + name="wheel-to-driver-command", + produces=DRIVER_COMMAND, + device_kind="wheel", + priority=priority, + consumes=( + UserInputCapability( + event_type="wheel_axis", + payload_fields=frozenset({"axis", "value"}), + ), + ), + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + self._steer = 0.0 + self._seen = False + + def convert( + self, user_inputs: UserInputs, window: TimeWindow + ) -> Mapping[str, Any] | None: + del window + for event in user_inputs.events: + if event.event_type == "wheel_axis" and event.payload["axis"] == "steer": + self._seen = True + self._steer = float(event.payload["value"]) + if not self._seen: + return None + return DRIVER_COMMAND.value( + { + "throttle": 0.0, + "brake": 0.0, + "steer": self._steer, + "stop": False, + "reverse": False, + } + ) + + +def _key(event_type: str, key: str, timestamp_s: float) -> UserInputEvent: + return UserInputEvent( + timestamp_s=timestamp_s, event_type=event_type, payload={"key": key} + ) + + +def _command(canonical: CanonicalInputs) -> Mapping[str, Any]: + assert DRIVER_COMMAND.name in canonical.per_step + return canonical.per_step[DRIVER_COMMAND.name] + + +# --- per-step conditioning ---------------------------------------------- + + +def test_keyboard_edges_become_canonical_driver_command() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs(events=(_key("key_down", "w", 0.1),)) + + canonical = canonicalizer.canonicalize( + inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert _command(canonical)["throttle"] == 1.0 + assert _command(canonical)["steer"] == 0.0 + assert canonical.metadata["canonical_sources"]["driver_command"] == "keyboard" + + +def test_key_aliases_are_normalized() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs(events=(_key("key_down", "ArrowLeft", 0.1),)) + + canonical = canonicalizer.canonicalize( + inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert _command(canonical)["steer"] == 1.0 + + +def test_held_key_still_emits_in_a_window_with_no_events() -> None: + """Edge-triggered HID must become level-triggered per-step conditioning.""" + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs(events=(_key("key_down", "w", 0.1),)) + canonicalizer.canonicalize(inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE) + + quiet = canonicalizer.canonicalize( + inputs, window=NEXT_WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert _command(quiet)["throttle"] == 1.0 + + +def test_key_release_returns_to_neutral() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs(events=(_key("key_down", "a", 0.1), _key("key_up", "a", 1.5))) + canonicalizer.canonicalize(inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE) + + released = canonicalizer.canonicalize( + inputs, window=NEXT_WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert _command(released)["steer"] == 0.0 + + +def test_reset_drops_device_state() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs(events=(_key("key_down", "w", 0.1),)) + canonicalizer.canonicalize(inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE) + + canonicalizer.reset() + after = canonicalizer.canonicalize( + UserInputs(), window=NEXT_WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert _command(after)["throttle"] == 0.0 + + +# --- global conditioning ------------------------------------------------ + + +def test_global_conditioning_is_emitted_only_when_it_changes() -> None: + """A quiet window must not look like a repeated update request.""" + canonicalizer = InputCanonicalizer( + [LatestEventToModality(modality=CONDITIONING_PROMPT, event_type="prompt_set")] + ) + inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.2, event_type="prompt_set", payload={"prompt": "rain"} + ), + ) + ) + + changed = canonicalizer.canonicalize( + inputs, window=WINDOW, source_schema=PROMPT_SOURCE + ) + quiet = canonicalizer.canonicalize( + inputs, window=NEXT_WINDOW, source_schema=PROMPT_SOURCE + ) + + assert changed.has_global_change + assert changed.global_conditioning["conditioning_prompt"]["prompt"] == "rain" + assert not quiet.has_global_change + + +def test_global_converter_rejects_a_per_step_modality() -> None: + with pytest.raises(ValueError, match="global conditioning"): + LatestEventToModality(modality=DRIVER_COMMAND, event_type="wheel_axis") + + +def test_global_and_per_step_land_in_separate_slots() -> None: + canonicalizer = InputCanonicalizer( + [ + KeyboardToDriverCommand(), + LatestEventToModality( + modality=CONDITIONING_PROMPT, event_type="prompt_set" + ), + ] + ) + source = UserInputSchema( + capabilities=KEYBOARD_SOURCE.capabilities + PROMPT_SOURCE.capabilities + ) + inputs = UserInputs( + events=( + _key("key_down", "w", 0.1), + UserInputEvent( + timestamp_s=0.2, event_type="prompt_set", payload={"prompt": "rain"} + ), + ) + ) + + canonical = canonicalizer.canonicalize(inputs, window=WINDOW, source_schema=source) + + assert set(canonical.per_step) == {"driver_command"} + assert set(canonical.global_conditioning) == {"conditioning_prompt"} + + +# --- device independence ------------------------------------------------ + + +def test_mapping_written_against_a_modality_accepts_a_keyboard() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + + compatibility = check_mapping_compatibility( + canonical_schema=canonicalizer.canonical_schema(KEYBOARD_SOURCE), + model_schema=STEERING_MODEL, + mapping_schema=STEERING_MAPPING, + ) + + assert compatibility.can_drive + + +def test_adding_a_device_needs_no_application_or_model_change() -> None: + """A wheel is one register() call; mapping and model schemas are untouched.""" + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + canonicalizer.register(WheelToDriverCommand()) + + compatibility = check_mapping_compatibility( + canonical_schema=canonicalizer.canonical_schema(WHEEL_SOURCE), + model_schema=STEERING_MODEL, + mapping_schema=STEERING_MAPPING, + ) + assert compatibility.can_drive + + canonical = canonicalizer.canonicalize( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.5, + event_type="wheel_axis", + payload={"axis": "steer", "value": -0.4}, + ), + ) + ), + window=WINDOW, + source_schema=WHEEL_SOURCE, + ) + assert _command(canonical)["steer"] == pytest.approx(-0.4) + + +def test_source_with_no_feedable_converter_supplies_no_modalities() -> None: + canonicalizer = InputCanonicalizer([WheelToDriverCommand()]) + + schema = canonicalizer.canonical_schema(KEYBOARD_SOURCE) + + assert schema.modalities == () + assert not schema.supports(DRIVER_COMMAND) + assert canonicalizer.unavailable_converters(KEYBOARD_SOURCE) + + +def test_highest_priority_device_wins_when_both_are_present() -> None: + canonicalizer = InputCanonicalizer( + [KeyboardToDriverCommand(), WheelToDriverCommand()] + ) + both = UserInputSchema( + capabilities=KEYBOARD_SOURCE.capabilities + WHEEL_SOURCE.capabilities + ) + + canonical = canonicalizer.canonicalize( + UserInputs( + events=( + _key("key_down", "a", 0.2), + UserInputEvent( + timestamp_s=0.5, + event_type="wheel_axis", + payload={"axis": "steer", "value": -0.4}, + ), + ) + ), + window=WINDOW, + source_schema=both, + ) + + assert canonical.metadata["canonical_sources"]["driver_command"] == "wheel" + assert _command(canonical)["steer"] == pytest.approx(-0.4) + + +def test_preempted_device_keeps_its_state_current() -> None: + """Keyboard state must not be stale when the wheel disappears.""" + canonicalizer = InputCanonicalizer( + [KeyboardToDriverCommand(), WheelToDriverCommand()] + ) + both = UserInputSchema( + capabilities=KEYBOARD_SOURCE.capabilities + WHEEL_SOURCE.capabilities + ) + inputs = UserInputs( + events=( + _key("key_down", "w", 0.2), + UserInputEvent( + timestamp_s=0.5, + event_type="wheel_axis", + payload={"axis": "steer", "value": -0.4}, + ), + ) + ) + preempted = canonicalizer.canonicalize(inputs, window=WINDOW, source_schema=both) + assert preempted.metadata["canonical_sources"]["driver_command"] == "wheel" + + keyboard_only = canonicalizer.canonicalize( + inputs, window=NEXT_WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert keyboard_only.metadata["canonical_sources"]["driver_command"] == "keyboard" + assert _command(keyboard_only)["throttle"] == 1.0 + + +# --- registry ----------------------------------------------------------- + + +def test_duplicate_converter_names_are_rejected() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + + with pytest.raises(ValueError, match="already registered"): + canonicalizer.register(KeyboardToDriverCommand()) + + +def test_converter_must_fill_the_declared_modality_payload() -> None: + modality = CanonicalModality( + name="steering_wheel", payload_fields=frozenset({"steer", "throttle"}) + ) + + with pytest.raises(ValueError, match="requires payload fields"): + modality.value({"steer": 0.0}) + + +def test_new_modality_is_a_registration_not_a_core_change() -> None: + pedals = CanonicalModality( + name="pedal_state", payload_fields=frozenset({"throttle"}) + ) + + class PedalsConverter: + schema = DeviceConverterSchema( + name="pedals", + produces=pedals, + device_kind="pedals", + consumes=( + UserInputCapability( + event_type="pedal_axis", + payload_fields=frozenset({"value"}), + ), + ), + ) + + def reset(self) -> None: + return None + + def convert( + self, user_inputs: UserInputs, window: TimeWindow + ) -> Mapping[str, Any] | None: + del window + if not user_inputs.events: + return None + return pedals.value( + {"throttle": float(user_inputs.events[-1].payload["value"])} + ) + + source = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="pedal_axis", payload_fields=frozenset({"value"}) + ), + ) + ) + canonicalizer = InputCanonicalizer([PedalsConverter()]) + + assert canonicalizer.canonical_schema(source).modalities == (pedals,) + canonical = canonicalizer.canonicalize( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.5, event_type="pedal_axis", payload={"value": 0.75} + ), + ) + ), + window=WINDOW, + source_schema=source, + ) + assert canonical.per_step["pedal_state"]["throttle"] == pytest.approx(0.75) + + +def test_replaying_the_same_windows_reproduces_the_same_canonical_inputs() -> None: + inputs = UserInputs(events=(_key("key_down", "w", 0.1), _key("key_down", "a", 1.2))) + + def run() -> list[dict[str, Any]]: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + return [ + dict( + _command( + canonicalizer.canonicalize( + inputs, window=window, source_schema=KEYBOARD_SOURCE + ) + ) + ) + for window in (WINDOW, NEXT_WINDOW) + ] + + assert run() == run() diff --git a/flashdreams/tests/test_runtime_input_mapping.py b/flashdreams/tests/test_runtime_input_mapping.py new file mode 100644 index 000000000..20deac189 --- /dev/null +++ b/flashdreams/tests/test_runtime_input_mapping.py @@ -0,0 +1,558 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for declarative input-mapping compatibility in the runtime API. + +These cover the T2/T3 contract: sources declare what user events they can +provide at payload granularity, models declare required and optional +initial/per-step inputs, and a mapping declares what it consumes and produces so +compatibility can be answered before expensive runtime initialization. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from flashdreams.runtime import ( + CONDITIONING_FRAME, + CONDITIONING_PROMPT, + DRIVER_COMMAND, + SESSION_START_ONLY, + CanonicalInputs, + CanonicalInputSchema, + IdentityInputMapping, + InferenceInput, + InferenceInputSchema, + InputField, + InputMappingSchema, + StepRequest, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, + check_mapping_compatibility, + check_mapping_set_compatibility, + combine_mapping_schemas, + undeclared_inference_input, +) + +pytestmark = pytest.mark.ci_cpu + +KEY_DOWN = UserInputCapability(event_type="key_down", payload_fields=frozenset({"key"})) +KEY_UP = UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})) +PROMPT_SET = UserInputCapability( + event_type="prompt_set", + semantic_type="text", + payload_fields=frozenset({"prompt"}), +) +FRAME_SET = UserInputCapability( + event_type="initial_frame_set", payload_fields=frozenset({"image"}) +) + +BROWSER_SOURCE = UserInputSchema( + capabilities=(KEY_DOWN, KEY_UP, PROMPT_SET, FRAME_SET), + description="browser webrtc client", +) + +CANONICAL_ALL = CanonicalInputSchema( + modalities=(CONDITIONING_PROMPT, CONDITIONING_FRAME, DRIVER_COMMAND) +) + +PROMPT_MAPPING = InputMappingSchema( + name="prompt", + consumes=(CONDITIONING_PROMPT,), + produces_global=(InputField(name="prompt", semantic_type="text"),), +) +FRAME_MAPPING = InputMappingSchema( + name="conditioning-frame", + consumes=(CONDITIONING_FRAME,), + produces_global=(InputField(name="global_conditioning_frame", required=False),), +) +STEERING_MAPPING = InputMappingSchema( + name="driver-command-to-steering", + consumes=(DRIVER_COMMAND,), + produces_step=(InputField(name="steering"),), +) + +DRIVING_MODEL = InferenceInputSchema( + global_fields=( + InputField(name="prompt", semantic_type="text", lifecycle="cache_init"), + InputField( + name="global_conditioning_frame", required=False, lifecycle="cache_init" + ), + ), + step_fields=(InputField(name="steering", lifecycle="step_input"),), +) + + +# --- user input events and windowing ------------------------------------ + + +def test_startup_values_are_represented_as_events() -> None: + inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, event_type="prompt_set", payload={"prompt": "drive"} + ), + UserInputEvent( + timestamp_s=0.5, event_type="key_down", payload={"key": "w"} + ), + ) + ) + + assert inputs.events[0].event_type == "prompt_set" + assert inputs.events[0].payload["prompt"] == "drive" + + +def test_windowing_is_half_open_and_deterministic() -> None: + inputs = UserInputs( + events=tuple( + UserInputEvent(timestamp_s=t, event_type="key_down", payload={"key": "w"}) + for t in (0.0, 0.5, 1.0, 1.5) + ) + ) + + windowed = inputs.window(TimeWindow(start_s=0.5, end_s=1.5)) + + assert [event.timestamp_s for event in windowed.events] == [0.5, 1.0] + + +def test_out_of_order_events_are_rejected() -> None: + with pytest.raises(ValueError, match="non-decreasing"): + UserInputs( + events=( + UserInputEvent(timestamp_s=1.0, event_type="key_down"), + UserInputEvent(timestamp_s=0.5, event_type="key_up"), + ) + ) + + +# --- user input schemas ------------------------------------------------- + + +def test_source_declares_capabilities_at_payload_granularity() -> None: + assert BROWSER_SOURCE.supports(KEY_DOWN) + assert not BROWSER_SOURCE.supports( + UserInputCapability( + event_type="key_down", payload_fields=frozenset({"key", "modifiers"}) + ) + ) + + +def test_bare_event_types_still_satisfy_payload_free_consumers() -> None: + """Coarse pre-capability schemas keep working against the finer query.""" + coarse = UserInputSchema(event_types=frozenset({"reset"})) + + assert coarse.supports(UserInputCapability(event_type="reset")) + assert not coarse.supports( + UserInputCapability(event_type="reset", payload_fields=frozenset({"reason"})) + ) + assert coarse.supports_event_types({"reset"}) + + +def test_capabilities_widen_declared_event_types() -> None: + assert "key_down" in BROWSER_SOURCE.declared_event_types() + assert BROWSER_SOURCE.supports_event_types({"key_down", "prompt_set"}) + + +def test_semantic_type_mismatch_blocks_capability_match() -> None: + source = UserInputSchema( + capabilities=( + UserInputCapability(event_type="prompt_set", semantic_type="embedding"), + ) + ) + + assert not source.supports( + UserInputCapability(event_type="prompt_set", semantic_type="text") + ) + + +def test_event_validation_reports_missing_payload_fields() -> None: + event = UserInputEvent(timestamp_s=0.0, event_type="key_down", payload={}) + + with pytest.raises(ValueError, match="missing required"): + BROWSER_SOURCE.validate_event(event) + + +def test_event_validation_rejects_undeclared_event_type() -> None: + event = UserInputEvent(timestamp_s=0.0, event_type="wheel_axis") + + with pytest.raises(ValueError, match="does not provide event type"): + BROWSER_SOURCE.validate_event(event) + + +# --- model input schemas ------------------------------------------------ + + +def test_model_declares_required_and_optional_fields_per_phase() -> None: + required = DRIVING_MODEL.required_fields() + optional = DRIVING_MODEL.optional_fields() + + assert {(phase, f.name) for phase, f in required} == { + ("global", "prompt"), + ("step", "steering"), + } + assert {(phase, f.name) for phase, f in optional} == { + ("global", "global_conditioning_frame") + } + + +def test_required_fields_can_be_filtered_by_phase() -> None: + step_only = DRIVING_MODEL.required_fields("step") + + assert [f.name for _, f in step_only] == ["steering"] + + +def test_field_lookup_is_phase_scoped() -> None: + assert DRIVING_MODEL.field_for(name="prompt", phase="global") is not None + assert DRIVING_MODEL.field_for(name="prompt", phase="step") is None + + +def test_invalid_phase_is_rejected() -> None: + bad_phase: Any = "final" + + with pytest.raises(ValueError, match="phase must be"): + DRIVING_MODEL.fields_for(bad_phase) + + +def test_inference_input_expose_payload_per_phase() -> None: + inputs = InferenceInput( + global_conditioning={"prompt": "drive"}, step={"steering": 0.25} + ) + + assert inputs.for_phase("global")["prompt"] == "drive" + assert inputs.for_phase("step")["steering"] == 0.25 + + +def test_lifecycle_and_update_policy_are_queryable_metadata() -> None: + field = InputField( + name="prompt", + update_policy="step_boundary", + lifecycle="cache_init", + metadata={"coordinates": "opencv_c2w"}, + ) + + assert field.update_policy == "step_boundary" + assert field.lifecycle == "cache_init" + assert field.metadata["coordinates"] == "opencv_c2w" + + +def test_metadata_is_excluded_from_field_equality() -> None: + plain = InputField(name="prompt") + annotated = InputField(name="prompt", metadata={"note": "hint"}) + + assert plain == annotated + + +# --- mapping compatibility ---------------------------------------------- + + +def test_compatible_source_model_and_mapping_can_drive() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CANONICAL_ALL, + model_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, FRAME_MAPPING, STEERING_MAPPING), + ) + + assert compatibility.can_drive + assert {(p, f.name) for p, f in compatibility.satisfied_required_model_fields} == { + ("global", "prompt"), + ("step", "steering"), + } + assert {(p, f.name) for p, f in compatibility.available_optional_model_fields} == { + ("global", "global_conditioning_frame") + } + + +def test_missing_required_model_field_blocks_the_run() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CANONICAL_ALL, + model_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING,), + ) + + assert not compatibility.can_drive + assert [f.name for _, f in compatibility.missing_required_model_fields] == [ + "steering" + ] + + +def test_missing_source_capability_is_reported_when_it_blocks() -> None: + prompt_only = CanonicalInputSchema(modalities=(CONDITIONING_PROMPT,)) + + compatibility = check_mapping_set_compatibility( + canonical_schema=prompt_only, + model_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), + ) + + assert not compatibility.can_drive + assert compatibility.unavailable_mapping_names == ("driver-command-to-steering",) + assert {m.name for m in compatibility.missing_modalities} == {"driver_command"} + + +def test_unfeedable_optional_mapping_degrades_instead_of_vetoing() -> None: + """Losing a mapping that fed only optional fields must not block the run.""" + no_frame_source = CanonicalInputSchema( + modalities=(CONDITIONING_PROMPT, DRIVER_COMMAND) + ) + + compatibility = check_mapping_set_compatibility( + canonical_schema=no_frame_source, + model_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, FRAME_MAPPING, STEERING_MAPPING), + ) + + assert compatibility.can_drive + assert compatibility.unavailable_mapping_names == ("conditioning-frame",) + # The dropped mapping's field must not be advertised as available. + assert compatibility.available_optional_model_fields == () + + +def test_optional_field_needs_mapping_support_to_be_available() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CANONICAL_ALL, + model_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), + ) + + assert compatibility.can_drive + assert compatibility.available_optional_model_fields == () + + +def test_lifecycle_disagreement_blocks_a_field_match() -> None: + model = InferenceInputSchema( + global_fields=(InputField(name="prompt", lifecycle="rollout_binding"),) + ) + mapping = InputMappingSchema( + name="prompt", + consumes=(CONDITIONING_PROMPT,), + produces_global=(InputField(name="prompt", lifecycle="cache_init"),), + ) + + compatibility = check_mapping_compatibility( + canonical_schema=CANONICAL_ALL, model_schema=model, mapping_schema=mapping + ) + + assert not compatibility.can_drive + + +def test_unspecified_lifecycle_stays_permissive() -> None: + model = InferenceInputSchema(global_fields=(InputField(name="prompt"),)) + + compatibility = check_mapping_compatibility( + canonical_schema=CANONICAL_ALL, + model_schema=model, + mapping_schema=PROMPT_MAPPING, + ) + + assert compatibility.can_drive + + +def test_raise_if_incompatible_names_both_failure_kinds() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CanonicalInputSchema(modalities=(CONDITIONING_PROMPT,)), + model_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), + ) + + with pytest.raises(ValueError) as excinfo: + compatibility.raise_if_incompatible() + + message = str(excinfo.value) + assert "missing canonical modalities" in message + assert "missing required model inputs" in message + + +def test_raise_if_incompatible_is_a_no_op_when_compatible() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CANONICAL_ALL, + model_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, FRAME_MAPPING, STEERING_MAPPING), + ) + + compatibility.raise_if_incompatible() + + +def test_check_mapping_compatibility_rejects_a_non_schema() -> None: + not_a_schema: Any = object() + + with pytest.raises(TypeError, match="InputMappingSchema"): + check_mapping_compatibility( + canonical_schema=CANONICAL_ALL, + model_schema=DRIVING_MODEL, + mapping_schema=not_a_schema, + ) + + +# --- mapping schema composition ----------------------------------------- + + +def test_combining_mappings_unions_their_surfaces() -> None: + combined = combine_mapping_schemas((PROMPT_MAPPING, STEERING_MAPPING)) + + assert {m.name for m in combined.consumes} == { + "conditioning_prompt", + "driver_command", + } + assert [f.name for f in combined.produces_global] == ["prompt"] + assert [f.name for f in combined.produces_step] == ["steering"] + + +def test_duplicate_declarations_collapse_and_merge_metadata() -> None: + first = InputMappingSchema( + name="a", + produces_global=(InputField(name="prompt", metadata={"source": "a"}),), + ) + second = InputMappingSchema( + name="b", + produces_global=( + InputField(name="prompt", metadata={"source": "b", "extra": "kept"}), + ), + ) + + combined = combine_mapping_schemas((first, second)) + + assert len(combined.produces_global) == 1 + metadata = combined.produces_global[0].metadata + assert metadata["source"] == "a" + assert metadata["extra"] == "kept" + + +def test_combine_rejects_non_schema_entries() -> None: + not_a_schema: Any = object() + + with pytest.raises(TypeError, match="InputMappingSchema"): + combine_mapping_schemas((PROMPT_MAPPING, not_a_schema)) + + +# --- declaration drift -------------------------------------------------- + + +def test_undeclared_inference_input_catches_schema_drift() -> None: + produced = InferenceInput( + global_conditioning={"prompt": "drive"}, step={"steering": 0.0} + ) + + undeclared = undeclared_inference_input(produced, PROMPT_MAPPING) + + assert undeclared == (("step", "steering"),) + + +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} + ) + + assert undeclared_inference_input(produced, combined) == () + + +# --- interoperability with the T1 envelope ------------------------------ + + +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} + ) + + mapped = mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=fixed, + request=StepRequest(step_index=0), + ) + + assert mapped.step["steering"] == 0.0 + + +def test_empty_mapping_set_cannot_satisfy_a_required_field() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CANONICAL_ALL, + model_schema=DRIVING_MODEL, + mapping_schemas=(), + ) + + assert not compatibility.can_drive + assert len(compatibility.missing_required_model_fields) == 2 + + +def test_model_with_no_requirements_is_always_drivable() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CanonicalInputSchema(), + model_schema=InferenceInputSchema(), + mapping_schemas=(), + ) + + assert compatibility.can_drive + + +# --- global conditioning updates vs reset ------------------------------- + + +def test_empty_global_slot_requests_no_update() -> None: + steady_state = InferenceInput(step={"steering": 0.25}) + + assert not steady_state.requests_global_update + + +def test_non_empty_global_slot_mid_rollout_is_an_update_request() -> None: + """Changing weather mid-run updates conditioning; it is not a reset.""" + updated = InferenceInput(step={"steering": 0.0}).with_global_update( + {"prompt": "heavy rain"} + ) + + assert updated.requests_global_update + assert updated.global_conditioning["prompt"] == "heavy rain" + assert updated.step["steering"] == 0.0 + + +def test_with_step_carries_the_global_slot_through() -> None: + started = InferenceInput(global_conditioning={"prompt": "drive"}) + + stepped = started.with_step({"steering": 0.5}) + + assert stepped.global_conditioning["prompt"] == "drive" + + +def test_without_global_update_clears_the_request() -> None: + started = InferenceInput( + global_conditioning={"prompt": "drive"}, step={"steering": 0.5} + ) + + steady_state = started.without_global_update() + + assert not steady_state.requests_global_update + assert steady_state.step["steering"] == 0.5 + + +def test_model_can_declare_conditioning_it_cannot_swap_mid_rollout() -> None: + schema = InferenceInputSchema( + global_fields=( + InputField(name="prompt", update_policy="step_boundary"), + InputField(name="scene_id", update_policy=SESSION_START_ONLY), + ) + ) + update = InferenceInput( + global_conditioning={"prompt": "heavy rain", "scene_id": "town_02"} + ) + + assert schema.unsupported_global_updates(update) == ("scene_id",) + + +def test_permissive_when_no_update_policy_is_declared() -> None: + schema = InferenceInputSchema(global_fields=(InputField(name="prompt"),)) + update = InferenceInput(global_conditioning={"prompt": "heavy rain"}) + + assert schema.unsupported_global_updates(update) == () + + +def test_undeclared_global_values_are_left_to_the_adapter() -> None: + schema = InferenceInputSchema(global_fields=(InputField(name="prompt"),)) + update = InferenceInput(global_conditioning={"mystery": 1}) + + assert schema.unsupported_global_updates(update) == () From 7a8c4d87826ba72544cad3dbeb6867cad29f82a1 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Tue, 4 Aug 2026 14:35:35 -0700 Subject: [PATCH 4/7] doc update --- ...inference_runtime_inputs_implementation.md | 530 ++++++------------ flashdreams/flashdreams/runtime/__init__.py | 4 +- flashdreams/flashdreams/runtime/mapping.py | 20 +- flashdreams/tests/test_runtime_canonical.py | 4 +- .../tests/test_runtime_input_mapping.py | 32 +- 5 files changed, 215 insertions(+), 375 deletions(-) diff --git a/docs/inference_runtime_inputs_implementation.md b/docs/inference_runtime_inputs_implementation.md index 465319b8c..33bc9ad7d 100644 --- a/docs/inference_runtime_inputs_implementation.md +++ b/docs/inference_runtime_inputs_implementation.md @@ -3,414 +3,252 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# Inference Runtime Inputs API Implementation Notes +# Inference Runtime Inputs Implementation Notes -This note documents the current T2/T3 implementation of the inference runtime -input contracts. It is meant as an evaluator's guide: what exists, how the -pieces fit together, what the compatibility query answers, and what is -intentionally still outside this layer. +This note documents the input layers of the experimental runtime API: what +exists, how the pieces fit together, what the compatibility query answers, and +what is intentionally still outside this layer. -Implementation lives in `flashdreams.inference`: +Implementation lives in `flashdreams.runtime`: -- `flashdreams/flashdreams/inference/__init__.py` -- `flashdreams/flashdreams/inference/inputs.py` -- `flashdreams/tests/test_inference_inputs.py` +- `flashdreams/flashdreams/runtime/inputs.py` — the input types and schemas +- `flashdreams/flashdreams/runtime/canonical.py` — raw device to canonical + modality conversion +- `flashdreams/flashdreams/runtime/mapping.py` — canonical to encoded mapping + and compatibility +- `flashdreams/tests/test_runtime_canonical.py` +- `flashdreams/tests/test_runtime_input_mapping.py` +- `flashdreams/tests/test_inference_runtime_api.py` — the T1 envelope tests, + including a reference loop that exercises all three layers -The current supported-model input inventory that informed this revision is in +The supported-model input inventory that informed this work is in `docs/inference_runtime_supported_inputs_inventory.md`. -## What This Implements - -The implementation covers the T2/T3 contract from -`docs/inference_runtime_api_design.md`: - -- `UserInputs` are represented as timestamped events. -- `UserInputSchema` describes what an app, transport, replay trace, or - benchmark source can provide. -- `ModelInputs` are semantic model-facing payloads split into initial and - per-step inputs. -- `ModelInputSchema` describes what a model or session requires or optionally - accepts. -- `InputMapper` is the conversion contract between user events and model inputs. -- Schema objects support open-ended `metadata` maps for lightweight hints that - future model adapters can expose without changing the core API. -- `check_mapping_compatibility()` answers whether a selected source can drive a - selected model through a selected mapper before expensive runtime setup. -- `check_mapping_set_compatibility()` answers the same question for a composed - set of mapper schemas, such as prompt plus first-frame plus live-control - mappings. - -This does not migrate LingBot, OmniDreams, WebRTC, CLI runners, or the standard -runtime/session loop. Those are T4+ and migration tasks. - -## User Inputs - -`UserInputEvent` is the primary user-input representation. Every user-facing -input is modeled as an event in session time, including static startup values. - -Examples: - -- `prompt_set` at `timestamp_s=0.0` -- `initial_frame_set` at `timestamp_s=0.0` -- `scene_selected` at `timestamp_s=0.0` -- `key_down` / `key_up` during a session -- `controller_axis` during a session -- `camera_pose` events in a replay trace - -`UserInputTrace` stores events in deterministic timestamp order and can slice a -`UserInputWindow` for one runtime step. Events with equal timestamps preserve -their original order. A `UserInputWindow` enforces its own invariant: every -event it holds must fall inside `[start_s, end_s]`, so a directly constructed -window cannot silently disagree with its bounds. +## The Three Layers -```python -from flashdreams.inference import UserInputEvent, UserInputTrace - -trace = UserInputTrace.from_events( - [ - UserInputEvent( - timestamp_s=0.0, - kind="prompt_set", - payload={"prompt": "drive forward"}, - ), - UserInputEvent( - timestamp_s=0.5, - kind="key_down", - payload={"key": "w"}, - ), - ] -) - -step_window = trace.window(start_s=0.0, end_s=1.0) +```text +UserInputs ──InputCanonicalizer──▶ CanonicalInputs ──InputMapping──▶ InferenceInput + raw canonicalized encoded +(device events) (device-independent) (what the session gets) ``` -Snapshots are intentionally not primary inputs. A mapper or runtime can derive -snapshots from a trace/window when a model wants snapshot-style controls. +| Layer | Type | Owner | Example | +| --- | --- | --- | --- | +| raw | `UserInputs` / `UserInputEvent` | transport, replay loader, benchmark driver | `key_down {"key": "w"}`, wheel axis reading | +| canonicalized | `CanonicalInputs` | device converters registered on `InputCanonicalizer` | `driver_command {throttle, brake, steer, ...}` | +| encoded | `InferenceInput` | the selected `InputMapping` | whatever the model's session consumes | -## User Input Schema +Applications and mappings consume `CanonicalInputs`. They never read raw device +events: `InputMapping.map_step_inputs` takes `canonical_inputs`, not +`user_inputs`, so this is enforced by the signature rather than by convention. +Adding a keyboard, gamepad, or wheel is an `InputCanonicalizer.register` call +that touches no application, mapping, or model code. -`UserInputSchema` is lightweight metadata for the source side. It answers: +## Conditioning Slots -- Which event kinds can this source provide? -- Which payload fields are present on those events? -- Is this source live, replayed, fixed, or otherwise identified? +Both the canonical and encoded layers split into two slots, and the split means +the same thing at each: -It does not assign model semantics. For example, `key_down` does not mean -steering or camera motion until a mapper says so. +- **global conditioning** — conditions the whole rollout: prompt, conditioning + frame, scene. Normally supplied at session start. +- **per-step conditioning** — needed to generate the next chunk or frame: + steering, HD map frames, camera trajectory. -```python -from flashdreams.inference import UserInputCapability, UserInputSchema - -browser_schema = UserInputSchema( - name="browser", - source_kind="live", - metadata={"transport": "webrtc"}, - capabilities=( - UserInputCapability( - event_kind="prompt_set", - payload_kind="text", - payload_fields=frozenset({"prompt"}), - ), - UserInputCapability( - event_kind="key_down", - payload_fields=frozenset({"key"}), - ), - UserInputCapability( - event_kind="key_up", - payload_fields=frozenset({"key"}), - ), - ), -) -``` - -The schema can validate that a source event has an event kind and payload fields -the source claims to provide. It does not validate tensor shape, image format, -camera coordinate system, or model-specific units. +`InputPhase` is `Literal["global", "step"]`. The axis names *which slot*, not +*when the value may arrive* — see the next section. -Capabilities and schemas also carry optional `metadata`. This is for query-time -hints such as source type, UI widget, file suffixes, units, coordinate frame, or -schema URI. Metadata is intentionally not part of compatibility matching. -Because metadata is an open-ended pass-through, its keys are validated (they -must be non-empty strings) but never rewritten, so a key round-trips unchanged. +## Global Conditioning Updates Are Not Resets -## Model Inputs - -`ModelInputs` is the model-facing payload container. It separates values needed -to start or reset a rollout from values needed for one generated step/chunk. +A non-empty global slot on a mid-rollout `InferenceInput` is an **update +request**. The session should apply it when the model supports doing so. +Resetting rollout state is a separate, explicit `InferenceSession.reset()` call. +The motivating case is changing prompt and conditioning frame mid-run to change +the weather in an Omnidreams rollout. ```python -from flashdreams.inference import ModelInputs +from flashdreams.runtime import InferenceInput -inputs = ModelInputs( - initial={"prompt": "drive forward", "first_frame": first_frame}, - step={"steering": 0.25}, -) +steady_state = InferenceInput(step={"steering": 0.25}) +assert not steady_state.requests_global_update + +changed_weather = steady_state.with_global_update({"prompt": "heavy rain"}) +assert changed_weather.requests_global_update ``` -`ModelInputSchema` declares the semantic fields a model or session expects. It -uses names like `prompt`, `first_frame`, `steering`, `camera_trajectory`, or -`hdmap_frames`, rather than generic modality-only keys. +Because `with_step()` carries the global slot through unchanged, use +`without_global_update()` for the steady-state case; otherwise every step looks +like an update request. + +Whether a value can actually be swapped mid-rollout is declared per field: ```python -from flashdreams.inference import ModelInputField, ModelInputSchema - -model_schema = ModelInputSchema( - name="driving-model", - fields=( - ModelInputField( - name="prompt", - phase="initial", - required=True, - payload_kind="text", - update_policy="step_boundary", - lifecycle="cache_init", - ), - ModelInputField( - name="steering", - phase="step", - required=True, - lifecycle="step_input", - ), - ModelInputField( - name="first_frame", - phase="initial", - required=False, - lifecycle="cache_init", - ), - ModelInputField( - name="camera_trajectory_c2w", - phase="initial", - required=False, - payload_kind="c2w_sequence", - lifecycle="rollout_binding", - metadata={"shape": "[F,4,4]", "coordinates": "opencv_c2w"}, - ), - ), - metadata={"model_family": "example-driving-model"}, +from flashdreams.runtime import SESSION_START_ONLY, InferenceInputSchema, InputField + +schema = InferenceInputSchema( + global_fields=( + InputField(name="prompt", update_policy="step_boundary"), + InputField(name="scene_id", update_policy=SESSION_START_ONLY), + ) ) +schema.unsupported_global_updates( + InferenceInput(global_conditioning={"prompt": "heavy rain", "scene_id": "town_02"}) +) +# ("scene_id",) ``` -`update_policy` is deliberately plain metadata. It lets a model advertise facts -such as "prompt updates can happen at step boundaries" without making this -schema layer responsible for implementing or deeply validating that behavior. - -`lifecycle` is also plain metadata. It lets a model distinguish initial values -used at different adapter moments, such as `runtime_config`, `cache_init`, -`rollout_binding`, `step_input`, or `session_update`. Compatibility requires -lifecycle agreement only when both the model field and mapper output specify a -lifecycle; otherwise simple schemas remain permissive. +`SESSION_START_ONLY` is the one reserved `update_policy` token. Everything else +in that vocabulary, and all of `lifecycle`, is open and adapter-owned; this layer +only carries it as queryable metadata. -Model input names, payload kinds, lifecycle labels, and metadata are -open-ended. They are not a FlashDreams-wide enum. A SANA-WM-like adapter can -declare fields such as `camera_trajectory_c2w`, `camera_intrinsics_vec4`, -`stage1_sampling`, or `streaming_chunking`; another model can declare different -semantic names. The adapter and mapper own the deep interpretation. +The canonical layer mirrors this. Per-step converters emit every window, because +per-step conditioning is level-triggered: a key held across a step emits no +events but still means full throttle. Global-conditioning converters emit *only* +when a value changed, so a quiet window does not look like a repeated update +request. -## Input Mappers +## Raw Inputs -`InputMapper` is a protocol with three responsibilities: +`UserInputEvent` carries `timestamp_s`, `event_type`, `payload`, `source`, and +`source_event_id`. `UserInputs` holds an ordered batch plus a `snapshot` and +`metadata`, and slices to a half-open `TimeWindow`: -- expose an `InputMapperSchema`; -- build initial `ModelInputs` from a `UserInputTrace`; -- build per-step `ModelInputs` from a `UserInputWindow`. +```python +from flashdreams.runtime import TimeWindow, UserInputEvent, UserInputs + +inputs = UserInputs( + events=( + UserInputEvent(timestamp_s=0.0, event_type="prompt_set", + payload={"prompt": "drive forward"}), + UserInputEvent(timestamp_s=0.5, event_type="key_down", payload={"key": "w"}), + ) +) +step_window = inputs.window(TimeWindow(start_s=0.0, end_s=1.0)) +``` -`InputMapperSchema` declares the mapper's compatibility surface: +`UserInputSchema` describes what a transport, replay trace, or benchmark driver +can provide. `event_types` declares only that an event type exists; +`UserInputCapability` additionally pins the payload fields it carries, so a +converter can require `key_down` events that actually have a `key`. A bare +`event_types` entry still satisfies any consumer needing no specific payload +fields, so schemas written before capabilities existed keep working. -- `consumes`: user event capabilities required by the mapper; -- `produces`: model input fields the mapper can produce. +## Canonical Modalities -Example: keyboard events can be mapped into steering for one model, camera -trajectory for another model, or ignored entirely. That meaning is owned by the -mapper, not by `UserInputEvent`. +A `CanonicalModality` is a device-independent input: a name, a phase, and the +payload fields it guarantees. Converters implement `DeviceConverter`, declaring +what raw capabilities they consume and which modality they produce. ```python -from flashdreams.inference import ( - InputMapperSchema, - ModelInputField, - UserInputCapability, +from flashdreams.runtime import ( + DRIVER_COMMAND, InputCanonicalizer, KeyboardToDriverCommand, TimeWindow, ) -keyboard_to_steering_schema = InputMapperSchema( - name="keyboard-to-steering", - consumes=( - UserInputCapability( - event_kind="key_down", - payload_fields=frozenset({"key"}), - ), - UserInputCapability( - event_kind="key_up", - payload_fields=frozenset({"key"}), - ), - ), - produces=( - ModelInputField(name="steering", phase="step"), - ), -) -``` - -For fixed runs that already have model-facing inputs, `StaticInputMapper` -provides the no-live-controls path. It consumes no user events and returns the -configured initial/per-step `ModelInputs`. +canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) +canonicalizer.register(WheelToDriverCommand()) # a wheel is one call -```python -from flashdreams.inference import ModelInputs, StaticInputMapper - -static_mapper = StaticInputMapper.from_inputs( - inputs=ModelInputs( - initial={"prompt": "fixed prompt"}, - step={"camera_trajectory": camera_poses}, - ), - name="fixed-scenario", +canonical = canonicalizer.canonicalize( + user_inputs, window=TimeWindow(start_s=0.0, end_s=1.0), source_schema=browser ) +canonical.per_step["driver_command"]["throttle"] ``` -A mapper schema is a hand-written declaration, so it can drift from what -`build_initial_inputs` / `build_step_inputs` actually return. -`undeclared_model_inputs()` reports payload keys a mapper produced but did not -declare, which keeps mapper tests honest about the compatibility surface. +Shipped modalities are `DRIVER_COMMAND` (step), `CONDITIONING_PROMPT` and +`CONDITIONING_FRAME` (global). `KeyboardToDriverCommand` reuses +`KeyboardState`/`normalize_key` from `flashdreams.serving.realtime.input` and +mirrors the semantics the Omnidreams interactive-drive keyboard backend already +has. `LatestEventToModality` is the general global-conditioning converter and +rejects a per-step modality at construction. -Mapper schemas can also be combined for compatibility checks. This matches the -current supported model inventory: a run may select one mapping for prompt -events, another for first-frame events, and another for keyboard or controller -events. +Converters are stateful, so feed windows in session order and call +`InputCanonicalizer.reset()` at a rollout boundary. Replaying the same window +sequence reproduces the same `CanonicalInputs`. -```python -from flashdreams.inference import check_mapping_set_compatibility +When several devices produce the same modality, the highest-priority one that +returned a value wins; `CanonicalInputs.metadata["canonical_sources"]` records +which device supplied each. Every feedable converter still sees each window, so +a preempted device's state stays current and unplugging the higher-priority +device does not resume from stale state. -compatibility = check_mapping_set_compatibility( - source_schema=browser_schema, - model_schema=lingbot_schema, - mapper_schemas=( - prompt_mapper.schema, - first_frame_mapper.schema, - keyboard_to_camera_mapper.schema, - ), -) -``` +## Mapping And Compatibility -## Compatibility Query +`InputMapping` is the canonical-to-encoded boundary. `InputMappingSchema` is its +declarative surface: `consumes` names canonical modalities; `produces_global` +and `produces_step` name the `InferenceInput` fields it can build. -`check_mapping_compatibility()` is the central query for T2/T3. Given a source -schema, model schema, and mapper schema, it reports: - -- whether the source can drive the model through this mapper; -- source capabilities the mapper needs but the source lacks; -- required model fields the mapper cannot produce; -- required model fields that are satisfied; -- optional model fields that can be enabled; -- mappers dropped because the source cannot feed them. +`InputMapping.validate()` raises, which fails a run late and cannot say *which* +optional model input a source would enable or *which* missing modality makes a +required one unreachable. `check_mapping_compatibility` answers those before +expensive runtime initialization: ```python -from flashdreams.inference import check_mapping_compatibility +from flashdreams.runtime import check_mapping_set_compatibility -compatibility = check_mapping_compatibility( - source_schema=browser_schema, - model_schema=model_schema, - mapper_schema=keyboard_to_steering_schema, +compatibility = check_mapping_set_compatibility( + canonical_schema=canonicalizer.canonical_schema(browser), + inference_input_schema=adapter.inference_input_schema, + mapping_schemas=(prompt_mapping, frame_mapping, steering_mapping), ) - if not compatibility.can_drive: compatibility.raise_if_incompatible() ``` -This is an early compatibility check. It is intended to fail before expensive -runtime initialization when the mismatch is obvious. It is not a guarantee that -the model run will succeed. - -Use `check_mapping_set_compatibility()` when the selected mapping is composed -from multiple mapper schemas. It returns the same `MappingCompatibility` report -for the composed mapping. - -Compatibility is evaluated per mapper rather than over a flattened bag of -capabilities, so each mapper keeps its own `consumes`/`produces` link. A mapper -the source cannot feed is dropped and reported in `unavailable_mapper_schemas`; -it costs only the model inputs that mapper produced. This means: - -- a dropped mapper that produced only optional fields degrades the run instead - of vetoing it, and those fields are correctly absent from - `available_optional_model_fields`; -- a dropped mapper that was the only producer of a required field still blocks, - and its unmet source capabilities are reported in - `missing_source_capabilities`; -- `satisfied_required_model_fields` and `available_optional_model_fields` name - only fields this source can really produce, not every field some mapper - declared. +`MappingCompatibility` reports `missing_modalities`, +`missing_required_model_fields`, `satisfied_required_model_fields`, +`available_optional_model_fields`, and `unavailable_mapping_schemas`. + +Compatibility is evaluated per mapping rather than over a flattened bag, so each +mapping keeps its own consumes/produces link. A mapping the source cannot feed +is dropped and reported, costing only the inputs it produced. So a dropped +mapping that fed only optional fields degrades the run instead of vetoing it, +and those fields are correctly absent from `available_optional_model_fields`; a +dropped mapping that was the only producer of a required field still blocks. + +Because a mapping consumes modalities rather than raw events, one mapping +written against `driver_command` works for a keyboard, a wheel, or any device +registered later, with no change to the mapping or the model schema. + +`undeclared_inference_inputs()` reports payload keys a mapping produced but did +not declare, which keeps hand-written schemas honest as the code drifts. ## What This Does Not Validate -The schemas intentionally avoid becoming a rich type system. The following -remain the responsibility of the model adapter, runtime, session, or mapper -implementation: +The schemas intentionally avoid becoming a rich type system. These remain the +responsibility of the model adapter, runtime, session, or mapping: -- tensor shape and dtype; -- image decode details; -- camera coordinate systems; -- pose and timestamp units; +- tensor shape and dtype, image decode details; +- camera coordinate systems, pose and timestamp units; - prompt-embedding swap mechanics; -- reset semantics; -- whether a specific model can actually apply an update policy at runtime; +- whether a model can actually apply a declared update policy at runtime; - deep validation of scene, HD map, or actor-state data. -The schema layer should be enough to answer "can this source plausibly drive -this model through this mapper?" It should not replace model-owned validation. - -## Current Tests - -The focused CPU tests are in `flashdreams/tests/test_inference_inputs.py`. -They cover: - -- deterministic event ordering and trace windowing; -- startup values represented as events; -- source capability declaration and basic event validation; -- required and optional model input declarations; -- prompt update metadata on a model field; -- lifecycle metadata on model fields and mapper outputs; -- open-ended schema metadata that stays queryable but does not constrain - compatibility; -- missing required model inputs; -- missing source capabilities; -- optional model inputs becoming available only when mapper support exists; -- composed mapper-set compatibility; -- graceful degradation when an optional mapper cannot be fed, and blocking when - a required one cannot; -- metadata merging when duplicate declarations collapse during mapper-set - combination; -- event hashability, window bound enforcement, and metadata key round-tripping; -- mappers producing only model inputs their schema declares; -- SANA-WM-like model-specific inputs without importing the SANA integration; -- fake prompt, initial-frame, keyboard-to-steering, and camera-trajectory - mappers; -- `StaticInputMapper` for fixed model-input scenarios. - -Targeted validation command: +The layer answers "can this source plausibly drive this model through this +mapping?" It does not replace model-owned validation. + +## Open Questions + +Tracked against the runtime API discussion, not yet settled: + +- **Alternative valid input combinations.** `InferenceInputSchema` has one flat + required set, so "accepts `{prompt}` OR `{prompt, conditioning_frame}`" cannot + be expressed. `MappingCompatibility.missing_required_model_fields` assumes a + single required set too. +- **Who owns encoding.** Whether text/image embedding happens app-side (a + `publish_input(modality) -> Tensor` handler) or model-side. This does not + change what `InputMappingSchema` declares — a mapping consumes canonical + modalities either way — but it decides whether `InferenceInput` payloads hold + tensors or canonical values. +- **`step()` returning a future**, for models with a dependency on their own + output. `InferenceSession.step()` is currently synchronous. +- **Loop ownership** — whether the application or the runtime/session drives the + main event loop, and whether inputs are queued and batched. + +## Validation ```bash -.venv/bin/pytest flashdreams/tests/test_inference_inputs.py -q -.venv/bin/ty check flashdreams/flashdreams/inference flashdreams/tests/test_inference_inputs.py -.venv/bin/python -m py_compile \ - flashdreams/flashdreams/inference/__init__.py \ - flashdreams/flashdreams/inference/inputs.py \ - flashdreams/tests/test_inference_inputs.py +.venv/bin/pytest flashdreams/tests/test_runtime_canonical.py \ + flashdreams/tests/test_runtime_input_mapping.py \ + flashdreams/tests/test_inference_runtime_api.py -q +.venv/bin/ty check flashdreams/flashdreams/runtime ``` -At the time this note was written, these targeted checks passed. - -## Evaluation Checklist - -Use this checklist to judge whether the implementation matches the T2/T3 design: - -- Are all `UserInputs` represented as timestamped events? -- Can a source declare what event capabilities it provides? -- Can a model declare required and optional initial/per-step inputs? -- Are model input names semantic rather than modality-only? -- Can new models add model-specific field names and metadata without changing - core dataclasses? -- Can a mapper declare what it consumes and what it produces? -- Can multiple mapper schemas be checked as one selected mapping surface? -- Does compatibility checking report both missing source capabilities and - missing required model inputs? -- Are optional model inputs reported separately from required inputs? -- Can model fields distinguish cache initialization, rollout binding, per-step - inputs, and active-session update support without deep tensor validation? -- Is deep model/tensor validation kept out of the lightweight schema layer? -- Is fixed input/model-input replay possible without live user events? +At the time of writing these pass: 87 tests, and `ty` is clean. diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index f1f6d0f48..b31582ec5 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -49,7 +49,7 @@ check_mapping_compatibility, check_mapping_set_compatibility, combine_mapping_schemas, - undeclared_inference_input, + undeclared_inference_inputs, ) from flashdreams.runtime.metrics import ( InMemoryMetricsRecorder, @@ -102,7 +102,7 @@ "StepRequest", "StepResult", "TimeWindow", - "undeclared_inference_input", + "undeclared_inference_inputs", "UserInputCapability", "UserInputEvent", "UserInputs", diff --git a/flashdreams/flashdreams/runtime/mapping.py b/flashdreams/flashdreams/runtime/mapping.py index 1bb41d0c2..94f481406 100644 --- a/flashdreams/flashdreams/runtime/mapping.py +++ b/flashdreams/flashdreams/runtime/mapping.py @@ -156,7 +156,7 @@ class MappingCompatibility: __hash__ = None canonical_schema: CanonicalInputSchema - model_schema: InferenceInputSchema + inference_input_schema: InferenceInputSchema mapping_schema: InputMappingSchema missing_modalities: tuple[CanonicalModality, ...] = () missing_required_model_fields: tuple[tuple[InputPhase, InputField], ...] = () @@ -256,7 +256,7 @@ def _merge(target: list[Any], value: Any) -> None: def _build_compatibility( *, canonical_schema: CanonicalInputSchema, - model_schema: InferenceInputSchema, + inference_input_schema: InferenceInputSchema, mapping_schemas: Sequence[InputMappingSchema], reported_schema: InputMappingSchema, ) -> MappingCompatibility: @@ -269,7 +269,7 @@ def _build_compatibility( unavailable.append(mapping_schema) usable = combine_mapping_schemas(feedable, name=reported_schema.name) - required = model_schema.required_fields() + required = inference_input_schema.required_fields() missing_required = tuple( (phase, input_field) for phase, input_field in required @@ -282,7 +282,7 @@ def _build_compatibility( ) available_optional = tuple( (phase, input_field) - for phase, input_field in model_schema.optional_fields() + for phase, input_field in inference_input_schema.optional_fields() if usable.can_produce(phase, input_field) ) @@ -303,7 +303,7 @@ def _build_compatibility( return MappingCompatibility( canonical_schema=canonical_schema, - model_schema=model_schema, + inference_input_schema=inference_input_schema, mapping_schema=reported_schema, missing_modalities=tuple(missing_modalities), missing_required_model_fields=missing_required, @@ -316,7 +316,7 @@ def _build_compatibility( def check_mapping_compatibility( *, canonical_schema: CanonicalInputSchema, - model_schema: InferenceInputSchema, + inference_input_schema: InferenceInputSchema, mapping_schema: InputMappingSchema, ) -> MappingCompatibility: """Check whether a user-input source can drive a model through a mapping.""" @@ -324,7 +324,7 @@ def check_mapping_compatibility( raise TypeError("mapping_schema must be an InputMappingSchema object.") return _build_compatibility( canonical_schema=canonical_schema, - model_schema=model_schema, + inference_input_schema=inference_input_schema, mapping_schemas=(mapping_schema,), reported_schema=mapping_schema, ) @@ -333,7 +333,7 @@ def check_mapping_compatibility( def check_mapping_set_compatibility( *, canonical_schema: CanonicalInputSchema, - model_schema: InferenceInputSchema, + inference_input_schema: InferenceInputSchema, mapping_schemas: Sequence[InputMappingSchema], name: str = "input-mapping-set", ) -> MappingCompatibility: @@ -345,13 +345,13 @@ def check_mapping_set_compatibility( mapping_schemas = tuple(mapping_schemas) return _build_compatibility( canonical_schema=canonical_schema, - model_schema=model_schema, + inference_input_schema=inference_input_schema, mapping_schemas=mapping_schemas, reported_schema=combine_mapping_schemas(mapping_schemas, name=name), ) -def undeclared_inference_input( +def undeclared_inference_inputs( inputs: InferenceInput, mapping_schema: InputMappingSchema, ) -> tuple[tuple[InputPhase, str], ...]: diff --git a/flashdreams/tests/test_runtime_canonical.py b/flashdreams/tests/test_runtime_canonical.py index 80637e15c..77d442ed9 100644 --- a/flashdreams/tests/test_runtime_canonical.py +++ b/flashdreams/tests/test_runtime_canonical.py @@ -262,7 +262,7 @@ def test_mapping_written_against_a_modality_accepts_a_keyboard() -> None: compatibility = check_mapping_compatibility( canonical_schema=canonicalizer.canonical_schema(KEYBOARD_SOURCE), - model_schema=STEERING_MODEL, + inference_input_schema=STEERING_MODEL, mapping_schema=STEERING_MAPPING, ) @@ -276,7 +276,7 @@ def test_adding_a_device_needs_no_application_or_model_change() -> None: compatibility = check_mapping_compatibility( canonical_schema=canonicalizer.canonical_schema(WHEEL_SOURCE), - model_schema=STEERING_MODEL, + inference_input_schema=STEERING_MODEL, mapping_schema=STEERING_MAPPING, ) assert compatibility.can_drive diff --git a/flashdreams/tests/test_runtime_input_mapping.py b/flashdreams/tests/test_runtime_input_mapping.py index 20deac189..ab66933da 100644 --- a/flashdreams/tests/test_runtime_input_mapping.py +++ b/flashdreams/tests/test_runtime_input_mapping.py @@ -36,7 +36,7 @@ check_mapping_compatibility, check_mapping_set_compatibility, combine_mapping_schemas, - undeclared_inference_input, + undeclared_inference_inputs, ) pytestmark = pytest.mark.ci_cpu @@ -253,7 +253,7 @@ def test_metadata_is_excluded_from_field_equality() -> None: def test_compatible_source_model_and_mapping_can_drive() -> None: compatibility = check_mapping_set_compatibility( canonical_schema=CANONICAL_ALL, - model_schema=DRIVING_MODEL, + inference_input_schema=DRIVING_MODEL, mapping_schemas=(PROMPT_MAPPING, FRAME_MAPPING, STEERING_MAPPING), ) @@ -270,7 +270,7 @@ def test_compatible_source_model_and_mapping_can_drive() -> None: def test_missing_required_model_field_blocks_the_run() -> None: compatibility = check_mapping_set_compatibility( canonical_schema=CANONICAL_ALL, - model_schema=DRIVING_MODEL, + inference_input_schema=DRIVING_MODEL, mapping_schemas=(PROMPT_MAPPING,), ) @@ -285,7 +285,7 @@ def test_missing_source_capability_is_reported_when_it_blocks() -> None: compatibility = check_mapping_set_compatibility( canonical_schema=prompt_only, - model_schema=DRIVING_MODEL, + inference_input_schema=DRIVING_MODEL, mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), ) @@ -302,7 +302,7 @@ def test_unfeedable_optional_mapping_degrades_instead_of_vetoing() -> None: compatibility = check_mapping_set_compatibility( canonical_schema=no_frame_source, - model_schema=DRIVING_MODEL, + inference_input_schema=DRIVING_MODEL, mapping_schemas=(PROMPT_MAPPING, FRAME_MAPPING, STEERING_MAPPING), ) @@ -315,7 +315,7 @@ def test_unfeedable_optional_mapping_degrades_instead_of_vetoing() -> None: def test_optional_field_needs_mapping_support_to_be_available() -> None: compatibility = check_mapping_set_compatibility( canonical_schema=CANONICAL_ALL, - model_schema=DRIVING_MODEL, + inference_input_schema=DRIVING_MODEL, mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), ) @@ -334,7 +334,9 @@ def test_lifecycle_disagreement_blocks_a_field_match() -> None: ) compatibility = check_mapping_compatibility( - canonical_schema=CANONICAL_ALL, model_schema=model, mapping_schema=mapping + canonical_schema=CANONICAL_ALL, + inference_input_schema=model, + mapping_schema=mapping, ) assert not compatibility.can_drive @@ -345,7 +347,7 @@ def test_unspecified_lifecycle_stays_permissive() -> None: compatibility = check_mapping_compatibility( canonical_schema=CANONICAL_ALL, - model_schema=model, + inference_input_schema=model, mapping_schema=PROMPT_MAPPING, ) @@ -355,7 +357,7 @@ def test_unspecified_lifecycle_stays_permissive() -> None: def test_raise_if_incompatible_names_both_failure_kinds() -> None: compatibility = check_mapping_set_compatibility( canonical_schema=CanonicalInputSchema(modalities=(CONDITIONING_PROMPT,)), - model_schema=DRIVING_MODEL, + inference_input_schema=DRIVING_MODEL, mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), ) @@ -370,7 +372,7 @@ def test_raise_if_incompatible_names_both_failure_kinds() -> None: def test_raise_if_incompatible_is_a_no_op_when_compatible() -> None: compatibility = check_mapping_set_compatibility( canonical_schema=CANONICAL_ALL, - model_schema=DRIVING_MODEL, + inference_input_schema=DRIVING_MODEL, mapping_schemas=(PROMPT_MAPPING, FRAME_MAPPING, STEERING_MAPPING), ) @@ -383,7 +385,7 @@ def test_check_mapping_compatibility_rejects_a_non_schema() -> None: with pytest.raises(TypeError, match="InputMappingSchema"): check_mapping_compatibility( canonical_schema=CANONICAL_ALL, - model_schema=DRIVING_MODEL, + inference_input_schema=DRIVING_MODEL, mapping_schema=not_a_schema, ) @@ -437,7 +439,7 @@ def test_undeclared_inference_input_catches_schema_drift() -> None: global_conditioning={"prompt": "drive"}, step={"steering": 0.0} ) - undeclared = undeclared_inference_input(produced, PROMPT_MAPPING) + undeclared = undeclared_inference_inputs(produced, PROMPT_MAPPING) assert undeclared == (("step", "steering"),) @@ -448,7 +450,7 @@ def test_declared_outputs_report_no_drift() -> None: global_conditioning={"prompt": "drive"}, step={"steering": 0.0} ) - assert undeclared_inference_input(produced, combined) == () + assert undeclared_inference_inputs(produced, combined) == () # --- interoperability with the T1 envelope ------------------------------ @@ -473,7 +475,7 @@ def test_identity_mapping_needs_no_declared_surface() -> None: def test_empty_mapping_set_cannot_satisfy_a_required_field() -> None: compatibility = check_mapping_set_compatibility( canonical_schema=CANONICAL_ALL, - model_schema=DRIVING_MODEL, + inference_input_schema=DRIVING_MODEL, mapping_schemas=(), ) @@ -484,7 +486,7 @@ def test_empty_mapping_set_cannot_satisfy_a_required_field() -> None: def test_model_with_no_requirements_is_always_drivable() -> None: compatibility = check_mapping_set_compatibility( canonical_schema=CanonicalInputSchema(), - model_schema=InferenceInputSchema(), + inference_input_schema=InferenceInputSchema(), mapping_schemas=(), ) From e71b2f61858b58998787d9c748d9d2b25fe54207 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Tue, 4 Aug 2026 16:03:14 -0700 Subject: [PATCH 5/7] doc updates --- docs/inference_runtime_api_design.md | 82 ++++--- ...ence_runtime_supported_inputs_inventory.md | 214 +++++++++--------- 2 files changed, 154 insertions(+), 142 deletions(-) diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md index 00f9395d7..24b537ac0 100644 --- a/docs/inference_runtime_api_design.md +++ b/docs/inference_runtime_api_design.md @@ -19,7 +19,7 @@ integration-specific runner code: - `InferenceConfig`: how the model and inference stack should run; - `UserInputs`: controls or events from an app, replay trace, or benchmark; -- `ModelInputs`: prompts, frames, videos, trajectories, maps, scene data, and +- `InferenceInput`: prompts, frames, videos, trajectories, maps, scene data, and other values required by a specific model; - input mapping: model/application-specific conversion from user-facing inputs into model-facing inputs; @@ -68,9 +68,9 @@ Initial scope: | 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. | +| T1 | Complete | Minimal API envelope and naming. | Partly. | T0. | `InferenceConfig`, `UserInputs`, `InferenceInput`, runtime/session, output target, and mapping boundaries are defined well enough for demos to use. | +| T2 | Complete | 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 | Complete | `CanonicalInputs`, `InferenceInput`, schemas, and mapping boundary. | Yes, after T1 direction is agreed. | T1. | Models can declare required global/per-step inputs, and mappings can convert canonical inputs into inference 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. | @@ -103,14 +103,16 @@ Main runtime flow: App / integration / benchmark / transport chooses how the run is driven and where output goes supplies run setup: - InferenceConfig + UserInputs + ModelInputs + output/metrics options + InferenceConfig + UserInputs + InferenceInput + output/metrics options | v ModelRunner / standard loop orchestrates validation, lifecycle, stepping, output, and metrics + canonicalizes raw UserInputs into device-independent CanonicalInputs, + so applications and mappings never read raw device events uses input mapping to: - validate that user/app inputs can drive the model - build initial and per-step ModelInputs during the run + validate that canonical inputs can drive the model + build global and per-step InferenceInput during the run | v InferenceRuntime @@ -151,7 +153,7 @@ Create InferenceRuntime from InferenceConfig | v Start InferenceSession A - initial ModelInputs: prompt/frame/scene/etc. + global conditioning: prompt/frame/scene/etc. per-session state: cache, current step, reset state step 0 -> step 1 -> ... -> done outputs -> Output target @@ -160,7 +162,7 @@ Start InferenceSession A | v Start InferenceSession B - new initial ModelInputs or replay scenario + new global conditioning or replay scenario independent cache/state step 0 -> step 1 -> ... -> done outputs -> Output target @@ -311,23 +313,41 @@ User inputs are not model inputs. A keyboard event does not have one universal meaning. One model may map it to pose segments, another to steering commands, and another may ignore it. -## ModelInputs +## CanonicalInputs And InferenceInput -`ModelInputs` describes the data the model or inference pipeline actually -requires. It should distinguish: +Inputs move through three layers: -- initial inputs: values needed to start or reset a rollout; -- per-step inputs: values needed for one generated chunk or frame window. +```text +UserInputs -> CanonicalInputs -> InferenceInput + raw canonicalized encoded +``` + +Raw device events are canonicalized into device-independent modalities before an +application sees them, so adding a keyboard, gamepad, or wheel is a converter +registration rather than an application change. `InferenceInput` is what an +`InferenceSession` actually receives. + +`InferenceInput` describes the data the model or inference pipeline actually +requires. Both it and `CanonicalInputs` distinguish two conditioning slots: + +- global conditioning: values that condition the whole rollout; +- per-step conditioning: values needed for one generated chunk or frame window. + +Examples of global conditioning include prompt, negative prompt, conditioning +frame, input video, scene id, HD map asset, camera calibration, initial camera +pose, seed, or model-specific fields. -Examples of initial model inputs include prompt, negative prompt, first frame, -input video, scene id, HD map asset, camera calibration, initial camera pose, -seed, or model-specific fields. +Global conditioning is normally supplied when a session starts, but a non-empty +global slot on a mid-rollout input is an update request rather than a reset; +resetting rollout state is a separate `InferenceSession.reset()` call. Whether a +given value can be swapped mid-rollout is declared per field by +`InputField.update_policy`. -Examples of per-step model inputs include frame timestamps, pose segments, +Examples of per-step conditioning include frame timestamps, pose segments, camera trajectory chunks, rendered HD map frames, conditioning video windows, control tensors, event markers, or model-specific fields. -Model input payloads should use semantic names, not only modality names. For +Inference input payloads should use semantic names, not only modality names. For example, a first frame and an HD map frame should be distinct inputs even if both are image-like values. @@ -342,13 +362,14 @@ future external adapters may need different semantic fields. Adding a new model should usually mean adding adapter-owned schema declarations and mappings, not changing a central FlashDreams enum. -For interactive runs, most `ModelInputs` will be initial values plus per-step -inputs produced by input mapping. For MP4 generation and benchmarking, the API +For interactive runs, most `InferenceInput` values will be global conditioning +plus per-step inputs produced by input mapping. For MP4 generation and benchmarking, the API should also support fixed per-step model inputs so runs can be deterministic. ## Schemas -The API should support lightweight `UserInputSchema` and `ModelInputSchema` +The API should support lightweight `UserInputSchema`, `CanonicalInputSchema`, +and `InferenceInputSchema` metadata. These schemas are not meant to be a rich type system or a replacement for @@ -370,7 +391,7 @@ event capabilities, semantic model fields, payload representation hints, and lifecycle labels. For simple CLI text-to-video or image-to-video runs, `UserInputSchema` can be -trivial or omitted because there may be no live controls. `ModelInputSchema` is +trivial or omitted because there may be no live controls. `InferenceInputSchema` is more important because each supported model still needs to declare the model-facing values it expects. @@ -434,7 +455,7 @@ unless the checkpoint already matches a supported generic adapter. ## Input Mapping Input mapping is required whenever `UserInputs` need to become per-step -`ModelInputs`. In the T1 envelope this boundary is represented by a separate +`InferenceInput`. 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 @@ -451,7 +472,7 @@ There are two separate moments to keep clear: model; - 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 + `InferenceInput` 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 @@ -650,14 +671,16 @@ registry, standard loop, concrete output modes, or model migrations: `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`. +- Raw inputs use `UserInputs`, canonicalized inputs use `CanonicalInputs`, and + model-facing inputs use `InferenceInput`. Both remain lightweight payload envelopes with shallow read-only mappings. -- `UserInputSchema` and `ModelInputSchema` stay intentionally small: they +- `UserInputSchema`, `CanonicalInputSchema`, and `InferenceInputSchema` 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 + it while preserving the `CanonicalInputs` to `InferenceInput` boundary. Simple fixed-input runs can use `IdentityInputMapping`. - Output handling is represented by `OutputTarget`; `NullOutputTarget` is the initial headless implementation. @@ -689,7 +712,8 @@ Proceed with the proposed split: - `InferenceConfig` for model/runtime execution; - `UserInputs` for app-facing controls and replay traces; -- `ModelInputs` for model-facing initial and per-step inputs; +- `CanonicalInputs` for device-independent application-facing inputs; +- `InferenceInput` for model-facing global and per-step conditioning; - input mapping for model/application-specific conversion; - runtime/session boundaries for lifecycle and stepping; - output targets for display, streaming, files, and benchmarks; diff --git a/docs/inference_runtime_supported_inputs_inventory.md b/docs/inference_runtime_supported_inputs_inventory.md index 42f1f450e..ebe9d853a 100644 --- a/docs/inference_runtime_supported_inputs_inventory.md +++ b/docs/inference_runtime_supported_inputs_inventory.md @@ -3,7 +3,7 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# Supported Model Input Inventory For T2/T3 +# Supported Model Input Inventory This note inventories the inputs used by the currently supported FlashDreams runners and interactive runtimes, plus the SANA-WM input surface on `main`, then @@ -139,10 +139,10 @@ The inventory changes the T2/T3 shape in four concrete ways. First, a selected mapping is often a composition. A LingBot-like run needs prompt mapping, first-frame mapping, and keyboard-to-camera mapping. Omnidreams may add scene selection, camera selection, and HDMap mapping. The implementation should -support checking a set of mapper schemas as one compatibility surface, while -still allowing a single mapper object when that is simpler. +support checking a set of mapping schemas as one compatibility surface, while +still allowing a single mapping object when that is simpler. -Second, `ModelInputSchema` needs a lightweight lifecycle tag in addition to the +Second, `InferenceInputSchema` needs a lightweight lifecycle tag in addition to the `initial` versus `step` phase. The phase answers when the value is needed at the standard-loop level. The lifecycle tag distinguishes where the model adapter uses it, such as: @@ -159,13 +159,13 @@ uses it, such as: such as LingBot text-event embedding swaps. The lifecycle tag is metadata, not a new deep type system. If both a model field -and mapper output specify lifecycle, compatibility should require them to agree. +and mapping output specify lifecycle, compatibility should require them to agree. If either side omits it, matching stays permissive for simple schemas. -Third, `payload_kind` should be treated as a representation hint rather than a +Third, `semantic_type` should be treated as a representation hint rather than a universal semantic type. For example, `prompt` may arrive as inline text or a -path but become prompt text or text embeddings; `first_frame` may arrive as a -path, URL, bytes, or decoded tensor; camera motion may arrive as keys, pose JSON, +path but become prompt text or text embeddings; the global conditioning frame +may arrive as a path, URL, bytes, or decoded tensor; camera motion may arrive as keys, pose JSON, Numpy arrays, or integrated tensors. The semantic input name is still the main contract. @@ -176,71 +176,85 @@ schema URI, units, coordinate frame, accepted file suffixes, cardinality hints, or update notes. Metadata should remain query information and should not become the compatibility type system. -Fifth, `UserInputSchema` should describe raw source capabilities, while -mapper schemas describe derived model-facing semantics. A browser may provide -`key_down`, `key_up`, `prompt_set`, and `initial_frame_set` events. Whether that -can drive `steering`, `camera_trajectory`, `driver_command`, or text embedding -updates depends on the selected mapper and model schema. - -## Revised T2/T3 Plan - -The implementation plan after this inventory is: - -1. Keep `UserInputEvent`, `UserInputTrace`, and `UserInputWindow` as the primary - event-based input API. Static startup values remain timestamp-zero events. -2. Keep `UserInputSchema` lightweight and source-facing. It declares event kinds, - payload representation hints, and payload fields that a source can provide. -3. Keep `ModelInputs` split into `initial` and `step` payload maps. This still - matches the standard loop's session-start and per-step moments. -4. Extend `ModelInputField` with optional lifecycle metadata so models can - distinguish runtime config, cache initialization, rollout binding, per-step - inputs, and supported active-session updates. -5. Keep `InputMapperSchema` as the mapping boundary, but add mapper-set - compatibility helpers for composed mappings. -6. Keep input names, payload kinds, lifecycle labels, and metadata open-ended. +Fifth, `UserInputSchema` describes raw source capabilities, `CanonicalModality` +describes what an application consumes, and mapping schemas describe derived +model-facing semantics. A browser may provide `key_down`, `key_up`, +`prompt_set`, and `initial_frame_set` events. Those become canonical modalities +such as `driver_command` or `conditioning_prompt`; whether they can then drive +`steering`, `camera_trajectory`, or text embedding updates depends on the +selected mapping and model schema. + +## Implemented T2/T3 Shape + +The implementation that came out of this inventory is: + +1. Keep `UserInputEvent` and `UserInputs` as the raw event API, sliced by a + half-open `TimeWindow`. Static startup values remain timestamp-zero events. +2. Keep `UserInputSchema` lightweight and source-facing. `event_types` declares + that an event type exists; `UserInputCapability` additionally pins the + payload fields it carries. +3. Add a canonical layer between raw and encoded. `CanonicalModality` names a + device-independent input and its conditioning phase; `InputCanonicalizer` + registers per-device converters and produces `CanonicalInputs`. Applications + and mappings consume canonical inputs and never read raw device events. +4. Split `InferenceInput` (formerly `ModelInputs`) into `global_conditioning` + and `step`. A non-empty global slot mid-rollout is an update request, not a + reset; `InputField.update_policy` declares whether the model can apply it. +5. Extend `InputField` with `update_policy`, `lifecycle`, and `metadata` so + models can distinguish runtime config, cache initialization, rollout binding, + per-step inputs, and supported active-session updates. +6. Keep `InputMappingSchema` as the canonical-to-encoded boundary, with + mapping-set compatibility helpers for composed mappings. +7. Keep input names, semantic types, lifecycle labels, and metadata open-ended. Adding a new model should usually mean adding adapter-owned schema - declarations and mappers, not changing the core input dataclasses. -7. Leave deep validation to model adapters, sessions, and mappers. The schema - layer should catch obvious source/mapping/model mismatches before expensive - runtime initialization, not validate every tensor and coordinate convention. + declarations and mappings, not changing the core input dataclasses. +8. Leave deep validation to model adapters, sessions, and mappings. The schema + layer catches obvious source/mapping/model mismatches before expensive + runtime initialization; it does not validate every tensor and coordinate + convention. + +See `docs/inference_runtime_inputs_implementation.md` for the resulting API. ## Extensibility Contract -The inventory above is not a vocabulary freeze. The core API should not contain -a closed enum of allowed model input names. New adapters can introduce semantic -field names that match the model boundary they own. +The inventory above is not a vocabulary freeze. The core API does not contain a +closed enum of allowed input names. New adapters can introduce semantic field +names that match the model boundary they own. Use these conventions when adding future model schemas: - Prefer semantic names over modality names, such as `camera_trajectory_c2w` instead of `array`, or `hdmap_frames` instead of `image`. -- Use `payload_kind` for a coarse representation hint, such as `path`, - `decoded_tensor`, `c2w_sequence`, `intrinsics_vec4_sequence`, `driver_command`, - or `embedding`. +- Use `semantic_type` for a coarse representation hint, such as `path`, + `decoded_tensor`, `c2w_sequence`, `intrinsics_vec4_sequence`, or `embedding`. - Use `lifecycle` to say where the adapter consumes the value, such as `runtime_config`, `cache_init`, `rollout_binding`, `step_input`, or `session_update`. +- Use `update_policy` to say when a value may change. `SESSION_START_ONLY` is + the one reserved token, meaning the value cannot be swapped mid-rollout. - Use `metadata` for query hints: units, coordinate frame, shape summary, accepted suffixes, schema URI, model family, value ranges, or cardinality. -- Keep deep validation in the adapter/mapper. The lightweight schemas answer +- Keep deep validation in the adapter/mapping. The lightweight schemas answer whether the selected source and mapping can plausibly drive the model before expensive initialization. ## Representative Schema Sketches -These sketches are not migration work for T4+, but they show that the current -T2/T3 primitives can describe the supported input surfaces. +These are not migration work for T4+, but they show that the current primitives +can describe the supported input surfaces. All use +`flashdreams.runtime.InferenceInputSchema` and `InputField`. ```python -lingbot_model = ModelInputSchema( - name="lingbot-world", - fields=( - ModelInputField("prompt", "initial", lifecycle="cache_init"), - ModelInputField("first_frame", "initial", lifecycle="cache_init"), - ModelInputField("camera_trajectory", "step", lifecycle="step_input"), - ModelInputField( - "text_embeddings", - "step", +lingbot_model = InferenceInputSchema( + description="lingbot-world", + global_fields=( + InputField(name="prompt", lifecycle="cache_init"), + InputField(name="global_conditioning_frame", lifecycle="cache_init"), + ), + step_fields=( + InputField(name="camera_trajectory", lifecycle="step_input"), + InputField( + name="text_embeddings", required=False, update_policy="step_boundary", lifecycle="session_update", @@ -250,84 +264,58 @@ lingbot_model = ModelInputSchema( ``` ```python -omnidreams_model = ModelInputSchema( - name="omnidreams", - fields=( - ModelInputField("prompts", "initial", lifecycle="cache_init"), - ModelInputField("first_frames", "initial", lifecycle="cache_init"), - ModelInputField("view_names", "initial", lifecycle="cache_init"), - ModelInputField("hdmap_frames", "step", lifecycle="step_input"), - ModelInputField( - "text_embeddings", - "initial", - required=False, - lifecycle="cache_init", - ), - ModelInputField( - "image_embeddings", - "initial", - required=False, - lifecycle="cache_init", - ), +omnidreams_model = InferenceInputSchema( + description="omnidreams", + global_fields=( + InputField(name="prompts", lifecycle="cache_init"), + InputField(name="global_conditioning_frames", lifecycle="cache_init"), + InputField(name="view_names", lifecycle="cache_init"), + InputField(name="text_embeddings", required=False, lifecycle="cache_init"), + InputField(name="image_embeddings", required=False, lifecycle="cache_init"), ), + step_fields=(InputField(name="hdmap_frames", lifecycle="step_input"),), ) ``` ```python -hy_worldplay_model = ModelInputSchema( - name="hy-worldplay", - fields=( - ModelInputField("prompt", "initial", lifecycle="cache_init"), - ModelInputField("first_frame", "initial", lifecycle="cache_init"), - ModelInputField("action_labels", "initial", lifecycle="rollout_binding"), - ModelInputField("camera_viewmats", "initial", lifecycle="rollout_binding"), - ModelInputField("camera_intrinsics", "initial", lifecycle="rollout_binding"), - ModelInputField("memory_config", "initial", lifecycle="rollout_binding"), +hy_worldplay_model = InferenceInputSchema( + description="hy-worldplay", + global_fields=( + InputField(name="prompt", lifecycle="cache_init"), + InputField(name="global_conditioning_frame", lifecycle="cache_init"), + InputField(name="action_labels", lifecycle="rollout_binding"), + InputField(name="camera_viewmats", lifecycle="rollout_binding"), + InputField(name="camera_intrinsics", lifecycle="rollout_binding"), + InputField(name="memory_config", lifecycle="rollout_binding"), ), ) ``` ```python -sana_wm_model = ModelInputSchema( - name="sana-wm", - fields=( - ModelInputField("prompt", "initial", lifecycle="cache_init"), - ModelInputField( - "negative_prompt", - "initial", - required=False, - lifecycle="cache_init", - ), - ModelInputField("first_frame", "initial", lifecycle="cache_init"), - ModelInputField( - "camera_trajectory_c2w", - "initial", - payload_kind="c2w_sequence", +sana_wm_model = InferenceInputSchema( + description="sana-wm", + global_fields=( + InputField(name="prompt", lifecycle="cache_init"), + InputField(name="negative_prompt", required=False, lifecycle="cache_init"), + InputField(name="global_conditioning_frame", lifecycle="cache_init"), + InputField( + name="camera_trajectory_c2w", + semantic_type="c2w_sequence", lifecycle="rollout_binding", metadata={"shape": "[F,4,4]", "coordinates": "opencv_c2w"}, ), - ModelInputField( - "camera_intrinsics_vec4", - "initial", + InputField( + name="camera_intrinsics_vec4", required=False, - payload_kind="intrinsics_vec4_sequence", + semantic_type="intrinsics_vec4_sequence", lifecycle="rollout_binding", metadata={"shape": "[F,4]"}, ), - ModelInputField( - "stage1_sampling", - "initial", - lifecycle="rollout_binding", - metadata={"fields": ("steps", "cfg_scale", "flow_shift", "seed")}, - ), - ModelInputField( - "streaming_chunking", - "initial", - required=False, - lifecycle="rollout_binding", - metadata={"fields": ("num_frame_per_block", "cached_blocks")}, - ), ), - metadata={"model_family": "sana-wm"}, ) ``` + +SANA-WM's `stage1_sampling` and `streaming_chunking` are deliberately absent +above. They describe how to run the model rather than what conditions it, so +they belong in `InferenceConfig`, not in an input schema. Flagged here because +the runner currently threads them alongside the conditioning inputs. From cb2056abf0fd22463c500149d38205b03be7f19a Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Tue, 4 Aug 2026 16:26:13 -0700 Subject: [PATCH 6/7] Update based on new diagrams --- docs/inference_runtime_api_design.md | 4 +- ...inference_runtime_inputs_implementation.md | 49 ++++++--- flashdreams/flashdreams/runtime/__init__.py | 6 -- flashdreams/flashdreams/runtime/canonical.py | 101 ++---------------- flashdreams/flashdreams/runtime/inputs.py | 53 +++------ .../tests/test_inference_runtime_api.py | 4 +- flashdreams/tests/test_runtime_canonical.py | 65 ++++------- .../tests/test_runtime_input_mapping.py | 71 +++++++----- 8 files changed, 116 insertions(+), 237 deletions(-) diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md index 24b537ac0..f70fbd890 100644 --- a/docs/inference_runtime_api_design.md +++ b/docs/inference_runtime_api_design.md @@ -108,10 +108,8 @@ App / integration / benchmark / transport v ModelRunner / standard loop orchestrates validation, lifecycle, stepping, output, and metrics - canonicalizes raw UserInputs into device-independent CanonicalInputs, - so applications and mappings never read raw device events uses input mapping to: - validate that canonical inputs can drive the model + validate that user/app inputs can drive the model build global and per-step InferenceInput during the run | v diff --git a/docs/inference_runtime_inputs_implementation.md b/docs/inference_runtime_inputs_implementation.md index 33bc9ad7d..a967d0ca8 100644 --- a/docs/inference_runtime_inputs_implementation.md +++ b/docs/inference_runtime_inputs_implementation.md @@ -44,6 +44,12 @@ events: `InputMapping.map_step_inputs` takes `canonical_inputs`, not Adding a keyboard, gamepad, or wheel is an `InputCanonicalizer.register` call that touches no application, mapping, or model code. +This path covers **live user control only**. Global conditioning is +application-owned data and reaches `InferenceInput` directly, without passing +through canonicalization or a device converter. An application that wants a +trigger key to swap the prompt reads that as ordinary canonical control input +and updates its own global conditioning in response. + ## Conditioning Slots Both the canonical and encoded layers split into two slots, and the split means @@ -100,11 +106,10 @@ schema.unsupported_global_updates( in that vocabulary, and all of `lifecycle`, is open and adapter-owned; this layer only carries it as queryable metadata. -The canonical layer mirrors this. Per-step converters emit every window, because -per-step conditioning is level-triggered: a key held across a step emits no -events but still means full throttle. Global-conditioning converters emit *only* -when a value changed, so a quiet window does not look like a repeated update -request. +Steady-state steps must leave the global slot empty; otherwise every step reads +as an update request. Converters emit every window, because live control is +level-triggered: a key held across a step emits no events but still means full +throttle. ## Raw Inputs @@ -134,8 +139,8 @@ fields, so schemas written before capabilities existed keep working. ## Canonical Modalities -A `CanonicalModality` is a device-independent input: a name, a phase, and the -payload fields it guarantees. Converters implement `DeviceConverter`, declaring +A `CanonicalModality` is a device-independent input: a name and the payload +fields it guarantees. Converters implement `DeviceConverter`, declaring what raw capabilities they consume and which modality they produce. ```python @@ -149,15 +154,13 @@ canonicalizer.register(WheelToDriverCommand()) # a wheel is one call canonical = canonicalizer.canonicalize( user_inputs, window=TimeWindow(start_s=0.0, end_s=1.0), source_schema=browser ) -canonical.per_step["driver_command"]["throttle"] +canonical.values["driver_command"]["throttle"] ``` -Shipped modalities are `DRIVER_COMMAND` (step), `CONDITIONING_PROMPT` and -`CONDITIONING_FRAME` (global). `KeyboardToDriverCommand` reuses +`DRIVER_COMMAND` is the one shipped modality. `KeyboardToDriverCommand` reuses `KeyboardState`/`normalize_key` from `flashdreams.serving.realtime.input` and mirrors the semantics the Omnidreams interactive-drive keyboard backend already -has. `LatestEventToModality` is the general global-conditioning converter and -rejects a per-step modality at construction. +has. Converters are stateful, so feed windows in session order and call `InputCanonicalizer.reset()` at a rollout boundary. Replaying the same window @@ -232,13 +235,25 @@ Tracked against the runtime API discussion, not yet settled: required set, so "accepts `{prompt}` OR `{prompt, conditioning_frame}`" cannot be expressed. `MappingCompatibility.missing_required_model_fields` assumes a single required set too. -- **Who owns encoding.** Whether text/image embedding happens app-side (a - `publish_input(modality) -> Tensor` handler) or model-side. This does not - change what `InputMappingSchema` declares — a mapping consumes canonical - modalities either way — but it decides whether `InferenceInput` payloads hold - tensors or canonical values. - **`step()` returning a future**, for models with a dependency on their own output. `InferenceSession.step()` is currently synchronous. +- **`Input System` ownership.** The diagrams show it pulling events and + supporting mock input, so the Application has-a input system. + `InputCanonicalizer` is currently a pure function over a supplied window and + owns no source. This is the one deferred item inside the input scope. + +## Owned Elsewhere + +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 + 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. +- **`Application`**, the class that has-a input system, input map, global + conditioning, session, and output target. T4. - **Loop ownership** — whether the application or the runtime/session drives the main event loop, and whether inputs are queued and batched. diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index b31582ec5..a27866893 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -8,14 +8,11 @@ """ from flashdreams.runtime.canonical import ( - CONDITIONING_FRAME, - CONDITIONING_PROMPT, DRIVER_COMMAND, DeviceConverter, DeviceConverterSchema, InputCanonicalizer, KeyboardToDriverCommand, - LatestEventToModality, ) from flashdreams.runtime.config import ExecutionBackend, InferenceConfig, Precision from flashdreams.runtime.inputs import ( @@ -67,8 +64,6 @@ "check_mapping_compatibility", "check_mapping_set_compatibility", "combine_mapping_schemas", - "CONDITIONING_FRAME", - "CONDITIONING_PROMPT", "DeclaresMappingSchema", "DeviceConverter", "DeviceConverterSchema", @@ -88,7 +83,6 @@ "InputMappingSchema", "InputPhase", "KeyboardToDriverCommand", - "LatestEventToModality", "MappingCompatibility", "MetricsRecorder", "ModelAdapter", diff --git a/flashdreams/flashdreams/runtime/canonical.py b/flashdreams/flashdreams/runtime/canonical.py index 5bd21fc32..f99ee8340 100644 --- a/flashdreams/flashdreams/runtime/canonical.py +++ b/flashdreams/flashdreams/runtime/canonical.py @@ -15,9 +15,9 @@ :meth:`InputCanonicalizer.reset` at a rollout boundary; replaying the same window sequence then reproduces the same canonical inputs. -Global-conditioning converters behave differently on purpose. They emit only -when a value actually changed in the window, so a downstream non-empty global -slot means "update this", not "re-apply it every step". +This layer covers live user control only. Global conditioning such as a prompt +or conditioning frame is application-owned and reaches ``InferenceInput`` +directly, without passing through canonicalization. """ from __future__ import annotations @@ -87,15 +87,13 @@ def convert( """Return the modality value for ``window``, or ``None`` if inactive. ``user_inputs`` is already filtered to ``window``. Returning ``None`` - lets a present-but-idle device yield to a lower-priority one, and lets a - global-conditioning converter stay silent when nothing changed. + lets a present-but-idle device yield to a lower-priority one. """ ... DRIVER_COMMAND = CanonicalModality( name="driver_command", - phase="step", payload_fields=frozenset({"throttle", "brake", "steer", "stop", "reverse"}), description=( "Normalized driving intent. throttle/brake are in [0, 1], steer is in " @@ -103,20 +101,6 @@ def convert( ), ) -CONDITIONING_PROMPT = CanonicalModality( - name="conditioning_prompt", - phase="global", - payload_fields=frozenset({"prompt"}), - description="Global conditioning prompt; may change mid-rollout.", -) - -CONDITIONING_FRAME = CanonicalModality( - name="conditioning_frame", - phase="global", - payload_fields=frozenset({"image"}), - description="Global conditioning frame; may change mid-rollout.", -) - class KeyboardToDriverCommand: """Convert keyboard edges into :data:`DRIVER_COMMAND` level state. @@ -193,70 +177,6 @@ def convert( ) -class LatestEventToModality: - """Emit a global-conditioning modality when its source event fires. - - Global conditioning is transient by design: this returns ``None`` in windows - where nothing changed, so downstream code only sees an update request when - the user actually changed something. - """ - - def __init__( - self, - *, - modality: CanonicalModality, - event_type: str, - payload_fields: frozenset[str] | None = None, - name: str | None = None, - device_kind: str | None = None, - priority: int = 0, - ) -> None: - if modality.phase != "global": - raise ValueError( - "LatestEventToModality is for global conditioning; " - f"{modality.name!r} declares phase {modality.phase!r}." - ) - self._modality = modality - self._event_type = event_type - fields = modality.payload_fields if payload_fields is None else payload_fields - self._schema = DeviceConverterSchema( - name=name or f"{event_type}-to-{modality.name}", - produces=modality, - device_kind=device_kind, - priority=priority, - consumes=( - UserInputCapability(event_type=event_type, payload_fields=fields), - ), - ) - - @property - def schema(self) -> DeviceConverterSchema: - return self._schema - - def reset(self) -> None: - return None - - def convert( - self, - user_inputs: UserInputs, - window: TimeWindow, - ) -> Mapping[str, Any] | None: - del window - latest = None - for event in user_inputs.events: - if event.event_type == self._event_type: - latest = event - if latest is None: - return None - return self._modality.value( - { - name: latest.payload[name] - for name in self._modality.payload_fields - if name in latest.payload - } - ) - - class InputCanonicalizer: """Registry of device converters plus the raw-to-canonical rewrite. @@ -353,22 +273,17 @@ def canonicalize( returned a value wins. """ windowed = user_inputs.window(window) - by_phase: dict[str, dict[str, Any]] = {"global": {}, "step": {}} + values: dict[str, Any] = {} sources: dict[str, str] = {} for converter in self.converters_for(source_schema): value = converter.convert(windowed, window) modality = converter.schema.produces - slot = by_phase[modality.phase] - if value is not None and modality.name not in slot: - slot[modality.name] = value + if value is not None and modality.name not in values: + values[modality.name] = value if converter.schema.device_kind is not None: sources[modality.name] = converter.schema.device_kind metadata: dict[str, Any] = {} if sources: metadata["canonical_sources"] = freeze_mapping(sources) - return CanonicalInputs( - global_conditioning=by_phase["global"], - per_step=by_phase["step"], - metadata=metadata, - ) + return CanonicalInputs(values=values, metadata=metadata) diff --git a/flashdreams/flashdreams/runtime/inputs.py b/flashdreams/flashdreams/runtime/inputs.py index cacab635f..f0be31bea 100644 --- a/flashdreams/flashdreams/runtime/inputs.py +++ b/flashdreams/flashdreams/runtime/inputs.py @@ -369,13 +369,12 @@ class CanonicalModality: modalities; they never read raw device events, so adding a new device is a converter registration rather than an application change. - ``phase`` says which conditioning slot the modality feeds: ``"global"`` for - one-shot rollout conditioning such as a prompt or conditioning frame, and - ``"step"`` for per-step conditioning such as steering. + 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. """ name: str - phase: InputPhase = "step" payload_fields: frozenset[str] = field(default_factory=frozenset) metadata: Mapping[str, Any] = field( default_factory=dict, @@ -387,7 +386,6 @@ class CanonicalModality: def __post_init__(self) -> None: if not self.name.strip(): raise ValueError("CanonicalModality.name must be non-empty.") - object.__setattr__(self, "phase", validate_phase(self.phase)) for payload_field in self.payload_fields: if not payload_field.strip(): raise ValueError("payload field names must be non-empty.") @@ -395,10 +393,8 @@ def __post_init__(self) -> None: def is_satisfied_by(self, provider: "CanonicalModality") -> bool: """Return whether ``provider`` can satisfy this consumed modality.""" - return ( - self.name == provider.name - and self.phase == provider.phase - and self.payload_fields.issubset(provider.payload_fields) + return self.name == provider.name and self.payload_fields.issubset( + provider.payload_fields ) def value(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: @@ -423,51 +419,26 @@ def supports(self, modality: CanonicalModality) -> bool: """Return whether this source can supply ``modality``.""" return any(modality.is_satisfied_by(provided) for provided in self.modalities) - def modalities_for(self, phase: InputPhase) -> tuple[CanonicalModality, ...]: - """Return supplied modalities feeding ``phase``.""" - validate_phase(phase) - return tuple( - modality for modality in self.modalities if modality.phase == phase - ) - @dataclass(frozen=True, kw_only=True, slots=True) class CanonicalInputs: - """Canonicalized user input for one step, split by conditioning slot. + """Canonicalized user input for one step, keyed by modality name. - ``per_step`` is level-triggered and normally present every step: a key held - down emits no events but still means full throttle. ``global_conditioning`` - is transient and populated only when a value actually changed in this - window, which is what makes a non-empty global slot downstream mean "update - this" rather than "re-apply every step". + 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. """ __hash__ = None - global_conditioning: Mapping[str, Any] = field(default_factory=dict) - per_step: Mapping[str, Any] = field(default_factory=dict) + values: Mapping[str, Any] = field(default_factory=dict) metadata: 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", freeze_mapping(self.per_step)) + object.__setattr__(self, "values", freeze_mapping(self.values)) object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) - @property - def has_global_change(self) -> bool: - """Return whether a global conditioning value changed in this window.""" - return bool(self.global_conditioning) - - def for_phase(self, phase: InputPhase) -> Mapping[str, Any]: - """Return the canonical payload for ``phase``.""" - return ( - self.global_conditioning - if validate_phase(phase) == "global" - else self.per_step - ) - @dataclass(frozen=True, kw_only=True, slots=True) class InferenceInput: diff --git a/flashdreams/tests/test_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py index 32fe7bdc2..edfafa634 100644 --- a/flashdreams/tests/test_inference_runtime_api.py +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -390,8 +390,10 @@ def _drive_two_step_session( or TimeWindow(start_s=0.0, end_s=_SESSION_HORIZON_S), source_schema=source_schema, ), + # The global slot stays empty in steady state. A mapping that + # sees ``canonical_inputs.has_global_change`` fills it via + # ``with_global_update`` to request a mid-rollout swap. inference_input=InferenceInput( - global_conditioning=initial_inputs.global_conditioning, step={"chunk_index": request.step_index}, ), request=request, diff --git a/flashdreams/tests/test_runtime_canonical.py b/flashdreams/tests/test_runtime_canonical.py index 77d442ed9..64ee85d75 100644 --- a/flashdreams/tests/test_runtime_canonical.py +++ b/flashdreams/tests/test_runtime_canonical.py @@ -16,17 +16,16 @@ import pytest from flashdreams.runtime import ( - CONDITIONING_PROMPT, DRIVER_COMMAND, CanonicalInputs, CanonicalModality, DeviceConverterSchema, + InferenceInput, InferenceInputSchema, InputCanonicalizer, InputField, InputMappingSchema, KeyboardToDriverCommand, - LatestEventToModality, TimeWindow, UserInputCapability, UserInputEvent, @@ -125,8 +124,8 @@ def _key(event_type: str, key: str, timestamp_s: float) -> UserInputEvent: def _command(canonical: CanonicalInputs) -> Mapping[str, Any]: - assert DRIVER_COMMAND.name in canonical.per_step - return canonical.per_step[DRIVER_COMMAND.name] + assert DRIVER_COMMAND.name in canonical.values + return canonical.values[DRIVER_COMMAND.name] # --- per-step conditioning ---------------------------------------------- @@ -194,64 +193,36 @@ def test_reset_drops_device_state() -> None: assert _command(after)["throttle"] == 0.0 -# --- global conditioning ------------------------------------------------ +# --- boundary: global conditioning is not canonicalized ----------------- -def test_global_conditioning_is_emitted_only_when_it_changes() -> None: - """A quiet window must not look like a repeated update request.""" - canonicalizer = InputCanonicalizer( - [LatestEventToModality(modality=CONDITIONING_PROMPT, event_type="prompt_set")] - ) +def test_canonical_inputs_carry_live_control_only() -> None: + """Global conditioning is application-owned and bypasses this layer.""" + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) inputs = UserInputs( events=( + _key("key_down", "w", 0.1), UserInputEvent( timestamp_s=0.2, event_type="prompt_set", payload={"prompt": "rain"} ), ) ) - changed = canonicalizer.canonicalize( - inputs, window=WINDOW, source_schema=PROMPT_SOURCE - ) - quiet = canonicalizer.canonicalize( - inputs, window=NEXT_WINDOW, source_schema=PROMPT_SOURCE + canonical = canonicalizer.canonicalize( + inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE ) - assert changed.has_global_change - assert changed.global_conditioning["conditioning_prompt"]["prompt"] == "rain" - assert not quiet.has_global_change - + assert set(canonical.values) == {"driver_command"} -def test_global_converter_rejects_a_per_step_modality() -> None: - with pytest.raises(ValueError, match="global conditioning"): - LatestEventToModality(modality=DRIVER_COMMAND, event_type="wheel_axis") - -def test_global_and_per_step_land_in_separate_slots() -> None: - canonicalizer = InputCanonicalizer( - [ - KeyboardToDriverCommand(), - LatestEventToModality( - modality=CONDITIONING_PROMPT, event_type="prompt_set" - ), - ] - ) - source = UserInputSchema( - capabilities=KEYBOARD_SOURCE.capabilities + PROMPT_SOURCE.capabilities - ) - inputs = UserInputs( - events=( - _key("key_down", "w", 0.1), - UserInputEvent( - timestamp_s=0.2, event_type="prompt_set", payload={"prompt": "rain"} - ), - ) +def test_application_supplies_global_conditioning_directly() -> None: + """A prompt swap reaches the session without touching canonicalization.""" + update = InferenceInput(step={"steering": 0.0}).with_global_update( + {"prompt": "heavy rain"} ) - canonical = canonicalizer.canonicalize(inputs, window=WINDOW, source_schema=source) - - assert set(canonical.per_step) == {"driver_command"} - assert set(canonical.global_conditioning) == {"conditioning_prompt"} + assert update.requests_global_update + assert update.global_conditioning["prompt"] == "heavy rain" # --- device independence ------------------------------------------------ @@ -434,7 +405,7 @@ def convert( window=WINDOW, source_schema=source, ) - assert canonical.per_step["pedal_state"]["throttle"] == pytest.approx(0.75) + assert canonical.values["pedal_state"]["throttle"] == pytest.approx(0.75) def test_replaying_the_same_windows_reproduces_the_same_canonical_inputs() -> None: diff --git a/flashdreams/tests/test_runtime_input_mapping.py b/flashdreams/tests/test_runtime_input_mapping.py index ab66933da..00cd9758f 100644 --- a/flashdreams/tests/test_runtime_input_mapping.py +++ b/flashdreams/tests/test_runtime_input_mapping.py @@ -16,12 +16,11 @@ import pytest from flashdreams.runtime import ( - CONDITIONING_FRAME, - CONDITIONING_PROMPT, DRIVER_COMMAND, SESSION_START_ONLY, CanonicalInputs, CanonicalInputSchema, + CanonicalModality, IdentityInputMapping, InferenceInput, InferenceInputSchema, @@ -57,18 +56,20 @@ description="browser webrtc client", ) -CANONICAL_ALL = CanonicalInputSchema( - modalities=(CONDITIONING_PROMPT, CONDITIONING_FRAME, DRIVER_COMMAND) +CAMERA_LOOK = CanonicalModality( + name="camera_look", payload_fields=frozenset({"yaw", "pitch"}) ) +CANONICAL_ALL = CanonicalInputSchema(modalities=(DRIVER_COMMAND, CAMERA_LOOK)) + +# Global conditioning is application-owned and does not come from a canonical +# modality, so this mapping consumes nothing and only declares what it produces. PROMPT_MAPPING = InputMappingSchema( name="prompt", - consumes=(CONDITIONING_PROMPT,), produces_global=(InputField(name="prompt", semantic_type="text"),), ) FRAME_MAPPING = InputMappingSchema( name="conditioning-frame", - consumes=(CONDITIONING_FRAME,), produces_global=(InputField(name="global_conditioning_frame", required=False),), ) STEERING_MAPPING = InputMappingSchema( @@ -76,15 +77,20 @@ consumes=(DRIVER_COMMAND,), produces_step=(InputField(name="steering"),), ) +LOOK_MAPPING = InputMappingSchema( + name="camera-look", + consumes=(CAMERA_LOOK,), + produces_step=(InputField(name="camera_delta", required=False),), +) DRIVING_MODEL = InferenceInputSchema( global_fields=( InputField(name="prompt", semantic_type="text", lifecycle="cache_init"), - InputField( - name="global_conditioning_frame", required=False, lifecycle="cache_init" - ), ), - step_fields=(InputField(name="steering", lifecycle="step_input"),), + step_fields=( + InputField(name="steering", lifecycle="step_input"), + InputField(name="camera_delta", required=False, lifecycle="step_input"), + ), ) @@ -195,9 +201,7 @@ def test_model_declares_required_and_optional_fields_per_phase() -> None: ("global", "prompt"), ("step", "steering"), } - assert {(phase, f.name) for phase, f in optional} == { - ("global", "global_conditioning_frame") - } + assert {(phase, f.name) for phase, f in optional} == {("step", "camera_delta")} def test_required_fields_can_be_filtered_by_phase() -> None: @@ -254,7 +258,7 @@ def test_compatible_source_model_and_mapping_can_drive() -> None: compatibility = check_mapping_set_compatibility( canonical_schema=CANONICAL_ALL, inference_input_schema=DRIVING_MODEL, - mapping_schemas=(PROMPT_MAPPING, FRAME_MAPPING, STEERING_MAPPING), + mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING, LOOK_MAPPING), ) assert compatibility.can_drive @@ -263,7 +267,7 @@ def test_compatible_source_model_and_mapping_can_drive() -> None: ("step", "steering"), } assert {(p, f.name) for p, f in compatibility.available_optional_model_fields} == { - ("global", "global_conditioning_frame") + ("step", "camera_delta") } @@ -281,10 +285,10 @@ def test_missing_required_model_field_blocks_the_run() -> None: def test_missing_source_capability_is_reported_when_it_blocks() -> None: - prompt_only = CanonicalInputSchema(modalities=(CONDITIONING_PROMPT,)) + no_wheel = CanonicalInputSchema(modalities=(CAMERA_LOOK,)) compatibility = check_mapping_set_compatibility( - canonical_schema=prompt_only, + canonical_schema=no_wheel, inference_input_schema=DRIVING_MODEL, mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), ) @@ -296,18 +300,16 @@ def test_missing_source_capability_is_reported_when_it_blocks() -> None: def test_unfeedable_optional_mapping_degrades_instead_of_vetoing() -> None: """Losing a mapping that fed only optional fields must not block the run.""" - no_frame_source = CanonicalInputSchema( - modalities=(CONDITIONING_PROMPT, DRIVER_COMMAND) - ) + no_look = CanonicalInputSchema(modalities=(DRIVER_COMMAND,)) compatibility = check_mapping_set_compatibility( - canonical_schema=no_frame_source, + canonical_schema=no_look, inference_input_schema=DRIVING_MODEL, - mapping_schemas=(PROMPT_MAPPING, FRAME_MAPPING, STEERING_MAPPING), + mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING, LOOK_MAPPING), ) assert compatibility.can_drive - assert compatibility.unavailable_mapping_names == ("conditioning-frame",) + assert compatibility.unavailable_mapping_names == ("camera-look",) # The dropped mapping's field must not be advertised as available. assert compatibility.available_optional_model_fields == () @@ -329,7 +331,6 @@ def test_lifecycle_disagreement_blocks_a_field_match() -> None: ) mapping = InputMappingSchema( name="prompt", - consumes=(CONDITIONING_PROMPT,), produces_global=(InputField(name="prompt", lifecycle="cache_init"),), ) @@ -356,7 +357,7 @@ def test_unspecified_lifecycle_stays_permissive() -> None: def test_raise_if_incompatible_names_both_failure_kinds() -> None: compatibility = check_mapping_set_compatibility( - canonical_schema=CanonicalInputSchema(modalities=(CONDITIONING_PROMPT,)), + canonical_schema=CanonicalInputSchema(modalities=(CAMERA_LOOK,)), inference_input_schema=DRIVING_MODEL, mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), ) @@ -396,10 +397,7 @@ def test_check_mapping_compatibility_rejects_a_non_schema() -> None: def test_combining_mappings_unions_their_surfaces() -> None: combined = combine_mapping_schemas((PROMPT_MAPPING, STEERING_MAPPING)) - assert {m.name for m in combined.consumes} == { - "conditioning_prompt", - "driver_command", - } + assert {m.name for m in combined.consumes} == {"driver_command"} assert [f.name for f in combined.produces_global] == ["prompt"] assert [f.name for f in combined.produces_step] == ["steering"] @@ -558,3 +556,18 @@ def test_undeclared_global_values_are_left_to_the_adapter() -> None: update = InferenceInput(global_conditioning={"mystery": 1}) assert schema.unsupported_global_updates(update) == () + + +def test_steady_state_steps_do_not_request_a_global_update() -> None: + """Carrying session-start conditioning forward would look like an update.""" + started = InferenceInput(global_conditioning={"prompt": "drive"}) + + steady_state = InferenceInput(step={"chunk_index": 1}) + + assert started.requests_global_update + assert not steady_state.requests_global_update + assert ( + not started.with_step({"chunk_index": 1}) + .without_global_update() + .requests_global_update + ) From b6ed51a91b151e6190e5e71fcedcac4578fee560 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Tue, 4 Aug 2026 16:55:00 -0700 Subject: [PATCH 7/7] Align closer to diagrams --- ...inference_runtime_inputs_implementation.md | 28 ++- flashdreams/flashdreams/runtime/__init__.py | 4 + flashdreams/flashdreams/runtime/canonical.py | 126 ++++++++++++-- flashdreams/tests/test_runtime_canonical.py | 163 ++++++++++++++++++ 4 files changed, 302 insertions(+), 19 deletions(-) diff --git a/docs/inference_runtime_inputs_implementation.md b/docs/inference_runtime_inputs_implementation.md index a967d0ca8..8d485768a 100644 --- a/docs/inference_runtime_inputs_implementation.md +++ b/docs/inference_runtime_inputs_implementation.md @@ -160,7 +160,24 @@ canonical.values["driver_command"]["throttle"] `DRIVER_COMMAND` is the one shipped modality. `KeyboardToDriverCommand` reuses `KeyboardState`/`normalize_key` from `flashdreams.serving.realtime.input` and mirrors the semantics the Omnidreams interactive-drive keyboard backend already -has. +has. Its key bindings are data (`DEFAULT_DRIVING_BINDINGS`), and the set of +tracked keys is derived from them, so a rebound layout cannot leave an action +unreachable. + +`ScriptedModality` is the mock/replay converter. It consumes no raw +capabilities, so a benchmark or test can author a scenario at the canonical +level without knowing any device vocabulary: + +```python +canonicalizer = InputCanonicalizer([ + ScriptedModality(modality=DRIVER_COMMAND, timeline=[(0.0, full_throttle)]), +]) +canonicalizer.canonicalize( + UserInputs(), window=step_window, source_schema=UserInputSchema() +) +``` + +Application code is identical between a real run and a scripted one. Converters are stateful, so feed windows in session order and call `InputCanonicalizer.reset()` at a rollout boundary. Replaying the same window @@ -237,10 +254,11 @@ Tracked against the runtime API discussion, not yet settled: single required set too. - **`step()` returning a future**, for models with a dependency on their own output. `InferenceSession.step()` is currently synchronous. -- **`Input System` ownership.** The diagrams show it pulling events and - supporting mock input, so the Application has-a input system. - `InputCanonicalizer` is currently a pure function over a supplied window and - owns no source. This is the one deferred item inside the input scope. +- **`Input System` ownership.** The diagrams show it pulling events, so the + Application owns an input system. `InputCanonicalizer` is currently a pure + function over a supplied window and owns no source. Whether it needs to grow + one depends on the loop-ownership decision. Mock input and key binding are + handled (`ScriptedModality`, `DEFAULT_DRIVING_BINDINGS`). ## Owned Elsewhere diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index a27866893..ab303c745 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -8,11 +8,13 @@ """ from flashdreams.runtime.canonical import ( + DEFAULT_DRIVING_BINDINGS, DRIVER_COMMAND, DeviceConverter, DeviceConverterSchema, InputCanonicalizer, KeyboardToDriverCommand, + ScriptedModality, ) from flashdreams.runtime.config import ExecutionBackend, InferenceConfig, Precision from flashdreams.runtime.inputs import ( @@ -65,6 +67,7 @@ "check_mapping_set_compatibility", "combine_mapping_schemas", "DeclaresMappingSchema", + "DEFAULT_DRIVING_BINDINGS", "DeviceConverter", "DeviceConverterSchema", "DRIVER_COMMAND", @@ -92,6 +95,7 @@ "OutputTarget", "Precision", "RuntimeMetricSample", + "ScriptedModality", "SESSION_START_ONLY", "StepRequest", "StepResult", diff --git a/flashdreams/flashdreams/runtime/canonical.py b/flashdreams/flashdreams/runtime/canonical.py index f99ee8340..55f333ce7 100644 --- a/flashdreams/flashdreams/runtime/canonical.py +++ b/flashdreams/flashdreams/runtime/canonical.py @@ -22,8 +22,9 @@ from __future__ import annotations -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field +from types import MappingProxyType from typing import Any, Protocol, runtime_checkable from flashdreams.runtime._utils import freeze_mapping @@ -36,11 +37,27 @@ UserInputs, UserInputSchema, ) -from flashdreams.serving.realtime.input import ( - DRIVING_SUPPORTED_KEYS, - KeyboardState, - normalize_key, +from flashdreams.serving.realtime.input import KeyboardState, normalize_key + +DriverBindings = Mapping[str, frozenset[str]] + +DEFAULT_DRIVING_BINDINGS: DriverBindings = MappingProxyType( + { + "throttle": frozenset({"w", "up"}), + "brake": frozenset({"s", "down"}), + "steer_left": frozenset({"a", "left"}), + "steer_right": frozenset({"d", "right"}), + "stop": frozenset({"space"}), + "reverse": frozenset(), + } ) +"""Default key bindings for :class:`KeyboardToDriverCommand`. + +Bindings are data so a layout can be rebound without editing the converter, and +so the set of tracked keys is derived from them rather than declared twice. +""" + +_DRIVER_ACTIONS = frozenset(DEFAULT_DRIVING_BINDINGS) @dataclass(frozen=True, kw_only=True, slots=True) @@ -114,11 +131,25 @@ def __init__( self, *, name: str = "keyboard-to-driver-command", - supported_keys: frozenset[str] = DRIVING_SUPPORTED_KEYS, + bindings: DriverBindings = DEFAULT_DRIVING_BINDINGS, priority: int = 0, ) -> None: - self._supported_keys = supported_keys - self._state = KeyboardState(supported_keys=supported_keys) + unknown = sorted(set(bindings) - _DRIVER_ACTIONS) + if unknown: + raise ValueError( + f"Unknown driver actions in bindings: {unknown}. " + f"Supported actions: {sorted(_DRIVER_ACTIONS)}." + ) + self._bindings = { + action: frozenset(normalize_key(key) for key in bindings.get(action, ())) + for action in _DRIVER_ACTIONS + } + # Tracked keys are derived, so they cannot drift from the bindings and + # silently make an action unreachable. + self._supported_keys = frozenset( + key for keys in self._bindings.values() for key in keys + ) + self._state = KeyboardState(supported_keys=self._supported_keys) self._schema = DeviceConverterSchema( name=name, produces=DRIVER_COMMAND, @@ -161,22 +192,89 @@ def convert( ) pressed = {normalize_key(key) for key in self._state.snapshot()} + + def held(action: str) -> bool: + return bool(self._bindings[action] & pressed) + steer = 0.0 - if {"a", "left"} & pressed: + if held("steer_left"): steer += 1.0 - if {"d", "right"} & pressed: + if held("steer_right"): steer -= 1.0 return DRIVER_COMMAND.value( { - "throttle": 1.0 if {"w", "up"} & pressed else 0.0, - "brake": 1.0 if {"s", "down"} & pressed else 0.0, + "throttle": 1.0 if held("throttle") else 0.0, + "brake": 1.0 if held("brake") else 0.0, "steer": steer, - "stop": "space" in pressed, - "reverse": False, + "stop": held("stop"), + "reverse": held("reverse"), } ) +class ScriptedModality: + """Emit pre-authored canonical values, for benchmarks, replay, and tests. + + Mocking input should not require knowing the raw device vocabulary. This + converter consumes no raw capabilities, so it is feedable by any source + -- including an empty :class:`UserInputSchema` -- and application code is + identical between a real run and a scripted one. + + ``timeline`` is ``(start_s, value)`` pairs. Values are level-triggered and + held until the next entry begins, matching how live converters behave. An + entry applies to a window once it has begun by the window's end, and + ``None`` is returned for windows before the first entry. + """ + + def __init__( + self, + *, + modality: CanonicalModality, + timeline: Sequence[tuple[float, Mapping[str, Any]]], + name: str | None = None, + device_kind: str | None = "scripted", + priority: int = 0, + ) -> None: + entries = tuple(sorted(timeline, key=lambda entry: entry[0])) + for start_s, value in entries: + if start_s < 0: + raise ValueError("timeline start_s must be >= 0.") + modality.value(value) + self._entries = tuple( + (start_s, modality.value(value)) for start_s, value in entries + ) + self._modality = modality + self._schema = DeviceConverterSchema( + name=name or f"scripted-{modality.name}", + produces=modality, + device_kind=device_kind, + priority=priority, + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + # The timeline is a pure function of the window, so replay is + # deterministic without any state to clear. + return None + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + del user_inputs + current: Mapping[str, Any] | None = None + for start_s, value in self._entries: + if start_s < window.end_s: + current = value + else: + break + return current + + class InputCanonicalizer: """Registry of device converters plus the raw-to-canonical rewrite. diff --git a/flashdreams/tests/test_runtime_canonical.py b/flashdreams/tests/test_runtime_canonical.py index 64ee85d75..1ad48d39e 100644 --- a/flashdreams/tests/test_runtime_canonical.py +++ b/flashdreams/tests/test_runtime_canonical.py @@ -26,6 +26,7 @@ InputField, InputMappingSchema, KeyboardToDriverCommand, + ScriptedModality, TimeWindow, UserInputCapability, UserInputEvent, @@ -425,3 +426,165 @@ def run() -> list[dict[str, Any]]: ] assert run() == run() + + +# --- key bindings ------------------------------------------------------- + + +def test_bindings_are_data_and_can_be_rebound() -> None: + """A layout change must not require editing the converter.""" + azerty = InputCanonicalizer( + [ + KeyboardToDriverCommand( + bindings={ + "throttle": frozenset({"z"}), + "brake": frozenset({"s"}), + "steer_left": frozenset({"q"}), + "steer_right": frozenset({"d"}), + "stop": frozenset({"space"}), + } + ) + ] + ) + + canonical = azerty.canonicalize( + UserInputs(events=(_key("key_down", "z", 0.1),)), + window=WINDOW, + source_schema=KEYBOARD_SOURCE, + ) + + assert _command(canonical)["throttle"] == 1.0 + + +def test_tracked_keys_are_derived_so_an_action_cannot_go_unreachable() -> None: + """Declaring bindings and tracked keys separately used to disagree.""" + converter = KeyboardToDriverCommand( + bindings={"stop": frozenset({"escape"}), "throttle": frozenset({"w"})} + ) + canonicalizer = InputCanonicalizer([converter]) + + canonical = canonicalizer.canonicalize( + UserInputs(events=(_key("key_down", "escape", 0.1),)), + window=WINDOW, + source_schema=KEYBOARD_SOURCE, + ) + + assert _command(canonical)["stop"] is True + + +def test_unknown_driver_action_is_rejected() -> None: + with pytest.raises(ValueError, match="Unknown driver actions"): + KeyboardToDriverCommand(bindings={"turbo": frozenset({"t"})}) + + +def test_reverse_is_bindable() -> None: + canonicalizer = InputCanonicalizer( + [KeyboardToDriverCommand(bindings={"reverse": frozenset({"r"})})] + ) + + canonical = canonicalizer.canonicalize( + UserInputs(events=(_key("key_down", "r", 0.1),)), + window=WINDOW, + source_schema=KEYBOARD_SOURCE, + ) + + assert _command(canonical)["reverse"] is True + + +# --- scripted / mock input ---------------------------------------------- + + +def _scripted() -> InputCanonicalizer: + return InputCanonicalizer( + [ + ScriptedModality( + modality=DRIVER_COMMAND, + timeline=[ + ( + 0.0, + { + "throttle": 1.0, + "brake": 0.0, + "steer": 0.0, + "stop": False, + "reverse": False, + }, + ), + ( + 2.0, + { + "throttle": 0.0, + "brake": 0.0, + "steer": 1.0, + "stop": False, + "reverse": False, + }, + ), + ], + ) + ] + ) + + +def test_mock_input_needs_no_raw_events_or_source_schema() -> None: + """Authoring a benchmark scenario must not require raw device vocabulary.""" + canonical = _scripted().canonicalize( + UserInputs(), window=WINDOW, source_schema=UserInputSchema() + ) + + assert _command(canonical)["throttle"] == 1.0 + + +def test_scripted_values_hold_until_the_next_entry() -> None: + canonicalizer = _scripted() + windows = [TimeWindow(start_s=t, end_s=t + 1.0) for t in (0.0, 1.0, 2.0)] + + steer = [ + _command( + canonicalizer.canonicalize( + UserInputs(), window=w, source_schema=UserInputSchema() + ) + )["steer"] + for w in windows + ] + + assert steer == [0.0, 0.0, 1.0] + + +def test_scripted_converter_is_silent_before_its_first_entry() -> None: + canonicalizer = InputCanonicalizer( + [ + ScriptedModality( + modality=CanonicalModality(name="late", payload_fields=frozenset()), + timeline=[(5.0, {})], + ) + ] + ) + + canonical = canonicalizer.canonicalize( + UserInputs(), window=WINDOW, source_schema=UserInputSchema() + ) + + assert canonical.values == {} + + +def test_scripted_timeline_is_validated_against_the_modality() -> None: + with pytest.raises(ValueError, match="requires payload fields"): + ScriptedModality(modality=DRIVER_COMMAND, timeline=[(0.0, {"throttle": 1.0})]) + + +def test_scripted_replay_is_deterministic() -> None: + def run() -> list[float]: + canonicalizer = _scripted() + return [ + _command( + canonicalizer.canonicalize( + UserInputs(), + window=TimeWindow(start_s=t, end_s=t + 1.0), + source_schema=UserInputSchema(), + ) + )["steer"] + for t in (0.0, 1.0, 2.0) + ] + + assert run() == run()