diff --git a/flashdreams/flashdreams/api_v2/README.md b/flashdreams/flashdreams/api_v2/README.md index c59911174..125ed0229 100644 --- a/flashdreams/flashdreams/api_v2/README.md +++ b/flashdreams/flashdreams/api_v2/README.md @@ -12,7 +12,7 @@ Protocols for the FlashDreams API. in `OutputSink.open`. - `user_input_event_data.py`: base type for event payloads. -`flashdreams.runtime_v2.session_runner.run_session` drives a session against a +`flashdreams.runtime_v2.session_runner.SessionRunner.run_session` drives a session against a window until the session reports `is_finished` or the window reports a close, or for a fixed number of steps a caller asks for. A caller holding an application uses `flashdreams.runtime_v2.application_runner.ApplicationRunner` to get there, @@ -66,7 +66,7 @@ Agreed design decisions. Change them by discussion. Threading --------- -`run_session` uses two threads, and every window runs that way. Generation is on +`SessionRunner.run_session` uses two threads, and every window runs that way. Generation is on the calling thread; the window gets a thread of its own, ticking at `frames_per_second_for_ui` to read input, call `ISession.step_ui`, and write finished results. A step that takes longer than one of those ticks does not hold @@ -84,7 +84,7 @@ a native window needs, and it keeps `IClientWindow` implementations free of locking. Writing happens on that thread too, so a window slower than generation leaves -results waiting. `run_session` bounds how many wait, with `max_pending`, and +results waiting. `SessionRunner` bounds how many wait, with `max_pending`, and `when_full` decides the rest: `WhenFull.BLOCK` holds generation back so every result is presented, which is what a file output wants, and `WhenFull.DROP_OLDEST` skips frames to keep latency down, which is what a realtime one wants. The caller @@ -113,7 +113,7 @@ Not built yet fast steps can finish several of them between polls and hand them all the same batch. Pacing generation is what would fix it. - `ApplicationRunner`: takes an `IApplication` and an `IClientWindow` and drives - the main loop. `run_session` is what exists today; it drives a session the + the main loop. `SessionRunner` is what exists today; it drives a session the caller already created. - `flashdreams-run`: a CLI that creates the requested kind of client window, loads an application module, and hands both to `ApplicationRunner`. Until it diff --git a/flashdreams/flashdreams/runtime_v2/application_runner.py b/flashdreams/flashdreams/runtime_v2/application_runner.py index 45ad72b84..73b7022b4 100644 --- a/flashdreams/flashdreams/runtime_v2/application_runner.py +++ b/flashdreams/flashdreams/runtime_v2/application_runner.py @@ -10,7 +10,7 @@ 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.runtime_v2.session_runner import run_session +from flashdreams.runtime_v2.session_runner import SessionRunner _LOGGER = logging.getLogger(__name__) """Logger for an application that could not be closed.""" @@ -45,7 +45,7 @@ def run( try: self._application.init(commandline_args) session = self._application.create_session(session_desc) - run_session(session, self._client_window) + SessionRunner(session, self._client_window).run_session() finally: _close_application( self._application, run_failed=sys.exc_info()[0] is not None diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/app.js b/flashdreams/flashdreams/runtime_v2/serving/web/app.js index 029e714e7..351934735 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/web/app.js +++ b/flashdreams/flashdreams/runtime_v2/serving/web/app.js @@ -5,9 +5,11 @@ const peer = new RTCPeerConnection(); const controls = peer.createDataChannel("controls"); peer.addTransceiver("video", {direction: "recvonly"}); +const video = document.getElementById("video"); + peer.ontrack = event => { - document.getElementById("video").srcObject = - event.streams[0] ?? new MediaStream([event.track]); + video.srcObject = event.streams[0] ?? new MediaStream([event.track]); + video.play().catch(error => console.error("Unable to play WebRTC video", error)); }; const send = payload => { @@ -36,7 +38,19 @@ document.getElementById("reset").onclick = () => { send({type: "reset"}); }; -window.addEventListener("beforeunload", () => send({type: "close"})); + +async function waitForIceGathering() { + if (peer.iceGatheringState === "complete") { + return; + } + await new Promise(resolve => { + peer.addEventListener("icegatheringstatechange", () => { + if (peer.iceGatheringState === "complete") { + resolve(); + } + }); + }); +} async function connect() { while (true) { @@ -47,6 +61,7 @@ async function connect() { await new Promise(resolve => setTimeout(resolve, 100)); } await peer.setLocalDescription(await peer.createOffer()); + await waitForIceGathering(); const response = await fetch("/api/webrtc/offer", { method: "POST", headers: {"content-type": "application/json"}, diff --git a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py index 68041879d..ea3f2bef0 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 @@ -35,6 +36,8 @@ _WEB_RESOURCES = files("flashdreams.runtime_v2.serving").joinpath("web") _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__) +_DISCONNECT_GRACE_SECONDS = 5.0 class _VideoTrack(MediaStreamTrack): @@ -130,6 +133,7 @@ def __init__( self._runner: web.AppRunner | None = None self._peer_connection: RTCPeerConnection | None = None self._video_track: _VideoTrack | None = None + self._disconnect_task: asyncio.Task[None] | None = None self._session_desc: SessionDesc | None = None self._session_start_ns: int | None = None self._closed = False @@ -159,8 +163,14 @@ def port(self) -> int: @property def url(self) -> str: - """Return the browser URL for this server.""" - return f"http://{self._host}:{self._port}/" + """Return a browser URL for this server's local host. + + ``0.0.0.0`` and ``::`` listen on every interface but are not browser + destinations. Map them to the local loopback address for the URL shown + by standalone applications. + """ + host = "127.0.0.1" if self._host in {"0.0.0.0", "::"} else self._host + return f"http://{host}:{self._port}/" def open(self, session_desc: SessionDesc) -> None: """Configure the server for one session's generated video. @@ -278,14 +288,17 @@ async def _start_server(self) -> None: async def _serve_browser(self, _: web.Request) -> web.Response: """Return the minimal browser client.""" + _LOGGER.info("WebRTC endpoint called: GET /") return web.Response(text=_BROWSER_PAGE, content_type="text/html") async def _serve_browser_script(self, _: web.Request) -> web.Response: """Return the browser client's JavaScript.""" + _LOGGER.info("WebRTC endpoint called: GET /app.js") 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.""" + _LOGGER.info("WebRTC endpoint called: GET /healthz") return web.json_response( { "open": self._session_desc is not None, @@ -295,14 +308,12 @@ async def _health(self, _: web.Request) -> web.Response: async def _offer(self, request: web.Request) -> web.Response: """Negotiate one browser peer connection.""" + _LOGGER.info("WebRTC endpoint called: %s %s", request.method, request.path) if self._closed: raise web.HTTPServiceUnavailable(reason="WebRTC server is closed.") 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.") - try: payload = await request.json() except (json.JSONDecodeError, web.HTTPException) as error: @@ -319,12 +330,11 @@ async def _offer(self, request: web.Request) -> web.Response: 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 + if self._peer_connection is peer_connection: + self._client_connected = True @channel.on("message") def on_message(message: Any) -> None: @@ -335,12 +345,12 @@ def on_message(message: Any) -> None: @channel.on("close") def on_close() -> None: - self._record_client_disconnect() + _LOGGER.info("WebRTC browser controls channel closed.") @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) try: await peer_connection.setRemoteDescription( @@ -350,17 +360,23 @@ async def on_connectionstatechange() -> None: await peer_connection.createAnswer() ) except Exception: - self._peer_connection = None - self._video_track = None await video_track.close() await peer_connection.close() raise local_description = peer_connection.localDescription if local_description is None: + await video_track.close() + await peer_connection.close() raise web.HTTPInternalServerError( reason="WebRTC peer did not create an answer." ) + active_peer = self._peer_connection + if active_peer is not None: + await self._release_peer_connection(active_peer, close_peer=True) + self._cancel_disconnect_timer() + self._peer_connection = peer_connection + self._video_track = video_track return web.json_response( {"sdp": local_description.sdp, "type": local_description.type} ) @@ -377,6 +393,7 @@ def _buffer_browser_message(self, raw_message: object) -> None: raise ValueError("Browser event must be a JSON object.") event_type = payload.get("type") + _LOGGER.info("WebRTC browser event received: %s", event_type) if event_type == "keyboard": key = payload.get("key") pressed = payload.get("pressed") @@ -416,14 +433,49 @@ def _append_event( # The callback stores it in WebRTCClientWindow’s thread-safe queue. callback(event) - def _record_client_disconnect(self) -> None: - """Buffer one close event when the active browser disconnects.""" - if not self._client_connected: + async def _release_peer_connection( + self, + peer_connection: RTCPeerConnection, + *, + close_peer: bool = False, + ) -> None: + """Release the active peer so a browser can reconnect. + + A browser refresh closes its data channel and peer connection. That ends + only this transport connection; it does not end the session running on + the server. + """ + if self._peer_connection is not peer_connection: return self._client_connected = False + self._peer_connection = None + track = self._video_track + self._video_track = None + if track is not None: + await track.close() + if close_peer: + await peer_connection.close() if not self._closed: + self._disconnect_task = asyncio.create_task( + self._close_session_after_disconnect() + ) + + async def _close_session_after_disconnect(self) -> None: + """End a session only when no browser reconnects during the grace period.""" + try: + await asyncio.sleep(_DISCONNECT_GRACE_SECONDS) + except asyncio.CancelledError: + return + if self._peer_connection is None and not self._closed: + _LOGGER.info("WebRTC reconnect grace period expired; closing session.") self._append_event(CloseUserInputEventData()) + def _cancel_disconnect_timer(self) -> None: + """Cancel a pending session close after a browser reconnects.""" + if self._disconnect_task is not None: + self._disconnect_task.cancel() + self._disconnect_task = None + async def _enqueue_frames( self, frames: tuple[np.ndarray[Any, np.dtype[np.uint8]], ...] ) -> None: @@ -434,6 +486,7 @@ async def _enqueue_frames( async def _shutdown(self) -> None: """Release async server resources on their owning loop.""" + self._cancel_disconnect_timer() peer_connection = self._peer_connection self._peer_connection = None track = self._video_track diff --git a/flashdreams/flashdreams/runtime_v2/session_runner.py b/flashdreams/flashdreams/runtime_v2/session_runner.py index 5d013930b..3e7074ebe 100644 --- a/flashdreams/flashdreams/runtime_v2/session_runner.py +++ b/flashdreams/flashdreams/runtime_v2/session_runner.py @@ -1,12 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Step loop connecting one session to one client window.""" +"""Concurrent loops connecting one session to one client window.""" import logging import queue import sys import threading +from dataclasses import dataclass, field from enum import Enum from flashdreams.api_v2.client_window import IClientWindow @@ -21,17 +22,13 @@ from flashdreams.runtime_v2.user_input_events import UserInputEvents _LOGGER = logging.getLogger(__name__) -"""Logger for results a run could not present.""" class WhenFull(Enum): - """What to do with a finished result when no room is left to hold it.""" + """What to do with a finished result when the result queue is full.""" BLOCK = "block" - """Hold generation back until the window catches up, presenting every result.""" - DROP_OLDEST = "drop_oldest" - """Discard the oldest waiting result, skipping frames to keep latency down.""" def _contains(events: UserInputEvents, event_type: type[UserInputEventData]) -> bool: @@ -42,15 +39,14 @@ def _contains(events: UserInputEvents, event_type: type[UserInputEventData]) -> def _close_session(session: ISession, *, run_failed: bool) -> None: - """Close a session, keeping its close from hiding an earlier failure. + """Close a session without hiding an earlier failure. Args: session: Session to close. - run_failed: Whether something has already failed the run. When it has, - a failing close is logged rather than raised over the top of it. + run_failed: Whether another failure already explains the run. Raises: - Whatever the session raises, when nothing has failed yet. + Exception: Whatever :meth:`ISession.close` raises when the run succeeded. """ try: session.close() @@ -62,268 +58,190 @@ def _close_session(session: ISession, *, run_failed: bool) -> None: ) -def run_session( - session: ISession, - window: IClientWindow, - *, - steps: int | None = None, - max_pending: int = 2, - when_full: WhenFull = WhenFull.BLOCK, -) -> 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. - - 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. - - 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 - down when the client restarts is still held after, and dropping the edge that - said so would lose that. A session that must not inherit what the abandoned - generation was given has to ignore events older than its reset itself. - - Writing happens on the I/O thread, so a window slower than generation leaves - results waiting. ``max_pending`` bounds how many wait, and ``when_full`` says - what to do about the next one. - - A run whose output is a file works the same way, driven against - :class:`~flashdreams.runtime_v2.mp4_client_window.Mp4ClientWindow`. That - window reports no input, so it never reports a close, and such a run ends on - ``steps`` or on the session saying it has finished. - - A session says so through :meth:`ISession.is_finished`, asked before every - step, which is how a model that knows its own length ends its own run. The - run ends at whichever comes first: that, ``steps``, or a close. - - A window that fails to close fails the run, because for a file that means the - encode did not finish. A close that fails after something else already has is - logged instead, though, since the earlier failure is what explains the run. +@dataclass +class SessionRunner: + """State shared by the UI and generation loops for one session. + + The UI loop owns the window and runs on the I/O thread. The generation loop + runs on the caller's thread. The event batch and pending-result queue are the + two handoffs between them. Args: session: Uninitialized session to drive. window: Client window supplying input events and presenting results. - steps: Most steps to run, counted across resets so a reset cannot extend - the run. ``None`` runs until the session finishes or the window - reports a close, which is what an interactive window does. - 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. + steps: Maximum number of steps across resets, or ``None`` to run until + the session finishes or the window closes. + max_pending: Maximum number of results waiting to be presented. + when_full: Behavior when ``max_pending`` results are already waiting. + """ - Raises: - ValueError: ``steps`` is negative, or ``max_pending`` is not positive. + session: ISession + window: IClientWindow + steps: int | None = None + max_pending: int = 2 + when_full: WhenFull = WhenFull.BLOCK + tick_seconds: float = field(init=False) + pending_results: queue.Queue[tuple[int, StepResult]] = field(init=False) + collected_events: list[UserInputEvent] = field(default_factory=list) + input_events_lock: threading.Lock = field(default_factory=threading.Lock) + ui_startup_complete: threading.Event = field(default_factory=threading.Event) + shutdown_requested: threading.Event = field(default_factory=threading.Event) + io_failure: list[Exception] = field(default_factory=list) + generation: int = 0 + dropped_for_space: int = 0 + discarded_at_reset: int = 0 + + def run_session(self) -> None: + """Initialize and run the session until it reaches an end condition. + + Raises: + ValueError: ``steps`` is negative, or ``max_pending`` is not positive. + """ + if self.steps is not None and self.steps < 0: + raise ValueError(f"steps must be >= 0 or None, got {self.steps}.") + if self.max_pending <= 0: + raise ValueError(f"max_pending must be > 0, got {self.max_pending}.") - Note: - ``step`` and ``step_ui`` run at the same time, so a session implementing - both must guard what they share. - """ - if steps is not None and steps < 0: - raise ValueError(f"steps must be >= 0 or None, got {steps}.") - if max_pending <= 0: - raise ValueError(f"max_pending must be > 0, got {max_pending}.") + try: + self.session.init() + except Exception: + _close_session(self.session, run_failed=True) + raise - # SessionDesc guarantees this is positive. - tick_seconds = 1.0 / session.session_desc.frames_per_second_for_ui + self.tick_seconds = 1.0 / self.session.session_desc.frames_per_second_for_ui + self.pending_results = queue.Queue(maxsize=self.max_pending) + # The runner has two major loops: the UI loop owns the window on its + # thread, while the calling thread runs the generation loop below. + io_thread = threading.Thread(target=self._run_ui_loop, name="flashdreams-io") + io_thread.start() + try: + self._run_generation_loop() + finally: + self.shutdown_requested.set() + io_thread.join() + run_failed = sys.exc_info()[0] is not None + if self.io_failure and run_failed: + _LOGGER.error( + "The window failed as well as the run, and this is that failure.", + exc_info=self.io_failure[0], + ) + _close_session( + self.session, + run_failed=run_failed or bool(self.io_failure), + ) - try: - session.init() - except Exception: - # A partly initialized session still holds whatever it managed to load. - _close_session(session, run_failed=True) - raise - - # Backpressure is all here. Finished results wait here for the I/O thread to - # write, and once max_pending of them are waiting, generation blocks in - # add_pending_result, or drops the oldest result when asked to instead. - # Each result carries the generation it was produced for, so a reset can tell - # what belongs to the run the client abandoned from what belongs to the new one. - pending_results: queue.Queue[tuple[int, StepResult]] = queue.Queue( - maxsize=max_pending - ) - generation = 0 - # Only the I/O thread may read the window, so input waits here for the next step. - collected_events: list[UserInputEvent] = [] - collected_events_lock = threading.Lock() - opened = threading.Event() - stop = threading.Event() - io_failure: list[Exception] = [] - # What never reached the window, reported once the run is over. - dropped_for_space = 0 - discarded_at_reset = 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. - """ - nonlocal discarded_at_reset - while True: - try: - result_generation, result = pending_results.get_nowait() - except queue.Empty: + if self.dropped_for_space: + _LOGGER.warning( + "Dropped %d results the window could not keep up with.", + self.dropped_for_space, + ) + if self.discarded_at_reset: + _LOGGER.info( + "Discarded %d results generated before a reset.", + self.discarded_at_reset, + ) + if self.io_failure: + raise self.io_failure[0] + + def _run_generation_loop(self) -> None: + """Continuously generate results until the run reaches an end condition.""" + self.ui_startup_complete.wait() + step_index = 0 + steps_run = 0 + while self.steps is None or steps_run < self.steps: + if self.io_failure or self.shutdown_requested.is_set(): return - if result_generation != generation: - discarded_at_reset += 1 - continue - window.write(result) - - def tick() -> None: - nonlocal generation - events = window.get_user_input_events() - 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. + events, result_generation = self._take_events() if _contains(events, ResetUserInputEventData): - 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): - stop.set() - session.step_ui(events) - present_pending_results() + self.session.reset() + step_index = 0 + if self.session.is_finished(): + return + result = self.session.step(step_index, events) + self.dropped_for_space += self._queue_result(result_generation, result) + step_index += 1 + steps_run += 1 - def run_io() -> None: + def _run_ui_loop(self) -> None: + """Continuously collect input and present results on the window thread.""" try: - window.open(session.session_desc) - # Collect once before generation starts, so the first step sees input - # the window already has. - tick() - opened.set() - # Stop as soon as the wait says to. Ticking once more would poll input - # the stopped run can no longer act on, and a reset in that poll would - # throw away the results it just finished. - while not stop.wait(tick_seconds): - tick() - # Present anything the final step produced after the last tick. - present_pending_results() + self.window.open(self.session.session_desc) + # A first tick ensures the first generation step sees queued input. + self._tick() + self.ui_startup_complete.set() + while not self.shutdown_requested.wait(self.tick_seconds): + self._tick() + self._present_results() except Exception as error: - io_failure.append(error) + self.io_failure.append(error) finally: - opened.set() + self.ui_startup_complete.set() try: - window.close() + self.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) - - def take_collected_events() -> tuple[UserInputEvents, int]: - """Take the input waiting for the next step, and the generation it is for.""" - with collected_events_lock: - events = UserInputEvents(list(collected_events)) - collected_events.clear() - return events, generation - - def add_pending_result(result_generation: int, result: StepResult) -> int: - """Hand one result to the I/O thread, applying ``when_full``. + self.io_failure.append(error) + + def _tick(self) -> None: + """Collect one input batch, update the UI, and present ready results.""" + events = self.window.get_user_input_events() + with self.input_events_lock: + self.collected_events.extend(events.get_events()) + # The UI loop observes reset first. Locking ties a result to either + # the generation before the reset or the one after it, never both. + if _contains(events, ResetUserInputEventData): + self.generation += 1 + if _contains(events, CloseUserInputEventData): + self.shutdown_requested.set() + self.session.step_ui(events) + self._present_results() + + def _take_events(self) -> tuple[UserInputEvents, int]: + """Take all events collected since the previous generation step.""" + with self.input_events_lock: + events = UserInputEvents(list(self.collected_events)) + self.collected_events.clear() + return events, self.generation + + def _present_results(self) -> None: + """Write every queued result that belongs to the current generation.""" + while True: + try: + result_generation, result = self.pending_results.get_nowait() + except queue.Empty: + return + if result_generation != self.generation: + self.discarded_at_reset += 1 + continue + self.window.write(result) - This is where backpressure reaches generation, since no room for a result - is the only thing that ever slows this thread down. + def _queue_result(self, result_generation: int, result: StepResult) -> int: + """Queue a result, applying the configured backpressure policy. Args: - result_generation: Generation this result was produced for. - result: Finished result to hand over. + result_generation: Generation that produced ``result``. + result: Finished result to present. Returns: - How many waiting results were dropped to make room. + The number of queued results dropped to make room. """ pending = (result_generation, result) - if when_full is WhenFull.DROP_OLDEST: - # Keep the newest and lose the stale, so a client behind a slow window - # sees the present rather than catching up through the backlog. + if self.when_full is WhenFull.DROP_OLDEST: dropped = 0 while True: try: - pending_results.put_nowait(pending) + self.pending_results.put_nowait(pending) return dropped except queue.Full: try: - pending_results.get_nowait() + self.pending_results.get_nowait() dropped += 1 except queue.Empty: continue - # Otherwise let the window set the pace: generation waits here until there - # is room, which is what an output that must keep every frame needs. Wait - # interruptibly, though, because once the I/O thread is gone nothing writes - # the waiting results and a plain put would never return. - while not (stop.is_set() or io_failure): + + while not (self.shutdown_requested.is_set() or self.io_failure): try: - pending_results.put(pending, timeout=tick_seconds) + self.pending_results.put(pending, timeout=self.tick_seconds) break except queue.Full: continue return 0 - - io_thread = threading.Thread(target=run_io, name="flashdreams-io") - io_thread.start() - try: - opened.wait() - step_index = 0 - # A caller passing steps gets exactly that many, which is what a test or a - # fixed-length output needs. Bounding on step_index would break that, since - # a reset takes it back to zero: each reset would extend the run, and - # enough of them would stop it ever ending. - steps_run = 0 - while steps is None or steps_run < steps: - if io_failure or stop.is_set(): - break - events, step_generation = take_collected_events() - if _contains(events, ResetUserInputEventData): - session.reset() - step_index = 0 - # After the reset, so a session starting over is asked about the new - # run rather than the one it just finished. - if session.is_finished(): - break - dropped_for_space += add_pending_result( - step_generation, session.step(step_index, events) - ) - step_index += 1 - steps_run += 1 - finally: - stop.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)) - - # 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: - _LOGGER.info( - "Discarded %d results generated before a reset.", discarded_at_reset - ) - - if io_failure: - raise io_failure[0] diff --git a/flashdreams/test_v2/test_mp4_client_window.py b/flashdreams/test_v2/test_mp4_client_window.py index 4d4c6a750..b335c5e3d 100644 --- a/flashdreams/test_v2/test_mp4_client_window.py +++ b/flashdreams/test_v2/test_mp4_client_window.py @@ -18,7 +18,7 @@ from flashdreams.api_v2.session import ISession from flashdreams.runtime_v2.mp4_client_window import Mp4ClientWindow from flashdreams.runtime_v2.session_desc import SessionDesc -from flashdreams.runtime_v2.session_runner import run_session +from flashdreams.runtime_v2.session_runner import SessionRunner 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 @@ -117,7 +117,7 @@ def test_a_run_writing_a_file_encodes_a_step_at_a_time(tmp_path: Path) -> None: session = FakeSession(_session_desc()) path = tmp_path / "clip.mp4" - run_session(session, Mp4ClientWindow(path), steps=3) + SessionRunner(session, Mp4ClientWindow(path), steps=3).run_session() assert _frame_count(path) == 3 * _FRAMES_PER_STEP assert [events.get_events() for events in session.observed_events] == [[], [], []] diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py index 7928aa412..0079f84cb 100644 --- a/flashdreams/test_v2/test_session_runner.py +++ b/flashdreams/test_v2/test_session_runner.py @@ -14,7 +14,7 @@ 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 SessionRunner, WhenFull from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import ( CloseUserInputEventData, @@ -34,6 +34,24 @@ """Logger the runner reports discarded results on.""" +def run_session( + session: ISession, + window: IClientWindow, + *, + steps: int | None = None, + max_pending: int = 2, + when_full: WhenFull = WhenFull.BLOCK, +) -> None: + """Run the public session-runner API with test-specific arguments.""" + SessionRunner( + session, + window, + steps=steps, + max_pending=max_pending, + when_full=when_full, + ).run_session() + + class CallLog: """Record calls made from either thread, with the thread that made them.""" diff --git a/flashdreams/test_v2/test_webrtc_client_window.py b/flashdreams/test_v2/test_webrtc_client_window.py index b7f119651..2a6c25395 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 logging import pytest import torch @@ -25,10 +26,15 @@ 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 ( + CloseUserInputEventData, + KeyboardUserInputEventData, +) from flashdreams.runtime_v2.video_tensor import VideoTensorLayout from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow +_WEBRTC_SERVER_LOGGER = "flashdreams.runtime_v2.serving.webrtc_server" + def _session_desc() -> SessionDesc: return SessionDesc( @@ -79,7 +85,10 @@ def on_track(track: MediaStreamTrack) -> None: @pytest.mark.asyncio -async def test_window_buffers_browser_events_until_drained() -> None: +async def test_window_buffers_browser_events_until_drained( + caplog: pytest.LogCaptureFixture, +) -> None: + caplog.set_level(logging.INFO, logger=_WEBRTC_SERVER_LOGGER) window = WebRTCClientWindow() peer: RTCPeerConnection | None = None try: @@ -124,6 +133,11 @@ 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() == [] + assert "WebRTC endpoint called: GET /healthz" in caplog.messages + assert "WebRTC endpoint called: GET /" in caplog.messages + assert "WebRTC endpoint called: GET /app.js" in caplog.messages + assert "WebRTC endpoint called: POST /api/webrtc/offer" in caplog.messages + assert caplog.messages.count("WebRTC browser event received: keyboard") == 2 finally: if peer is not None: await peer.close() @@ -138,6 +152,12 @@ async def test_write_delivers_a_video_frame_to_the_browser() -> None: window.open(_session_desc()) peer, _, video_track = await _connect_browser(window) track = await asyncio.wait_for(video_track, timeout=5) + async with ClientSession() as client: + async with client.post( + f"{window.server.url}api/webrtc/offer", + json={"type": "offer"}, + ) as response: + assert response.status == 400 window.write( StepResult( @@ -158,3 +178,60 @@ async def test_write_delivers_a_video_frame_to_the_browser() -> None: if peer is not None: await peer.close() window.close() + + +@pytest.mark.asyncio +async def test_window_closes_the_session_after_a_browser_disconnect() -> None: + window = WebRTCClientWindow() + peer: RTCPeerConnection | None = None + try: + window.open(_session_desc()) + peer, _, _ = await _connect_browser(window) + await peer.close() + + for _ in range(60): + events = window.get_user_input_events().get_events() + if any( + isinstance(event.get_event_data(), CloseUserInputEventData) + for event in events + ): + break + await asyncio.sleep(0.1) + else: + pytest.fail("Browser disconnect did not close the session.") + finally: + if peer is not None: + await peer.close() + window.close() + + +@pytest.mark.asyncio +async def test_window_accepts_a_browser_reconnect() -> None: + window = WebRTCClientWindow() + first_peer: RTCPeerConnection | None = None + second_peer: RTCPeerConnection | None = None + try: + window.open(_session_desc()) + first_peer, _, _ = await _connect_browser(window) + await first_peer.close() + + second_peer, _, video_track = await _connect_browser(window) + track = await asyncio.wait_for(video_track, timeout=5) + window.write( + StepResult( + step_index=0, + output=torch.full((2, 3, 16, 16), 17, dtype=torch.uint8), + frame_count=2, + output_layout=VideoTensorLayout.tchw, + metrics={}, + ) + ) + + frame = await asyncio.wait_for(track.recv(), timeout=5) + assert isinstance(frame, VideoFrame) + finally: + if first_peer is not None: + await first_peer.close() + if second_peer is not None: + await second_peer.close() + window.close() diff --git a/integrations_v2/red_screen/README.md b/integrations_v2/red_screen/README.md index 29bc2cd14..455971526 100644 --- a/integrations_v2/red_screen/README.md +++ b/integrations_v2/red_screen/README.md @@ -7,7 +7,7 @@ SPDX-License-Identifier: Apache-2.0 Smallest end-to-end application on the v2 API. It holds no model: a session emits red frames controlled by activation and intensity keys. It runs the whole path — -`IApplication`, `ISession`, `run_session`, `IClientWindow` — on CPU. +`IApplication`, `ISession`, `SessionRunner`, `IClientWindow` — on CPU. ## What it demonstrates @@ -15,7 +15,7 @@ red frames controlled by activation and intensity keys. It runs the whole path else. This one never names `IClientWindow`, `InputSource` or `OutputSink`. - A session is created from a `SessionDesc` with no client window involved, so it works before any client connects. It reports what it resolved to in - `ISession.session_desc`, which `run_session` hands to `IClientWindow.open`. + `ISession.session_desc`, which `SessionRunner` hands to `IClientWindow.open`. - An application can reject a description it cannot honour: this one raises on any layout other than `bcthw`. - The runner drives the loop and hands the session each `step_index`, so the @@ -53,7 +53,7 @@ The same application can be driven directly without WebRTC: ```python from flashdreams.runtime_v2.session_desc import SessionDesc -from flashdreams.runtime_v2.session_runner import run_session +from flashdreams.runtime_v2.session_runner import SessionRunner from flashdreams.runtime_v2.video_tensor import VideoTensorLayout from red_screen import create_app @@ -66,7 +66,7 @@ session = app.create_session( video_height=16, ) ) -run_session(session, my_client_window, steps=4) +SessionRunner(session, my_client_window, steps=4).run_session() app.close() ``` 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..d082a46cc 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 @@ -13,7 +13,7 @@ from flashdreams.api_v2.client_window import IClientWindow from flashdreams.api_v2.session import ISession 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 SessionRunner, WhenFull from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import ( KeyboardUserInputEventData, @@ -118,7 +118,7 @@ def _run( session = app.create_session(_session_desc()) window = ScriptedClientWindow([initial_events] if initial_events else None) try: - run_session(session, window, steps=steps) + SessionRunner(session, window, steps=steps).run_session() finally: app.close() return window @@ -211,7 +211,13 @@ def test_red_screen_turns_red_for_a_key_pressed_during_the_run() -> None: # previous one, and every tick polls input before it presents. That makes the # key reach a step rather than depending on how the threads are scheduled. try: - run_session(session, window, steps=3, max_pending=1, when_full=WhenFull.BLOCK) + SessionRunner( + session, + window, + steps=3, + max_pending=1, + when_full=WhenFull.BLOCK, + ).run_session() finally: app.close()