From 33f94f3f6ac4ac2edcfe32da395814e3500d1f4f Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 21 Aug 2026 01:21:48 +0000 Subject: [PATCH 1/6] Refactor v2 application session lifecycle --- flashdreams/flashdreams/api_v2/application.py | 18 ++-- .../runtime_v2/application_runner.py | 90 ++++++++++++++----- flashdreams/flashdreams/runtime_v2/cli.py | 57 +++++------- .../flashdreams/runtime_v2/session_desc.py | 49 +++++++++- flashdreams/flashdreams/t2v_v2/application.py | 2 +- flashdreams/flashdreams/t2v_v2/testing.py | 2 +- .../test_v2/test_application_runner.py | 64 +++++++------ flashdreams/test_v2/test_cli.py | 34 +++++++ flashdreams/test_v2/test_t2v_application.py | 4 +- .../color_fade/tests/test_color_fade.py | 22 +++-- integrations_v2/red_screen/red_screen/app.py | 11 ++- .../tests/test_stand_in_model.py | 2 +- .../tests/test_stand_in_model.py | 2 +- .../tests/test_stand_in_model.py | 2 +- .../tests/test_stand_in_model.py | 2 +- .../t2v_wan21/tests/test_stand_in_model.py | 2 +- 16 files changed, 251 insertions(+), 112 deletions(-) diff --git a/flashdreams/flashdreams/api_v2/application.py b/flashdreams/flashdreams/api_v2/application.py index 6a3c88002..45f2e8eb9 100644 --- a/flashdreams/flashdreams/api_v2/application.py +++ b/flashdreams/flashdreams/api_v2/application.py @@ -27,17 +27,19 @@ def init(self, commandline_args: Sequence[str]) -> None: """Parse application arguments and validate startup state.""" ... - def session_desc(self) -> SessionDesc | None: - """Return the description of a session this application would generate. + def default_session_desc(self) -> SessionDesc | None: + """Return this initialized application's default session description. - A caller has to describe a session before there is one to describe, and - only the application knows what its model was trained for. Asked before - :meth:`init`, so describing a session costs nothing. + The application owns its model's output requirements and defaults, such + as its layout, preferred dimensions, and frame rate. The runtime asks + after :meth:`init`, then applies the caller's explicit requests to this + default before calling :meth:`create_session`. That method may reject + or further resolve the request to satisfy the model's requirements. The + created session's ``session_desc`` is authoritative when a window opens. Returns: - The session to create when nobody asks for another, or ``None``, - the default, from an application that generates whatever it is - asked for. Its caller describes the session instead. + The default session description, or ``None`` when the application + has no output requirements and accepts the runtime's defaults. """ return None diff --git a/flashdreams/flashdreams/runtime_v2/application_runner.py b/flashdreams/flashdreams/runtime_v2/application_runner.py index 883362abe..ebc20c91f 100644 --- a/flashdreams/flashdreams/runtime_v2/application_runner.py +++ b/flashdreams/flashdreams/runtime_v2/application_runner.py @@ -9,7 +9,8 @@ from flashdreams.api_v2.application import IApplication from flashdreams.api_v2.client_window import IClientWindow -from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.api_v2.session import ISession +from flashdreams.runtime_v2.session_desc import SessionDesc, SessionDescRequest from flashdreams.runtime_v2.session_runner import run_session _LOGGER = logging.getLogger(__name__) @@ -17,47 +18,90 @@ class ApplicationRunner: - """Create and run one application session against one client window.""" + """Hold one initialized application and run its sessions.""" - def __init__(self, application: IApplication, client_window: IClientWindow) -> None: + def __init__(self, application: IApplication) -> None: """ Args: - application: Long-lived application that creates the session. - client_window: Window that supplies input and presents generated output. + application: Long-lived application that creates sessions. """ self._application = application - self._client_window = client_window + self._initialized = False + self._initialization_attempted = False + self._closed = False + + def init(self, commandline_args: Sequence[str] = ()) -> None: + """Initialize the application once. + + Args: + commandline_args: Arguments owned and parsed by the application. + + Raises: + RuntimeError: The application has already been initialized or closed. + """ + if self._closed: + raise RuntimeError("ApplicationRunner is closed.") + if self._initialization_attempted: + raise RuntimeError("ApplicationRunner is already initialized.") + self._initialization_attempted = True + self._application.init(commandline_args) + self._initialized = True + + def create_session(self, session_desc_request: SessionDescRequest) -> ISession: + """Create a session for one request against the initialized application. - def run( - self, session_desc: SessionDesc, commandline_args: Sequence[str] = () + Args: + session_desc_request: Explicit overrides to apply to the + application's initialized default description. + + Returns: + A new, uninitialized session. + + Raises: + RuntimeError: The application has not been initialized or is closed. + """ + if self._closed: + raise RuntimeError("ApplicationRunner is closed.") + if not self._initialized: + raise RuntimeError("ApplicationRunner.init() must run first.") + default = self._application.default_session_desc() or SessionDesc() + return self._application.create_session(session_desc_request.resolve(default)) + + def run_session( + self, + session_desc_request: SessionDescRequest, + client_window: IClientWindow, ) -> None: - """Initialize the application, create one session, and run it. + """Create and run one session against ``client_window``. The run ends when the window reports a close or the session reports that it has finished. - The application is closed before this method returns or raises. - - The window is closed too when the run never starts, since ``run_session`` - is what otherwise owns it, and a window may already be serving a client - before the application has loaded anything. + The session and window are closed before this method returns or raises. + The application remains initialized, so callers can run another session + without reloading its shared state. Call :meth:`close` when no further + sessions are needed. Args: - session_desc: Output shape and timing requested for the session. - commandline_args: Arguments owned and parsed by the application. + session_desc_request: Explicit overrides to apply to the + application's initialized default description. + client_window: Window that supplies input and presents generated output. """ run_started = False try: - self._application.init(commandline_args) - session = self._application.create_session(session_desc) + session = self.create_session(session_desc_request) run_started = True - run_session(session, self._client_window) + run_session(session, client_window) finally: if not run_started: - _close_client_window(self._client_window) - _close_application( - self._application, run_failed=sys.exc_info()[0] is not None - ) + _close_client_window(client_window) + + def close(self) -> None: + """Release the application and the state it shares across sessions.""" + if self._closed: + return + self._closed = True + _close_application(self._application, run_failed=sys.exc_info()[0] is not None) def _close_client_window(client_window: IClientWindow) -> None: diff --git a/flashdreams/flashdreams/runtime_v2/cli.py b/flashdreams/flashdreams/runtime_v2/cli.py index 438c019cb..175580c5e 100644 --- a/flashdreams/flashdreams/runtime_v2/cli.py +++ b/flashdreams/flashdreams/runtime_v2/cli.py @@ -5,8 +5,8 @@ ``flashdreams-run-v2`` finds an application by slug, gives it the arguments after ``--``, and hands it to :class:`ApplicationRunner` along with the window -``--mode`` asked for. The session it asks for is the one the application says it -would generate, with whatever the frame arguments here override. +``--mode`` asked for. The runner initializes the application, then resolves its +session default with whatever frame arguments this command overrides. What the modes are, and what each one takes, belongs to :mod:`flashdreams.runtime_v2.client_window_factory`. Nothing here reads an @@ -16,10 +16,7 @@ import argparse import sys from collections.abc import Sequence -from dataclasses import replace -from typing import Any -from flashdreams.api_v2.application import IApplication from flashdreams.runtime_v2.application_registry import ( create_application, registered_application_slugs, @@ -29,7 +26,7 @@ add_client_window_arguments, client_window_mode, ) -from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.session_desc import SessionDescRequest from flashdreams.runtime_v2.video_tensor import VideoTensorLayout _ARGUMENT_SEPARATOR = "--" @@ -55,13 +52,17 @@ def entrypoint(argv: Sequence[str] | None = None) -> None: # Before the window, so a slug this cannot run costs nothing to find out. application = create_application(parsed.slug) - session_desc = _session_desc(application, parsed) - window = mode.create(parsed) - _report(mode.starting(window)) - # Nothing here says how long the run is: a session reports itself finished, - # and a window ends the run when its client goes away. - ApplicationRunner(application, window).run(session_desc, application_args) - _report(mode.finished(window)) + runner = ApplicationRunner(application) + try: + runner.init(application_args) + window = mode.create(parsed) + _report(mode.starting(window)) + # Nothing here says how long the run is: a session reports itself finished, + # and a window ends the run when its client goes away. + runner.run_session(_session_desc_request(parsed), window) + _report(mode.finished(window)) + finally: + runner.close() def split_arguments(arguments: Sequence[str]) -> tuple[list[str], list[str]]: @@ -132,25 +133,11 @@ def _add_session_arguments(parser: argparse.ArgumentParser) -> None: ) -def _session_desc( - application: IApplication, parsed_args: argparse.Namespace -) -> SessionDesc: - """Return the session to ask for: the application's, with the arguments on top. - - An application describing no session of its own gets the arguments alone, - over :class:`SessionDesc`'s own defaults. - """ - asked_for: dict[str, Any] = { - field: value - for field, value in ( - ("output_layout", parsed_args.layout), - ("frames_per_second_for_step", parsed_args.fps), - ("video_width", parsed_args.pixel_width), - ("video_height", parsed_args.pixel_height), - ) - if value is not None - } - described = application.session_desc() - if described is None: - return SessionDesc(**asked_for) - return replace(described, **asked_for) +def _session_desc_request(parsed_args: argparse.Namespace) -> SessionDescRequest: + """Return the output properties the command line explicitly requested.""" + return SessionDescRequest( + output_layout=parsed_args.layout, + frames_per_second_for_step=parsed_args.fps, + video_width=parsed_args.pixel_width, + video_height=parsed_args.pixel_height, + ) diff --git a/flashdreams/flashdreams/runtime_v2/session_desc.py b/flashdreams/flashdreams/runtime_v2/session_desc.py index efeb7dbed..7333cb846 100644 --- a/flashdreams/flashdreams/runtime_v2/session_desc.py +++ b/flashdreams/flashdreams/runtime_v2/session_desc.py @@ -4,7 +4,7 @@ """Description of the session a runtime asks an application for.""" import math -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any from flashdreams.runtime_v2.video_tensor import VideoTensorLayout @@ -56,3 +56,50 @@ def __post_init__(self) -> None: raise ValueError("SessionDesc.video_width must be > 0 when set.") if self.video_height <= 0: raise ValueError("SessionDesc.video_height must be > 0 when set.") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class SessionDescRequest: + """The parts of a session description a caller explicitly requests. + + Leaving a field unset preserves the application's initialized default. The + runtime resolves this request before it asks the application to create a + session. + """ + + output_layout: VideoTensorLayout | None = None + frames_per_second_for_ui: int | None = None + frames_per_second_for_step: int | None = None + video_width: int | None = None + video_height: int | None = None + metadata: dict[str, Any] | None = None + + def resolve(self, default: SessionDesc) -> SessionDesc: + """Return ``default`` with every explicitly requested field replaced.""" + return replace( + default, + output_layout=( + default.output_layout + if self.output_layout is None + else self.output_layout + ), + frames_per_second_for_ui=( + default.frames_per_second_for_ui + if self.frames_per_second_for_ui is None + else self.frames_per_second_for_ui + ), + frames_per_second_for_step=( + default.frames_per_second_for_step + if self.frames_per_second_for_step is None + else self.frames_per_second_for_step + ), + video_width=( + default.video_width if self.video_width is None else self.video_width + ), + video_height=( + default.video_height + if self.video_height is None + else self.video_height + ), + metadata=default.metadata if self.metadata is None else self.metadata, + ) diff --git a/flashdreams/flashdreams/t2v_v2/application.py b/flashdreams/flashdreams/t2v_v2/application.py index f21ad8364..26f24e483 100644 --- a/flashdreams/flashdreams/t2v_v2/application.py +++ b/flashdreams/flashdreams/t2v_v2/application.py @@ -136,7 +136,7 @@ def init(self, commandline_args: Sequence[str]) -> None: total_blocks=args.total_blocks, ) - def session_desc(self) -> SessionDesc: + def default_session_desc(self) -> SessionDesc: """Return the description of a session this application uses. A model generates its best video at the size and rate it was trained diff --git a/flashdreams/flashdreams/t2v_v2/testing.py b/flashdreams/flashdreams/t2v_v2/testing.py index 4751ee907..9c545ed56 100644 --- a/flashdreams/flashdreams/t2v_v2/testing.py +++ b/flashdreams/flashdreams/t2v_v2/testing.py @@ -132,7 +132,7 @@ def check_t2v_model_impl( application.init(commandline_args) try: if session_desc is None: - session_desc = application.session_desc() + session_desc = application.default_session_desc() run_session( application.create_session(session_desc), _InspectingClientWindow(inspector), diff --git a/flashdreams/test_v2/test_application_runner.py b/flashdreams/test_v2/test_application_runner.py index e77b73732..8b872b1f1 100644 --- a/flashdreams/test_v2/test_application_runner.py +++ b/flashdreams/test_v2/test_application_runner.py @@ -14,7 +14,7 @@ from flashdreams.api_v2.client_window import IClientWindow from flashdreams.api_v2.session import ISession from flashdreams.runtime_v2.application_runner import ApplicationRunner -from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.session_desc import SessionDesc, SessionDescRequest from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import ( CloseUserInputEventData, @@ -136,8 +136,8 @@ def get_user_input_events(self) -> UserInputEvents: return UserInputEvents([]) -def _session_desc() -> SessionDesc: - return SessionDesc( +def _session_desc_request() -> SessionDescRequest: + return SessionDescRequest( output_layout=VideoTensorLayout.bcthw, frames_per_second_for_ui=100, frames_per_second_for_step=30, @@ -146,45 +146,51 @@ def _session_desc() -> SessionDesc: ) -def test_application_runner_drives_complete_lifecycle() -> None: +def test_application_runner_keeps_the_application_open_for_another_session() -> None: calls: list[str] = [] - application = _Application(calls) - window = _Window(calls) - - ApplicationRunner(application, window).run(_session_desc(), ["--model-option"]) - - assert window.results == [] + application = _Application(calls, session_length=1) + runner = ApplicationRunner(application) + first_window = _SilentWindow(calls) + second_window = _SilentWindow(calls) + + runner.init(["--model-option"]) + runner.run_session(_session_desc_request(), first_window) + runner.run_session(_session_desc_request(), second_window) + + assert [result.step_index for result in first_window.results] == [0] + assert [result.step_index for result in second_window.results] == [0] + assert calls.count("application.create_session") == 2 + assert calls.count("application.close") == 0 assert calls[0:3] == [ "application.init(['--model-option'])", "application.create_session", "session.init", ] - assert calls[-3:] == ["window.close", "session.close", "application.close"] + runner.close() + assert calls[-1] == "application.close" -def test_application_runner_closes_both_when_the_run_never_starts() -> None: - """The window is closed by the loop, which a failure here never reaches, and - a window may already be serving a client by then.""" +def test_application_runner_closes_the_window_when_a_session_cannot_start() -> None: calls: list[str] = [] - application = _Application(calls, fail_to_init=True) + runner = ApplicationRunner(_Application(calls)) + window = _Window(calls) - with pytest.raises(RuntimeError, match="application init failed"): - ApplicationRunner(application, _Window(calls)).run(_session_desc()) + with pytest.raises(RuntimeError, match=r"init\(\) must run first"): + runner.run_session(_session_desc_request(), window) - assert calls == ["application.init([])", "window.close", "application.close"] + assert calls == ["window.close"] -def test_application_runner_ends_a_run_a_window_cannot_end() -> None: - """A window with no client never reports a close, so the session ends it.""" +def test_application_runner_rejects_a_second_initialization() -> None: calls: list[str] = [] - window = _SilentWindow(calls) + runner = ApplicationRunner(_Application(calls)) - ApplicationRunner(_Application(calls, session_length=3), window).run( - _session_desc() - ) + runner.init() + with pytest.raises(RuntimeError, match="already initialized"): + runner.init() + runner.close() - assert [result.step_index for result in window.results] == [0, 1, 2] - assert calls[-3:] == ["window.close", "session.close", "application.close"] + assert calls == ["application.init([])", "application.close"] def test_application_runner_reports_the_run_rather_than_the_close( @@ -192,9 +198,13 @@ def test_application_runner_reports_the_run_rather_than_the_close( ) -> None: calls: list[str] = [] application = _Application(calls, fail_to_init=True, fail_to_close=True) + runner = ApplicationRunner(application) with caplog.at_level(logging.ERROR, logger=_RUNNER_LOGGER): with pytest.raises(RuntimeError, match="application init failed"): - ApplicationRunner(application, _Window(calls)).run(_session_desc()) + try: + runner.init() + finally: + runner.close() assert "application close failed" in caplog.text diff --git a/flashdreams/test_v2/test_cli.py b/flashdreams/test_v2/test_cli.py index ff65f7e8c..8cd62499e 100644 --- a/flashdreams/test_v2/test_cli.py +++ b/flashdreams/test_v2/test_cli.py @@ -98,6 +98,28 @@ def create_session(self, session_desc: SessionDesc) -> ISession: return OneStepSession(session_desc) +class InitializedDescriptionApplication(IApplication): + """Application whose default output width comes from its own arguments.""" + + def __init__(self) -> None: + self._width: int | None = None + self.asked_for: SessionDesc | None = None + + def init(self, commandline_args: Sequence[str]) -> None: + if len(commandline_args) != 2 or commandline_args[0] != "--width": + raise ValueError("--width is required.") + self._width = int(commandline_args[1]) + + def default_session_desc(self) -> SessionDesc: + if self._width is None: + raise RuntimeError("init() must run before default_session_desc().") + return SessionDesc(video_width=self._width) + + def create_session(self, session_desc: SessionDesc) -> ISession: + self.asked_for = session_desc + return OneStepSession(session_desc) + + class OneStepSession(ISession): """A session generating one frame and reporting itself finished.""" @@ -437,6 +459,18 @@ def test_a_model_generates_what_it_was_trained_for_unless_asked_otherwise( assert window.session_desc.video_height == pipeline.height +def test_application_arguments_resolve_the_default_session_before_creation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + application = InitializedDescriptionApplication() + _install(monkeypatch, application, RecordingWindow()) + + cli.entrypoint(["stub", "--mode", "webrtc", "--", "--width", "64"]) + + assert application.asked_for is not None + assert application.asked_for.video_width == 64 + + ## The command itself diff --git a/flashdreams/test_v2/test_t2v_application.py b/flashdreams/test_v2/test_t2v_application.py index 0efebf7da..9984087d9 100644 --- a/flashdreams/test_v2/test_t2v_application.py +++ b/flashdreams/test_v2/test_t2v_application.py @@ -136,7 +136,7 @@ def _rollout_length(app: T2VApplication) -> int: An application answers no such question: the length reaches a run through the session, which reports itself finished at the end of it. """ - session = app.create_session(app.session_desc()) + session = app.create_session(app.default_session_desc()) assert isinstance(session, RecordingRollout) return session.blocks_to_generate @@ -228,7 +228,7 @@ def test_a_model_describes_the_clip_it_was_trained_for_without_loading() -> None describe a session before it can ask for one.""" app = T2VApplication(defaults=_defaults()) - desc = app.session_desc() + desc = app.default_session_desc() assert desc == _session_desc() assert _pipeline_config(app).setup_count == 0 diff --git a/integrations_v2/color_fade/color_fade/tests/test_color_fade.py b/integrations_v2/color_fade/color_fade/tests/test_color_fade.py index dc7e4acc3..c0f2ee068 100644 --- a/integrations_v2/color_fade/color_fade/tests/test_color_fade.py +++ b/integrations_v2/color_fade/color_fade/tests/test_color_fade.py @@ -16,7 +16,7 @@ from flashdreams.api_v2.session import ISession from flashdreams.runtime_v2.application_runner import ApplicationRunner from flashdreams.runtime_v2.mp4_client_window import Mp4ClientWindow -from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.session_desc import SessionDesc, SessionDescRequest from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_events import UserInputEvents from flashdreams.runtime_v2.video_tensor import VideoTensorLayout @@ -260,10 +260,22 @@ def test_a_run_writes_the_whole_fade_to_an_mp4(tmp_path: Path) -> None: frames_per_step = 5 # No step count: the session knows how long its fade is and ends the run. - ApplicationRunner(create_app(), Mp4ClientWindow(path)).run( - _session_desc(width=_PLAYABLE_WIDTH, height=_PLAYABLE_HEIGHT), - ["--seconds", str(_SECONDS), "--frames-per-step", str(frames_per_step)], - ) + runner = ApplicationRunner(create_app()) + try: + runner.init( + ["--seconds", str(_SECONDS), "--frames-per-step", str(frames_per_step)] + ) + runner.run_session( + SessionDescRequest( + output_layout=VideoTensorLayout.bcthw, + frames_per_second_for_step=_FRAMES_PER_SECOND, + video_width=_PLAYABLE_WIDTH, + video_height=_PLAYABLE_HEIGHT, + ), + Mp4ClientWindow(path), + ) + finally: + runner.close() frames = _decode(path, width=_PLAYABLE_WIDTH, height=_PLAYABLE_HEIGHT) assert len(frames) == _STEPS_FOR_THE_FADE * frames_per_step diff --git a/integrations_v2/red_screen/red_screen/app.py b/integrations_v2/red_screen/red_screen/app.py index 7a4e64892..63cb78b15 100644 --- a/integrations_v2/red_screen/red_screen/app.py +++ b/integrations_v2/red_screen/red_screen/app.py @@ -14,7 +14,7 @@ from flashdreams.api_v2.session import ISession from flashdreams.runtime_v2.application_runner import ApplicationRunner from flashdreams.runtime_v2.client_window_factory import create_client_window -from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.session_desc import SessionDesc, SessionDescRequest from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import KeyboardUserInputEventData from flashdreams.runtime_v2.user_input_events import UserInputEvents @@ -205,6 +205,7 @@ def main(commandline_args: Sequence[str] | None = None) -> int: window = create_client_window(args) app = create_app() + runner = ApplicationRunner(app) if isinstance(window, WebRTCClientWindow): print(f"Open {window.server.url} in a browser.", flush=True) try: @@ -213,20 +214,22 @@ def main(commandline_args: Sequence[str] | None = None) -> int: # TODO: in production, commandline argument parsing and IClientWindow creation should be done by flashdreams-run, a CLI tool # basically, we need to generailze this main function to be shared by all applications - ApplicationRunner(app, window).run( - SessionDesc( + runner.init(application_args) + runner.run_session( + SessionDescRequest( output_layout=VideoTensorLayout.bcthw, frames_per_second_for_ui=args.fps, frames_per_second_for_step=args.fps, video_width=args.width, video_height=args.height, ), - application_args, + window, ) except KeyboardInterrupt: return 130 finally: window.close() + runner.close() return 0 diff --git a/integrations_v2/t2v_causal_forcing/t2v_causal_forcing/tests/test_stand_in_model.py b/integrations_v2/t2v_causal_forcing/t2v_causal_forcing/tests/test_stand_in_model.py index 2ef24d233..80d8f7f76 100644 --- a/integrations_v2/t2v_causal_forcing/t2v_causal_forcing/tests/test_stand_in_model.py +++ b/integrations_v2/t2v_causal_forcing/t2v_causal_forcing/tests/test_stand_in_model.py @@ -26,7 +26,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None: integration already ships rather than written down again.""" app = CausalForcingT2VApplication(pipeline_config=FakeT2VPipelineConfig()) - desc = app.session_desc() + desc = app.default_session_desc() assert (desc.video_width, desc.video_height) == ( RUNNER_WAN21_T2V_1PT3B_CHUNKWISE.pixel_width, diff --git a/integrations_v2/t2v_cosmos_predict2/t2v_cosmos_predict2/tests/test_stand_in_model.py b/integrations_v2/t2v_cosmos_predict2/t2v_cosmos_predict2/tests/test_stand_in_model.py index 303dd5e35..31fc223fa 100644 --- a/integrations_v2/t2v_cosmos_predict2/t2v_cosmos_predict2/tests/test_stand_in_model.py +++ b/integrations_v2/t2v_cosmos_predict2/t2v_cosmos_predict2/tests/test_stand_in_model.py @@ -26,7 +26,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None: integration already ships rather than written down again.""" app = CosmosPredict2T2VApplication(pipeline_config=FakeT2VPipelineConfig()) - desc = app.session_desc() + desc = app.default_session_desc() assert (desc.video_width, desc.video_height) == ( RUNNER_COSMOS2_T2V_2B_720P.pixel_width, diff --git a/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py b/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py index 669aba346..75900757b 100644 --- a/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py +++ b/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py @@ -26,7 +26,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None: integration already ships rather than written down again.""" app = FastvideoCausalWan22T2VApplication(pipeline_config=FakeT2VPipelineConfig()) - desc = app.session_desc() + desc = app.default_session_desc() assert (desc.video_width, desc.video_height) == ( RUNNER_WAN22_T2V_14B.pixel_width, diff --git a/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py b/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py index 67b5e0b8e..0ca793e64 100644 --- a/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py +++ b/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py @@ -36,7 +36,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None: integration already ships rather than written down again.""" app = SelfForcingT2VApplication(pipeline_config=FakeT2VPipelineConfig()) - desc = app.session_desc() + desc = app.default_session_desc() assert (desc.video_width, desc.video_height) == ( RUNNER_WAN21_T2V_1PT3B.pixel_width, diff --git a/integrations_v2/t2v_wan21/t2v_wan21/tests/test_stand_in_model.py b/integrations_v2/t2v_wan21/t2v_wan21/tests/test_stand_in_model.py index 26ee7140b..f3a337b81 100644 --- a/integrations_v2/t2v_wan21/t2v_wan21/tests/test_stand_in_model.py +++ b/integrations_v2/t2v_wan21/t2v_wan21/tests/test_stand_in_model.py @@ -27,7 +27,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None: for a model that does not roll out does not carry one.""" app = Wan21T2VApplication(pipeline_config=FakeT2VPipelineConfig()) - desc = app.session_desc() + desc = app.default_session_desc() assert (desc.video_width, desc.video_height) == ( RUNNER_WAN21_T2V_1PT3B_480P.pixel_width, From cb95cbad0a3ca96fd9ea76ec1be292ac989b2b28 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 21 Aug 2026 06:43:31 +0000 Subject: [PATCH 2/6] Add persistent WebRTC sessions for v2 text-to-video --- .../flashdreams/api_v2/client_window.py | 11 +- .../runtime_v2/application_runner.py | 105 +++++++---- flashdreams/flashdreams/runtime_v2/cli.py | 13 +- .../runtime_v2/client_window_factory.py | 4 + .../flashdreams/runtime_v2/serving/web/app.js | 37 +++- .../runtime_v2/serving/web/index.html | 3 + .../runtime_v2/serving/webrtc_server.py | 146 +++++++++++---- .../flashdreams/runtime_v2/session_runner.py | 166 ++++++++++++++---- .../runtime_v2/user_input_event.py | 20 ++- .../runtime_v2/webrtc_client_window.py | 2 +- flashdreams/flashdreams/t2v_v2/application.py | 32 ++-- .../test_v2/test_application_runner.py | 105 +++++++++++ flashdreams/test_v2/test_cli.py | 42 ++++- .../test_v2/test_client_window_factory.py | 2 + flashdreams/test_v2/test_session_runner.py | 47 +++++ flashdreams/test_v2/test_t2v_application.py | 37 +++- .../test_v2/test_webrtc_client_window.py | 78 +++++++- integrations_v2/t2v_self_forcing/README.md | 13 ++ 18 files changed, 728 insertions(+), 135 deletions(-) diff --git a/flashdreams/flashdreams/api_v2/client_window.py b/flashdreams/flashdreams/api_v2/client_window.py index 8befe1c02..07d65d4dc 100644 --- a/flashdreams/flashdreams/api_v2/client_window.py +++ b/flashdreams/flashdreams/api_v2/client_window.py @@ -13,14 +13,17 @@ class IClientWindow(InputSource, OutputSink, ABC): """Handle application input and output for one client window. The runtime opens the window with the session's description, then reads input - and writes results until the run ends, and closes it then. A window stays open - across a session reset. + and writes results until the run ends. A window stays open across a session + reset. When the client asks for a replacement session, the runtime closes the + old session and opens the same window with the replacement's description. A + session-serving runner can also leave the window open between sessions while + it waits for another client request. A window does not describe the output shape. The session does, and the window is given that description in :meth:`OutputSink.open`. - One thread makes every call on a window, so an implementation needs no - locking except when its backend delivers input from another thread. + One I/O thread at a time makes every call on a window, so an implementation + needs no locking except when its backend delivers input from another thread. Created by the runtime, never by an application. """ diff --git a/flashdreams/flashdreams/runtime_v2/application_runner.py b/flashdreams/flashdreams/runtime_v2/application_runner.py index ebc20c91f..e73f744fa 100644 --- a/flashdreams/flashdreams/runtime_v2/application_runner.py +++ b/flashdreams/flashdreams/runtime_v2/application_runner.py @@ -9,9 +9,8 @@ from flashdreams.api_v2.application import IApplication from flashdreams.api_v2.client_window import IClientWindow -from flashdreams.api_v2.session import ISession from flashdreams.runtime_v2.session_desc import SessionDesc, SessionDescRequest -from flashdreams.runtime_v2.session_runner import run_session +from flashdreams.runtime_v2.session_runner import run_session, wait_for_new_session _LOGGER = logging.getLogger(__name__) """Logger for an application or window that could not be closed.""" @@ -47,35 +46,23 @@ def init(self, commandline_args: Sequence[str] = ()) -> None: self._application.init(commandline_args) self._initialized = True - def create_session(self, session_desc_request: SessionDescRequest) -> ISession: - """Create a session for one request against the initialized application. - - Args: - session_desc_request: Explicit overrides to apply to the - application's initialized default description. - - Returns: - A new, uninitialized session. - - Raises: - RuntimeError: The application has not been initialized or is closed. - """ - if self._closed: - raise RuntimeError("ApplicationRunner is closed.") - if not self._initialized: - raise RuntimeError("ApplicationRunner.init() must run first.") - default = self._application.default_session_desc() or SessionDesc() - return self._application.create_session(session_desc_request.resolve(default)) - def run_session( self, session_desc_request: SessionDescRequest, client_window: IClientWindow, + *, + serve_sessions: bool = False, ) -> None: - """Create and run one session against ``client_window``. + """Create sessions against ``client_window`` until the run ends. - The run ends when the window reports a close or the session reports that - it has finished. + Normally the run ends when the window reports a close or the session + reports that it has finished. A replacement description returned by the + session loop starts another session after the current one has closed. + + With ``serve_sessions``, the window opens before any session exists and + this method waits for a session description. It returns to that waiting + state whenever a session finishes or its browser disconnects. The server + therefore remains available until the process interrupts this method. The session and window are closed before this method returns or raises. The application remains initialized, so callers can run another session @@ -86,15 +73,63 @@ def run_session( session_desc_request: Explicit overrides to apply to the application's initialized default description. client_window: Window that supplies input and presents generated output. + serve_sessions: Keep the window running and create sessions only in + response to client requests. """ - run_started = False + if serve_sessions: + self._serve_sessions(session_desc_request, client_window) + return + try: - session = self.create_session(session_desc_request) - run_started = True - run_session(session, client_window) - finally: - if not run_started: + next_session_desc = self._resolve_session_desc(session_desc_request) + except Exception: + _close_client_window(client_window) + raise + while True: + try: + session = self._application.create_session(next_session_desc) + except Exception: _close_client_window(client_window) + raise + next_session_desc = run_session(session, client_window) + if next_session_desc is None: + return + + def _resolve_session_desc( + self, session_desc_request: SessionDescRequest + ) -> SessionDesc: + """Resolve one request against the initialized application's default.""" + if self._closed: + raise RuntimeError("ApplicationRunner is closed.") + if not self._initialized: + raise RuntimeError("ApplicationRunner.init() must run first.") + default = self._application.default_session_desc() or SessionDesc() + return session_desc_request.resolve(default) + + def _serve_sessions( + self, + session_desc_request: SessionDescRequest, + client_window: IClientWindow, + ) -> None: + """Keep one client window available for browser-requested sessions.""" + current_session_desc = self._resolve_session_desc(session_desc_request) + try: + client_window.open(current_session_desc) + next_session_desc: SessionDesc | None = None + while True: + if next_session_desc is None: + next_session_desc = wait_for_new_session( + client_window, current_session_desc + ) + session = self._application.create_session(next_session_desc) + current_session_desc = session.session_desc + next_session_desc = run_session( + session, + client_window, + keep_window_open=True, + ) + finally: + _close_client_window(client_window) def close(self) -> None: """Release the application and the state it shares across sessions.""" @@ -105,16 +140,16 @@ def close(self) -> None: def _close_client_window(client_window: IClientWindow) -> None: - """Close a window the run never reached, so what it was serving goes with it. + """Close a window during runner cleanup without hiding an active failure. - The run has already failed by the time this is called, so a failure here is - logged rather than raised over the top of it. + This runs after session creation fails or a persistent run is interrupted, + so a failure here is logged rather than raised over the top of it. """ try: client_window.close() except Exception: _LOGGER.exception( - "The client window failed to close after a run that never started." + "The client window failed to close while the runner was stopping." ) diff --git a/flashdreams/flashdreams/runtime_v2/cli.py b/flashdreams/flashdreams/runtime_v2/cli.py index 175580c5e..4e5ac8cd8 100644 --- a/flashdreams/flashdreams/runtime_v2/cli.py +++ b/flashdreams/flashdreams/runtime_v2/cli.py @@ -57,9 +57,16 @@ def entrypoint(argv: Sequence[str] | None = None) -> None: runner.init(application_args) window = mode.create(parsed) _report(mode.starting(window)) - # Nothing here says how long the run is: a session reports itself finished, - # and a window ends the run when its client goes away. - runner.run_session(_session_desc_request(parsed), window) + # Nothing here says how long a session is: the application reports that. + # A serving mode stays up between sessions; a file mode runs just one. + try: + runner.run_session( + _session_desc_request(parsed), + window, + serve_sessions=mode.serves_sessions, + ) + except KeyboardInterrupt: + return _report(mode.finished(window)) finally: runner.close() diff --git a/flashdreams/flashdreams/runtime_v2/client_window_factory.py b/flashdreams/flashdreams/runtime_v2/client_window_factory.py index 155354296..0500b6c65 100644 --- a/flashdreams/flashdreams/runtime_v2/client_window_factory.py +++ b/flashdreams/flashdreams/runtime_v2/client_window_factory.py @@ -27,6 +27,9 @@ class ClientWindowMode(ABC): name: str """What ``--mode`` calls this.""" + serves_sessions: bool = False + """Whether the window stays available between application sessions.""" + def add_arguments(self, parser: argparse.ArgumentParser) -> None: """Add the arguments this mode takes and no other does.""" @@ -98,6 +101,7 @@ class _WebRTCMode(ClientWindowMode): """Stream the run to a browser.""" name = "webrtc" + serves_sessions = True def add_arguments(self, parser: argparse.ArgumentParser) -> None: parser.add_argument( diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/app.js b/flashdreams/flashdreams/runtime_v2/serving/web/app.js index 029e714e7..d64248929 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/web/app.js +++ b/flashdreams/flashdreams/runtime_v2/serving/web/app.js @@ -4,6 +4,15 @@ const peer = new RTCPeerConnection(); const controls = peer.createDataChannel("controls"); peer.addTransceiver("video", {direction: "recvonly"}); +const newSessionButton = document.getElementById("new-session"); + +controls.addEventListener("open", () => { + newSessionButton.disabled = false; +}); + +controls.addEventListener("close", () => { + newSessionButton.disabled = true; +}); peer.ontrack = event => { document.getElementById("video").srcObject = @@ -36,6 +45,17 @@ document.getElementById("reset").onclick = () => { send({type: "reset"}); }; +newSessionButton.onclick = () => { + const promptInput = document.getElementById("prompt"); + if (!promptInput.reportValidity()) { + return; + } + send({ + type: "new_session", + metadata: {prompt: promptInput.value}, + }); +}; + window.addEventListener("beforeunload", () => send({type: "close"})); async function connect() { @@ -47,11 +67,18 @@ async function connect() { await new Promise(resolve => setTimeout(resolve, 100)); } await peer.setLocalDescription(await peer.createOffer()); - const response = await fetch("/api/webrtc/offer", { - method: "POST", - headers: {"content-type": "application/json"}, - body: JSON.stringify(peer.localDescription), - }); + let response; + while (true) { + response = await fetch("/api/webrtc/offer", { + method: "POST", + headers: {"content-type": "application/json"}, + body: JSON.stringify(peer.localDescription), + }); + if (response.status !== 409) { + break; + } + await new Promise(resolve => setTimeout(resolve, 100)); + } if (!response.ok) { throw new Error(await response.text()); } diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/index.html b/flashdreams/flashdreams/runtime_v2/serving/web/index.html index 2fb07c96e..eb8b145b6 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/web/index.html +++ b/flashdreams/flashdreams/runtime_v2/serving/web/index.html @@ -11,6 +11,9 @@ + + + diff --git a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py index 68041879d..b7e5b9112 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py +++ b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py @@ -27,6 +27,7 @@ from flashdreams.runtime_v2.user_input_event import ( CloseUserInputEventData, KeyboardUserInputEventData, + NewSessionUserInputEventData, ResetUserInputEventData, UserInputEvent, ) @@ -46,9 +47,10 @@ def __init__(self, frames_per_second: int) -> None: super().__init__() self._frames_per_second = frames_per_second self._time_base = Fraction(1, frames_per_second) - self._frames: asyncio.Queue[np.ndarray[Any, np.dtype[np.uint8]] | None] = ( - asyncio.Queue() - ) + self._frames: asyncio.Queue[ + tuple[int, np.ndarray[Any, np.dtype[np.uint8]]] | None + ] = asyncio.Queue() + self._session_generation = 0 self._next_frame_time: float | None = None self._pts = 0 self._closed = False @@ -60,29 +62,46 @@ async def enqueue( if self._closed: return for frame in frames: - await self._frames.put(frame) + await self._frames.put((self._session_generation, frame)) + + async def start_session(self) -> None: + """Discard frames queued by the session being replaced.""" + self._session_generation += 1 + self._next_frame_time = None + while True: + try: + self._frames.get_nowait() + except asyncio.QueueEmpty: + return async def recv(self) -> VideoFrame: """Return the next generated frame when aiortc requests one.""" if self._closed: raise MediaStreamError - frame = await self._frames.get() - if frame is None: - raise MediaStreamError - - loop = asyncio.get_running_loop() - now = loop.time() - if self._next_frame_time is None: - self._next_frame_time = now - else: - self._next_frame_time += 1.0 / self._frames_per_second - await asyncio.sleep(max(0.0, self._next_frame_time - now)) - - video_frame = VideoFrame.from_ndarray(frame, format="rgb24") - video_frame.pts = self._pts - video_frame.time_base = self._time_base - self._pts += 1 - return video_frame + while True: + item = await self._frames.get() + if item is None: + raise MediaStreamError + session_generation, frame = item + if session_generation != self._session_generation: + continue + + loop = asyncio.get_running_loop() + now = loop.time() + if self._next_frame_time is None: + self._next_frame_time = now + else: + self._next_frame_time += 1.0 / self._frames_per_second + await asyncio.sleep(max(0.0, self._next_frame_time - now)) + if session_generation != self._session_generation: + self._next_frame_time = None + continue + + video_frame = VideoFrame.from_ndarray(frame, format="rgb24") + video_frame.pts = self._pts + video_frame.time_base = self._time_base + self._pts += 1 + return video_frame async def close(self) -> None: """Stop the track and release a pending receiver.""" @@ -163,22 +182,38 @@ def url(self) -> str: return f"http://{self._host}:{self._port}/" def open(self, session_desc: SessionDesc) -> None: - """Configure the server for one session's generated video. + """Configure the video format this server's sessions generate. Args: session_desc: Resolved dimensions, frame rate, and tensor layout. + A session-serving runner calls this before its first session so the + browser can connect and request one. Each session opens it again while + the browser remains connected. Its stream format must match the first + call because the existing media track keeps its negotiated frame rate. + Raises: - RuntimeError: The server is closed or already open. + RuntimeError: The server is closed or has no input callback. + ValueError: A replacement changes the active stream format. """ if self._closed: raise RuntimeError("Cannot open a closed WebRTC server.") - if self._session_desc is not None: - raise RuntimeError("WebRTC server is already open.") if self._input_callback is None: raise RuntimeError("Register an input callback before opening WebRTC.") + current_desc = self._session_desc + if current_desc is not None and not _same_stream_format( + current_desc, session_desc + ): + raise ValueError( + "A replacement WebRTC session must keep the original output " + "layout, frame rate, width, and height." + ) self._session_desc = session_desc self._session_start_ns = time.monotonic_ns() + track = self._video_track + loop = self._loop + if track is not None and loop is not None: + asyncio.run_coroutine_threadsafe(track.start_session(), loop).result() def register_input_callback( self, callback: Callable[[UserInputEvent], None] @@ -285,7 +320,7 @@ async def _serve_browser_script(self, _: web.Request) -> web.Response: return web.Response(text=_BROWSER_SCRIPT, content_type="text/javascript") async def _health(self, _: web.Request) -> web.Response: - """Report whether the server has an open session and client.""" + """Report whether the server can negotiate video and has a client.""" return web.json_response( { "open": self._session_desc is not None, @@ -300,8 +335,16 @@ async def _offer(self, request: web.Request) -> web.Response: session_desc = self._session_desc if session_desc is None: raise web.HTTPConflict(reason="WebRTC server is not open.") - if self._peer_connection is not None: - raise web.HTTPConflict(reason="A WebRTC client is already connected.") + existing_peer = self._peer_connection + if existing_peer is not None: + if existing_peer.connectionState in { + "failed", + "disconnected", + "closed", + }: + await self._release_peer_connection(existing_peer) + else: + raise web.HTTPConflict(reason="A WebRTC client is already connected.") try: payload = await request.json() @@ -341,6 +384,7 @@ def on_close() -> None: async def on_connectionstatechange() -> None: if peer_connection.connectionState in {"failed", "disconnected", "closed"}: self._record_client_disconnect() + await self._release_peer_connection(peer_connection) try: await peer_connection.setRemoteDescription( @@ -387,11 +431,17 @@ def _buffer_browser_message(self, raw_message: object) -> None: event_data = KeyboardUserInputEventData(key=key, pressed=pressed) elif event_type == "reset": event_data = ResetUserInputEventData() + elif event_type == "new_session": + metadata = payload.get("metadata") + if not isinstance(metadata, dict): + raise ValueError("New-session event requires metadata object.") + event_data = NewSessionUserInputEventData(metadata=metadata) elif event_type == "close": event_data = CloseUserInputEventData() else: raise ValueError( - "Browser event type must be 'keyboard', 'reset', or 'close'." + "Browser event type must be 'keyboard', 'reset', 'new_session', " + "or 'close'." ) self._append_event(event_data) @@ -400,6 +450,7 @@ def _append_event( event_data: ( KeyboardUserInputEventData | ResetUserInputEventData + | NewSessionUserInputEventData | CloseUserInputEventData ), ) -> None: @@ -412,8 +463,8 @@ def _append_event( callback = self._input_callback if callback is None: raise RuntimeError("WebRTC input callback is not registered.") - # Pass that UserInputEvent to the callback. - # The callback stores it in WebRTCClientWindow’s thread-safe queue. + # Every event takes the same callback path into WebRTCClientWindow's + # thread-safe queue; the server loop never handles lifecycle events. callback(event) def _record_client_disconnect(self) -> None: @@ -432,22 +483,47 @@ async def _enqueue_frames( if track is not None: await track.enqueue(frames) - async def _shutdown(self) -> None: - """Release async server resources on their owning loop.""" - peer_connection = self._peer_connection + async def _release_peer_connection( + self, peer_connection: RTCPeerConnection + ) -> None: + """Release one disconnected peer so a refreshed page can reconnect.""" + if self._peer_connection is not peer_connection: + return + self._record_client_disconnect() self._peer_connection = None track = self._video_track self._video_track = None if track is not None: await track.close() - if peer_connection is not None: + if peer_connection.connectionState != "closed": await peer_connection.close() + + async def _shutdown(self) -> None: + """Release async server resources on their owning loop.""" + peer_connection = self._peer_connection + if peer_connection is not None: + await self._release_peer_connection(peer_connection) + else: + track = self._video_track + self._video_track = None + if track is not None: + await track.close() runner = self._runner self._runner = None if runner is not None: await runner.cleanup() +def _same_stream_format(first: SessionDesc, second: SessionDesc) -> bool: + """Return whether two sessions can use the same WebRTC media track.""" + return ( + first.output_layout is second.output_layout + and first.frames_per_second_for_ui == second.frames_per_second_for_ui + and first.video_width == second.video_width + and first.video_height == second.video_height + ) + + def _result_to_rgb_frames( result: StepResult, session_desc: SessionDesc ) -> tuple[np.ndarray[Any, np.dtype[np.uint8]], ...]: diff --git a/flashdreams/flashdreams/runtime_v2/session_runner.py b/flashdreams/flashdreams/runtime_v2/session_runner.py index 5d013930b..54679d105 100644 --- a/flashdreams/flashdreams/runtime_v2/session_runner.py +++ b/flashdreams/flashdreams/runtime_v2/session_runner.py @@ -7,14 +7,18 @@ import queue import sys import threading +import time +from dataclasses import replace from enum import Enum from flashdreams.api_v2.client_window import IClientWindow from flashdreams.api_v2.session import ISession from flashdreams.api_v2.user_input_event_data import UserInputEventData +from flashdreams.runtime_v2.session_desc import SessionDesc from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import ( CloseUserInputEventData, + NewSessionUserInputEventData, ResetUserInputEventData, UserInputEvent, ) @@ -41,6 +45,50 @@ def _contains(events: UserInputEvents, event_type: type[UserInputEventData]) -> ) +def _next_session_desc( + events: UserInputEvents, current: SessionDesc +) -> SessionDesc | None: + """Return the latest session requested by ``events``, if there is one. + + A new-session request replaces application-owned metadata while preserving + the resolved video and timing settings of the current session. The + application interprets and validates the metadata when it creates the next + session. + """ + for event in reversed(events.get_events()): + event_data = event.get_event_data() + if isinstance(event_data, NewSessionUserInputEventData): + return replace(current, metadata=dict(event_data.metadata)) + return None + + +def wait_for_new_session( + window: IClientWindow, current_session_desc: SessionDesc +) -> SessionDesc: + """Wait until ``window`` requests a new session. + + This polls the window at its configured UI rate and converts the latest + buffered new-session event into a complete session description. Other + events, including a browser close, are ignored: no session exists to receive + them, and a persistent server must remain ready for a refreshed page. + + Args: + window: Open client window supplying buffered input events. + current_session_desc: Resolved settings to preserve for the next session. + + Returns: + A complete description for the requested session. + """ + tick_seconds = 1.0 / current_session_desc.frames_per_second_for_ui + while True: + next_session_desc = _next_session_desc( + window.get_user_input_events(), current_session_desc + ) + if next_session_desc is not None: + return next_session_desc + time.sleep(tick_seconds) + + def _close_session(session: ISession, *, run_failed: bool) -> None: """Close a session, keeping its close from hiding an earlier failure. @@ -69,24 +117,28 @@ def run_session( steps: int | None = None, max_pending: int = 2, when_full: WhenFull = WhenFull.BLOCK, -) -> None: + keep_window_open: bool = False, +) -> SessionDesc | None: """Drive one session against one client window. Runs on two threads. The calling thread initializes the session and calls ``step`` for each index, with the input collected since the previous step. A second thread owns the window: it opens it, ticks at ``frames_per_second_for_ui`` to read input, call ``step_ui`` and write - whatever generation has finished, then closes it. A slow step therefore does - not hold up input or output. Only the I/O thread touches the window, which is - what a native window needs, and the window and session are always closed, - including on failure. + whatever generation has finished, then closes it unless a replacement was + requested or ``keep_window_open`` is set. A slow step therefore does not hold + up input or output. Only the I/O thread touches the window, which is what a + native window needs. The session is always closed, and a failure always + closes the window. The window ends the run by reporting a :class:`CloseUserInputEventData`, and restarts it by reporting a :class:`ResetUserInputEventData`, which resets the - session and takes the step index back to zero. The window stays open. Nothing - from the abandoned generation is presented: each result carries the generation - it was produced for, so results already waiting and a step that was still - running when the reset arrived are both dropped rather than written. + session and takes the step index back to zero. A + :class:`NewSessionUserInputEventData` instead ends and cleans up this session, + leaves the window open, and becomes the complete session description returned + for its replacement. Nothing from an abandoned generation is presented: each + result carries the generation it was produced for, so results already waiting + and a step that was still running when the request arrived are both dropped. Input is not split at a reset: the batch carrying it reaches the first step afterwards whole, earlier events included. Events are edges, so a key held @@ -120,6 +172,13 @@ def run_session( max_pending: How many finished results may wait to be written. when_full: What to do with a result when ``max_pending`` are already waiting. + keep_window_open: Leave the window available after normal completion or + a client close, so an application runner can wait for another session. + + Returns: + The complete description requested for the next session, or ``None`` + when the run ended without one. The window remains open when a + description is returned or ``keep_window_open`` is set. Raises: ValueError: ``steps`` is negative, or ``max_pending`` is not positive. @@ -157,44 +216,56 @@ def run_session( collected_events_lock = threading.Lock() opened = threading.Event() stop = threading.Event() + io_stopped = threading.Event() + finish_io = threading.Event() + window_stays_open = False io_failure: list[Exception] = [] + next_session_desc: SessionDesc | None = None # What never reached the window, reported once the run is over. dropped_for_space = 0 - discarded_at_reset = 0 + discarded_at_restart = 0 def present_pending_results() -> None: """Write every waiting result to the window, oldest first. Because each tick writes all of them, results only pile up when writing itself is slower than generation, not merely because the UI rate is lower. - Results the client reset away from are dropped here rather than written, - which is also what frees the room they were holding. + Results the client reset or replaced away from are dropped here rather + than written, which is also what frees the room they were holding. """ - nonlocal discarded_at_reset + nonlocal discarded_at_restart while True: try: result_generation, result = pending_results.get_nowait() except queue.Empty: return if result_generation != generation: - discarded_at_reset += 1 + discarded_at_restart += 1 continue window.write(result) def tick() -> None: - nonlocal generation + nonlocal generation, next_session_desc events = window.get_user_input_events() + requested_session_desc = _next_session_desc(events, session.session_desc) with collected_events_lock: collected_events.extend(events.get_events()) # Move on to the next generation from here, since this thread sees the - # reset first. Under the lock, so a step already picking up its input - # either belongs to the generation being abandoned or to the new one, - # never to neither. - if _contains(events, ResetUserInputEventData): + # reset or replacement first. Under the lock, so a step already + # picking up its input either belongs to the generation being + # abandoned or to the new one, never to neither. + if ( + _contains(events, ResetUserInputEventData) + or requested_session_desc is not None + ): generation += 1 # Stop from here rather than waiting for the step loop to notice, so a # slow step does not delay a client that has gone away. if _contains(events, CloseUserInputEventData): + next_session_desc = None + stop.set() + elif requested_session_desc is not None: + next_session_desc = requested_session_desc stop.set() session.step_ui(events) present_pending_results() @@ -217,14 +288,17 @@ def run_io() -> None: io_failure.append(error) finally: opened.set() - try: - window.close() - except Exception as error: - # Closing is where a sink finishes the writes it was holding, so - # swallowing this would report a run as complete when the output - # never landed. An open that raised part way through gets closed - # here too, since it still holds whatever it had acquired. - io_failure.append(error) + io_stopped.set() + finish_io.wait() + if not window_stays_open: + try: + window.close() + except Exception as error: + # Closing is where a sink finishes the writes it was holding, + # so swallowing this would report a run as complete when the + # output never landed. An open that raised part way through + # gets closed here too, since it may hold acquired resources. + io_failure.append(error) def take_collected_events() -> tuple[UserInputEvents, int]: """Take the input waiting for the next step, and the generation it is for.""" @@ -301,29 +375,59 @@ def add_pending_result(result_generation: int, result: StepResult) -> int: steps_run += 1 finally: stop.set() + io_stopped.wait() + # Another session can own the same window, but only after this one has + # released everything it holds. If cleanup fails, tell the I/O thread to + # close the window and report the failure instead of preserving it. + run_failed = sys.exc_info()[0] is not None + session_close_attempted = False + session_close_error: Exception | None = None + if ( + (next_session_desc is not None or keep_window_open) + and not run_failed + and not io_failure + ): + session_close_attempted = True + try: + session.close() + except Exception as error: + session_close_error = error + else: + window_stays_open = True + finish_io.set() io_thread.join() # A failure here is what the run reports, since a window failure stops # generation rather than raising through it: both places this thread can # be sitting give up once io_failure is set, so a run that reports a # window failure got there without failing itself. The two are only ever # both set by failing independently, and then this is the one raised. - run_failed = sys.exc_info()[0] is not None if io_failure and run_failed: _LOGGER.error( "The window failed as well as the run, and this is that failure.", exc_info=io_failure[0], ) - _close_session(session, run_failed=run_failed or bool(io_failure)) + if session_close_error is not None and io_failure: + _LOGGER.error( + "The window failed as well as the session cleanup, and this is " + "that failure.", + exc_info=io_failure[0], + ) + if not session_close_attempted: + _close_session(session, run_failed=run_failed or bool(io_failure)) + if session_close_error is not None: + raise session_close_error # A log line is the only report of these: a caller cannot count them. if dropped_for_space: _LOGGER.warning( "Dropped %d results the window could not keep up with.", dropped_for_space ) - if discarded_at_reset: + if discarded_at_restart: _LOGGER.info( - "Discarded %d results generated before a reset.", discarded_at_reset + "Discarded %d results generated before a reset or replacement.", + discarded_at_restart, ) if io_failure: raise io_failure[0] + return next_session_desc diff --git a/flashdreams/flashdreams/runtime_v2/user_input_event.py b/flashdreams/flashdreams/runtime_v2/user_input_event.py index fdbd4edad..dbb894b32 100644 --- a/flashdreams/flashdreams/runtime_v2/user_input_event.py +++ b/flashdreams/flashdreams/runtime_v2/user_input_event.py @@ -4,7 +4,7 @@ """User input events, each a timestamp plus the data for one input modality.""" from dataclasses import dataclass -from typing import ClassVar +from typing import Any from numpy import uint64 @@ -67,6 +67,24 @@ def get_type_name(cls) -> str: return "reset" +@dataclass(frozen=True, slots=True, eq=False) +class NewSessionUserInputEventData(UserInputEventData): + """The client asked the application to replace the current session. + + The runtime carries ``metadata`` to the application without interpreting + it. The application owns the requirements for its sessions and validates + the values when it creates the replacement. + """ + + metadata: dict[str, Any] + """Application-specific values requested for the replacement session.""" + + @classmethod + def get_type_name(cls) -> str: + """Return the event type name.""" + return "new_session" + + # Below are stubbed input event data implementations for the sake of future implementation. @dataclass(frozen=True, slots=True, eq=False) class MouseUserInputEventData(UserInputEventData): diff --git a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py index 28936d1fc..5761e7059 100644 --- a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py +++ b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py @@ -47,7 +47,7 @@ def handle_input(event: UserInputEvent) -> None: self.server.register_input_callback(handle_input) def open(self, session_desc: SessionDesc) -> None: - """Implement ``OutputSink.open`` by configuring WebRTC output. + """Configure WebRTC output for waiting or running a session. Args: session_desc: Resolved dimensions, frame rate, and tensor layout. diff --git a/flashdreams/flashdreams/t2v_v2/application.py b/flashdreams/flashdreams/t2v_v2/application.py index 26f24e483..51161b52e 100644 --- a/flashdreams/flashdreams/t2v_v2/application.py +++ b/flashdreams/flashdreams/t2v_v2/application.py @@ -23,8 +23,8 @@ class T2VSessionConfig: """What one command line resolved to, shared by every session it creates.""" - prompt: str - """Text every session generates from.""" + prompt: str | None + """Default text for a session, when its request does not provide one.""" device: str """Device the pipeline is built on.""" @@ -65,21 +65,26 @@ def pipeline_config(self) -> Any: return self._pipeline_config def init(self, commandline_args: Sequence[str]) -> None: - """Parse what to generate, how much of it, and where. + """Parse the default prompt, rollout length, and device. Not what size or rate to generate at: that describes the session, which the caller asks for. The model is not loaded here either. Raises: - ValueError: No prompt was given, or the rollout length is not one - this model can generate. + ValueError: An explicitly provided prompt is empty, or the rollout + length is not one this model can generate. """ parser = argparse.ArgumentParser( prog="flashdreams-run-v2 SLUG --", description="Generate video from text.", ) parser.add_argument( - "--prompt", default="", help="Text to generate from. Required." + "--prompt", + default=None, + help=( + "Default text to generate from. A client may instead provide " + "a prompt for each session." + ), ) parser.add_argument( "--device", @@ -117,8 +122,8 @@ def init(self, commandline_args: Sequence[str]) -> None: self._configure_argument_parser(parser) args = parser.parse_args(list(commandline_args)) - if not args.prompt.strip(): - raise ValueError("--prompt is required, and cannot be empty.") + if args.prompt is not None and not args.prompt.strip(): + raise ValueError("--prompt cannot be empty.") self._validate_total_blocks(args.total_blocks) self._apply_parsed_arguments(args) @@ -155,7 +160,8 @@ def create_session(self, session_desc: SessionDesc) -> ISession: Raises: RuntimeError: :meth:`init` has not run yet. - ValueError: The description asks for output this cannot generate. + ValueError: The description asks for output this cannot generate, + or its prompt metadata is not a non-empty string. """ config = self._config if config is None: @@ -165,11 +171,17 @@ def create_session(self, session_desc: SessionDesc) -> ISession: # Before loading rather than after: a checkpoint of several gigabytes is # a long wait for a layout this was never going to accept. self._validate_layout(session_desc) + prompt = session_desc.metadata.get("prompt", config.prompt) + if not isinstance(prompt, str) or not prompt.strip(): + raise ValueError( + "A session prompt is required: pass --prompt or set " + "SessionDesc.metadata['prompt'] to a non-empty string." + ) if self._pipeline is None: self._pipeline = self._pipeline_config.setup().to(config.device).eval() self._validate_frame_size(session_desc, self._pipeline) return self.session_type( - self._pipeline, config.prompt, session_desc, config.total_blocks + self._pipeline, prompt, session_desc, config.total_blocks ) def close(self) -> None: diff --git a/flashdreams/test_v2/test_application_runner.py b/flashdreams/test_v2/test_application_runner.py index 8b872b1f1..8ae2afa0d 100644 --- a/flashdreams/test_v2/test_application_runner.py +++ b/flashdreams/test_v2/test_application_runner.py @@ -4,6 +4,7 @@ """CPU tests for the v2 application runner.""" import logging +import threading from collections.abc import Sequence import pytest @@ -18,6 +19,7 @@ from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import ( CloseUserInputEventData, + NewSessionUserInputEventData, UserInputEvent, ) from flashdreams.runtime_v2.user_input_events import UserInputEvents @@ -82,6 +84,7 @@ def __init__( self._fail_to_init = fail_to_init self._fail_to_close = fail_to_close self._session_length = session_length + self.created_session_descs: list[SessionDesc] = [] def init(self, commandline_args: Sequence[str]) -> None: self._calls.append(f"application.init({list(commandline_args)!r})") @@ -90,6 +93,7 @@ def init(self, commandline_args: Sequence[str]) -> None: def create_session(self, session_desc: SessionDesc) -> ISession: self._calls.append("application.create_session") + self.created_session_descs.append(session_desc) return _Session(session_desc, self._calls, length=self._session_length) def close(self) -> None: @@ -136,6 +140,44 @@ def get_user_input_events(self) -> UserInputEvents: return UserInputEvents([]) +class _ScriptedWindow(_Window): + """Report one scripted event batch each time the runner polls.""" + + def __init__(self, calls: list[str], events: list[UserInputEvents]) -> None: + super().__init__(calls) + self._events = list(events) + + def get_user_input_events(self) -> UserInputEvents: + if self._events: + return self._events.pop(0) + return UserInputEvents([]) + + +class _ServingWindow(_Window): + """Request two sessions, then interrupt the persistent runner.""" + + def __init__(self, calls: list[str]) -> None: + super().__init__(calls) + self._prompts = ["A cat surfing", "A dog snowboarding"] + + def get_user_input_events(self) -> UserInputEvents: + if threading.current_thread() is not threading.main_thread(): + return UserInputEvents([]) + completed_sessions = self._calls.count("session.close") + if completed_sessions == len(self._prompts): + raise KeyboardInterrupt + return UserInputEvents( + [ + UserInputEvent( + timestamp=uint64(completed_sessions), + event_data=NewSessionUserInputEventData( + metadata={"prompt": self._prompts[completed_sessions]} + ), + ) + ] + ) + + def _session_desc_request() -> SessionDescRequest: return SessionDescRequest( output_layout=VideoTensorLayout.bcthw, @@ -170,6 +212,69 @@ def test_application_runner_keeps_the_application_open_for_another_session() -> assert calls[-1] == "application.close" +def test_application_runner_replaces_a_session_from_window_metadata() -> None: + calls: list[str] = [] + application = _Application(calls) + runner = ApplicationRunner(application) + new_session = UserInputEvents( + [ + UserInputEvent( + timestamp=uint64(0), + event_data=NewSessionUserInputEventData( + metadata={"prompt": "A dog snowboarding"} + ), + ) + ] + ) + close = UserInputEvents( + [ + UserInputEvent( + timestamp=uint64(1), + event_data=CloseUserInputEventData(), + ) + ] + ) + window = _ScriptedWindow(calls, [new_session, close]) + + runner.init() + runner.run_session(_session_desc_request(), window) + runner.close() + + assert len(application.created_session_descs) == 2 + assert application.created_session_descs[0].metadata == {} + assert application.created_session_descs[1].metadata == { + "prompt": "A dog snowboarding" + } + assert calls.count("session.close") == 2 + assert calls.count("window.open") == 2 + assert calls.count("window.close") == 1 + first_creation = calls.index("application.create_session") + second_creation = calls.index("application.create_session", first_creation + 1) + assert calls.index("session.close") < second_creation + + +def test_application_runner_serves_sessions_until_it_is_interrupted() -> None: + calls: list[str] = [] + application = _Application(calls, session_length=1) + runner = ApplicationRunner(application) + window = _ServingWindow(calls) + + runner.init() + with pytest.raises(KeyboardInterrupt): + runner.run_session(_session_desc_request(), window, serve_sessions=True) + + assert [desc.metadata for desc in application.created_session_descs] == [ + {"prompt": "A cat surfing"}, + {"prompt": "A dog snowboarding"}, + ] + assert calls.count("session.close") == 2 + assert calls.count("window.open") == 3 + assert calls.count("window.close") == 1 + assert calls.count("application.close") == 0 + runner.close() + assert calls[-1] == "application.close" + + def test_application_runner_closes_the_window_when_a_session_cannot_start() -> None: calls: list[str] = [] runner = ApplicationRunner(_Application(calls)) diff --git a/flashdreams/test_v2/test_cli.py b/flashdreams/test_v2/test_cli.py index 8cd62499e..8a0e25874 100644 --- a/flashdreams/test_v2/test_cli.py +++ b/flashdreams/test_v2/test_cli.py @@ -277,9 +277,16 @@ def _write_application_module( class StubMode(ClientWindowMode): """A mode handing the command a window the test can look inside.""" - def __init__(self, name: str, window: IClientWindow) -> None: + def __init__( + self, + name: str, + window: IClientWindow, + *, + serves_sessions: bool = False, + ) -> None: self.name = name self._window = window + self.serves_sessions = serves_sessions def create(self, parsed_args: argparse.Namespace) -> IClientWindow: del parsed_args @@ -400,15 +407,44 @@ def test_nothing_is_measured_unless_a_run_asks( assert list(tmp_path.glob("*.json")) == [] -def test_an_application_that_will_not_start_reports_why( +def test_a_one_shot_run_without_a_prompt_reports_why( monkeypatch: pytest.MonkeyPatch, ) -> None: _install(monkeypatch, StubT2VApplication(_stand_in()), RecordingWindow()) - with pytest.raises(ValueError, match="--prompt is required"): + with pytest.raises(ValueError, match="session prompt is required"): cli.entrypoint(["stub", "--mode", "webrtc"]) +def test_a_browser_server_starts_without_a_command_line_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class InterruptingWindow(RecordingWindow): + def __init__(self) -> None: + super().__init__() + self.closed = False + + def get_user_input_events(self) -> UserInputEvents: + raise KeyboardInterrupt + + def close(self) -> None: + self.closed = True + + pipeline = _stand_in() + window = InterruptingWindow() + _install(monkeypatch, StubT2VApplication(pipeline)) + monkeypatch.setattr( + cli, + "client_window_mode", + lambda name: StubMode(name, window, serves_sessions=True), + ) + + cli.entrypoint(["stub", "--mode", "webrtc"]) + + assert window.closed + assert pipeline.device is None + + ## Describing the session to run diff --git a/flashdreams/test_v2/test_client_window_factory.py b/flashdreams/test_v2/test_client_window_factory.py index b7e675cbb..c77466d7e 100644 --- a/flashdreams/test_v2/test_client_window_factory.py +++ b/flashdreams/test_v2/test_client_window_factory.py @@ -36,6 +36,7 @@ def test_a_run_goes_to_a_file_unless_it_says_otherwise(tmp_path: Path) -> None: assert parsed.mode == "mp4" assert isinstance(window, Mp4ClientWindow) + assert client_window_mode("mp4").serves_sessions is False def test_a_file_run_with_nowhere_to_write_says_so() -> None: @@ -74,5 +75,6 @@ def test_a_browser_run_is_told_where_to_connect(self) -> None: try: assert isinstance(window, WebRTCClientWindow) assert mode.starting(window) == f"Open {window.server.url} in a browser." + assert mode.serves_sessions is True finally: window.close() diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py index 7928aa412..d0d8faca4 100644 --- a/flashdreams/test_v2/test_session_runner.py +++ b/flashdreams/test_v2/test_session_runner.py @@ -5,6 +5,7 @@ import logging import threading +from dataclasses import replace import pytest import torch @@ -19,6 +20,7 @@ from flashdreams.runtime_v2.user_input_event import ( CloseUserInputEventData, KeyboardUserInputEventData, + NewSessionUserInputEventData, ResetUserInputEventData, UserInputEvent, ) @@ -326,6 +328,39 @@ def test_run_session_stops_when_the_window_reports_a_close() -> None: assert log.calls[-2:] == ["window.close", "session.close"] +def test_run_session_returns_the_next_session_desc_after_cleanup() -> None: + log = CallLog() + session_desc = _session_desc() + session = FakeSession(session_desc, log) + request = NewSessionUserInputEventData(metadata={"prompt": "A cat surfing"}) + window = RecordingClientWindow(log, [_lifecycle_event(request)]) + + returned = run_session(session, window, steps=None) + + assert returned == replace(session_desc, metadata={"prompt": "A cat surfing"}) + assert "session.close" in log.calls + assert "window.close" not in log.calls + assert "session.step(0)" not in log.calls + + +def test_run_session_closes_the_window_when_replacement_cleanup_fails() -> None: + log = CallLog() + session = FakeSession(_session_desc(), log, fail_to_close=True) + window = RecordingClientWindow( + log, + [ + _lifecycle_event( + NewSessionUserInputEventData(metadata={"prompt": "A cat surfing"}) + ) + ], + ) + + with pytest.raises(RuntimeError, match="session close failed"): + run_session(session, window, steps=None) + + assert log.calls[-2:] == ["session.close", "window.close"] + + def test_run_session_resets_the_session_and_the_step_index() -> None: log = CallLog() session = FakeSession(_session_desc(), log) @@ -354,6 +389,18 @@ def test_run_session_stops_when_the_session_says_it_has_finished() -> None: assert [result.step_index for result in window.results] == [0, 1] +def test_run_session_can_leave_the_window_open_after_completion() -> None: + log = CallLog() + session = FiniteSession(_session_desc(), log, length=1) + window = RecordingClientWindow(log) + + returned = run_session(session, window, keep_window_open=True) + + assert returned is None + assert "session.close" in log.calls + assert "window.close" not in log.calls + + def test_run_session_ends_at_whichever_comes_first() -> None: """A caller can ask for fewer steps than the session would generate.""" log = CallLog() diff --git a/flashdreams/test_v2/test_t2v_application.py b/flashdreams/test_v2/test_t2v_application.py index 9984087d9..53b275d8b 100644 --- a/flashdreams/test_v2/test_t2v_application.py +++ b/flashdreams/test_v2/test_t2v_application.py @@ -8,6 +8,7 @@ """ import argparse +from dataclasses import replace from types import SimpleNamespace from typing import Any @@ -88,6 +89,7 @@ def __init__( self, pipeline: Any, prompt: str, session_desc: SessionDesc, total_blocks: int ) -> None: super().__init__(pipeline, prompt, session_desc, total_blocks) + self.prompt = prompt self.blocks_to_generate = total_blocks @@ -240,12 +242,18 @@ def test_the_rollout_length_can_be_overridden() -> None: assert _rollout_length(app) == 3 -def test_a_run_needs_something_to_generate_from() -> None: +def test_an_application_can_initialize_before_a_session_has_a_prompt() -> None: app = T2VApplication(defaults=_defaults()) - with pytest.raises(ValueError, match="--prompt is required"): - app.init([]) - with pytest.raises(ValueError, match="--prompt is required"): + app.init([]) + with pytest.raises(ValueError, match="session prompt is required"): + app.create_session(_session_desc()) + + +def test_an_explicit_command_line_prompt_cannot_be_empty() -> None: + app = T2VApplication(defaults=_defaults()) + + with pytest.raises(ValueError, match="--prompt cannot be empty"): app.init(["--prompt", " "]) @@ -277,6 +285,27 @@ def test_the_model_loads_once_and_every_session_shares_it() -> None: assert first is not second +def test_session_metadata_can_replace_the_command_line_prompt() -> None: + app = _application() + + session = app.create_session( + replace(_session_desc(), metadata={"prompt": "A dog snowboarding"}) + ) + + assert isinstance(session, RecordingRollout) + assert session.prompt == "A dog snowboarding" + + +@pytest.mark.parametrize("prompt", ["", 7]) +def test_a_replacement_prompt_must_be_non_empty_text(prompt: object) -> None: + app = _application() + + with pytest.raises(ValueError, match="session prompt is required"): + app.create_session(replace(_session_desc(), metadata={"prompt": prompt})) + + assert _pipeline_config(app).setup_count == 0 + + def test_closing_the_application_releases_the_model() -> None: app = _application() config = _pipeline_config(app) diff --git a/flashdreams/test_v2/test_webrtc_client_window.py b/flashdreams/test_v2/test_webrtc_client_window.py index b7f119651..6e46a23ef 100644 --- a/flashdreams/test_v2/test_webrtc_client_window.py +++ b/flashdreams/test_v2/test_webrtc_client_window.py @@ -5,6 +5,7 @@ import asyncio import json +from dataclasses import replace import pytest import torch @@ -25,7 +26,10 @@ from flashdreams.runtime_v2.session_desc import SessionDesc from flashdreams.runtime_v2.step_result import StepResult -from flashdreams.runtime_v2.user_input_event import KeyboardUserInputEventData +from flashdreams.runtime_v2.user_input_event import ( + KeyboardUserInputEventData, + NewSessionUserInputEventData, +) from flashdreams.runtime_v2.video_tensor import VideoTensorLayout from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow @@ -94,25 +98,38 @@ async def test_window_buffers_browser_events_until_drained() -> None: browser_page = await response.text() assert response.status == 200 assert 'id="activate"' in browser_page + assert 'id="prompt"' in browser_page + assert 'id="new-session" type="button" disabled' in browser_page assert '' in browser_page async with client.get(f"{window.server.url}app.js") as response: browser_script = await response.text() assert response.status == 200 assert 'key: "r", pressed: activationPressed' in browser_script + assert 'type: "new_session"' in browser_script + assert "metadata: {prompt: promptInput.value}" in browser_script + assert "response.status !== 409" in browser_script window.open(_session_desc()) peer, channel, _ = await _connect_browser(window) channel.send(json.dumps({"type": "keyboard", "key": "w", "pressed": True})) channel.send(json.dumps({"type": "keyboard", "key": "w", "pressed": False})) + channel.send( + json.dumps( + { + "type": "new_session", + "metadata": {"prompt": "A dog snowboarding"}, + } + ) + ) events = [] for _ in range(100): events.extend(window.get_user_input_events().get_events()) - if len(events) == 2: + if len(events) == 3: break await asyncio.sleep(0.01) - assert len(events) == 2 + assert len(events) == 3 keyboard_events = [ data for event in events @@ -122,8 +139,51 @@ async def test_window_buffers_browser_events_until_drained() -> None: ("w", True), ("w", False), ] + new_session_events = [ + data + for event in events + if isinstance(data := event.get_event_data(), NewSessionUserInputEventData) + ] + assert [event.metadata for event in new_session_events] == [ + {"prompt": "A dog snowboarding"} + ] assert events[0].get_timestamp() <= events[1].get_timestamp() assert window.get_user_input_events().get_events() == [] + + await peer.close() + peer = None + async with ClientSession() as client: + for _ in range(100): + async with client.get(f"{window.server.url}healthz") as response: + health = await response.json() + if not health["client_connected"]: + break + await asyncio.sleep(0.01) + assert health == {"open": True, "client_connected": False} + + peer, channel, _ = await _connect_browser(window) + channel.send( + json.dumps( + { + "type": "new_session", + "metadata": {"prompt": "A fox in a forest"}, + } + ) + ) + refreshed_events = [] + for _ in range(100): + refreshed_events.extend(window.get_user_input_events().get_events()) + if any( + isinstance(event.get_event_data(), NewSessionUserInputEventData) + for event in refreshed_events + ): + break + await asyncio.sleep(0.01) + assert [ + event.get_event_data().metadata + for event in refreshed_events + if isinstance(event.get_event_data(), NewSessionUserInputEventData) + ] == [{"prompt": "A fox in a forest"}] finally: if peer is not None: await peer.close() @@ -154,6 +214,18 @@ async def test_write_delivers_a_video_frame_to_the_browser() -> None: pixels = frame.to_ndarray(format="rgb24") assert pixels.shape == (16, 16, 3) assert abs(float(pixels.mean()) - 17.0) <= 2.0 + + window.open(replace(_session_desc(), metadata={"prompt": "replacement"})) + window.write( + StepResult( + step_index=0, + output=torch.full((2, 3, 16, 16), 211, dtype=torch.uint8), + frame_count=2, + output_layout=VideoTensorLayout.tchw, + metrics={}, + ) + ) + assert peer.connectionState == "connected" finally: if peer is not None: await peer.close() diff --git a/integrations_v2/t2v_self_forcing/README.md b/integrations_v2/t2v_self_forcing/README.md index 84afb0ae6..e62676758 100644 --- a/integrations_v2/t2v_self_forcing/README.md +++ b/integrations_v2/t2v_self_forcing/README.md @@ -20,6 +20,19 @@ flashdreams-run-v2 t2v-self-forcing --output-path clip.mp4 \ Arguments after `--` go to the application, and `flashdreams-run-v2 t2v-self-forcing -- --help` lists them. +To stream one clip to a browser instead of writing an MP4: + +```bash +uv run --project integrations_v2/t2v_self_forcing flashdreams-run-v2 \ + t2v-self-forcing --mode webrtc +``` + +Open the printed URL, enter a prompt, and select **New session**. The server +cleans up each completed rollout and waits for another prompt without reloading +the model. Closing or refreshing the page ends the active rollout but not the +server; the browser can reconnect and request another session. Stop it with +Ctrl-C. + `--total-blocks` is how many autoregressive blocks to generate, and the run ends when the session has generated them. The first block decodes 9 frames and every block after it 12, at 16 frames per second, so seven blocks is about four and a From 91ca5597e275dc8b7e9779c90161a38c2070ee4a Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 21 Aug 2026 07:00:05 +0000 Subject: [PATCH 3/6] Keep WebRTC session requests ready while connecting --- .../flashdreams/runtime_v2/serving/web/app.js | 28 +++++++++++++------ .../runtime_v2/serving/web/index.html | 2 +- .../test_v2/test_webrtc_client_window.py | 4 ++- 3 files changed, 24 insertions(+), 10 deletions(-) diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/app.js b/flashdreams/flashdreams/runtime_v2/serving/web/app.js index d64248929..9e25d6542 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/web/app.js +++ b/flashdreams/flashdreams/runtime_v2/serving/web/app.js @@ -5,9 +5,21 @@ const peer = new RTCPeerConnection(); const controls = peer.createDataChannel("controls"); peer.addTransceiver("video", {direction: "recvonly"}); const newSessionButton = document.getElementById("new-session"); +let pendingNewSession = null; + +const send = payload => { + if (controls.readyState === "open") { + controls.send(JSON.stringify(payload)); + } +}; controls.addEventListener("open", () => { newSessionButton.disabled = false; + if (pendingNewSession !== null) { + send(pendingNewSession); + pendingNewSession = null; + } + newSessionButton.textContent = "New session"; }); controls.addEventListener("close", () => { @@ -19,12 +31,6 @@ peer.ontrack = event => { event.streams[0] ?? new MediaStream([event.track]); }; -const send = payload => { - if (controls.readyState === "open") { - controls.send(JSON.stringify(payload)); - } -}; - window.addEventListener("keydown", event => { send({type: "keyboard", key: event.key, pressed: true}); }); @@ -50,10 +56,16 @@ newSessionButton.onclick = () => { if (!promptInput.reportValidity()) { return; } - send({ + const request = { type: "new_session", metadata: {prompt: promptInput.value}, - }); + }; + if (controls.readyState === "open") { + send(request); + } else { + pendingNewSession = request; + newSessionButton.textContent = "Opening..."; + } }; window.addEventListener("beforeunload", () => send({type: "close"})); diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/index.html b/flashdreams/flashdreams/runtime_v2/serving/web/index.html index eb8b145b6..0cd2798a8 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/web/index.html +++ b/flashdreams/flashdreams/runtime_v2/serving/web/index.html @@ -13,7 +13,7 @@ - + diff --git a/flashdreams/test_v2/test_webrtc_client_window.py b/flashdreams/test_v2/test_webrtc_client_window.py index 6e46a23ef..f5e2bd632 100644 --- a/flashdreams/test_v2/test_webrtc_client_window.py +++ b/flashdreams/test_v2/test_webrtc_client_window.py @@ -99,7 +99,7 @@ async def test_window_buffers_browser_events_until_drained() -> None: assert response.status == 200 assert 'id="activate"' in browser_page assert 'id="prompt"' in browser_page - assert 'id="new-session" type="button" disabled' in browser_page + assert 'id="new-session" type="button">' in browser_page assert '' in browser_page async with client.get(f"{window.server.url}app.js") as response: browser_script = await response.text() @@ -107,6 +107,8 @@ async def test_window_buffers_browser_events_until_drained() -> None: assert 'key: "r", pressed: activationPressed' in browser_script assert 'type: "new_session"' in browser_script assert "metadata: {prompt: promptInput.value}" in browser_script + assert "pendingNewSession = request" in browser_script + assert 'newSessionButton.textContent = "Opening..."' in browser_script assert "response.status !== 409" in browser_script window.open(_session_desc()) From 779147279224ebb8b70884fb2ffc379c85e60cbe Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 21 Aug 2026 07:00:05 +0000 Subject: [PATCH 4/6] Preload text-to-video models during application init --- flashdreams/flashdreams/t2v_v2/application.py | 30 +++++++++---------- flashdreams/test_v2/test_cli.py | 4 +-- flashdreams/test_v2/test_t2v_application.py | 29 +++++++++++------- .../tests/test_stand_in_model.py | 12 ++++++-- .../tests/test_stand_in_model.py | 13 +++++--- 5 files changed, 54 insertions(+), 34 deletions(-) diff --git a/flashdreams/flashdreams/t2v_v2/application.py b/flashdreams/flashdreams/t2v_v2/application.py index 51161b52e..c9d3e143d 100644 --- a/flashdreams/flashdreams/t2v_v2/application.py +++ b/flashdreams/flashdreams/t2v_v2/application.py @@ -41,8 +41,9 @@ class T2VApplication(IApplication): command line for all of them. An integration supplies :class:`T2VApplicationDefaults` and inherits the rest. - The model is loaded once, on the first session, and shared by every session - after it, since loading reads a checkpoint of several gigabytes. + :meth:`init` loads the model once after resolving its command-line options. + The application then keeps that model resident and shares it with every + session, since loading reads a checkpoint of several gigabytes. """ session_type: type[T2VSession] = T2VSession @@ -65,10 +66,11 @@ def pipeline_config(self) -> Any: return self._pipeline_config def init(self, commandline_args: Sequence[str]) -> None: - """Parse the default prompt, rollout length, and device. + """Parse application options and load the shared model. Not what size or rate to generate at: that describes the session, which - the caller asks for. The model is not loaded here either. + the caller asks for. Loading happens after device, compilation, seed, + and integration-specific options have been resolved. Raises: ValueError: An explicitly provided prompt is empty, or the rollout @@ -135,11 +137,14 @@ def init(self, commandline_args: Sequence[str]) -> None: self._pipeline_config = self._apply_seed_override( self._pipeline_config, args.seed ) - self._config = T2VSessionConfig( + config = T2VSessionConfig( prompt=args.prompt, device=args.device, total_blocks=args.total_blocks, ) + pipeline = self._pipeline_config.setup().to(config.device).eval() + self._config = config + self._pipeline = pipeline def default_session_desc(self) -> SessionDesc: """Return the description of a session this application uses. @@ -156,7 +161,7 @@ def default_session_desc(self) -> SessionDesc: ) def create_session(self, session_desc: SessionDesc) -> ISession: - """Create one uninitialized session, loading the model if needed. + """Create one uninitialized session against the resident model. Raises: RuntimeError: :meth:`init` has not run yet. @@ -164,12 +169,11 @@ def create_session(self, session_desc: SessionDesc) -> ISession: or its prompt metadata is not a non-empty string. """ config = self._config - if config is None: + pipeline = self._pipeline + if config is None or pipeline is None: raise RuntimeError( f"{type(self).__name__}.init() must run before create_session()." ) - # Before loading rather than after: a checkpoint of several gigabytes is - # a long wait for a layout this was never going to accept. self._validate_layout(session_desc) prompt = session_desc.metadata.get("prompt", config.prompt) if not isinstance(prompt, str) or not prompt.strip(): @@ -177,12 +181,8 @@ def create_session(self, session_desc: SessionDesc) -> ISession: "A session prompt is required: pass --prompt or set " "SessionDesc.metadata['prompt'] to a non-empty string." ) - if self._pipeline is None: - self._pipeline = self._pipeline_config.setup().to(config.device).eval() - self._validate_frame_size(session_desc, self._pipeline) - return self.session_type( - self._pipeline, prompt, session_desc, config.total_blocks - ) + self._validate_frame_size(session_desc, pipeline) + return self.session_type(pipeline, prompt, session_desc, config.total_blocks) def close(self) -> None: """Release the model, and whatever memory it was holding.""" diff --git a/flashdreams/test_v2/test_cli.py b/flashdreams/test_v2/test_cli.py index 8a0e25874..5605e7ddb 100644 --- a/flashdreams/test_v2/test_cli.py +++ b/flashdreams/test_v2/test_cli.py @@ -13,7 +13,6 @@ import shutil from collections.abc import Sequence from pathlib import Path -from typing import Any import pytest import torch @@ -442,7 +441,8 @@ def close(self) -> None: cli.entrypoint(["stub", "--mode", "webrtc"]) assert window.closed - assert pipeline.device is None + assert pipeline.device == "cpu" + assert pipeline.eval_count == 1 ## Describing the session to run diff --git a/flashdreams/test_v2/test_t2v_application.py b/flashdreams/test_v2/test_t2v_application.py index 53b275d8b..7db998bfa 100644 --- a/flashdreams/test_v2/test_t2v_application.py +++ b/flashdreams/test_v2/test_t2v_application.py @@ -242,10 +242,13 @@ def test_the_rollout_length_can_be_overridden() -> None: assert _rollout_length(app) == 3 -def test_an_application_can_initialize_before_a_session_has_a_prompt() -> None: - app = T2VApplication(defaults=_defaults()) +def test_an_application_preloads_before_a_session_has_a_prompt() -> None: + defaults = _defaults() + app = T2VApplication(defaults=defaults) app.init([]) + + assert defaults.pipeline_config.setup_count == 1 with pytest.raises(ValueError, match="session prompt is required"): app.create_session(_session_desc()) @@ -272,10 +275,14 @@ def test_no_session_is_created_before_the_application_is_told_what_to_do() -> No ## Loading the model -def test_the_model_loads_once_and_every_session_shares_it() -> None: - app = _application() +def test_initialization_loads_the_model_once_for_every_session() -> None: + app = ApplicationUnderTest(defaults=_defaults()) config = _pipeline_config(app) + assert config.setup_count == 0 + + app.init(["--prompt", _PROMPT]) + first = app.create_session(_session_desc()) second = app.create_session(_session_desc()) @@ -303,13 +310,12 @@ def test_a_replacement_prompt_must_be_non_empty_text(prompt: object) -> None: with pytest.raises(ValueError, match="session prompt is required"): app.create_session(replace(_session_desc(), metadata={"prompt": prompt})) - assert _pipeline_config(app).setup_count == 0 + assert _pipeline_config(app).setup_count == 1 def test_closing_the_application_releases_the_model() -> None: app = _application() config = _pipeline_config(app) - app.create_session(_session_desc()) app.close() @@ -319,15 +325,14 @@ def test_closing_the_application_releases_the_model() -> None: ## What a model will not generate -def test_a_layout_the_model_does_not_emit_is_refused_before_it_loads() -> None: - """A checkpoint of several gigabytes is a long wait for a certain refusal.""" +def test_a_layout_the_model_does_not_emit_is_refused_before_a_session_starts() -> None: app = _application() config = _pipeline_config(app) with pytest.raises(ValueError, match="only produces tchw output"): app.create_session(_session_desc(VideoTensorLayout.bcthw)) - assert config.setup_count == 0 + assert config.setup_count == 1 @pytest.mark.parametrize("width,height", [(130, 64), (128, 60)]) @@ -437,9 +442,11 @@ def test_a_seed_reaches_the_model_where_a_model_keeps_one() -> None: """Straight onto the config the pipeline is built from, so a model that draws its own noise draws the same noise twice.""" diffusion_model = SimpleNamespace(seed=42) - defaults = _defaults( - pipeline_config=SimpleNamespace(diffusion_model=diffusion_model) + pipeline_config = SimpleNamespace( + diffusion_model=diffusion_model, + setup=lambda: FakePipeline(), ) + defaults = _defaults(pipeline_config=pipeline_config) app = T2VApplication(defaults=defaults) app.init(["--prompt", _PROMPT, "--seed", "7"]) diff --git a/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py b/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py index 75900757b..dea1f679c 100644 --- a/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py +++ b/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py @@ -8,12 +8,14 @@ ``test_real_model.py``. """ +import copy + import pytest from fastvideo_causal_wan22.config import RUNNER_WAN22_T2V_14B from t2v_fastvideo_causal_wan22 import FastvideoCausalWan22T2VApplication from flashdreams.runtime_v2.video_tensor import VideoTensorLayout -from flashdreams.t2v_v2.testing import FakeT2VPipelineConfig +from flashdreams.t2v_v2.testing import FakeT2VPipeline, FakeT2VPipelineConfig pytestmark = pytest.mark.ci_cpu @@ -44,7 +46,13 @@ def test_compilation_is_turned_off_for_both_noise_level_transformers() -> None: model splits denoising across two transformers, and the shared override reaches only one of them, so it is overridden here. """ - app = FastvideoCausalWan22T2VApplication() + pipeline_config = copy.deepcopy(RUNNER_WAN22_T2V_14B.pipeline) + + def load_stand_in(_: object) -> FakeT2VPipeline: + return FakeT2VPipeline() + + pipeline_config._target = load_stand_in + app = FastvideoCausalWan22T2VApplication(pipeline_config=pipeline_config) app.init(["--prompt", _PROMPT, "--no-compile"]) diff --git a/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py b/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py index 0ca793e64..fb5b66630 100644 --- a/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py +++ b/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py @@ -9,6 +9,7 @@ ``test_real_model.py``. """ +import copy import shutil from pathlib import Path @@ -48,10 +49,14 @@ def test_the_model_says_what_it_generates_without_being_told() -> None: def test_compilation_can_be_turned_off_for_a_run() -> None: - """Run against the real config rather than a stand-in, since what this - covers is the override landing where this model keeps the setting. No model - is loaded to answer it.""" - app = SelfForcingT2VApplication() + """Apply the override to the real config while loading a stand-in model.""" + pipeline_config = copy.deepcopy(RUNNER_WAN21_T2V_1PT3B.pipeline) + + def load_stand_in(_: object) -> FakeT2VPipeline: + return FakeT2VPipeline() + + pipeline_config._target = load_stand_in + app = SelfForcingT2VApplication(pipeline_config=pipeline_config) app.init(["--prompt", _PROMPT, "--no-compile"]) From 0c313deb6dcb893a4dc68757b5ca08329231ce7b Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 21 Aug 2026 08:14:56 +0000 Subject: [PATCH 5/6] Simplify and harden v2 WebRTC lifecycle --- docs/source/developer_guides/index.rst | 8 + .../developer_guides/v2_webrtc_lifecycle.md | 153 +++++++++++++ flashdreams/flashdreams/api_v2/application.py | 9 +- flashdreams/flashdreams/api_v2/session.py | 6 +- .../runtime_v2/application_runner.py | 8 +- flashdreams/flashdreams/runtime_v2/cli.py | 16 +- .../runtime_v2/client_window_factory.py | 4 +- .../flashdreams/runtime_v2/serving/web/app.js | 15 +- .../runtime_v2/serving/web/index.html | 2 +- .../runtime_v2/serving/webrtc_server.py | 195 +++++++++++------ .../flashdreams/runtime_v2/session_desc.py | 12 +- .../flashdreams/runtime_v2/session_runner.py | 105 +++++---- .../runtime_v2/user_input_event.py | 2 +- .../runtime_v2/webrtc_client_window.py | 15 +- flashdreams/flashdreams/t2v_v2/application.py | 8 +- .../test_v2/test_application_runner.py | 27 ++- flashdreams/test_v2/test_cli.py | 32 ++- .../test_v2/test_client_window_factory.py | 2 +- flashdreams/test_v2/test_session_runner.py | 181 +++++++++++++++- flashdreams/test_v2/test_t2v_application.py | 22 +- .../test_v2/test_webrtc_client_window.py | 205 ++++++++++++++++-- .../color_fade/tests/test_color_fade.py | 2 +- integrations_v2/red_screen/red_screen/app.py | 4 +- .../tests/test_stand_in_model.py | 3 +- .../tests/test_stand_in_model.py | 3 +- 25 files changed, 856 insertions(+), 183 deletions(-) create mode 100644 docs/source/developer_guides/v2_webrtc_lifecycle.md diff --git a/docs/source/developer_guides/index.rst b/docs/source/developer_guides/index.rst index 8a8874d8b..fe2943ef0 100644 --- a/docs/source/developer_guides/index.rst +++ b/docs/source/developer_guides/index.rst @@ -41,6 +41,13 @@ Developer Guides How public runner names are registered, parsed, matched to manifests, and dispatched to integration-owned demo launch modes. + .. grid-item-card:: V2 WebRTC lifecycle + :link: v2_webrtc_lifecycle + :link-type: doc + + Ownership and data flow across the CLI, application runner, persistent + browser window, resident model, and sequential sessions. + .. grid-item-card:: Add a new method :link: new_integration :link-type: doc @@ -70,6 +77,7 @@ generated clip, see :doc:`/quickstart/index`. inference_pipeline_overview config_system runner_slugs + v2_webrtc_lifecycle new_integration local_benchmarks diff --git a/docs/source/developer_guides/v2_webrtc_lifecycle.md b/docs/source/developer_guides/v2_webrtc_lifecycle.md new file mode 100644 index 000000000..8afa48a18 --- /dev/null +++ b/docs/source/developer_guides/v2_webrtc_lifecycle.md @@ -0,0 +1,153 @@ + + +# V2 WebRTC application and session lifecycle + +The v2 WebRTC path keeps an application and its model alive while it creates +one session at a time from browser requests. A session owns one rollout; it +does not own the server or reload the model. + +The Self-Forcing entry command remains: + +```bash +uv run --project integrations_v2/t2v_self_forcing flashdreams-run-v2 \ + t2v-self-forcing --mode webrtc --port 8080 +``` + +Open the printed URL, enter a prompt, and select **New session**. Closing or +refreshing the page ends the active rollout, but the process and loaded model +remain available for the next request. Ctrl-C ends the application. + +## Ownership + +| Component | Owns | Does not own | +| --- | --- | --- | +| `cli.py` | Argument parsing, application and window construction, final runner cleanup | Model state, session state, browser events | +| `ApplicationRunner` | The initialized application and the serial application-level session loop | Browser event types or per-step scheduling | +| `run_session()` | Exactly one session, its generation loop, its I/O thread, and session cleanup | The application or a persistent window after a successful handoff | +| `WebRTCClientWindow` | The WebRTC server and the thread-safe browser-event queue | Application or session state | +| `WebRTCServer` | HTTP, one active peer connection, its data channel, and its media track | Runtime lifecycle decisions | +| `T2VApplication` | The loaded model pipeline shared by every session | A rollout cache or browser connection | +| `T2VSession` | One prompt and one rollout cache | Model loading or server lifetime | + +The CLI constructs the window, then transfers its cleanup responsibility to +the runner. While the server is idle, the runner's calling thread opens and +polls the window. During a session, `run_session()` gives the window to one +dedicated I/O thread. The runner does not touch it again until that thread has +stopped and the old session has closed. This is a sequential ownership +handoff, not concurrent access. + +## Data flow + +1. `cli.py` selects the `t2v-self-forcing` application and the `webrtc` window + mode. Arguments after `--` belong to the application. +2. `ApplicationRunner.init()` calls `T2VApplication.init()`. The pipeline is + set up, moved to its device, evaluated, and retained by the application. +3. The WebRTC mode creates `WebRTCClientWindow`, which starts its server thread + and exposes the browser URL. +4. `ApplicationRunner.run()` resolves the CLI's partial `SessionDescRequest` + against the initialized application's default `SessionDesc`. +5. In serving mode, the runner opens the window with that resolved stream + format before a session exists. This lets the browser negotiate WebRTC and + send its first request without loading a session cache first. +6. The server validates every data-channel message and invokes the callback + registered by `WebRTCClientWindow`. The callback only appends the event to a + thread-safe queue. +7. While idle, `wait_for_new_session()` drains that queue. It translates the + latest close/new-session transition into a complete `SessionDesc` and + returns only that description. `ApplicationRunner` therefore never parses a + `UserInputEvent`. +8. The application validates the description and creates a `T2VSession` over + its resident pipeline. `T2VSession.init()` creates only the per-rollout + prompt/cache state. +9. `run_session()` starts one I/O thread. That thread opens the same window for + the actual session boundary, drains input each UI tick, calls `step_ui()`, + and writes completed results. The calling thread runs model steps. +10. A new-session event stops that rollout and returns the requested next + `SessionDesc`. A close or natural completion returns no replacement. In + serving mode the runner keeps the window open, closes the old session, and + either starts the replacement or waits for another browser request. +11. Ctrl-C closes the peer/server and then the application. Releasing the + application drops the one resident pipeline after every session cache has + already been released. + +## Why the WebRTC window is opened twice + +The two calls mark different boundaries: + +- The first call prepares the fixed stream format so a browser can connect + while no session exists. +- The per-session call discards source frames queued by the previous rollout. + Event timestamps retain one window-lifetime monotonic origin, so opening a + replacement cannot reorder buffered events from the old and new browser. + +They do not create two servers or two peer connections. A connected peer is +reused. The application must either accept the already-resolved `SessionDesc` +or reject it; it cannot silently change width, height, layout, or playback rate +after the browser stream was prepared. + +## Lifecycle event ordering + +Close and new-session events can arrive in one drained batch during a page +refresh. Their order is meaningful, so the latest transition wins: + +- close, then new session: start the new browser's request; +- new session, then close: cancel the request because that browser went away. + +Any close or replacement also advances the session generation. A model step +that was already running may finish, but its result is tagged with the old +generation and is not written. The WebRTC media track similarly clears frames +that it has not yet handed to the encoder. At most one already-encoded frame +can still be in flight; a transport cannot recall a frame it has sent. + +## WebRTC connection lifecycle + +Offer negotiation is serialized, and only one browser is admitted. Data +channel and peer callbacks capture the peer that registered them, so callbacks +from an old page cannot close or mutate its replacement. Closing the active +data channel releases the peer immediately, allowing a refreshed page to +negotiate without an indefinite series of HTTP 409 responses. + +The new-session button starts enabled and reads **Opening...** while signaling +is in progress. A click during that interval keeps the latest valid prompt in +the page and sends it as soon as the data channel opens. The UI therefore does +not make model loading or a brief reconnect look like an unavailable action. + +The media track uses `frames_per_second_for_step`, the generated video's +playback rate. `frames_per_second_for_ui` controls only how often the runtime +polls input and presents completed work. + +The media source queue is currently unbounded. T2V rollouts are finite, and a +session replacement clears frames that have not reached the encoder, so its +size remains bounded by a rollout in the current application. A future +long-running producer will need an explicit block-or-drop policy chosen for +that application's latency requirements. + +## Why these boundaries are useful + +- The core generation pipeline and checkpoint load once and stay resident + across browser refreshes and sequential prompts. The existing Wan pipeline + may release and reload its one-shot text encoder after prompt encoding to fit + within GPU memory; it does not reload the diffusion model between sessions. +- Session cleanup is complete before the next session is created, so rollout + caches cannot overlap accidentally. +- WebRTC code transports validated events but never makes application + lifecycle decisions. +- `ApplicationRunner.run()` visibly owns the possibly multi-session run, while + `run_session()` visibly owns exactly one rollout. +- Failure paths keep the first useful exception while still releasing the + session I/O thread, WebRTC server thread, window, and partially initialized + model. + +`SessionDescRequest` remains separate from `SessionDesc` on purpose: the former +means “only the fields the CLI explicitly supplied,” while the latter is the +fully resolved contract shared by the application, session, and window. No +additional host, session-runner class, or WebRTC-specific lifecycle wrapper is +needed. + +The bundled browser page is currently prompt-oriented because text-to-video is +the application that needs client-created sessions today. A generic UI schema +should be introduced only when another application has a concrete, different +request format. diff --git a/flashdreams/flashdreams/api_v2/application.py b/flashdreams/flashdreams/api_v2/application.py index 45f2e8eb9..4bde6b0f8 100644 --- a/flashdreams/flashdreams/api_v2/application.py +++ b/flashdreams/flashdreams/api_v2/application.py @@ -33,9 +33,9 @@ def default_session_desc(self) -> SessionDesc | None: The application owns its model's output requirements and defaults, such as its layout, preferred dimensions, and frame rate. The runtime asks after :meth:`init`, then applies the caller's explicit requests to this - default before calling :meth:`create_session`. That method may reject - or further resolve the request to satisfy the model's requirements. The - created session's ``session_desc`` is authoritative when a window opens. + default before calling :meth:`create_session`. That method accepts the + resolved description or rejects it; it does not silently change the + stream the runtime already asked the client window to prepare for. Returns: The default session description, or ``None`` when the application @@ -51,8 +51,7 @@ def create_session(self, session_desc: SessionDesc) -> ISession: session_desc: Session the runtime is asking for. Returns: - A session for ``session_desc``, resolved to what this application can - actually produce. + A session that produces ``session_desc``. Raises: ValueError: The application cannot honour ``session_desc``. diff --git a/flashdreams/flashdreams/api_v2/session.py b/flashdreams/flashdreams/api_v2/session.py index 9b91a1e74..c328363c9 100644 --- a/flashdreams/flashdreams/api_v2/session.py +++ b/flashdreams/flashdreams/api_v2/session.py @@ -26,9 +26,11 @@ class ISession(ABC): @abstractmethod def init(self) -> None: - """Load the model and anything else this run needs. + """Prepare the state owned by this run. - Must not do client I/O, since this can run before a client connects. + Shared model loading belongs to the application. This method prepares + per-session state such as an encoded prompt or KV cache. It must not do + client I/O, since this can run before a client connects. """ ... diff --git a/flashdreams/flashdreams/runtime_v2/application_runner.py b/flashdreams/flashdreams/runtime_v2/application_runner.py index e73f744fa..b6c1b1ea4 100644 --- a/flashdreams/flashdreams/runtime_v2/application_runner.py +++ b/flashdreams/flashdreams/runtime_v2/application_runner.py @@ -46,7 +46,7 @@ def init(self, commandline_args: Sequence[str] = ()) -> None: self._application.init(commandline_args) self._initialized = True - def run_session( + def run( self, session_desc_request: SessionDescRequest, client_window: IClientWindow, @@ -82,13 +82,13 @@ def run_session( try: next_session_desc = self._resolve_session_desc(session_desc_request) - except Exception: + except BaseException: _close_client_window(client_window) raise while True: try: session = self._application.create_session(next_session_desc) - except Exception: + except BaseException: _close_client_window(client_window) raise next_session_desc = run_session(session, client_window) @@ -112,8 +112,8 @@ def _serve_sessions( client_window: IClientWindow, ) -> None: """Keep one client window available for browser-requested sessions.""" - current_session_desc = self._resolve_session_desc(session_desc_request) try: + current_session_desc = self._resolve_session_desc(session_desc_request) client_window.open(current_session_desc) next_session_desc: SessionDesc | None = None while True: diff --git a/flashdreams/flashdreams/runtime_v2/cli.py b/flashdreams/flashdreams/runtime_v2/cli.py index 4e5ac8cd8..14d7386d9 100644 --- a/flashdreams/flashdreams/runtime_v2/cli.py +++ b/flashdreams/flashdreams/runtime_v2/cli.py @@ -14,6 +14,7 @@ """ import argparse +import logging import sys from collections.abc import Sequence @@ -36,6 +37,8 @@ command also has, so the split is stated rather than guessed. """ +_LOGGER = logging.getLogger(__name__) + def entrypoint(argv: Sequence[str] | None = None) -> None: """Run the command, reporting where to watch what it generates.""" @@ -56,11 +59,20 @@ def entrypoint(argv: Sequence[str] | None = None) -> None: try: runner.init(application_args) window = mode.create(parsed) - _report(mode.starting(window)) + try: + _report(mode.starting(window)) + except BaseException: + try: + window.close() + except Exception: + _LOGGER.exception( + "The client window failed to close after startup failed." + ) + raise # Nothing here says how long a session is: the application reports that. # A serving mode stays up between sessions; a file mode runs just one. try: - runner.run_session( + runner.run( _session_desc_request(parsed), window, serve_sessions=mode.serves_sessions, diff --git a/flashdreams/flashdreams/runtime_v2/client_window_factory.py b/flashdreams/flashdreams/runtime_v2/client_window_factory.py index 0500b6c65..8decc572b 100644 --- a/flashdreams/flashdreams/runtime_v2/client_window_factory.py +++ b/flashdreams/flashdreams/runtime_v2/client_window_factory.py @@ -117,8 +117,8 @@ def create(self, parsed_args: argparse.Namespace) -> IClientWindow: def starting(self, client_window: IClientWindow) -> str | None: """Return where to connect, which nobody can guess when the port is free.""" - server = cast("WebRTCClientWindow", client_window).server - return f"Open {server.url} in a browser." + window = cast("WebRTCClientWindow", client_window) + return f"Open {window.url} in a browser." _MODES: tuple[ClientWindowMode, ...] = (_Mp4Mode(), _WebRTCMode()) diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/app.js b/flashdreams/flashdreams/runtime_v2/serving/web/app.js index 9e25d6542..8ed145b0f 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/web/app.js +++ b/flashdreams/flashdreams/runtime_v2/serving/web/app.js @@ -5,6 +5,7 @@ const peer = new RTCPeerConnection(); const controls = peer.createDataChannel("controls"); peer.addTransceiver("video", {direction: "recvonly"}); const newSessionButton = document.getElementById("new-session"); +const promptInput = document.getElementById("prompt"); let pendingNewSession = null; const send = payload => { @@ -24,6 +25,7 @@ controls.addEventListener("open", () => { controls.addEventListener("close", () => { newSessionButton.disabled = true; + newSessionButton.textContent = "Disconnected"; }); peer.ontrack = event => { @@ -32,10 +34,16 @@ peer.ontrack = event => { }; window.addEventListener("keydown", event => { + if (event.target === promptInput) { + return; + } send({type: "keyboard", key: event.key, pressed: true}); }); window.addEventListener("keyup", event => { + if (event.target === promptInput) { + return; + } send({type: "keyboard", key: event.key, pressed: false}); }); @@ -52,7 +60,6 @@ document.getElementById("reset").onclick = () => { }; newSessionButton.onclick = () => { - const promptInput = document.getElementById("prompt"); if (!promptInput.reportValidity()) { return; } @@ -97,4 +104,8 @@ async function connect() { await peer.setRemoteDescription(await response.json()); } -connect(); +connect().catch(error => { + newSessionButton.disabled = true; + newSessionButton.textContent = "Connection failed"; + console.error(error); +}); diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/index.html b/flashdreams/flashdreams/runtime_v2/serving/web/index.html index 0cd2798a8..bd37b26c1 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/web/index.html +++ b/flashdreams/flashdreams/runtime_v2/serving/web/index.html @@ -13,7 +13,7 @@ - + diff --git a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py index b7e5b9112..4e89190ec 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py +++ b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py @@ -7,6 +7,7 @@ import asyncio import json +import logging import socket import threading import time @@ -18,7 +19,12 @@ import numpy as np import torch from aiohttp import web -from aiortc import MediaStreamTrack, RTCPeerConnection, RTCSessionDescription +from aiortc import ( + MediaStreamTrack, + RTCDataChannel, + RTCPeerConnection, + RTCSessionDescription, +) from aiortc.mediastreams import MediaStreamError from av import VideoFrame @@ -37,6 +43,8 @@ _BROWSER_PAGE = _WEB_RESOURCES.joinpath("index.html").read_text(encoding="utf-8") _BROWSER_SCRIPT = _WEB_RESOURCES.joinpath("app.js").read_text(encoding="utf-8") +_LOGGER = logging.getLogger(__name__) + class _VideoTrack(MediaStreamTrack): """Video track whose frames are supplied by the server.""" @@ -113,7 +121,11 @@ async def close(self) -> None: class WebRTCServer: - """Own the HTTP, signaling, input buffering, and media transport.""" + """Own HTTP signaling and media transport for one browser at a time. + + Validated browser events leave through the registered callback. The client + window owns the thread-safe queue that buffers them for the runtime loop. + """ def __init__( self, @@ -147,10 +159,12 @@ def __init__( self._startup_error: BaseException | None = None self._loop: asyncio.AbstractEventLoop | None = None self._runner: web.AppRunner | None = None + self._offer_lock: asyncio.Lock | None = None self._peer_connection: RTCPeerConnection | None = None + self._control_channel: RTCDataChannel | None = None self._video_track: _VideoTrack | None = None self._session_desc: SessionDesc | None = None - self._session_start_ns: int | None = None + self._event_origin_ns: int | None = None self._closed = False self._client_connected = False self._thread = threading.Thread( @@ -160,6 +174,8 @@ def __init__( ) self._thread.start() if not self._started.wait(startup_timeout_seconds): + self._closed = True + self._stop_server_thread() raise TimeoutError("WebRTC server did not start before the timeout.") if self._startup_error is not None: raise RuntimeError( @@ -209,7 +225,11 @@ def open(self, session_desc: SessionDesc) -> None: "layout, frame rate, width, and height." ) self._session_desc = session_desc - self._session_start_ns = time.monotonic_ns() + if self._event_origin_ns is None: + # UserInputEvents sorts by timestamp. Keep one origin for this + # window's lifetime so opening a replacement session cannot reorder + # an older browser's close after a newer browser's request. + self._event_origin_ns = time.monotonic_ns() track = self._video_track loop = self._loop if track is not None and loop is not None: @@ -262,8 +282,24 @@ def close(self) -> None: if loop is None: return future = asyncio.run_coroutine_threadsafe(self._shutdown(), loop) - future.result(timeout=self._startup_timeout_seconds) - loop.call_soon_threadsafe(loop.stop) + try: + future.result(timeout=self._startup_timeout_seconds) + except BaseException: + try: + self._stop_server_thread() + except Exception: + _LOGGER.exception( + "The WebRTC server thread also failed to stop during cleanup." + ) + raise + else: + self._stop_server_thread() + + def _stop_server_thread(self) -> None: + """Stop and join the server thread after startup or shutdown.""" + loop = self._loop + if loop is not None and not loop.is_closed(): + loop.call_soon_threadsafe(loop.stop) self._thread.join(timeout=self._startup_timeout_seconds) if self._thread.is_alive(): raise TimeoutError("WebRTC server did not stop before the timeout.") @@ -288,6 +324,7 @@ def _run_server(self) -> None: async def _start_server(self) -> None: """Create and bind the standalone aiohttp application.""" + self._offer_lock = asyncio.Lock() app = web.Application() app.router.add_get("/", self._serve_browser) app.router.add_get("/app.js", self._serve_browser_script) @@ -335,16 +372,6 @@ async def _offer(self, request: web.Request) -> web.Response: session_desc = self._session_desc if session_desc is None: raise web.HTTPConflict(reason="WebRTC server is not open.") - existing_peer = self._peer_connection - if existing_peer is not None: - if existing_peer.connectionState in { - "failed", - "disconnected", - "closed", - }: - await self._release_peer_connection(existing_peer) - else: - raise web.HTTPConflict(reason="A WebRTC client is already connected.") try: payload = await request.json() @@ -359,55 +386,84 @@ async def _offer(self, request: web.Request) -> web.Response: reason="WebRTC offer requires string sdp and type." ) - peer_connection = RTCPeerConnection() - video_track = _VideoTrack(session_desc.frames_per_second_for_ui) - peer_connection.addTrack(video_track) - self._peer_connection = peer_connection - self._video_track = video_track - - @peer_connection.on("datachannel") - def on_datachannel(channel: Any) -> None: - self._client_connected = True - - @channel.on("message") - def on_message(message: Any) -> None: - try: - self._buffer_browser_message(message) - except ValueError as error: - channel.send(json.dumps({"type": "error", "message": str(error)})) - - @channel.on("close") - def on_close() -> None: - self._record_client_disconnect() - - @peer_connection.on("connectionstatechange") - async def on_connectionstatechange() -> None: - if peer_connection.connectionState in {"failed", "disconnected", "closed"}: - self._record_client_disconnect() - await self._release_peer_connection(peer_connection) + offer_lock = self._offer_lock + if offer_lock is None: + raise web.HTTPServiceUnavailable(reason="WebRTC server is not ready.") + async with offer_lock: + existing_peer = self._peer_connection + if existing_peer is not None: + control_channel = self._control_channel + if existing_peer.connectionState in { + "failed", + "disconnected", + "closed", + } or ( + control_channel is not None + and control_channel.readyState == "closed" + ): + await self._release_peer_connection(existing_peer) + else: + raise web.HTTPConflict( + reason="A WebRTC client is already connected." + ) + + peer_connection = RTCPeerConnection() + video_track = _VideoTrack(session_desc.frames_per_second_for_step) + peer_connection.addTrack(video_track) + self._peer_connection = peer_connection + self._video_track = video_track + + @peer_connection.on("datachannel") + def on_datachannel(channel: RTCDataChannel) -> None: + if self._peer_connection is not peer_connection: + channel.close() + return + self._control_channel = channel + self._client_connected = True + + @channel.on("message") + def on_message(message: Any) -> None: + if self._peer_connection is not peer_connection: + return + try: + self._buffer_browser_message(message) + except ValueError as error: + channel.send( + json.dumps({"type": "error", "message": str(error)}) + ) + + @channel.on("close") + async def on_close() -> None: + await self._release_peer_connection(peer_connection) + + @peer_connection.on("connectionstatechange") + async def on_connectionstatechange() -> None: + if peer_connection.connectionState in { + "failed", + "disconnected", + "closed", + }: + await self._release_peer_connection(peer_connection) - try: - await peer_connection.setRemoteDescription( - RTCSessionDescription(sdp=sdp, type=offer_type) - ) - await peer_connection.setLocalDescription( - await peer_connection.createAnswer() - ) - except Exception: - self._peer_connection = None - self._video_track = None - await video_track.close() - await peer_connection.close() - raise + try: + await peer_connection.setRemoteDescription( + RTCSessionDescription(sdp=sdp, type=offer_type) + ) + await peer_connection.setLocalDescription( + await peer_connection.createAnswer() + ) + local_description = peer_connection.localDescription + if local_description is None: + raise web.HTTPInternalServerError( + reason="WebRTC peer did not create an answer." + ) + except BaseException: + await self._release_peer_connection(peer_connection) + raise - local_description = peer_connection.localDescription - if local_description is None: - raise web.HTTPInternalServerError( - reason="WebRTC peer did not create an answer." + return web.json_response( + {"sdp": local_description.sdp, "type": local_description.type} ) - return web.json_response( - {"sdp": local_description.sdp, "type": local_description.type} - ) def _buffer_browser_message(self, raw_message: object) -> None: """Validate and append one data-channel message.""" @@ -455,10 +511,10 @@ def _append_event( ), ) -> None: """Timestamp and buffer one validated browser event.""" - session_start_ns = self._session_start_ns - if session_start_ns is None: + event_origin_ns = self._event_origin_ns + if event_origin_ns is None: return - timestamp_us = np.uint64((time.monotonic_ns() - session_start_ns) // 1_000) + timestamp_us = np.uint64((time.monotonic_ns() - event_origin_ns) // 1_000) event = UserInputEvent(timestamp=timestamp_us, event_data=event_data) callback = self._input_callback if callback is None: @@ -467,9 +523,9 @@ def _append_event( # thread-safe queue; the server loop never handles lifecycle events. callback(event) - def _record_client_disconnect(self) -> None: + def _record_client_disconnect(self, peer_connection: RTCPeerConnection) -> None: """Buffer one close event when the active browser disconnects.""" - if not self._client_connected: + if self._peer_connection is not peer_connection or not self._client_connected: return self._client_connected = False if not self._closed: @@ -489,8 +545,9 @@ async def _release_peer_connection( """Release one disconnected peer so a refreshed page can reconnect.""" if self._peer_connection is not peer_connection: return - self._record_client_disconnect() + self._record_client_disconnect(peer_connection) self._peer_connection = None + self._control_channel = None track = self._video_track self._video_track = None if track is not None: @@ -518,7 +575,7 @@ def _same_stream_format(first: SessionDesc, second: SessionDesc) -> bool: """Return whether two sessions can use the same WebRTC media track.""" return ( first.output_layout is second.output_layout - and first.frames_per_second_for_ui == second.frames_per_second_for_ui + and first.frames_per_second_for_step == second.frames_per_second_for_step and first.video_width == second.video_width and first.video_height == second.video_height ) diff --git a/flashdreams/flashdreams/runtime_v2/session_desc.py b/flashdreams/flashdreams/runtime_v2/session_desc.py index 7333cb846..076359e4c 100644 --- a/flashdreams/flashdreams/runtime_v2/session_desc.py +++ b/flashdreams/flashdreams/runtime_v2/session_desc.py @@ -14,9 +14,9 @@ class SessionDesc: """Description of a session, passed to create one and to open a window on it. - The runtime fills this in to ask an application for a session, and the - session reports back what it resolved to. The same description then - configures the client window through ``OutputSink.open``. + The runtime resolves this before asking an application for a session. The + application either accepts it or raises; the same description configures + the client window through ``OutputSink.open``. """ output_layout: VideoTensorLayout = VideoTensorLayout.tchw @@ -26,7 +26,7 @@ class SessionDesc: """Rate to read input and present finished results at, in frames per second.""" frames_per_second_for_step: int = 30 - """Rate to generate at, in frames per second. Nothing paces by it yet.""" + """Playback rate of generated frames, in frames per second.""" video_width: int = 1280 """Output video width in pixels.""" @@ -97,9 +97,7 @@ def resolve(self, default: SessionDesc) -> SessionDesc: default.video_width if self.video_width is None else self.video_width ), video_height=( - default.video_height - if self.video_height is None - else self.video_height + default.video_height if self.video_height is None else self.video_height ), metadata=default.metadata if self.metadata is None else self.metadata, ) diff --git a/flashdreams/flashdreams/runtime_v2/session_runner.py b/flashdreams/flashdreams/runtime_v2/session_runner.py index 54679d105..ab070d2e0 100644 --- a/flashdreams/flashdreams/runtime_v2/session_runner.py +++ b/flashdreams/flashdreams/runtime_v2/session_runner.py @@ -45,21 +45,31 @@ def _contains(events: UserInputEvents, event_type: type[UserInputEventData]) -> ) -def _next_session_desc( +def _latest_session_transition( events: UserInputEvents, current: SessionDesc -) -> SessionDesc | None: - """Return the latest session requested by ``events``, if there is one. +) -> tuple[bool, SessionDesc | None]: + """Return the latest close or replacement requested by ``events``. A new-session request replaces application-owned metadata while preserving the resolved video and timing settings of the current session. The application interprets and validates the metadata when it creates the next - session. + session. A close after that request cancels it; a request after a close wins. + + Returns: + Whether the batch contains a close or replacement, and the replacement + description. The description is ``None`` when close is latest. """ - for event in reversed(events.get_events()): + found = False + next_session_desc = None + for event in events.get_events(): event_data = event.get_event_data() - if isinstance(event_data, NewSessionUserInputEventData): - return replace(current, metadata=dict(event_data.metadata)) - return None + if isinstance(event_data, CloseUserInputEventData): + found = True + next_session_desc = None + elif isinstance(event_data, NewSessionUserInputEventData): + found = True + next_session_desc = replace(current, metadata=dict(event_data.metadata)) + return found, next_session_desc def wait_for_new_session( @@ -81,7 +91,7 @@ def wait_for_new_session( """ tick_seconds = 1.0 / current_session_desc.frames_per_second_for_ui while True: - next_session_desc = _next_session_desc( + _, next_session_desc = _latest_session_transition( window.get_user_input_events(), current_session_desc ) if next_session_desc is not None: @@ -102,7 +112,7 @@ def _close_session(session: ISession, *, run_failed: bool) -> None: """ try: session.close() - except Exception: + except BaseException: if not run_failed: raise _LOGGER.exception( @@ -197,9 +207,15 @@ def run_session( try: session.init() - except Exception: + except BaseException: # A partly initialized session still holds whatever it managed to load. _close_session(session, run_failed=True) + try: + window.close() + except Exception: + _LOGGER.exception( + "The window failed to close after session initialization failed." + ) raise # Backpressure is all here. Finished results wait here for the I/O thread to @@ -247,24 +263,20 @@ def present_pending_results() -> None: def tick() -> None: nonlocal generation, next_session_desc events = window.get_user_input_events() - requested_session_desc = _next_session_desc(events, session.session_desc) + transition_found, requested_session_desc = _latest_session_transition( + events, session.session_desc + ) with collected_events_lock: collected_events.extend(events.get_events()) # Move on to the next generation from here, since this thread sees the # reset or replacement first. Under the lock, so a step already # picking up its input either belongs to the generation being # abandoned or to the new one, never to neither. - if ( - _contains(events, ResetUserInputEventData) - or requested_session_desc is not None - ): + if _contains(events, ResetUserInputEventData) or transition_found: generation += 1 # Stop from here rather than waiting for the step loop to notice, so a # slow step does not delay a client that has gone away. - if _contains(events, CloseUserInputEventData): - next_session_desc = None - stop.set() - elif requested_session_desc is not None: + if transition_found: next_session_desc = requested_session_desc stop.set() session.step_ui(events) @@ -375,27 +387,42 @@ def add_pending_result(result_generation: int, result: StepResult) -> int: steps_run += 1 finally: stop.set() - io_stopped.wait() - # Another session can own the same window, but only after this one has - # released everything it holds. If cleanup fails, tell the I/O thread to - # close the window and report the failure instead of preserving it. - run_failed = sys.exc_info()[0] is not None + run_failed = False session_close_attempted = False session_close_error: Exception | None = None - if ( - (next_session_desc is not None or keep_window_open) - and not run_failed - and not io_failure - ): - session_close_attempted = True + try: try: - session.close() - except Exception as error: - session_close_error = error - else: - window_stays_open = True - finish_io.set() - io_thread.join() + io_stopped.wait() + # Another session can own the same window, but only after this + # one has released everything it holds. If cleanup fails, tell + # the I/O thread to close the window and report the failure. + run_failed = sys.exc_info()[0] is not None + if ( + (next_session_desc is not None or keep_window_open) + and not run_failed + and not io_failure + ): + session_close_attempted = True + try: + session.close() + except Exception as error: + session_close_error = error + else: + window_stays_open = True + finally: + # The I/O thread waits for the cleanup decision before it can + # close the window. Always release it, including when cleanup + # is interrupted. + finish_io.set() + io_thread.join() + finally: + # An interrupt while waiting for the I/O thread skips the normal + # cleanup below. Close the session here so that path cannot leak it. + if not session_close_attempted: + _close_session( + session, + run_failed=sys.exc_info()[0] is not None or bool(io_failure), + ) # A failure here is what the run reports, since a window failure stops # generation rather than raising through it: both places this thread can # be sitting give up once io_failure is set, so a run that reports a @@ -412,8 +439,6 @@ def add_pending_result(result_generation: int, result: StepResult) -> int: "that failure.", exc_info=io_failure[0], ) - if not session_close_attempted: - _close_session(session, run_failed=run_failed or bool(io_failure)) if session_close_error is not None: raise session_close_error diff --git a/flashdreams/flashdreams/runtime_v2/user_input_event.py b/flashdreams/flashdreams/runtime_v2/user_input_event.py index dbb894b32..64fe9acb3 100644 --- a/flashdreams/flashdreams/runtime_v2/user_input_event.py +++ b/flashdreams/flashdreams/runtime_v2/user_input_event.py @@ -151,7 +151,7 @@ class UserInputEvent: """User input event.""" timestamp: uint64 - """Timestamp in microseconds since the start of the session.""" + """Microseconds since the input source's stable monotonic time origin.""" event_data: UserInputEventData """Event data.""" diff --git a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py index 5761e7059..4e5d1660b 100644 --- a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py +++ b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py @@ -34,7 +34,7 @@ def __init__( startup_timeout_seconds: Maximum time to wait for server startup. """ self._input_events: queue.SimpleQueue[UserInputEvent] = queue.SimpleQueue() - self.server = WebRTCServer( + self._server = WebRTCServer( host=host, port=port, startup_timeout_seconds=startup_timeout_seconds, @@ -44,7 +44,12 @@ def handle_input(event: UserInputEvent) -> None: """Buffer one backend event for the ``InputSource`` protocol.""" self._input_events.put(event) - self.server.register_input_callback(handle_input) + self._server.register_input_callback(handle_input) + + @property + def url(self) -> str: + """Return the URL at which a browser can open this window.""" + return self._server.url def open(self, session_desc: SessionDesc) -> None: """Configure WebRTC output for waiting or running a session. @@ -52,7 +57,7 @@ def open(self, session_desc: SessionDesc) -> None: Args: session_desc: Resolved dimensions, frame rate, and tensor layout. """ - self.server.open(session_desc) + self._server.open(session_desc) def get_user_input_events(self) -> UserInputEvents: """Implement ``InputSource.get_user_input_events`` for browser input. @@ -73,8 +78,8 @@ def write(self, result: StepResult) -> None: Args: result: Generated frames matching the opened session. """ - self.server.write(result) + self._server.write(result) def close(self) -> None: """Implement ``OutputSink.close`` by releasing WebRTC resources.""" - self.server.close() + self._server.close() diff --git a/flashdreams/flashdreams/t2v_v2/application.py b/flashdreams/flashdreams/t2v_v2/application.py index c9d3e143d..0060b59bb 100644 --- a/flashdreams/flashdreams/t2v_v2/application.py +++ b/flashdreams/flashdreams/t2v_v2/application.py @@ -142,7 +142,13 @@ def init(self, commandline_args: Sequence[str]) -> None: device=args.device, total_blocks=args.total_blocks, ) - pipeline = self._pipeline_config.setup().to(config.device).eval() + # Take ownership as soon as setup returns. If moving or evaluating the + # model fails, ApplicationRunner.close() can still release what loaded. + pipeline = self._pipeline_config.setup() + self._pipeline = pipeline + pipeline = pipeline.to(config.device) + self._pipeline = pipeline + pipeline = pipeline.eval() self._config = config self._pipeline = pipeline diff --git a/flashdreams/test_v2/test_application_runner.py b/flashdreams/test_v2/test_application_runner.py index 8ae2afa0d..3660a5902 100644 --- a/flashdreams/test_v2/test_application_runner.py +++ b/flashdreams/test_v2/test_application_runner.py @@ -196,8 +196,8 @@ def test_application_runner_keeps_the_application_open_for_another_session() -> second_window = _SilentWindow(calls) runner.init(["--model-option"]) - runner.run_session(_session_desc_request(), first_window) - runner.run_session(_session_desc_request(), second_window) + runner.run(_session_desc_request(), first_window) + runner.run(_session_desc_request(), second_window) assert [result.step_index for result in first_window.results] == [0] assert [result.step_index for result in second_window.results] == [0] @@ -237,7 +237,7 @@ def test_application_runner_replaces_a_session_from_window_metadata() -> None: window = _ScriptedWindow(calls, [new_session, close]) runner.init() - runner.run_session(_session_desc_request(), window) + runner.run(_session_desc_request(), window) runner.close() assert len(application.created_session_descs) == 2 @@ -261,7 +261,7 @@ def test_application_runner_serves_sessions_until_it_is_interrupted() -> None: runner.init() with pytest.raises(KeyboardInterrupt): - runner.run_session(_session_desc_request(), window, serve_sessions=True) + runner.run(_session_desc_request(), window, serve_sessions=True) assert [desc.metadata for desc in application.created_session_descs] == [ {"prompt": "A cat surfing"}, @@ -281,11 +281,28 @@ def test_application_runner_closes_the_window_when_a_session_cannot_start() -> N window = _Window(calls) with pytest.raises(RuntimeError, match=r"init\(\) must run first"): - runner.run_session(_session_desc_request(), window) + runner.run(_session_desc_request(), window) assert calls == ["window.close"] +def test_serving_closes_the_window_when_the_session_request_is_invalid() -> None: + calls: list[str] = [] + runner = ApplicationRunner(_Application(calls)) + window = _Window(calls) + runner.init() + + with pytest.raises(ValueError, match="frames_per_second_for_ui"): + runner.run( + SessionDescRequest(frames_per_second_for_ui=0), + window, + serve_sessions=True, + ) + + assert calls == ["application.init([])", "window.close"] + runner.close() + + def test_application_runner_rejects_a_second_initialization() -> None: calls: list[str] = [] runner = ApplicationRunner(_Application(calls)) diff --git a/flashdreams/test_v2/test_cli.py b/flashdreams/test_v2/test_cli.py index 5605e7ddb..6a5c26b7e 100644 --- a/flashdreams/test_v2/test_cli.py +++ b/flashdreams/test_v2/test_cli.py @@ -154,6 +154,7 @@ class RecordingWindow(IClientWindow): def __init__(self) -> None: self.results: list[StepResult] = [] + self.closed = False def get_user_input_events(self) -> UserInputEvents: return UserInputEvents([]) @@ -165,7 +166,7 @@ def write(self, result: StepResult) -> None: self.results.append(result) def close(self) -> None: - return + self.closed = True ## Splitting the command line @@ -406,13 +407,13 @@ def test_nothing_is_measured_unless_a_run_asks( assert list(tmp_path.glob("*.json")) == [] -def test_a_one_shot_run_without_a_prompt_reports_why( +def test_a_non_serving_run_without_a_prompt_reports_why( monkeypatch: pytest.MonkeyPatch, ) -> None: _install(monkeypatch, StubT2VApplication(_stand_in()), RecordingWindow()) with pytest.raises(ValueError, match="session prompt is required"): - cli.entrypoint(["stub", "--mode", "webrtc"]) + cli.entrypoint(["stub", "--mode", "mp4"]) def test_a_browser_server_starts_without_a_command_line_prompt( @@ -530,6 +531,31 @@ def test_the_run_goes_to_the_window_the_mode_asked_for( assert len(window.results) == _TOTAL_BLOCKS +def test_a_window_is_closed_when_reporting_its_start_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FailingStartMode(StubMode): + def starting(self, client_window: IClientWindow) -> str | None: + del client_window + raise RuntimeError("cannot report window") + + pipeline = _stand_in() + application = StubT2VApplication(pipeline) + window = RecordingWindow() + _install(monkeypatch, application) + monkeypatch.setattr( + cli, + "client_window_mode", + lambda name: FailingStartMode(name, window), + ) + + with pytest.raises(RuntimeError, match="cannot report window"): + cli.entrypoint(["stub", "--mode", "webrtc"]) + + assert window.closed + assert pipeline.closed + + def test_the_command_needs_somewhere_to_write() -> None: with pytest.raises(SystemExit): cli.entrypoint(["stub"]) diff --git a/flashdreams/test_v2/test_client_window_factory.py b/flashdreams/test_v2/test_client_window_factory.py index c77466d7e..46ce8267b 100644 --- a/flashdreams/test_v2/test_client_window_factory.py +++ b/flashdreams/test_v2/test_client_window_factory.py @@ -74,7 +74,7 @@ def test_a_browser_run_is_told_where_to_connect(self) -> None: window = mode.create(_parsed(["--mode", "webrtc"])) try: assert isinstance(window, WebRTCClientWindow) - assert mode.starting(window) == f"Open {window.server.url} in a browser." + assert mode.starting(window) == f"Open {window.url} in a browser." assert mode.serves_sessions is True finally: window.close() diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py index d0d8faca4..e30de98cd 100644 --- a/flashdreams/test_v2/test_session_runner.py +++ b/flashdreams/test_v2/test_session_runner.py @@ -15,7 +15,11 @@ from flashdreams.api_v2.session import ISession from flashdreams.api_v2.user_input_event_data import UserInputEventData from flashdreams.runtime_v2.session_desc import SessionDesc -from flashdreams.runtime_v2.session_runner import WhenFull, run_session +from flashdreams.runtime_v2.session_runner import ( + WhenFull, + run_session, + wait_for_new_session, +) from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import ( CloseUserInputEventData, @@ -343,6 +347,71 @@ def test_run_session_returns_the_next_session_desc_after_cleanup() -> None: assert "session.step(0)" not in log.calls +def test_the_latest_close_or_new_session_event_wins() -> None: + session_desc = _session_desc() + replacement = NewSessionUserInputEventData(metadata={"prompt": "new browser"}) + + close_then_new = RecordingClientWindow( + CallLog(), + [ + UserInputEvents( + [ + UserInputEvent(uint64(0), CloseUserInputEventData()), + UserInputEvent(uint64(1), replacement), + ] + ) + ], + ) + returned = run_session( + FakeSession(session_desc, CallLog()), close_then_new, steps=None + ) + + assert returned == replace(session_desc, metadata={"prompt": "new browser"}) + + new_then_close = RecordingClientWindow( + CallLog(), + [ + UserInputEvents( + [ + UserInputEvent(uint64(0), replacement), + UserInputEvent(uint64(1), CloseUserInputEventData()), + ] + ) + ], + ) + returned = run_session( + FakeSession(session_desc, CallLog()), new_then_close, steps=None + ) + + assert returned is None + + +def test_wait_for_new_session_ignores_a_request_cancelled_by_close() -> None: + session_desc = _session_desc() + replacement = NewSessionUserInputEventData(metadata={"prompt": "new browser"}) + window = RecordingClientWindow( + CallLog(), + [ + UserInputEvents( + [ + UserInputEvent(uint64(0), replacement), + UserInputEvent(uint64(1), CloseUserInputEventData()), + ] + ), + UserInputEvents( + [ + UserInputEvent(uint64(2), CloseUserInputEventData()), + UserInputEvent(uint64(3), replacement), + ] + ), + ], + ) + + returned = wait_for_new_session(window, session_desc) + + assert returned == replace(session_desc, metadata={"prompt": "new browser"}) + + def test_run_session_closes_the_window_when_replacement_cleanup_fails() -> None: log = CallLog() session = FakeSession(_session_desc(), log, fail_to_close=True) @@ -361,6 +430,24 @@ def test_run_session_closes_the_window_when_replacement_cleanup_fails() -> None: assert log.calls[-2:] == ["session.close", "window.close"] +def test_run_session_releases_the_io_thread_when_cleanup_is_interrupted() -> None: + log = CallLog() + + class InterruptedCleanupSession(FakeSession): + def close(self) -> None: + super().close() + raise KeyboardInterrupt + + session = InterruptedCleanupSession(_session_desc(), log) + window = RecordingClientWindow(log) + + with pytest.raises(KeyboardInterrupt): + run_session(session, window, steps=0, keep_window_open=True) + + assert log.calls[-2:] == ["session.close", "window.close"] + assert not any(thread.name == _IO_THREAD_NAME for thread in threading.enumerate()) + + def test_run_session_resets_the_session_and_the_step_index() -> None: log = CallLog() session = FakeSession(_session_desc(), log) @@ -426,7 +513,7 @@ def test_run_session_lets_a_reset_restart_a_finished_session() -> None: assert "session.reset" in log.calls -def test_run_session_closes_a_session_that_failed_to_init() -> None: +def test_run_session_closes_a_session_and_window_that_failed_to_init() -> None: log = CallLog() class FailingSession(FakeSession): @@ -440,9 +527,65 @@ def init(self) -> None: with pytest.raises(RuntimeError, match="init failed"): run_session(session, window, steps=1) - # A session that got halfway through starting still has to be released, and - # the window is never opened for a session that cannot run. - assert log.calls == ["session.init", "session.close"] + # A session that got halfway through starting still has to be released. The + # constructor-owned window may already hold resources even though open was + # never called, so it is closed too. + assert log.calls == ["session.init", "session.close", "window.close"] + + +def test_run_session_closes_a_session_whose_init_is_interrupted() -> None: + log = CallLog() + + class InterruptedSession(FakeSession): + def init(self) -> None: + super().init() + raise KeyboardInterrupt + + session = InterruptedSession(_session_desc(), log) + + with pytest.raises(KeyboardInterrupt): + run_session(session, RecordingClientWindow(log), steps=1) + + assert log.calls == ["session.init", "session.close", "window.close"] + + +def test_run_session_closes_the_session_when_io_shutdown_is_interrupted( + monkeypatch: pytest.MonkeyPatch, +) -> None: + real_event = threading.Event + events_created = 0 + + class InterruptingEvent(real_event): + """Interrupt the caller waiting for the runner's I/O-stopped event.""" + + def __init__(self, *, interrupt_wait: bool = False) -> None: + super().__init__() + self._interrupt_wait = interrupt_wait + + def wait(self, timeout: float | None = None) -> bool: + if ( + self._interrupt_wait + and threading.current_thread() is threading.main_thread() + ): + raise KeyboardInterrupt + return super().wait(timeout) + + def event_factory() -> InterruptingEvent: + nonlocal events_created + events_created += 1 + # run_session creates opened, stop, then io_stopped in this order. + return InterruptingEvent(interrupt_wait=events_created == 3) + + monkeypatch.setattr(threading, "Event", event_factory) + log = CallLog() + + with pytest.raises(KeyboardInterrupt): + run_session( + FakeSession(_session_desc(), log), RecordingClientWindow(log), steps=0 + ) + + assert log.calls[-2:] == ["window.close", "session.close"] + assert not any(thread.name == _IO_THREAD_NAME for thread in threading.enumerate()) def test_run_session_gives_the_step_after_a_reset_the_whole_batch() -> None: @@ -522,6 +665,34 @@ def get_user_input_events(self) -> UserInputEvents: assert [result.step_index for result in window.results] == [0] +def test_run_session_drops_a_result_generated_during_disconnect() -> None: + log = CallLog() + disconnect_reported = threading.Event() + + class SlowStep(FakeSession): + def step(self, step_index: int, events: UserInputEvents) -> StepResult: + del events + disconnect_reported.wait() + return super().step(step_index, UserInputEvents([])) + + class DisconnectingWindow(RecordingClientWindow): + def get_user_input_events(self) -> UserInputEvents: + events = super().get_user_input_events() + if events.get_events(): + disconnect_reported.set() + return events + + session = SlowStep(_session_desc(), log) + window = DisconnectingWindow( + log, + [UserInputEvents([]), _lifecycle_event(CloseUserInputEventData())], + ) + + run_session(session, window, steps=None) + + assert window.results == [] + + def test_run_session_presents_every_result_when_blocking() -> None: log = CallLog() session = FakeSession(_session_desc(), log) diff --git a/flashdreams/test_v2/test_t2v_application.py b/flashdreams/test_v2/test_t2v_application.py index 7db998bfa..4d58cad9c 100644 --- a/flashdreams/test_v2/test_t2v_application.py +++ b/flashdreams/test_v2/test_t2v_application.py @@ -73,8 +73,8 @@ def close(self) -> None: class FakePipelineConfig: """Record how often the model was loaded, which is the expensive part.""" - def __init__(self) -> None: - self.pipeline = FakePipeline() + def __init__(self, pipeline: FakePipeline | None = None) -> None: + self.pipeline = pipeline if pipeline is not None else FakePipeline() self.setup_count = 0 def setup(self) -> FakePipeline: @@ -322,6 +322,24 @@ def test_closing_the_application_releases_the_model() -> None: assert config.pipeline.closed +def test_failed_model_initialization_remains_owned_for_cleanup() -> None: + class FailingPipeline(FakePipeline): + def eval(self) -> "FakePipeline": + super().eval() + raise RuntimeError("eval failed") + + pipeline = FailingPipeline() + app = T2VApplication( + defaults=_defaults(pipeline_config=FakePipelineConfig(pipeline)) + ) + + with pytest.raises(RuntimeError, match="eval failed"): + app.init(["--prompt", _PROMPT]) + app.close() + + assert pipeline.closed + + ## What a model will not generate diff --git a/flashdreams/test_v2/test_webrtc_client_window.py b/flashdreams/test_v2/test_webrtc_client_window.py index f5e2bd632..c853cb3c4 100644 --- a/flashdreams/test_v2/test_webrtc_client_window.py +++ b/flashdreams/test_v2/test_webrtc_client_window.py @@ -5,6 +5,7 @@ import asyncio import json +import threading from dataclasses import replace import pytest @@ -27,6 +28,7 @@ from flashdreams.runtime_v2.session_desc import SessionDesc from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import ( + CloseUserInputEventData, KeyboardUserInputEventData, NewSessionUserInputEventData, ) @@ -37,8 +39,8 @@ def _session_desc() -> SessionDesc: return SessionDesc( output_layout=VideoTensorLayout.tchw, - frames_per_second_for_ui=30, - frames_per_second_for_step=30, + frames_per_second_for_ui=60, + frames_per_second_for_step=20, video_width=16, video_height=16, ) @@ -46,7 +48,11 @@ def _session_desc() -> SessionDesc: async def _connect_browser( window: WebRTCClientWindow, -) -> tuple[RTCPeerConnection, RTCDataChannel, asyncio.Future[MediaStreamTrack]]: +) -> tuple[ + RTCPeerConnection, + RTCDataChannel, + asyncio.Future[MediaStreamTrack], +]: peer = RTCPeerConnection() channel = peer.createDataChannel("controls") peer.addTransceiver("video", direction="recvonly") @@ -65,16 +71,23 @@ def on_track(track: MediaStreamTrack) -> None: video_track.set_result(track) await peer.setLocalDescription(await peer.createOffer()) + deadline = asyncio.get_running_loop().time() + 5 async with ClientSession() as client: - async with client.post( - f"{window.server.url}api/webrtc/offer", - json={ - "sdp": peer.localDescription.sdp, - "type": peer.localDescription.type, - }, - ) as response: - assert response.status == 200 - answer = await response.json() + while True: + async with client.post( + f"{window.url}api/webrtc/offer", + json={ + "sdp": peer.localDescription.sdp, + "type": peer.localDescription.type, + }, + ) as response: + if response.status != 409: + assert response.status == 200 + answer = await response.json() + break + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError("The previous browser did not disconnect.") + await asyncio.sleep(0.01) await peer.setRemoteDescription( RTCSessionDescription(sdp=answer["sdp"], type=answer["type"]) ) @@ -88,20 +101,23 @@ async def test_window_buffers_browser_events_until_drained() -> None: peer: RTCPeerConnection | None = None try: async with ClientSession() as client: - async with client.get(f"{window.server.url}healthz") as response: + async with client.get(f"{window.url}healthz") as response: assert response.status == 200 assert await response.json() == { "open": False, "client_connected": False, } - async with client.get(window.server.url) as response: + async with client.get(window.url) as response: browser_page = await response.text() assert response.status == 200 assert 'id="activate"' in browser_page assert 'id="prompt"' in browser_page assert 'id="new-session" type="button">' in browser_page + assert ( + 'id="new-session" type="button">Opening...' in browser_page + ) assert '' in browser_page - async with client.get(f"{window.server.url}app.js") as response: + async with client.get(f"{window.url}app.js") as response: browser_script = await response.text() assert response.status == 200 assert 'key: "r", pressed: activationPressed' in browser_script @@ -109,6 +125,7 @@ async def test_window_buffers_browser_events_until_drained() -> None: assert "metadata: {prompt: promptInput.value}" in browser_script assert "pendingNewSession = request" in browser_script assert 'newSessionButton.textContent = "Opening..."' in browser_script + assert "event.target === promptInput" in browser_script assert "response.status !== 409" in browser_script window.open(_session_desc()) @@ -152,17 +169,23 @@ async def test_window_buffers_browser_events_until_drained() -> None: assert events[0].get_timestamp() <= events[1].get_timestamp() assert window.get_user_input_events().get_events() == [] + # Make the old timestamp epoch observably older than a fresh offer, then + # reopen between the old browser's close and the refreshed request. The + # drained FIFO order must survive that session boundary. + await asyncio.sleep(1) await peer.close() peer = None async with ClientSession() as client: - for _ in range(100): - async with client.get(f"{window.server.url}healthz") as response: + deadline = asyncio.get_running_loop().time() + 5 + while True: + async with client.get(f"{window.url}healthz") as response: health = await response.json() if not health["client_connected"]: break + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError("The closed browser was not released.") await asyncio.sleep(0.01) - assert health == {"open": True, "client_connected": False} - + window.open(replace(_session_desc(), metadata={"prompt": "next"})) peer, channel, _ = await _connect_browser(window) channel.send( json.dumps( @@ -182,13 +205,114 @@ async def test_window_buffers_browser_events_until_drained() -> None: break await asyncio.sleep(0.01) assert [ - event.get_event_data().metadata + event_data.metadata for event in refreshed_events - if isinstance(event.get_event_data(), NewSessionUserInputEventData) + if isinstance( + event_data := event.get_event_data(), + NewSessionUserInputEventData, + ) ] == [{"prompt": "A fox in a forest"}] + lifecycle_types = [ + type(event.get_event_data()) + for event in refreshed_events + if isinstance( + event.get_event_data(), + (CloseUserInputEventData, NewSessionUserInputEventData), + ) + ] + assert lifecycle_types == [ + CloseUserInputEventData, + NewSessionUserInputEventData, + ] finally: if peer is not None: await peer.close() + await asyncio.sleep(0.05) + window.close() + + +@pytest.mark.asyncio +async def test_overlapping_offer_connects_after_active_browser_releases() -> None: + window = WebRTCClientWindow() + first = RTCPeerConnection() + second = RTCPeerConnection() + try: + window.open(_session_desc()) + for peer in (first, second): + peer.createDataChannel("controls") + peer.addTransceiver("video", direction="recvonly") + await peer.setLocalDescription(await peer.createOffer()) + + async with ClientSession() as client: + + async def offer( + peer: RTCPeerConnection, + ) -> tuple[int, dict[str, str] | None]: + async with client.post( + f"{window.url}api/webrtc/offer", + json={ + "sdp": peer.localDescription.sdp, + "type": peer.localDescription.type, + }, + ) as response: + answer = await response.json() if response.status == 200 else None + return response.status, answer + + responses = await asyncio.gather(offer(first), offer(second)) + + assert sorted(status for status, _ in responses) == [200, 409] + admitted_index = next( + index for index, (status, _) in enumerate(responses) if status == 200 + ) + rejected_index = 1 - admitted_index + peers = (first, second) + admitted = peers[admitted_index] + rejected = peers[rejected_index] + answer = responses[admitted_index][1] + assert answer is not None + await admitted.setRemoteDescription( + RTCSessionDescription(sdp=answer["sdp"], type=answer["type"]) + ) + deadline = asyncio.get_running_loop().time() + 5 + while admitted.connectionState != "connected": + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError("The admitted peer did not connect.") + await asyncio.sleep(0.01) + + # The refreshed browser overlaps the old one and must be rejected + # while that old peer remains active. + status, _ = await offer(rejected) + assert status == 409 + + # Once the old page releases its peer, retrying that same pending + # offer must succeed within a bound. + await admitted.close() + deadline = asyncio.get_running_loop().time() + 5 + replacement_answer = None + while replacement_answer is None: + status, replacement_answer = await offer(rejected) + if status != 409: + assert status == 200 + break + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError("The refreshed browser did not connect.") + await asyncio.sleep(0.01) + assert replacement_answer is not None + await rejected.setRemoteDescription( + RTCSessionDescription( + sdp=replacement_answer["sdp"], + type=replacement_answer["type"], + ) + ) + deadline = asyncio.get_running_loop().time() + 5 + while rejected.connectionState != "connected": + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError("The refreshed browser did not connect.") + await asyncio.sleep(0.01) + finally: + await first.close() + await second.close() + await asyncio.sleep(0.05) window.close() @@ -227,8 +351,47 @@ async def test_write_delivers_a_video_frame_to_the_browser() -> None: metrics={}, ) ) + # One frame may already be inside the encoder when the replacement + # starts. The track discards the rest of the old source queue, so the + # replacement must arrive immediately after that in-flight frame. + replacement_pixels = None + for _ in range(2): + replacement_frame = await asyncio.wait_for(track.recv(), timeout=5) + assert isinstance(replacement_frame, VideoFrame) + pixels = replacement_frame.to_ndarray(format="rgb24") + if float(pixels.mean()) > 100: + replacement_pixels = pixels + break + assert replacement_pixels is not None + assert abs(float(replacement_pixels.mean()) - 211.0) <= 2.0 assert peer.connectionState == "connected" finally: if peer is not None: await peer.close() + await asyncio.sleep(0.05) window.close() + + +def test_an_open_peer_keeps_its_media_format_across_sessions() -> None: + window = WebRTCClientWindow() + try: + session_desc = _session_desc() + window.open(session_desc) + + # UI polling is a runtime concern and does not change the media track. + window.open(replace(session_desc, frames_per_second_for_ui=30)) + with pytest.raises(ValueError, match="original output"): + window.open(replace(session_desc, frames_per_second_for_step=10)) + finally: + window.close() + + +def test_closing_a_window_stops_its_server_thread() -> None: + window = WebRTCClientWindow() + + window.close() + window.close() + + assert not any( + thread.name == "flashdreams-webrtc" for thread in threading.enumerate() + ) diff --git a/integrations_v2/color_fade/color_fade/tests/test_color_fade.py b/integrations_v2/color_fade/color_fade/tests/test_color_fade.py index c0f2ee068..eaf46acf0 100644 --- a/integrations_v2/color_fade/color_fade/tests/test_color_fade.py +++ b/integrations_v2/color_fade/color_fade/tests/test_color_fade.py @@ -265,7 +265,7 @@ def test_a_run_writes_the_whole_fade_to_an_mp4(tmp_path: Path) -> None: runner.init( ["--seconds", str(_SECONDS), "--frames-per-step", str(frames_per_step)] ) - runner.run_session( + runner.run( SessionDescRequest( output_layout=VideoTensorLayout.bcthw, frames_per_second_for_step=_FRAMES_PER_SECOND, diff --git a/integrations_v2/red_screen/red_screen/app.py b/integrations_v2/red_screen/red_screen/app.py index 63cb78b15..28be4a69b 100644 --- a/integrations_v2/red_screen/red_screen/app.py +++ b/integrations_v2/red_screen/red_screen/app.py @@ -207,7 +207,7 @@ def main(commandline_args: Sequence[str] | None = None) -> int: app = create_app() runner = ApplicationRunner(app) if isinstance(window, WebRTCClientWindow): - print(f"Open {window.server.url} in a browser.", flush=True) + print(f"Open {window.url} in a browser.", flush=True) try: # ApplicationRunner is a FlashDreams runtime component that takes an IApplication instance, a IClientWindow instance, # and drives the main loop. @@ -215,7 +215,7 @@ def main(commandline_args: Sequence[str] | None = None) -> int: # TODO: in production, commandline argument parsing and IClientWindow creation should be done by flashdreams-run, a CLI tool # basically, we need to generailze this main function to be shared by all applications runner.init(application_args) - runner.run_session( + runner.run( SessionDescRequest( output_layout=VideoTensorLayout.bcthw, frames_per_second_for_ui=args.fps, diff --git a/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py b/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py index dea1f679c..5c7697210 100644 --- a/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py +++ b/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py @@ -9,6 +9,7 @@ """ import copy +from typing import Any import pytest from fastvideo_causal_wan22.config import RUNNER_WAN22_T2V_14B @@ -46,7 +47,7 @@ def test_compilation_is_turned_off_for_both_noise_level_transformers() -> None: model splits denoising across two transformers, and the shared override reaches only one of them, so it is overridden here. """ - pipeline_config = copy.deepcopy(RUNNER_WAN22_T2V_14B.pipeline) + pipeline_config: Any = copy.deepcopy(RUNNER_WAN22_T2V_14B.pipeline) def load_stand_in(_: object) -> FakeT2VPipeline: return FakeT2VPipeline() diff --git a/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py b/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py index fb5b66630..11ce0c015 100644 --- a/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py +++ b/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py @@ -12,6 +12,7 @@ import copy import shutil from pathlib import Path +from typing import Any import pytest from self_forcing.config import RUNNER_WAN21_T2V_1PT3B @@ -50,7 +51,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None: def test_compilation_can_be_turned_off_for_a_run() -> None: """Apply the override to the real config while loading a stand-in model.""" - pipeline_config = copy.deepcopy(RUNNER_WAN21_T2V_1PT3B.pipeline) + pipeline_config: Any = copy.deepcopy(RUNNER_WAN21_T2V_1PT3B.pipeline) def load_stand_in(_: object) -> FakeT2VPipeline: return FakeT2VPipeline() From 1df0053208d3c6cb392059dce59ef1b12990b506 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 21 Aug 2026 17:30:32 +0000 Subject: [PATCH 6/6] Unify v2 application session execution Signed-off-by: Gangzheng Tong --- .../developer_guides/v2_webrtc_lifecycle.md | 23 +-- .../flashdreams/api_v2/client_window.py | 30 +++- .../runtime_v2/application_runner.py | 89 ++++------- flashdreams/flashdreams/runtime_v2/cli.py | 9 +- .../runtime_v2/client_window_factory.py | 10 +- .../flashdreams/runtime_v2/session_runner.py | 16 +- .../runtime_v2/webrtc_client_window.py | 7 + .../test_v2/test_application_runner.py | 147 +++++++++++++++++- flashdreams/test_v2/test_cli.py | 13 +- .../test_v2/test_client_window_factory.py | 4 +- flashdreams/test_v2/test_session_runner.py | 18 +++ .../test_v2/test_webrtc_client_window.py | 1 + integrations_v2/red_screen/red_screen/app.py | 12 +- .../red_screen/tests/test_red_screen.py | 59 +++++++ 14 files changed, 327 insertions(+), 111 deletions(-) diff --git a/docs/source/developer_guides/v2_webrtc_lifecycle.md b/docs/source/developer_guides/v2_webrtc_lifecycle.md index 8afa48a18..973c75f59 100644 --- a/docs/source/developer_guides/v2_webrtc_lifecycle.md +++ b/docs/source/developer_guides/v2_webrtc_lifecycle.md @@ -32,12 +32,14 @@ remain available for the next request. Ctrl-C ends the application. | `T2VApplication` | The loaded model pipeline shared by every session | A rollout cache or browser connection | | `T2VSession` | One prompt and one rollout cache | Model loading or server lifetime | -The CLI constructs the window, then transfers its cleanup responsibility to -the runner. While the server is idle, the runner's calling thread opens and -polls the window. During a session, `run_session()` gives the window to one -dedicated I/O thread. The runner does not touch it again until that thread has -stopped and the old session has closed. This is a sequential ownership -handoff, not concurrent access. +The CLI constructs the window, then transfers its lifetime cleanup +responsibility to the runner. While the server is idle, the runner's calling +thread opens and polls the window. During a session, `run_session()` lends the +window to one dedicated I/O thread, which performs the meaningful close when +the window should end. The runner does not touch it again until that thread has +stopped and the old session has closed. Its final, idempotent close is a fallback +for setup failures and interrupted handoffs. This is sequential access, not +concurrent access. ## Data flow @@ -49,9 +51,10 @@ handoff, not concurrent access. and exposes the browser URL. 4. `ApplicationRunner.run()` resolves the CLI's partial `SessionDescRequest` against the initialized application's default `SessionDesc`. -5. In serving mode, the runner opens the window with that resolved stream - format before a session exists. This lets the browser negotiate WebRTC and - send its first request without loading a session cache first. +5. A window whose `keeps_open_between_sessions` capability is true opens with + that resolved stream format before a session exists. This lets the browser + negotiate WebRTC and send its first request without loading a session cache + first. Other windows start the resolved session immediately. 6. The server validates every data-channel message and invokes the callback registered by `WebRTCClientWindow`. The callback only appends the event to a thread-safe queue. @@ -67,7 +70,7 @@ handoff, not concurrent access. and writes completed results. The calling thread runs model steps. 10. A new-session event stops that rollout and returns the requested next `SessionDesc`. A close or natural completion returns no replacement. In - serving mode the runner keeps the window open, closes the old session, and + a persistent window the runner keeps it open, closes the old session, and either starts the replacement or waits for another browser request. 11. Ctrl-C closes the peer/server and then the application. Releasing the application drops the one resident pipeline after every session cache has diff --git a/flashdreams/flashdreams/api_v2/client_window.py b/flashdreams/flashdreams/api_v2/client_window.py index 07d65d4dc..b617d6d22 100644 --- a/flashdreams/flashdreams/api_v2/client_window.py +++ b/flashdreams/flashdreams/api_v2/client_window.py @@ -3,7 +3,7 @@ """Client window abstract interface.""" -from abc import ABC +from abc import ABC, abstractmethod from .input_source import InputSource from .output_sink import OutputSink @@ -16,14 +16,34 @@ class IClientWindow(InputSource, OutputSink, ABC): and writes results until the run ends. A window stays open across a session reset. When the client asks for a replacement session, the runtime closes the old session and opens the same window with the replacement's description. A - session-serving runner can also leave the window open between sessions while - it waits for another client request. + persistent window can also remain open between sessions while the runtime + waits for another client request. A window does not describe the output shape. The session does, and the window is given that description in :meth:`OutputSink.open`. - One I/O thread at a time makes every call on a window, so an implementation - needs no locking except when its backend delivers input from another thread. + One runtime thread at a time makes every call on a window, so an + implementation needs no locking except when its backend delivers input from + another thread. Created by the runtime, never by an application. """ + + keeps_open_between_sessions: bool = False + """Whether sessions start on demand and the window persists between them. + + When false, the runtime starts the resolved initial session immediately and + returns after it ends unless the client requested a replacement. When true, + the runtime opens the window before creating a session, waits for a client + request, and returns to waiting after completion or disconnection. + """ + + @abstractmethod + def close(self) -> None: + """Release this window's resources. + + This must be safe before :meth:`OutputSink.open` and after an earlier + call. The session loop performs the meaningful close on its I/O thread; + the application runner calls it again as a lifetime-cleanup fallback. + """ + ... diff --git a/flashdreams/flashdreams/runtime_v2/application_runner.py b/flashdreams/flashdreams/runtime_v2/application_runner.py index b6c1b1ea4..3b108ff7b 100644 --- a/flashdreams/flashdreams/runtime_v2/application_runner.py +++ b/flashdreams/flashdreams/runtime_v2/application_runner.py @@ -50,19 +50,14 @@ def run( self, session_desc_request: SessionDescRequest, client_window: IClientWindow, - *, - serve_sessions: bool = False, ) -> None: """Create sessions against ``client_window`` until the run ends. - Normally the run ends when the window reports a close or the session - reports that it has finished. A replacement description returned by the - session loop starts another session after the current one has closed. - - With ``serve_sessions``, the window opens before any session exists and - this method waits for a session description. It returns to that waiting - state whenever a session finishes or its browser disconnects. The server - therefore remains available until the process interrupts this method. + The window decides whether to start immediately with the resolved + description or stay open and wait for a client request. A replacement + description returned by the session loop starts another session after + the current one has closed. A persistent window returns to waiting when + a session finishes or its client disconnects. The session and window are closed before this method returns or raises. The application remains initialized, so callers can run another session @@ -73,27 +68,32 @@ def run( session_desc_request: Explicit overrides to apply to the application's initialized default description. client_window: Window that supplies input and presents generated output. - serve_sessions: Keep the window running and create sessions only in - response to client requests. """ - if serve_sessions: - self._serve_sessions(session_desc_request, client_window) - return - try: - next_session_desc = self._resolve_session_desc(session_desc_request) - except BaseException: - _close_client_window(client_window) - raise - while True: - try: + persistent_window = client_window.keeps_open_between_sessions + current_session_desc = self._resolve_session_desc(session_desc_request) + if persistent_window: + client_window.open(current_session_desc) + next_session_desc = wait_for_new_session( + client_window, current_session_desc + ) + else: + next_session_desc = current_session_desc + + while next_session_desc is not None: session = self._application.create_session(next_session_desc) - except BaseException: - _close_client_window(client_window) - raise - next_session_desc = run_session(session, client_window) - if next_session_desc is None: - return + current_session_desc = session.session_desc + next_session_desc = run_session( + session, + client_window, + keep_window_open=persistent_window, + ) + if persistent_window and next_session_desc is None: + next_session_desc = wait_for_new_session( + client_window, current_session_desc + ) + finally: + _close_client_window(client_window) def _resolve_session_desc( self, session_desc_request: SessionDescRequest @@ -106,31 +106,6 @@ def _resolve_session_desc( default = self._application.default_session_desc() or SessionDesc() return session_desc_request.resolve(default) - def _serve_sessions( - self, - session_desc_request: SessionDescRequest, - client_window: IClientWindow, - ) -> None: - """Keep one client window available for browser-requested sessions.""" - try: - current_session_desc = self._resolve_session_desc(session_desc_request) - client_window.open(current_session_desc) - next_session_desc: SessionDesc | None = None - while True: - if next_session_desc is None: - next_session_desc = wait_for_new_session( - client_window, current_session_desc - ) - session = self._application.create_session(next_session_desc) - current_session_desc = session.session_desc - next_session_desc = run_session( - session, - client_window, - keep_window_open=True, - ) - finally: - _close_client_window(client_window) - def close(self) -> None: """Release the application and the state it shares across sessions.""" if self._closed: @@ -140,10 +115,12 @@ def close(self) -> None: def _close_client_window(client_window: IClientWindow) -> None: - """Close a window during runner cleanup without hiding an active failure. + """Ensure a window closes without hiding the run's meaningful result. - This runs after session creation fails or a persistent run is interrupted, - so a failure here is logged rather than raised over the top of it. + ``IClientWindow.close`` is idempotent. The session loop normally performs + the first close on its I/O thread so file-finalization errors can fail the + run. This is the application runner's fallback for setup failures, + persistent windows, and interrupted ownership handoffs. """ try: client_window.close() diff --git a/flashdreams/flashdreams/runtime_v2/cli.py b/flashdreams/flashdreams/runtime_v2/cli.py index 14d7386d9..b958aa042 100644 --- a/flashdreams/flashdreams/runtime_v2/cli.py +++ b/flashdreams/flashdreams/runtime_v2/cli.py @@ -69,14 +69,9 @@ def entrypoint(argv: Sequence[str] | None = None) -> None: "The client window failed to close after startup failed." ) raise - # Nothing here says how long a session is: the application reports that. - # A serving mode stays up between sessions; a file mode runs just one. + # Nothing here says how long a session or window lives: each reports that. try: - runner.run( - _session_desc_request(parsed), - window, - serve_sessions=mode.serves_sessions, - ) + runner.run(_session_desc_request(parsed), window) except KeyboardInterrupt: return _report(mode.finished(window)) diff --git a/flashdreams/flashdreams/runtime_v2/client_window_factory.py b/flashdreams/flashdreams/runtime_v2/client_window_factory.py index 8decc572b..9e5edf400 100644 --- a/flashdreams/flashdreams/runtime_v2/client_window_factory.py +++ b/flashdreams/flashdreams/runtime_v2/client_window_factory.py @@ -27,9 +27,6 @@ class ClientWindowMode(ABC): name: str """What ``--mode`` calls this.""" - serves_sessions: bool = False - """Whether the window stays available between application sessions.""" - def add_arguments(self, parser: argparse.ArgumentParser) -> None: """Add the arguments this mode takes and no other does.""" @@ -101,7 +98,6 @@ class _WebRTCMode(ClientWindowMode): """Stream the run to a browser.""" name = "webrtc" - serves_sessions = True def add_arguments(self, parser: argparse.ArgumentParser) -> None: parser.add_argument( @@ -113,7 +109,11 @@ def create(self, parsed_args: argparse.Namespace) -> IClientWindow: # Imported here so a run writing a file needs none of the serving stack. from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow - return WebRTCClientWindow(host=parsed_args.host, port=parsed_args.port) + return WebRTCClientWindow( + host=parsed_args.host, + port=parsed_args.port, + keeps_open_between_sessions=True, + ) def starting(self, client_window: IClientWindow) -> str | None: """Return where to connect, which nobody can guess when the port is free.""" diff --git a/flashdreams/flashdreams/runtime_v2/session_runner.py b/flashdreams/flashdreams/runtime_v2/session_runner.py index ab070d2e0..016645016 100644 --- a/flashdreams/flashdreams/runtime_v2/session_runner.py +++ b/flashdreams/flashdreams/runtime_v2/session_runner.py @@ -359,8 +359,20 @@ def add_pending_result(result_generation: int, result: StepResult) -> int: continue return 0 - io_thread = threading.Thread(target=run_io, name="flashdreams-io") - io_thread.start() + try: + io_thread = threading.Thread(target=run_io, name="flashdreams-io") + io_thread.start() + except BaseException: + # Starting the I/O thread is part of starting the session. Nothing else + # can own either object when startup fails this early. + _close_session(session, run_failed=True) + try: + window.close() + except Exception: + _LOGGER.exception( + "The window failed to close after I/O thread startup failed." + ) + raise try: opened.wait() step_index = 0 diff --git a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py index 4e5d1660b..df4387dcb 100644 --- a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py +++ b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py @@ -16,12 +16,16 @@ class WebRTCClientWindow(IClientWindow): """Implement ``IClientWindow`` with WebRTC input and presentation.""" + keeps_open_between_sessions = False + """Default direct construction to one immediate session.""" + def __init__( self, *, host: str = "127.0.0.1", port: int = 0, startup_timeout_seconds: float = 10.0, + keeps_open_between_sessions: bool = False, ) -> None: """Create the WebRTC backend. @@ -32,7 +36,10 @@ def __init__( host: Interface on which the HTTP server listens. port: Listening port. Zero asks the operating system to choose one. startup_timeout_seconds: Maximum time to wait for server startup. + keeps_open_between_sessions: Wait for browser-requested sessions and + keep serving after each one. False runs one session immediately. """ + self.keeps_open_between_sessions = keeps_open_between_sessions self._input_events: queue.SimpleQueue[UserInputEvent] = queue.SimpleQueue() self._server = WebRTCServer( host=host, diff --git a/flashdreams/test_v2/test_application_runner.py b/flashdreams/test_v2/test_application_runner.py index 3660a5902..783658ea9 100644 --- a/flashdreams/test_v2/test_application_runner.py +++ b/flashdreams/test_v2/test_application_runner.py @@ -14,6 +14,7 @@ from flashdreams.api_v2.application import IApplication from flashdreams.api_v2.client_window import IClientWindow from flashdreams.api_v2.session import ISession +from flashdreams.runtime_v2 import application_runner as application_runner_module from flashdreams.runtime_v2.application_runner import ApplicationRunner from flashdreams.runtime_v2.session_desc import SessionDesc, SessionDescRequest from flashdreams.runtime_v2.step_result import StepResult @@ -78,11 +79,13 @@ def __init__( *, fail_to_init: bool = False, fail_to_close: bool = False, + fail_to_create_at: int | None = None, session_length: int | None = None, ) -> None: self._calls = calls self._fail_to_init = fail_to_init self._fail_to_close = fail_to_close + self._fail_to_create_at = fail_to_create_at self._session_length = session_length self.created_session_descs: list[SessionDesc] = [] @@ -93,6 +96,8 @@ def init(self, commandline_args: Sequence[str]) -> None: def create_session(self, session_desc: SessionDesc) -> ISession: self._calls.append("application.create_session") + if self._fail_to_create_at == len(self.created_session_descs): + raise RuntimeError("session creation failed") self.created_session_descs.append(session_desc) return _Session(session_desc, self._calls, length=self._session_length) @@ -107,6 +112,7 @@ def __init__(self, calls: list[str]) -> None: self._calls = calls self.results: list[StepResult] = [] self._reported_close = False + self._closed = False def get_user_input_events(self) -> UserInputEvents: if not self._reported_close: @@ -130,6 +136,9 @@ def write(self, result: StepResult) -> None: self._calls.append(f"window.write({result.step_index})") def close(self) -> None: + if self._closed: + return + self._closed = True self._calls.append("window.close") @@ -156,6 +165,8 @@ def get_user_input_events(self) -> UserInputEvents: class _ServingWindow(_Window): """Request two sessions, then interrupt the persistent runner.""" + keeps_open_between_sessions = True + def __init__(self, calls: list[str]) -> None: super().__init__(calls) self._prompts = ["A cat surfing", "A dog snowboarding"] @@ -253,7 +264,7 @@ def test_application_runner_replaces_a_session_from_window_metadata() -> None: assert calls.index("session.close") < second_creation -def test_application_runner_serves_sessions_until_it_is_interrupted() -> None: +def test_application_runner_keeps_a_persistent_window_between_sessions() -> None: calls: list[str] = [] application = _Application(calls, session_length=1) runner = ApplicationRunner(application) @@ -261,7 +272,7 @@ def test_application_runner_serves_sessions_until_it_is_interrupted() -> None: runner.init() with pytest.raises(KeyboardInterrupt): - runner.run(_session_desc_request(), window, serve_sessions=True) + runner.run(_session_desc_request(), window) assert [desc.metadata for desc in application.created_session_descs] == [ {"prompt": "A cat surfing"}, @@ -275,6 +286,130 @@ def test_application_runner_serves_sessions_until_it_is_interrupted() -> None: assert calls[-1] == "application.close" +def test_persistent_window_returns_to_waiting_after_a_client_disconnect() -> None: + class DisconnectingWindow(_ScriptedWindow): + keeps_open_between_sessions = True + + def get_user_input_events(self) -> UserInputEvents: + if not self._events: + raise KeyboardInterrupt + return super().get_user_input_events() + + def lifecycle_event( + timestamp: int, + event_data: CloseUserInputEventData | NewSessionUserInputEventData, + ) -> UserInputEvents: + return UserInputEvents( + [UserInputEvent(timestamp=uint64(timestamp), event_data=event_data)] + ) + + calls: list[str] = [] + application = _Application(calls) + window = DisconnectingWindow( + calls, + [ + lifecycle_event( + 0, NewSessionUserInputEventData(metadata={"prompt": "first"}) + ), + lifecycle_event(1, CloseUserInputEventData()), + lifecycle_event( + 2, NewSessionUserInputEventData(metadata={"prompt": "second"}) + ), + lifecycle_event(3, CloseUserInputEventData()), + ], + ) + runner = ApplicationRunner(application) + runner.init() + + with pytest.raises(KeyboardInterrupt): + runner.run(_session_desc_request(), window) + + assert [desc.metadata for desc in application.created_session_descs] == [ + {"prompt": "first"}, + {"prompt": "second"}, + ] + assert calls.count("session.close") == 2 + assert calls.count("window.open") == 3 + assert calls.count("window.close") == 1 + runner.close() + + +@pytest.mark.parametrize("persistent", [False, True]) +def test_application_runner_closes_the_window_when_session_creation_fails( + persistent: bool, +) -> None: + calls: list[str] = [] + application = _Application(calls, fail_to_create_at=0) + window = _ServingWindow(calls) if persistent else _Window(calls) + runner = ApplicationRunner(application) + runner.init() + + with pytest.raises(RuntimeError, match="session creation failed"): + runner.run(_session_desc_request(), window) + + assert calls.count("window.close") == 1 + runner.close() + + +def test_application_runner_closes_the_window_when_replacement_creation_fails() -> None: + calls: list[str] = [] + application = _Application(calls, fail_to_create_at=1) + window = _ScriptedWindow( + calls, + [ + UserInputEvents( + [ + UserInputEvent( + timestamp=uint64(0), + event_data=NewSessionUserInputEventData( + metadata={"prompt": "replacement"} + ), + ) + ] + ) + ], + ) + runner = ApplicationRunner(application) + runner.init() + + with pytest.raises(RuntimeError, match="session creation failed"): + runner.run(_session_desc_request(), window) + + assert calls.count("session.close") == 1 + assert calls.count("window.close") == 1 + runner.close() + + +def test_persistent_window_closes_if_session_handoff_is_interrupted( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def interrupt_after_cleanup( + session: ISession, + window: IClientWindow, + *, + keep_window_open: bool, + ) -> SessionDesc | None: + del window + assert keep_window_open + session.close() + raise KeyboardInterrupt + + monkeypatch.setattr( + application_runner_module, "run_session", interrupt_after_cleanup + ) + calls: list[str] = [] + runner = ApplicationRunner(_Application(calls)) + window = _ServingWindow(calls) + runner.init() + + with pytest.raises(KeyboardInterrupt): + runner.run(_session_desc_request(), window) + + assert calls.count("session.close") == 1 + assert calls.count("window.close") == 1 + runner.close() + + def test_application_runner_closes_the_window_when_a_session_cannot_start() -> None: calls: list[str] = [] runner = ApplicationRunner(_Application(calls)) @@ -289,15 +424,11 @@ def test_application_runner_closes_the_window_when_a_session_cannot_start() -> N def test_serving_closes_the_window_when_the_session_request_is_invalid() -> None: calls: list[str] = [] runner = ApplicationRunner(_Application(calls)) - window = _Window(calls) + window = _ServingWindow(calls) runner.init() with pytest.raises(ValueError, match="frames_per_second_for_ui"): - runner.run( - SessionDescRequest(frames_per_second_for_ui=0), - window, - serve_sessions=True, - ) + runner.run(SessionDescRequest(frames_per_second_for_ui=0), window) assert calls == ["application.init([])", "window.close"] runner.close() diff --git a/flashdreams/test_v2/test_cli.py b/flashdreams/test_v2/test_cli.py index 6a5c26b7e..2b288438b 100644 --- a/flashdreams/test_v2/test_cli.py +++ b/flashdreams/test_v2/test_cli.py @@ -277,16 +277,9 @@ def _write_application_module( class StubMode(ClientWindowMode): """A mode handing the command a window the test can look inside.""" - def __init__( - self, - name: str, - window: IClientWindow, - *, - serves_sessions: bool = False, - ) -> None: + def __init__(self, name: str, window: IClientWindow) -> None: self.name = name self._window = window - self.serves_sessions = serves_sessions def create(self, parsed_args: argparse.Namespace) -> IClientWindow: del parsed_args @@ -420,6 +413,8 @@ def test_a_browser_server_starts_without_a_command_line_prompt( monkeypatch: pytest.MonkeyPatch, ) -> None: class InterruptingWindow(RecordingWindow): + keeps_open_between_sessions = True + def __init__(self) -> None: super().__init__() self.closed = False @@ -436,7 +431,7 @@ def close(self) -> None: monkeypatch.setattr( cli, "client_window_mode", - lambda name: StubMode(name, window, serves_sessions=True), + lambda name: StubMode(name, window), ) cli.entrypoint(["stub", "--mode", "webrtc"]) diff --git a/flashdreams/test_v2/test_client_window_factory.py b/flashdreams/test_v2/test_client_window_factory.py index 46ce8267b..a0892f1c8 100644 --- a/flashdreams/test_v2/test_client_window_factory.py +++ b/flashdreams/test_v2/test_client_window_factory.py @@ -36,7 +36,7 @@ def test_a_run_goes_to_a_file_unless_it_says_otherwise(tmp_path: Path) -> None: assert parsed.mode == "mp4" assert isinstance(window, Mp4ClientWindow) - assert client_window_mode("mp4").serves_sessions is False + assert window.keeps_open_between_sessions is False def test_a_file_run_with_nowhere_to_write_says_so() -> None: @@ -75,6 +75,6 @@ def test_a_browser_run_is_told_where_to_connect(self) -> None: try: assert isinstance(window, WebRTCClientWindow) assert mode.starting(window) == f"Open {window.url} in a browser." - assert mode.serves_sessions is True + assert window.keeps_open_between_sessions is True finally: window.close() diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py index e30de98cd..58bff6c69 100644 --- a/flashdreams/test_v2/test_session_runner.py +++ b/flashdreams/test_v2/test_session_runner.py @@ -549,6 +549,24 @@ def init(self) -> None: assert log.calls == ["session.init", "session.close", "window.close"] +def test_run_session_closes_the_session_when_the_io_thread_cannot_start( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fail_to_start(thread: threading.Thread) -> None: + del thread + raise RuntimeError("thread start failed") + + monkeypatch.setattr(threading.Thread, "start", fail_to_start) + log = CallLog() + + with pytest.raises(RuntimeError, match="thread start failed"): + run_session( + FakeSession(_session_desc(), log), RecordingClientWindow(log), steps=1 + ) + + assert log.calls == ["session.init", "session.close", "window.close"] + + def test_run_session_closes_the_session_when_io_shutdown_is_interrupted( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/flashdreams/test_v2/test_webrtc_client_window.py b/flashdreams/test_v2/test_webrtc_client_window.py index c853cb3c4..2e1e299fe 100644 --- a/flashdreams/test_v2/test_webrtc_client_window.py +++ b/flashdreams/test_v2/test_webrtc_client_window.py @@ -100,6 +100,7 @@ async def test_window_buffers_browser_events_until_drained() -> None: window = WebRTCClientWindow() peer: RTCPeerConnection | None = None try: + assert window.keeps_open_between_sessions is False async with ClientSession() as client: async with client.get(f"{window.url}healthz") as response: assert response.status == 200 diff --git a/integrations_v2/red_screen/red_screen/app.py b/integrations_v2/red_screen/red_screen/app.py index 28be4a69b..620f92c79 100644 --- a/integrations_v2/red_screen/red_screen/app.py +++ b/integrations_v2/red_screen/red_screen/app.py @@ -13,7 +13,6 @@ from flashdreams.api_v2.application import IApplication from flashdreams.api_v2.session import ISession from flashdreams.runtime_v2.application_runner import ApplicationRunner -from flashdreams.runtime_v2.client_window_factory import create_client_window from flashdreams.runtime_v2.session_desc import SessionDesc, SessionDescRequest from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import KeyboardUserInputEventData @@ -203,17 +202,16 @@ def main(commandline_args: Sequence[str] | None = None) -> int: if application_args[:1] == ["--"]: application_args = application_args[1:] - window = create_client_window(args) + window = WebRTCClientWindow( + host=args.host, + port=args.port, + keeps_open_between_sessions=False, + ) app = create_app() runner = ApplicationRunner(app) if isinstance(window, WebRTCClientWindow): print(f"Open {window.url} in a browser.", flush=True) try: - # ApplicationRunner is a FlashDreams runtime component that takes an IApplication instance, a IClientWindow instance, - # and drives the main loop. - - # TODO: in production, commandline argument parsing and IClientWindow creation should be done by flashdreams-run, a CLI tool - # basically, we need to generailze this main function to be shared by all applications runner.init(application_args) runner.run( SessionDescRequest( diff --git a/integrations_v2/red_screen/red_screen/tests/test_red_screen.py b/integrations_v2/red_screen/red_screen/tests/test_red_screen.py index d6f666764..c4d30f503 100644 --- a/integrations_v2/red_screen/red_screen/tests/test_red_screen.py +++ b/integrations_v2/red_screen/red_screen/tests/test_red_screen.py @@ -8,6 +8,7 @@ import pytest import torch from numpy import uint64 +from red_screen import app as red_screen_app from red_screen import create_app from flashdreams.api_v2.client_window import IClientWindow @@ -16,6 +17,7 @@ from flashdreams.runtime_v2.session_runner import WhenFull, run_session from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import ( + CloseUserInputEventData, KeyboardUserInputEventData, UserInputEvent, ) @@ -135,6 +137,63 @@ def _new_session() -> ISession: ## Tests +def test_webrtc_entrypoint_runs_its_initial_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class EntrypointWindow(IClientWindow): + def __init__( + self, + *, + host: str, + port: int, + keeps_open_between_sessions: bool, + ) -> None: + assert host == "127.0.0.1" + assert port == 8080 + assert keeps_open_between_sessions is False + self.keeps_open_between_sessions = keeps_open_between_sessions + self.url = "http://127.0.0.1:8080" + self.opened = False + self.closed = False + self._polled = False + windows.append(self) + + def get_user_input_events(self) -> UserInputEvents: + if self._polled: + raise RuntimeError("entrypoint waited for a session request") + self._polled = True + return UserInputEvents( + [ + UserInputEvent( + timestamp=uint64(0), + event_data=CloseUserInputEventData(), + ) + ] + ) + + def open(self, session_desc: SessionDesc) -> None: + del session_desc + self.opened = True + + def write(self, result: StepResult) -> None: + del result + + def close(self) -> None: + self.closed = True + + windows: list[EntrypointWindow] = [] + monkeypatch.setattr(red_screen_app, "WebRTCClientWindow", EntrypointWindow) + + exit_code = red_screen_app.main( + ["--port", "8080", "--width", "2", "--height", "2", "--fps", "100"] + ) + + assert exit_code == 0 + assert len(windows) == 1 + assert windows[0].opened + assert windows[0].closed + + def test_red_screen_holds_red_between_key_edges() -> None: # Key down at step 0 and up at step 2. The step in between carries no events, # so it exercises held state rather than a repeated key-down.