feat: Allow Python Engine to end stream before final

[kthui]

Overview:

Allow the Python engine to raise GeneratorExit exception while generating response, which will end the stream without final, so the request is migrated by the frontend if migration is enabled.

Details:

Should be merged together with #2280

Where should the reviewer start?

Start with the backend docs, and then move to the Python-Rust binding changes.

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

N/A

Summary by CodeRabbit

• Bug Fixes
  • Improved handling of Python generator exits in async stream processing, resulting in clearer error messages when streams end unexpectedly.
  • Enhanced error detection to prevent sending a final completion message if a stream ends prematurely, ensuring more accurate stream termination behavior.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Walkthrough

A new error variant for Python's GeneratorExit was added to the Rust engine, allowing for explicit detection and handling of generator termination events. The runtime's push handler was updated to recognize this error, log a warning, and adjust stream completion behavior. Trait bounds were also updated to include MaybeError.

Changes

Cohort / File(s)	Change Summary
Python GeneratorExit Handling in Rust Engine lib/bindings/python/rust/engine.rs	Added PyGeneratorExit(String) variant to ResponseProcessingError; updated error matching and async stream logic to handle Python GeneratorExit distinctly.
Push Handler Stream Termination Handling lib/runtime/src/pipeline/network/ingress/push_handler.rs	Extended trait bound for U to include MaybeError; added logic to detect and handle "Stream ended before generation completed" error, logging a warning and altering finalization.

Sequence Diagram(s)

sequenceDiagram
    participant PythonAsyncGen as Python Async Generator
    participant RustEngine as Rust Engine
    participant PushHandler as Push Handler

    PythonAsyncGen->>RustEngine: Yields item or raises exception
    RustEngine->>RustEngine: Detects exception type
    alt Exception is GeneratorExit
        RustEngine->>PushHandler: Return PyGeneratorExit error
        PushHandler->>PushHandler: Log warning, set flag, break loop
    else Other Exception
        RustEngine->>PushHandler: Return PythonException error
        PushHandler->>PushHandler: Handle as generic error
    end

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~15 minutes

Poem

A rabbit hopped through Rust and Py,
Saw generators wave goodbye.
Now with care, the stream detects
When Python’s exit interjects.
Warnings logged, the flow is neat,
No endless loops, just clean retreat.
🐇✨

Note

⚡️ Unit Test Generation is now available in beta!

Learn more here, or try it out under "Finishing Touches" below.

---

📜 Recent review details

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between dbb4caa and 727080c.

📒 Files selected for processing (2)

• lib/bindings/python/rust/engine.rs (3 hunks)
• lib/runtime/src/pipeline/network/ingress/push_handler.rs (3 hunks)

🧰 Additional context used

🧠 Learnings (5)

📓 Common learnings

Learnt from: ryanolson
PR: ai-dynamo/dynamo#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: kthui
PR: ai-dynamo/dynamo#1424
File: lib/runtime/src/pipeline/network/egress/push_router.rs:204-209
Timestamp: 2025-06-13T22:07:24.843Z
Learning: The codebase uses async-nats version 0.40, not the older nats crate. Error handling should use async_nats::error::Error variants, not nats::Error variants.

📚 Learning: the asyncenginecontextprovider trait in lib/runtime/src/engine.rs was intentionally changed from `se...

Learnt from: ryanolson
PR: ai-dynamo/dynamo#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.

Applied to files:

• lib/bindings/python/rust/engine.rs
• lib/runtime/src/pipeline/network/ingress/push_handler.rs

📚 Learning: the codebase uses async-nats version 0.40, not the older nats crate. error handling should use async...

Learnt from: kthui
PR: ai-dynamo/dynamo#1424
File: lib/runtime/src/pipeline/network/egress/push_router.rs:204-209
Timestamp: 2025-06-13T22:07:24.843Z
Learning: The codebase uses async-nats version 0.40, not the older nats crate. Error handling should use async_nats::error::Error variants, not nats::Error variants.

Applied to files:

• lib/bindings/python/rust/engine.rs
• lib/runtime/src/pipeline/network/ingress/push_handler.rs

📚 Learning: in lib/llm/src/kv_router/scoring.rs, the user prefers to keep the panic behavior when calculating lo...

Learnt from: PeaBrane
PR: ai-dynamo/dynamo#1285
File: lib/llm/src/kv_router/scoring.rs:58-63
Timestamp: 2025-05-30T06:38:09.630Z
Learning: In lib/llm/src/kv_router/scoring.rs, the user prefers to keep the panic behavior when calculating load_avg and variance with empty endpoints rather than adding guards for division by zero. They want the code to fail fast on this error condition.

Applied to files:

• lib/runtime/src/pipeline/network/ingress/push_handler.rs

📚 Learning: in async-nats, the "no responders" error is represented as async_nats::error::requesterrorkind::nore...

Learnt from: kthui
PR: ai-dynamo/dynamo#1424
File: lib/runtime/src/pipeline/network/egress/push_router.rs:204-209
Timestamp: 2025-06-13T22:32:05.022Z
Learning: In async-nats, the "no responders" error is represented as async_nats::error::RequestErrorKind::NoResponders. Use err.downcast_ref::<async_nats::error::RequestError>() and then check req_err.kind() against RequestErrorKind::NoResponders to handle this error properly.

Applied to files:

• lib/runtime/src/pipeline/network/ingress/push_handler.rs

🔇 Additional comments (6)

lib/bindings/python/rust/engine.rs (3)

137-139: LGTM: New error variant for Python GeneratorExit.

The addition of PyGeneratorExit(String) as a distinct error variant allows for specific handling of Python generator exit conditions, which is appropriate for the feature described in the PR objectives.

---

231-233: LGTM: Specific handling for generator exit.

The match arm correctly returns a distinct message "Stream ended before generation completed" for the PyGeneratorExit variant, which enables downstream components to identify and handle this specific termination condition appropriately.

---

285-295: LGTM: Proper detection of Python GeneratorExit exception.

The implementation correctly:

• Uses the GIL to safely check the Python exception type
• Distinguishes between GeneratorExit and other Python exceptions
• Returns the appropriate error variant based on the exception type

This enables the feature to differentiate between normal Python exceptions and intentional generator termination.

lib/runtime/src/pipeline/network/ingress/push_handler.rs (3)

17-17: LGTM: Import for MaybeError trait.

The import is necessary to support the new trait bound on the generic type U.

---

109-109: LGTM: Added MaybeError trait bound.

The addition of MaybeError trait bound to U is necessary to enable checking for errors in the response stream, which supports the new generator exit handling logic.

---

224-231: LGTM: Specific handling for stream termination.

The implementation correctly:

• Checks for the specific error message that indicates generator exit
• Logs an appropriate warning
• Sets send_complete_final to false to prevent sending the final completion message
• Breaks early from the loop

This coordinates well with the Python engine changes to handle streams that end before completion due to generator exit.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[tedzhouhk]

tested to be working

[rmccorm4]

One question, but not blocking approval: https://github.com/ai-dynamo/dynamo/pull/2270/files#r2253082585

[kthui]

Merging this PR now to unblock #2280, will improve on error type handling with DIS-364.