fix: close partially consumed Python LLM streams#465
Conversation
Signed-off-by: Will Killian <wkillian@nvidia.com>
github-actions
Bot
added
size:M
WalkthroughStreaming LLM plumbing now supports cooperative cancellation and explicit completion signaling across Rust and Python async iterators. ChangesStreaming lifecycle coordination
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PythonConsumer
participant PyLlmStream
participant forward_stream_to_channel
participant AsyncIterator
PythonConsumer->>PyLlmStream: consume stream items
PyLlmStream->>forward_stream_to_channel: receive forwarded items
forward_stream_to_channel->>AsyncIterator: await next item
PythonConsumer->>PyLlmStream: call aclose
PyLlmStream->>forward_stream_to_channel: signal cancellation
forward_stream_to_channel->>AsyncIterator: call aclose
forward_stream_to_channel->>PyLlmStream: signal closed
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/python/src/py_callable.rs (1)
228-235: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winEnsure receiver shutdown and errors finalize the Python iterator.
crates/python/src/py_callable.rs#L228-L235: race error delivery against cancellation and callclose_async_iterbefore every terminal exit.crates/python/tests/coverage/py_callable_coverage_tests.rs#L387-L399: use an observableaclose()implementation and assert receiver shutdown invokes it.🤖 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 `@crates/python/src/py_callable.rs` around lines 228 - 235, Update the receiver loop in py_callable.rs around the terminal branches to race error delivery against cancellation, call close_async_iter before every terminal exit including receiver shutdown and send failures, and preserve the existing break behavior. In crates/python/tests/coverage/py_callable_coverage_tests.rs lines 387-399, replace the iterator’s no-op close behavior with an observable aclose() implementation and assert that receiver shutdown invokes it.Source: Path instructions
🤖 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 `@crates/python/tests/coverage/py_callable_coverage_tests.rs`:
- Around line 387-399: Update the test around forward_async_iter to give the
Python iterator an observable aclose() method, then assert that aclose() is
invoked after the receiver is dropped. Keep the existing dropped-receiver setup
and await behavior, but verify the iterator cleanup side effect rather than only
confirming that forwarding returns.
---
Outside diff comments:
In `@crates/python/src/py_callable.rs`:
- Around line 228-235: Update the receiver loop in py_callable.rs around the
terminal branches to race error delivery against cancellation, call
close_async_iter before every terminal exit including receiver shutdown and send
failures, and preserve the existing break behavior. In
crates/python/tests/coverage/py_callable_coverage_tests.rs lines 387-399,
replace the iterator’s no-op close behavior with an observable aclose()
implementation and assert that receiver shutdown invokes it.
🪄 Autofix (Beta)
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: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: d9a757f0-b569-48c9-a917-277eb89e1846
📒 Files selected for processing (10)
crates/python/src/py_api/mod.rscrates/python/src/py_callable.rscrates/python/src/py_types/core.rscrates/python/tests/coverage/coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rspython/nemo_relay/_native.pyipython/nemo_relay/llm.pypython/tests/test_llm.py
📜 Review details
⏰ Context from checks skipped due to timeout. (2)
- GitHub Check: Check / Run
- GitHub Check: Preview docs
🧰 Additional context used
📓 Path-based instructions (18)
{crates/python/src/py_api/mod.rs,python/nemo_relay/**/*.py,python/nemo_relay/**/*.pyi,go/nemo_relay/**/*.go,crates/node/src/api/**/*.rs}
📄 CodeRabbit inference engine (.agents/skills/add-binding-feature/SKILL.md)
Update the language-native bindings for every exposed surface in Python, Go, and Node.js.
Files:
python/nemo_relay/_native.pyipython/nemo_relay/llm.pycrates/python/src/py_api/mod.rs
{python/nemo_relay/**/*.py,python/nemo_relay/**/*.pyi,go/nemo_relay/**/*.go}
📄 CodeRabbit inference engine (.agents/skills/add-binding-feature/SKILL.md)
Update language wrapper helpers such as Python wrapper modules, Python type stubs, and Go shorthand packages when the new behavior belongs in those helper layers.
Files:
python/nemo_relay/_native.pyipython/nemo_relay/llm.py
**/*
📄 CodeRabbit inference engine (.agents/skills/validate-change/SKILL.md)
**/*: Format changed files with the language-native formatter before the final lint/test pass.
If dynamic plugin behavior changed, usemaintain-dynamic-pluginsand include the native SDK, worker protocol, Python SDK, docs, packaging, and Codecov surfaces in the validation plan.
If code changes alter APIs, bindings, commands, paths, packaging behavior, observability/adaptive semantics, or documented best practices, update any dependent maintainer or consumer skills in the same branch.
During iteration, preferuv run pre-commit run --files <changed files...>.
Before review or handoff, runuv run pre-commit run --all-files.
Files:
python/nemo_relay/_native.pyicrates/python/tests/coverage/coverage_tests.rspython/nemo_relay/llm.pycrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rspython/tests/test_llm.pycrates/python/src/py_api/mod.rscrates/python/src/py_types/core.rscrates/python/src/py_callable.rs
python/nemo_relay/**/*
⚙️ CodeRabbit configuration file
python/nemo_relay/**/*: Review Python wrapper changes for typed API consistency, contextvars-based scope isolation, async behavior, and parity with the native extension.
Stubs and runtime implementations should stay aligned.
Files:
python/nemo_relay/_native.pyipython/nemo_relay/llm.py
**/*.rs
📄 CodeRabbit inference engine (.agents/skills/prepare-pr/SKILL.md)
**/*.rs: Any Rust change must runjust test-rust
Any Rust change must runcargo fmt --all
Any Rust change must runcargo clippy --workspace --all-targets -- -D warnings
**/*.rs: Runcargo fmt --allfor all FFI work since it is Rust work
Runjust test-rustto validate FFI changes
Runcargo clippy --workspace --all-targets -- -D warningsto enforce strict linting on FFI workWhen Rust files changed as part of Go work, also run
cargo fmt --all,just test-rust, andcargo clippy --workspace --all-targets -- -D warnings
**/*.rs: Runcargo fmt --allwhen Rust files are changed as part of Node work
Runcargo clippy --workspace --all-targets -- -D warningswhen Rust files are changed as part of Node work
Runjust test-rustwhen Rust files are changed as part of Node workWhen changing the core Rust runtime or Rust-facing API surface, format Rust code with
cargo fmt(rustfmt defaults), keepcargo clippy -- -D warningsclean, and satisfycargo deny checkperdeny.toml.
**/*.rs: If any Rust code changed, always runjust test-rust.
If any Rust code changed, also runcargo fmt --all.
If any Rust code changed, also runcargo clippy --workspace --all-targets -- -D warnings.
For Rust changes headed for review, runcargo fmt --allandcargo clippy --workspace --all-targets -- -D warningseven if relying on pre-commit.
Files:
crates/python/tests/coverage/coverage_tests.rscrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rscrates/python/src/py_api/mod.rscrates/python/src/py_types/core.rscrates/python/src/py_callable.rs
**/*.{rs,py}
📄 CodeRabbit inference engine (AGENTS.md)
Follow binding naming conventions in Rust and Python: use
snake_case.
Files:
crates/python/tests/coverage/coverage_tests.rspython/nemo_relay/llm.pycrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rspython/tests/test_llm.pycrates/python/src/py_api/mod.rscrates/python/src/py_types/core.rscrates/python/src/py_callable.rs
**/*.{rs,py,js,mjs,cjs,ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{rs,py,js,mjs,cjs,ts,tsx}: UseJson = serde_json::Valuein Rust-facing runtime APIs where the existing code expects JSON payloads.
UseResult<T>withFlowErrorin core runtime paths, and keep errors explicit and binding-appropriate at the wrapper layer.
Keep async behavior on the existing tokio-based model; bindings should preserve callback and future lifetimes rather than blocking or hiding async work unexpectedly.
Files:
crates/python/tests/coverage/coverage_tests.rspython/nemo_relay/llm.pycrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rspython/tests/test_llm.pycrates/python/src/py_api/mod.rscrates/python/src/py_types/core.rscrates/python/src/py_callable.rs
**/*.{rs,py,go,js,ts,c,h}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
Use language-appropriate naming conventions: Rust
snake_case, C FFI exports prefixednemo_relay_, GoPascalCase, Node.jscamelCase, and Pythonsnake_case.
Files:
crates/python/tests/coverage/coverage_tests.rspython/nemo_relay/llm.pycrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rspython/tests/test_llm.pycrates/python/src/py_api/mod.rscrates/python/src/py_types/core.rscrates/python/src/py_callable.rs
**/*.{rs,go,js,ts}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
Add the SPDX license header to all Rust, Go, JavaScript, and TypeScript source files using the corresponding
//comment form.
Files:
crates/python/tests/coverage/coverage_tests.rscrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rscrates/python/src/py_api/mod.rscrates/python/src/py_types/core.rscrates/python/src/py_callable.rs
**/*.{rs,py,go,js,ts}
📄 CodeRabbit inference engine (.agents/skills/validate-change/SKILL.md)
If a language surface changed, always run that language's test target even when Rust core did not change.
Files:
crates/python/tests/coverage/coverage_tests.rspython/nemo_relay/llm.pycrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rspython/tests/test_llm.pycrates/python/src/py_api/mod.rscrates/python/src/py_types/core.rscrates/python/src/py_callable.rs
**/*.{rs,py,js,ts,tsx,go,java,kt,swift}
📄 CodeRabbit inference engine (.agents/skills/add-middleware/SKILL.md)
Add tests covering registration and duplicate names, deregistration and missing names, priority ordering, callback failure policy, scope-local inheritance and cleanup, event payload semantics, immutable mark and scope fields, and parity across affected bindings.
Files:
crates/python/tests/coverage/coverage_tests.rspython/nemo_relay/llm.pycrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rspython/tests/test_llm.pycrates/python/src/py_api/mod.rscrates/python/src/py_types/core.rscrates/python/src/py_callable.rs
crates/{python,ffi,node}/**/*
⚙️ CodeRabbit configuration file
crates/{python,ffi,node}/**/*: Treat binding changes as public API changes. Check for parity with the other language bindings, FFI ownership/lifetime safety,
callback error propagation, stable type conversion, and consistent async/stream semantics.
Flag changes that update one binding without corresponding tests or documentation for the same surface elsewhere.
Files:
crates/python/tests/coverage/coverage_tests.rscrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rscrates/python/src/py_api/mod.rscrates/python/src/py_types/core.rscrates/python/src/py_callable.rs
{crates/**/tests/**,python/tests/**,go/nemo_relay/**/*_test.go}
⚙️ CodeRabbit configuration file
{crates/**/tests/**,python/tests/**,go/nemo_relay/**/*_test.go}: Tests should cover the behavior promised by the changed API surface, including error paths and cross-request isolation where relevant.
Prefer assertions on lifecycle events, scope stacks, middleware ordering, and binding parity over shallow smoke tests.
Files:
crates/python/tests/coverage/coverage_tests.rscrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rspython/tests/test_llm.py
python/nemo_relay/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Python wrapper modules live under
python/nemo_relay/, and the native extension is built fromcrates/pythonwithmaturin.
Files:
python/nemo_relay/llm.py
**/*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
**/*.py: When changing the Python wrapper package, tests, or docs tooling, lint with Ruff (E,F,W,I), format with Ruff formatter (120-character lines, double quotes), and passtytype checking.
Add the SPDX license header to all Python source files using the#comment form.
Files:
python/nemo_relay/llm.pypython/tests/test_llm.py
{crates/**/src/**/*.rs,python/**/*.py}
📄 CodeRabbit inference engine (.agents/skills/maintain-dynamic-plugins/SKILL.md)
Do not add tests under
src; Rust tests belong in cratetests/trees, and Python SDK tests belong underpython/tests.
Files:
python/nemo_relay/llm.pypython/tests/test_llm.pycrates/python/src/py_api/mod.rscrates/python/src/py_types/core.rscrates/python/src/py_callable.rs
**/*.{py,go,js,ts}
📄 CodeRabbit inference engine (.agents/skills/maintain-observability/SKILL.md)
Keep Python, Go, and Node.js config objects and subscriber/exporter methods aligned so all bindings expose the same logical knobs and semantics.
Files:
python/nemo_relay/llm.pypython/tests/test_llm.py
python/tests/**/*.py
📄 CodeRabbit inference engine (.agents/skills/test-python-binding/SKILL.md)
python/tests/**/*.py: Pytest is used to run tests.
Do not add@pytest.mark.asyncioto any test; async tests are automatically detected and run by the async runner.
Do not add a-> Nonereturn type annotation to test functions.
When mocking a class, do not define a new class; useunittest.mock.MagicMockorunittest.mock.AsyncMock, with thespecconstructor argument when necessary.
Name mocked classes with themockprefix, notfake.
Prefer pytest fixtures over helper methods.
Do not repeat fixtures; if a fixture is needed in multiple test files, place it in aconftest.pyfile.
When creating a fixture, use@pytest.fixture(name="<fixture_name>"[, scope="<scope>"])and define the fixture function asdef <fixture_name>_fixture() -> <return_type>:; only specifyscopewhen it is notfunction.
Preferpytest.mark.parametrizeover creating individual tests for different input types.
Files:
python/tests/test_llm.py
🪛 Ruff (0.15.21)
python/tests/test_llm.py
[warning] 520-520: Missing return type annotation for private function stream_func
(ANN202)
[warning] 520-520: Unused function argument: request
(ARG001)
[warning] 531-531: Unused lambda argument: chunk
(ARG005)
[warning] 532-532: Prefer dict over useless lambda
Replace with lambda with dict
(PIE807)
🔇 Additional comments (10)
crates/python/tests/coverage/coverage_tests.rs (1)
223-226: LGTM!crates/python/tests/coverage/py_api_coverage_tests.rs (1)
822-825: LGTM!crates/python/tests/coverage/py_types_coverage_tests.rs (1)
789-796: LGTM!Also applies to: 826-833, 851-858
python/nemo_relay/_native.pyi (1)
936-938: LGTM!python/nemo_relay/llm.py (1)
317-319: LGTM!python/tests/test_llm.py (1)
6-6: LGTM!Also applies to: 516-540
crates/python/src/py_callable.rs (1)
26-26: LGTM!Also applies to: 35-38, 126-227, 241-259, 270-278, 474-486
crates/python/tests/coverage/py_callable_coverage_tests.rs (1)
357-385: LGTM!crates/python/src/py_api/mod.rs (1)
103-126: LGTM!Also applies to: 882-894
crates/python/src/py_types/core.rs (1)
4-5: LGTM!Also applies to: 28-31, 40-80
| let (tx, rx) = tokio::sync::mpsc::channel(1); | ||
| drop(rx); | ||
| let (_cancel, cancel_rx) = tokio::sync::watch::channel(false); | ||
| forward_async_iter( | ||
| Arc::new(Python::attach(|py| { | ||
| value_iter_cls | ||
| .call1(py, (dropped_payload.bind(py),)) | ||
| .unwrap() | ||
| })), | ||
| tx, | ||
| cancel_rx, | ||
| ) | ||
| .await; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Assert that dropping the receiver closes the iterator.
This test only verifies that forwarding returns, so it misses the cleanup regression at crates/python/src/py_callable.rs Lines 228-235. Give the iterator an observable aclose() and assert it was invoked.
As per path instructions, “Tests should cover the behavior promised by the changed API surface, including error paths.”
🤖 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 `@crates/python/tests/coverage/py_callable_coverage_tests.rs` around lines 387
- 399, Update the test around forward_async_iter to give the Python iterator an
observable aclose() method, then assert that aclose() is invoked after the
receiver is dropped. Keep the existing dropped-receiver setup and await
behavior, but verify the iterator cleanup side effect rather than only
confirming that forwarding returns.
Source: Path instructions
| pyo3_async_runtimes::tokio::future_into_py(py, async move { | ||
| cancel.send_replace(true); | ||
| while !*closed.borrow() { | ||
| closed.changed().await.map_err(|_| { | ||
| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>( | ||
| "stream close task ended before releasing the native stream", | ||
| ) | ||
| })?; | ||
| } | ||
| Ok(()) | ||
| }) | ||
| } | ||
| } | ||
|
|
There was a problem hiding this comment.
Nice cleanup. This closes and drains buffered chunks before aclose() returns, so the next read stops consistently.
| pyo3_async_runtimes::tokio::future_into_py(py, async move { | |
| cancel.send_replace(true); | |
| while !*closed.borrow() { | |
| closed.changed().await.map_err(|_| { | |
| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>( | |
| "stream close task ended before releasing the native stream", | |
| ) | |
| })?; | |
| } | |
| Ok(()) | |
| }) | |
| } | |
| } | |
| let cancel = self.cancel.clone(); | |
| let mut closed = self.closed.clone(); | |
| let receiver = Arc::clone(&self.receiver); | |
| pyo3_async_runtimes::tokio::future_into_py(py, async move { | |
| cancel.send_replace(true); | |
| while !*closed.borrow() { | |
| closed.changed().await.map_err(|_| { | |
| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>( | |
| "stream close task ended before releasing the native stream", | |
| ) | |
| })?; | |
| } | |
| let mut receiver = receiver.lock().await; | |
| receiver.close(); | |
| while receiver.try_recv().is_ok() {} | |
| Ok(()) | |
| }) |
ericevans-nv
left a comment
ericevans-nv
left a comment
There was a problem hiding this comment.
The new stream-close API is a valuable improvement, but the closure acknowledgment is currently nondeterministic. Repeating the new regression test caused stream.aclose() to return while the underlying generator’s finally block had not completed; run 18 of a 20-run loop timed out waiting for cleanup.
Please make aclose() wait for the full downstream Python-generator cleanup—not only the outer forwarding task’s cancellation—and ensure cleanup failures can be represented appropriately. Once that lifecycle is deterministic, the existing regression should pass consistently under repetition.
| } | ||
| } | ||
| } | ||
| closed.send_replace(true); |
There was a problem hiding this comment.
aclose() waits for this closed flag, but it is published before the owned stream is dropped and before downstream Python-generator cleanup is acknowledged. Repeating the new regression test caused aclose() to return while the generator’s finally block never completed within the timeout. Could the completion signal represent the full underlying iterator cleanup before aclose() returns?
ericevans-nv
left a comment
ericevans-nv
left a comment
There was a problem hiding this comment.
Approved. The intermittent close-acknowledgment behavior identified in review is a non-critical reliability concern and can be handled as follow-up work.
3 participants
Overview
Add explicit closure for partially consumed Python LLM streams so producer tasks do not remain blocked.
Details
LlmStream.aclose()API and cancellation propagation through the native stream bridge.Where should the reviewer start?
Start with
crates/python/src/py_callable.rs, which owns the async-generator bridge and cancellation path; then reviewpython/tests/test_llm.pyfor the regression scenario.Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Summary by CodeRabbit
New Features
aclose()support for LLM streaming results, allowing partially consumed streams to be stopped and resources released.Documentation
aclose()when stopping before completion.Tests