Fix streaming run key colliding across sequential runs

.changeset/violet-clocks-reply.md

@@ -0,0 +1,5 @@
+---
+"gradio": patch
+---
+
+fix:Fix streaming run key colliding across sequential runs

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,
                     )

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

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:

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, {})

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
 

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: