Skip to content

[WIP] fix: correct callback event firing in streaming loops - #736

Draft
bhimrazy wants to merge 1 commit into
Lightning-AI:mainfrom
bhimrazy:fix/735-streaming-loop-callback-events
Draft

[WIP] fix: correct callback event firing in streaming loops#736
bhimrazy wants to merge 1 commit into
Lightning-AI:mainfrom
bhimrazy:fix/735-streaming-loop-callback-events

Conversation

@bhimrazy

@bhimrazy bhimrazy commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

What does this PR do?

Fixes #735.

The sync streaming loop fired AFTER_PREDICT twice and never fired the decode-request events.

While fixing it, a second issue surfaced: predict and encode_response return generators, so no work has happened when they return — it all runs interleaved as the stream is consumed. Firing the "after" events at that point meant PredictionTimeLogger reported Prediction took 0.00 seconds on every streaming request.

So all three streaming paths now fire:

BEFORE_DECODE_REQUEST → AFTER_DECODE_REQUEST → BEFORE_PREDICT
  → BEFORE_ENCODE_RESPONSE → [stream consumed] → AFTER_PREDICT → AFTER_ENCODE_RESPONSE

each exactly once.

Note: the async and BatchedStreamingLoop changes go beyond the reported issue — same bug, but they shift event timing for existing users.

The sync streaming loop fired AFTER_PREDICT twice and never fired the
decode-request events. Separately, all three streaming paths fired the
"after" events before the stream was consumed - since predict and
encode_response return generators, no work has happened at that point,
so AFTER_PREDICT always measured ~0s.

Add the decode-request events and move AFTER_PREDICT /
AFTER_ENCODE_RESPONSE to after the stream is exhausted.

Fixes Lightning-AI#735
@codecov-commenter

codecov-commenter commented Aug 26, 2026

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 60.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 85%. Comparing base (4c9fecc) to head (6146f25).
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files
@@         Coverage Diff         @@
##           main   #736   +/-   ##
===================================
- Coverage    85%    85%   -0%     
===================================
  Files        39     39           
  Lines      3285   3286    +1     
===================================
- Hits       2781   2780    -1     
- Misses      504    506    +2     
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Sync streaming loop fires AFTER_PREDICT twice and skips decode-request callback events

2 participants