diff --git a/.github/workflows/omnidreams-demo-runtime.yml b/.github/workflows/omnidreams-demo-runtime.yml
new file mode 100644
index 000000000..1f494fa83
--- /dev/null
+++ b/.github/workflows/omnidreams-demo-runtime.yml
@@ -0,0 +1,328 @@
+name: OmniDreams Demo Runtime
+
+on:
+ push:
+ branches:
+ - main
+ - "pull-request/[0-9]+"
+ paths:
+ - ".github/workflows/omnidreams-demo-runtime.yml"
+ - "pyproject.toml"
+ - "uv.lock"
+ - "flashdreams/pyproject.toml"
+ - "flashdreams/flashdreams/core/**"
+ - "flashdreams/flashdreams/infra/**"
+ - "flashdreams/flashdreams/runtime/**"
+ - "flashdreams/flashdreams/serving/**"
+ - "flashdreams/flashdreams/recipes/taehv/**"
+ - "flashdreams/flashdreams/recipes/wan/**"
+ - "integrations/omnidreams/**"
+ - "integrations/lingbot/**"
+ workflow_dispatch:
+
+permissions:
+ contents: read
+
+jobs:
+ demo-runtime:
+ name: null, precomputed MP4, and Ludus MP4
+ runs-on: linux-amd64-gpu-rtxpro6000-latest-2
+ timeout-minutes: 180
+ defaults:
+ run:
+ shell: bash
+ container:
+ image: nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04
+ options: --gpus all
+ env:
+ UV_PROJECT_ENVIRONMENT: /tmp/flashdreams-venv
+ UV_LINK_MODE: copy
+ UV_PYTHON: "3.10"
+ MAX_JOBS: 8
+ ARTIFACT_DIR: artifacts/omnidreams_demo_runtime
+ NULL_BLOCKS: "10"
+ PRECOMPUTED_BLOCKS: "75"
+ LUDUS_BLOCKS: "76"
+ FPS: "30"
+ EXAMPLE_DATA_UUID: 239560dc-33d1-11ef-9720-00044bcbccac
+ LUDUS_SCENE_UUID: 0d404ff7-2b66-498c-b047-1ed8cded60d4
+ LUDUS_TRACE: integrations/omnidreams/omnidreams/demo/traces/ludus_forward_sweep_60s.json
+ EXPECTED_WIDTH: "1280"
+ EXPECTED_HEIGHT: "704"
+ MIN_DURATION_SECONDS: "18"
+ MAX_DURATION_SECONDS: "22"
+ steps:
+ - name: Detect GPU architecture
+ id: gpu-arch
+ run: |
+ nvidia-smi
+ compute_cap=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -1 | tr -d '[:space:]')
+ arch=$(echo "${compute_cap}" | tr -d '.')
+ echo "arch=${arch}" >> "$GITHUB_OUTPUT"
+ echo "Detected GPU compute capability: ${compute_cap} -> sm_${arch}"
+
+ - name: Checkout
+ uses: actions/checkout@v4
+
+ - name: Install system dependencies
+ run: |
+ apt-get update -qq
+ DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends \
+ python3 python3-dev python3-venv \
+ ffmpeg \
+ gcc g++ ninja-build \
+ libnccl-dev \
+ curl git ca-certificates jq unzip
+ rm -rf /var/lib/apt/lists/*
+
+ - name: Setup proxy cache
+ uses: nv-gha-runners/setup-proxy-cache@main
+
+ - name: Setup uv
+ uses: astral-sh/setup-uv@v6
+ with:
+ enable-cache: true
+ cache-suffix: "omnidreams-demo-runtime-sm${{ steps.gpu-arch.outputs.arch }}"
+ prune-cache: false
+
+ - name: Install dependencies
+ env:
+ NVTE_CUDA_ARCHS: ${{ steps.gpu-arch.outputs.arch }}
+ run: |
+ uv venv --clear
+ uv sync --locked --extra dev
+
+ - name: Verify GPU availability
+ run: nvidia-smi
+
+ - name: Run OmniDreams demo modes
+ env:
+ HF_TOKEN: ${{ secrets.HF_TOKEN }}
+ run: |
+ set -uo pipefail
+
+ log_dir="${ARTIFACT_DIR}/logs"
+ output_dir="${ARTIFACT_DIR}/outputs"
+ summary="${ARTIFACT_DIR}/summary.md"
+ status_file="${ARTIFACT_DIR}/command-status.env"
+ mkdir -p "${log_dir}" "${output_dir}"
+ : > "${status_file}"
+
+ odemo() {
+ uv run --no-sync --package flashdreams-omnidreams omnidreams-demo "$@"
+ }
+
+ run_demo() {
+ local name="$1"
+ shift
+ local log="${log_dir}/${name}.log"
+
+ {
+ printf '$'
+ printf ' %q' "$@"
+ printf '\n\n'
+ "$@"
+ } 2>&1 | tee "${log}"
+
+ local rc="${PIPESTATUS[0]}"
+ echo "${name}=${rc}" >> "${status_file}"
+ echo "${name} exit code: ${rc}" | tee -a "${summary}"
+ return 0
+ }
+
+ {
+ echo "# OmniDreams Demo Runtime CI"
+ echo
+ echo "| Mode | Expected blocks | Output |"
+ echo "| --- | ---: | --- |"
+ echo "| null | ${NULL_BLOCKS} | none |"
+ echo "| precomputed MP4 | ${PRECOMPUTED_BLOCKS} | omnidreams-demo-precomputed-20s.mp4 |"
+ echo "| Ludus MP4 | ${LUDUS_BLOCKS} | omnidreams-demo-ludus-20s.mp4 |"
+ echo
+ echo "## Command Status"
+ } > "${summary}"
+
+ run_demo null \
+ odemo replay \
+ --output-mode null \
+ --device cuda:0 \
+ --total-blocks "${NULL_BLOCKS}"
+
+ run_demo precomputed-mp4 \
+ odemo replay \
+ --device cuda:0 \
+ --example-data \
+ --example-data-uuid "${EXAMPLE_DATA_UUID}" \
+ --total-blocks "${PRECOMPUTED_BLOCKS}" \
+ --fps "${FPS}" \
+ --output "${output_dir}/omnidreams-demo-precomputed-20s.mp4"
+
+ run_demo ludus-mp4 \
+ odemo replay \
+ --conditioning-mode ludus-scene-driving \
+ --keyboard-trace "${LUDUS_TRACE}" \
+ --device cuda:0 \
+ --scene-uuid "${LUDUS_SCENE_UUID}" \
+ --seed 42 \
+ --total-blocks "${LUDUS_BLOCKS}" \
+ --output "${output_dir}/omnidreams-demo-ludus-20s.mp4"
+
+ - name: Validate OmniDreams demo artifacts
+ run: |
+ set -euo pipefail
+
+ log_dir="${ARTIFACT_DIR}/logs"
+ output_dir="${ARTIFACT_DIR}/outputs"
+ probe_dir="${ARTIFACT_DIR}/ffprobe"
+ summary="${ARTIFACT_DIR}/summary.md"
+ status_file="${ARTIFACT_DIR}/command-status.env"
+ mkdir -p "${probe_dir}"
+
+ status_of() {
+ awk -F= -v name="$1" '$1 == name { print $2 }' "${status_file}"
+ }
+
+ assert_exit_zero() {
+ local name="$1"
+ local rc
+ rc="$(status_of "${name}")"
+ if [ "${rc}" != "0" ]; then
+ echo "${name} command failed with exit code ${rc}" >&2
+ exit 1
+ fi
+ }
+
+ assert_clean_log() {
+ local log="$1"
+ if grep -En "ERROR|Traceback|Exception|status=failed|Run failed|failed run" "${log}"; then
+ echo "failure marker found in ${log}" >&2
+ exit 1
+ fi
+ }
+
+ assert_log_contains() {
+ local log="$1"
+ local pattern="$2"
+ local label="$3"
+ if ! grep -Eq "${pattern}" "${log}"; then
+ echo "expected ${label} in ${log}" >&2
+ exit 1
+ fi
+ }
+
+ validate_mp4() {
+ local mode="$1"
+ local mp4="$2"
+ local metadata="${probe_dir}/${mode}.json"
+
+ if [ ! -s "${mp4}" ]; then
+ echo "expected non-empty MP4 at ${mp4}" >&2
+ exit 1
+ fi
+
+ ffprobe \
+ -v error \
+ -select_streams v:0 \
+ -show_entries stream=width,height,r_frame_rate,avg_frame_rate,nb_frames,duration:format=duration \
+ -of json \
+ "${mp4}" > "${metadata}"
+
+ local stream_count width height duration
+ stream_count="$(jq '.streams | length' "${metadata}")"
+ width="$(jq -r '.streams[0].width // ""' "${metadata}")"
+ height="$(jq -r '.streams[0].height // ""' "${metadata}")"
+ duration="$(jq -r '.streams[0].duration // .format.duration // "0"' "${metadata}")"
+
+ if [ "${stream_count}" -lt 1 ]; then
+ echo "ffprobe found no video stream in ${mp4}" >&2
+ exit 1
+ fi
+
+ if [ "${width}" != "${EXPECTED_WIDTH}" ] || [ "${height}" != "${EXPECTED_HEIGHT}" ]; then
+ echo "unexpected ${mode} resolution ${width}x${height}; expected ${EXPECTED_WIDTH}x${EXPECTED_HEIGHT}" >&2
+ exit 1
+ fi
+
+ awk \
+ -v duration="${duration}" \
+ -v min_duration="${MIN_DURATION_SECONDS}" \
+ -v max_duration="${MAX_DURATION_SECONDS}" \
+ 'BEGIN {
+ if ((duration + 0) < min_duration || (duration + 0) > max_duration) {
+ exit 1
+ }
+ }' || {
+ echo "unexpected ${mode} duration ${duration}s; expected ${MIN_DURATION_SECONDS}-${MAX_DURATION_SECONDS}s" >&2
+ exit 1
+ }
+ }
+
+ null_log="${log_dir}/null.log"
+ precomputed_log="${log_dir}/precomputed-mp4.log"
+ ludus_log="${log_dir}/ludus-mp4.log"
+
+ assert_exit_zero null
+ assert_exit_zero precomputed-mp4
+ assert_exit_zero ludus-mp4
+
+ assert_clean_log "${null_log}"
+ assert_clean_log "${precomputed_log}"
+ assert_clean_log "${ludus_log}"
+
+ assert_log_contains "${null_log}" "AR 9 encode" "null final AR block"
+ assert_log_contains "${null_log}" "OmniDreams demo replay step 9 frames=" "null final replay step"
+ assert_log_contains "${null_log}" "Loaded OmniDreams demo HDMaps shape=.*views=1" "null precomputed HDMaps"
+
+ assert_log_contains "${precomputed_log}" "AR 74 encode" "precomputed final AR block"
+ assert_log_contains "${precomputed_log}" "OmniDreams demo replay step 74 frames=" "precomputed final replay step"
+ assert_log_contains "${precomputed_log}" "Loaded OmniDreams demo HDMaps shape=.*views=1" "precomputed HDMaps"
+
+ assert_log_contains "${ludus_log}" "AR 75 encode" "Ludus final AR block"
+ assert_log_contains "${ludus_log}" "OmniDreams demo replay step 75 frames=" "Ludus final replay step"
+ assert_log_contains "${ludus_log}" "ludus_backend=cuda" "Ludus CUDA backend"
+ assert_log_contains "${ludus_log}" "trace_events=[1-9][0-9]*" "nonzero Ludus trace events"
+
+ validate_mp4 precomputed-mp4 "${output_dir}/omnidreams-demo-precomputed-20s.mp4"
+ validate_mp4 ludus-mp4 "${output_dir}/omnidreams-demo-ludus-20s.mp4"
+
+ {
+ echo
+ echo "## Validation"
+ echo
+ echo "- All commands exited zero."
+ echo "- Logs contained expected final AR blocks and provider markers."
+ echo "- MP4 outputs were non-empty and passed ffprobe stream checks."
+ } >> "${summary}"
+
+ - name: Trim uv cache for upload
+ if: always()
+ run: |
+ cache_dir="${UV_CACHE_DIR:-/github/home/.cache/uv}"
+ echo "=== Cache size before trim ==="
+ du -sh "${cache_dir}" || true
+ du -sh "${cache_dir}"/*/ 2>/dev/null || true
+
+ rm -rf "${cache_dir}/wheels-v6"
+ rm -rf "${cache_dir}/archive-v0"
+
+ find "${cache_dir}/git-v0/checkouts" \
+ \( -name "build" -o -name "*.egg-info" -o -name "__pycache__" \) \
+ -type d -exec rm -rf {} + 2>/dev/null || true
+
+ rm -rf "${cache_dir}/sdists-v9/editable"
+
+ echo ""
+ echo "=== Cache size after trim ==="
+ du -sh "${cache_dir}" || true
+ du -sh "${cache_dir}"/*/ 2>/dev/null || true
+ echo ""
+ echo "=== Cached built wheels (sdists-v9) ==="
+ find "${cache_dir}/sdists-v9" -name "*.whl" -exec ls -lh {} \; 2>/dev/null || true
+
+ - name: Upload demo runtime artifacts
+ if: always()
+ uses: actions/upload-artifact@v4
+ with:
+ name: omnidreams-demo-runtime
+ path: ${{ env.ARTIFACT_DIR }}
+ if-no-files-found: ignore
diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py
index fb6eb4b05..6e205823c 100644
--- a/flashdreams/flashdreams/runtime/__init__.py
+++ b/flashdreams/flashdreams/runtime/__init__.py
@@ -52,14 +52,20 @@
from flashdreams.runtime.metrics import (
InMemoryMetricsRecorder,
MetricsRecorder,
+ MetricsSnapshot,
NullMetricsRecorder,
RuntimeMetricSample,
)
from flashdreams.runtime.output import NullOutputTarget, OutputArtifact, OutputTarget
from flashdreams.runtime.runner import run_inference_session
-from flashdreams.runtime.types import StepRequest, StepResult
+from flashdreams.runtime.types import (
+ StepRequest,
+ StepRequirements,
+ StepResult,
+ step_requirements_from_request,
+)
from flashdreams.runtime.video_output import Mp4VideoOutputTarget
-from flashdreams.runtime.worker import ThreadAffineRuntimeWorker
+from flashdreams.runtime.worker import ModelExecutionWorker, ThreadAffineRuntimeWorker
__all__ = [
"CanonicalInputs",
@@ -90,7 +96,9 @@
"KeyboardToDriverCommand",
"MappingCompatibility",
"MetricsRecorder",
+ "MetricsSnapshot",
"ModelAdapter",
+ "ModelExecutionWorker",
"Mp4VideoOutputTarget",
"NullMetricsRecorder",
"NullOutputTarget",
@@ -100,10 +108,12 @@
"RuntimeMetricSample",
"ScriptedModality",
"StepRequest",
+ "StepRequirements",
"StepResult",
"TimeWindow",
"ThreadAffineRuntimeWorker",
"run_inference_session",
+ "step_requirements_from_request",
"undeclared_inference_inputs",
"UserInputCapability",
"UserInputEvent",
diff --git a/flashdreams/flashdreams/runtime/config.py b/flashdreams/flashdreams/runtime/config.py
index 4b8752f13..f1c0c2a0c 100644
--- a/flashdreams/flashdreams/runtime/config.py
+++ b/flashdreams/flashdreams/runtime/config.py
@@ -61,6 +61,9 @@ class InferenceConfig:
cache_policy: str | None = None
"""Optional cache policy selector; ``None`` leaves the choice to the adapter."""
+ seed: int | None = None
+ """Optional seed used when resolving deterministic demo/runtime behavior."""
+
runtime_options: Mapping[str, Any] = field(default_factory=dict)
"""Adapter/backend-specific runtime options."""
@@ -70,6 +73,11 @@ class InferenceConfig:
def __post_init__(self) -> None:
if not self.model_id.strip():
raise ValueError("InferenceConfig.model_id must be non-empty.")
+ if self.seed is not None:
+ if isinstance(self.seed, bool) or not isinstance(self.seed, int):
+ raise TypeError("InferenceConfig.seed must be an integer.")
+ if self.seed < 0:
+ raise ValueError("InferenceConfig.seed must be >= 0.")
object.__setattr__(
self, "runtime_options", freeze_mapping(self.runtime_options)
)
diff --git a/flashdreams/flashdreams/runtime/demo/__init__.py b/flashdreams/flashdreams/runtime/demo/__init__.py
index 7a0535556..60187237c 100644
--- a/flashdreams/flashdreams/runtime/demo/__init__.py
+++ b/flashdreams/flashdreams/runtime/demo/__init__.py
@@ -3,11 +3,70 @@
"""Experimental shared demo API above the inference runtime API."""
-from flashdreams.runtime.demo.outputs import build_output_target
-from flashdreams.runtime.demo.replay import run_replay_demo
+from flashdreams.runtime.demo.drivers import (
+ CLEANUP_TIMEOUT_S,
+ BatchSessionDriver,
+ DriverInvariantError,
+ RealtimeSessionDriver,
+ run_demo_session,
+ run_demo_session_async,
+ shielded_session_cleanup,
+ uncancel_current_task,
+)
+from flashdreams.runtime.demo.host import (
+ ModelWarmupPlan,
+ RuntimeHost,
+ WarmupSessionInputs,
+)
+from flashdreams.runtime.demo.outputs import (
+ Mp4OutputSink,
+ NullOutputSink,
+ OutputDecision,
+ OutputSink,
+ SessionInfo,
+ build_output_sink,
+ build_output_target,
+)
+from flashdreams.runtime.demo.pipeline import StepOutcome, StepPipeline
+from flashdreams.runtime.demo.replay import OutputSinkFactory, run_replay_demo
+from flashdreams.runtime.demo.run_modes import (
+ AsyncSessionDriver,
+ BenchmarkErrorPolicy,
+ DefaultErrorPolicy,
+ ErrorAction,
+ InMemorySessionMetricsRecorder,
+ MetricsSnapshot,
+ Mp4ErrorPolicy,
+ NativeWindowErrorPolicy,
+ NoopTransportService,
+ NullErrorPolicy,
+ RunContext,
+ RunMode,
+ RunModeCapabilities,
+ RunModeWarmup,
+ RunResult,
+ RunSummary,
+ SessionDriver,
+ SessionEdges,
+ SingleSessionAdmissionPolicy,
+ WebRTCErrorPolicy,
+ build_model_warmup_plan,
+ warmup_run_context,
+)
+from flashdreams.runtime.demo.session_inputs import (
+ BatchInputSource,
+ ControlDecision,
+ InputSource,
+ ModelInputProvider,
+ PreparedStep,
+ ProviderCapabilities,
+ RealtimeInputSource,
+ UserInputWindow,
+)
from flashdreams.runtime.demo.spec import (
DemoAdapter,
DemoSpec,
+ ModelWarmupAdapter,
Mp4OutputSpec,
NullOutputSpec,
OutputSpec,
@@ -15,16 +74,105 @@
WebRTCAppResources,
WebRTCOutputSpec,
)
+from flashdreams.runtime.demo.timing import (
+ SPARSE_KEY_SEGMENTS_METADATA_KEY,
+ ActivationPolicy,
+ ActivationResult,
+ ActivationSignal,
+ AlwaysActiveActivationPolicy,
+ CatchUpDecision,
+ CatchUpPolicy,
+ DeterministicClock,
+ KeyboardRealtimeInputSource,
+ RealtimeClock,
+ RealtimeWindowResult,
+ ResamplerRealtimeClock,
+ SignalActivationPolicy,
+ input_frame_count_from_request,
+)
+from flashdreams.runtime.demo.validation import (
+ ResolvedRunCapabilities,
+ resolve_run_capabilities,
+ validate_resolved_run,
+)
__all__ = [
+ "BatchInputSource",
+ "BatchSessionDriver",
+ "CLEANUP_TIMEOUT_S",
+ "ControlDecision",
+ "DefaultErrorPolicy",
"DemoAdapter",
"DemoSpec",
+ "DriverInvariantError",
+ "ErrorAction",
+ "AsyncSessionDriver",
+ "ActivationPolicy",
+ "ActivationResult",
+ "ActivationSignal",
+ "AlwaysActiveActivationPolicy",
+ "BenchmarkErrorPolicy",
+ "InMemorySessionMetricsRecorder",
+ "InputSource",
+ "CatchUpDecision",
+ "CatchUpPolicy",
+ "DeterministicClock",
+ "MetricsSnapshot",
+ "ModelWarmupAdapter",
+ "ModelWarmupPlan",
+ "ModelInputProvider",
+ "KeyboardRealtimeInputSource",
+ "Mp4ErrorPolicy",
+ "Mp4OutputSink",
"Mp4OutputSpec",
+ "NativeWindowErrorPolicy",
+ "NoopTransportService",
"NullOutputSpec",
+ "NullOutputSink",
+ "NullErrorPolicy",
+ "OutputDecision",
+ "OutputSinkFactory",
"OutputSpec",
+ "OutputSink",
"PreparedScenario",
+ "PreparedStep",
+ "ProviderCapabilities",
+ "RealtimeInputSource",
+ "RealtimeClock",
+ "RealtimeSessionDriver",
+ "RealtimeWindowResult",
+ "ResolvedRunCapabilities",
+ "ResamplerRealtimeClock",
+ "RunContext",
+ "RunMode",
+ "RunModeCapabilities",
+ "RunModeWarmup",
+ "RunResult",
+ "RunSummary",
+ "RuntimeHost",
+ "SessionEdges",
+ "SessionDriver",
+ "SessionInfo",
+ "SignalActivationPolicy",
+ "SingleSessionAdmissionPolicy",
+ "SPARSE_KEY_SEGMENTS_METADATA_KEY",
+ "StepOutcome",
+ "StepPipeline",
+ "UserInputWindow",
+ "WarmupSessionInputs",
"WebRTCAppResources",
+ "WebRTCErrorPolicy",
"WebRTCOutputSpec",
+ "build_output_sink",
"build_output_target",
+ "build_model_warmup_plan",
+ "input_frame_count_from_request",
+ "resolve_run_capabilities",
+ "run_demo_session",
+ "run_demo_session_async",
"run_replay_demo",
+ "shielded_session_cleanup",
+ "uncancel_current_task",
+ "validate_resolved_run",
+ "warmup_run_context",
]
diff --git a/flashdreams/flashdreams/runtime/demo/app.py b/flashdreams/flashdreams/runtime/demo/app.py
index b84e59163..d44dd81bb 100644
--- a/flashdreams/flashdreams/runtime/demo/app.py
+++ b/flashdreams/flashdreams/runtime/demo/app.py
@@ -6,6 +6,7 @@
from __future__ import annotations
import argparse
+import sys
from abc import ABC, abstractmethod
from typing import Any
@@ -29,10 +30,18 @@ def main(self, argv: list[str] | None = None) -> None:
configure_logging()
args = self.parse_args(argv)
if args.command == "replay":
- run_replay_demo(
+ result = run_replay_demo(
spec=self.replay_spec(args),
adapter=self.replay_adapter(),
)
+ if result.status != "completed":
+ reason = result.reason or (
+ str(result.error) if result.error is not None else None
+ )
+ if reason is None:
+ reason = f"Replay demo ended with status {result.status!r}."
+ print(reason, file=sys.stderr)
+ raise SystemExit(1)
return
if args.command == "webrtc":
context = initialize_cuda_distributed(
diff --git a/flashdreams/flashdreams/runtime/demo/drivers.py b/flashdreams/flashdreams/runtime/demo/drivers.py
new file mode 100644
index 000000000..9bdd40db8
--- /dev/null
+++ b/flashdreams/flashdreams/runtime/demo/drivers.py
@@ -0,0 +1,1011 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Session drivers and helpers for demo runtime vertical slices."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+from typing import Any, cast
+
+from flashdreams.runtime.interfaces import InferenceSession
+from flashdreams.runtime.types import (
+ StepRequest,
+ StepRequirements,
+ step_requirements_from_request,
+)
+
+from .host import RuntimeHost
+from .outputs import SessionInfo
+from .pipeline import StepPipeline
+from .run_modes import (
+ DriverStatus,
+ RunContext,
+ RunMode,
+ RunResult,
+ SessionEdges,
+ SessionReservation,
+)
+from .session_inputs import BatchInputSource, ModelInputProvider
+from .spec import DemoAdapter, DemoSpec, PreparedScenario
+from .timing import ActivationPolicy, RealtimeClock
+from .validation import resolve_run_capabilities, validate_resolved_run
+
+CLEANUP_TIMEOUT_S = 30.0
+_MODEL_CLEANUP_FAILED_REASON = "model-affine cleanup failed"
+_MODEL_CLEANUP_TIMED_OUT_REASON = "model-affine cleanup timed out"
+
+
+class DriverInvariantError(RuntimeError):
+ """A driver invariant was violated; this is a driver bug, not a run result."""
+
+
+class BatchSessionDriver:
+ """Minimal finite-session driver for Phase 2 fake-model coverage."""
+
+ def run_one_session(
+ self,
+ *,
+ host: RuntimeHost,
+ provider: ModelInputProvider,
+ session_edges: SessionEdges,
+ pipeline: StepPipeline,
+ ) -> RunResult:
+ session: InferenceSession | None = None
+ final_status: DriverStatus = "completed"
+ final_reason: str | None = None
+ final_error: Exception | None = None
+ invariant_closed = False
+ setup_ok = False
+ try:
+ try:
+ initial_input = host.call(provider.prepare_initial_input)
+ session = host.call(host.start_session, initial_input)
+ session_info = host.call(_session_info, session)
+ session_edges.output_sink.open(session_info)
+ setup_ok = True
+ except Exception as exc:
+ action = session_edges.error_policy.handle_setup_error(exc)
+ if action.drop_chunk or action.result_status == "completed":
+ raise DriverInvariantError(
+ "Setup failures must resolve to failed or skipped."
+ ) from exc
+ session_edges.metrics.record_error(exc, action)
+ final_status = action.result_status
+ final_reason = str(exc)
+ final_error = exc if action.result_status == "failed" else None
+
+ input_source = cast(BatchInputSource, session_edges.input_source)
+ while setup_ok:
+ if session is None:
+ raise DriverInvariantError("setup_ok was set without a session.")
+ try:
+ if session_edges.input_source.is_finished():
+ break
+ request = _next_step_requirements(host=host, session=session)
+ if request is None:
+ break
+ user_window = input_source.next_window(request)
+ outcome = host.call(
+ pipeline.execute_step,
+ request=request,
+ user_window=user_window,
+ provider=provider,
+ session=session,
+ output=session_edges.output_sink,
+ metrics=session_edges.metrics,
+ )
+ if outcome.control.reset:
+ host.call(session.reset, outcome.control.reset_input)
+ if not outcome.control.provider_already_reset:
+ host.call(provider.reset, outcome.control.reset_input)
+ continue
+ if outcome.control.close_session:
+ break
+ if outcome.output.should_stop:
+ break
+ except DriverInvariantError:
+ raise
+ except Exception as exc:
+ action = session_edges.error_policy.handle(exc)
+ session_edges.metrics.record_error(exc, action)
+ if action.drop_chunk:
+ continue
+ final_status = action.result_status
+ final_reason = str(exc)
+ final_error = exc if action.result_status == "failed" else None
+ break
+ except DriverInvariantError as exc:
+ if session is not None:
+ _close_on_host_best_effort(
+ host=host,
+ close=session.close,
+ session_edges=session_edges,
+ )
+ _close_on_host_best_effort(
+ host=host,
+ close=provider.close,
+ session_edges=session_edges,
+ )
+ session_edges.close_result(
+ status="failed",
+ reason=str(exc),
+ error=exc,
+ )
+ invariant_closed = True
+ raise
+ except Exception as exc:
+ final_status = "failed"
+ final_reason = str(exc)
+ final_error = exc
+ finally:
+ if not invariant_closed:
+ if session is not None:
+ _close_on_host_best_effort(
+ host=host,
+ close=session.close,
+ session_edges=session_edges,
+ )
+ _close_on_host_best_effort(
+ host=host,
+ close=provider.close,
+ session_edges=session_edges,
+ )
+
+ return session_edges.close_result(
+ status=final_status,
+ reason=final_reason,
+ error=final_error,
+ )
+
+
+class RealtimeSessionDriver:
+ """Async realtime session driver built on shared Phase 5 primitives."""
+
+ cleanup_timeout_s: float
+
+ def __init__(self, *, cleanup_timeout_s: float = CLEANUP_TIMEOUT_S) -> None:
+ if cleanup_timeout_s <= 0:
+ raise ValueError("cleanup_timeout_s must be > 0.")
+ self.cleanup_timeout_s = float(cleanup_timeout_s)
+
+ async def run_one_session(
+ self,
+ *,
+ host: RuntimeHost,
+ provider: ModelInputProvider,
+ session_edges: SessionEdges,
+ pipeline: StepPipeline,
+ ) -> RunResult:
+ session: InferenceSession | None = None
+ final_status: DriverStatus = "completed"
+ final_reason: str | None = None
+ final_error: Exception | None = None
+ setup_ok = False
+ generation = 0
+ first_step_started = False
+ invariant_error: DriverInvariantError | None = None
+ try:
+ activation, clock = _realtime_activation_and_clock(session_edges)
+ input_source = _realtime_input_source(session_edges)
+ activation_result = await activation.wait_until_active(clock)
+ if not activation_result.activated:
+ final_status = "not_activated"
+ final_reason = activation_result.reason
+ elif not session_edges.transport.is_active():
+ final_status = "not_activated"
+ final_reason = "transport closed before first step"
+ else:
+ try:
+ initial_input = await host.call_async(
+ provider.prepare_initial_input
+ )
+ session = await host.call_async(host.start_session, initial_input)
+ session_info = await host.call_async(_session_info, session)
+ session_edges.output_sink.open(session_info)
+ session_edges.output_sink.begin_generation(generation)
+ setup_ok = True
+ except Exception as exc:
+ action = session_edges.error_policy.handle_setup_error(exc)
+ if action.drop_chunk or action.result_status == "completed":
+ raise DriverInvariantError(
+ "Setup failures must resolve to failed or skipped."
+ ) from exc
+ session_edges.metrics.record_error(exc, action)
+ final_status = action.result_status
+ final_reason = str(exc)
+ final_error = exc if action.result_status == "failed" else None
+
+ while setup_ok:
+ if session is None:
+ raise DriverInvariantError("setup_ok was set without a session.")
+ if not session_edges.transport.is_active():
+ if not first_step_started:
+ final_status = "not_activated"
+ final_reason = "transport closed before first step"
+ break
+ try:
+ request = await _next_step_requirements_async(
+ host=host,
+ session=session,
+ )
+ if request is None:
+ break
+ window_result = await input_source.next_realtime_window(
+ request=request,
+ clock=clock,
+ )
+ session_edges.metrics.record_catch_up(window_result.catch_up)
+ if (
+ not session_edges.transport.is_active()
+ and not first_step_started
+ ):
+ final_status = "not_activated"
+ final_reason = "transport closed before first step"
+ break
+ outcome = await host.call_async(
+ pipeline.execute_step,
+ request=request,
+ user_window=window_result.window,
+ provider=provider,
+ session=session,
+ output=session_edges.output_sink,
+ metrics=session_edges.metrics,
+ )
+ first_step_started = True
+ if outcome.control.reset:
+ await host.call_async(
+ session.reset,
+ outcome.control.reset_input,
+ )
+ if not outcome.control.provider_already_reset:
+ await host.call_async(
+ provider.reset,
+ outcome.control.reset_input,
+ )
+ generation += 1
+ session_edges.output_sink.begin_generation(generation)
+ continue
+ if outcome.control.close_session:
+ break
+ if outcome.output.should_stop:
+ break
+ if outcome.output.backpressure_s > 0:
+ await clock.apply_backpressure(outcome.output.backpressure_s)
+ except DriverInvariantError:
+ raise
+ except Exception as exc:
+ action = session_edges.error_policy.handle(exc)
+ session_edges.metrics.record_error(exc, action)
+ if action.close_session:
+ final_status = action.result_status
+ final_reason = str(exc)
+ final_error = exc if action.result_status == "failed" else None
+ break
+ if action.drop_chunk:
+ continue
+ final_status = "failed"
+ final_reason = str(exc)
+ final_error = exc
+ break
+ except asyncio.CancelledError:
+ uncancel_current_task()
+ final_status = "cancelled"
+ final_reason = (
+ "cancelled before first step" if session is None else "cancelled"
+ )
+ final_error = None
+ except DriverInvariantError as exc:
+ invariant_error = exc
+ final_status = "failed"
+ final_reason = str(exc)
+ final_error = exc
+ except Exception as exc:
+ final_status = "failed"
+ final_reason = str(exc)
+ final_error = exc
+
+ result = await shielded_session_cleanup(
+ host=host,
+ session=session,
+ provider=provider,
+ session_edges=session_edges,
+ status=final_status,
+ reason=final_reason,
+ error=final_error,
+ timeout_s=self.cleanup_timeout_s,
+ )
+ if invariant_error is not None:
+ raise invariant_error
+ return result
+
+
+def _mark_host_cleanup_failed(host: RuntimeHost, exc: Exception | None = None) -> None:
+ host.mark_unhealthy(_MODEL_CLEANUP_FAILED_REASON, exc)
+
+
+def run_demo_session(
+ *,
+ context: RunContext,
+ spec: DemoSpec,
+ scenario: PreparedScenario,
+ adapter: DemoAdapter,
+ run_mode: RunMode,
+ pipeline: StepPipeline,
+ reservation: SessionReservation | None = None,
+) -> RunResult:
+ """Run one prepared demo session through a selected run mode."""
+ if reservation is None:
+ reservation = context.admission.try_reserve()
+ if reservation is None:
+ result = RunResult.rejected(reason="busy")
+ context.run_metrics.record_session(result)
+ return result
+
+ provider: Any | None = None
+ session_edges: SessionEdges | None = None
+ driver_started = False
+ try:
+ create_provider = getattr(adapter, "create_model_input_provider")
+ provider = context.host.call(create_provider, spec, scenario)
+ run_mode.validate_session(
+ spec=spec,
+ scenario=scenario,
+ adapter=adapter,
+ provider=provider,
+ )
+ session_edges = run_mode.create_session_edges(
+ context=context,
+ spec=spec,
+ scenario=scenario,
+ provider=provider,
+ adapter=adapter,
+ )
+ resolved_capabilities = resolve_run_capabilities(
+ spec=spec,
+ provider=provider,
+ session_edges=session_edges,
+ )
+ validate_resolved_run(
+ spec=spec,
+ adapter=adapter,
+ provider=provider,
+ run_mode=run_mode,
+ session_edges=session_edges,
+ resolved=resolved_capabilities,
+ )
+ if session_edges.is_closed:
+ raise DriverInvariantError(
+ "RunMode returned already closed SessionEdges; session edges "
+ "must not be reused."
+ )
+ driver = run_mode.select_driver()
+ driver_started = True
+ result = _run_sync_driver(
+ driver=driver,
+ host=context.host,
+ provider=provider,
+ session_edges=session_edges,
+ pipeline=pipeline,
+ )
+ context.run_metrics.record_session(result)
+ return result
+ except DriverInvariantError as exc:
+ _record_run_session_error(context, exc)
+ if provider is not None and not driver_started:
+ _close_partial_provider_sync(
+ context=context,
+ provider=provider,
+ session_edges=session_edges,
+ )
+ if session_edges is not None and (
+ driver_started or not session_edges.is_closed
+ ):
+ result = session_edges.close_result(
+ status="failed",
+ reason=str(exc),
+ error=exc,
+ )
+ context.run_metrics.record_session(result)
+ raise
+ except Exception as exc:
+ _record_run_session_error(context, exc)
+ if provider is not None and not driver_started:
+ _close_partial_provider_sync(
+ context=context,
+ provider=provider,
+ session_edges=session_edges,
+ )
+ if session_edges is not None and (
+ driver_started or not session_edges.is_closed
+ ):
+ result = session_edges.close_result(
+ status="failed",
+ reason=str(exc),
+ error=exc,
+ )
+ else:
+ result = RunResult(status="failed", reason=str(exc), error=exc)
+ context.run_metrics.record_session(result)
+ return result
+ finally:
+ reservation.release()
+
+
+async def run_demo_session_async(
+ *,
+ context: RunContext,
+ spec: DemoSpec,
+ scenario: PreparedScenario,
+ adapter: DemoAdapter,
+ run_mode: RunMode,
+ pipeline: StepPipeline,
+ reservation: SessionReservation | None = None,
+) -> RunResult:
+ """Run one prepared async/realtime demo session through a selected run mode."""
+ if reservation is None:
+ reservation = context.admission.try_reserve()
+ if reservation is None:
+ result = RunResult.rejected(reason="busy")
+ context.run_metrics.record_session(result)
+ return result
+
+ provider: Any | None = None
+ session_edges: SessionEdges | None = None
+ try:
+ try:
+ create_provider = getattr(adapter, "create_model_input_provider")
+ provider = await context.host.call_async(create_provider, spec, scenario)
+ run_mode.validate_session(
+ spec=spec,
+ scenario=scenario,
+ adapter=adapter,
+ provider=provider,
+ )
+ session_edges = run_mode.create_session_edges(
+ context=context,
+ spec=spec,
+ scenario=scenario,
+ provider=provider,
+ adapter=adapter,
+ )
+ resolved_capabilities = resolve_run_capabilities(
+ spec=spec,
+ provider=provider,
+ session_edges=session_edges,
+ )
+ validate_resolved_run(
+ spec=spec,
+ adapter=adapter,
+ provider=provider,
+ run_mode=run_mode,
+ session_edges=session_edges,
+ resolved=resolved_capabilities,
+ )
+ if session_edges.is_closed:
+ raise DriverInvariantError(
+ "RunMode returned already closed SessionEdges; session edges "
+ "must not be reused."
+ )
+ driver = run_mode.select_driver()
+ result = await _run_async_driver(
+ driver=driver,
+ host=context.host,
+ provider=provider,
+ session_edges=session_edges,
+ pipeline=pipeline,
+ )
+ context.run_metrics.record_session(result)
+ return result
+ except asyncio.CancelledError:
+ uncancel_current_task()
+ result = await _close_partial_session_async(
+ context=context,
+ provider=provider,
+ session_edges=session_edges,
+ status="cancelled",
+ reason="cancelled during session assembly",
+ error=None,
+ close_provider=_needs_partial_provider_cleanup(session_edges),
+ )
+ context.run_metrics.record_session(result)
+ return result
+ except DriverInvariantError as exc:
+ _record_run_session_error(context, exc)
+ should_record_session = session_edges is not None
+ result = await _close_partial_session_async(
+ context=context,
+ provider=provider,
+ session_edges=session_edges,
+ status="failed",
+ reason=str(exc),
+ error=exc,
+ close_provider=_needs_partial_provider_cleanup(session_edges),
+ )
+ if should_record_session:
+ context.run_metrics.record_session(result)
+ raise
+ except Exception as exc:
+ _record_run_session_error(context, exc)
+ result = await _close_partial_session_async(
+ context=context,
+ provider=provider,
+ session_edges=session_edges,
+ status="failed",
+ reason=str(exc),
+ error=exc,
+ close_provider=_needs_partial_provider_cleanup(session_edges),
+ )
+ context.run_metrics.record_session(result)
+ return result
+ finally:
+ reservation.release()
+
+
+def _session_info(session: InferenceSession) -> SessionInfo:
+ session_info = getattr(session, "session_info", None)
+ if not callable(session_info):
+ return SessionInfo()
+ value = session_info()
+ if not isinstance(value, SessionInfo):
+ raise TypeError(
+ "session.session_info() must return SessionInfo, "
+ f"got {type(value).__name__}."
+ )
+ return value
+
+
+def _next_step_requirements(
+ *,
+ host: RuntimeHost,
+ session: InferenceSession,
+) -> StepRequirements | None:
+ next_requirements = getattr(session, "next_step_requirements", None)
+ if callable(next_requirements):
+ return _coerce_step_requirements(host.call(next_requirements))
+
+ next_request = getattr(session, "next_step_request", None)
+ if not callable(next_request):
+ raise TypeError(
+ "InferenceSession must provide next_step_requirements() or "
+ "legacy next_step_request()."
+ )
+ return _coerce_step_requirements(host.call(next_request))
+
+
+async def _next_step_requirements_async(
+ *,
+ host: RuntimeHost,
+ session: InferenceSession,
+) -> StepRequirements | None:
+ next_requirements = getattr(session, "next_step_requirements", None)
+ if callable(next_requirements):
+ return _coerce_step_requirements(await host.call_async(next_requirements))
+
+ next_request = getattr(session, "next_step_request", None)
+ if not callable(next_request):
+ raise TypeError(
+ "InferenceSession must provide next_step_requirements() or "
+ "legacy next_step_request()."
+ )
+ return _coerce_step_requirements(await host.call_async(next_request))
+
+
+def _coerce_step_requirements(value: object) -> StepRequirements | None:
+ if value is None:
+ return None
+ if isinstance(value, StepRequirements):
+ return value
+ if isinstance(value, StepRequest):
+ return step_requirements_from_request(value)
+ raise TypeError(
+ "Session next-step method must return StepRequirements, legacy "
+ f"StepRequest, or None; got {type(value).__name__}."
+ )
+
+
+def _close_safely(close: Any, session_edges: SessionEdges) -> bool:
+ try:
+ close()
+ except Exception as exc:
+ session_edges.record_cleanup_error(exc)
+ return False
+ return True
+
+
+def _close_on_host_best_effort(
+ *,
+ host: RuntimeHost,
+ close: Any,
+ session_edges: SessionEdges,
+) -> bool:
+ try:
+ cleanup_succeeded = host.call(_close_safely, close, session_edges)
+ except Exception as exc:
+ # If the host/worker is already unavailable, do not fall back to calling
+ # model-affine cleanup directly on the caller thread. Record the loss and
+ # let close_result finalize output, transport, and metrics.
+ session_edges.record_cleanup_error(exc)
+ _mark_host_cleanup_failed(host, exc)
+ return False
+ if not cleanup_succeeded:
+ _mark_host_cleanup_failed(host)
+ return False
+ return True
+
+
+def _close_partial_provider_sync(
+ *,
+ context: RunContext,
+ provider: Any,
+ session_edges: SessionEdges | None,
+) -> None:
+ if session_edges is not None:
+ _close_on_host_best_effort(
+ host=context.host,
+ close=provider.close,
+ session_edges=session_edges,
+ )
+ return
+ try:
+ cleanup_succeeded = context.host.call(
+ _close_run_provider_safely,
+ provider.close,
+ context,
+ )
+ except Exception as exc:
+ _record_run_cleanup_error(context, exc)
+ _mark_host_cleanup_failed(context.host, exc)
+ return
+ if not cleanup_succeeded:
+ _mark_host_cleanup_failed(context.host)
+
+
+def _close_run_provider_safely(close: Any, context: RunContext) -> bool:
+ try:
+ close()
+ except Exception as exc:
+ _record_run_cleanup_error(context, exc)
+ return False
+ return True
+
+
+async def shielded_session_cleanup(
+ *,
+ host: RuntimeHost,
+ session: InferenceSession | None,
+ provider: ModelInputProvider,
+ session_edges: SessionEdges,
+ status: DriverStatus,
+ reason: str | None,
+ error: Exception | None,
+ timeout_s: float = CLEANUP_TIMEOUT_S,
+) -> RunResult:
+ """Close realtime session resources exactly once without leaking cancellation."""
+
+ if timeout_s <= 0:
+ session_edges.record_cleanup_error(ValueError("timeout_s must be > 0."))
+ return session_edges.close_result(status=status, reason=reason, error=error)
+
+ async def cleanup() -> RunResult:
+ unhealthy_reason = await _close_model_resources_async(
+ host=host,
+ session=session,
+ provider=provider,
+ session_edges=session_edges,
+ timeout_s=timeout_s,
+ )
+ if unhealthy_reason is not None:
+ host.mark_unhealthy(unhealthy_reason)
+ return session_edges.close_result(
+ status=status,
+ reason=reason,
+ error=error,
+ )
+
+ try:
+ cleanup_task = asyncio.create_task(cleanup())
+ except RuntimeError as exc:
+ session_edges.record_cleanup_error(exc)
+ return session_edges.close_result(status=status, reason=reason, error=error)
+
+ session_edges.cleanup_tasks.add(cleanup_task)
+ try:
+ while not cleanup_task.done():
+ try:
+ await asyncio.shield(cleanup_task)
+ except asyncio.CancelledError:
+ uncancel_current_task()
+ continue
+ except Exception:
+ break
+ return _cleanup_result(cleanup_task, session_edges, status, reason, error)
+ finally:
+ session_edges.cleanup_tasks.discard(cleanup_task)
+
+
+def _run_sync_driver(
+ *,
+ driver: object,
+ host: RuntimeHost,
+ provider: ModelInputProvider,
+ session_edges: SessionEdges,
+ pipeline: StepPipeline,
+) -> RunResult:
+ run_one_session = getattr(driver, "run_one_session", None)
+ if not callable(run_one_session):
+ raise TypeError(
+ "RunMode.select_driver() must return an object with run_one_session(...)."
+ )
+ result = run_one_session(
+ host=host,
+ provider=provider,
+ session_edges=session_edges,
+ pipeline=pipeline,
+ )
+ if inspect.isawaitable(result):
+ raise TypeError(
+ "run_demo_session(...) requires a synchronous session driver; "
+ "use run_demo_session_async(...) for async drivers."
+ )
+ if not isinstance(result, RunResult):
+ raise TypeError(
+ "Session driver run_one_session(...) must return RunResult, "
+ f"got {type(result).__name__}."
+ )
+ return result
+
+
+async def _run_async_driver(
+ *,
+ driver: object,
+ host: RuntimeHost,
+ provider: ModelInputProvider,
+ session_edges: SessionEdges,
+ pipeline: StepPipeline,
+) -> RunResult:
+ run_one_session = getattr(driver, "run_one_session", None)
+ if not callable(run_one_session):
+ raise TypeError(
+ "RunMode.select_driver() must return an object with run_one_session(...)."
+ )
+ result = run_one_session(
+ host=host,
+ provider=provider,
+ session_edges=session_edges,
+ pipeline=pipeline,
+ )
+ if not inspect.isawaitable(result):
+ raise TypeError("run_demo_session_async(...) requires an async session driver.")
+ resolved = await result
+ if not isinstance(resolved, RunResult):
+ raise TypeError(
+ "Async session driver run_one_session(...) must return RunResult, "
+ f"got {type(resolved).__name__}."
+ )
+ return resolved
+
+
+async def _close_partial_session_async(
+ *,
+ context: RunContext,
+ provider: Any | None,
+ session_edges: SessionEdges | None,
+ status: DriverStatus,
+ reason: str | None,
+ error: Exception | None,
+ close_provider: bool,
+) -> RunResult:
+ if provider is not None and close_provider and session_edges is not None:
+ return await shielded_session_cleanup(
+ host=context.host,
+ session=None,
+ provider=provider,
+ session_edges=session_edges,
+ status=status,
+ reason=reason,
+ error=error,
+ )
+ if provider is not None and close_provider:
+ await _close_provider_async(
+ context=context,
+ provider=provider,
+ session_edges=session_edges,
+ )
+ if session_edges is not None:
+ return session_edges.close_result(status=status, reason=reason, error=error)
+ return RunResult(status=status, reason=reason, error=error)
+
+
+def _needs_partial_provider_cleanup(session_edges: SessionEdges | None) -> bool:
+ return session_edges is None or not session_edges.is_closed
+
+
+async def _close_provider_async(
+ *,
+ context: RunContext,
+ provider: Any,
+ session_edges: SessionEdges | None,
+) -> None:
+ try:
+ close_task = asyncio.create_task(context.host.call_async(provider.close))
+ except RuntimeError as close_exc:
+ _record_provider_cleanup_error(
+ context=context,
+ session_edges=session_edges,
+ exc=close_exc,
+ )
+ return
+
+ while not close_task.done():
+ try:
+ await asyncio.shield(close_task)
+ except asyncio.CancelledError:
+ uncancel_current_task()
+ continue
+ except Exception:
+ break
+
+ try:
+ await close_task
+ except asyncio.CancelledError:
+ uncancel_current_task()
+ _record_provider_cleanup_error(
+ context=context,
+ session_edges=session_edges,
+ exc=RuntimeError("provider cleanup was cancelled"),
+ )
+ except Exception as close_exc:
+ _record_provider_cleanup_error(
+ context=context,
+ session_edges=session_edges,
+ exc=close_exc,
+ )
+
+
+def _record_provider_cleanup_error(
+ *,
+ context: RunContext,
+ session_edges: SessionEdges | None,
+ exc: Exception,
+) -> None:
+ if session_edges is not None:
+ session_edges.record_cleanup_error(exc)
+ else:
+ _record_run_cleanup_error(context, exc)
+ # Partial async assembly may only have a provider to close. If that
+ # model-affine cleanup fails, quarantine the host instead of admitting a new
+ # session onto a worker that may still own model resources.
+ _mark_host_cleanup_failed(context.host, exc)
+
+
+def _record_run_cleanup_error(context: RunContext, exc: Exception) -> None:
+ try:
+ context.run_metrics.record_cleanup_error(exc)
+ except Exception:
+ return
+
+
+def _record_run_session_error(context: RunContext, exc: Exception) -> None:
+ try:
+ context.run_metrics.record_session_error(exc)
+ except Exception:
+ return
+
+
+async def _close_model_resources_async(
+ *,
+ host: RuntimeHost,
+ session: InferenceSession | None,
+ provider: ModelInputProvider,
+ session_edges: SessionEdges,
+ timeout_s: float,
+) -> str | None:
+ try:
+ resources_closed = await asyncio.wait_for(
+ host.call_async(
+ _close_model_resources_safely,
+ session.close if session is not None else None,
+ provider.close,
+ session_edges,
+ ),
+ timeout=timeout_s,
+ )
+ except asyncio.TimeoutError as exc:
+ # Keep provider cleanup ordered behind session cleanup on the model worker.
+ # A timed-out session close may still hold CUDA/Triton state, so running
+ # provider cleanup on another thread or replacing the worker is unsafe.
+ # The caller marks the host unhealthy so future sessions reject instead.
+ session_edges.record_orphaned_cleanup(exc)
+ return _MODEL_CLEANUP_TIMED_OUT_REASON
+ except Exception as exc:
+ session_edges.record_cleanup_error(exc)
+ return _MODEL_CLEANUP_FAILED_REASON
+ if not resources_closed:
+ return _MODEL_CLEANUP_FAILED_REASON
+ return None
+
+
+def _close_model_resources_safely(
+ session_close: Any | None,
+ provider_close: Any,
+ session_edges: SessionEdges,
+) -> bool:
+ resources_closed = True
+ # Session and provider close are intentionally ordered on the model worker.
+ # If session close hangs, timeout handling records orphaned cleanup and
+ # quarantines the host rather than moving provider close to another thread.
+ if session_close is not None:
+ resources_closed = _close_safely(session_close, session_edges)
+ return _close_safely(provider_close, session_edges) and resources_closed
+
+
+def _cleanup_result(
+ cleanup_task: asyncio.Task[RunResult],
+ session_edges: SessionEdges,
+ status: DriverStatus,
+ reason: str | None,
+ error: Exception | None,
+) -> RunResult:
+ if cleanup_task.done() and not cleanup_task.cancelled():
+ exc = cleanup_task.exception()
+ if exc is None:
+ return cleanup_task.result()
+ if isinstance(exc, Exception):
+ session_edges.record_cleanup_error(exc)
+ else:
+ session_edges.record_cleanup_error(
+ RuntimeError(f"cleanup failed with {type(exc).__name__}")
+ )
+ return session_edges.close_result(status=status, reason=reason, error=error)
+
+
+def _realtime_activation_and_clock(
+ session_edges: SessionEdges,
+) -> tuple[ActivationPolicy, RealtimeClock]:
+ activation = session_edges.activation
+ if activation is None:
+ raise DriverInvariantError(
+ "RealtimeSessionDriver requires SessionEdges.activation."
+ )
+ clock = session_edges.clock
+ if not isinstance(clock, RealtimeClock):
+ raise DriverInvariantError("RealtimeSessionDriver requires a RealtimeClock.")
+ return activation, clock
+
+
+def _realtime_input_source(session_edges: SessionEdges) -> Any:
+ input_source = session_edges.input_source
+ next_realtime_window = getattr(input_source, "next_realtime_window", None)
+ if not callable(next_realtime_window):
+ raise DriverInvariantError(
+ "RealtimeSessionDriver requires a RealtimeInputSource."
+ )
+ return input_source
+
+
+def uncancel_current_task() -> None:
+ task = asyncio.current_task()
+ if task is None:
+ return
+ uncancel = getattr(task, "uncancel", None)
+ if not callable(uncancel):
+ return
+ cancelling = getattr(task, "cancelling", None)
+ if not callable(cancelling):
+ return
+ while cancelling():
+ uncancel()
+
+
+__all__ = [
+ "BatchSessionDriver",
+ "CLEANUP_TIMEOUT_S",
+ "DriverInvariantError",
+ "RealtimeSessionDriver",
+ "run_demo_session",
+ "run_demo_session_async",
+ "shielded_session_cleanup",
+ "uncancel_current_task",
+]
diff --git a/flashdreams/flashdreams/runtime/demo/host.py b/flashdreams/flashdreams/runtime/demo/host.py
new file mode 100644
index 000000000..6a6e93ce5
--- /dev/null
+++ b/flashdreams/flashdreams/runtime/demo/host.py
@@ -0,0 +1,174 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Runtime host and model-execution boundary for shared demos."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import TypeVar
+
+from flashdreams.runtime._utils import freeze_mapping
+from flashdreams.runtime.inputs import InferenceInput
+from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession
+from flashdreams.runtime.worker import ModelExecutionWorker
+
+_T = TypeVar("_T")
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class WarmupSessionInputs:
+ """Inputs used to warm one temporary runtime session."""
+
+ initial_input: InferenceInput
+ step_inputs: Sequence[InferenceInput] = ()
+
+ def __post_init__(self) -> None:
+ object.__setattr__(self, "step_inputs", tuple(self.step_inputs))
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class ModelWarmupPlan:
+ """Host-owned model warmup plan built by a demo adapter or run mode."""
+
+ sessions: Sequence[WarmupSessionInputs] = ()
+ measured: bool = False
+ metadata: Mapping[str, object] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ object.__setattr__(self, "sessions", tuple(self.sessions))
+ object.__setattr__(self, "metadata", freeze_mapping(self.metadata))
+
+
+class RuntimeHost:
+ """Own one runtime and the worker used for model-affine calls."""
+
+ def __init__(
+ self,
+ runtime: InferenceRuntime,
+ *,
+ worker: ModelExecutionWorker | None = None,
+ is_control_rank: bool = True,
+ worker_loop: Callable[[], None] | None = None,
+ ) -> None:
+ self._runtime = runtime
+ self._worker = worker or ModelExecutionWorker()
+ self._is_control_rank = is_control_rank
+ self._worker_loop = worker_loop
+ self._healthy = True
+ self._closed = False
+ self._unhealthy_reason: str | None = None
+ self._unhealthy_error: Exception | None = None
+
+ @property
+ def runtime(self) -> InferenceRuntime:
+ """Return the hosted runtime."""
+ return self._runtime
+
+ @property
+ def worker(self) -> ModelExecutionWorker:
+ """Return the host's model-execution worker."""
+ return self._worker
+
+ @property
+ def is_control_rank(self) -> bool:
+ """Whether this process owns run modes, providers, sinks, and metrics."""
+ return self._is_control_rank
+
+ @property
+ def is_healthy(self) -> bool:
+ """Return whether admission should continue accepting sessions."""
+ return self._healthy and not self._closed
+
+ @property
+ def unhealthy_reason(self) -> str | None:
+ """Return the first latched unhealthy reason, if any."""
+ return self._unhealthy_reason
+
+ @property
+ def unhealthy_error(self) -> Exception | None:
+ """Return the first latched unhealthy error, if any."""
+ return self._unhealthy_error
+
+ def mark_unhealthy(
+ self,
+ reason: str = "marked unhealthy",
+ error: Exception | None = None,
+ ) -> None:
+ """Latch the host as unhealthy without overwriting the first reason."""
+ if not self._healthy:
+ return
+ self._healthy = False
+ self._unhealthy_reason = reason
+ self._unhealthy_error = error
+
+ def preload(self) -> None:
+ """Initialize optional distributed state and preload runtime resources."""
+ self._call_optional_runtime_hook("initialize_distributed")
+ self._call_optional_runtime_hook("preload")
+
+ def warmup(self, plan: ModelWarmupPlan | None = None) -> None:
+ """Run warmup sessions through the same worker boundary as real sessions."""
+ plan = plan or ModelWarmupPlan()
+ for warmup_session in plan.sessions:
+ session = self.call(self.start_session, warmup_session.initial_input)
+ try:
+ for step_input in warmup_session.step_inputs:
+ self.call(session.step, step_input)
+ finally:
+ self.call(session.close)
+
+ def call(self, func: Callable[..., _T], /, *args: object, **kwargs: object) -> _T:
+ """Run one model-affine callable synchronously on the worker."""
+ self._require_open()
+ return self._worker.call_blocking(func, *args, **kwargs)
+
+ async def call_async(
+ self,
+ func: Callable[..., _T],
+ /,
+ *args: object,
+ **kwargs: object,
+ ) -> _T:
+ """Run model-affine work without blocking realtime event loops."""
+ self._require_open()
+ return await self._worker.call(func, *args, **kwargs)
+
+ def start_session(self, inputs: InferenceInput) -> InferenceSession:
+ """Start one inference session through the hosted runtime."""
+ self._require_open()
+ return self._runtime.start_session(inputs)
+
+ def run_worker_loop(self) -> None:
+ """Serve control-rank work on non-control ranks until runtime shutdown."""
+ worker_loop = self._worker_loop
+ if worker_loop is None:
+ worker_loop = getattr(self._runtime, "run_worker_loop", None)
+ if worker_loop is None:
+ worker_loop = getattr(self._runtime, "wait_for_termination", None)
+ if callable(worker_loop):
+ worker_loop()
+
+ def close(self) -> None:
+ """Close runtime-owned state and stop the model-execution worker."""
+ if self._closed:
+ return
+ try:
+ self._worker.call_blocking(self._runtime.close)
+ self._call_optional_runtime_hook("close_distributed")
+ finally:
+ self._closed = True
+ self._worker.close_blocking()
+
+ def _call_optional_runtime_hook(self, name: str) -> None:
+ hook = getattr(self._runtime, name, None)
+ if callable(hook):
+ self.call(hook)
+
+ def _require_open(self) -> None:
+ if self._closed:
+ raise RuntimeError("runtime host is closed")
+
+
+__all__ = ["ModelWarmupPlan", "RuntimeHost", "WarmupSessionInputs"]
diff --git a/flashdreams/flashdreams/runtime/demo/outputs.py b/flashdreams/flashdreams/runtime/demo/outputs.py
index 421ec3bb4..866acbb00 100644
--- a/flashdreams/flashdreams/runtime/demo/outputs.py
+++ b/flashdreams/flashdreams/runtime/demo/outputs.py
@@ -1,18 +1,270 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""Shared demo output-target construction."""
+"""Shared demo output contracts and output construction."""
from __future__ import annotations
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
from pathlib import Path
+from typing import Literal, Protocol, runtime_checkable
-from flashdreams.runtime.output import NullOutputTarget, OutputTarget
+from flashdreams.infra.postprocess import VideoTensorLayout
+from flashdreams.infra.runner_io import (
+ DEFAULT_RUNNER_INSTALL_HINT,
+ write_video_tensor,
+)
+from flashdreams.infra.video_output import VideoResultCollector, prepare_video_for_mp4
+from flashdreams.runtime._utils import freeze_mapping
+from flashdreams.runtime.output import NullOutputTarget, OutputArtifact, OutputTarget
+from flashdreams.runtime.types import StepResult
from flashdreams.runtime.video_output import Mp4VideoOutputTarget, VideoWriter
from .spec import Mp4OutputSpec, NullOutputSpec, OutputSpec, WebRTCOutputSpec
+@dataclass(frozen=True, kw_only=True, slots=True)
+class SessionInfo:
+ """Output-facing metadata known after session setup."""
+
+ output_layout: str | None = None
+ steady_output_frame_count: int | None = None
+ metadata: Mapping[str, object] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ if self.output_layout is not None and not self.output_layout.strip():
+ raise ValueError("SessionInfo.output_layout must be non-empty when set.")
+ if (
+ self.steady_output_frame_count is not None
+ and self.steady_output_frame_count < 0
+ ):
+ raise ValueError(
+ "SessionInfo.steady_output_frame_count must be >= 0 when set."
+ )
+ object.__setattr__(self, "metadata", freeze_mapping(self.metadata))
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class OutputDecision:
+ """Flow-control decision returned by an output sink after one step."""
+
+ should_stop: bool = False
+ dropped: bool = False
+ drop_policy: Literal["none", "drop_newest", "drop_oldest"] = "none"
+ backpressure_s: float = 0.0
+ metadata: Mapping[str, object] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ if self.drop_policy not in {"none", "drop_newest", "drop_oldest"}:
+ raise ValueError(f"Unsupported drop_policy={self.drop_policy!r}.")
+ if self.backpressure_s < 0:
+ raise ValueError("OutputDecision.backpressure_s must be >= 0.")
+ object.__setattr__(self, "metadata", freeze_mapping(self.metadata))
+
+
+@runtime_checkable
+class OutputSink(Protocol):
+ """Consumes generated session outputs for a demo run mode."""
+
+ produces_artifacts: bool
+
+ def open(self, session_info: SessionInfo) -> None:
+ """Prepare output resources for a session."""
+ ...
+
+ def begin_generation(self, generation: int) -> None:
+ """Start an output generation, discarding stale live output if needed."""
+ ...
+
+ def write(self, result: StepResult) -> OutputDecision:
+ """Consume one generated result and return output flow-control state."""
+ ...
+
+ def close(self) -> Sequence[OutputArtifact]:
+ """Finalize output resources and return produced artifacts."""
+ ...
+
+
+@dataclass(slots=True)
+class NullOutputSink:
+ """Output sink for headless runs and fake-model vertical-slice tests."""
+
+ store_results: bool = False
+ produces_artifacts: bool = False
+ output_count: int = field(default=0, init=False)
+ results: list[Mapping[str, object]] = field(default_factory=list, init=False)
+ opened: bool = field(default=False, init=False)
+ closed: bool = field(default=False, init=False)
+ session_info: SessionInfo | None = field(default=None, init=False)
+ generation: int | None = field(default=None, init=False)
+
+ def open(self, session_info: SessionInfo) -> None:
+ self.session_info = session_info
+ self.output_count = 0
+ self.results.clear()
+ self.opened = True
+ self.closed = False
+
+ def begin_generation(self, generation: int) -> None:
+ if generation < 0:
+ raise ValueError("generation must be >= 0.")
+ self.generation = generation
+
+ def write(self, result: StepResult) -> OutputDecision:
+ if not self.opened or self.closed:
+ raise RuntimeError("Cannot write to a closed output sink.")
+ self.output_count += 1
+ if self.store_results:
+ self.results.append(_result_record(result))
+ return OutputDecision()
+
+ def close(self) -> Sequence[OutputArtifact]:
+ self.closed = True
+ return ()
+
+
+@dataclass(slots=True)
+class Mp4OutputSink:
+ """MP4 artifact sink for shared demo drivers."""
+
+ output_path: Path
+ fps: int | float
+ output_layout: VideoTensorLayout = "bvtchw"
+ writer: VideoWriter = field(default=write_video_tensor, repr=False)
+ install_hint: str = DEFAULT_RUNNER_INSTALL_HINT
+ move_to_cpu: bool = True
+ enabled: bool = True
+ produces_artifacts: bool = True
+ _opened: bool = field(default=False, init=False, repr=False)
+ _closed: bool = field(default=True, init=False, repr=False)
+ _collector: VideoResultCollector | None = field(
+ default=None,
+ init=False,
+ repr=False,
+ )
+ _artifacts: tuple[OutputArtifact, ...] | None = field(
+ default=None,
+ init=False,
+ repr=False,
+ )
+ session_info: SessionInfo | None = field(default=None, init=False)
+
+ def __post_init__(self) -> None:
+ if float(self.fps) <= 0:
+ raise ValueError("Mp4OutputSink.fps must be > 0.")
+ self.output_path = Path(self.output_path)
+
+ def open(self, session_info: SessionInfo) -> None:
+ self.session_info = session_info
+ self._collector = VideoResultCollector(
+ output_layout=self.output_layout,
+ enabled=self.enabled,
+ move_to_cpu=self.move_to_cpu,
+ )
+ self._artifacts = None
+ self._opened = True
+ self._closed = False
+
+ def begin_generation(self, generation: int) -> None:
+ if generation < 0:
+ raise ValueError("generation must be >= 0.")
+
+ def write(self, result: StepResult) -> OutputDecision:
+ if not self._opened or self._closed or self._collector is None:
+ raise RuntimeError("Cannot write to a closed output sink.")
+ if result.layout is None:
+ raise TypeError("Mp4OutputSink requires a video StepResult with layout.")
+ if result.layout != self.output_layout:
+ raise ValueError(
+ "Mp4OutputSink received layout "
+ f"{result.layout!r}; expected {self.output_layout!r}."
+ )
+ self._collector.add(result)
+ return OutputDecision()
+
+ def close(self) -> Sequence[OutputArtifact]:
+ if self._artifacts is not None:
+ return self._artifacts
+ if self._collector is None:
+ self._opened = False
+ self._closed = True
+ self._artifacts = ()
+ return self._artifacts
+
+ collector = self._collector
+ self._collector = None
+ self._opened = False
+ self._closed = True
+ video = collector.finish()
+ if video is None:
+ self._artifacts = ()
+ return self._artifacts
+ writable_video, writable_layout = prepare_video_for_mp4(
+ video,
+ layout=self.output_layout,
+ )
+ path = self.writer(
+ writable_video,
+ self.output_path,
+ fps=self.fps,
+ layout=writable_layout,
+ install_hint=self.install_hint,
+ )
+ self._artifacts = (
+ OutputArtifact(
+ kind="video/mp4",
+ uri=str(path),
+ metadata={
+ "fps": self.fps,
+ "source_layout": self.output_layout,
+ "shape": tuple(int(dim) for dim in video.shape),
+ "stats_history": tuple(collector.stats_history),
+ },
+ ),
+ )
+ return self._artifacts
+
+
+def build_output_sink(
+ output: OutputSpec,
+ *,
+ mp4_writer: VideoWriter | None = None,
+) -> OutputSink:
+ """Build a shared demo output sink from a demo output spec."""
+ if isinstance(output, NullOutputSpec):
+ return NullOutputSink(store_results=output.store_results)
+ if isinstance(output, Mp4OutputSpec):
+ writer = mp4_writer or write_video_tensor
+ return Mp4OutputSink(
+ output_path=Path(output.path),
+ fps=output.fps,
+ output_layout=output.output_layout,
+ writer=writer,
+ move_to_cpu=output.move_to_cpu,
+ )
+ if isinstance(output, WebRTCOutputSpec):
+ raise ValueError("WebRTC output requires a realtime transport sink.")
+ raise TypeError(f"Unsupported demo output spec: {type(output).__name__}.")
+
+
+def _result_record(result: StepResult) -> Mapping[str, object]:
+ record: dict[str, object] = {
+ "step_index": result.step_index,
+ "frame_count": result.frame_count,
+ "metrics": dict(result.metrics),
+ "metadata": dict(result.metadata),
+ }
+ if result.layout is not None:
+ record["layout"] = result.layout
+ if result.output_window is not None:
+ record["output_window"] = (
+ result.output_window.start_s,
+ result.output_window.end_s,
+ )
+ return freeze_mapping(record)
+
+
def build_output_target(
output: OutputSpec,
*,
@@ -42,4 +294,12 @@ def build_output_target(
raise TypeError(f"Unsupported demo output spec: {type(output).__name__}.")
-__all__ = ["build_output_target"]
+__all__ = [
+ "Mp4OutputSink",
+ "NullOutputSink",
+ "OutputDecision",
+ "OutputSink",
+ "SessionInfo",
+ "build_output_sink",
+ "build_output_target",
+]
diff --git a/flashdreams/flashdreams/runtime/demo/pipeline.py b/flashdreams/flashdreams/runtime/demo/pipeline.py
new file mode 100644
index 000000000..5a9da08b0
--- /dev/null
+++ b/flashdreams/flashdreams/runtime/demo/pipeline.py
@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared per-step pipeline for demo session drivers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from flashdreams.runtime.interfaces import InferenceSession
+from flashdreams.runtime.types import StepRequirements, StepResult
+
+from .outputs import OutputDecision, OutputSink
+from .run_modes import SessionMetricsRecorder
+from .session_inputs import ControlDecision, ModelInputProvider, UserInputWindow
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class StepOutcome:
+ """Combined output and control result from one shared model step."""
+
+ output: OutputDecision = field(default_factory=OutputDecision)
+ control: ControlDecision = field(default_factory=ControlDecision)
+
+
+class StepPipeline:
+ """Shared invariant for provider conversion, model step, output, and metrics."""
+
+ def execute_step(
+ self,
+ *,
+ request: StepRequirements,
+ user_window: UserInputWindow,
+ provider: ModelInputProvider,
+ session: InferenceSession,
+ output: OutputSink,
+ metrics: SessionMetricsRecorder,
+ ) -> StepOutcome:
+ prepared = provider.prepare_step(
+ request=request,
+ user_window=user_window,
+ )
+ if prepared.control.reset or prepared.control.close_session:
+ metrics.record_control(
+ request=request,
+ user_window=user_window,
+ control=prepared.control,
+ )
+ return StepOutcome(control=prepared.control)
+ if prepared.inference_input is None:
+ raise RuntimeError("ModelInputProvider returned no inference input.")
+
+ result = session.step(prepared.inference_input)
+ if not isinstance(result, StepResult):
+ raise TypeError(
+ "InferenceSession.step must return StepResult, "
+ f"got {type(result).__name__}."
+ )
+ decision = output.write(result)
+ if not isinstance(decision, OutputDecision):
+ raise TypeError(
+ "OutputSink.write must return OutputDecision, "
+ f"got {type(decision).__name__}."
+ )
+ metrics.record_step(
+ request=request,
+ user_window=user_window,
+ inference_input=prepared.inference_input,
+ result=result,
+ decision=decision,
+ )
+ return StepOutcome(output=decision)
+
+
+__all__ = ["StepOutcome", "StepPipeline"]
diff --git a/flashdreams/flashdreams/runtime/demo/replay.py b/flashdreams/flashdreams/runtime/demo/replay.py
index 18b873254..fa2f507a9 100644
--- a/flashdreams/flashdreams/runtime/demo/replay.py
+++ b/flashdreams/flashdreams/runtime/demo/replay.py
@@ -5,28 +5,72 @@
from __future__ import annotations
+import math
from collections.abc import Callable, Sequence
+from flashdreams.runtime.canonical import InputCanonicalizer
+from flashdreams.runtime.config import InferenceConfig
+from flashdreams.runtime.inputs import (
+ CanonicalInputs,
+ CanonicalInputSchema,
+ InferenceInput,
+ InferenceInputSchema,
+ TimeWindow,
+ UserInputs,
+ UserInputSchema,
+)
+from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession
+from flashdreams.runtime.mapping import (
+ DeclaresMappingSchema,
+ InputMapping,
+ check_mapping_compatibility,
+)
from flashdreams.runtime.metrics import MetricsRecorder, NullMetricsRecorder
from flashdreams.runtime.output import OutputArtifact, OutputTarget
-from flashdreams.runtime.runner import run_inference_session
+from flashdreams.runtime.types import (
+ StepRequest,
+ StepRequirements,
+ StepResult,
+ step_requirements_from_request,
+)
-from .outputs import build_output_target
-from .spec import DemoAdapter, DemoSpec, OutputSpec, WebRTCOutputSpec
+from .drivers import BatchSessionDriver, run_demo_session
+from .host import ModelWarmupPlan, RuntimeHost
+from .outputs import OutputDecision, OutputSink, build_output_sink, build_output_target
+from .pipeline import StepPipeline
+from .run_modes import (
+ Mp4ErrorPolicy,
+ NullErrorPolicy,
+ RunContext,
+ RunModeCapabilities,
+ RunResult,
+ SessionEdges,
+ SingleSessionAdmissionPolicy,
+)
+from .session_inputs import PreparedStep, ProviderCapabilities, UserInputWindow
+from .spec import (
+ DemoAdapter,
+ DemoSpec,
+ OutputSpec,
+ PreparedScenario,
+ WebRTCOutputSpec,
+)
OutputTargetFactory = Callable[[OutputSpec], OutputTarget]
InferenceSessionRunner = Callable[..., Sequence[OutputArtifact]]
+OutputSinkFactory = Callable[[OutputSpec], OutputSink]
def run_replay_demo(
*,
spec: DemoSpec,
adapter: DemoAdapter,
- output_target_factory: OutputTargetFactory = build_output_target,
+ output_target_factory: OutputTargetFactory | None = None,
+ output_sink_factory: OutputSinkFactory = build_output_sink,
metrics: MetricsRecorder | None = None,
- runner: InferenceSessionRunner = run_inference_session,
-) -> tuple[OutputArtifact, ...]:
- """Run one prepared demo scenario through the shared runtime runner."""
+ runner: InferenceSessionRunner | None = None,
+) -> RunResult:
+ """Run one prepared replay scenario through the shared batch demo path."""
_require_supported_mode(
mode=spec.input_mode,
supported=adapter.supported_input_modes(),
@@ -55,12 +99,87 @@ def run_replay_demo(
if spec.config is None:
raise RuntimeError("DemoSpec.config was not initialized.")
+ if runner is not None:
+ return _run_replay_demo_with_compat_runner(
+ spec=spec,
+ adapter=adapter,
+ prepared=prepared,
+ mapping=mapping,
+ output_target_factory=output_target_factory or build_output_target,
+ metrics=metrics,
+ runner=runner or _default_inference_session_runner(),
+ )
+
+ if output_target_factory is not None:
+ output_sink_factory = _output_target_sink_factory(output_target_factory)
+
+ return _run_replay_demo_with_run_mode(
+ spec=spec,
+ adapter=adapter,
+ prepared=prepared,
+ mapping=mapping,
+ output_sink_factory=output_sink_factory,
+ metrics=metrics,
+ )
+
+
+def _output_target_sink_factory(
+ output_target_factory: OutputTargetFactory,
+) -> OutputSinkFactory:
+ def create_output_sink(output_spec: OutputSpec) -> "_OutputTargetSink":
+ return _OutputTargetSink(output_target_factory(output_spec))
+
+ return create_output_sink
+
+
+class _OutputTargetSink:
+ produces_artifacts = True
+
+ def __init__(self, output: OutputTarget) -> None:
+ self._output = output
+ self._closed = True
+ self._artifacts: tuple[OutputArtifact, ...] | None = None
+
+ def open(self, session_info: object) -> None:
+ del session_info
+ self._output.open()
+ self._closed = False
+ self._artifacts = None
+
+ def begin_generation(self, generation: int) -> None:
+ del generation
+
+ def write(self, result: StepResult) -> OutputDecision:
+ self._output.write(result)
+ return OutputDecision()
+
+ def close(self) -> Sequence[OutputArtifact]:
+ if self._artifacts is not None:
+ return self._artifacts
+ if self._closed:
+ self._artifacts = ()
+ return self._artifacts
+ self._closed = True
+ self._artifacts = tuple(self._output.close())
+ return self._artifacts
+
+
+def _run_replay_demo_with_compat_runner(
+ *,
+ spec: DemoSpec,
+ adapter: DemoAdapter,
+ prepared: "PreparedScenario",
+ mapping: InputMapping,
+ output_target_factory: OutputTargetFactory,
+ metrics: MetricsRecorder | None,
+ runner: InferenceSessionRunner,
+) -> RunResult:
output = output_target_factory(spec.output)
metrics_recorder = metrics or NullMetricsRecorder()
- return tuple(
+ artifacts = tuple(
runner(
adapter=adapter,
- config=spec.config,
+ config=_require_config(spec),
mapping=mapping,
canonicalizer=prepared.canonicalizer,
source_schema=prepared.source_schema,
@@ -70,6 +189,422 @@ def run_replay_demo(
metrics=metrics_recorder,
)
)
+ return RunResult(status="completed", artifacts=artifacts)
+
+
+def _run_replay_demo_with_run_mode(
+ *,
+ spec: DemoSpec,
+ adapter: DemoAdapter,
+ prepared: "PreparedScenario",
+ mapping: InputMapping,
+ output_sink_factory: OutputSinkFactory,
+ metrics: MetricsRecorder | None,
+) -> RunResult:
+ config = _require_config(spec)
+ _validate_replay_mapping(
+ adapter=adapter,
+ config=config,
+ mapping=mapping,
+ source_schema=prepared.source_schema,
+ canonicalizer=prepared.canonicalizer,
+ )
+ request_state = _ReplayStepRequestState()
+ runtime = _ReplayRuntimeAdapter(
+ runtime=adapter.create_runtime(config),
+ request_state=request_state,
+ )
+ host = RuntimeHost(runtime)
+ mode = _ReplayRunMode(
+ request_state=request_state,
+ output_sink_factory=output_sink_factory,
+ run_metrics=metrics or NullMetricsRecorder(),
+ )
+ replay_adapter = _ReplayProviderAdapter(
+ adapter=adapter,
+ mapping=mapping,
+ request_state=request_state,
+ )
+ context = mode.create_run_context(
+ spec=spec,
+ adapter=replay_adapter,
+ host=host,
+ model_warmup_plan=ModelWarmupPlan(),
+ )
+ try:
+ return run_demo_session(
+ context=context,
+ spec=spec,
+ scenario=prepared,
+ adapter=replay_adapter,
+ run_mode=mode,
+ pipeline=StepPipeline(),
+ )
+ finally:
+ context.close()
+ host.close()
+
+
+class _ReplayRunMode:
+ name = "replay"
+ capabilities = RunModeCapabilities(
+ requires_finite_input=True,
+ supports_artifacts=True,
+ )
+
+ def __init__(
+ self,
+ *,
+ request_state: "_ReplayStepRequestState",
+ output_sink_factory: OutputSinkFactory,
+ run_metrics: MetricsRecorder,
+ ) -> None:
+ self._request_state = request_state
+ self._output_sink_factory = output_sink_factory
+ self._run_metrics = run_metrics
+
+ def validate_run(self, *, spec: DemoSpec, adapter: DemoAdapter) -> None:
+ del spec, adapter
+
+ def validate_session(
+ self,
+ *,
+ spec: DemoSpec,
+ scenario: "PreparedScenario",
+ adapter: DemoAdapter,
+ provider: object,
+ ) -> None:
+ del spec, scenario, adapter, provider
+
+ def create_run_context(
+ self,
+ *,
+ spec: DemoSpec,
+ adapter: DemoAdapter,
+ host: RuntimeHost,
+ model_warmup_plan: ModelWarmupPlan,
+ ) -> RunContext:
+ del spec, adapter
+ return RunContext(
+ host=host,
+ run_metrics=self._run_metrics,
+ admission=SingleSessionAdmissionPolicy(
+ health_check=lambda: host.is_healthy
+ ),
+ model_warmup_plan=model_warmup_plan,
+ )
+
+ def create_session_edges(
+ self,
+ *,
+ context: RunContext,
+ spec: DemoSpec,
+ scenario: "PreparedScenario",
+ provider: object,
+ adapter: DemoAdapter,
+ ) -> SessionEdges:
+ del provider, adapter
+ return SessionEdges(
+ input_source=_ReplayBatchInputSource(
+ scenario=scenario,
+ request_state=self._request_state,
+ ),
+ output_sink=self._output_sink_factory(spec.output),
+ cleanup_tasks=context.cleanup_tasks,
+ error_policy=(
+ Mp4ErrorPolicy() if spec.output.mode == "mp4" else NullErrorPolicy()
+ ),
+ )
+
+ def select_driver(self) -> BatchSessionDriver:
+ return BatchSessionDriver()
+
+
+class _ReplayProviderAdapter:
+ def __init__(
+ self,
+ *,
+ adapter: DemoAdapter,
+ mapping: InputMapping,
+ request_state: "_ReplayStepRequestState",
+ ) -> None:
+ self._adapter = adapter
+ self._mapping = mapping
+ self._request_state = request_state
+
+ @property
+ def model_id(self) -> str:
+ return self._adapter.model_id
+
+ @property
+ def inference_input_schema(self) -> InferenceInputSchema:
+ return self._adapter.inference_input_schema
+
+ @property
+ def canonical_input_schema(self) -> CanonicalInputSchema | None:
+ return self._adapter.canonical_input_schema
+
+ def default_input_mapping(self) -> InputMapping | None:
+ return self._adapter.default_input_mapping()
+
+ def supported_input_modes(self) -> tuple[str, ...]:
+ return self._adapter.supported_input_modes()
+
+ def supported_output_modes(self) -> tuple[str, ...]:
+ return self._adapter.supported_output_modes()
+
+ def validate_config(self, config: InferenceConfig) -> None:
+ self._adapter.validate_config(config)
+
+ def create_runtime(self, config: InferenceConfig) -> InferenceRuntime:
+ return self._adapter.create_runtime(config)
+
+ def prepare_scenario(self, spec: DemoSpec) -> "PreparedScenario":
+ return self._adapter.prepare_scenario(spec)
+
+ def create_model_input_provider(
+ self,
+ spec: DemoSpec,
+ scenario: "PreparedScenario",
+ ) -> object:
+ create_provider = getattr(self._adapter, "create_model_input_provider", None)
+ if callable(create_provider):
+ return create_provider(spec, scenario)
+ return _ReplayMappingModelInputProvider(
+ adapter=self._adapter,
+ scenario=scenario,
+ mapping=self._mapping,
+ request_state=self._request_state,
+ )
+
+
+class _ReplayRuntimeAdapter:
+ def __init__(
+ self,
+ *,
+ runtime: InferenceRuntime,
+ request_state: "_ReplayStepRequestState",
+ ) -> None:
+ self._runtime = runtime
+ self._request_state = request_state
+
+ def start_session(self, inputs: InferenceInput) -> InferenceSession:
+ return _ReplaySessionAdapter(
+ session=self._runtime.start_session(inputs),
+ request_state=self._request_state,
+ )
+
+ def close(self) -> None:
+ self._runtime.close()
+
+
+class _ReplaySessionAdapter:
+ def __init__(
+ self,
+ *,
+ session: InferenceSession,
+ request_state: "_ReplayStepRequestState",
+ ) -> None:
+ self._session = session
+ self._request_state = request_state
+
+ def next_step_requirements(self) -> StepRequirements | None:
+ next_requirements = getattr(self._session, "next_step_requirements", None)
+ if callable(next_requirements):
+ value = next_requirements()
+ self._request_state.clear()
+ return value
+
+ request = self._session.next_step_request()
+ if request is None:
+ self._request_state.clear()
+ return None
+ self._request_state.store(request)
+ return step_requirements_from_request(
+ request,
+ allow_user_input_window=True,
+ )
+
+ def next_step_request(self) -> StepRequest | None:
+ return self._session.next_step_request()
+
+ def step(self, inputs: InferenceInput) -> StepResult:
+ return self._session.step(inputs)
+
+ def reset(self, inputs: InferenceInput | None = None) -> None:
+ self._session.reset(inputs)
+
+ def close(self) -> None:
+ self._session.close()
+
+
+class _ReplayStepRequestState:
+ def __init__(self) -> None:
+ self._request: StepRequest | None = None
+
+ def store(self, request: StepRequest) -> None:
+ self._request = request
+
+ def request_for_window(self, step_index: int) -> StepRequest | None:
+ request = self._request
+ if request is None:
+ return None
+ if request.step_index != step_index:
+ raise RuntimeError(
+ "Replay input source request mismatch: "
+ f"expected step {request.step_index}, got {step_index}."
+ )
+ return request
+
+ def consume_for_step(self, request: StepRequirements) -> StepRequest:
+ legacy_request = self.request_for_window(request.step_index)
+ if legacy_request is not None:
+ self._request = None
+ return legacy_request
+ return StepRequest(
+ step_index=request.step_index,
+ inference_input_schema=request.inference_input_schema,
+ metadata=request.metadata,
+ )
+
+ def clear(self) -> None:
+ self._request = None
+
+
+class _ReplayBatchInputSource:
+ is_finite = True
+ is_deterministic = True
+
+ def __init__(
+ self,
+ *,
+ scenario: "PreparedScenario",
+ request_state: _ReplayStepRequestState,
+ ) -> None:
+ self.user_input_schema = scenario.source_schema
+ self._user_inputs = scenario.user_inputs
+ self._request_state = request_state
+
+ def is_finished(self) -> bool:
+ return False
+
+ def next_window(self, request: StepRequirements) -> UserInputWindow:
+ legacy_request = self._request_state.request_for_window(request.step_index)
+ window = (
+ legacy_request.user_input_window if legacy_request is not None else None
+ )
+ if window is None:
+ window = _all_user_inputs_window(self._user_inputs)
+ return UserInputWindow(
+ start_s=window.start_s,
+ end_s=window.end_s,
+ inputs=self._user_inputs,
+ )
+
+
+class _ReplayMappingModelInputProvider:
+ def __init__(
+ self,
+ *,
+ adapter: DemoAdapter,
+ scenario: "PreparedScenario",
+ mapping: InputMapping,
+ request_state: _ReplayStepRequestState,
+ ) -> None:
+ self.capabilities = ProviderCapabilities(
+ supports_recorded_input=True,
+ deterministic_given_inputs=True,
+ user_input_schema=scenario.source_schema,
+ inference_input_schema=adapter.inference_input_schema,
+ )
+ self._scenario = scenario
+ self._mapping = mapping
+ self._request_state = request_state
+ self._step_base_inputs = InferenceInput(
+ step=scenario.initial_inputs.step,
+ metadata=scenario.initial_inputs.metadata,
+ )
+
+ def prepare_initial_input(self) -> InferenceInput:
+ self._scenario.canonicalizer.reset()
+ return self._mapping.map_global_conditioning_inputs(
+ canonical_inputs=CanonicalInputs(),
+ inference_input=self._scenario.initial_inputs,
+ )
+
+ def prepare_step(
+ self,
+ *,
+ request: StepRequirements,
+ user_window: UserInputWindow,
+ ) -> PreparedStep:
+ legacy_request = self._request_state.consume_for_step(request)
+ canonical_inputs = self._scenario.canonicalizer.canonicalize(
+ self._scenario.user_inputs,
+ window=TimeWindow(start_s=user_window.start_s, end_s=user_window.end_s),
+ source_schema=self._scenario.source_schema,
+ )
+ return PreparedStep(
+ inference_input=self._mapping.map_step_inputs(
+ canonical_inputs=canonical_inputs,
+ inference_input=self._step_base_inputs,
+ request=legacy_request,
+ )
+ )
+
+ def reset(self, inputs: InferenceInput | None = None) -> None:
+ del inputs
+ self._scenario.canonicalizer.reset()
+
+ def close(self) -> None:
+ return None
+
+
+def _validate_replay_mapping(
+ *,
+ adapter: DemoAdapter,
+ config: InferenceConfig,
+ mapping: InputMapping,
+ source_schema: UserInputSchema,
+ canonicalizer: InputCanonicalizer,
+) -> None:
+ adapter.validate_config(config)
+ canonical_schema = canonicalizer.canonical_schema(source_schema)
+ if isinstance(mapping, DeclaresMappingSchema):
+ compatibility = check_mapping_compatibility(
+ canonical_schema=canonical_schema,
+ inference_input_schema=adapter.inference_input_schema,
+ mapping_schema=mapping.mapping_schema,
+ )
+ compatibility.raise_if_incompatible()
+ mapping.validate(
+ canonical_schema=canonical_schema,
+ inference_input_schema=adapter.inference_input_schema,
+ )
+
+
+def _require_config(spec: DemoSpec) -> InferenceConfig:
+ if spec.config is None:
+ raise RuntimeError("DemoSpec.config was not initialized.")
+ return spec.config
+
+
+def _all_user_inputs_window(user_inputs: UserInputs) -> TimeWindow:
+ if not user_inputs.events:
+ return TimeWindow(start_s=0.0, end_s=3600.0)
+ return TimeWindow(
+ start_s=0.0,
+ end_s=max(
+ 3600.0,
+ math.nextafter(user_inputs.events[-1].timestamp_s, math.inf),
+ ),
+ )
+
+
+def _default_inference_session_runner() -> InferenceSessionRunner:
+ from flashdreams.runtime.runner import run_inference_session
+
+ return run_inference_session
def _require_supported_mode(
@@ -88,6 +623,7 @@ def _require_supported_mode(
__all__ = [
"InferenceSessionRunner",
+ "OutputSinkFactory",
"OutputTargetFactory",
"run_replay_demo",
]
diff --git a/flashdreams/flashdreams/runtime/demo/run_modes.py b/flashdreams/flashdreams/runtime/demo/run_modes.py
new file mode 100644
index 000000000..cb93e4e7f
--- /dev/null
+++ b/flashdreams/flashdreams/runtime/demo/run_modes.py
@@ -0,0 +1,535 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Run/session result and policy helpers for demo session drivers."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from threading import Lock
+from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable
+
+from flashdreams.runtime._utils import freeze_mapping
+from flashdreams.runtime.metrics import (
+ InMemoryMetricsRecorder,
+ MetricsRecorder,
+ MetricsSnapshot,
+)
+from flashdreams.runtime.output import OutputArtifact
+
+from .host import ModelWarmupPlan, WarmupSessionInputs
+from .outputs import OutputSink
+
+if TYPE_CHECKING:
+ from .host import RuntimeHost
+ from .pipeline import StepPipeline
+ from .session_inputs import InputSource, ModelInputProvider
+ from .spec import DemoAdapter, DemoSpec, PreparedScenario
+ from .timing import ActivationPolicy, DeterministicClock, RealtimeClock
+
+SessionStatus = Literal[
+ "completed",
+ "failed",
+ "skipped",
+ "cancelled",
+ "rejected",
+ "not_activated",
+]
+
+DriverStatus = Literal[
+ "completed",
+ "failed",
+ "skipped",
+ "cancelled",
+ "not_activated",
+]
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class RunResult:
+ """Outcome of one demo session."""
+
+ __hash__ = None
+
+ status: SessionStatus
+ artifacts: Sequence[OutputArtifact] = ()
+ metrics: MetricsSnapshot | None = None
+ reason: str | None = None
+ error: Exception | None = None
+
+ @classmethod
+ def rejected(cls, reason: str) -> "RunResult":
+ """Admission refused the session. The only no-session result helper."""
+ return cls(status="rejected", reason=reason)
+
+ def __post_init__(self) -> None:
+ object.__setattr__(self, "artifacts", tuple(self.artifacts))
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class RunSummary:
+ """Summary for a run context after one or more sessions."""
+
+ metrics: MetricsSnapshot
+ sessions: Sequence[RunResult] = ()
+
+ def __post_init__(self) -> None:
+ object.__setattr__(self, "sessions", tuple(self.sessions))
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class ErrorAction:
+ """Driver policy decision for an operational error."""
+
+ close_session: bool = True
+ drop_chunk: bool = False
+ continue_next_scenario: bool = False
+ result_status: Literal["completed", "failed", "skipped"] = "failed"
+
+
+class DefaultErrorPolicy:
+ """Default policy: operational errors fail the current session."""
+
+ def handle_setup_error(self, exc: Exception) -> ErrorAction:
+ del exc
+ return ErrorAction(result_status="failed")
+
+ def handle(self, exc: Exception) -> ErrorAction:
+ del exc
+ return ErrorAction(result_status="failed")
+
+
+@runtime_checkable
+class ErrorPolicy(Protocol):
+ """Maps driver-observed exceptions to session outcomes."""
+
+ def handle_setup_error(self, exc: Exception) -> ErrorAction: ...
+
+ def handle(self, exc: Exception) -> ErrorAction: ...
+
+
+class Mp4ErrorPolicy(DefaultErrorPolicy):
+ """Abort MP4 sessions on setup or step errors."""
+
+
+class NullErrorPolicy(DefaultErrorPolicy):
+ """Abort headless/null sessions on setup or step errors."""
+
+
+class NativeWindowErrorPolicy(DefaultErrorPolicy):
+ """Abort native-window sessions unless a future UI policy overrides it."""
+
+
+class BenchmarkErrorPolicy(DefaultErrorPolicy):
+ """Close failed scenarios while letting benchmark loops continue."""
+
+ def handle_setup_error(self, exc: Exception) -> ErrorAction:
+ del exc
+ return ErrorAction(result_status="failed", continue_next_scenario=True)
+
+ def handle(self, exc: Exception) -> ErrorAction:
+ del exc
+ return ErrorAction(result_status="failed", continue_next_scenario=True)
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class WebRTCErrorPolicy:
+ """Drop configured recoverable realtime errors, otherwise close the session."""
+
+ recoverable_exception_types: tuple[type[Exception], ...] = ()
+
+ def handle_setup_error(self, exc: Exception) -> ErrorAction:
+ del exc
+ return ErrorAction(result_status="failed")
+
+ def handle(self, exc: Exception) -> ErrorAction:
+ if self.recoverable_exception_types and isinstance(
+ exc, self.recoverable_exception_types
+ ):
+ return ErrorAction(
+ close_session=False,
+ drop_chunk=True,
+ result_status="failed",
+ )
+ return ErrorAction(result_status="failed")
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class RunModeCapabilities:
+ """Run-mode requirements and output/transport capabilities."""
+
+ realtime: bool = False
+ requires_finite_input: bool = False
+ supports_backpressure: bool = False
+ supports_interactive_events: bool = False
+ supports_artifacts: bool = False
+
+
+SessionMetricsRecorder = MetricsRecorder
+InMemorySessionMetricsRecorder = InMemoryMetricsRecorder
+
+
+class NoopTransportService:
+ """Idempotent placeholder transport for batch sessions."""
+
+ def __init__(self) -> None:
+ self.closed = False
+
+ def is_active(self) -> bool:
+ return not self.closed
+
+ def close(self) -> None:
+ self.closed = True
+
+
+@runtime_checkable
+class TransportService(Protocol):
+ """Per-session transport lifecycle hook."""
+
+ def is_active(self) -> bool: ...
+
+ def close(self) -> None: ...
+
+
+@runtime_checkable
+class SessionReservation(Protocol):
+ """Admission reservation for one session."""
+
+ def release(self) -> None: ...
+
+
+class SingleSessionAdmissionPolicy:
+ """Atomic single-session admission policy."""
+
+ def __init__(self, *, health_check: Any | None = None) -> None:
+ self._lock = Lock()
+ self._reserved = False
+ self._health_check = health_check
+
+ def try_reserve(self) -> SessionReservation | None:
+ with self._lock:
+ if self._reserved or not self._is_healthy():
+ return None
+ self._reserved = True
+ return _SingleSessionReservation(self)
+
+ def _release(self) -> None:
+ with self._lock:
+ self._reserved = False
+
+ def _is_healthy(self) -> bool:
+ if self._health_check is None:
+ return True
+ return bool(self._health_check())
+
+
+class _SingleSessionReservation:
+ def __init__(self, policy: SingleSessionAdmissionPolicy) -> None:
+ self._policy = policy
+ self._released = False
+ self.release_count = 0
+
+ def release(self) -> None:
+ if self._released:
+ return
+ self._released = True
+ self.release_count += 1
+ self._policy._release()
+
+
+@runtime_checkable
+class AdmissionPolicy(Protocol):
+ """Atomically reserves session capacity or rejects."""
+
+ def try_reserve(self) -> SessionReservation | None: ...
+
+
+@runtime_checkable
+class SessionDriver(Protocol):
+ """Synchronous one-session driver selected by a run mode."""
+
+ def run_one_session(
+ self,
+ *,
+ host: "RuntimeHost",
+ provider: "ModelInputProvider",
+ session_edges: "SessionEdges",
+ pipeline: "StepPipeline",
+ ) -> RunResult: ...
+
+
+@runtime_checkable
+class AsyncSessionDriver(Protocol):
+ """Async one-session driver selected by realtime run modes."""
+
+ async def run_one_session(
+ self,
+ *,
+ host: "RuntimeHost",
+ provider: "ModelInputProvider",
+ session_edges: "SessionEdges",
+ pipeline: "StepPipeline",
+ ) -> RunResult: ...
+
+
+@dataclass(slots=True)
+class RunContext:
+ """Run-scoped services shared by one or more demo sessions."""
+
+ host: "RuntimeHost"
+ run_metrics: SessionMetricsRecorder
+ admission: AdmissionPolicy
+ model_warmup_plan: ModelWarmupPlan = field(default_factory=ModelWarmupPlan)
+ services: Mapping[str, object] = field(default_factory=dict)
+ cleanup_tasks: set[asyncio.Task[RunResult]] = field(default_factory=set)
+
+ def __post_init__(self) -> None:
+ self.services = freeze_mapping(self.services)
+
+ def close(self) -> RunSummary:
+ if self.cleanup_tasks:
+ raise RuntimeError(
+ "Pending session cleanup tasks; async runs must await close_async()."
+ )
+ for service in self.services.values():
+ close = getattr(service, "close", None)
+ if callable(close):
+ try:
+ close()
+ except Exception as exc:
+ self.run_metrics.record_cleanup_error(exc)
+ return RunSummary(
+ metrics=self.run_metrics.close(),
+ sessions=tuple(getattr(self.run_metrics, "sessions", ())),
+ )
+
+ async def close_async(self) -> RunSummary:
+ while self.cleanup_tasks:
+ pending = tuple(self.cleanup_tasks)
+ await asyncio.gather(*pending, return_exceptions=True)
+ self.cleanup_tasks.difference_update(pending)
+ return self.close()
+
+
+@dataclass(slots=True)
+class SessionEdges:
+ """Per-session input/output/policy bundle consumed by drivers."""
+
+ input_source: "InputSource"
+ output_sink: OutputSink
+ cleanup_tasks: set[asyncio.Task[RunResult]]
+ metrics: SessionMetricsRecorder = field(
+ default_factory=InMemorySessionMetricsRecorder
+ )
+ error_policy: ErrorPolicy = field(default_factory=DefaultErrorPolicy)
+ transport: TransportService = field(default_factory=NoopTransportService)
+ clock: "RealtimeClock | DeterministicClock | None" = None
+ activation: "ActivationPolicy | None" = None
+ _closed_result: RunResult | None = field(default=None, init=False, repr=False)
+
+ @property
+ def is_closed(self) -> bool:
+ """Return whether ``close_result(...)`` has already finalized this session."""
+ return self._closed_result is not None
+
+ def record_cleanup_error(self, exc: Exception) -> None:
+ """Record a cleanup error without letting metrics failures block teardown."""
+ try:
+ self.metrics.record_cleanup_error(exc)
+ except Exception:
+ return
+
+ def record_orphaned_cleanup(self, exc: Exception) -> None:
+ """Record timed-out worker cleanup without blocking teardown."""
+ try:
+ self.metrics.record_orphaned_cleanup(exc)
+ except Exception:
+ return
+
+ def close_result(
+ self,
+ *,
+ status: DriverStatus = "completed",
+ reason: str | None = None,
+ error: Exception | None = None,
+ ) -> RunResult:
+ """Idempotently close output, transport, and metrics once."""
+ if self._closed_result is not None:
+ return self._closed_result
+
+ artifacts: Sequence[OutputArtifact] = ()
+ try:
+ artifacts = tuple(self.output_sink.close())
+ except Exception as exc:
+ self.record_cleanup_error(exc)
+ try:
+ self.transport.close()
+ except Exception as exc:
+ self.record_cleanup_error(exc)
+ try:
+ metrics = self.metrics.close()
+ except Exception as exc:
+ metrics = MetricsSnapshot(errors=(f"metrics.close failed: {exc}",))
+ self._closed_result = RunResult(
+ status=status,
+ artifacts=artifacts,
+ metrics=metrics,
+ reason=reason,
+ error=error,
+ )
+ return self._closed_result
+
+
+@runtime_checkable
+class RunMode(Protocol):
+ """Run/session construction strategy consumed by shared helpers."""
+
+ name: str
+ capabilities: RunModeCapabilities
+
+ def validate_run(
+ self,
+ *,
+ spec: "DemoSpec",
+ adapter: "DemoAdapter",
+ ) -> None: ...
+
+ def validate_session(
+ self,
+ *,
+ spec: "DemoSpec",
+ scenario: "PreparedScenario",
+ adapter: "DemoAdapter",
+ provider: "ModelInputProvider",
+ ) -> None: ...
+
+ def create_run_context(
+ self,
+ *,
+ spec: "DemoSpec",
+ adapter: "DemoAdapter",
+ host: "RuntimeHost",
+ model_warmup_plan: ModelWarmupPlan,
+ ) -> RunContext: ...
+
+ def create_session_edges(
+ self,
+ *,
+ context: RunContext,
+ spec: "DemoSpec",
+ scenario: "PreparedScenario",
+ provider: "ModelInputProvider",
+ adapter: "DemoAdapter",
+ ) -> SessionEdges: ...
+
+ def select_driver(self) -> SessionDriver | AsyncSessionDriver: ...
+
+
+@runtime_checkable
+class RunModeWarmup(Protocol):
+ """Optional run-mode warmup for output or transport services."""
+
+ def warmup_context(
+ self,
+ *,
+ context: RunContext,
+ spec: "DemoSpec",
+ scenario: "PreparedScenario",
+ adapter: "DemoAdapter",
+ ) -> None: ...
+
+
+def build_model_warmup_plan(
+ *,
+ host: "RuntimeHost",
+ adapter: "DemoAdapter",
+ spec: "DemoSpec",
+ scenario: "PreparedScenario",
+) -> ModelWarmupPlan:
+ """Build a host-owned warmup plan through the model-affine worker."""
+
+ create_sessions = getattr(adapter, "create_model_warmup_sessions", None)
+ if create_sessions is None:
+ return ModelWarmupPlan()
+ if not callable(create_sessions):
+ raise TypeError(
+ "Demo adapter create_model_warmup_sessions attribute must be callable."
+ )
+ sessions = host.call(create_sessions, spec, scenario)
+ return ModelWarmupPlan(sessions=_coerce_warmup_sessions(sessions))
+
+
+def warmup_run_context(
+ *,
+ context: RunContext,
+ spec: "DemoSpec",
+ scenario: "PreparedScenario",
+ adapter: "DemoAdapter",
+ run_mode: object,
+) -> None:
+ """Run model warmup, then optional output/transport warmup for a context."""
+
+ context.host.warmup(context.model_warmup_plan)
+ warmup_context = getattr(run_mode, "warmup_context", None)
+ if warmup_context is None:
+ return
+ if not callable(warmup_context):
+ raise TypeError("RunMode.warmup_context attribute must be callable.")
+ warmup_context(
+ context=context,
+ spec=spec,
+ scenario=scenario,
+ adapter=adapter,
+ )
+
+
+def _coerce_warmup_sessions(value: object) -> tuple[WarmupSessionInputs, ...]:
+ if not isinstance(value, Sequence):
+ raise TypeError(
+ "Demo adapter create_model_warmup_sessions(...) must return a sequence "
+ f"of WarmupSessionInputs, got {type(value).__name__}."
+ )
+ sessions: list[WarmupSessionInputs] = []
+ for session in value:
+ if not isinstance(session, WarmupSessionInputs):
+ raise TypeError(
+ "Demo adapter create_model_warmup_sessions(...) must return only "
+ f"WarmupSessionInputs, got {type(session).__name__}."
+ )
+ sessions.append(session)
+ return tuple(sessions)
+
+
+__all__ = [
+ "AdmissionPolicy",
+ "AsyncSessionDriver",
+ "BenchmarkErrorPolicy",
+ "DefaultErrorPolicy",
+ "DriverStatus",
+ "ErrorAction",
+ "ErrorPolicy",
+ "InMemorySessionMetricsRecorder",
+ "MetricsSnapshot",
+ "Mp4ErrorPolicy",
+ "NativeWindowErrorPolicy",
+ "NoopTransportService",
+ "NullErrorPolicy",
+ "RunContext",
+ "RunMode",
+ "RunModeCapabilities",
+ "RunModeWarmup",
+ "RunResult",
+ "RunSummary",
+ "SessionEdges",
+ "SessionDriver",
+ "SessionMetricsRecorder",
+ "SessionReservation",
+ "SessionStatus",
+ "SingleSessionAdmissionPolicy",
+ "TransportService",
+ "WebRTCErrorPolicy",
+ "build_model_warmup_plan",
+ "warmup_run_context",
+]
diff --git a/flashdreams/flashdreams/runtime/demo/session_inputs.py b/flashdreams/flashdreams/runtime/demo/session_inputs.py
new file mode 100644
index 000000000..dda42995c
--- /dev/null
+++ b/flashdreams/flashdreams/runtime/demo/session_inputs.py
@@ -0,0 +1,184 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Input-source and model-input-provider contracts for demo sessions."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+from flashdreams.runtime._utils import freeze_mapping
+from flashdreams.runtime.inputs import (
+ InferenceInput,
+ InferenceInputSchema,
+ UserInputs,
+ UserInputSchema,
+)
+from flashdreams.runtime.types import StepRequirements
+
+if TYPE_CHECKING:
+ from .timing import RealtimeClock, RealtimeWindowResult
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class ProviderCapabilities:
+ """Model-provider capabilities used to validate run-mode compatibility."""
+
+ supports_realtime_clock: bool = False
+ supports_recorded_input: bool = False
+ supports_reset: bool = False
+ deterministic_given_inputs: bool = False
+ user_input_schema: UserInputSchema = field(default_factory=UserInputSchema)
+ inference_input_schema: InferenceInputSchema = field(
+ default_factory=InferenceInputSchema
+ )
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class ControlDecision:
+ """Provider-authored control request for the current session."""
+
+ reset: bool = False
+ close_session: bool = False
+ reset_input: InferenceInput | None = None
+ provider_already_reset: bool = False
+ reason: str | None = None
+
+ def __post_init__(self) -> None:
+ if self.reason is not None and not self.reason.strip():
+ raise ValueError("ControlDecision.reason must be non-empty when set.")
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class UserInputWindow:
+ """User/app inputs selected by a driver for one model step."""
+
+ __hash__ = None
+
+ start_s: float
+ end_s: float
+ frame_times: Sequence[float] = ()
+ inputs: UserInputs = field(default_factory=UserInputs)
+ control: ControlDecision | None = None
+ metadata: Mapping[str, object] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ if not math.isfinite(self.start_s) or self.start_s < 0:
+ raise ValueError("UserInputWindow.start_s must be finite and >= 0.")
+ if not math.isfinite(self.end_s) or self.end_s < self.start_s:
+ raise ValueError("UserInputWindow.end_s must be finite and >= start_s.")
+ previous = -math.inf
+ for frame_time in self.frame_times:
+ if not math.isfinite(float(frame_time)):
+ raise ValueError("UserInputWindow.frame_times must be finite.")
+ if float(frame_time) < previous:
+ raise ValueError(
+ "UserInputWindow.frame_times must be sorted in ascending order."
+ )
+ previous = float(frame_time)
+ object.__setattr__(
+ self, "frame_times", tuple(float(t) for t in self.frame_times)
+ )
+ object.__setattr__(self, "metadata", freeze_mapping(self.metadata))
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class PreparedStep:
+ """Model-facing input plus optional provider-authored control decision."""
+
+ __hash__ = None
+
+ inference_input: InferenceInput | None = None
+ control: ControlDecision = field(default_factory=ControlDecision)
+
+
+@runtime_checkable
+class InputSource(Protocol):
+ """Facts common to every demo session input source."""
+
+ is_finite: bool
+ is_deterministic: bool
+ user_input_schema: UserInputSchema
+
+ def is_finished(self) -> bool:
+ """Return whether the driver should stop requesting windows."""
+ ...
+
+
+@runtime_checkable
+class BatchInputSource(InputSource, Protocol):
+ """Finite input source consumed by the batch driver."""
+
+ def next_window(self, request: StepRequirements) -> UserInputWindow:
+ """Return the next batch input window for ``request``."""
+ ...
+
+
+@runtime_checkable
+class RealtimeInputSource(InputSource, Protocol):
+ """Realtime input source consumed by a future realtime driver."""
+
+ async def next_realtime_window(
+ self,
+ *,
+ request: StepRequirements,
+ clock: "RealtimeClock",
+ ) -> "RealtimeWindowResult":
+ """Return the next realtime window result.
+
+ The concrete realtime result shape lands with the realtime clock phase.
+ Keeping this protocol separate now prevents batch sources from stubbing
+ async behavior they never serve.
+ """
+ ...
+
+
+@runtime_checkable
+class ModelInputProvider(Protocol):
+ """Model-owned conversion from user windows into model-facing inputs."""
+
+ capabilities: ProviderCapabilities
+
+ def prepare_initial_input(self) -> InferenceInput:
+ """Prepare session-global model inputs."""
+ ...
+
+ def prepare_step(
+ self,
+ *,
+ request: StepRequirements,
+ user_window: UserInputWindow,
+ ) -> PreparedStep:
+ """Prepare one model step from a driver-owned user input window."""
+ ...
+
+ def reset(self, inputs: InferenceInput | None = None) -> None:
+ """Reset provider-owned session state.
+
+ Implementations must be idempotent so driver cleanup and reset control
+ paths can safely converge after failures.
+ """
+ ...
+
+ def close(self) -> None:
+ """Release provider-owned resources.
+
+ Implementations must be idempotent and tolerate cleanup after partial
+ setup or earlier reset failures.
+ """
+ ...
+
+
+__all__ = [
+ "BatchInputSource",
+ "ControlDecision",
+ "InputSource",
+ "ModelInputProvider",
+ "PreparedStep",
+ "ProviderCapabilities",
+ "RealtimeInputSource",
+ "UserInputWindow",
+]
diff --git a/flashdreams/flashdreams/runtime/demo/spec.py b/flashdreams/flashdreams/runtime/demo/spec.py
index 6ba652f38..d7a6a309b 100644
--- a/flashdreams/flashdreams/runtime/demo/spec.py
+++ b/flashdreams/flashdreams/runtime/demo/spec.py
@@ -5,7 +5,7 @@
from __future__ import annotations
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any, Literal, Protocol, TypeAlias
@@ -18,6 +18,8 @@
from flashdreams.runtime.interfaces import ModelAdapter
from flashdreams.runtime.mapping import InputMapping
+from .host import WarmupSessionInputs
+
@dataclass(frozen=True, kw_only=True, slots=True)
class NullOutputSpec:
@@ -170,9 +172,22 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario:
...
+class ModelWarmupAdapter(Protocol):
+ """Optional adapter hook for model-affine runtime warmup inputs."""
+
+ def create_model_warmup_sessions(
+ self,
+ spec: DemoSpec,
+ scenario: PreparedScenario,
+ ) -> Sequence[WarmupSessionInputs]:
+ """Return temporary synthetic or loopback sessions for model warmup."""
+ ...
+
+
__all__ = [
"DemoAdapter",
"DemoSpec",
+ "ModelWarmupAdapter",
"Mp4OutputSpec",
"NullOutputSpec",
"OutputSpec",
diff --git a/flashdreams/flashdreams/runtime/demo/timing.py b/flashdreams/flashdreams/runtime/demo/timing.py
new file mode 100644
index 000000000..d48e1e213
--- /dev/null
+++ b/flashdreams/flashdreams/runtime/demo/timing.py
@@ -0,0 +1,383 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Realtime activation, clock, and input-window primitives for demo run modes."""
+
+from __future__ import annotations
+
+import asyncio
+import math
+import time
+from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Literal, Protocol, runtime_checkable
+
+from flashdreams.runtime.inputs import UserInputs, UserInputSchema
+from flashdreams.runtime.types import StepRequirements
+
+from .session_inputs import UserInputWindow
+
+CatchUpPolicy = Literal["drop", "fold", "compress"]
+
+SPARSE_KEY_SEGMENTS_METADATA_KEY = "sparse_key_segments"
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class CatchUpDecision:
+ """How a realtime clock bounded stale virtual input time."""
+
+ skipped_s: float = 0.0
+ skipped_windows: int = 0
+ input_policy: CatchUpPolicy | None = None
+ reason: str | None = None
+
+ def __post_init__(self) -> None:
+ if not math.isfinite(self.skipped_s) or self.skipped_s < 0.0:
+ raise ValueError("CatchUpDecision.skipped_s must be finite and >= 0.")
+ if self.skipped_windows < 0:
+ raise ValueError("CatchUpDecision.skipped_windows must be >= 0.")
+ if self.input_policy not in {None, "drop", "fold", "compress"}:
+ raise ValueError(
+ f"Unsupported catch-up input_policy={self.input_policy!r}."
+ )
+ if self.reason is not None and not self.reason.strip():
+ raise ValueError("CatchUpDecision.reason must be non-empty when set.")
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class RealtimeWindowResult:
+ """Realtime input window plus any catch-up decision that preceded it."""
+
+ window: UserInputWindow
+ catch_up: CatchUpDecision = field(default_factory=CatchUpDecision)
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class ActivationResult:
+ """Result of waiting for a realtime activation gate."""
+
+ activated: bool
+ reason: str | None = None
+
+ def __post_init__(self) -> None:
+ if self.reason is not None and not self.reason.strip():
+ raise ValueError("ActivationResult.reason must be non-empty when set.")
+
+
+@runtime_checkable
+class DeterministicClock(Protocol):
+ """Clock facts for finite deterministic run modes."""
+
+ is_realtime: bool
+ is_deterministic: bool
+
+
+@runtime_checkable
+class RealtimeClock(Protocol):
+ """Realtime virtual clock used by realtime drivers and input sources."""
+
+ is_realtime: bool
+ is_deterministic: bool
+
+ def now(self) -> float: ...
+
+ def anchor(self, wall_time_s: float) -> None: ...
+
+ async def wait_until_window_end(self, end_s: float) -> None: ...
+
+ async def apply_backpressure(self, requested_s: float) -> None: ...
+
+ def catch_up(
+ self,
+ *,
+ request: StepRequirements,
+ max_lag_s: float,
+ policy: CatchUpPolicy,
+ ) -> CatchUpDecision: ...
+
+
+@runtime_checkable
+class ActivationPolicy(Protocol):
+ """Wait until a realtime session should start generating."""
+
+ timeout_s: float | None
+
+ async def wait_until_active(
+ self,
+ clock: RealtimeClock | DeterministicClock,
+ ) -> ActivationResult: ...
+
+
+@runtime_checkable
+class ActivationSignal(Protocol):
+ """Event-like object accepted by ``SignalActivationPolicy``."""
+
+ def is_set(self) -> bool: ...
+
+ async def wait(self) -> object: ...
+
+
+@dataclass(slots=True)
+class AlwaysActiveActivationPolicy:
+ """Activation policy for batch/null modes or already-ready realtime modes."""
+
+ timeout_s: float | None = None
+ anchor_clock: bool = False
+
+ async def wait_until_active(
+ self,
+ clock: RealtimeClock | DeterministicClock,
+ ) -> ActivationResult:
+ _anchor_if_realtime(clock, anchor=self.anchor_clock)
+ return ActivationResult(activated=True)
+
+
+@dataclass(slots=True)
+class SignalActivationPolicy:
+ """Activate when any supplied signal fires, with optional timeout."""
+
+ signals: Sequence[ActivationSignal]
+ timeout_s: float | None = None
+ timeout_reason: str = "activation timed out"
+ anchor_clock: bool = True
+
+ def __post_init__(self) -> None:
+ if not self.signals:
+ raise ValueError("SignalActivationPolicy.signals must be non-empty.")
+ self.signals = tuple(self.signals)
+ if self.timeout_s is not None and self.timeout_s <= 0.0:
+ raise ValueError("SignalActivationPolicy.timeout_s must be > 0 when set.")
+ if not self.timeout_reason.strip():
+ raise ValueError("SignalActivationPolicy.timeout_reason must be non-empty.")
+
+ async def wait_until_active(
+ self,
+ clock: RealtimeClock | DeterministicClock,
+ ) -> ActivationResult:
+ if any(signal.is_set() for signal in self.signals):
+ _anchor_if_realtime(clock, anchor=self.anchor_clock)
+ return ActivationResult(activated=True)
+
+ tasks = [asyncio.create_task(signal.wait()) for signal in self.signals]
+ try:
+ done, pending = await asyncio.wait(
+ tasks,
+ timeout=self.timeout_s,
+ return_when=asyncio.FIRST_COMPLETED,
+ )
+ if not done:
+ return ActivationResult(
+ activated=False,
+ reason=self.timeout_reason,
+ )
+ for task in done:
+ task.result()
+ _anchor_if_realtime(clock, anchor=self.anchor_clock)
+ return ActivationResult(activated=True)
+ finally:
+ for task in tasks:
+ if not task.done():
+ task.cancel()
+ await asyncio.gather(*tasks, return_exceptions=True)
+
+
+SparseKeySegment = tuple[float, float, frozenset[str]]
+
+
+class _RealtimeTimeline(Protocol):
+ dt: float
+ next_chunk_start_v: float
+
+
+class _SparseInputResampler(_RealtimeTimeline, Protocol):
+ def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: ...
+
+ def reset(self, *, start_v: float) -> None: ...
+
+ def sample_chunk(
+ self,
+ num_frames: int,
+ ) -> tuple[Sequence[SparseKeySegment], Sequence[float]]: ...
+
+
+@dataclass(slots=True)
+class ResamplerRealtimeClock:
+ """Realtime clock that reuses ``KeyboardResampler``'s virtual timeline."""
+
+ resampler: _RealtimeTimeline
+ now_fn: Callable[[], float] = time.monotonic
+ sleep_fn: Callable[[float], Awaitable[None]] = asyncio.sleep
+ is_realtime: bool = True
+ is_deterministic: bool = False
+ _pending_backpressure_s: float = field(default=0.0, init=False, repr=False)
+
+ @property
+ def pending_backpressure_s(self) -> float:
+ return self._pending_backpressure_s
+
+ def now(self) -> float:
+ return float(self.now_fn())
+
+ def anchor(self, wall_time_s: float) -> None:
+ if not math.isfinite(wall_time_s):
+ raise ValueError("wall_time_s must be finite.")
+ self.resampler.next_chunk_start_v = float(wall_time_s)
+ self._pending_backpressure_s = 0.0
+
+ async def wait_until_window_end(self, end_s: float) -> None:
+ if not math.isfinite(end_s):
+ raise ValueError("end_s must be finite.")
+ delay_s = float(end_s) - self.now()
+ if delay_s > 0.0:
+ await self.sleep_fn(delay_s)
+
+ async def apply_backpressure(self, requested_s: float) -> None:
+ if not math.isfinite(requested_s) or requested_s < 0.0:
+ raise ValueError("requested_s must be finite and >= 0.")
+ self._pending_backpressure_s += float(requested_s)
+
+ def catch_up(
+ self,
+ *,
+ request: StepRequirements,
+ max_lag_s: float,
+ policy: CatchUpPolicy,
+ ) -> CatchUpDecision:
+ if policy != "fold":
+ raise NotImplementedError(
+ f"Catch-up policy {policy!r} has no existing resampler analog yet."
+ )
+ if not math.isfinite(max_lag_s) or max_lag_s < 0.0:
+ raise ValueError("max_lag_s must be finite and >= 0.")
+
+ input_frame_count = input_frame_count_from_request(request)
+ chunk_duration_s = input_frame_count * float(self.resampler.dt)
+ if chunk_duration_s <= 0.0:
+ raise ValueError("Realtime resampler dt must produce a positive window.")
+
+ effective_now_s = self.now() + self._pending_backpressure_s
+ self._pending_backpressure_s = 0.0
+ current_start_s = float(self.resampler.next_chunk_start_v)
+ lag_s = effective_now_s - (current_start_s + chunk_duration_s)
+ if lag_s <= max_lag_s:
+ return CatchUpDecision()
+
+ latest_start_s = effective_now_s - chunk_duration_s
+ if latest_start_s <= current_start_s:
+ return CatchUpDecision()
+
+ skipped_s = latest_start_s - current_start_s
+ skipped_windows = max(1, math.ceil(skipped_s / chunk_duration_s))
+ self.resampler.next_chunk_start_v = latest_start_s
+ return CatchUpDecision(
+ skipped_s=skipped_s,
+ skipped_windows=skipped_windows,
+ input_policy=policy,
+ reason="lag exceeded max_lag_s",
+ )
+
+
+@dataclass(slots=True)
+class KeyboardRealtimeInputSource:
+ """Realtime input source backed by the existing keyboard resampler."""
+
+ resampler: _SparseInputResampler
+ max_lag_s: float | None = None
+ catch_up_policy: CatchUpPolicy = "fold"
+ is_finite: bool = False
+ is_deterministic: bool = False
+ user_input_schema: UserInputSchema = field(default_factory=UserInputSchema)
+
+ def __post_init__(self) -> None:
+ if self.max_lag_s is not None and (
+ not math.isfinite(self.max_lag_s) or self.max_lag_s < 0.0
+ ):
+ raise ValueError(
+ "KeyboardRealtimeInputSource.max_lag_s must be finite and >= 0."
+ )
+ if self.catch_up_policy != "fold":
+ raise NotImplementedError(
+ f"Catch-up policy {self.catch_up_policy!r} has no existing "
+ "KeyboardResampler analog yet."
+ )
+
+ def is_finished(self) -> bool:
+ return False
+
+ def on_edge(self, *, arrival_t: float, event: str, key: str) -> None:
+ self.resampler.on_edge(arrival_t=arrival_t, event=event, key=key)
+
+ def reset(self, *, start_v: float) -> None:
+ self.resampler.reset(start_v=start_v)
+
+ async def next_realtime_window(
+ self,
+ *,
+ request: StepRequirements,
+ clock: RealtimeClock,
+ ) -> RealtimeWindowResult:
+ input_frame_count = input_frame_count_from_request(request)
+ chunk_duration_s = input_frame_count * self.resampler.dt
+ window_end_s = self.resampler.next_chunk_start_v + chunk_duration_s
+ await clock.wait_until_window_end(window_end_s)
+ catch_up = clock.catch_up(
+ request=request,
+ max_lag_s=self.max_lag_s
+ if self.max_lag_s is not None
+ else chunk_duration_s,
+ policy=self.catch_up_policy,
+ )
+ start_s = self.resampler.next_chunk_start_v
+ segments, frame_times = self.resampler.sample_chunk(input_frame_count)
+ end_s = self.resampler.next_chunk_start_v
+ window = UserInputWindow(
+ start_s=start_s,
+ end_s=end_s,
+ frame_times=tuple(frame_times),
+ inputs=UserInputs(),
+ metadata={SPARSE_KEY_SEGMENTS_METADATA_KEY: tuple(segments)},
+ )
+ return RealtimeWindowResult(window=window, catch_up=catch_up)
+
+
+def input_frame_count_from_request(request: StepRequirements) -> int:
+ """Return the positive input frame count declared by a step requirement."""
+
+ value = request.input_frame_count
+ if isinstance(value, bool) or not isinstance(value, int):
+ raise ValueError("StepRequirements.input_frame_count must be an integer.")
+ parsed = value
+ if parsed <= 0:
+ raise ValueError("StepRequirements.input_frame_count must be > 0.")
+ return parsed
+
+
+def _anchor_if_realtime(
+ clock: RealtimeClock | DeterministicClock,
+ *,
+ anchor: bool,
+) -> None:
+ if not anchor or not getattr(clock, "is_realtime", False):
+ return
+ now = getattr(clock, "now", None)
+ clock_anchor = getattr(clock, "anchor", None)
+ if callable(now) and callable(clock_anchor):
+ clock_anchor(float(now()))
+
+
+__all__ = [
+ "ActivationPolicy",
+ "ActivationResult",
+ "ActivationSignal",
+ "AlwaysActiveActivationPolicy",
+ "CatchUpDecision",
+ "CatchUpPolicy",
+ "DeterministicClock",
+ "KeyboardRealtimeInputSource",
+ "RealtimeClock",
+ "RealtimeWindowResult",
+ "ResamplerRealtimeClock",
+ "SPARSE_KEY_SEGMENTS_METADATA_KEY",
+ "SignalActivationPolicy",
+ "input_frame_count_from_request",
+]
diff --git a/flashdreams/flashdreams/runtime/demo/validation.py b/flashdreams/flashdreams/runtime/demo/validation.py
new file mode 100644
index 000000000..c5058b148
--- /dev/null
+++ b/flashdreams/flashdreams/runtime/demo/validation.py
@@ -0,0 +1,200 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Capability resolution and validation for shared demo runs."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from flashdreams.runtime.inputs import UserInputCapability, UserInputSchema
+
+from .run_modes import RunMode, RunModeCapabilities, SessionEdges
+from .session_inputs import (
+ BatchInputSource,
+ ModelInputProvider,
+ ProviderCapabilities,
+ RealtimeInputSource,
+)
+from .spec import DemoAdapter, DemoSpec
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class ResolvedRunCapabilities:
+ """Capabilities of one concrete provider/run-mode/session-edges pairing."""
+
+ finite: bool
+ deterministic: bool
+ realtime: bool
+ resettable: bool
+ produces_artifacts: bool
+
+
+def resolve_run_capabilities(
+ *,
+ spec: DemoSpec,
+ provider: ModelInputProvider,
+ session_edges: SessionEdges,
+) -> ResolvedRunCapabilities:
+ """Resolve concrete run capabilities from provider, edges, and config."""
+
+ provider_capabilities = _provider_capabilities(provider)
+ clock = session_edges.clock
+ realtime = bool(getattr(clock, "is_realtime", False))
+ deterministic_clock = (
+ bool(getattr(clock, "is_deterministic", False)) if clock is not None else True
+ )
+ config = spec.config
+ seeded = config is not None and config.seed is not None
+ return ResolvedRunCapabilities(
+ finite=bool(session_edges.input_source.is_finite),
+ deterministic=(
+ provider_capabilities.deterministic_given_inputs
+ and bool(session_edges.input_source.is_deterministic)
+ and deterministic_clock
+ and seeded
+ ),
+ realtime=realtime,
+ resettable=provider_capabilities.supports_reset,
+ produces_artifacts=bool(session_edges.output_sink.produces_artifacts),
+ )
+
+
+def validate_resolved_run(
+ *,
+ spec: DemoSpec,
+ adapter: DemoAdapter,
+ provider: ModelInputProvider,
+ run_mode: RunMode,
+ session_edges: SessionEdges,
+ resolved: ResolvedRunCapabilities,
+) -> None:
+ """Reject structurally incompatible provider/input/run-mode combinations."""
+
+ del spec, adapter
+ provider_capabilities = _provider_capabilities(provider)
+ run_mode_capabilities = _run_mode_capabilities(run_mode)
+ _validate_input_source_shape(
+ run_mode_capabilities=run_mode_capabilities,
+ session_edges=session_edges,
+ resolved=resolved,
+ )
+ _validate_provider_modes(
+ provider_capabilities=provider_capabilities,
+ run_mode_capabilities=run_mode_capabilities,
+ resolved=resolved,
+ )
+ _validate_user_input_schema(
+ provider_schema=provider_capabilities.user_input_schema,
+ source_schema=_input_source_user_input_schema(session_edges.input_source),
+ )
+
+
+def _validate_input_source_shape(
+ *,
+ run_mode_capabilities: RunModeCapabilities,
+ session_edges: SessionEdges,
+ resolved: ResolvedRunCapabilities,
+) -> None:
+ input_source = session_edges.input_source
+ if run_mode_capabilities.realtime:
+ if not isinstance(input_source, RealtimeInputSource):
+ raise ValueError("Realtime run modes require a RealtimeInputSource.")
+ if session_edges.clock is None or not resolved.realtime:
+ raise ValueError("Realtime run modes require a realtime clock.")
+ return
+ if not isinstance(input_source, BatchInputSource):
+ raise ValueError("Batch run modes require a BatchInputSource.")
+ if resolved.realtime:
+ raise ValueError("Batch run modes cannot use a realtime clock.")
+
+
+def _validate_provider_modes(
+ *,
+ provider_capabilities: ProviderCapabilities,
+ run_mode_capabilities: RunModeCapabilities,
+ resolved: ResolvedRunCapabilities,
+) -> None:
+ if run_mode_capabilities.realtime and not (
+ provider_capabilities.supports_realtime_clock
+ ):
+ raise ValueError("Provider does not support realtime input.")
+ if run_mode_capabilities.requires_finite_input:
+ if not resolved.finite:
+ raise ValueError("Run mode requires finite input.")
+ if not provider_capabilities.supports_recorded_input:
+ raise ValueError("Provider does not support recorded input.")
+ if resolved.produces_artifacts and not run_mode_capabilities.supports_artifacts:
+ raise ValueError("Run mode does not support artifact output.")
+ if (
+ run_mode_capabilities.supports_interactive_events
+ and not provider_capabilities.supports_realtime_clock
+ ):
+ raise ValueError("Interactive run mode requires realtime provider support.")
+
+
+def _validate_user_input_schema(
+ *,
+ provider_schema: UserInputSchema,
+ source_schema: UserInputSchema,
+) -> None:
+ missing = _missing_capabilities(
+ required=provider_schema.declared_capabilities(),
+ provided=source_schema,
+ )
+ if missing:
+ names = ", ".join(
+ f"{capability.event_type}[{','.join(sorted(capability.payload_fields))}]"
+ for capability in missing
+ )
+ raise ValueError(
+ f"Input source does not satisfy provider raw user input schema: {names}."
+ )
+
+
+def _missing_capabilities(
+ *,
+ required: Sequence[UserInputCapability],
+ provided: UserInputSchema,
+) -> tuple[UserInputCapability, ...]:
+ return tuple(
+ capability for capability in required if not provided.supports(capability)
+ )
+
+
+def _provider_capabilities(provider: ModelInputProvider) -> ProviderCapabilities:
+ capabilities = getattr(provider, "capabilities", None)
+ if not isinstance(capabilities, ProviderCapabilities):
+ raise TypeError(
+ "ModelInputProvider.capabilities must be a ProviderCapabilities "
+ f"instance, got {type(capabilities).__name__}."
+ )
+ return capabilities
+
+
+def _run_mode_capabilities(run_mode: RunMode) -> RunModeCapabilities:
+ capabilities = getattr(run_mode, "capabilities", None)
+ if not isinstance(capabilities, RunModeCapabilities):
+ raise TypeError(
+ "RunMode.capabilities must be a RunModeCapabilities instance, "
+ f"got {type(capabilities).__name__}."
+ )
+ return capabilities
+
+
+def _input_source_user_input_schema(input_source: object) -> UserInputSchema:
+ schema = getattr(input_source, "user_input_schema", None)
+ if not isinstance(schema, UserInputSchema):
+ raise TypeError(
+ "InputSource.user_input_schema must be a UserInputSchema instance, "
+ f"got {type(schema).__name__}."
+ )
+ return schema
+
+
+__all__ = [
+ "ResolvedRunCapabilities",
+ "resolve_run_capabilities",
+ "validate_resolved_run",
+]
diff --git a/flashdreams/flashdreams/runtime/metrics.py b/flashdreams/flashdreams/runtime/metrics.py
index 4286204f6..8a8fd6b1b 100644
--- a/flashdreams/flashdreams/runtime/metrics.py
+++ b/flashdreams/flashdreams/runtime/metrics.py
@@ -6,7 +6,8 @@
from __future__ import annotations
import math
-from collections.abc import Mapping
+from collections import Counter, defaultdict
+from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from typing import Any, Protocol, runtime_checkable
@@ -45,9 +46,31 @@ def __post_init__(self) -> None:
object.__setattr__(self, "metadata", freeze_mapping(self.metadata))
+@dataclass(frozen=True, kw_only=True, slots=True)
+class MetricsSnapshot:
+ """Closed session or run metrics summary."""
+
+ counters: Mapping[str, int | float] = field(default_factory=dict)
+ timings: Mapping[str, Sequence[float]] = field(default_factory=dict)
+ session_statuses: Sequence[str] = ()
+ errors: Sequence[str] = ()
+
+ def __post_init__(self) -> None:
+ object.__setattr__(self, "counters", freeze_mapping(self.counters))
+ object.__setattr__(
+ self,
+ "timings",
+ freeze_mapping(
+ {key: tuple(values) for key, values in self.timings.items()}
+ ),
+ )
+ object.__setattr__(self, "session_statuses", tuple(self.session_statuses))
+ object.__setattr__(self, "errors", tuple(self.errors))
+
+
@runtime_checkable
class MetricsRecorder(Protocol):
- """Collector for runtime metrics."""
+ """Collector for runtime, session, and run metrics."""
def record(self, sample: RuntimeMetricSample) -> None:
"""Record one metric sample."""
@@ -64,7 +87,53 @@ def record_timing(
"""Record one timing sample in seconds."""
...
- def close(self) -> None:
+ def record_step(
+ self,
+ *,
+ request: object,
+ user_window: object,
+ inference_input: object,
+ result: object,
+ decision: object,
+ ) -> None:
+ """Record one successful model step."""
+ ...
+
+ def record_control(
+ self,
+ *,
+ request: object,
+ user_window: object,
+ control: object,
+ ) -> None:
+ """Record one provider-authored control decision."""
+ ...
+
+ def record_error(self, exc: Exception, action: object) -> None:
+ """Record a driver-observed operational error."""
+ ...
+
+ def record_catch_up(self, decision: object) -> None:
+ """Record a realtime catch-up decision."""
+ ...
+
+ def record_cleanup_error(self, exc: Exception) -> None:
+ """Record a cleanup failure without interrupting teardown."""
+ ...
+
+ def record_orphaned_cleanup(self, exc: Exception) -> None:
+ """Record cleanup that timed out and is still queued."""
+ ...
+
+ def record_session(self, result: object) -> None:
+ """Record one closed session result."""
+ ...
+
+ def record_session_error(self, exc: Exception) -> None:
+ """Record diagnostic session assembly failure detail."""
+ ...
+
+ def close(self) -> MetricsSnapshot:
"""Finalize metric collection."""
...
@@ -74,6 +143,14 @@ class InMemoryMetricsRecorder:
"""Simple metrics recorder useful for tests, smoke runs, and adapters."""
samples: list[RuntimeMetricSample] = field(default_factory=list)
+ step_count: int = 0
+ control_count: int = 0
+ catch_up_count: int = 0
+ errors: list[str] = field(default_factory=list)
+ cleanup_errors: list[str] = field(default_factory=list)
+ orphaned_cleanup_errors: list[str] = field(default_factory=list)
+ session_errors: list[str] = field(default_factory=list)
+ sessions: list[object] = field(default_factory=list)
closed: bool = False
def record(self, sample: RuntimeMetricSample) -> None:
@@ -100,8 +177,96 @@ def record_timing(
)
)
- def close(self) -> None:
+ def record_step(
+ self,
+ *,
+ request: object,
+ user_window: object,
+ inference_input: object,
+ result: object,
+ decision: object,
+ ) -> None:
+ del request, user_window, inference_input, result, decision
+ if not self.closed:
+ self.step_count += 1
+
+ def record_control(
+ self,
+ *,
+ request: object,
+ user_window: object,
+ control: object,
+ ) -> None:
+ del request, user_window, control
+ if not self.closed:
+ self.control_count += 1
+
+ def record_error(self, exc: Exception, action: object) -> None:
+ del action
+ if not self.closed:
+ self.errors.append(str(exc))
+
+ def record_catch_up(self, decision: object) -> None:
+ del decision
+ if not self.closed:
+ self.catch_up_count += 1
+
+ def record_cleanup_error(self, exc: Exception) -> None:
+ if not self.closed:
+ self.cleanup_errors.append(str(exc))
+
+ def record_orphaned_cleanup(self, exc: Exception) -> None:
+ if not self.closed:
+ self.orphaned_cleanup_errors.append(str(exc))
+
+ def record_session(self, result: object) -> None:
+ if not self.closed:
+ self.sessions.append(result)
+
+ def record_session_error(self, exc: Exception) -> None:
+ if not self.closed:
+ self.session_errors.append(str(exc))
+
+ def close(self) -> MetricsSnapshot:
self.closed = True
+ return self.snapshot()
+
+ def snapshot(self) -> MetricsSnapshot:
+ timings: defaultdict[str, list[float]] = defaultdict(list)
+ for sample in self.samples:
+ if sample.category == "timing":
+ timings[sample.name].append(float(sample.value))
+ session_statuses = tuple(
+ str(getattr(result, "status", "unknown")) for result in self.sessions
+ )
+ session_status_counts = Counter(session_statuses)
+ return MetricsSnapshot(
+ counters={
+ "samples": len(self.samples),
+ "steps": self.step_count,
+ "controls": self.control_count,
+ "catch_ups": self.catch_up_count,
+ "sessions": len(self.sessions),
+ "errors": len(self.errors),
+ "cleanup_errors": len(self.cleanup_errors),
+ "orphaned_cleanup_errors": len(self.orphaned_cleanup_errors),
+ "session_errors": len(self.session_errors),
+ **{
+ f"sessions.{status}": count
+ for status, count in sorted(session_status_counts.items())
+ },
+ },
+ timings=timings,
+ session_statuses=session_statuses,
+ errors=tuple(
+ (
+ *self.errors,
+ *self.cleanup_errors,
+ *self.orphaned_cleanup_errors,
+ *self.session_errors,
+ )
+ ),
+ )
class NullMetricsRecorder:
@@ -120,5 +285,52 @@ def record_timing(
) -> None:
del name, duration_s, step_index, metadata
- def close(self) -> None:
- return None
+ def record_step(
+ self,
+ *,
+ request: object,
+ user_window: object,
+ inference_input: object,
+ result: object,
+ decision: object,
+ ) -> None:
+ del request, user_window, inference_input, result, decision
+
+ def record_control(
+ self,
+ *,
+ request: object,
+ user_window: object,
+ control: object,
+ ) -> None:
+ del request, user_window, control
+
+ def record_error(self, exc: Exception, action: object) -> None:
+ del exc, action
+
+ def record_catch_up(self, decision: object) -> None:
+ del decision
+
+ def record_cleanup_error(self, exc: Exception) -> None:
+ del exc
+
+ def record_orphaned_cleanup(self, exc: Exception) -> None:
+ del exc
+
+ def record_session(self, result: object) -> None:
+ del result
+
+ def record_session_error(self, exc: Exception) -> None:
+ del exc
+
+ def close(self) -> MetricsSnapshot:
+ return MetricsSnapshot()
+
+
+__all__ = [
+ "InMemoryMetricsRecorder",
+ "MetricsRecorder",
+ "MetricsSnapshot",
+ "NullMetricsRecorder",
+ "RuntimeMetricSample",
+]
diff --git a/flashdreams/flashdreams/runtime/runner.py b/flashdreams/flashdreams/runtime/runner.py
index c73a4f67c..f9a83411f 100644
--- a/flashdreams/flashdreams/runtime/runner.py
+++ b/flashdreams/flashdreams/runtime/runner.py
@@ -6,6 +6,8 @@
from __future__ import annotations
import math
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any
from flashdreams.runtime.canonical import InputCanonicalizer
from flashdreams.runtime.config import InferenceConfig
@@ -13,6 +15,7 @@
CanonicalInputs,
CanonicalInputSchema,
InferenceInput,
+ InferenceInputSchema,
TimeWindow,
UserInputs,
UserInputSchema,
@@ -29,7 +32,18 @@
)
from flashdreams.runtime.metrics import MetricsRecorder
from flashdreams.runtime.output import OutputArtifact, OutputTarget
-from flashdreams.runtime.types import StepResult
+from flashdreams.runtime.types import (
+ StepRequest,
+ StepRequirements,
+ StepResult,
+ step_requirements_from_request,
+)
+
+if TYPE_CHECKING:
+ from flashdreams.runtime.demo.host import RuntimeHost
+ from flashdreams.runtime.demo.outputs import OutputDecision, SessionInfo
+ from flashdreams.runtime.demo.run_modes import RunResult
+ from flashdreams.runtime.demo.session_inputs import PreparedStep, UserInputWindow
_DEFAULT_SESSION_HORIZON_S = 3600.0
@@ -46,73 +60,567 @@ def run_inference_session(
output: OutputTarget,
metrics: MetricsRecorder,
) -> tuple[OutputArtifact, ...]:
- """Run one sequential inference session through the standard loop.
+ """Run one sequential inference session through the shared batch driver.
- This v0 loop intentionally handles one adapter/runtime/session, one selected
- input mapping, one replay/live input batch, one output target, and one
- metrics recorder. It is synchronous and owns only orchestration.
+ The signature and failure semantics remain compatible with the original
+ replay runner while the implementation delegates the step loop to the shared
+ demo runtime pipeline.
"""
- runtime: InferenceRuntime | None = None
- session: InferenceSession | None = None
- output_opened = False
- output_artifacts: tuple[OutputArtifact, ...] = ()
+ return _run_inference_session_with_shared_batch(
+ adapter=adapter,
+ config=config,
+ mapping=mapping,
+ canonicalizer=canonicalizer,
+ source_schema=source_schema,
+ user_inputs=user_inputs,
+ initial_inputs=initial_inputs,
+ output=output,
+ metrics=metrics,
+ )
+
+
+def _run_inference_session_with_shared_batch(
+ *,
+ adapter: ModelAdapter,
+ config: InferenceConfig,
+ mapping: InputMapping,
+ canonicalizer: InputCanonicalizer,
+ source_schema: UserInputSchema,
+ user_inputs: UserInputs,
+ initial_inputs: InferenceInput,
+ output: OutputTarget,
+ metrics: MetricsRecorder,
+) -> tuple[OutputArtifact, ...]:
+ lifecycle = _LegacyBatchLifecycle()
+ request_state = _LegacyStepRequestState()
+ runtime = _LegacyLazyRuntime(
+ adapter=adapter,
+ config=config,
+ lifecycle=lifecycle,
+ request_state=request_state,
+ )
+ from flashdreams.runtime.demo.host import RuntimeHost
+
+ host = RuntimeHost(runtime)
+ metrics_recorder = _LegacyRunnerMetricsRecorder(metrics)
+ output_sink = _LegacyOutputTargetSink(
+ output=output,
+ lifecycle=lifecycle,
+ host=host,
+ )
primary_error: BaseException | None = None
try:
- adapter.validate_config(config)
- canonical_schema = canonicalizer.canonical_schema(source_schema)
- _check_declared_mapping_compatibility(
- mapping=mapping,
- canonical_schema=canonical_schema,
+ _validate_legacy_runner_inputs(
adapter=adapter,
+ config=config,
+ mapping=mapping,
+ canonicalizer=canonicalizer,
+ source_schema=source_schema,
)
- mapping.validate(
- canonical_schema=canonical_schema,
+ provider = _LegacyMappedModelInputProvider(
+ mapping=mapping,
+ canonicalizer=canonicalizer,
+ source_schema=source_schema,
+ user_inputs=user_inputs,
+ initial_inputs=initial_inputs,
inference_input_schema=adapter.inference_input_schema,
+ request_state=request_state,
)
- canonicalizer.reset()
- mapped_initial_inputs = mapping.map_global_conditioning_inputs(
- canonical_inputs=CanonicalInputs(),
- inference_input=initial_inputs,
+ input_source = _LegacyBatchInputSource(
+ source_schema=source_schema,
+ user_inputs=user_inputs,
+ request_state=request_state,
)
- runtime = adapter.create_runtime(config)
- session = runtime.start_session(mapped_initial_inputs)
- output.open()
- output_opened = True
- step_base_inputs = InferenceInput(
- step=initial_inputs.step,
- metadata=initial_inputs.metadata,
+ result = _run_shared_batch_session(
+ host=host,
+ provider=provider,
+ input_source=input_source,
+ output_sink=output_sink,
+ metrics=metrics_recorder,
)
-
- while (request := session.next_step_request()) is not None:
- step_inputs = mapping.map_step_inputs(
- canonical_inputs=canonicalizer.canonicalize(
- user_inputs,
- window=request.user_input_window
- or _all_user_inputs_window(user_inputs),
- source_schema=source_schema,
- ),
- inference_input=step_base_inputs,
- request=request,
- )
- result = session.step(step_inputs)
- output.write(result)
- _record_timing_metrics(metrics, result)
+ _raise_legacy_runner_error(
+ result=result,
+ output_sink=output_sink,
+ metrics=metrics_recorder,
+ )
+ return tuple(result.artifacts)
except BaseException as exc:
primary_error = exc
+ if not metrics_recorder.closed:
+ _close_metrics_suppressing_secondary(metrics_recorder)
raise
finally:
- cleanup_error, output_artifacts = _close_run_resources(
- output=output if output_opened else None,
- session=session,
- runtime=runtime,
+ try:
+ host.close()
+ except BaseException:
+ if primary_error is None:
+ raise
+
+
+def _validate_legacy_runner_inputs(
+ *,
+ adapter: ModelAdapter,
+ config: InferenceConfig,
+ mapping: InputMapping,
+ canonicalizer: InputCanonicalizer,
+ source_schema: UserInputSchema,
+) -> CanonicalInputSchema:
+ adapter.validate_config(config)
+ canonical_schema = canonicalizer.canonical_schema(source_schema)
+ _check_declared_mapping_compatibility(
+ mapping=mapping,
+ canonical_schema=canonical_schema,
+ adapter=adapter,
+ )
+ mapping.validate(
+ canonical_schema=canonical_schema,
+ inference_input_schema=adapter.inference_input_schema,
+ )
+ return canonical_schema
+
+
+def _run_shared_batch_session(
+ *,
+ host: RuntimeHost,
+ provider: "_LegacyMappedModelInputProvider",
+ input_source: "_LegacyBatchInputSource",
+ output_sink: "_LegacyOutputTargetSink",
+ metrics: "_LegacyRunnerMetricsRecorder",
+) -> RunResult:
+ from flashdreams.runtime.demo.drivers import BatchSessionDriver
+ from flashdreams.runtime.demo.pipeline import StepPipeline
+ from flashdreams.runtime.demo.run_modes import SessionEdges
+
+ return BatchSessionDriver().run_one_session(
+ host=host,
+ provider=provider,
+ session_edges=SessionEdges(
+ input_source=input_source,
+ output_sink=output_sink,
+ cleanup_tasks=set(),
metrics=metrics,
+ ),
+ pipeline=StepPipeline(),
+ )
+
+
+class _LegacyLazyRuntime:
+ """Create the legacy runtime only after global inputs are mapped."""
+
+ def __init__(
+ self,
+ *,
+ adapter: ModelAdapter,
+ config: InferenceConfig,
+ lifecycle: "_LegacyBatchLifecycle",
+ request_state: "_LegacyStepRequestState",
+ ) -> None:
+ self._adapter = adapter
+ self._config = config
+ self._lifecycle = lifecycle
+ self._request_state = request_state
+
+ def start_session(self, inputs: InferenceInput) -> InferenceSession:
+ runtime = self._lifecycle.runtime
+ if runtime is None:
+ runtime = self._adapter.create_runtime(self._config)
+ self._lifecycle.set_runtime(runtime)
+ session = runtime.start_session(inputs)
+ self._lifecycle.set_session(session)
+ return _LegacySessionAdapter(
+ session=session,
+ request_state=self._request_state,
+ )
+
+ def close(self) -> None:
+ self._lifecycle.close_runtime_direct()
+
+
+class _LegacyBatchLifecycle:
+ """Own legacy model resources whose close order is caller-visible."""
+
+ def __init__(self) -> None:
+ self.runtime: InferenceRuntime | None = None
+ self.session: InferenceSession | None = None
+ self._session_close_attempted = False
+ self._runtime_close_attempted = False
+
+ def set_runtime(self, runtime: InferenceRuntime) -> None:
+ self.runtime = runtime
+
+ def set_session(self, session: InferenceSession) -> None:
+ self.session = session
+ self._session_close_attempted = False
+
+ def close_session_via_host(self, host: RuntimeHost) -> None:
+ host.call(self.close_session_direct)
+
+ def close_runtime_via_host(self, host: RuntimeHost) -> None:
+ host.call(self.close_runtime_direct)
+
+ def close_session_direct(self) -> None:
+ if self.session is None or self._session_close_attempted:
+ return
+ self._session_close_attempted = True
+ self.session.close()
+
+ def close_runtime_direct(self) -> None:
+ if self.runtime is None or self._runtime_close_attempted:
+ return
+ self._runtime_close_attempted = True
+ self.runtime.close()
+
+
+class _LegacySessionAdapter:
+ """Expose old sessions through the new StepRequirements boundary."""
+
+ def __init__(
+ self,
+ *,
+ session: InferenceSession,
+ request_state: "_LegacyStepRequestState",
+ ) -> None:
+ self._session = session
+ self._request_state = request_state
+
+ def next_step_requirements(self) -> StepRequirements | None:
+ request = self._session.next_step_request()
+ if request is None:
+ self._request_state.clear()
+ return None
+ self._request_state.store(request)
+ return step_requirements_from_request(
+ request,
+ allow_user_input_window=True,
+ )
+
+ def next_step_request(self) -> StepRequest | None:
+ return self._session.next_step_request()
+
+ def session_info(self) -> SessionInfo:
+ from flashdreams.runtime.demo.outputs import SessionInfo
+
+ session_info = getattr(self._session, "session_info", None)
+ if not callable(session_info):
+ return SessionInfo()
+ value = session_info()
+ if not isinstance(value, SessionInfo):
+ raise TypeError(
+ "session.session_info() must return SessionInfo, "
+ f"got {type(value).__name__}."
+ )
+ return value
+
+ def step(self, inputs: InferenceInput) -> StepResult:
+ return self._session.step(inputs)
+
+ def reset(self, inputs: InferenceInput | None = None) -> None:
+ self._session.reset(inputs)
+
+ def close(self) -> None:
+ return None
+
+
+class _LegacyStepRequestState:
+ """Share the current legacy request between the session, source, and provider."""
+
+ def __init__(self) -> None:
+ self._request: StepRequest | None = None
+
+ def store(self, request: StepRequest) -> None:
+ self._request = request
+
+ def require_for_window(self, step_index: int) -> StepRequest:
+ request = self._request
+ if request is None:
+ raise RuntimeError("Legacy input source has no active step request.")
+ if request.step_index != step_index:
+ raise RuntimeError(
+ "Legacy input source request mismatch: "
+ f"expected step {request.step_index}, got {step_index}."
+ )
+ return request
+
+ def consume_for_step(self, step_index: int) -> StepRequest:
+ request = self.require_for_window(step_index)
+ self._request = None
+ return request
+
+ def clear(self) -> None:
+ self._request = None
+
+
+class _LegacyBatchInputSource:
+ is_finite = True
+ is_deterministic = True
+
+ def __init__(
+ self,
+ *,
+ source_schema: UserInputSchema,
+ user_inputs: UserInputs,
+ request_state: _LegacyStepRequestState,
+ ) -> None:
+ self.user_input_schema = source_schema
+ self._user_inputs = user_inputs
+ self._request_state = request_state
+
+ def is_finished(self) -> bool:
+ return False
+
+ def next_window(self, request: StepRequirements) -> UserInputWindow:
+ from flashdreams.runtime.demo.session_inputs import UserInputWindow
+
+ legacy_request = self._request_state.require_for_window(request.step_index)
+ window = legacy_request.user_input_window or _all_user_inputs_window(
+ self._user_inputs
+ )
+ return UserInputWindow(
+ start_s=window.start_s,
+ end_s=window.end_s,
+ inputs=self._user_inputs,
+ )
+
+
+class _LegacyMappedModelInputProvider:
+ def __init__(
+ self,
+ *,
+ mapping: InputMapping,
+ canonicalizer: InputCanonicalizer,
+ source_schema: UserInputSchema,
+ user_inputs: UserInputs,
+ initial_inputs: InferenceInput,
+ inference_input_schema: InferenceInputSchema,
+ request_state: _LegacyStepRequestState,
+ ) -> None:
+ from flashdreams.runtime.demo.session_inputs import ProviderCapabilities
+
+ self.capabilities = ProviderCapabilities(
+ supports_recorded_input=True,
+ deterministic_given_inputs=True,
+ user_input_schema=source_schema,
+ inference_input_schema=inference_input_schema,
+ )
+ self._mapping = mapping
+ self._canonicalizer = canonicalizer
+ self._source_schema = source_schema
+ self._user_inputs = user_inputs
+ self._initial_inputs = initial_inputs
+ self._request_state = request_state
+ self._step_base_inputs = InferenceInput(
+ step=initial_inputs.step,
+ metadata=initial_inputs.metadata,
+ )
+
+ def prepare_initial_input(self) -> InferenceInput:
+ self._canonicalizer.reset()
+ return self._mapping.map_global_conditioning_inputs(
+ canonical_inputs=CanonicalInputs(),
+ inference_input=self._initial_inputs,
)
- if cleanup_error is not None and primary_error is None:
+
+ def prepare_step(
+ self,
+ *,
+ request: StepRequirements,
+ user_window: UserInputWindow,
+ ) -> PreparedStep:
+ from flashdreams.runtime.demo.session_inputs import PreparedStep
+
+ legacy_request = self._request_state.consume_for_step(request.step_index)
+ canonical_inputs = self._canonicalizer.canonicalize(
+ self._user_inputs,
+ window=TimeWindow(start_s=user_window.start_s, end_s=user_window.end_s),
+ source_schema=self._source_schema,
+ )
+ return PreparedStep(
+ inference_input=self._mapping.map_step_inputs(
+ canonical_inputs=canonical_inputs,
+ inference_input=self._step_base_inputs,
+ request=legacy_request,
+ )
+ )
+
+ def reset(self, inputs: InferenceInput | None = None) -> None:
+ del inputs
+ self._canonicalizer.reset()
+
+ def close(self) -> None:
+ return None
+
+
+class _LegacyOutputTargetSink:
+ produces_artifacts = True
+
+ def __init__(
+ self,
+ *,
+ output: OutputTarget,
+ lifecycle: _LegacyBatchLifecycle,
+ host: RuntimeHost,
+ ) -> None:
+ self._output = output
+ self._lifecycle = lifecycle
+ self._host = host
+ self._opened = False
+ self._closed = False
+ self._artifacts: tuple[OutputArtifact, ...] = ()
+ self.cleanup_error: BaseException | None = None
+
+ def open(self, session_info: SessionInfo) -> None:
+ del session_info
+ self._output.open()
+ self._opened = True
+ self._closed = False
+ self._artifacts = ()
+ self.cleanup_error = None
+
+ def begin_generation(self, generation: int) -> None:
+ del generation
+
+ def write(self, result: StepResult) -> OutputDecision:
+ from flashdreams.runtime.demo.outputs import OutputDecision
+
+ self._output.write(result)
+ return OutputDecision()
+
+ def close(self) -> Sequence[OutputArtifact]:
+ if self._closed:
+ if self.cleanup_error is not None:
+ raise self.cleanup_error
+ return self._artifacts
+
+ self._closed = True
+ cleanup_error: BaseException | None = None
+ artifacts: tuple[OutputArtifact, ...] = ()
+
+ def remember_error(exc: BaseException) -> None:
+ nonlocal cleanup_error
+ if cleanup_error is None:
+ cleanup_error = exc
+
+ if self._opened:
+ try:
+ artifacts = tuple(self._output.close())
+ except BaseException as exc:
+ remember_error(exc)
+ try:
+ self._lifecycle.close_session_via_host(self._host)
+ except BaseException as exc:
+ remember_error(exc)
+ try:
+ self._lifecycle.close_runtime_via_host(self._host)
+ except BaseException as exc:
+ remember_error(exc)
+
+ self._artifacts = artifacts
+ self.cleanup_error = cleanup_error
+ if cleanup_error is not None:
raise cleanup_error
+ return self._artifacts
+
+
+class _LegacyRunnerMetricsRecorder:
+ def __init__(self, metrics: MetricsRecorder) -> None:
+ self._metrics = metrics
+ self.closed = False
+ self.close_error: BaseException | None = None
- return output_artifacts
+ def record(self, sample: Any) -> None:
+ self._metrics.record(sample)
+
+ def record_timing(
+ self,
+ name: str,
+ duration_s: float,
+ *,
+ step_index: int | None = None,
+ metadata: Mapping[str, Any] | None = None,
+ ) -> None:
+ self._metrics.record_timing(
+ name,
+ duration_s,
+ step_index=step_index,
+ metadata=metadata,
+ )
+
+ def record_step(
+ self,
+ *,
+ request: object,
+ user_window: object,
+ inference_input: object,
+ result: object,
+ decision: object,
+ ) -> None:
+ del request, user_window, inference_input, decision
+ if isinstance(result, StepResult):
+ _record_timing_metrics(self._metrics, result)
+
+ def record_control(
+ self,
+ *,
+ request: object,
+ user_window: object,
+ control: object,
+ ) -> None:
+ del request, user_window, control
+
+ def record_error(self, exc: Exception, action: object) -> None:
+ del exc, action
+
+ def record_catch_up(self, decision: object) -> None:
+ del decision
+
+ def record_cleanup_error(self, exc: Exception) -> None:
+ del exc
+
+ def record_orphaned_cleanup(self, exc: Exception) -> None:
+ del exc
+
+ def record_session(self, result: object) -> None:
+ del result
+
+ def record_session_error(self, exc: Exception) -> None:
+ del exc
+
+ def close(self) -> Any:
+ if self.closed:
+ if self.close_error is not None:
+ raise self.close_error
+ return None
+ self.closed = True
+ try:
+ return self._metrics.close()
+ except BaseException as exc:
+ self.close_error = exc
+ raise
+
+
+def _raise_legacy_runner_error(
+ *,
+ result: RunResult,
+ output_sink: _LegacyOutputTargetSink,
+ metrics: _LegacyRunnerMetricsRecorder,
+) -> None:
+ if result.error is not None:
+ raise result.error
+ if output_sink.cleanup_error is not None:
+ raise output_sink.cleanup_error
+ if metrics.close_error is not None:
+ raise metrics.close_error
+
+
+def _close_metrics_suppressing_secondary(
+ metrics: _LegacyRunnerMetricsRecorder,
+) -> None:
+ try:
+ metrics.close()
+ except BaseException:
+ return
def _check_declared_mapping_compatibility(
@@ -158,45 +666,4 @@ def _record_timing_metrics(
)
-def _close_run_resources(
- *,
- output: OutputTarget | None,
- session: InferenceSession | None,
- runtime: InferenceRuntime | None,
- metrics: MetricsRecorder,
-) -> tuple[BaseException | None, tuple[OutputArtifact, ...]]:
- cleanup_error: BaseException | None = None
- artifacts: tuple[OutputArtifact, ...] = ()
-
- def remember_error(exc: BaseException) -> None:
- nonlocal cleanup_error
- if cleanup_error is None:
- cleanup_error = exc
-
- if output is not None:
- try:
- artifacts = tuple(output.close())
- except BaseException as exc:
- remember_error(exc)
-
- if session is not None:
- try:
- session.close()
- except BaseException as exc:
- remember_error(exc)
-
- if runtime is not None:
- try:
- runtime.close()
- except BaseException as exc:
- remember_error(exc)
-
- try:
- metrics.close()
- except BaseException as exc:
- remember_error(exc)
-
- return cleanup_error, artifacts
-
-
__all__ = ["run_inference_session"]
diff --git a/flashdreams/flashdreams/runtime/types.py b/flashdreams/flashdreams/runtime/types.py
index 4130a925f..49ee958a0 100644
--- a/flashdreams/flashdreams/runtime/types.py
+++ b/flashdreams/flashdreams/runtime/types.py
@@ -13,6 +13,65 @@
from flashdreams.runtime._utils import freeze_mapping
from flashdreams.runtime.inputs import InferenceInputSchema, TimeWindow
+_STEP_REQUIREMENTS_INPUT_COUNT_METADATA_KEY = "input_frame_count"
+_STEP_REQUIREMENTS_STEADY_OUTPUT_COUNT_METADATA_KEY = "steady_output_frame_count"
+_STEP_REQUIREMENTS_USER_INPUT_METADATA_KEYS = frozenset(
+ {
+ "input_window",
+ "user_input",
+ "user_input_window",
+ "user_inputs",
+ }
+)
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class StepRequirements:
+ """Model-authored per-step requirements consumed by shared demo drivers."""
+
+ __hash__ = None
+
+ step_index: int
+ input_frame_count: int = 1
+ steady_output_frame_count: int | None = None
+ inference_input_schema: InferenceInputSchema | None = None
+ metadata: Mapping[str, Any] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ if isinstance(self.step_index, bool) or not isinstance(self.step_index, int):
+ raise TypeError("StepRequirements.step_index must be an integer.")
+ if self.step_index < 0:
+ raise ValueError("StepRequirements.step_index must be >= 0.")
+ if isinstance(self.input_frame_count, bool) or not isinstance(
+ self.input_frame_count, int
+ ):
+ raise TypeError("StepRequirements.input_frame_count must be an integer.")
+ if self.input_frame_count <= 0:
+ raise ValueError("StepRequirements.input_frame_count must be > 0.")
+ if self.steady_output_frame_count is not None:
+ if isinstance(self.steady_output_frame_count, bool) or not isinstance(
+ self.steady_output_frame_count, int
+ ):
+ raise TypeError(
+ "StepRequirements.steady_output_frame_count must be an integer."
+ )
+ if self.steady_output_frame_count < 0:
+ raise ValueError(
+ "StepRequirements.steady_output_frame_count must be >= 0."
+ )
+ user_input_keys = sorted(
+ key
+ for key in self.metadata
+ if key in _STEP_REQUIREMENTS_USER_INPUT_METADATA_KEYS
+ )
+ if user_input_keys:
+ joined = ", ".join(user_input_keys)
+ raise ValueError(
+ "StepRequirements.metadata must not include driver-owned user "
+ f"input keys: {joined}."
+ )
+ object.__setattr__(self, "metadata", freeze_mapping(self.metadata))
+
@dataclass(frozen=True, kw_only=True, slots=True)
class StepRequest:
@@ -36,4 +95,36 @@ def __post_init__(self) -> None:
object.__setattr__(self, "metadata", freeze_mapping(self.metadata))
-__all__ = ["StepRequest", "StepResult"]
+def step_requirements_from_request(
+ request: StepRequest,
+ *,
+ allow_user_input_window: bool = False,
+) -> StepRequirements:
+ """Adapt a legacy ``StepRequest`` that did not carry driver-owned inputs."""
+
+ if request.user_input_window is not None and not allow_user_input_window:
+ raise ValueError(
+ "StepRequest.user_input_window cannot be adapted to StepRequirements; "
+ "user input windows are driver-owned."
+ )
+ metadata = dict(request.metadata)
+ input_frame_count = metadata.pop(_STEP_REQUIREMENTS_INPUT_COUNT_METADATA_KEY, 1)
+ steady_output_frame_count = metadata.pop(
+ _STEP_REQUIREMENTS_STEADY_OUTPUT_COUNT_METADATA_KEY,
+ None,
+ )
+ return StepRequirements(
+ step_index=request.step_index,
+ input_frame_count=input_frame_count,
+ steady_output_frame_count=steady_output_frame_count,
+ inference_input_schema=request.inference_input_schema,
+ metadata=metadata,
+ )
+
+
+__all__ = [
+ "StepRequest",
+ "StepRequirements",
+ "StepResult",
+ "step_requirements_from_request",
+]
diff --git a/flashdreams/flashdreams/runtime/worker.py b/flashdreams/flashdreams/runtime/worker.py
index af4576c09..e523b6e6a 100644
--- a/flashdreams/flashdreams/runtime/worker.py
+++ b/flashdreams/flashdreams/runtime/worker.py
@@ -6,15 +6,17 @@
from __future__ import annotations
import asyncio
+import threading
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Callable, TypeVar, cast
import torch
_T = TypeVar("_T")
+_EXECUTOR_FUTURE_POLL_INTERVAL_S = 0.01
-class ThreadAffineRuntimeWorker:
+class ModelExecutionWorker:
"""Run ordered runtime lifecycle calls on one owned OS thread.
CUDA graphs, Triton launchers, and some backend contexts are thread-local.
@@ -38,14 +40,23 @@ def __init__(
thread_name_prefix=thread_name,
initializer=self._initialize_thread,
)
+ self._state_lock = threading.Lock()
self._accepting = True
self._closed = False
- self._close_lock = asyncio.Lock()
+ self._thread_id: int | None = None
@property
def closed(self) -> bool:
return self._closed
+ @property
+ def worker_thread_id(self) -> int | None:
+ return self._thread_id
+
+ @property
+ def is_worker_thread(self) -> bool:
+ return self._thread_id == threading.get_ident()
+
async def call(
self,
func: Callable[..., _T],
@@ -54,11 +65,11 @@ async def call(
**kwargs: Any,
) -> _T:
"""Run one callable after all previously submitted worker calls."""
- if not self._accepting:
- raise RuntimeError("runtime worker is closed")
+ self._require_not_worker_thread()
+ self._require_accepting()
future = self._submit(func, args, kwargs)
try:
- return await asyncio.shield(future)
+ return await _await_executor_future(future)
except asyncio.CancelledError:
future.add_done_callback(_consume_exception)
raise
@@ -71,21 +82,32 @@ def call_blocking(
**kwargs: Any,
) -> _T:
"""Run one callable from synchronous code on the owned worker thread."""
- if not self._accepting:
- raise RuntimeError("runtime worker is closed")
+ self._require_not_worker_thread()
+ self._require_accepting()
future = self._executor.submit(_invoke, func, args, kwargs)
return cast(_T, future.result())
async def close(self) -> None:
"""Drain submitted work and stop accepting lifecycle calls."""
- async with self._close_lock:
- if self._closed:
- return
- self._accepting = False
- barrier = self._submit(_noop, (), {})
- await asyncio.shield(barrier)
- self._executor.shutdown(wait=True, cancel_futures=False)
- self._closed = True
+ self._require_not_worker_thread()
+ if not self._begin_close():
+ return
+ try:
+ barrier = asyncio.wrap_future(self._executor.submit(_noop))
+ await _await_executor_future(barrier)
+ finally:
+ self._finish_close()
+
+ def close_blocking(self) -> None:
+ """Synchronous close for non-async setup and teardown paths."""
+ self._require_not_worker_thread()
+ if not self._begin_close():
+ return
+ try:
+ barrier = self._executor.submit(_noop)
+ barrier.result()
+ finally:
+ self._finish_close()
def _submit(
self,
@@ -97,9 +119,36 @@ def _submit(
return loop.run_in_executor(self._executor, _invoke, func, args, kwargs)
def _initialize_thread(self) -> None:
+ self._thread_id = threading.get_ident()
if self._device is not None and self._device.type == "cuda":
torch.cuda.set_device(self._device)
+ def _require_accepting(self) -> None:
+ if not self._accepting:
+ raise RuntimeError("runtime worker is closed")
+
+ def _require_not_worker_thread(self) -> None:
+ if self.is_worker_thread:
+ raise RuntimeError(
+ "Cannot dispatch to the model execution worker from its own "
+ "thread; call the function directly."
+ )
+
+ def _begin_close(self) -> bool:
+ with self._state_lock:
+ if self._closed:
+ return False
+ self._accepting = False
+ return True
+
+ def _finish_close(self) -> None:
+ self._executor.shutdown(wait=True, cancel_futures=False)
+ with self._state_lock:
+ self._closed = True
+
+
+ThreadAffineRuntimeWorker = ModelExecutionWorker
+
def _invoke(
func: Callable[..., _T],
@@ -113,9 +162,19 @@ def _noop() -> None:
return
+async def _await_executor_future(future: asyncio.Future[_T]) -> _T:
+ """Await an executor future without relying on a single cross-thread wakeup."""
+ while not future.done():
+ await asyncio.wait(
+ {future},
+ timeout=_EXECUTOR_FUTURE_POLL_INTERVAL_S,
+ )
+ return future.result()
+
+
def _consume_exception(future: asyncio.Future[Any]) -> None:
if not future.cancelled():
future.exception()
-__all__ = ["ThreadAffineRuntimeWorker"]
+__all__ = ["ModelExecutionWorker", "ThreadAffineRuntimeWorker"]
diff --git a/flashdreams/flashdreams/serving/realtime/timing.py b/flashdreams/flashdreams/serving/realtime/timing.py
index 9891cfcdc..30147f97d 100644
--- a/flashdreams/flashdreams/serving/realtime/timing.py
+++ b/flashdreams/flashdreams/serving/realtime/timing.py
@@ -11,6 +11,8 @@
from threading import Lock
from typing import Protocol
+from flashdreams.runtime.metrics import MetricsRecorder
+
TraceComponentValue = str | int | float | bool | None
@@ -402,6 +404,36 @@ def summarize_chunk_history(chunks: Iterable[ChunkTimes]) -> RecentTimingSummary
)
+def record_chunk_timing_metrics(
+ metrics: MetricsRecorder,
+ chunk: ChunkTimes,
+) -> None:
+ """Record available chunk timing durations to a session metrics recorder."""
+
+ _record_stage_timing_metrics(
+ metrics,
+ prefix="realtime.chunk",
+ durations_ms=chunk.stage_durations_ms(),
+ step_index=chunk.chunk_index,
+ )
+
+
+def record_video_model_timing_metrics(
+ metrics: MetricsRecorder,
+ timings: VideoModelTimings,
+ *,
+ chunk_index: int | None = None,
+) -> None:
+ """Record backend-visible video model stage durations to session metrics."""
+
+ _record_stage_timing_metrics(
+ metrics,
+ prefix="realtime.model",
+ durations_ms=timings.stage_durations_ms(),
+ step_index=chunk_index,
+ )
+
+
class RollingChunkTimingSummary:
def __init__(self, capacity: int) -> None:
self._chunks: deque[ChunkTimes] = deque(maxlen=capacity)
@@ -518,6 +550,24 @@ def _add_optional_trace_range(
)
+def _record_stage_timing_metrics(
+ metrics: MetricsRecorder,
+ *,
+ prefix: str,
+ durations_ms: Mapping[str, float],
+ step_index: int | None,
+) -> None:
+ for stage_name, duration_ms in durations_ms.items():
+ try:
+ metrics.record_timing(
+ f"{prefix}.{stage_name}",
+ float(duration_ms) / 1000.0,
+ step_index=step_index,
+ )
+ except Exception:
+ return
+
+
def _summarize_values(values: list[float]) -> StageDurationSummary:
ordered = sorted(values)
count = len(ordered)
diff --git a/flashdreams/flashdreams/serving/webrtc/encoders.py b/flashdreams/flashdreams/serving/webrtc/encoders.py
index 5eea69991..b07d9324c 100644
--- a/flashdreams/flashdreams/serving/webrtc/encoders.py
+++ b/flashdreams/flashdreams/serving/webrtc/encoders.py
@@ -20,7 +20,7 @@
import importlib.util
from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, runtime_checkable
import torch
from aiortc import MediaStreamTrack
@@ -69,6 +69,20 @@ class VideoEncoder(Protocol):
def create_track(self, *, maxsize: int) -> BufferedVideoTrack | NVENCVideoTrack: ...
+ def prepare_chunk_payload(
+ self,
+ result: StepResult,
+ track: MediaStreamTrack,
+ ) -> object: ...
+
+ async def deliver_prepared_chunk(
+ self,
+ payload: object,
+ track: MediaStreamTrack,
+ *,
+ force_keyframe: bool = False,
+ ) -> ChunkDeliveryResult: ...
+
async def deliver_chunk(
self,
result: StepResult,
@@ -110,10 +124,24 @@ def create_track(self, *, maxsize: int) -> BufferedVideoTrack:
return BufferedVideoTrack(fps=self.fps, maxsize=maxsize)
- async def deliver_chunk(
+ def prepare_chunk_payload(
self,
result: StepResult,
track: MediaStreamTrack,
+ ) -> tuple[object, ...]:
+ from flashdreams.serving.webrtc.media import BufferedVideoTrack
+
+ if not isinstance(track, BufferedVideoTrack):
+ raise TypeError(
+ "DefaultRTCEncoder requires a BufferedVideoTrack; got "
+ f"{type(track).__name__}. Create it via encoder.create_track()."
+ )
+ return track.prepare_result_frames(result)
+
+ async def deliver_prepared_chunk(
+ self,
+ payload: object,
+ track: MediaStreamTrack,
*,
force_keyframe: bool = False,
) -> ChunkDeliveryResult:
@@ -128,7 +156,9 @@ async def deliver_chunk(
"DefaultRTCEncoder requires a BufferedVideoTrack; got "
f"{type(track).__name__}. Create it via encoder.create_track()."
)
- enqueued = await track.enqueue_result(result)
+ if not isinstance(payload, tuple):
+ raise TypeError("DefaultRTCEncoder payload must be a tuple of RGB frames.")
+ enqueued = await track.enqueue_frames(cast(Any, payload))
return ChunkDeliveryResult(
backend=self.backend,
num_frames=enqueued,
@@ -136,6 +166,19 @@ async def deliver_chunk(
encode_ms=0.0,
)
+ async def deliver_chunk(
+ self,
+ result: StepResult,
+ track: MediaStreamTrack,
+ *,
+ force_keyframe: bool = False,
+ ) -> ChunkDeliveryResult:
+ return await self.deliver_prepared_chunk(
+ self.prepare_chunk_payload(result, track),
+ track,
+ force_keyframe=force_keyframe,
+ )
+
def close(self) -> None:
return
diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py
index a7c6c87a4..4291cd637 100644
--- a/flashdreams/flashdreams/serving/webrtc/manager.py
+++ b/flashdreams/flashdreams/serving/webrtc/manager.py
@@ -10,10 +10,10 @@
import inspect
import json
from collections import deque
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
from collections.abc import Set as AbstractSet
from dataclasses import dataclass, field, replace
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, TypeVar, cast
from aiortc import (
RTCConfiguration,
@@ -23,12 +23,37 @@
)
from loguru import logger
+from flashdreams.runtime.demo import (
+ SPARSE_KEY_SEGMENTS_METADATA_KEY,
+ DemoSpec,
+ InMemorySessionMetricsRecorder,
+ ModelInputProvider,
+ PreparedScenario,
+ PreparedStep,
+ ProviderCapabilities,
+ ResamplerRealtimeClock,
+ RunContext,
+ RunResult,
+ RuntimeHost,
+ SessionEdges,
+ SessionInfo,
+ SingleSessionAdmissionPolicy,
+ StepPipeline,
+ UserInputWindow,
+ WebRTCErrorPolicy,
+ WebRTCOutputSpec,
+ run_demo_session_async,
+)
from flashdreams.runtime.inputs import (
+ CanonicalInputSchema,
InferenceInput,
+ InferenceInputSchema,
TimeWindow,
UserInputEvent,
UserInputs,
+ UserInputSchema,
)
+from flashdreams.runtime.mapping import InputMapping
from flashdreams.runtime.types import StepRequest, StepResult
from flashdreams.serving.realtime.input import (
DEFAULT_SUPPORTED_KEYS,
@@ -37,7 +62,9 @@
)
from flashdreams.serving.webrtc.encoders import (
DefaultRTCEncoder,
+ EncoderBackend,
VideoEncoder,
+ select_encoder,
)
from flashdreams.serving.webrtc.media import BufferedVideoTrack, NVENCVideoTrack
from flashdreams.serving.webrtc.messages import (
@@ -55,6 +82,18 @@
WebRTCSessionRuntime,
)
from flashdreams.serving.webrtc.server import SessionBusyError
+from flashdreams.serving.webrtc.services import (
+ WEBRTC_SKIPPED_INPUTS_METADATA_KEY,
+ WEBRTC_SKIPPED_WINDOW_METADATA_KEY,
+ WEBRTC_USER_INPUT_SCHEMA,
+ ThreadSafeWebRTCOutputBridge,
+ WebRTCActivationPolicy,
+ WebRTCChunkDelivery,
+ WebRTCInputSource,
+ WebRTCOutputSink,
+ WebRTCRunMode,
+ WebRTCTransportService,
+)
from flashdreams.serving.webrtc.warmup import (
run_loopback_warmup_session,
wait_for_ice_gathering_complete,
@@ -78,8 +117,12 @@
"""Maximum unconsumed raw events kept for an ``InferenceSession`` step."""
_RELEASE_USER_EVENT_TYPES = frozenset({"key_up"})
_KEY_USER_EVENT_TYPES = frozenset({"key_down", "key_up"})
+_SESSION_INPUT_KEY = "webrtc_session_input"
+_STEP_REQUEST_KEY = "webrtc_step_request"
+_SEGMENTS_KEY = "webrtc_segments"
+_FRAME_TIMES_KEY = "webrtc_frame_times"
-_RuntimeT = TypeVar("_RuntimeT", bound=WebRTCSessionRuntime)
+_RuntimeT = TypeVar("_RuntimeT")
_RuntimeConfigT = TypeVar("_RuntimeConfigT", bound=WebRTCRuntimeConfig)
@@ -140,6 +183,421 @@ def _stat_int(stats: Mapping[str, float | int], name: str) -> int:
return int(round(_stat_float(stats, name)))
+def _runtime_drives_inference_session(runtime: Any) -> bool:
+ return callable(getattr(runtime, "start_inference_session", None))
+
+
+def _run_on_event_loop(loop: asyncio.AbstractEventLoop, awaitable: Any) -> Any:
+ """Run one legacy async WebRTC runtime call from a RuntimeHost worker."""
+ return asyncio.run_coroutine_threadsafe(awaitable, loop).result()
+
+
+def _step_request_from_requirements(
+ request: Any,
+ *,
+ window: TimeWindow,
+) -> StepRequest:
+ metadata = dict(getattr(request, "metadata", {}))
+ metadata["input_frame_count"] = request.input_frame_count
+ steady_output_frame_count = getattr(request, "steady_output_frame_count", None)
+ if steady_output_frame_count is not None:
+ metadata["steady_output_frame_count"] = steady_output_frame_count
+ return StepRequest(
+ step_index=request.step_index,
+ inference_input_schema=getattr(request, "inference_input_schema", None),
+ user_input_window=window,
+ metadata=metadata,
+ )
+
+
+def _encoder_backend_from_config(value: object) -> EncoderBackend:
+ backend = str(value)
+ if backend not in {"auto", "default", "nvenc"}:
+ raise ValueError(
+ f"encoder_backend must be 'auto', 'default', or 'nvenc', got {backend!r}."
+ )
+ return cast(EncoderBackend, backend)
+
+
+def _gpu_id_from_device_spec(device_spec: str) -> int:
+ if not device_spec.startswith("cuda"):
+ return 0
+ _prefix, separator, index = device_spec.partition(":")
+ if not separator or not index:
+ return 0
+ try:
+ return int(index)
+ except ValueError:
+ return 0
+
+
+class _LegacyWebRTCRuntimeAdapter:
+ """Shared compatibility adapter from old async WebRTC runtimes to RuntimeHost."""
+
+ def __init__(self, *, runtime: Any, loop: asyncio.AbstractEventLoop) -> None:
+ self._runtime = runtime
+ self._loop = loop
+
+ def reset_for_new_session(self, session_input: Any = None) -> None:
+ _run_on_event_loop(
+ self._loop,
+ self._runtime.reset_for_new_session(session_input=session_input),
+ )
+
+ def start_session(self, inputs: InferenceInput) -> "_LegacyWebRTCSessionAdapter":
+ del inputs
+ inference_session = None
+ if _runtime_drives_inference_session(self._runtime):
+ inference_session = _run_on_event_loop(
+ self._loop,
+ self._runtime.start_inference_session(),
+ )
+ return _LegacyWebRTCSessionAdapter(
+ runtime=self._runtime,
+ inference_session=inference_session,
+ loop=self._loop,
+ )
+
+ def close(self) -> None:
+ # The underlying async runtime is owned by BaseWebRTCSessionManager and
+ # closed from shutdown(); RuntimeHost only owns this adapter's worker.
+ return
+
+
+class _LegacyWebRTCSessionAdapter:
+ """RuntimeHost-facing session view over a legacy WebRTC runtime/session."""
+
+ def __init__(
+ self,
+ *,
+ runtime: Any,
+ inference_session: Any | None,
+ loop: asyncio.AbstractEventLoop,
+ ) -> None:
+ self._runtime = runtime
+ self._inference_session = inference_session
+ self._loop = loop
+
+ def session_info(self) -> SessionInfo:
+ steady_frames: int | None = None
+ try:
+ steady_frames = int(self._runtime.peek_steady_output_num_frames())
+ except Exception:
+ steady_frames = None
+ return SessionInfo(steady_output_frame_count=steady_frames)
+
+ def next_step_request(self) -> StepRequest | None:
+ if self._inference_session is not None:
+ return self._inference_session.next_step_request()
+ return self._runtime.next_step_request()
+
+ def step(self, inputs: InferenceInput) -> StepResult:
+ if self._inference_session is not None:
+ result = self._inference_session.step(inputs)
+ else:
+ result = _run_on_event_loop(
+ self._loop,
+ self._runtime.step(
+ request=inputs.step[_STEP_REQUEST_KEY],
+ segments=list(inputs.step[_SEGMENTS_KEY]),
+ frame_times=list(inputs.step[_FRAME_TIMES_KEY]),
+ ),
+ )
+ request = inputs.step[_STEP_REQUEST_KEY]
+ if result.step_index != request.step_index:
+ raise RuntimeError(
+ "Runtime result step does not match its request: "
+ f"requested {request.step_index}, got {result.step_index}."
+ )
+ if not isinstance(result, StepResult):
+ raise TypeError(
+ "WebRTC session steps must produce StepResult, got "
+ f"{type(result).__name__}."
+ )
+ return result
+
+ def reset(self, inputs: InferenceInput | None = None) -> None:
+ session_input = None
+ if inputs is not None:
+ session_input = inputs.global_conditioning.get(_SESSION_INPUT_KEY)
+ _run_on_event_loop(
+ self._loop,
+ self._runtime.reset_for_new_session(session_input=session_input),
+ )
+ if self._inference_session is not None:
+ self._inference_session = _run_on_event_loop(
+ self._loop,
+ self._runtime.start_inference_session(),
+ )
+
+ def close(self) -> None:
+ close = getattr(self._inference_session, "close", None)
+ if callable(close):
+ close()
+
+
+class _LegacyWebRTCModelInputProvider:
+ """Shared provider used until model-specific WebRTC providers land."""
+
+ def __init__(self, *, runtime: Any, session_input: Any = None) -> None:
+ self._runtime = runtime
+ self._session_input = session_input
+ self._uses_inference_session = _runtime_drives_inference_session(runtime)
+ self._session_input_state_advanced = False
+ self.capabilities = ProviderCapabilities(
+ supports_realtime_clock=True,
+ supports_reset=True,
+ deterministic_given_inputs=False,
+ user_input_schema=self._user_input_schema(),
+ )
+
+ def prepare_initial_input(self) -> InferenceInput:
+ if self._session_input is None:
+ return InferenceInput()
+ return InferenceInput(
+ global_conditioning={_SESSION_INPUT_KEY: self._session_input}
+ )
+
+ def prepare_step(
+ self,
+ *,
+ request: Any,
+ user_window: UserInputWindow,
+ ) -> PreparedStep:
+ if self._uses_inference_session:
+ return PreparedStep(
+ inference_input=self._prepare_inference_session_step(
+ request=request,
+ user_window=user_window,
+ )
+ )
+ return PreparedStep(
+ inference_input=self._prepare_segment_step(
+ request=request,
+ user_window=user_window,
+ )
+ )
+
+ def reset(self, inputs: InferenceInput | None = None) -> None:
+ del inputs
+ self._session_input_state_advanced = False
+
+ def close(self) -> None:
+ return
+
+ def _user_input_schema(self) -> UserInputSchema:
+ schema = getattr(self._runtime, "input_source_schema", None)
+ if isinstance(schema, UserInputSchema):
+ return schema
+ return WEBRTC_USER_INPUT_SCHEMA
+
+ def _prepare_inference_session_step(
+ self,
+ *,
+ request: Any,
+ user_window: UserInputWindow,
+ ) -> InferenceInput:
+ self._advance_skipped_input_state(user_window)
+ window_start = user_window.start_s
+ if request.step_index == 0 and not self._session_input_state_advanced:
+ window_start = 0.0
+ window = TimeWindow(start_s=window_start, end_s=user_window.end_s)
+ canonical_inputs = self._runtime.input_canonicalizer.canonicalize(
+ user_window.inputs,
+ window=window,
+ source_schema=self._runtime.input_source_schema,
+ )
+ mapping = self._runtime.input_mapping
+ inference_input = InferenceInput(
+ metadata={
+ **dict(user_window.metadata),
+ "frame_times": tuple(user_window.frame_times),
+ "window_start_s": window.start_s,
+ "window_end_s": window.end_s,
+ }
+ )
+ return mapping.map_step_inputs(
+ canonical_inputs=canonical_inputs,
+ inference_input=inference_input,
+ request=_step_request_from_requirements(request, window=window),
+ )
+
+ def _advance_skipped_input_state(self, user_window: UserInputWindow) -> None:
+ skipped_inputs = user_window.metadata.get(WEBRTC_SKIPPED_INPUTS_METADATA_KEY)
+ skipped_window = user_window.metadata.get(WEBRTC_SKIPPED_WINDOW_METADATA_KEY)
+ if not isinstance(skipped_inputs, UserInputs):
+ return
+ if not isinstance(skipped_window, tuple) or len(skipped_window) != 2:
+ return
+ start_value, end_value = skipped_window
+ if not isinstance(start_value, int | float) or not isinstance(
+ end_value,
+ int | float,
+ ):
+ return
+ start_s = float(start_value)
+ end_s = float(end_value)
+ if end_s <= start_s:
+ return
+ self._runtime.input_canonicalizer.canonicalize(
+ skipped_inputs,
+ window=TimeWindow(start_s=start_s, end_s=end_s),
+ source_schema=self._runtime.input_source_schema,
+ )
+ self._session_input_state_advanced = True
+
+ @staticmethod
+ def _prepare_segment_step(
+ *,
+ request: Any,
+ user_window: UserInputWindow,
+ ) -> InferenceInput:
+ segments = user_window.metadata.get(SPARSE_KEY_SEGMENTS_METADATA_KEY)
+ if not isinstance(segments, tuple):
+ raise RuntimeError("WebRTC user window is missing resampled key segments.")
+ window = TimeWindow(start_s=user_window.start_s, end_s=user_window.end_s)
+ return InferenceInput(
+ step={
+ _STEP_REQUEST_KEY: _step_request_from_requirements(
+ request,
+ window=window,
+ ),
+ _SEGMENTS_KEY: tuple(segments),
+ _FRAME_TIMES_KEY: tuple(user_window.frame_times),
+ }
+ )
+
+
+class _LegacyWebRTCDemoAdapter:
+ """Minimal adapter for the shared helper while WebRTC providers migrate."""
+
+ model_id: str
+ inference_input_schema = InferenceInputSchema()
+ canonical_input_schema = CanonicalInputSchema()
+
+ def __init__(
+ self,
+ *,
+ runtime: Any,
+ identity: str,
+ session_input: Any = None,
+ ) -> None:
+ self._runtime = runtime
+ self.model_id = identity
+ self._session_input = session_input
+
+ def supported_input_modes(self) -> tuple[str, ...]:
+ return ("webrtc",)
+
+ def supported_output_modes(self) -> tuple[str, ...]:
+ return ("webrtc",)
+
+ def default_input_mapping(self) -> InputMapping | None:
+ return None
+
+ def validate_config(self, config: Any) -> None:
+ if config.model_id != self.model_id:
+ raise ValueError(
+ f"Expected WebRTC model_id={self.model_id!r}, got {config.model_id!r}."
+ )
+
+ def create_runtime(self, config: Any) -> Any:
+ self.validate_config(config)
+ return self._runtime
+
+ def prepare_scenario(self, spec: Any) -> PreparedScenario:
+ del spec
+ return PreparedScenario(initial_inputs=self._initial_inputs())
+
+ def create_model_input_provider(
+ self,
+ spec: Any,
+ scenario: PreparedScenario,
+ ) -> _LegacyWebRTCModelInputProvider:
+ del spec, scenario
+ return _LegacyWebRTCModelInputProvider(
+ runtime=self._runtime,
+ session_input=self._session_input,
+ )
+
+ def _initial_inputs(self) -> InferenceInput:
+ if self._session_input is None:
+ return InferenceInput()
+ return InferenceInput(
+ global_conditioning={_SESSION_INPUT_KEY: self._session_input}
+ )
+
+
+class _ManagedWebRTCSessionEdgeFactory:
+ """Build shared realtime edges for one negotiated peer connection."""
+
+ def __init__(
+ self,
+ *,
+ manager: "BaseWebRTCSessionManager[Any, Any]",
+ managed_session: "ManagedWebRTCSession",
+ loop: asyncio.AbstractEventLoop,
+ ) -> None:
+ self._manager = manager
+ self._managed_session = managed_session
+ self._loop = loop
+
+ def create_session_edges(
+ self,
+ *,
+ context: RunContext,
+ spec: Any,
+ scenario: PreparedScenario,
+ provider: ModelInputProvider,
+ adapter: Any,
+ ) -> SessionEdges:
+ del spec, scenario, provider, adapter
+ input_source = self._managed_session.input_source
+ transport = self._managed_session.transport
+ if input_source is None or transport is None:
+ raise RuntimeError("Managed WebRTC session is missing shared edges.")
+ bridge = ThreadSafeWebRTCOutputBridge(
+ loop=self._loop,
+ video_encoder=self._managed_session.video_encoder,
+ video_track=self._managed_session.video_track,
+ on_chunk_delivery=self._on_chunk_delivery,
+ on_error=self._on_delivery_error,
+ )
+ return SessionEdges(
+ input_source=input_source,
+ output_sink=WebRTCOutputSink(bridge=bridge),
+ cleanup_tasks=context.cleanup_tasks,
+ metrics=InMemorySessionMetricsRecorder(),
+ error_policy=WebRTCErrorPolicy(),
+ transport=transport,
+ clock=ResamplerRealtimeClock(
+ resampler=self._managed_session.resampler,
+ now_fn=self._loop.time,
+ sleep_fn=asyncio.sleep,
+ ),
+ activation=WebRTCActivationPolicy(
+ input_source=input_source,
+ transport=transport,
+ ),
+ )
+
+ def _on_chunk_delivery(self, chunk: WebRTCChunkDelivery) -> None:
+ self._manager._handle_shared_chunk_delivery(
+ managed_session=self._managed_session,
+ chunk=chunk,
+ )
+
+ def _on_delivery_error(self, exc: BaseException) -> None:
+ self._manager._handle_shared_delivery_error(
+ managed_session=self._managed_session,
+ exc=exc,
+ )
+ if self._manager.fatal_generation_errors:
+ self._loop.call_soon_threadsafe(
+ lambda: asyncio.create_task(self._manager.close_active_session())
+ )
+
+
@dataclass(slots=True)
class ManagedWebRTCSession:
"""Per-session state for the single active WebRTC peer connection."""
@@ -152,6 +610,9 @@ class ManagedWebRTCSession:
control_channel: Any | None = None
generation_task: asyncio.Task[Any] | None = None
first_action_received: asyncio.Event = field(default_factory=asyncio.Event)
+ input_source: WebRTCInputSource | None = None
+ transport: WebRTCTransportService | None = None
+ reservation: Any | None = None
pending_action_arrivals: deque[float] = field(default_factory=deque)
inference_session: Any | None = None
"""Active ``InferenceSession``; ``None`` means call ``runtime.generate_chunk``."""
@@ -189,6 +650,11 @@ async def close(self) -> None:
self.generation_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await self.generation_task
+ if self.generation_task is None or self.generation_task.done():
+ reservation = self.reservation
+ self.reservation = None
+ if reservation is not None:
+ reservation.release()
self.generation_task = None
await self.video_track.close()
@@ -212,6 +678,11 @@ def __init__(
supported_control_keys: AbstractSet[str] | None = None,
fatal_generation_errors: bool = False,
client_liveness_timeout_s: float = DEFAULT_CLIENT_LIVENESS_TIMEOUT_S,
+ shared_host: RuntimeHost | None = None,
+ shared_adapter: Any | None = None,
+ shared_spec: DemoSpec | None = None,
+ shared_scenario: PreparedScenario | None = None,
+ shared_pipeline_factory: Callable[[], StepPipeline] | None = None,
) -> None:
if client_liveness_timeout_s <= 0:
raise ValueError("client_liveness_timeout_s must be > 0")
@@ -234,6 +705,15 @@ def __init__(
self._preload_lock = asyncio.Lock()
self._session_lock = asyncio.Lock()
self._pending_session_input: Any = None
+ self._shared_runtime_adapter: _LegacyWebRTCRuntimeAdapter | None = None
+ self._shared_host: RuntimeHost | None = shared_host
+ self._owns_shared_host = shared_host is not None
+ self._shared_context: RunContext | None = None
+ self._shared_adapter = shared_adapter
+ self._shared_spec = shared_spec
+ self._shared_scenario = shared_scenario
+ self._shared_pipeline_factory = shared_pipeline_factory
+ self._shared_video_encoder: VideoEncoder | None = None
@property
def pending_session_input(self) -> Any:
@@ -296,8 +776,11 @@ def _positive_float_runtime_value(value: Any, *, label: str) -> float:
return parsed
def _runtime_input_fps(self, runtime: Any) -> float:
+ peek_input_fps = getattr(runtime, "peek_input_fps", None)
+ if not callable(peek_input_fps):
+ return float(self.fps)
return self._positive_float_runtime_value(
- runtime.peek_input_fps(),
+ peek_input_fps(),
label="peek_input_fps",
)
@@ -315,9 +798,22 @@ def _runtime_next_step_request(self, runtime: Any) -> tuple[StepRequest, int]:
return request, input_num_frames
def _runtime_steady_output_num_frames(self, runtime: Any) -> int:
+ peek_output_frames = getattr(runtime, "peek_steady_output_num_frames", None)
+ if callable(peek_output_frames):
+ return self._positive_int_runtime_value(
+ peek_output_frames(),
+ label="peek_steady_output_num_frames",
+ )
+ pipeline = getattr(runtime, "pipeline", None)
+ get_num_frames = getattr(pipeline, "get_num_frames", None)
+ if callable(get_num_frames):
+ return self._positive_int_runtime_value(
+ get_num_frames(1),
+ label="pipeline.get_num_frames(1)",
+ )
return self._positive_int_runtime_value(
- runtime.peek_steady_output_num_frames(),
- label="peek_steady_output_num_frames",
+ 1,
+ label="fallback steady output frame count",
)
def _resolve_video_encoder(self) -> VideoEncoder:
@@ -330,10 +826,60 @@ def _resolve_video_encoder(self) -> VideoEncoder:
transparently get the software path without having to opt in.
"""
encoder = getattr(self._runtime, "video_encoder", None)
+ if encoder is None:
+ encoder = self._shared_video_encoder
if encoder is None:
encoder = DefaultRTCEncoder(fps=self.fps)
return encoder
+ def _shared_run_context(self, loop: asyncio.AbstractEventLoop) -> RunContext:
+ if self._shared_context is not None:
+ return self._shared_context
+ host = self._shared_host
+ if host is None:
+ runtime_adapter = _LegacyWebRTCRuntimeAdapter(
+ runtime=self._runtime,
+ loop=loop,
+ )
+ host = RuntimeHost(runtime_adapter)
+ self._shared_runtime_adapter = runtime_adapter
+ self._shared_host = host
+ self._shared_context = RunContext(
+ host=host,
+ run_metrics=InMemorySessionMetricsRecorder(),
+ admission=SingleSessionAdmissionPolicy(
+ health_check=lambda: host.is_healthy
+ ),
+ )
+ return self._shared_context
+
+ def _shared_demo_spec(self) -> DemoSpec:
+ return DemoSpec(
+ model_id=self.identity,
+ input_mode="webrtc",
+ output=WebRTCOutputSpec(
+ fps=self.fps,
+ video_width=self.runtime_config.video_width,
+ video_height=self.runtime_config.video_height,
+ warmup_chunks=self.runtime_config.warmup_chunks,
+ warmup_timeout_s=self.runtime_config.warmup_timeout_s,
+ client_liveness_timeout_s=self.client_liveness_timeout_s,
+ ),
+ )
+
+ async def _reset_runtime_for_session(
+ self,
+ *,
+ context: RunContext,
+ session_input: Any,
+ ) -> None:
+ reset = getattr(context.host.runtime, "reset_for_new_session", None)
+ if not callable(reset):
+ if self._shared_adapter is not None:
+ return
+ raise RuntimeError("WebRTC runtime adapter cannot reset sessions.")
+ await context.host.call_async(reset, session_input)
+
def _prefer_h264_video_codec(self, *, transceiver: Any) -> None:
"""Constrain the transceiver's codec preferences to H.264 variants.
@@ -496,7 +1042,7 @@ async def _handle_event_message(
@staticmethod
def _drives_inference_session(runtime: Any) -> bool:
"""Return whether ``runtime`` should be driven through ``InferenceSession``."""
- return callable(getattr(runtime, "start_inference_session", None))
+ return _runtime_drives_inference_session(runtime)
def _record_user_event(
self,
@@ -769,14 +1315,46 @@ def is_runtime_ready(self) -> bool:
async def preload_runtime(self) -> None:
async with self._preload_lock:
if not self._runtime_ready:
- await self._runtime.initialize()
+ initialize = getattr(self._runtime, "initialize", None)
+ if callable(initialize):
+ result = initialize()
+ if inspect.isawaitable(result):
+ await result
+ elif self._shared_host is not None:
+ await asyncio.to_thread(self._shared_host.preload)
self._runtime_ready = True
+ self._initialize_shared_video_encoder()
if not self._warmup_complete:
await self._run_loopback_warmup_session(
num_chunks=self.runtime_config.warmup_chunks
)
self._warmup_complete = True
+ def _initialize_shared_video_encoder(self) -> None:
+ if self._shared_video_encoder is not None:
+ return
+ if getattr(self._runtime, "video_encoder", None) is not None:
+ return
+ encoder_backend = getattr(self.runtime_config, "encoder_backend", None)
+ if encoder_backend is None:
+ return
+ backend = _encoder_backend_from_config(encoder_backend)
+ device_spec = str(getattr(self.runtime_config, "device", ""))
+ device_type = device_spec.split(":", maxsplit=1)[0]
+ if device_type != "cuda" and backend == "auto":
+ backend = "default"
+ if device_type != "cuda" and backend == "nvenc":
+ raise RuntimeError("encoder_backend='nvenc' requires a CUDA device.")
+ self._shared_video_encoder = select_encoder(
+ backend=backend,
+ width=self.runtime_config.video_width,
+ height=self.runtime_config.video_height,
+ fps=self.fps,
+ bitrate=int(getattr(self.runtime_config, "encoder_bitrate_bps", 6_000_000)),
+ gpu_id=_gpu_id_from_device_spec(device_spec),
+ gop=int(getattr(self.runtime_config, "encoder_gop", self.fps)),
+ )
+
async def create_answer(self, *, offer_sdp: str, offer_type: str) -> dict[str, str]:
if not self._runtime_ready or not self._warmup_complete:
await self.preload_runtime()
@@ -808,46 +1386,61 @@ async def _create_answer_with_runtime_ready_locked(
if not self._runtime_ready:
raise RuntimeError("Runtime is not initialized.")
- await self._runtime.reset_for_new_session(session_input=session_input)
-
- peer_connection = RTCPeerConnection(rtc_configuration)
- # Bounded queue sized to one *steady-state* chunk so the producer
- # is throttled to the consumer's drain rate. AR step 0 emits fewer
- # frames than steady state; sizing to it would force a per-chunk
- # stall, so we size to the steady-state count.
- num_frames = self._runtime_steady_output_num_frames(self._runtime)
- video_encoder = self._resolve_video_encoder()
- video_track = video_encoder.create_track(maxsize=num_frames)
- # Use ``addTransceiver`` (not ``addTrack``) so we can constrain the
- # SDP m-line's codec list via ``setCodecPreferences`` when the
- # encoder emits pre-encoded H.264 packets.
- video_transceiver = peer_connection.addTransceiver(
- video_track,
- direction="sendonly",
- )
- if video_encoder.prefers_codec == "h264":
- self._prefer_h264_video_codec(transceiver=video_transceiver)
- # Start the resampler's virtual clock at 0; the real anchor is set
- # in the ``on_datachannel`` handler so chunk 0's window starts when
- # input can actually arrive.
- resampler = self._make_resampler_at_fps(
- start_v=0.0,
- fps=self._runtime_input_fps(self._runtime),
- )
loop = asyncio.get_running_loop()
- managed_session = ManagedWebRTCSession(
- runtime=self._runtime,
- video_track=video_track,
- video_encoder=video_encoder,
- peer_connection=peer_connection,
- resampler=resampler,
- last_client_message_at=loop.time(),
- )
- session_runtime: Any = self._runtime
- if self._drives_inference_session(session_runtime):
- managed_session.inference_session = (
- await session_runtime.start_inference_session()
+ context = self._shared_run_context(loop)
+ reservation = context.admission.try_reserve()
+ if reservation is None:
+ raise SessionBusyError(self.busy_message)
+ try:
+ await self._reset_runtime_for_session(
+ context=context,
+ session_input=session_input,
+ )
+ except Exception:
+ reservation.release()
+ raise
+
+ try:
+ peer_connection = RTCPeerConnection(rtc_configuration)
+ # Bounded queue sized to one *steady-state* chunk so the producer
+ # is throttled to the consumer's drain rate. AR step 0 emits fewer
+ # frames than steady state; sizing to it would force a per-chunk
+ # stall, so we size to the steady-state count.
+ num_frames = self._runtime_steady_output_num_frames(self._runtime)
+ video_encoder = self._resolve_video_encoder()
+ video_track = video_encoder.create_track(maxsize=num_frames)
+ # Use ``addTransceiver`` (not ``addTrack``) so we can constrain the
+ # SDP m-line's codec list via ``setCodecPreferences`` when the
+ # encoder emits pre-encoded H.264 packets.
+ video_transceiver = peer_connection.addTransceiver(
+ video_track,
+ direction="sendonly",
+ )
+ if video_encoder.prefers_codec == "h264":
+ self._prefer_h264_video_codec(transceiver=video_transceiver)
+ # Start the resampler's virtual clock at 0; the real anchor is set
+ # in the ``on_datachannel`` handler so chunk 0's window starts when
+ # input can actually arrive.
+ resampler = self._make_resampler_at_fps(
+ start_v=0.0,
+ fps=self._runtime_input_fps(self._runtime),
)
+ input_source = WebRTCInputSource(resampler=resampler)
+ transport = WebRTCTransportService(loop=loop)
+ managed_session = ManagedWebRTCSession(
+ runtime=self._runtime,
+ video_track=video_track,
+ video_encoder=video_encoder,
+ peer_connection=peer_connection,
+ resampler=resampler,
+ input_source=input_source,
+ transport=transport,
+ reservation=reservation,
+ last_client_message_at=loop.time(),
+ )
+ except Exception:
+ reservation.release()
+ raise
self._active_session = managed_session
if enable_liveness_watchdog:
managed_session.liveness_task = asyncio.create_task(
@@ -858,10 +1451,13 @@ async def _create_answer_with_runtime_ready_locked(
def on_datachannel(channel: Any) -> None:
managed_session.control_channel = channel
# Re-anchor the resampler at channel open. The real
- # virtual-clock anchor happens in ``_generation_worker`` once
- # the first keyboard event arrives.
+ # virtual-clock anchor happens in ``WebRTCActivationPolicy`` once
+ # the first browser event activates the shared realtime driver.
channel_open_v = asyncio.get_running_loop().time()
- managed_session.resampler.reset(start_v=channel_open_v)
+ if managed_session.input_source is not None:
+ managed_session.input_source.reset(start_v=channel_open_v)
+ else:
+ managed_session.resampler.reset(start_v=channel_open_v)
@channel.on("message")
def on_message(message: Any) -> None:
@@ -872,15 +1468,21 @@ def on_message(message: Any) -> None:
)
)
- # Spawn the generation worker once the channel is wired up so
+ # Spawn the shared realtime session once the channel is wired up so
# ``chunk_done`` notifications have a channel to land on.
managed_session.generation_task = asyncio.create_task(
- self._generation_worker(managed_session=managed_session)
+ self._run_realtime_driver_session(
+ managed_session=managed_session,
+ context=context,
+ session_input=session_input,
+ )
)
@channel.on("close")
def on_close() -> None:
logger.info("Control data channel closed; closing active session.")
+ if managed_session.transport is not None:
+ managed_session.transport.disconnect("data channel closed")
asyncio.create_task(self.close_active_session())
@peer_connection.on("connectionstatechange")
@@ -993,15 +1595,37 @@ async def _client_liveness_watchdog(
async def shutdown(self) -> None:
await self.close_active_session()
- await self._runtime.close()
+ if self._shared_context is not None:
+ await self._shared_context.close_async()
+ if self._shared_host is not None:
+ await asyncio.to_thread(self._shared_host.close)
+ self._shared_context = None
+ self._shared_host = None
+ self._shared_runtime_adapter = None
+ if self._shared_video_encoder is not None:
+ self._shared_video_encoder.close()
+ self._shared_video_encoder = None
+ if not self._owns_shared_host:
+ close = getattr(self._runtime, "close", None)
+ if callable(close):
+ result = close()
+ if inspect.isawaitable(result):
+ await result
self._runtime_ready = False
self._warmup_complete = False
def wait_for_termination(self) -> None:
- self._runtime.wait_for_termination()
+ wait = getattr(self._runtime, "wait_for_termination", None)
+ if callable(wait):
+ wait()
+ return
+ if self._shared_host is not None:
+ self._shared_host.run_worker_loop()
def send_exit_signal(self) -> None:
- self._runtime.send_exit_signal()
+ send = getattr(self._runtime, "send_exit_signal", None)
+ if callable(send):
+ send()
async def _handle_datachannel_message(
self,
@@ -1013,6 +1637,10 @@ async def _handle_datachannel_message(
if channel is None or managed_session.closed:
return
managed_session.last_client_message_at = asyncio.get_running_loop().time()
+ if managed_session.transport is not None:
+ managed_session.transport.mark_client_message(
+ managed_session.last_client_message_at
+ )
if not isinstance(raw_message, str):
self._send_json(channel, make_error_payload("Expected text payload."))
@@ -1029,6 +1657,12 @@ async def _handle_datachannel_message(
channel, make_error_payload("Payload must be a JSON object.")
)
return
+ if managed_session.input_source is not None:
+ await self._handle_shared_datachannel_payload(
+ managed_session=managed_session,
+ payload=payload,
+ )
+ return
message_type = str(payload.get("type", "")).strip().lower()
if message_type == MESSAGE_TYPE_HEARTBEAT:
return
@@ -1106,6 +1740,202 @@ async def _handle_datachannel_message(
# user actually interacts. Idempotent once already set.
managed_session.first_action_received.set()
+ async def _handle_shared_datachannel_payload(
+ self,
+ *,
+ managed_session: ManagedWebRTCSession,
+ payload: dict[str, Any],
+ ) -> None:
+ channel = managed_session.control_channel
+ input_source = managed_session.input_source
+ if channel is None or input_source is None:
+ return
+ message_type = str(payload.get("type", "")).strip().lower()
+ if message_type == MESSAGE_TYPE_HEARTBEAT:
+ return
+ if message_type == MESSAGE_TYPE_DISCONNECT:
+ logger.info("Client requested disconnect; closing active session.")
+ if managed_session.transport is not None:
+ managed_session.transport.disconnect("client disconnected")
+ await self.close_active_session()
+ return
+ if message_type == MESSAGE_TYPE_EVENT:
+ handled = self._record_shared_event_payload(
+ managed_session=managed_session,
+ payload=payload,
+ )
+ if handled:
+ managed_session.first_action_received.set()
+ return
+ result = input_source.handle_browser_payload(
+ payload,
+ timestamp_s=asyncio.get_running_loop().time(),
+ )
+ if result.kind == "error":
+ self._send_json(channel, make_error_payload(result.error or "Bad input."))
+ return
+ if result.activated:
+ managed_session.first_action_received.set()
+
+ def _record_shared_event_payload(
+ self,
+ *,
+ managed_session: ManagedWebRTCSession,
+ payload: dict[str, Any],
+ ) -> bool:
+ channel = managed_session.control_channel
+ input_source = managed_session.input_source
+ if channel is None or input_source is None:
+ return False
+ event_id = str(payload.get("event_id", payload.get("id", ""))).strip()
+ state = str(payload.get("state", "trigger")).strip().lower() or "trigger"
+ clear_states = {"clear", "release", "off", "none"}
+ if not event_id and state not in clear_states:
+ self._send_json(
+ channel,
+ make_error_payload(
+ (
+ "Event payload must include non-empty 'event_id' "
+ "unless state clears the active event."
+ ),
+ ),
+ )
+ return False
+ clears = state in clear_states
+ try:
+ event_payload = self._validate_user_event_payload(
+ managed_session=managed_session,
+ event_type="text_event",
+ payload={
+ "event_id": None if clears else event_id,
+ "state": state,
+ },
+ )
+ active_event_id = event_payload.get("event_id")
+ source_event_id = None if active_event_id is None else str(active_event_id)
+ input_source.record_user_event(
+ timestamp_s=asyncio.get_running_loop().time(),
+ event_type="text_event",
+ payload=event_payload,
+ source_event_id=source_event_id,
+ )
+ except Exception as exc:
+ self._send_json(channel, make_error_payload(str(exc)))
+ return False
+ active_event_id = event_payload.get("event_id")
+ ack_event_id = None if active_event_id is None else str(active_event_id)
+ self._send_json(
+ channel,
+ make_event_ack_payload(
+ event_id=ack_event_id,
+ state=str(event_payload.get("state", state)),
+ result={"active_event_id": ack_event_id},
+ ),
+ )
+ return True
+
+ async def _run_realtime_driver_session(
+ self,
+ *,
+ managed_session: ManagedWebRTCSession,
+ context: RunContext,
+ session_input: Any,
+ ) -> None:
+ adapter = self._shared_adapter
+ spec = self._shared_spec
+ scenario = self._shared_scenario
+ if adapter is None or spec is None:
+ adapter = _LegacyWebRTCDemoAdapter(
+ runtime=self._runtime,
+ identity=self.identity,
+ session_input=session_input,
+ )
+ spec = self._shared_demo_spec()
+ if scenario is None:
+ scenario = adapter.prepare_scenario(spec)
+ run_mode = WebRTCRunMode(
+ edge_factory=_ManagedWebRTCSessionEdgeFactory(
+ manager=self,
+ managed_session=managed_session,
+ loop=asyncio.get_running_loop(),
+ )
+ )
+ try:
+ result = await run_demo_session_async(
+ context=context,
+ spec=spec,
+ scenario=scenario,
+ adapter=adapter,
+ run_mode=run_mode,
+ pipeline=(
+ self._shared_pipeline_factory()
+ if self._shared_pipeline_factory is not None
+ else StepPipeline()
+ ),
+ reservation=managed_session.reservation,
+ )
+ if result.status == "completed":
+ logger.info("Shared WebRTC session completed.")
+ else:
+ logger.warning(
+ "Shared WebRTC session ended with status={} reason={}",
+ result.status,
+ result.reason,
+ )
+ if result.status != "completed" and result.reason:
+ channel = managed_session.control_channel
+ if channel is not None:
+ self._send_json(channel, make_error_payload(result.reason))
+ finally:
+ managed_session.reservation = None
+ if self._active_session is managed_session:
+ await self.close_active_session()
+
+ def _handle_shared_chunk_delivery(
+ self,
+ *,
+ managed_session: ManagedWebRTCSession,
+ chunk: WebRTCChunkDelivery,
+ ) -> None:
+ channel = managed_session.control_channel
+ if channel is None or managed_session.closed:
+ return
+ delivery = chunk.delivery
+ enqueued_frames = int(getattr(delivery, "num_frames", chunk.frame_count))
+ encode_ms = float(getattr(delivery, "encode_ms", 0.0))
+ play_ms = chunk.frame_count * 1000.0 / managed_session.video_track.fps
+ queue_depth = managed_session.video_track.qsize()
+ self._send_json(
+ channel,
+ make_chunk_done_payload(
+ chunk_index=chunk.step_index,
+ num_frames=chunk.frame_count,
+ enqueued_frames=enqueued_frames,
+ fps=managed_session.video_track.fps,
+ width=self.runtime_config.video_width,
+ height=self.runtime_config.video_height,
+ model=self.identity,
+ gen_ms=_stat_ms(chunk.metrics, "model_step_s"),
+ enqueue_ms=encode_ms,
+ play_ms=play_ms,
+ queue_depth=queue_depth,
+ lag_ms=0.0,
+ control_latency_ms=None,
+ consumed_actions=0,
+ extra=chunk.metadata,
+ ),
+ )
+
+ def _handle_shared_delivery_error(
+ self,
+ *,
+ managed_session: ManagedWebRTCSession,
+ exc: BaseException,
+ ) -> None:
+ channel = managed_session.control_channel
+ if channel is not None:
+ self._send_json(channel, make_error_payload(str(exc)))
+
async def _generation_worker(
self, *, managed_session: ManagedWebRTCSession
) -> None:
diff --git a/flashdreams/flashdreams/serving/webrtc/media.py b/flashdreams/flashdreams/serving/webrtc/media.py
index 470569132..25fa09437 100644
--- a/flashdreams/flashdreams/serving/webrtc/media.py
+++ b/flashdreams/flashdreams/serving/webrtc/media.py
@@ -81,16 +81,37 @@ def maxsize(self) -> int:
def qsize(self) -> int:
return self._frames.qsize()
- async def enqueue_result(self, result: StepResult) -> int:
+ def prepare_result_frames(self, result: StepResult) -> tuple[np.ndarray, ...]:
+ if self._closed:
+ return ()
+ return tuple(self._frame_converter(result))
+
+ async def enqueue_frames(self, frames: Sequence[np.ndarray]) -> int:
if self._closed:
return 0
- frames = await asyncio.to_thread(self._frame_converter, result)
for i, frame in enumerate(frames):
if self._closed:
return i
await self._frames.put(frame)
return len(frames)
+ async def enqueue_result(self, result: StepResult) -> int:
+ if self._closed:
+ return 0
+ frames = await asyncio.to_thread(self.prepare_result_frames, result)
+ return await self.enqueue_frames(frames)
+
+ async def flush(self) -> None:
+ """Drop queued frames while keeping the RTP timestamp sequence alive."""
+ if self._closed:
+ return
+ while True:
+ try:
+ self._frames.get_nowait()
+ except asyncio.QueueEmpty:
+ break
+ self._next_deadline_s = None
+
async def recv(self) -> VideoFrame:
if self._closed:
raise MediaStreamError
@@ -234,6 +255,17 @@ def enqueue_encoded_packet_nowait(self, packet: Packet) -> bool:
self._packets.put_nowait(packet)
return True
+ async def flush(self) -> None:
+ """Drop queued encoded packets while preserving the open media track."""
+ if self._closed:
+ return
+ while True:
+ try:
+ self._packets.get_nowait()
+ except asyncio.QueueEmpty:
+ break
+ self._next_deadline_s = None
+
async def recv(self) -> Packet:
if self._closed:
raise MediaStreamError
diff --git a/flashdreams/flashdreams/serving/webrtc/nvenc.py b/flashdreams/flashdreams/serving/webrtc/nvenc.py
index 3a4cc93dc..94a52200f 100644
--- a/flashdreams/flashdreams/serving/webrtc/nvenc.py
+++ b/flashdreams/flashdreams/serving/webrtc/nvenc.py
@@ -23,6 +23,7 @@
import contextlib
import time
from collections.abc import Callable
+from dataclasses import dataclass
from fractions import Fraction
from typing import TYPE_CHECKING, Any
@@ -30,6 +31,7 @@
from aiortc import MediaStreamTrack
from av.packet import Packet
from loguru import logger
+from torch import Tensor
from flashdreams.runtime import StepResult
from flashdreams.serving.webrtc.encoders import ChunkDeliveryResult
@@ -57,6 +59,13 @@
_RTP_VIDEO_CLOCK = 90_000
+@dataclass(frozen=True, slots=True)
+class NVENCChunkPayload:
+ """Encoder-owned CUDA frames prepared before async delivery is scheduled."""
+
+ frames: Tensor
+
+
def _payload_contains_nal_type(payload: bytes, nal_type: int) -> bool:
"""Scan an Annex-B H.264 payload for the presence of a specific NAL type."""
i = 0
@@ -240,6 +249,35 @@ def create_track(self, *, maxsize: int) -> NVENCVideoTrack:
return NVENCVideoTrack(fps=self.fps, maxsize=maxsize)
+ def prepare_chunk_payload(
+ self,
+ result: StepResult,
+ track: MediaStreamTrack,
+ ) -> NVENCChunkPayload:
+ from flashdreams.serving.webrtc.media import NVENCVideoTrack
+
+ if not isinstance(track, NVENCVideoTrack):
+ raise TypeError(
+ "PyNvHardwareEncoder requires an NVENCVideoTrack; got "
+ f"{type(track).__name__}. Create it via encoder.create_track()."
+ )
+ return NVENCChunkPayload(frames=_result_to_abgr_frames(result))
+
+ async def deliver_prepared_chunk(
+ self,
+ payload: object,
+ track: MediaStreamTrack,
+ *,
+ force_keyframe: bool = False,
+ ) -> ChunkDeliveryResult:
+ if not isinstance(payload, NVENCChunkPayload):
+ raise TypeError("PyNvHardwareEncoder payload must be an NVENCChunkPayload.")
+ return await self._deliver_prepared_frames(
+ payload.frames,
+ track,
+ force_keyframe=force_keyframe,
+ )
+
async def deliver_chunk(
self,
result: StepResult,
@@ -297,6 +335,63 @@ def _stream(packet: Packet) -> None:
encode_ms=encode_ms,
)
+ async def _deliver_prepared_frames(
+ self,
+ frames: Tensor,
+ track: MediaStreamTrack,
+ *,
+ force_keyframe: bool = False,
+ ) -> ChunkDeliveryResult:
+ from flashdreams.serving.webrtc.media import NVENCVideoTrack
+
+ if not isinstance(track, NVENCVideoTrack):
+ raise TypeError(
+ "PyNvHardwareEncoder requires an NVENCVideoTrack; got "
+ f"{type(track).__name__}. Create it via encoder.create_track()."
+ )
+ loop = asyncio.get_running_loop()
+ emitted = 0
+ enqueued = 0
+
+ def _stream(packet: Packet) -> None:
+ nonlocal emitted, enqueued
+ emitted += 1
+ enqueue = track.enqueue_encoded_packet(packet)
+ try:
+ future = asyncio.run_coroutine_threadsafe(
+ enqueue,
+ loop,
+ )
+ except RuntimeError:
+ enqueue.close()
+ return
+ try:
+ accepted = future.result()
+ except Exception:
+ return
+ if accepted:
+ enqueued += 1
+
+ _num_frames, num_keyframes, encode_ms = await asyncio.to_thread(
+ self.encode_frames_sync,
+ frames,
+ force_keyframe=force_keyframe,
+ on_packet=_stream,
+ )
+ if enqueued < emitted:
+ logger.debug(
+ "NVENC track closed while enqueueing encoded chunk; "
+ "enqueued {} of {} packet(s).",
+ enqueued,
+ emitted,
+ )
+ return ChunkDeliveryResult(
+ backend=self.backend,
+ num_frames=enqueued,
+ num_keyframes=num_keyframes,
+ encode_ms=encode_ms,
+ )
+
def encode_chunk_sync(
self,
result: StepResult,
@@ -311,6 +406,20 @@ def encode_chunk_sync(
just to get access to the emitted packets.
"""
frames = _result_to_abgr_frames(result)
+ return self.encode_frames_sync(
+ frames,
+ force_keyframe=force_keyframe,
+ on_packet=on_packet,
+ )
+
+ def encode_frames_sync(
+ self,
+ frames: Tensor,
+ *,
+ force_keyframe: bool = False,
+ on_packet: Callable[[Packet], None] | None = None,
+ ) -> tuple[int, int, float]:
+ """Encode preconverted ``ABGR`` frames for prepared async delivery."""
if not frames.is_cuda:
raise ValueError("expected CUDA tensor for hardware encode path")
num_frames = frames.shape[0]
diff --git a/flashdreams/flashdreams/serving/webrtc/runtime.py b/flashdreams/flashdreams/serving/webrtc/runtime.py
index 70c16fb38..32c7d4365 100644
--- a/flashdreams/flashdreams/serving/webrtc/runtime.py
+++ b/flashdreams/flashdreams/serving/webrtc/runtime.py
@@ -37,6 +37,7 @@ class WebRTCControlSignal(IntEnum):
CLOSE = 3
EVENT = 4
SESSION_STEP = 5
+ SESSION_CLOSE = 6
EXIT = 99
diff --git a/flashdreams/flashdreams/serving/webrtc/server.py b/flashdreams/flashdreams/serving/webrtc/server.py
index bf257701d..7a2c6694d 100644
--- a/flashdreams/flashdreams/serving/webrtc/server.py
+++ b/flashdreams/flashdreams/serving/webrtc/server.py
@@ -91,10 +91,12 @@ async def healthz(request: web.Request) -> web.StreamResponse:
)
async def ui_config(_: web.Request) -> web.StreamResponse:
- adapter_module = None
+ payload: dict[str, str | None] = {"adapter_module": None}
if model_web_dir is not None and (model_web_dir / "adapter.js").is_file():
- adapter_module = "/model-static/adapter.js?v=model-ui-v1"
- return web.json_response({"adapter_module": adapter_module})
+ payload["adapter_module"] = "/model-static/adapter.js?v=model-ui-v2"
+ if model_web_dir is not None and (model_web_dir / "adapter.css").is_file():
+ payload["model_stylesheet"] = "/model-static/adapter.css?v=model-ui-v2"
+ return web.json_response(payload)
async def on_startup(app: web.Application) -> None:
manager = app[SESSION_MANAGER_KEY]
diff --git a/flashdreams/flashdreams/serving/webrtc/services.py b/flashdreams/flashdreams/serving/webrtc/services.py
new file mode 100644
index 000000000..3ce38e4be
--- /dev/null
+++ b/flashdreams/flashdreams/serving/webrtc/services.py
@@ -0,0 +1,1161 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Session-edge services for the shared WebRTC demo run mode.
+
+These classes are the Phase 12 decomposition layer: they translate WebRTC
+transport facts into the shared demo runtime contracts without owning model
+execution. The production manager still uses its legacy execution hook until
+the realtime-driver adoption phase.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import inspect
+import json
+import math
+import threading
+from collections import deque
+from collections.abc import Callable, Coroutine, Mapping, MutableSet, Sequence
+from concurrent.futures import Future
+from dataclasses import dataclass, field
+from typing import Any, Literal, Protocol, runtime_checkable
+
+from flashdreams.runtime import (
+ StepRequirements,
+ StepResult,
+ UserInputCapability,
+ UserInputEvent,
+ UserInputs,
+ UserInputSchema,
+)
+from flashdreams.runtime._utils import freeze_mapping
+from flashdreams.runtime.demo import (
+ AsyncSessionDriver,
+ DemoAdapter,
+ DemoSpec,
+ InMemorySessionMetricsRecorder,
+ ModelInputProvider,
+ ModelWarmupPlan,
+ OutputDecision,
+ PreparedScenario,
+ RealtimeSessionDriver,
+ RealtimeWindowResult,
+ RunContext,
+ RunModeCapabilities,
+ RunResult,
+ RuntimeHost,
+ SessionEdges,
+ SessionInfo,
+ SingleSessionAdmissionPolicy,
+ StepPipeline,
+ UserInputWindow,
+ WebRTCErrorPolicy,
+ input_frame_count_from_request,
+ run_demo_session_async,
+)
+from flashdreams.runtime.demo.timing import (
+ SPARSE_KEY_SEGMENTS_METADATA_KEY,
+ ActivationResult,
+ CatchUpPolicy,
+ DeterministicClock,
+ RealtimeClock,
+)
+
+from .messages import (
+ MESSAGE_TYPE_ACTION,
+ MESSAGE_TYPE_DISCONNECT,
+ MESSAGE_TYPE_EVENT,
+ MESSAGE_TYPE_HEARTBEAT,
+)
+from .server import SessionBusyError
+
+WebRTCMessageKind = Literal[
+ "action",
+ "disconnect",
+ "event",
+ "heartbeat",
+ "error",
+]
+
+WebRTCDropPolicy = Literal["none", "drop_newest", "drop_oldest"]
+
+_CLEAR_EVENT_STATES = frozenset({"clear", "release", "off", "none"})
+WEBRTC_SKIPPED_INPUTS_METADATA_KEY = "webrtc_skipped_inputs"
+WEBRTC_SKIPPED_WINDOW_METADATA_KEY = "webrtc_skipped_window"
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class WebRTCMessageResult:
+ """Result of translating one browser data-channel message."""
+
+ kind: WebRTCMessageKind
+ activated: bool = False
+ error: str | None = None
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class WebRTCOfferRequest:
+ """Browser SDP offer passed to the shared offer/session handler."""
+
+ sdp: str
+ type: str
+
+ def __post_init__(self) -> None:
+ if not self.sdp.strip():
+ raise ValueError("WebRTCOfferRequest.sdp must be non-empty.")
+ if not self.type.strip():
+ raise ValueError("WebRTCOfferRequest.type must be non-empty.")
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class WebRTCOutputBridgeDecision:
+ """Immediate delivery decision from a nonblocking WebRTC output bridge."""
+
+ accepted: bool = True
+ should_stop: bool = False
+ dropped: bool = False
+ drop_policy: WebRTCDropPolicy = "none"
+ backpressure_s: float = 0.0
+ metadata: Mapping[str, object] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ if self.drop_policy not in {"none", "drop_newest", "drop_oldest"}:
+ raise ValueError(f"Unsupported drop_policy={self.drop_policy!r}.")
+ if not math.isfinite(self.backpressure_s) or self.backpressure_s < 0.0:
+ raise ValueError(
+ "WebRTCOutputBridgeDecision.backpressure_s must be finite and >= 0."
+ )
+ object.__setattr__(self, "metadata", freeze_mapping(self.metadata))
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class WebRTCChunkDelivery:
+ """Completed WebRTC chunk delivery plus the model chunk summary."""
+
+ delivery: object
+ step_index: int
+ frame_count: int
+ generation: int
+ force_keyframe: bool
+ metadata: Mapping[str, object] = field(default_factory=dict)
+ metrics: Mapping[str, float | int] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ object.__setattr__(self, "metadata", freeze_mapping(self.metadata))
+ object.__setattr__(self, "metrics", freeze_mapping(self.metrics))
+
+
+@runtime_checkable
+class WebRTCOfferAnswerer(Protocol):
+ """Creates an SDP answer after the shared session task has been scheduled."""
+
+ async def create_answer(
+ self,
+ *,
+ offer: WebRTCOfferRequest,
+ session_task: asyncio.Task[RunResult],
+ ) -> Mapping[str, str]: ...
+
+
+@runtime_checkable
+class BlockingPreparationService(Protocol):
+ """Runs blocking scenario preparation outside the aiohttp event loop."""
+
+ async def run(
+ self,
+ func: Callable[..., object],
+ *args: object,
+ **kwargs: object,
+ ) -> object: ...
+
+
+@runtime_checkable
+class WebRTCOutputBridge(Protocol):
+ """Thread-safe bridge from model-worker output writes to WebRTC delivery."""
+
+ def begin_generation(self, generation: int) -> None: ...
+
+ def submit_chunk(
+ self,
+ result: StepResult,
+ *,
+ generation: int,
+ force_keyframe: bool = False,
+ ) -> WebRTCOutputBridgeDecision: ...
+
+ def close(self) -> None: ...
+
+
+@runtime_checkable
+class WebRTCSessionEdgeFactory(Protocol):
+ """Builds per-peer shared session edges on the WebRTC control rank."""
+
+ def create_session_edges(
+ self,
+ *,
+ context: RunContext,
+ spec: DemoSpec,
+ scenario: PreparedScenario,
+ provider: ModelInputProvider,
+ adapter: DemoAdapter,
+ ) -> SessionEdges: ...
+
+
+class AsyncioBlockingPreparationService:
+ """Default blocking-prep service backed by ``asyncio.to_thread``."""
+
+ async def run(
+ self,
+ func: Callable[..., object],
+ *args: object,
+ **kwargs: object,
+ ) -> object:
+ return await asyncio.to_thread(func, *args, **kwargs)
+
+
+class WebRTCTransportService:
+ """Idempotent per-peer transport lifecycle for realtime session edges."""
+
+ def __init__(
+ self,
+ *,
+ loop: asyncio.AbstractEventLoop | None = None,
+ on_close: Callable[[str | None], None] | None = None,
+ ) -> None:
+ self._loop = loop
+ self._on_close = on_close
+ self._closed_signal = _ThreadSafeActivationSignal(loop=loop)
+ self._lock = threading.Lock()
+ self._closed = False
+ self._close_reason: str | None = None
+ self._close_count = 0
+ self.last_client_message_at: float | None = None
+
+ @property
+ def close_count(self) -> int:
+ """Number of effective transport closes, after idempotency."""
+
+ return self._close_count
+
+ @property
+ def close_reason(self) -> str | None:
+ return self._close_reason
+
+ @property
+ def closed_signal(self) -> "_ThreadSafeActivationSignal":
+ return self._closed_signal
+
+ def mark_client_message(self, timestamp_s: float) -> None:
+ if not math.isfinite(timestamp_s) or timestamp_s < 0.0:
+ raise ValueError("timestamp_s must be finite and >= 0.")
+ self.last_client_message_at = float(timestamp_s)
+
+ def is_active(self) -> bool:
+ return not self._closed
+
+ def disconnect(self, reason: str = "client disconnected") -> None:
+ self.close(reason=reason)
+
+ def close(self, reason: str | None = None) -> None:
+ callback: Callable[[str | None], None] | None = None
+ with self._lock:
+ if self._closed:
+ return
+ self._closed = True
+ self._close_reason = reason
+ self._close_count += 1
+ callback = self._on_close
+ self._closed_signal.set()
+ if callback is not None:
+ callback(reason)
+
+
+@dataclass(slots=True)
+class WebRTCActivationPolicy:
+ """Activate on the first browser action/event or stop on disconnect."""
+
+ input_source: "WebRTCInputSource"
+ transport: WebRTCTransportService
+ timeout_s: float | None = None
+ timeout_reason: str = "activation timed out"
+ anchor_clock: bool = True
+
+ def __post_init__(self) -> None:
+ if self.timeout_s is not None and self.timeout_s <= 0.0:
+ raise ValueError("timeout_s must be > 0 when set.")
+ if not self.timeout_reason.strip():
+ raise ValueError("timeout_reason must be non-empty.")
+
+ async def wait_until_active(
+ self,
+ clock: RealtimeClock | DeterministicClock,
+ ) -> ActivationResult:
+ if self.input_source.activation_signal.is_set():
+ self._anchor(clock)
+ return ActivationResult(activated=True)
+ if not self.transport.is_active():
+ return ActivationResult(
+ activated=False,
+ reason=self.transport.close_reason or "transport closed",
+ )
+
+ activation_task = asyncio.create_task(
+ self.input_source.activation_signal.wait()
+ )
+ closed_task = asyncio.create_task(self.transport.closed_signal.wait())
+ try:
+ done, pending = await asyncio.wait(
+ {activation_task, closed_task},
+ timeout=self.timeout_s,
+ return_when=asyncio.FIRST_COMPLETED,
+ )
+ if not done:
+ return ActivationResult(
+ activated=False,
+ reason=self.timeout_reason,
+ )
+ for task in done:
+ task.result()
+ if not self.transport.is_active():
+ return ActivationResult(
+ activated=False,
+ reason=self.transport.close_reason or "transport closed",
+ )
+ self._anchor(clock)
+ return ActivationResult(activated=True)
+ finally:
+ for task in (activation_task, closed_task):
+ if not task.done():
+ task.cancel()
+ await asyncio.gather(
+ activation_task,
+ closed_task,
+ return_exceptions=True,
+ )
+
+ def _anchor(self, clock: RealtimeClock | DeterministicClock) -> None:
+ if not self.anchor_clock or not clock.is_realtime:
+ return
+ now = getattr(clock, "now", None)
+ anchor = getattr(clock, "anchor", None)
+ if callable(now) and callable(anchor):
+ anchor(now())
+
+
+@dataclass(slots=True)
+class WebRTCInputSource:
+ """Realtime source fed by browser data-channel events."""
+
+ resampler: Any
+ max_lag_s: float | None = None
+ catch_up_policy: CatchUpPolicy = "fold"
+ user_input_schema: UserInputSchema = field(
+ default_factory=lambda: WEBRTC_USER_INPUT_SCHEMA
+ )
+ is_finite: bool = False
+ is_deterministic: bool = False
+ _activation_signal: "_ThreadSafeActivationSignal" = field(
+ default_factory=lambda: _ThreadSafeActivationSignal(),
+ init=False,
+ repr=False,
+ )
+ _events: deque[UserInputEvent] = field(
+ default_factory=deque,
+ init=False,
+ repr=False,
+ )
+
+ def __post_init__(self) -> None:
+ if self.max_lag_s is not None and (
+ not math.isfinite(self.max_lag_s) or self.max_lag_s < 0.0
+ ):
+ raise ValueError("max_lag_s must be finite and >= 0.")
+ if self.catch_up_policy != "fold":
+ raise NotImplementedError(
+ f"Catch-up policy {self.catch_up_policy!r} has no WebRTC analog yet."
+ )
+
+ @property
+ def activation_signal(self) -> "_ThreadSafeActivationSignal":
+ return self._activation_signal
+
+ def is_finished(self) -> bool:
+ return False
+
+ def reset(self, *, start_v: float) -> None:
+ self.resampler.reset(start_v=start_v)
+ self._events.clear()
+ self._activation_signal.clear()
+
+ def handle_browser_message(
+ self,
+ raw_message: object,
+ *,
+ timestamp_s: float,
+ ) -> WebRTCMessageResult:
+ """Translate one browser data-channel message into typed user inputs."""
+
+ if not isinstance(raw_message, str):
+ return WebRTCMessageResult(kind="error", error="Expected text payload.")
+ try:
+ payload = json.loads(raw_message)
+ except json.JSONDecodeError:
+ return WebRTCMessageResult(kind="error", error="Invalid JSON payload.")
+ if not isinstance(payload, dict):
+ return WebRTCMessageResult(
+ kind="error",
+ error="Payload must be a JSON object.",
+ )
+ return self.handle_browser_payload(payload, timestamp_s=timestamp_s)
+
+ def handle_browser_payload(
+ self,
+ payload: Mapping[str, object],
+ *,
+ timestamp_s: float,
+ ) -> WebRTCMessageResult:
+ message_type = str(payload.get("type", "")).strip().lower()
+ if message_type == MESSAGE_TYPE_HEARTBEAT:
+ return WebRTCMessageResult(kind="heartbeat")
+ if message_type == MESSAGE_TYPE_DISCONNECT:
+ return WebRTCMessageResult(kind="disconnect")
+ if message_type == MESSAGE_TYPE_EVENT:
+ return self._record_text_event(payload, timestamp_s=timestamp_s)
+ if message_type == MESSAGE_TYPE_ACTION:
+ action_payload = payload.get("action", payload)
+ if not isinstance(action_payload, Mapping):
+ return WebRTCMessageResult(
+ kind="error",
+ error="'action' must be an object.",
+ )
+ return self._record_action(
+ {str(key): value for key, value in action_payload.items()},
+ timestamp_s=timestamp_s,
+ )
+ return WebRTCMessageResult(
+ kind="error",
+ error=(
+ "Unsupported message type, expected "
+ "'action', 'event', 'heartbeat', or 'disconnect'."
+ ),
+ )
+
+ def record_user_event(
+ self,
+ *,
+ timestamp_s: float,
+ event_type: str,
+ payload: Mapping[str, object],
+ source_event_id: str | None = None,
+ activate: bool = True,
+ ) -> None:
+ event = UserInputEvent(
+ timestamp_s=timestamp_s,
+ event_type=event_type,
+ payload=dict(payload),
+ source="webrtc",
+ source_event_id=source_event_id,
+ )
+ self.user_input_schema.validate_event(event)
+ self._events.append(event)
+ if activate:
+ self._activation_signal.set()
+
+ async def next_realtime_window(
+ self,
+ *,
+ request: StepRequirements,
+ clock: RealtimeClock,
+ ) -> RealtimeWindowResult:
+ input_frame_count = input_frame_count_from_request(request)
+ chunk_duration_s = input_frame_count * float(self.resampler.dt)
+ if chunk_duration_s <= 0.0:
+ raise ValueError("Realtime resampler dt must produce a positive window.")
+
+ window_end_s = float(self.resampler.next_chunk_start_v) + chunk_duration_s
+ await clock.wait_until_window_end(window_end_s)
+ pre_catch_up_start_s = float(self.resampler.next_chunk_start_v)
+ catch_up = clock.catch_up(
+ request=request,
+ max_lag_s=self.max_lag_s
+ if self.max_lag_s is not None
+ else chunk_duration_s,
+ policy=self.catch_up_policy,
+ )
+ start_s = float(self.resampler.next_chunk_start_v)
+ segments, frame_times = self.resampler.sample_chunk(input_frame_count)
+ end_s = float(self.resampler.next_chunk_start_v)
+ metadata: dict[str, object] = {
+ SPARSE_KEY_SEGMENTS_METADATA_KEY: tuple(segments),
+ }
+ if start_s > pre_catch_up_start_s:
+ metadata[WEBRTC_SKIPPED_INPUTS_METADATA_KEY] = UserInputs(
+ events=self._events_for_window(pre_catch_up_start_s, start_s)
+ )
+ metadata[WEBRTC_SKIPPED_WINDOW_METADATA_KEY] = (
+ pre_catch_up_start_s,
+ start_s,
+ )
+ window = RealtimeWindowResult(
+ window=_user_input_window(
+ start_s=start_s,
+ end_s=end_s,
+ frame_times=tuple(frame_times),
+ inputs=UserInputs(events=self._events_for_window(start_s, end_s)),
+ metadata=metadata,
+ ),
+ catch_up=catch_up,
+ )
+ self._prune_events(before_s=start_s)
+ return window
+
+ def _record_action(
+ self,
+ payload: Mapping[str, object],
+ *,
+ timestamp_s: float,
+ ) -> WebRTCMessageResult:
+ event = str(payload.get("event", "")).strip().lower()
+ if event == "step":
+ self._activation_signal.set()
+ return WebRTCMessageResult(kind="action", activated=True)
+ if event not in {"keydown", "keyup"}:
+ return WebRTCMessageResult(
+ kind="error",
+ error=f"Unsupported event={event!r}; expected 'keydown' or 'keyup'.",
+ )
+ key = str(payload.get("key", "")).strip()
+ if not key:
+ return WebRTCMessageResult(
+ kind="error",
+ error="Action payload must include non-empty 'key'.",
+ )
+ self.resampler.on_edge(arrival_t=timestamp_s, event=event, key=key)
+ self.record_user_event(
+ timestamp_s=timestamp_s,
+ event_type="key_down" if event == "keydown" else "key_up",
+ payload={"key": key},
+ )
+ return WebRTCMessageResult(kind="action", activated=True)
+
+ def _record_text_event(
+ self,
+ payload: Mapping[str, object],
+ *,
+ timestamp_s: float,
+ ) -> WebRTCMessageResult:
+ state = str(payload.get("state", "trigger")).strip().lower() or "trigger"
+ event_id = str(payload.get("event_id", payload.get("id", ""))).strip()
+ clears = state in _CLEAR_EVENT_STATES
+ if not event_id and not clears:
+ return WebRTCMessageResult(
+ kind="error",
+ error=(
+ "Event payload must include non-empty 'event_id' unless state "
+ "clears the active event."
+ ),
+ )
+ active_event_id = None if clears else event_id
+ self.record_user_event(
+ timestamp_s=timestamp_s,
+ event_type="text_event",
+ payload={"event_id": active_event_id, "state": state},
+ source_event_id=active_event_id,
+ )
+ return WebRTCMessageResult(kind="event", activated=True)
+
+ def _events_for_window(
+ self,
+ start_s: float,
+ end_s: float,
+ ) -> tuple[UserInputEvent, ...]:
+ return tuple(
+ sorted(
+ (
+ event
+ for event in self._events
+ if start_s <= event.timestamp_s < end_s
+ ),
+ key=lambda event: event.timestamp_s,
+ )
+ )
+
+ def _prune_events(self, *, before_s: float) -> None:
+ self._events = deque(
+ event for event in self._events if event.timestamp_s >= before_s
+ )
+
+
+class WebRTCOutputSink:
+ """Output sink that schedules WebRTC media delivery without blocking."""
+
+ produces_artifacts = False
+
+ def __init__(self, *, bridge: WebRTCOutputBridge) -> None:
+ self._bridge = bridge
+ self._opened = False
+ self._closed = True
+ self._bridge_closed = False
+ self._generation = 0
+ self._force_keyframe = False
+ self.session_info: SessionInfo | None = None
+
+ def open(self, session_info: SessionInfo) -> None:
+ self.session_info = session_info
+ self._opened = True
+ self._closed = False
+ self._generation = 0
+ self._force_keyframe = True
+ self._bridge.begin_generation(0)
+
+ def begin_generation(self, generation: int) -> None:
+ if generation < 0:
+ raise ValueError("generation must be >= 0.")
+ self._generation = generation
+ self._force_keyframe = True
+ self._bridge.begin_generation(generation)
+
+ def write(self, result: StepResult) -> OutputDecision:
+ if not self._opened or self._closed:
+ raise RuntimeError("Cannot write to a closed output sink.")
+ decision = self._bridge.submit_chunk(
+ result,
+ generation=self._generation,
+ force_keyframe=self._force_keyframe,
+ )
+ self._force_keyframe = False
+ return OutputDecision(
+ should_stop=decision.should_stop,
+ dropped=decision.dropped,
+ drop_policy=decision.drop_policy,
+ backpressure_s=decision.backpressure_s,
+ metadata=decision.metadata,
+ )
+
+ def close(self) -> Sequence[Any]:
+ if self._bridge_closed:
+ return ()
+ self._closed = True
+ self._opened = False
+ self._bridge.close()
+ self._bridge_closed = True
+ return ()
+
+
+class ThreadSafeWebRTCOutputBridge:
+ """Schedule async encoder delivery from any thread without blocking writes."""
+
+ def __init__(
+ self,
+ *,
+ loop: asyncio.AbstractEventLoop,
+ video_encoder: Any,
+ video_track: Any,
+ max_pending_chunks: int = 2,
+ close_track: bool = True,
+ on_delivery: Callable[[object], None] | None = None,
+ on_chunk_delivery: Callable[[WebRTCChunkDelivery], None] | None = None,
+ on_error: Callable[[BaseException], None] | None = None,
+ ) -> None:
+ if max_pending_chunks <= 0:
+ raise ValueError("max_pending_chunks must be > 0.")
+ self._loop = loop
+ self._video_encoder = video_encoder
+ self._video_track = video_track
+ self._max_pending_chunks = max_pending_chunks
+ self._close_track = close_track
+ self._on_delivery = on_delivery
+ self._on_chunk_delivery = on_chunk_delivery
+ self._on_error = on_error
+ self._pending: dict[Future[WebRTCChunkDelivery], int] = {}
+ self._lock = threading.Lock()
+ self._closed = False
+ self._generation = 0
+
+ @property
+ def pending_count(self) -> int:
+ with self._lock:
+ return len(self._pending)
+
+ def begin_generation(self, generation: int) -> None:
+ if generation < 0:
+ raise ValueError("generation must be >= 0.")
+ with self._lock:
+ if self._closed or generation <= self._generation:
+ return
+ self._generation = generation
+ stale = tuple(
+ future
+ for future, future_generation in self._pending.items()
+ if future_generation < generation
+ )
+ for future in stale:
+ future.cancel()
+ self._schedule_track_flush()
+
+ def submit_chunk(
+ self,
+ result: StepResult,
+ *,
+ generation: int,
+ force_keyframe: bool = False,
+ ) -> WebRTCOutputBridgeDecision:
+ prepare = getattr(self._video_encoder, "prepare_chunk_payload", None)
+ deliver = getattr(self._video_encoder, "deliver_prepared_chunk", None)
+ if not callable(prepare) or not callable(deliver):
+ raise TypeError(
+ "ThreadSafeWebRTCOutputBridge requires a video encoder with "
+ "prepare_chunk_payload(...) and deliver_prepared_chunk(...)."
+ )
+ with self._lock:
+ if self._closed:
+ return WebRTCOutputBridgeDecision(
+ accepted=False,
+ should_stop=True,
+ dropped=True,
+ drop_policy="drop_newest",
+ metadata={"reason": "closed"},
+ )
+ if generation < self._generation:
+ return WebRTCOutputBridgeDecision(
+ accepted=False,
+ dropped=True,
+ drop_policy="drop_newest",
+ metadata={"reason": "stale generation"},
+ )
+ if len(self._pending) >= self._max_pending_chunks:
+ return WebRTCOutputBridgeDecision(
+ accepted=False,
+ dropped=True,
+ drop_policy="drop_newest",
+ metadata={"reason": "pending queue full"},
+ )
+ payload = prepare(result, self._video_track)
+ chunk = WebRTCChunkDelivery(
+ delivery=None,
+ step_index=result.step_index,
+ frame_count=result.frame_count,
+ generation=generation,
+ force_keyframe=force_keyframe,
+ metadata=result.metadata,
+ metrics=result.metrics,
+ )
+ with self._lock:
+ if self._closed:
+ return WebRTCOutputBridgeDecision(
+ accepted=False,
+ should_stop=True,
+ dropped=True,
+ drop_policy="drop_newest",
+ metadata={"reason": "closed"},
+ )
+ if generation < self._generation:
+ return WebRTCOutputBridgeDecision(
+ accepted=False,
+ dropped=True,
+ drop_policy="drop_newest",
+ metadata={"reason": "stale generation"},
+ )
+ if len(self._pending) >= self._max_pending_chunks:
+ return WebRTCOutputBridgeDecision(
+ accepted=False,
+ dropped=True,
+ drop_policy="drop_newest",
+ metadata={"reason": "pending queue full"},
+ )
+ future = asyncio.run_coroutine_threadsafe(
+ self._deliver(
+ payload,
+ chunk=chunk,
+ generation=generation,
+ force_keyframe=force_keyframe,
+ ),
+ self._loop,
+ )
+ self._pending[future] = generation
+ future.add_done_callback(self._on_done)
+
+ return WebRTCOutputBridgeDecision(
+ accepted=True,
+ backpressure_s=self._track_backpressure_s(),
+ )
+
+ def close(self) -> None:
+ with self._lock:
+ if self._closed:
+ return
+ self._closed = True
+ pending = tuple(self._pending)
+ for future in pending:
+ future.cancel()
+ if self._close_track:
+ self._schedule_track_close()
+
+ async def _deliver(
+ self,
+ payload: object,
+ *,
+ chunk: WebRTCChunkDelivery,
+ generation: int,
+ force_keyframe: bool,
+ ) -> WebRTCChunkDelivery:
+ with self._lock:
+ if self._closed or generation < self._generation:
+ raise asyncio.CancelledError
+ delivery = await self._video_encoder.deliver_prepared_chunk(
+ payload,
+ self._video_track,
+ force_keyframe=force_keyframe,
+ )
+ with self._lock:
+ if self._closed or generation < self._generation:
+ stale_after_delivery = True
+ else:
+ stale_after_delivery = False
+ if stale_after_delivery:
+ self._schedule_track_flush()
+ raise asyncio.CancelledError
+ return WebRTCChunkDelivery(
+ delivery=delivery,
+ step_index=chunk.step_index,
+ frame_count=chunk.frame_count,
+ generation=chunk.generation,
+ force_keyframe=chunk.force_keyframe,
+ metadata=chunk.metadata,
+ metrics=chunk.metrics,
+ )
+
+ def _on_done(self, future: Future[WebRTCChunkDelivery]) -> None:
+ with self._lock:
+ self._pending.pop(future, None)
+ if future.cancelled():
+ return
+ try:
+ result = future.result()
+ except BaseException as exc:
+ if self._on_error is not None:
+ self._on_error(exc)
+ return
+ if self._on_delivery is not None:
+ self._on_delivery(result.delivery)
+ if self._on_chunk_delivery is not None:
+ self._on_chunk_delivery(result)
+
+ def _track_backpressure_s(self) -> float:
+ qsize = getattr(self._video_track, "qsize", None)
+ fps = getattr(self._video_track, "fps", None) or getattr(
+ self._video_encoder,
+ "fps",
+ None,
+ )
+ if not callable(qsize) or fps is None:
+ return 0.0
+ try:
+ queue_depth = int(qsize())
+ frames_per_second = float(fps)
+ except (TypeError, ValueError):
+ return 0.0
+ if frames_per_second <= 0.0:
+ return 0.0
+ return max(0.0, queue_depth / frames_per_second)
+
+ def _schedule_track_close(self) -> None:
+ close = getattr(self._video_track, "close", None)
+ if not callable(close):
+ return
+ try:
+ result = close()
+ if inspect.isawaitable(result):
+ asyncio.run_coroutine_threadsafe(result, self._loop)
+ except BaseException as exc:
+ if self._on_error is not None:
+ self._on_error(exc)
+
+ def _schedule_track_flush(self) -> None:
+ flush = getattr(self._video_track, "flush", None)
+ if not callable(flush):
+ return
+ try:
+ result = flush()
+ if inspect.isawaitable(result):
+ asyncio.run_coroutine_threadsafe(result, self._loop)
+ except BaseException as exc:
+ if self._on_error is not None:
+ self._on_error(exc)
+
+
+class WebRTCRunMode:
+ """Shared realtime run mode that delegates peer-specific edges to WebRTC."""
+
+ name = "webrtc"
+ capabilities = RunModeCapabilities(
+ realtime=True,
+ supports_backpressure=True,
+ supports_interactive_events=True,
+ )
+
+ def __init__(
+ self,
+ *,
+ edge_factory: WebRTCSessionEdgeFactory,
+ blocking_preparation: BlockingPreparationService | None = None,
+ driver: AsyncSessionDriver | None = None,
+ error_policy: WebRTCErrorPolicy | None = None,
+ ) -> None:
+ self._edge_factory = edge_factory
+ self._blocking_preparation = (
+ blocking_preparation or AsyncioBlockingPreparationService()
+ )
+ self._driver = driver or RealtimeSessionDriver()
+ self._error_policy = error_policy or WebRTCErrorPolicy()
+
+ @property
+ def blocking_preparation(self) -> BlockingPreparationService:
+ return self._blocking_preparation
+
+ @property
+ def error_policy(self) -> WebRTCErrorPolicy:
+ return self._error_policy
+
+ def validate_run(self, *, spec: DemoSpec, adapter: DemoAdapter) -> None:
+ del adapter
+ if spec.output.mode != "webrtc":
+ raise ValueError("WebRTCRunMode requires WebRTC output.")
+
+ def validate_session(
+ self,
+ *,
+ spec: DemoSpec,
+ scenario: PreparedScenario,
+ adapter: DemoAdapter,
+ provider: ModelInputProvider,
+ ) -> None:
+ del spec, scenario, adapter
+ if not provider.capabilities.supports_realtime_clock:
+ raise ValueError("WebRTC providers must support realtime clocks.")
+
+ def create_run_context(
+ self,
+ *,
+ spec: DemoSpec,
+ adapter: DemoAdapter,
+ host: RuntimeHost,
+ model_warmup_plan: ModelWarmupPlan,
+ ) -> RunContext:
+ del spec, adapter
+ services: dict[str, object] = {}
+ if host.is_control_rank:
+ services["blocking_preparation"] = self._blocking_preparation
+ return RunContext(
+ host=host,
+ run_metrics=InMemorySessionMetricsRecorder(),
+ admission=SingleSessionAdmissionPolicy(
+ health_check=lambda: host.is_control_rank and host.is_healthy
+ ),
+ model_warmup_plan=model_warmup_plan,
+ services=services,
+ )
+
+ def create_session_edges(
+ self,
+ *,
+ context: RunContext,
+ spec: DemoSpec,
+ scenario: PreparedScenario,
+ provider: ModelInputProvider,
+ adapter: DemoAdapter,
+ ) -> SessionEdges:
+ if not context.host.is_control_rank:
+ raise RuntimeError("WebRTC session edges are control-rank only.")
+ edges = self._edge_factory.create_session_edges(
+ context=context,
+ spec=spec,
+ scenario=scenario,
+ provider=provider,
+ adapter=adapter,
+ )
+ if not isinstance(edges, SessionEdges):
+ raise TypeError(
+ "WebRTC edge factory must return SessionEdges, "
+ f"got {type(edges).__name__}."
+ )
+ return edges
+
+ def select_driver(self) -> AsyncSessionDriver:
+ return self._driver
+
+
+class WebRTCSessionOfferHandler:
+ """Reserve, prepare, and launch one WebRTC session before SDP negotiation."""
+
+ def __init__(
+ self,
+ *,
+ context: RunContext,
+ spec: DemoSpec,
+ adapter: DemoAdapter,
+ run_mode: WebRTCRunMode,
+ answerer: WebRTCOfferAnswerer,
+ pipeline: StepPipeline | None = None,
+ session_helper: Callable[..., Coroutine[Any, Any, RunResult]] | None = None,
+ busy_message: str = "Another WebRTC session is already active.",
+ session_tasks: MutableSet[asyncio.Task[RunResult]] | None = None,
+ ) -> None:
+ self._context = context
+ self._spec = spec
+ self._adapter = adapter
+ self._run_mode = run_mode
+ self._answerer = answerer
+ self._pipeline = pipeline or StepPipeline()
+ self._session_helper = session_helper or run_demo_session_async
+ self._busy_message = busy_message
+ self._session_tasks = session_tasks if session_tasks is not None else set()
+
+ async def handle_offer(
+ self,
+ *,
+ offer_sdp: str,
+ offer_type: str,
+ ) -> Mapping[str, str]:
+ if not self._context.host.is_control_rank:
+ raise RuntimeError("WebRTC offers are handled only on the control rank.")
+ reservation = self._context.admission.try_reserve()
+ if reservation is None:
+ raise SessionBusyError(self._busy_message)
+
+ task: asyncio.Task[RunResult] | None = None
+ try:
+ scenario = await self._run_blocking_prepare(self._spec)
+ task = asyncio.create_task(
+ self._session_helper(
+ context=self._context,
+ spec=self._spec,
+ scenario=scenario,
+ adapter=self._adapter,
+ run_mode=self._run_mode,
+ pipeline=self._pipeline,
+ reservation=reservation,
+ )
+ )
+ self._track_task(task)
+ answer = await self._answerer.create_answer(
+ offer=WebRTCOfferRequest(sdp=offer_sdp, type=offer_type),
+ session_task=task,
+ )
+ return dict(answer)
+ except Exception:
+ if task is not None:
+ task.cancel()
+ with contextlib.suppress(asyncio.CancelledError):
+ await task
+ reservation.release()
+ raise
+
+ async def _run_blocking_prepare(self, spec: DemoSpec) -> PreparedScenario:
+ service = self._run_mode.blocking_preparation
+ result = await service.run(self._adapter.prepare_scenario, spec)
+ if not isinstance(result, PreparedScenario):
+ raise TypeError(
+ "DemoAdapter.prepare_scenario must return PreparedScenario, "
+ f"got {type(result).__name__}."
+ )
+ return result
+
+ def _track_task(self, task: asyncio.Task[RunResult]) -> None:
+ self._session_tasks.add(task)
+ task.add_done_callback(self._discard_task)
+
+ def _discard_task(self, task: asyncio.Task[RunResult]) -> None:
+ self._session_tasks.discard(task)
+ if task.cancelled():
+ return
+ with contextlib.suppress(Exception):
+ task.exception()
+
+
+WEBRTC_USER_INPUT_SCHEMA = UserInputSchema(
+ capabilities=(
+ UserInputCapability(
+ event_type="key_down",
+ input_modality="keyboard",
+ payload_fields=frozenset({"key"}),
+ ),
+ UserInputCapability(
+ event_type="key_up",
+ input_modality="keyboard",
+ payload_fields=frozenset({"key"}),
+ ),
+ UserInputCapability(
+ event_type="text_event",
+ input_modality="text",
+ payload_fields=frozenset({"event_id", "state"}),
+ ),
+ ),
+ description="browser WebRTC data-channel events",
+)
+
+
+def _user_input_window(
+ *,
+ start_s: float,
+ end_s: float,
+ frame_times: Sequence[float],
+ inputs: UserInputs,
+ metadata: Mapping[str, object],
+) -> UserInputWindow:
+ return UserInputWindow(
+ start_s=start_s,
+ end_s=end_s,
+ frame_times=frame_times,
+ inputs=inputs,
+ metadata=metadata,
+ )
+
+
+class _ThreadSafeActivationSignal:
+ def __init__(self, *, loop: asyncio.AbstractEventLoop | None = None) -> None:
+ self._loop = loop
+ self._event = asyncio.Event()
+
+ def is_set(self) -> bool:
+ return self._event.is_set()
+
+ async def wait(self) -> object:
+ return await self._event.wait()
+
+ def set(self) -> None:
+ if self._event.is_set():
+ return
+ if self._loop is not None and self._loop.is_running():
+ self._loop.call_soon_threadsafe(self._event.set)
+ return
+ self._event.set()
+
+ def clear(self) -> None:
+ self._event.clear()
+
+
+__all__ = [
+ "AsyncioBlockingPreparationService",
+ "BlockingPreparationService",
+ "ThreadSafeWebRTCOutputBridge",
+ "WEBRTC_USER_INPUT_SCHEMA",
+ "WEBRTC_SKIPPED_INPUTS_METADATA_KEY",
+ "WEBRTC_SKIPPED_WINDOW_METADATA_KEY",
+ "WebRTCActivationPolicy",
+ "WebRTCInputSource",
+ "WebRTCMessageResult",
+ "WebRTCOfferAnswerer",
+ "WebRTCOfferRequest",
+ "WebRTCOutputBridge",
+ "WebRTCOutputBridgeDecision",
+ "WebRTCChunkDelivery",
+ "WebRTCOutputSink",
+ "WebRTCRunMode",
+ "WebRTCSessionEdgeFactory",
+ "WebRTCSessionOfferHandler",
+ "WebRTCTransportService",
+]
diff --git a/flashdreams/flashdreams/serving/webrtc/warmup.py b/flashdreams/flashdreams/serving/webrtc/warmup.py
index 05f3d1c10..5f59c4a86 100644
--- a/flashdreams/flashdreams/serving/webrtc/warmup.py
+++ b/flashdreams/flashdreams/serving/webrtc/warmup.py
@@ -13,6 +13,8 @@
from aiortc.mediastreams import MediaStreamError
from loguru import logger as loguru_logger
+from .messages import MESSAGE_TYPE_CHUNK_DONE, MESSAGE_TYPE_ERROR
+
class CreateAnswerCallback(Protocol):
async def __call__(self, *, offer_sdp: str, offer_type: str) -> dict[str, str]: ...
@@ -47,14 +49,30 @@ async def run_loopback_warmup_session(
client_peer.addTransceiver("video", direction="recvonly")
channel_open = asyncio.Event()
warmup_done = asyncio.Event()
+ warmup_failure: str | None = None
received_chunks = 0
drain_tasks: set[asyncio.Task[Any]] = set()
heartbeat_task: asyncio.Task[Any] | None = None
+ def fail_warmup(reason: str) -> None:
+ nonlocal warmup_failure
+ if warmup_done.is_set():
+ return
+ warmup_failure = reason
+ warmup_done.set()
+
@control_channel.on("open")
def on_open() -> None:
channel_open.set()
+ @control_channel.on("close")
+ def on_close() -> None:
+ if received_chunks < num_chunks:
+ fail_warmup(
+ f"{label} loopback warmup data channel closed before warmup "
+ f"completed ({received_chunks}/{num_chunks} chunk(s))."
+ )
+
@control_channel.on("message")
def on_message(message: Any) -> None:
nonlocal received_chunks
@@ -64,7 +82,14 @@ def on_message(message: Any) -> None:
payload = json.loads(message)
except json.JSONDecodeError:
return
- if not isinstance(payload, dict) or payload.get("type") != "chunk_done":
+ if not isinstance(payload, dict):
+ return
+ message_type = payload.get("type")
+ if message_type == MESSAGE_TYPE_ERROR:
+ message_text = str(payload.get("message", "unknown error"))
+ fail_warmup(f"{label} loopback warmup failed: {message_text}")
+ return
+ if message_type != MESSAGE_TYPE_CHUNK_DONE:
return
received_chunks += 1
logger.info(
@@ -109,6 +134,8 @@ def on_track(track: Any) -> None:
for action_payload in action_payloads:
control_channel.send(json.dumps(action_payload))
await asyncio.wait_for(warmup_done.wait(), timeout=warmup_timeout_s)
+ if warmup_failure is not None:
+ raise RuntimeError(warmup_failure)
finally:
if heartbeat_task is not None:
heartbeat_task.cancel()
diff --git a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py
index 735179a4a..0c62855b5 100644
--- a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py
+++ b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py
@@ -83,13 +83,15 @@ def do_HEAD(self) -> None:
def _serve_ui_config(self) -> bool:
if urlsplit(self.path).path != "/api/ui/config":
return False
- adapter_module = (
- "/model-static/adapter.js?v=model-ui-v1"
- if self.model_web_dir is not None
- and (self.model_web_dir / "adapter.js").is_file()
- else None
- )
- payload = json.dumps({"adapter_module": adapter_module}).encode("utf-8")
+ ui_config: dict[str, str | None] = {"adapter_module": None}
+ if self.model_web_dir is not None:
+ if (self.model_web_dir / "adapter.js").is_file():
+ ui_config["adapter_module"] = "/model-static/adapter.js?v=model-ui-v2"
+ if (self.model_web_dir / "adapter.css").is_file():
+ ui_config["model_stylesheet"] = (
+ "/model-static/adapter.css?v=model-ui-v2"
+ )
+ payload = json.dumps(ui_config).encode("utf-8")
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.html b/flashdreams/flashdreams/serving/webrtc/web/request_session.html
index ee42c82a0..ad158656e 100644
--- a/flashdreams/flashdreams/serving/webrtc/web/request_session.html
+++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.html
@@ -87,6 +87,6 @@ Client Logs
-
+