diff --git a/.changeset/violet-clocks-reply.md b/.changeset/violet-clocks-reply.md new file mode 100644 index 0000000000..7078c803a4 --- /dev/null +++ b/.changeset/violet-clocks-reply.md @@ -0,0 +1,5 @@ +--- +"gradio": patch +--- + +fix:Fix streaming run key colliding across sequential runs diff --git a/gradio/blocks.py b/gradio/blocks.py index 456b163fb5..895be2cc28 100644 --- a/gradio/blocks.py +++ b/gradio/blocks.py @@ -13,6 +13,7 @@ import sys import threading import time +import uuid import warnings import weakref import webbrowser @@ -1115,6 +1116,13 @@ def __init__( self.max_threads = 40 self.pending_streams = defaultdict(dict) self.pending_diff_streams = defaultdict(dict) + # Per-run keys for streaming outputs, held weakly against the iterator + # so that a finished run's key goes away with it. The iterators are + # what call_function hands back, a generator, an async generator or a + # SyncToAsyncIterator, all weak-referenceable and hashed by identity. + self._stream_run_ids: weakref.WeakKeyDictionary[Any, str] = ( + weakref.WeakKeyDictionary() + ) self.show_error = True self.fill_height = fill_height self.fill_width = fill_width @@ -2124,20 +2132,60 @@ async def postprocess_data( return output + def _stream_run_key(self, iterator: Any) -> str: + """Return the key of the streaming run that `iterator` is driving. + + The key goes into the playlist URL, so it has to hold for every chunk of + a run and never repeat. `id()` only holds while the object is alive, so + the key is held weakly against the iterator and dies with it instead. + """ + run = self._stream_run_ids.get(iterator) + if run is None: + run = uuid.uuid4().hex + self._stream_run_ids[iterator] = run + return run + + def _drop_run_streams(self, session_hash: str | None, iterator: Any) -> None: + """Close out the streaming state of the run `iterator` was driving. + + For a run that reaches no final chunk: it raised, was cancelled, or its + client went away. Its streams are ended but stay, since the playlist is + fetched after the run ends; its diff state goes, nothing reads it again. + """ + if session_hash is None or iterator is None: + return + run = self._stream_run_ids.get(iterator) + if run is None: + return + for stream in self.pending_streams.get(session_hash, {}).get(run, {}).values(): + stream.end_stream() + self._pop_run_diffs(session_hash, run) + + def _pop_run_diffs(self, session_hash: str, run: str) -> None: + """Drop a run's diff state, and its session's dict if that leaves it empty.""" + runs = self.pending_diff_streams.get(session_hash) + if runs is None: + return + runs.pop(run, None) + if not runs: + del self.pending_diff_streams[session_hash] + async def handle_streaming_outputs( self, block_fn: BlockFunction, data: list, session_hash: str | None, - run: int | None, + run: str | None, root_path: str | None = None, final: bool = False, ) -> list: if session_hash is None or run is None: return data - if run not in self.pending_streams[session_hash]: - self.pending_streams[session_hash][run] = {} - stream_run: dict[int, MediaStream] = self.pending_streams[session_hash][run] + # Filed only once an output opens a stream, so a run with no streaming + # output never touches this dict + stream_run: dict[int, MediaStream] = self.pending_streams.get( + session_hash, {} + ).get(run, {}) for i, block in enumerate(block_fn.outputs): output_id = block._id @@ -2165,10 +2213,10 @@ async def handle_streaming_outputs( desired_output_format = None if orig_name := output_data.get("orig_name"): desired_output_format = Path(orig_name).suffix[1:] + stream_run = self.pending_streams[session_hash].setdefault(run, {}) stream_run[output_id] = MediaStream( desired_output_format=desired_output_format ) - stream_run[output_id] await stream_run[output_id].add_segment(binary_data) output_data = await processing_utils.async_move_files_to_cache( @@ -2189,7 +2237,7 @@ def handle_streaming_diffs( block_fn: BlockFunction, data: list, session_hash: str | None, - run: int | None, + run: str | None, final: bool, simple_format: bool = False, ) -> list: @@ -2218,7 +2266,7 @@ def handle_streaming_diffs( data[i] = utils.diff(prev_chunk, data[i]) if final: - del self.pending_diff_streams[session_hash][run] + self._pop_run_diffs(session_hash, run) return data @@ -2357,7 +2405,11 @@ async def process_api( data = processing_utils.add_root_url(data, root_path, None) is_generating, iterator = result["is_generating"], result["iterator"] if is_generating or was_generating: - run = id(old_iterator) if was_generating else id(iterator) + run = ( + self._stream_run_key(old_iterator if was_generating else iterator) + if session_hash is not None + else None + ) async with trace_phase("streaming_diff"): data = await self.handle_streaming_outputs( block_fn, @@ -2367,11 +2419,14 @@ async def process_api( root_path=root_path, final=not is_generating, ) + # Diff state serves the later chunks of a run, which can + # only be fetched under an event id. A call without one + # gets full values, which is what its clients expect. data = self.handle_streaming_diffs( block_fn, data, session_hash=session_hash, - run=run, + run=run if event_id is not None else None, final=not is_generating, simple_format=simple_format, ) diff --git a/gradio/queueing.py b/gradio/queueing.py index 53b1961740..8257152701 100644 --- a/gradio/queueing.py +++ b/gradio/queueing.py @@ -1088,6 +1088,18 @@ async def process_events( # without putting the `events` into `self.active_jobs`. # https://github.com/gradio-app/gradio/blob/f09aea34d6bd18c1e2fef80c86ab2476a6d1dd83/gradio/routes.py#L594-L596 pass + app = self.server_app + if app is not None: + blocks = app.get_blocks() + for event in events: + # A run that raised, was cancelled or lost its client reaches + # no final chunk, so close out its streams here, while the + # iterator that keys them is still stored (a finished run's + # is already None). /cancel awaits this task before dropping + # that iterator itself. + blocks._drop_run_streams( + event.session_hash, app.iterators.get(event._id) + ) for event in events: # Always reset the state of the iterator # If the job finished successfully, this has no effect diff --git a/gradio/route_utils.py b/gradio/route_utils.py index 87b364ede1..a17935a016 100644 --- a/gradio/route_utils.py +++ b/gradio/route_utils.py @@ -435,13 +435,7 @@ async def call_process_api( raise output except BaseException: iterator = app.iterators.get(event_id) if event_id is not None else None - if iterator is not None: # close off any streams that are still open - run_id = id(iterator) - pending_streams: dict[int, MediaStream] = ( - app.get_blocks().pending_streams.get(session_hash, {}).get(run_id, {}) - ) - for stream in pending_streams.values(): - stream.end_stream() + app.get_blocks()._drop_run_streams(session_hash, iterator) raise if batch_in_single_out: diff --git a/gradio/routes.py b/gradio/routes.py index 396437f1a4..799d2a8c6d 100644 --- a/gradio/routes.py +++ b/gradio/routes.py @@ -1217,7 +1217,7 @@ async def _(event_id: str): return {"msg": "success"} @router.get("/stream/{session_hash}/{run}/{component_id}/playlist.m3u8") - async def _(session_hash: str, run: int, component_id: int): + async def _(session_hash: str, run: str, component_id: int): stream: route_utils.MediaStream | None = ( app.get_blocks() .pending_streams.get(session_hash, {}) @@ -1247,7 +1247,7 @@ async def _(session_hash: str, run: int, component_id: int): @router.get("/stream/{session_hash}/{run}/{component_id}/{segment_id}.{ext}") async def _( - session_hash: str, run: int, component_id: int, segment_id: str, ext: str + session_hash: str, run: str, component_id: int, segment_id: str, ext: str ): if ext not in ["aac", "ts"]: return Response(status_code=400, content="Unsupported file extension") @@ -1272,7 +1272,7 @@ async def _( return Response(content=segment["data"], media_type="video/MP2T") @router.get("/stream/{session_hash}/{run}/{component_id}/playlist-file") - async def _(session_hash: str, run: int, component_id: int): + async def _(session_hash: str, run: str, component_id: int): stream: route_utils.MediaStream | None = ( app.get_blocks() .pending_streams.get(session_hash, {}) @@ -1366,6 +1366,8 @@ async def iterator(): if session_hash in app.state_holder.session_data: app.state_holder.session_data[session_hash].is_closed = True caching.clear_session_caches(session_hash) + # Streams only; diff state is dropped by the queue when + # the run ends for run in ( app.get_blocks() .pending_streams.pop(session_hash, {}) diff --git a/gradio/utils.py b/gradio/utils.py index 9fe2ad3dc0..a8cce3c5e7 100644 --- a/gradio/utils.py +++ b/gradio/utils.py @@ -178,6 +178,7 @@ def swap_blocks(self, demo: "Blocks"): # when the app was reloaded continue to send diffs, not full values demo.pending_streams = self.running_app.blocks.pending_streams demo.pending_diff_streams = self.running_app.blocks.pending_diff_streams + demo._stream_run_ids = self.running_app.blocks._stream_run_ids demo.allowed_paths = self.running_app.blocks.allowed_paths demo.blocked_paths = self.running_app.blocks.blocked_paths diff --git a/test/test_blocks.py b/test/test_blocks.py index 2792a328f7..f18bce2fa2 100644 --- a/test/test_blocks.py +++ b/test/test_blocks.py @@ -1,5 +1,6 @@ import asyncio import copy +import gc import io import json import os @@ -25,7 +26,7 @@ from PIL import Image import gradio as gr -from gradio import blocks, helpers +from gradio import blocks, helpers, processing_utils from gradio.context import LocalContext from gradio.data_classes import GradioModel, GradioRootModel from gradio.events import SelectData @@ -1660,9 +1661,62 @@ def gen(): assert event_id in app.iterators_to_reset +requires_ffmpeg = pytest.mark.skipif( + not processing_utils.ffmpeg_installed(), reason="ffmpeg not installed" +) + + +def streaming_audio_demo(): + """A Blocks whose button streams two audio chunks into a streaming Audio.""" + chunk = ( + pathlib.Path(__file__).parent / "test_files" / "audio_sample.wav" + ).read_bytes() + + def stream(): + yield chunk + yield chunk + + with gr.Blocks() as demo: + audio = gr.Audio(streaming=True) + gr.Button().click(stream, None, audio) + + return demo, next(iter(demo.fns.values())), audio + + +async def drive_streaming_run(demo, block_fn, session_hash, event_id): + """Run a generator to completion the way the queue does, one call per chunk. + + Returns the playlist URL from the first chunk. + """ + iterator, url = None, None + for _ in range(10): + output = await demo.process_api( + block_fn=block_fn, + inputs=[], + state=None, + iterator=iterator, + session_hash=session_hash, + event_id=event_id, + ) + iterator = output["iterator"] + if url is None: + url = output["data"][0]["url"] + if not output["is_generating"]: + # The final chunk carries the whole value again, so its URL has to + # still be the first chunk's. A key that changed mid-run would have + # opened a second stream under a second URL. The chunks in between + # carry a diff, so there is nothing to compare there. + final = output["data"][0] + assert isinstance(final, dict) and final.get("url") == url, ( + "the run key changed mid-run" + ) + return url + raise AssertionError("the generator never stopped generating") + + class TestHandleStreamingOutputs: @pytest.mark.asyncio - async def test_final_chunk_with_no_open_stream_is_a_no_op(self): + async def test_final_chunk_with_no_open_stream_leaves_nothing_behind(self): # The session's streams may be gone by the time the final chunk lands — # a disconnect drops them — and a generator that sends only prop updates # never opens one at all. Neither should cost the caller their output. @@ -1673,11 +1727,80 @@ async def test_final_chunk_with_no_open_stream_is_a_no_op(self): block_fn = next(iter(demo.fns.values())) data = await demo.handle_streaming_outputs( - block_fn, [b"final"], session_hash="s", run=0, final=True + block_fn, [b"final"], session_hash="s", run="run-1", final=True ) assert data == [b"final"] - assert demo.pending_streams["s"][0] == {} + # and the run leaves nothing behind, since nothing can fetch it + assert "s" not in demo.pending_streams + + @pytest.mark.asyncio + async def test_sequential_runs_get_distinct_run_keys(self): + # A later run used to land on the address a finished run had just freed + # and inherit the stream that run had already ended. + # See https://github.com/gradio-app/gradio/issues/13809 + def stream(): + yield None + yield None + + with gr.Blocks() as demo: + audio = gr.Audio(streaming=True) + gr.Button().click(stream, None, audio) + + block_fn = next(iter(demo.fns.values())) + urls = [ + await drive_streaming_run(demo, block_fn, "s", f"event-{i}") + for i in range(100) + ] + + assert len(set(urls)) == 100 + + def test_run_keys_are_released_with_their_iterator(self): + # Holding the keys strongly would also give every run its own key, by + # keeping every iterator alive, which is a leak rather than a fix. + from gradio.utils import SyncToAsyncIterator + + demo = gr.Blocks() + for _ in range(100): + iterator = SyncToAsyncIterator(iter([1]), None) + demo._stream_run_key(iterator) + del iterator + + gc.collect() + assert len(demo._stream_run_ids) == 0 + + @requires_ffmpeg + @pytest.mark.asyncio + async def test_each_run_gets_its_own_stream(self): + demo, block_fn, audio = streaming_audio_demo() + + first = await drive_streaming_run(demo, block_fn, "s", "event-1") + second = await drive_streaming_run(demo, block_fn, "s", "event-2") + + streams = demo.pending_streams["s"] + assert first != second + assert {first, second} == { + f"{API_PREFIX}/stream/s/{key}/{audio._id}/playlist.m3u8" for key in streams + } + # two segments each, so neither run appended to the other's stream + assert [len(streams[key][audio._id].segments) for key in streams] == [2, 2] + + @requires_ffmpeg + @pytest.mark.asyncio + async def test_runs_are_keyed_by_iterator_not_event_id(self): + # An event id is not a run id. Cancelling an event drops + # `app.iterators[event_id]`, so the next call for that event id starts + # the generator over, and keying the run on the event id would hand the + # restart the streams the first run had already ended. This drives + # `process_api` directly, so it pins the keying, not the cancel path. + demo, block_fn, audio = streaming_audio_demo() + + first = await drive_streaming_run(demo, block_fn, "s", "event-1") + second = await drive_streaming_run(demo, block_fn, "s", "event-1") + + streams = demo.pending_streams["s"] + assert first != second + assert [len(streams[key][audio._id].segments) for key in streams] == [2, 2] class TestGetAPIInfo: diff --git a/test/test_routes.py b/test/test_routes.py index f135d44a75..25611695f8 100644 --- a/test/test_routes.py +++ b/test/test_routes.py @@ -11,7 +11,7 @@ import time from contextlib import asynccontextmanager, closing from pathlib import Path -from threading import Thread +from threading import Event, Thread from types import SimpleNamespace from unittest.mock import patch from urllib.parse import parse_qs, unquote, urlparse @@ -39,8 +39,10 @@ Textbox, close_all, oauth, + route_utils, routes, ) +from gradio.data_classes import PredictBodyInternal from gradio.oauth import _generate_redirect_uri, _redirect_to_target from gradio.route_utils import ( API_PREFIX, @@ -58,6 +60,20 @@ ) +async def drive_aborted_run(app, fn, chunks=2): + """Drive a generator that raises on its second chunk through the route.""" + for _ in range(chunks): + await route_utils.call_process_api( + app=app, + body=PredictBodyInternal( + data=[], session_hash="s", event_id="e1", request=None + ), + gr_request=gr.Request(), + fn=fn, + root_path="", + ) + + @pytest.fixture() def test_client(): io = Interface(lambda x: x + x, "text", "text", api_name="predict") @@ -116,7 +132,7 @@ def test_audio_stream_playlist_uses_stable_target_duration(self): {"data": b"second", "duration": 0.5, "extension": ".aac"} ) ) - demo.pending_streams["session"][0] = {audio._id: stream} + demo.pending_streams["session"]["0"] = {audio._id: stream} response = TestClient(app).get( f"{API_PREFIX}/stream/session/0/{audio._id}/playlist.m3u8" @@ -1039,6 +1055,196 @@ def test_monitoring_link_disabled(self): response = client.get("/monitoring/summary") assert response.status_code == 403 + def test_stream_playlist_route_takes_a_string_run_key(self): + with Blocks() as demo: + audio = gr.Audio(streaming=True) + + app = routes.App.create_app(demo) + client = TestClient(app) + run = "2b0e1cb7f4a34d4b9f2b6d1c8a7e5f30" + demo.pending_streams["session"] = {run: {audio._id: MediaStream()}} + + response = client.get( + f"{API_PREFIX}/stream/session/{run}/{audio._id}/playlist.m3u8" + ) + assert response.status_code == 200 + assert response.text.startswith("#EXTM3U") + + @pytest.mark.parametrize("suffix", ["segment-1.aac", "playlist-file"]) + def test_stream_routes_do_not_reject_a_string_run_key(self, suffix): + # A `run: int` here would 422 on every uuid key, which would break the + # segments the player fetches and the download button, not the playlist. + with Blocks() as demo: + audio = gr.Audio(streaming=True) + + app = routes.App.create_app(demo) + client = TestClient(app) + run = "2b0e1cb7f4a34d4b9f2b6d1c8a7e5f30" + + response = client.get(f"{API_PREFIX}/stream/session/{run}/{audio._id}/{suffix}") + assert response.status_code == 404 + + def test_an_aborted_run_ends_its_streams_and_drops_its_diffs(self): + # A run that raises never reaches its final chunk, so the exception + # path has to end its streams and drop its diff state itself. + def stream(): + yield None + raise RuntimeError("boom") + + with Blocks() as demo: + audio = gr.Audio(streaming=True) + gr.Button().click(stream, None, audio) + + app = routes.App.create_app(demo) + fn = next(iter(demo.fns.values())) + + with pytest.raises(RuntimeError): + asyncio.run(drive_aborted_run(app, fn)) + + # `.get`, not a subscript: both dicts are defaultdicts, so reading a + # missing session would create it and pass the assertion below. + runs = demo.pending_streams.get("s") + assert runs is not None and len(runs) == 1 + run = next(iter(runs)) + + assert runs[run][audio._id].ended is True + assert "s" not in demo.pending_diff_streams + + def test_a_run_with_no_streaming_output_files_no_stream_entry(self): + # The entry is filed when an output opens a stream, not up front, so a + # generator with no streaming output leaves nothing behind however it + # ends. + def stream(): + yield "chunk one" + raise RuntimeError("boom") + + with Blocks() as demo: + box = gr.Textbox() + gr.Button().click(stream, None, box) + + app = routes.App.create_app(demo) + fn = next(iter(demo.fns.values())) + + with pytest.raises(RuntimeError): + asyncio.run(drive_aborted_run(app, fn)) + + assert "s" not in demo.pending_streams + + def test_a_call_with_no_event_id_keeps_no_diff_state(self): + # Diff state serves the later chunks of a run, and without an event id + # there is no way to fetch them: /run/{api_name} makes one call and + # returns. Keeping it would leave an entry per call for the life of the + # process. + def stream(): + yield "chunk one" + yield "chunk two" + + with Blocks() as demo: + box = gr.Textbox() + gr.Button().click(stream, None, box) + + app = routes.App.create_app(demo) + fn = next(iter(demo.fns.values())) + + async def one_call(): + return await route_utils.call_process_api( + app=app, + body=PredictBodyInternal(data=[], session_hash="s", request=None), + gr_request=gr.Request(), + fn=fn, + root_path="", + ) + + output = asyncio.run(one_call()) + + assert output["data"] == ["chunk one"] + assert "s" not in demo.pending_diff_streams + assert "s" not in demo.pending_streams + + def test_a_run_whose_client_goes_away_drops_its_diff_state(self): + # A client that closes the tab mid-run does not raise, and no final + # chunk arrives, so neither of the other two cleanups runs. The queue + # closes the run out when the event ends, heartbeat or not: the stream + # is ended and kept, the diff state goes. + stop = Event() + + def stream(): + while not stop.is_set(): + time.sleep(0.1) + yield None + + with Blocks() as demo: + audio = gr.Audio(streaming=True) + gr.Button().click(stream, None, audio) + + _, local_url, _ = demo.launch(prevent_thread_lock=True) + try: + with httpx.Client(base_url=local_url, timeout=30) as client: + join = client.post( + f"{API_PREFIX}/queue/join", + json={"data": [], "fn_index": 0, "session_hash": "s"}, + ) + assert join.status_code == 200 + with client.stream( + "GET", f"{API_PREFIX}/queue/data", params={"session_hash": "s"} + ) as sse: + for line in sse.iter_lines(): + if "process_generating" in line: + break + # Checked while the connection is still open, so the run + # cannot have been torn down yet + assert len(demo.pending_streams.get("s", {})) == 1 + assert len(demo.pending_diff_streams.get("s", {})) == 1 + + for _ in range(150): + if "s" not in demo.pending_diff_streams: + break + time.sleep(0.1) + assert "s" not in demo.pending_diff_streams + runs = demo.pending_streams.get("s", {}) + assert len(runs) == 1 + assert next(iter(runs.values()))[audio._id].ended is True + finally: + stop.set() + demo.close() + + def test_a_finished_run_keeps_its_stream_for_its_client(self): + # A client that saw the run finish fetches the playlist afterwards, so + # the stream has to stay, ended, once the run is closed out. + def stream(): + for _ in range(3): + yield None + + with Blocks() as demo: + audio = gr.Audio(streaming=True) + gr.Button().click(stream, None, audio) + + _, local_url, _ = demo.launch(prevent_thread_lock=True) + try: + with httpx.Client(base_url=local_url, timeout=30) as client: + join = client.post( + f"{API_PREFIX}/queue/join", + json={"data": [], "fn_index": 0, "session_hash": "s"}, + ) + assert join.status_code == 200 + with client.stream( + "GET", f"{API_PREFIX}/queue/data", params={"session_hash": "s"} + ) as sse: + for line in sse.iter_lines(): + if "process_completed" in line: + break + + runs = demo.pending_streams.get("s", {}) + assert len(runs) == 1 + run = next(iter(runs)) + assert runs[run][audio._id].ended is True + playlist = client.get( + f"{API_PREFIX}/stream/s/{run}/{audio._id}/playlist.m3u8" + ) + assert playlist.status_code == 200 + finally: + demo.close() + def test_api_listener(connect): with gr.Blocks() as demo: