Skip to content

[Fix] Reject untrusted DistServe fields on Decode OpenAI APIs - #4969

Open
gokay-ai wants to merge 1 commit into
InternLM:mainfrom
gokay-ai:cursor/fix-distserve-migration-request-d1ec
Open

gokay-ai wants to merge 1 commit into
InternLM:mainfrom
gokay-ai:cursor/fix-distserve-migration-request-d1ec

Conversation

@gokay-ai

Copy link
Copy Markdown

Motivation

Decode /v1/completions and /v1/chat/completions accepted client migration_request after Pydantic shape checks only. A syntactically valid but semantically bad payload can raise in the Decode migration loop. EngineLoop waits with asyncio.FIRST_EXCEPTION, so that one request can stop the Decode engine until restart (/health unhealthy, lasting 503s).

Fixes #4965

Modification

  • Public OpenAI completions/chat endpoints reject DistServe-only fields (migration_request, with_cache, preserve_cache) unless the request is Prefill/Decode replica traffic marked by the DistServe proxy.
  • DistServe proxy sets x-lmdeploy-distserve-proxy on replica POSTs so legitimate PD traffic still works.
  • Decode migration failures are isolated to the offending request instead of cancelling the whole EngineLoop task group.

BC-breaking

External clients can no longer send DistServe-only fields on /v1/completions or /v1/chat/completions. DistServe through the proxy is unchanged.

Checklist

  1. Pre-commit / lint to be run in CI.
  2. Unit tests added for the public-path reject path and migration isolation.
  3. N/A
  4. N/A

Copilot AI lite review requested due to automatic review settings September 14, 2026 15:27

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

return str(value).strip()


def is_trusted_distserve_replica_request(raw_request, engine_config) -> bool:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the fix. I verified that the original reproducer from #4965 now results in a request-level error rather than terminating the engine. Health checks and subsequent generation remained successful.

is_trusted_distserve_replica_request() checks the receiving server’s role and a fixed, client-controlled header. This identifies the intended request path, but does not authenticate the caller.
If P/D replicas are internal-only services protected by network access controls, this is a reasonable deployment model, consistent with vLLM’s documented network-isolation requirements.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for verifying. Agreed — this is meant for DistServe P/D replicas behind network isolation (same model as vLLM’s inter-node guidance), not as a substitute for authenticating arbitrary public clients. Happy to adjust further if you want a stronger internal-auth check on top.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is enough. Note that there are conflicts with the latest main. After resolving it, the pr is ready to merge.

@cursor
cursor Bot force-pushed the cursor/fix-distserve-migration-request-d1ec branch from 817259a to 810824f Compare September 16, 2026 06:28
@@ -0,0 +1,211 @@
# Copyright (c) OpenMMLab. All rights reserved.
"""Reject untrusted DistServe fields on the public OpenAI path."""

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test file adds little value. I suggest removing it.

assert context.async_engine.call_count == 0


def test_chat_completions_still_serves_after_untrusted_migration_request(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test adds little value and can be removed.

Public /v1/completions and /v1/chat/completions accepted client
migration_request after Pydantic shape checks only. A bad payload
could raise in the Decode migration loop and tear down EngineLoop.

Reject DistServe-only fields on the public path unless the DistServe
proxy marks Prefill/Decode replica traffic. Isolate per-request
migration failures so they cannot cancel the engine task group.

Rebased onto latest main: use Scheduler.get_block_tables after the
scheduler ownership refactor.

Fixes InternLM#4965

Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com>
@cursor
cursor Bot force-pushed the cursor/fix-distserve-migration-request-d1ec branch from 810824f to e4bcd8a Compare September 16, 2026 14:48
@gokay-ai

Copy link
Copy Markdown
Author

Removed the two test files per @lvhan028.

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.

[Bug] Untrusted migration_request can terminate the DistServe Decode EngineLoop and cause persistent denial of service

4 participants