feat(request cancellation): pycontext, propagating the is_stopped into python land.
#2158
Conversation
michaelfeil
requested review from
a team,
GuanLuo,
PeaBrane,
alec-flowers,
biswapanda,
grahamking,
ishandhanani,
jthomson04,
kkranen,
nnshah1,
paulhendricks,
piotrm-nvidia,
ptarasiewiczNV,
rmccorm4,
ryanolson,
tanmayv25,
tedzhouhk,
tmonty12 and
whoisj
as code owners
|
👋 Hi michaelfeil! Thank you for contributing to ai-dynamo/dynamo. Just a reminder: The 🚀 |
github-actions
bot
added
the
external-contribution
michaelfeil
changed the title
is_stopped into python land.
WalkthroughA new Python context object ( Changes
Sequence Diagram(s)sequenceDiagram
participant RustEngine
participant PythonGenerator
participant Context
RustEngine->>Context: Create PyContext from AsyncEngineContext
RustEngine->>PythonGenerator: Call generator(request, context)
PythonGenerator->>Context: Access context.id(), is_stopped(), etc.
PythonGenerator-->>RustEngine: Yield/return results
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~18 minutes Poem
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
lib/bindings/python/rust/context.rs (2)
7-7: Fix typo in comment."jsut" should be "just".
-// Not all methods of the AsyncEngineContext are exposed, jsut the primary ones for tracing + cancellation. +// Not all methods of the AsyncEngineContext are exposed, just the primary ones for tracing + cancellation.
32-34: Fix missing semicolon in stop_generating method.For consistency with Rust conventions, void methods should end with a semicolon.
fn stop_generating(&self) { - self.inner.stop_generating(); + self.inner.stop_generating(); }
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
examples/runtime/hello_world/hello_world.py(1 hunks)lib/bindings/python/rust/context.rs(1 hunks)lib/bindings/python/rust/engine.rs(3 hunks)lib/bindings/python/rust/lib.rs(2 hunks)lib/bindings/python/src/dynamo/runtime/__init__.py(3 hunks)
🧰 Additional context used
🧠 Learnings (5)
lib/bindings/python/rust/lib.rs (1)
Learnt from: ryanolson
PR: #1919
File: lib/runtime/src/engine.rs:168-168
Timestamp: 2025-07-14T21:25:56.930Z
Learning: The AsyncEngineContextProvider trait in lib/runtime/src/engine.rs was intentionally changed from Send + Sync + Debug to Send + Debug because the Sync bound was overly constraining. The trait should only require Send + Debug as designed.
examples/runtime/hello_world/hello_world.py (1)
Learnt from: nnshah1
PR: #1444
File: tests/fault_tolerance/utils/metrics.py:30-32
Timestamp: 2025-07-01T13:55:03.940Z
Learning: The @dynamo_worker() decorator in the dynamo codebase returns a wrapper that automatically injects the runtime parameter before calling the wrapped function. This means callers only need to provide the non-runtime parameters, while the decorator handles injecting the runtime argument automatically. For example, a function with signature async def get_metrics(runtime, log_dir) decorated with @dynamo_worker() can be called as get_metrics(log_dir) because the decorator wrapper injects the runtime parameter.
lib/bindings/python/rust/engine.rs (2)
Learnt from: ryanolson
PR: #1919
File: lib/runtime/src/engine.rs:168-168
Timestamp: 2025-07-14T21:25:56.930Z
Learning: The AsyncEngineContextProvider trait in lib/runtime/src/engine.rs was intentionally changed from Send + Sync + Debug to Send + Debug because the Sync bound was overly constraining. The trait should only require Send + Debug as designed.
Learnt from: PeaBrane
PR: #1236
File: lib/llm/src/mocker/engine.rs:140-161
Timestamp: 2025-06-17T00:50:44.845Z
Learning: In Rust async code, when an Arc<Mutex<_>> is used solely to transfer ownership of a resource (like a channel receiver) into a spawned task rather than for sharing between multiple tasks, holding the mutex lock across an await is not problematic since there's no actual contention.
lib/bindings/python/src/dynamo/runtime/__init__.py (1)
Learnt from: nnshah1
PR: #1444
File: tests/fault_tolerance/utils/metrics.py:30-32
Timestamp: 2025-07-01T13:55:03.940Z
Learning: The @dynamo_worker() decorator in the dynamo codebase returns a wrapper that automatically injects the runtime parameter before calling the wrapped function. This means callers only need to provide the non-runtime parameters, while the decorator handles injecting the runtime argument automatically. For example, a function with signature async def get_metrics(runtime, log_dir) decorated with @dynamo_worker() can be called as get_metrics(log_dir) because the decorator wrapper injects the runtime parameter.
lib/bindings/python/rust/context.rs (1)
Learnt from: ryanolson
PR: #1919
File: lib/runtime/src/engine.rs:168-168
Timestamp: 2025-07-14T21:25:56.930Z
Learning: The AsyncEngineContextProvider trait in lib/runtime/src/engine.rs was intentionally changed from Send + Sync + Debug to Send + Debug because the Sync bound was overly constraining. The trait should only require Send + Debug as designed.
🪛 GitHub Actions: Pre Merge Validation of (ai-dynamo/dynamo/refs/pull/2158/merge) by michaelfeil.
examples/runtime/hello_world/hello_world.py
[error] 18-26: isort formatting check failed. Imports were reordered. Please run 'isort' to fix import order.
lib/bindings/python/src/dynamo/runtime/__init__.py
[error] 15-90: isort and black formatting checks failed. File was modified to fix import order and code style. Please run 'isort' and 'black' to fix formatting issues.
🪛 GitHub Actions: Copyright Checks
lib/bindings/python/rust/context.rs
[error] 1-1: Invalid or missing copyright header detected by copyright checker. File is out of compliance with required SPDX header.
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: Build and Test - vllm
- GitHub Check: pre-merge-rust (lib/runtime/examples)
- GitHub Check: pre-merge-rust (.)
🔇 Additional comments (10)
lib/bindings/python/rust/lib.rs (2)
46-46: LGTM!Correctly declares the new context module.
104-104: LGTM!Properly registers the PyContext class in the Python module, making it accessible from Python code.
examples/runtime/hello_world/hello_world.py (1)
30-31: LGTM!The function signature correctly accepts the context parameter and the logging demonstrates proper usage of the context ID. This is a good example of how to use the new context functionality.
lib/bindings/python/rust/engine.rs (3)
18-18: LGTM!Correctly imports the PyContext module for use in the engine.
167-167: LGTM!The context cloning is safe due to the Arc wrapper and enables passing the context to Python.
182-184: LGTM!The PyContext creation and argument passing to the Python generator is correctly implemented. The context is properly wrapped and passed as the second argument, enabling Python code to access the context functionality.
lib/bindings/python/src/dynamo/runtime/__init__.py (3)
32-32: LGTM!Correctly imports and aliases PyContext as Context for Python use.
69-69: LGTM!Properly uses function signature introspection to detect if the function accepts a context parameter.
75-80: LGTM!The context detection and injection logic is well-implemented. It correctly:
- Detects if the last argument is a Context instance
- Separates the context from other arguments
- Injects the context into kwargs if the function accepts it
- Preserves the existing request processing logic
lib/bindings/python/rust/context.rs (1)
44-63: LGTM!The async_is_stopped method is well-implemented with proper:
- Input validation for timeout bounds
- Timeout handling using tokio::time::timeout
- Proper async-to-Python conversion using pyo3_async_runtimes
- Correct logic for checking stopped/killed state
|
@kthui As discussed, I moved the Ready to be merged / please squash and merge. |
kthui
left a comment
kthui
left a comment
There was a problem hiding this comment.
Looks good to me! I like the idea you keep the signature check in Rust.
|
@nnshah1 can you also help review this PR? We need a reviewer with devops and Python permission on the Cargo.lock and .py files. The change passed TRT-LLM and SGLang 1 GPU E2E test on pipeline #33470282 |
rmccorm4
left a comment
rmccorm4
left a comment
There was a problem hiding this comment.
Approving to unblock as this has been sitting a couple days - approved by Jacky as PIC and assuming signoff from @ryanolson @nnshah1 from discussions
|
Congratulations @michaelfeil on getting your first PR merged into main! This feature updates some of the deeper Dynamo interfaces. I believe your future PRs will be much smoother now that you've successfully implemented this in-depth feature. |
…nto python land. (#2158) Signed-off-by: Hannah Zhang <[email protected]>
5 participants
Overview:
Details:
Where should the reviewer start?
Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Summary by CodeRabbit