Lightning-AI / Lightning-AI/LitServe
Sync streaming loop fires AFTER_PREDICT twice and skips decode-request callback events
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.9k
- Forks
- 304
- Avg merge
- 3d 13h
- Merged PRs (30d)
- 6
Description
## Summary
In the synchronous streaming loop (`StreamingLoop.run_streaming_loop`), lifecycle callback events are fired incorrectly: `AFTER_PREDICT` is triggered **twice** per request, and `BEFORE_DECODE_REQUEST` / `AFTER_DECODE_REQUEST` are never triggered at all. Every other execution path in the codebase fires each event exactly once.
*This report is based on static analysis of master; I have not executed the code.*
## Location
- File: `src/litserve/loops/streaming_loops.py`
- Function: `StreamingLoop.run_streaming_loop`
- Contrast: `src/litserve/loops/simple_loops.py::run_loop` and `_process_single_request`, and `streaming_loops.py::_process_streaming_request`
## Problem
`run_streaming_loop` does:
```python
callback_runner.trigger_event(EventTypes.BEFORE_PREDICT.value, lit_api=lit_api)
y_gen = _inject_context(context, lit_api.predict, x)
callback_runner.trigger_event(EventTypes.AFTER_PREDICT.value, lit_api=lit_api) # (1) fired immediately
...
for y_enc in y_enc_gen:
...
self.put_response(..., LitAPIStatus.FINISH_STREAMING, ...)
callback_runner.trigger_event(EventTypes.AFTER_PREDICT.value, lit_api=lit_api) # (2) fired AGAIN
callback_runner.trigger_event(EventTypes.AFTER_ENCODE_RESPONSE.value, lit_api=lit_api)
```
1. `AFTER_PREDICT` fires twice per request. Any user callback registered on `AFTER_PREDICT` (e.g. a metric/logging callback) runs two times for every streamed request.
2. The decode phase events are missing: `decode_request` is invoked via `_inject_context` without triggering `BEFORE_DECODE_REQUEST`/`AFTER_DECODE_REQUEST`, so callbacks listening for those events silently never run in this loop.
For comparison, the regular sync loop (`simple_loops.run_loop`) fires the full sequence exactly once:
```python
BEFORE_DECODE_REQUEST → decode_request → AFTER_DECODE_REQUEST
→ BEFORE_PREDICT → predict → AFTER_PREDICT
→ BEFORE_ENCODE_RESPONSE → encode_response → AFTER_ENCODE_RESPONSE
```
and the async paths (`_process_streaming_request`, `_process_single_request`) also fire each event once.
Additionally, the first `AFTER_PREDICT` (1) is semantically premature: `lit_api.predict` returns a generator at that point, so no prediction work has happened yet — the meaningful "after predict" point is when the stream has been consumed.
## Trigger / Reproduction
Register any callback on `AFTER_PREDICT` (or on the decode-request events) and serve a streaming LitAPI with the default synchronous worker setup; observe callbacks firing twice / not at all respectively. (Static-analysis finding — not verified by running the server.)
## Expected Behavior
Each request triggers the same event sequence exactly once in all loops: decode events around `decode_request`, one `AFTER_PREDICT` after prediction completes, etc.
## Actual Behavior
Sync streaming requests trigger `AFTER_PREDICT` twice and never trigger the decode-phase events.
## Impact
Callbacks (metrics, logging, custom hooks) behave inconsistently depending on whether an API streams and which loop executes: double-counted metrics or missing instrumentation for streaming endpoints.
## Suggested Direction
Mirror `run_loop`'s event placement: add the decode-phase events around `decode_request`, keep a single `AFTER_PREDICT` (most consistently placed right before `BEFORE_ENCODE_RESPONSE`), and keep `AFTER_ENCODE_RESPONSE` after the stream finishes.
## Evidence
Side-by-side source of `run_streaming_loop` vs `run_loop`/`_process_streaming_request` as shown above; both files live under `src/litserve/loops/`.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/litserve/loops/streaming_loops.py at StreamingLoop.run_streaming_loop, then compare its callback placement with src/litserve/loops/simple_loops.py::run_loop and _process_single_request, plus streaming_loops.py::_process_streaming_request. Reproduce or inspect the synchronous streaming path and verify that decode events fire around decoding, AFTER_PREDICT fires once after prediction, and the completed stream retains the expected response events.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 84/100