Skip to content

[Fix] Preserve SessionServer traces across response failures - #2059

Closed
matrix72c wants to merge 1 commit into
InternLM:mainfrom
matrix72c:fix/session-server-untraceable-response
Closed

[Fix] Preserve SessionServer traces across response failures#2059
matrix72c wants to merge 1 commit into
InternLM:mainfrom
matrix72c:fix/session-server-untraceable-response

Conversation

@matrix72c

@matrix72c matrix72c commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR prevents trace-enabled SessionServer generations from being silently lost when the downstream client disconnects before response headers are prepared or when an upstream successful response cannot be recorded.

It also keeps upstream error responses out of generation trace processing and returns those errors unchanged to the client.

The endpoint-classification and proxy allowlist changes originally included in this PR have been split into #2061. This PR now contains only trace-response reliability behavior and is independently based on upstream/main.

Problem

Disconnect before response preparation

The streaming path already stopped writing after a downstream disconnect while continuing to drain the upstream response for trace completion. However, response.prepare(request) ran before that protected state was established. If the client disconnected before headers were prepared, the exception escaped and the upstream response was no longer drained, leaving the training trace incomplete.

Successful but untraceable response

For a trace-enabled non-streaming request, a JSON array or non-JSON 2xx response previously left response_data unset and returned the successful upstream response without calling on_response. The client could observe or record a successful turn while the trace store had no corresponding node.

Upstream error payloads

Non-2xx responses are capability or request errors, not successful generations. They should remain visible to the caller in their original form and must not enter response cleaning or trace hooks.

Implementation

  • Wrap downstream stream preparation in a small helper that treats connection-reset failures as a dead downstream client.
    • The proxy still drains the upstream stream.
    • Subsequent writes are skipped.
    • A complete trace can still be parsed and recorded.
  • Derive trace_response from both the request trace flag and the upstream HTTP status.
    • Only successful generation responses are buffered and passed to on_response.
    • Upstream 4xx/error payloads bypass response mutation and trace hooks.
  • Validate trace-enabled successful non-stream responses before recording.
    • JSON objects continue to the existing response hook.
    • JSON arrays and non-JSON bodies produce a controlled native-format 500 error instead of silently omitting the trace.
    • Structurally malformed generation objects still reach on_response, whose existing fail-closed error handling remains authoritative.

No endpoint allowlist, count_tokens, hello, or management-route policy is introduced here; those changes are isolated in #2061.

Tests

Added focused local aiohttp fake-upstream tests covering:

  • downstream connection loss before response headers are prepared;
  • successful stream preparation;
  • malformed generation envelopes reaching the response hook and failing closed;
  • JSON-list and non-JSON successful responses producing controlled errors; and
  • upstream 4xx bodies being returned unchanged without invoking on_response.

Validation performed on this topic alone:

pytest -q tests/rl/test_session_server_trace.py
6 passed

ruff check xtuner/v1/rl/rollout/session_server.py tests/rl/test_session_server_trace.py
All checks passed

git diff --check upstream/main...fix/session-server-untraceable-response

This topic merges cleanly with #2061 in the local integration branch; the combined endpoint and trace suites pass all 12 tests.

Scope

  • No request endpoint classification or forwarding policy changes.
  • No LMDeploy changes.
  • No public API changes.

@matrix72c
matrix72c force-pushed the fix/session-server-untraceable-response branch from ee51a92 to 4104f70 Compare September 1, 2026 09:06
@matrix72c matrix72c changed the title Fix SessionServer endpoint classification and trace handling Preserve SessionServer traces across response failures Sep 1, 2026
@matrix72c matrix72c closed this Sep 2, 2026
@matrix72c matrix72c changed the title Preserve SessionServer traces across response failures [Fix] Preserve SessionServer traces across response failures Sep 2, 2026
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.

1 participant