Fix streaming run key colliding across sequential runs - #13811
Conversation
The per-run key for a streaming output was `id(iterator)`. It keys `pending_streams[session_hash][run]` and goes straight into the playlist URL the frontend receives. `id()` is only unique among objects that are alive at the same time, and nothing keeps a finished run's iterator alive: `pending_streams` holds `MediaStream` objects keyed by the id, not the iterator, and `app.iterators[event_id]` is set to None once the generator is exhausted. So the next run's iterator can land on the address the previous one just freed. When it does, `first_chunk` is False on the new run's first chunk, so segments are appended to a `MediaStream` that already had `end_stream()` called on it, and the playlist URL is byte-identical to the previous run's, which is the only signal the player has that a new stream started. The new run is silent. The run key is now the event id: a uuid4 per queue event, the same for every chunk of a run, already the key of `app.iterators`. The three `/stream/...` routes take `run` as a str accordingly.
🪼 branch checks and previews
Install Gradio from this PR pip install https://huggingface.co/buckets/gradio/pypi-previews/resolve/e1b32b7d7debb1acaf41bf45425a568df5d0ca51/gradio-6.26.0-py3-none-any.whlInstall Gradio Python Client from this PR pip install "gradio-client @ git+https://github.com/gradio-app/gradio@e1b32b7d7debb1acaf41bf45425a568df5d0ca51#subdirectory=client/python"Import Gradio JS Client from this PR via CDN import { Client } from "https://huggingface.co/buckets/gradio/npm-previews/resolve/e1b32b7d7debb1acaf41bf45425a568df5d0ca51/browser.js"; |
🦄 change detectedThis Pull Request includes changes to the following packages.
|
There was a problem hiding this comment.
Pull request overview
This PR fixes a backend streaming identifier bug by replacing the per-run key derived from id(iterator) (which can be reused by CPython) with a stable event-scoped identifier, preventing run-key collisions that can corrupt pending_streams state and reuse playlist URLs across sequential streaming runs.
Changes:
- Use
event_id(with a fallback) as the streaming run key instead ofid(iterator)when building playlist URLs and indexingpending_streams. - Update the
/stream/{session_hash}/{run}/...routes to acceptrunas astrrather than anint. - Add/adjust tests to validate playlist routing and that sequential streaming runs produce distinct playlist URLs and independent
MediaStreamsegment sets.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
gradio/blocks.py |
Switches streaming run key generation to use event_id (and updates type hints for run). |
gradio/route_utils.py |
Updates stream-closure logic on exceptions to look up pending streams by event_id. |
gradio/routes.py |
Changes stream route path parameter typing for run from int to str. |
test/test_blocks.py |
Updates existing streaming tests for run as str and adds coverage for per-event run isolation. |
test/test_routes.py |
Adds a route-level test ensuring the playlist endpoint works with an event-id run key. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
id(iterator)A run driven straight through `process_api`, without going through the queue, has no event id, and the previous fallback minted a fresh uuid on every chunk. Over HTTP that cannot happen, because `restore_session_state` hands back no iterator without an event id, but a direct caller can resume a run and would then get a new playlist URL per chunk. Key those runs off the iterator instead, through a `WeakKeyDictionary` so the entry goes away with the iterator and a reused address never becomes a second run's key.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
.changeset/violet-clocks-reply.md:5
- This PR adds a hand-authored Changeset file. Repo guidance says not to commit .changeset/*.md because a GitHub Action generates it from the PR title, and an existing file will override the autogenerated changelog entry.
---
"gradio": minor
---
feat:Fix streaming run key colliding across sequential runs
An event id identifies an event, not a generator run, and one event can drive two runs. Cancelling a streaming event deletes its entry from app.iterators and adds the id to app.iterators_to_reset, so the next call for that event gets iterator=None and process_api starts the generator over under the same id. The restarted run then inherits the streams the cancel had already ended, and its playlist URL does not change, which leaves the player on a source it is done with. The iterator is the only object that is one-to-one with a run, so key the run on that. id() is not usable because it only holds for objects that are alive at the same time, which is what #13809 is about, so hold a uuid4 against the iterator in a WeakKeyDictionary instead. The entry dies with the iterator, so a later iterator that lands on the same address gets a fresh key. swap_blocks carries pending_streams and pending_diff_streams across a hot reload; it now carries the run keys that index them too. The exception path in call_process_api goes back through app.iterators to find the run, as it did before the event id was used.
There was a problem hiding this comment.
🟡 Changes recommended
_stream_run_key()’s TypeError fallback can generate a new run key per chunk (breaking streaming for non-weakref iterators), and the PR includes a hand-authored changeset file that conflicts with repository contribution rules.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 7/7 changed files
- Comments generated: 2
- Review effort level: Lite
| except TypeError: | ||
| # Not every object an async generator function can return supports | ||
| # weak references, so such a run gets a throwaway key. | ||
| run = uuid.uuid4().hex | ||
| return run |
There was a problem hiding this comment.
Half of this holds. The suggested cache-on-the-iterator step does not: an object that cannot be weak-referenced is one whose __slots__ omits __weakref__, and such an object rejects attribute assignment too, so caching the key on it fails for exactly the case it is meant to cover.
The concern about the fallback was fair, though, and the fallback is gone rather than reworked. This branch is not reachable through gradio's own paths: call_function produces either a utils.SyncToAsyncIterator or a real async generator, and both support weak references. So the only thing the branch decided was what happens if that stops being true, and every candidate answer was worse than raising. id(iterator) would have silently restored the collision this PR fixes, while contradicting the "never repeat" line in the docstring right above it, and a fresh uuid per chunk would restart the player on every chunk. An iterator that cannot be keyed now fails loudly at that line instead.
| --- | ||
| "gradio": patch | ||
| --- | ||
|
|
||
| fix:Fix streaming run key colliding across sequential runs |
The docstring and the two test comments restated each other and the mechanism the PR description already covers. Keep the non-obvious why in the helper's docstring and leave the tests to point at the issue. The TypeError comment now says why a throwaway key is the best available answer rather than only that the branch exists: an object that takes no weak reference offers no stable key we could hold without leaking it.
Two problems the third self-review round found. handle_streaming_outputs files an entry under the run key for every generator, streaming outputs or not, and only drops it when the session disconnects. Reused keys used to cap that by accident, so making keys unique turned a self-limiting leak into a linear one: 500 sequential runs of a plain text generator left 165 entries before and 500 after. A run that ends without having opened a stream now drops its entry, which nothing could fetch anyway, and the same 500 runs leave none. Entries holding a stream still stay, since the playlist is fetched after the run ends. The TypeError fallback minted a fresh uuid per call, so an iterator that takes no weak reference got a different key on every chunk and its playback restarted each time. The iterator is alive for the whole run, so its address is stable there; reuse across runs is still possible but that is what main already does. The test that named the collision could not fail: it asserted 200 uuid4 hexes were distinct, which they are by construction, and it passed just as well with a strong dict. Replace it with one that drives process_api and counts run keys, which fails 80 == 100 on main, and one that asserts the keys are released with their iterators. Both run without ffmpeg, so the suite keeps real coverage on a machine that has none. The cancel test's end_stream() loop was dead, since driving the generator to exhaustion already ends every stream. Drop it and name the test for what it checks: two runs under one event id.
The TypeError fallback contradicted the docstring right above it: it promised the key would never repeat and then returned str(id()), which is the repeat this whole change is about. Nothing reachable triggers the branch, since call_function only ever produces a SyncToAsyncIterator or a real async generator and both take weak references, so the only question it answers is what happens if that stops being true. A stack trace at this line beats silently going back to address-keyed runs. Removing it is also what makes the exception path in call_process_api safe to turn into a plain lookup. Minting a key there wrote to _stream_run_ids while another exception was in flight, to find streams that by definition do not exist, and ran user __hash__ that could replace the exception being handled. With no fallback left, any iterator that reaches the lookup has already succeeded as a key on an earlier chunk, so reading it cannot fail in a way the happy path would not have failed first. Also pop the empty entry instead of deleting it, since the proof that the key is present rests on there being no await in between, and assert the run keys are distinct rather than counting retained entries, which pinned a retention policy the test does not own. The segment and playlist-file routes now have coverage too: all three 422 on every uuid key if run goes back to int, and only one of them was tested.
A cancelled run never reaches its final chunk, so nothing ends its MediaStreams and the player goes on polling a playlist that never gets its #EXT-X-ENDLIST. Its pending_diff_streams entry, the last full postprocessed output, is stranded too, and that is the one stream dict the disconnect handler does not clear, so it stays for the life of the process. /cancel is the last place that still holds the iterator, and with it the run key, so do both there: 300 cancelled runs leave 300 entries before and none after. This is not a regression from the key change, since app.iterators keeps the iterator alive across a cancel and no address was being reused there either. It is just not fixable without something that identifies the run. The lookup that needs is also the one call_process_api's handler was making by reaching into the table directly, so give it a name. Two callers want to find a run without minting one for an iterator that has none, which is what a get-or-create helper does if a caller forgets. The comment there claimed any iterator reaching the lookup had already succeeded as a key; a call with an event id and no session hash stores one that never opened a run, so say that instead. The route test asserted only "not 422", which a 500 would satisfy; both routes 404 on a missing entry, so assert that.
The cleanup added to /cancel cannot fire. The route awaits cancel_tasks first, which gathers the cancelled tasks to completion, and their finally calls Queue.reset_iterators, which deletes app.iterators[event_id]. By the time the route resumes the event id is gone and the block is skipped. Worse, it broke /cancel for events that had already finished. call_process_api assigns app.iterators[event_id] after every call, so the last chunk leaves None under a key that is never removed. A cancel for such an event reached the new lookup, and WeakKeyDictionary.get(None) raises rather than returning None, so the request went from 200 to 500. gradio_client's Job.cancel posts there unconditionally once a job is cancellable. Ending the streams was redundant in any case: call_process_api's exception path already does it, which is why the streams looked closed. The test missed all of this because it drove process_api directly and filled app.iterators by hand, so no task existed for cancel_tasks to match and the block ran in a state the server never reaches. The diff-state leak the cleanup was aiming at is not a regression from this PR, and closing it properly means covering the error path and sibling runs too, so it belongs elsewhere. Keep the lookup accessor, since the exception path in call_process_api still wants a lookup that does not mint, but underscore it: Blocks is public API and this is not.
A run that is cancelled or raises never reaches the final chunk that would delete its pending_diff_streams entry, and that entry holds the last full postprocessed output of every component: a whole Chatbot history for a chat app. Colliding keys used to reclaim it, because a later run inherited the entry and deleted it on its own final chunk, so 300 aborted runs left 155 entries. Distinct keys leave 300, and pending_diff_streams is the one stream dict the disconnect handler never clears, so they outlive the session. This is the leak the backed-out commit was aiming at, at a place that actually runs. call_process_api's handler already ends the run's streams on the way out and has the session hash and the run key in hand, so it drops the diff entry there too. The same 300 runs now leave none. A first-chunk abort still slips through, since the iterator is not in app.iterators yet, but that is unchanged from before. Also capture the session's dict once in handle_streaming_outputs. Reaching back through the defaultdict after the awaits in the output loop can recreate a session that the disconnect handler removed in between, leaving an entry nothing will ever clear. Two test gaps behind this. Nothing covered the abort path, so add a test that drives call_process_api into it. And drive_streaming_run only recorded the first chunk's URL, so a key that churned per chunk - the regression bcdd127 fixed - was caught only by an ffmpeg test; it now checks the final chunk's URL against the first.
There was a problem hiding this comment.
Very nice @hysts! LGTM, with one small finding:
Aborted non-media generators still leak one empty pending_streams entry per run.
handle_streaming_outputs() eagerly creates pending_streams[session][run] = {} before determining whether the generator has any streaming media output (blocks.py). Normal completion removes that empty entry, but the exception path only ends existing media streams and removes pending_diff_streams; it never removes an empty stream entry (route_utils.py).
Therefore, a text generator that yields once and then raises—or is cancelled—leaves:
handle_streaming_outputs files an entry under the run key for every generator, before it knows whether any of its outputs stream. The final chunk drops it when the run never opened one, but a run that raises or is cancelled never reaches that chunk, so the entry stayed until the session disconnected. Reused keys used to cap this the way they capped the finished-run case: 200 aborted runs of a text generator leave 200 entries under distinct keys, against 154 to 157 across repeat runs under the old id(iterator) key. The exception path already ends the run's streams and drops its diff state, so drop the stream entry there too when it holds no stream, which takes the same 200 runs back to none. A real cancel reaches that handler as well, before the cancelled task's finally gets to Queue.reset_iterators, so it is covered by the same line: one entry before, none after. Entries that do hold a stream still stay, since the playlist is fetched after the run ends. Found by abidlabs in review.
The comment on test_two_runs_under_one_event_id_get_separate_streams gave the same cause for a cancel restart that was already taken out of the PR description as inaccurate. The only writer of app.iterators_to_reset sits inside `if body.event_id in app.iterators`, which Queue.reset_iterators has already emptied from the cancelled task's own finally by the time /cancel resumes, so that block does not run. What restarts the generator is the delete itself, after which restore_session_state hands back iterator=None. The test's conclusion is unaffected. The exception handler reached through blocks.pending_streams twice. Capture the session's runs once, the way handle_streaming_outputs already does. Drop a dead assignment in the sibling test while there.
Both abort tests looped five times over a generator that raises on its second chunk, so the third pass was unreachable and the five read as if the count mattered. They also carried the same driver twice. One module-level helper, bounded at the two chunks the run actually has, covers both. The exception handler asked the app for its blocks twice. Bind them once, above the lookup that wants them first.
|
Good catch, fixed in f45f280 and tidied since. You are right that the entry is filed before the run is known to have any streaming output, and that the exception path did not drop it. Measured with 100 aborted runs of a text generator that yields once and then raises, each driven through a queue join and an SSE stream: 100 empty entries before, none after. The old The fix sits next to the two cleanups that handler already does, since the entry holds no stream and nothing can fetch it: if not pending_streams:
stream_runs.pop(run_id, None)A cancel is covered by the same line. Three cases stay out of reach, and all three are blind spots the handler already had. A run that dies during its first chunk, since the iterator is stored only once Entries that do hold a stream still stay, since the playlist is fetched after the run ends. |
A client that closes the tab mid-run leaves nothing to hang a cleanup
on. clean_events marks the event not alive, process_events drops out at
`if not awake_events: return` without raising, so call_process_api's
handler never runs, and no final chunk arrives to make
handle_streaming_diffs delete its own entry. pending_diff_streams is
the one stream dict with no session-level cleanup, so the entry stays
for the life of the process, holding the last full postprocessed value
of every output component.
Reused keys used to reclaim most of that: a later run in the same
session landed on the freed address and deleted the entry on its own
final chunk. 300 runs interrupted this way left 152 entries before the
key change and 300 after, so this is one the change does make worse.
The `/run/{api_name}` case measured earlier is not, at 189 against 200,
because a whole request cycle allocates in between and the addresses
churn.
Both dicts now go together, in one place named for what it does, which
the disconnect handler calls instead of ending the streams inline. The
same 300 runs leave none.
Also give _stream_run_ids its type parameters, since the run key
contract is what the routes interpolate into a URL, and stop the abort
tests from reading these defaultdicts by subscript, which created the
session key they meant to assert was gone.
Found by review.
The comment beside the diff-state pop said this handler was the only place that dropped it and that nothing cleared the dict on disconnect. Both stopped being true one commit ago, when the disconnect handler started dropping it. Say what is left that is still true, which is that the final chunk never arrives, and say it once for the two pops instead of once each. That handler had grown twelve comment lines around thirteen of code, against about five percent for the file as a whole. It is down to six. The line restating that the loop below it closes open streams goes with them.
#13806 landed a playlist test that files its stream under an int run key and fetches it back through the route. On this branch the route's `run` is a `str`, so the path segment no longer parses to the int the dict is keyed by and the fetch 404s. The key in the test is a string now. What it asserts, that the playlist carries a stable target duration, is untouched.
4092317 dropped a dead session's diff state from the heartbeat disconnect handler. That handler only runs for an app that opens a heartbeat, which takes gr.State, an unload or stream listener or a per-session cache. The issue's app, a button driving a streaming gr.Audio, has none of those, so for the app this PR is about the cleanup never ran. The place every queued event does pass through is the finally in Queue.process_events, whether the run finished, raised, was cancelled or lost its client, and it already resets the event's iterator there. Close out the run's streams first, while the iterator that keys them is still stored, through one Blocks helper that the exception path in call_process_api now calls as well. The heartbeat handler goes back to what it was. The number that justified 4092317 was also wrong. It came from a loop that popped app.iterators by hand with nothing allocating in between, so id() reuse reclaimed half the entries. Driven through a real connection, with a queue join and an SSE stream between runs, 100 interrupted runs leave 98 entries under the old key and 100 under this PR's, so the leak is pre-existing and all but unchanged by the key switch. It is closed here anyway, since the entry holds the last full output of every component for the life of the process. The same 100 runs now leave none. The disconnect test drives a launched app over a real connection and fails without the queueing.py change.
handle_streaming_outputs filed pending_streams[session][run] for every
generator up front and two other places removed it again when it turned
out to hold nothing: the final chunk, and _drop_run_streams for a run
that never got there. File it where a MediaStream is first constructed
instead. A run with no streaming output now never touches the dict, so
there is nothing to remove on any path, and the two removals go.
A run whose client went away kept its streams, on the rule that a
playlist is fetched after the run ends. That rule is for a client that
is still there. event.alive is False only after clean_events, which
only the /queue/data disconnect reaches, so the queue's finally now
passes orphaned=not event.alive and the helper drops the run's entry
outright in that case. Without it an abandoned audio run on an app with
no heartbeat pinned its segments for the life of the process, since the
heartbeat handler never runs for such an app.
The heartbeat handler now drops a session's pending_diff_streams beside
pending_streams. A generator driven through /run/{api_name} with a
session hash never enters the queue, so nothing else reaches its diff
state when the session goes away, and the per-session dict itself was
never removed.
Two orderings hold this together, so the finally now says so. /cancel
awaits this task before dropping the iterator itself, and both cleanup
sites look the run up through app.iterators. And a finished run is
still alive at the finally, since nothing awaits between its completion
message and that point; the server, not the client, closes the SSE
stream once a session has no pending events left.
The disconnect test now streams into gr.Audio, so it covers the dropped
stream, and it checks the run is filed while the connection is still
open rather than racing the teardown after closing it. A new test drives
a streaming run to completion over a real connection and fetches its
playlist afterwards, the half of the rule the orphan drop must not
break.
Found by self-review.
6bedc27 had the heartbeat disconnect handler drop a session's pending_diff_streams beside its streams. A heartbeat can die while the /queue/data stream that carries a run survives, and the handler does not stop the run, so its next chunk found no diff state, went out as a full value, and the client applied it as a diff and failed the run. That line goes; the queue's finally already drops a run's diff state when the run ends. handle_streaming_outputs re-filed the dict it had read before awaiting stream_output. A session torn down during that await left the local holding an already-ended stream, which setdefault then put back under the same run key, and later chunks appended to it. Read the entry fresh at the point of use instead. Both dicts kept an empty per-session dict once its last run was gone, one per visitor session for the life of the process on an app with no heartbeat. A run's entry and a session dict it leaves empty now go together, through one small helper used by the finally path and the final-chunk path alike. The finally reaches the shared dicts through the queue's own Blocks, which cannot raise the way app.get_blocks() can during teardown, and its comment no longer argues about finished runs: their iterator is already None by then. Tests: the final-chunk test asserted through a defaultdict subscript that created the very entry it checked, and the disconnect test's generator could finish on its own before a slow runner noticed the disconnect. The first now checks the session key is absent; the second gates the generator on an event the test sets on the way out. Found by self-review.
6bedc27 dropped the stream entry of a run whose client had gone, keyed on event.alive. That flag says the session's /queue/data stream is connected, not that nobody wants the run's output: the JS client aborts that stream on any unexpected client-side exception, and a proxy can cut it, with the tab and its player still there. The player then 404s on its next playlist poll instead of finishing on the ended playlist, and so does the download link. And on an app with no heartbeat every finished run's stream stays anyway, so the drop bounded nothing. Every cleanup path now does the same two things: end the run's streams and leave them fetchable, drop its diff state. The session dict pruning stays, for the diff table it still applies to. The finally reaches the shared dicts through app.get_blocks() again, matching the rest of process_events; the queue's own Blocks reference is never updated across a hot reload and only worked through the aliasing swap_blocks sets up. The heartbeat handler says why it leaves diff state alone. Tests: the disconnect test now asserts the stream is ended and kept and the diff state gone; the finished-run test says what it pins; the one-event-id test is named for what it drives, process_api directly, not the cancel path. Found by self-review.
handle_streaming_diffs filed an entry for every generator run, but its
later chunks are only fetchable under an event id: restore_session_state
returns no iterator without one. A /run/{api_name} call, which has none,
makes one call and returns, so every such call left an entry holding the
full postprocessed output for the life of the process, as did every
.stream() chunk after the first, which arrives without an id. Under the
old key a collision reclaimed one in twenty of those; under this PR's
none. process_api now passes no run to the diff bookkeeping when there
is no event id. The first chunk's output is unchanged, since the first
chunk of a run never carried a diff.
The queue's finally looked the session up on event.data, which
/stream/{event_id} replaces with the client's payload; a caller that
leaves session_hash out of it would have skipped the cleanup. Event
keeps the session hash it was created with, and push writes that same
value into the body, so the finally reads it from the event.
Also record on _stream_run_ids what its weak keys rest on, and shorten
the heartbeat handler's note to what it does rather than why.
Found by self-review.
The queue's finally checked the server app once per event, inside the
loop that also resets iterators, and that reset raises on the same None
the check guards against. The close-out now has its own small loop with
the check outside it, and the reset loop is back to what it was.
The diff bookkeeping comment now says what a call without an event id
gets, full values, since that is what /run/{api_name} callers and the
.stream() client read.
The two tests that launch a server used httpx's default five second
read timeout and gave the disconnect five seconds to be noticed. Both
now say 30 and 15 seconds, so a slow runner reads as a slow runner.
Found by self-review.
|
@abidlabs since your approval I have pushed several follow-up commits, most of them closing leaks that your finding led me to, so the diff has moved and I would like to ask for another look. The net change against what you reviewed:
Two things I tried and took back: dropping diff state in the heartbeat handler (a heartbeat can drop while the run's own connection survives, and the run's next chunk then breaks the client), and dropping an abandoned run's stream entry outright (SSE liveness is not "nobody wants the output", and on an app with no heartbeat every finished run's stream stays anyway). The one known limit, stated in the description: on an app with no heartbeat a run's |
A generator called straight through /gradio_api/run/{api_name} yields
once and is dropped: only the queue continues one, and it always carries
an event id, so without one `restore_session_state` hands back no
iterator at all. The run therefore reaches neither a final chunk nor the
exception path, and its streams stayed open.
That cost segment bytes before this branch. It costs a live ffmpeg
process now, one per abandoned run, held for as long as the session is:
twenty such calls measured at twenty processes and about 800 MiB. Ending
the streams there releases each encoder and gives the playlist the
#EXT-X-ENDLIST it is owed, since the audio that exists is all there will
ever be.
This does not touch the MediaStream itself, which stays in
pending_streams on this path exactly as it does on main. That retention
was looked at in the review of #13811 and left alone; nothing here
re-opens it.
Description
The per-run key for a streaming output was
id(iterator). It is the key ofpending_streams[session_hash][run]and it goes straight into the playlist URL the frontend receives,{session_hash}/{run}/{output_id}/playlist.m3u8.id()is only unique among objects that are alive at the same time, and nothing keeps a finished run's iterator alive:pending_streamsstoresMediaStreamobjects keyed by the id, not the iterator, andapp.iterators[event_id]is set toNoneas soon as the generator is exhausted. So the next run's iterator can land on the address the previous one just freed.When that happens,
run not in self.pending_streams[session_hash]is False, sofirst_chunkis False and the new run's segments are appended to aMediaStreamthat already hadend_stream()called on it. On the client side the playlist URL is what tells the player a new stream has started, so an identical URL means the player keeps the source it already has and the new run is silent, the same symptom as #13807.The issue filed this as latent, but it fires in ordinary use. For a sync generator function the object being identified is not the generator itself but the
utils.SyncToAsyncIteratorwrappercall_functionbuilds around it, and those are all the same small size, so the allocator hands the address straight back. Driving sequential streaming runs throughprocess_apiin one process, withpending_streamscounting the distinct keys:The first dozen or so runs are fine, which is why this does not show up while trying it out by hand, and then the rate climbs. Through a real connection, with a queue join and an SSE stream between runs, more gets allocated in between and the rate is lower, 14 collisions in 300 sequential runs, about one run in twenty, which is still well within a session of trying the demo. After the change both loops give a distinct key for every run, 2000 of 2000 and 300 of 300.
The fix
The iterator is the only object here that is one-to-one with a generator run, so it stays the thing that identifies the run. What changes is the identifier. The run key is now a
uuid4().hexheld in aWeakKeyDictionarykeyed by the iterator, so the entry dies with the iterator and a later iterator that lands on the same address gets a fresh key instead of inheriting the old run's streams.Reloader.swap_blocksalready carriespending_streamsandpending_diff_streamsacross a hot reload so that an in-flight generator keeps its stream; it now carries the run keys that index them too.One consequence needed handling.
handle_streaming_outputsfiled an entry under the run key for every generator, whether or not any of its outputs stream, and only the session disconnect ever removed it. Reused keys trimmed that a little by accident, 300 sequential runs of a plain text generator through a real connection left 286 entries, and distinct keys would make it one per run. The entry is now filed where aMediaStreamis first constructed, so a run with no streaming output never touches the dict, however it ends: the same 300 runs leave nothing, 100 runs of a text generator that raises after its first chunk go from 98 entries to nothing, and the same holds for a generator driven through/run/{api_name}, where 189 of 200 calls used to leave one. Entries that do hold a stream stay, since the playlist is fetched after the run ends. Those cost no more than they did: a colliding run did not free anything, it appended its segments to the previous run's stream, so 60 real audio runs hold the same 754,920 bytes before and after this change and only the number of wrapper objects differs. A generator that yieldsNoneinto a streaming component still opens an empty stream, at about a hundred bytes a run.pending_diff_streamsneeded the same treatment for a different reason. A run that is cancelled or raises never reaches the final chunk that would drop its entry, and that entry holds the last full postprocessed output of every component, a wholegr.Chatbothistory for a chat app. Colliding keys reclaimed a little of it, since a later run that landed on the same key inherited the entry and deleted it on its own final chunk, but through a real connection 100 aborted runs left 98 entries either way.call_process_apialready ends the run's streams on the way out, so it drops the diff entry there too, which takes the same 100 runs to none. Diff state is only of use to a run whose later chunks can be fetched, which takes an event id, so a call without one files none:/run/{api_name}makes one call and returns, and a.stream()chunk after the first arrives without an id, and each such call used to leave an entry behind for the life of the process.That covers a run that raises or is cancelled, but not one the client walks away from.
clean_eventsmarks the event not alive,process_eventsdrops out without raising, and no final chunk arrives, so neither of those places runs and the entry stays for the life of the process. This one the key change does not make worse: driven through a real connection, with a queue join and an SSE stream between runs, the addresses churn and reuse reclaimed almost nothing, 98 entries after 100 interrupted runs under the old key against 100 under this one. It is the same leak either way, a wholegr.Chatbothistory per closed tab for a chat app, so it is closed here too. The heartbeat disconnect handler was the obvious place, but a heartbeat only opens for an app withgr.State, anunloadorstreamlistener or a per-session cache, and the issue's app, a button driving a streaminggr.Audio, has none of those.Queue.process_eventshas onefinallythat every queued event passes through, whether it finished, raised, was cancelled or lost its client, and it already resets the event's iterator there. It now closes out the run's streams first, while the iterator that keys them is still stored, through the sameBlockshelper the exception path uses. Its streams are ended and stay, like those of a run that raised, since a client that is still there fetches the playlist after the run ends; only the diff state goes, and a session dict that this leaves empty goes with its last run, where before every visitor session that ran a generator left one behind for the life of the process. The same 100 runs leave no diff state, andtest_a_run_whose_client_goes_away_drops_its_diff_statedrives it through a real connection into a streaminggr.Audio. What this PR does not do is bound the streams themselves. Nothing in gradio signals the end of a session that opened no heartbeat, so on an app like the issue's every run'sMediaStream, finished or not, stays until the process ends, as it did before. That needs an eviction policy and is left for a follow-up. The same goes for a stream that nothing ends: one opened by a call with no event id,/run/{api_name}or a.stream()chunk after the first, or by a first chunk that failed after opening it, stays un-ended as before.Why not the event id
The obvious alternative is to key the run on
event_id. It is already at hand inprocess_api, it is auuid4().hexper event, and it is already the key ofapp.iterators. I wrote it that way first, and it is wrong: an event id identifies an event, and one event can drive two generator runs.Cancelling a streaming event drops
app.iterators[event_id](Queue.reset_iterators, reached from the cancelled task's ownfinally), so the nextcall_process_apifor that event getsiterator=Noneback fromrestore_session_stateandprocess_apistarts the generator over. Meanwhile the cancelled run's streams have already been ended bycall_process_api's exception path. Keying on the event id hands the restarted run those same streams:id(iterator)event_idThat is the original symptom coming back through a different door, so
test_runs_are_keyed_by_iterator_not_event_idcovers it.Route signature
The three
/stream/...routes tookrunas anint, so they now take astr, and the keys ofpending_streamsandpending_diff_streamsare those strings too. Nothing on the frontend parses the run out of the URL, it only rewritesplaylist.m3u8toplaylist-file, so there is no frontend change.id(iterator)was the run key from the start, in #5077, and #7102 later reused the same key forpending_diff_streams. Neither change needed the key to be an address; there was just no run identifier at hand at the time.Closes: #13809
AI Disclosure
We encourage the use of AI tooling in creating PRs, but the any non-trivial use of AI needs be disclosed. E.g. if you used Claude to write a first draft, you should mention that. Trivial tab-completion doesn't need to be disclosed. You should self-review all PRs, especially if they were generated with AI.
🎯 PRs Should Target Issues
Before your create a PR, please check to see if there is an existing issue for this change. If not, please create an issue before you create this PR, unless the fix is very small.
Not adhering to this guideline will result in the PR being closed.
Testing and Formatting Your Code
PRs will only be merged if tests pass on CI. We recommend at least running the backend tests locally, please set up your Gradio environment locally and run the backed tests:
bash scripts/run_backend_tests.shPlease run these bash scripts to automatically format your code:
bash scripts/format_backend.sh, and (if you made any changes to non-Python files)bash scripts/format_frontend.sh