fix(python): stop the async stream dropping its terminal sentinel - #623
Open
ayaangazali wants to merge 1 commit into
Open
fix(python): stop the async stream dropping its terminal sentinel#623ayaangazali wants to merge 1 commit into
ayaangazali wants to merge 1 commit into
Conversation
_AsyncBridge signals completion and failure with a bare put_nowait posted
to the loop thread:
self._loop.call_soon_threadsafe(self._q.put_nowait, _DONE)
The queue is bounded (maxsize defaults to 64), so when it is saturated at
that moment put_nowait raises QueueFull inside the callback, where the
loop's exception handler logs and swallows it. The sentinel is gone, and
aiter_tokens is left awaiting get() with nothing else to post, so the
async for never ends.
The token path already guards against exactly this: _enqueue catches
QueueFull and reschedules itself until the consumer drains. The sync
bridge is safe too, because it uses a blocking put. The async terminal
path was the only one that could lose its message.
Route both sentinels through the same reschedule, bailing out when stop
is set since the consumer is gone by then.
Signed-off-by: ayaangazali <ayaangazali.work@gmail.com>
📝 WalkthroughWalkthroughThe async streaming bridge now retries completion and error sentinel insertion when its bounded queue is full. Sentinel delivery stops after the bridge stops. A regression test verifies completion delivery after queue saturation. ChangesAsync streaming bridge
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
sdk/runanywhere-python/tests/test_streaming.py (1)
11-11: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover
_ERRORunder queue saturation.Line 11 imports
_DONE, and Lines 145-168 invoke only_on_done. Add a matching full-queue case that calls_on_errorand asserts_ERROR. Otherwise, a regression in failure signaling can still leaveaiter_tokenswaiting indefinitely.Also applies to: 143-168
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sdk/runanywhere-python/tests/test_streaming.py` at line 11, Extend the streaming tests around _AsyncBridge and the existing _on_done full-queue case to cover _on_error when the queue is saturated: invoke the error callback, then assert that the queued sentinel is _ERROR and that aiter_tokens does not remain waiting. Reuse the imported _ERROR symbol and preserve the existing completion-case assertions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sdk/runanywhere-python/runanywhere/_streaming.py`:
- Around line 211-226: Update _enqueue_sentinel and the related _enqueue retry
path to await queue capacity or use a cancellable delayed retry instead of
immediately rescheduling on asyncio.QueueFull. Ensure pending sentinel/token
delivery is cancelled when _stop is set or the bridge stops, while preserving
delivery when the consumer remains active.
---
Nitpick comments:
In `@sdk/runanywhere-python/tests/test_streaming.py`:
- Line 11: Extend the streaming tests around _AsyncBridge and the existing
_on_done full-queue case to cover _on_error when the queue is saturated: invoke
the error callback, then assert that the queued sentinel is _ERROR and that
aiter_tokens does not remain waiting. Reuse the imported _ERROR symbol and
preserve the existing completion-case assertions.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 2e194f46-2d80-49e2-b3ec-402b4dff1036
📒 Files selected for processing (2)
sdk/runanywhere-python/runanywhere/_streaming.pysdk/runanywhere-python/tests/test_streaming.py
Comment on lines
+211
to
+226
| def _enqueue_sentinel(self, sentinel: object) -> None: | ||
| """Run on the loop thread: push a terminal sentinel, rescheduling while the queue | ||
| is full so it is never dropped. | ||
|
|
||
| A bare ``put_nowait`` raises :class:`asyncio.QueueFull` inside the loop callback, | ||
| where it is swallowed by the exception handler and the consumer waits on ``get()`` | ||
| forever. ``_enqueue`` already reschedules tokens for the same reason, and the sync | ||
| bridge gets this for free from its blocking ``put``. Stop means the consumer is | ||
| gone, so there is nobody left to hand the sentinel to. | ||
| """ | ||
| if self._stop.is_set(): | ||
| return | ||
| try: | ||
| self._q.put_nowait(sentinel) | ||
| except asyncio.QueueFull: | ||
| self._loop.call_soon(self._enqueue_sentinel, sentinel) |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Files:\n'
git ls-files | rg 'sdk/runanywhere-python/runanywhere/_streaming\.py|sdk/runanywhere-python/runanywhere/__init__\.py$|sdk/runanywhere-python/runanywhere/VERSION' || true
printf '\nFile outline:\n'
ast-grep outline sdk/runanywhere-python/runanywhere/_streaming.py --view expanded || true
printf '\nRelevant section:\n'
cat -n sdk/runanywhere-python/runanywhere/_streaming.py | sed -n '1,280p'
printf '\nSearch close/stop/usages:\n'
rg -n "_enqueue_sentinel|_stop|close\\(|_q\\.|call_soon|QueueFull|backpressure|stop" sdk/runanywhere-python/runanywhere/_streaming.pyRepository: RunanywhereAI/runanywhere-sdks
Length of output: 18529
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Python queue behavior probe:\n'
python3 - <<'PY'
import asyncio, queue
async def immediate_reschedule_queuefull():
ev = asyncio.Event()
calls = 0
def callback():
nonlocal calls
calls += 1
q = queue.Queue(maxsize=1)
q.put_nowait("x")
try:
q.put_nowait("sentinel")
except queue.Full:
# Same callback scheduled immediately when queue remains full.
asyncio.get_running_loop().call_soon(callback)
return
ev.set()
# Limit iterations to avoid an unbounded run in sandbox.
for _ in range(100):
callback()
print({"calls": calls, "ev_set": ev.is_set(), "full": queue.Queue(maxsize=1).put_nowait("x") or (queue.Queue(maxsize=1).full() if True else None)})
def run():
ev = asyncio.Event()
calls = 0
def callback():
nonlocal calls
calls += 1
q = queue.Queue(maxsize=1)
q.put_nowait("x")
try:
q.put_nowait("sentinel")
except queue.Full:
loop.call_soon(callback)
return
ev.set()
loop = asyncio.new_event_loop()
loop.call_soon(callback)
while loop._ready and calls < 200:
loop._run_once()
loop.close()
print({"calls": calls, "ev_set": ev.is_set(), "full": queue.Queue(maxsize=1).full()})
run()
PY
printf '\nasyncio.Queue.put_nowait behavior probe:\n'
python3 - <<'PY'
import asyncio
async def main():
q = asyncio.Queue(maxsize=1)
await q.put("x")
try:
q.put_nowait("sentinel")
except Exception as ex:
print(type(ex).__name__, str(ex), isinstance(ex, asyncio.QueueFull))
print("full:", q.full())
q.get_nowait()
print("after put_nowait:", await q.get())
asyncio.run(main())
PY
printf '\nCheck asyncio.Queue.put and wait_for semantics from installed stdlib source if unavailable:\n'
python3 - <<'PY'
import asyncio, inspect, sys
print(sys.version)
print(inspect.getsource(asyncio.Queue.put).replace("\n", "\n"))
print("--- get ---")
print(inspect.getsource(asyncio.Queue.get))
PYRepository: RunanywhereAI/runanywhere-sdks
Length of output: 160
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Bounded immediate reschedule probe with step limit:\n'
python3 - <<'PY'
import asyncio
import queue
def immediate_reschedule_queuefull():
loop = asyncio.new_event_loop()
calls = 0
def callback():
nonlocal calls
calls += 1
q = queue.Queue(maxsize=1)
q.put_nowait("x")
try:
q.put_nowait("sentinel")
except queue.Full:
loop.call_soon(callback)
return
q.put_nowait("sentinel")
loop.call_soon(callback)
while loop._ready and calls <= 5:
loop._run_once()
loop.close()
print({"calls": calls, "sentinel_received": calls > 1})
immediate_reschedule_queuefull()
python_version = __import__("sys").version
print({"python_version": python_version})
import asyncio, inspect
print({"has_getsource": bool(inspect.getsource(asyncio.Queue.put))})
print(inspect.getsource(asyncio.Queue.put))
print("--- get ---")
print(inspect.getsource(asyncio.Queue.get))
PYRepository: RunanywhereAI/runanywhere-sdks
Length of output: 2506
Wait instead of immediately rescheduling full-queue retries.
_enqueue() and _enqueue_sentinel() both call put_nowait() inside the event-loop callback and reschedule with call_soon on asyncio.QueueFull. The queue becomes full, _stop stays clear, and the callback is scheduled again before any consumer action can remove items. If the queue stays full, this repeatedly schedules the same callback. Use await _q.put(sentinel) or a cancellable delayed retry, and cancel pending delivery when the bridge stops.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sdk/runanywhere-python/runanywhere/_streaming.py` around lines 211 - 226,
Update _enqueue_sentinel and the related _enqueue retry path to await queue
capacity or use a cancellable delayed retry instead of immediately rescheduling
on asyncio.QueueFull. Ensure pending sentinel/token delivery is cancelled when
_stop is set or the bridge stops, while preserving delivery when the consumer
remains active.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
_AsyncBridgesignals both completion and failure with a bareput_nowaitposted to the loop thread:The queue is bounded (
maxsizedefaults to 64). When it is saturated at that moment,put_nowaitraisesasyncio.QueueFullinside the loop callback, where the event loop's exception handler logs it and moves on. The sentinel is simply gone, and nothing else will post one.aiter_tokensis then parked onawait bridge.get()forever, so the caller'sasync fornever ends and the worker thread is never joined.Why this is an oversight rather than a choice
Every sibling path in the same file already handles it:
_enqueuecatchesQueueFulland reschedules itself until the consumer drains. That is the file's own backpressure mechanism, so the author knew the queue fills._on_done/_on_erroruse a blockingself._q.put(...), which waits for space and cannot lose the sentinel.The async terminal path was the only one that could drop its message. This routes both sentinels through the same reschedule, bailing out when
_stopis set, since by then the consumer has gone and there is nobody to hand it to.The trigger is ordinary: a consumer slower than the producer, which is the normal case for a local model feeding UI or I/O work per token.
Testing
From
sdk/runanywhere-python(Python 3.9, pytest 8.4.2):test_package_import.py::test_version_string, which raisesPackageNotFoundError: runanywherebecause the package is not pip-installed in my environment. It fails identically with my change stashed, so it is unrelated.I first wrote this test at the
aiter_tokenslevel and it passed without the fix, because a promptly draining consumer frees a slot before the sentinel lands. It only reproduces deterministically by saturating the queue and driving_on_donedirectly, so that is what the committed test does.🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Tests