Lightning-AI / Lightning-AI/LitServe

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

Open
#735 0 comments 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.