Skip to content

fix(distserve): coalesce PD connection warmup to prevent Proxy OOM - #4980

Draft
gokay-ai wants to merge 2 commits into
InternLM:mainfrom
gokay-ai:cursor/pd-warmup-conn-coalesce-96e8
Draft

gokay-ai wants to merge 2 commits into
InternLM:mainfrom
gokay-ai:cursor/pd-warmup-conn-coalesce-96e8

Conversation

@gokay-ai

Copy link
Copy Markdown

Motivation

Fixes #4972.

/distserve/connection_warmup gathered connect() for every Prefill×Decode pair with no in-flight coalescing or waiter timeout. Slow or stuck handshakes piled waiting_conn items and wait_for_conn tasks until the Proxy was OOM-killed (exit 137).

Modification

  • Per-link single-flight in PDConnectionPool: later connect() calls join the existing Connecting event instead of enqueueing another handshake or wait_for_conn task.
  • Bound connect() waiters with LMDEPLOY_PD_CONN_WAIT_TIMEOUT (default 60s; 0 restores unlimited wait).
  • Cap waiting_conn as a safety net.
  • Make warmup idempotent: concurrent /distserve/connection_warmup callers share one in-flight gather, with a semaphore limiting concurrent handshakes.
  • Add unit tests for coalescing and waiter timeout.

BC-breaking (Optional)

No API break. DistServe-only. Operators who need unbounded PD handshake waits can set LMDEPLOY_PD_CONN_WAIT_TIMEOUT=0.

Checklist

  1. Pre-commit or other linting tools are used to fix the potential lint issues.
  2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
  3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
  4. The documentation has been modified accordingly, like docstring or example tutorials.

Fixes InternLM#4972

Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com>
# Bound waiters so connect()/warmup cannot hang forever when a handshake is stuck.
# Set LMDEPLOY_PD_CONN_WAIT_TIMEOUT=0 to restore unlimited wait.
DEFAULT_CONNECT_WAIT_TIMEOUT = 60.0
CONNECT_WAIT_TIMEOUT = _optional_float_env('LMDEPLOY_PD_CONN_WAIT_TIMEOUT', DEFAULT_CONNECT_WAIT_TIMEOUT)

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.

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.

Moved LMDEPLOY_PD_CONN_WAIT_TIMEOUT into lmdeploy/pytorch/envs.py so get_all_envs() registers it.

Move LMDEPLOY_PD_CONN_WAIT_TIMEOUT into lmdeploy/pytorch/envs.py so
get_all_envs() can propagate it, and apply docformatter on touched files.

Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com>
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] Proxy connection warmup can accumulate unbounded PD connection wait tasks and OOM

2 participants