[Bug] Proxy connection warmup can accumulate unbounded PD connection wait tasks and OOM
- Dominant language
- Python
- Stars
- 8.1k
- Forks
- 748
- Avg merge
- 6d 2h
- Merged PRs (30d)
- 54
Description
### Checklist
- [x] 1. I have searched related issues but cannot get the expected help.
- [x] 2. The bug has not been fixed in the latest version.
- [x] 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback.
### Describe the bug
In LMDeploy v0.17.0 DistServe mode, repeated calls to Proxy `/distserve/connection_warmup` can accumulate unbounded pending tasks during PD connection establishment.
With a limited but realistic topology of 10 Prefill nodes and 10 Decode nodes, each warmup covers 100 Prefill-Decode connection pairs. If the backend endpoints are slow or never complete the connection handshake, repeated warmup requests create many pending `PDConnectionPool.connect()` and `wait_for_conn()` tasks for the same 100 connection pairs. In my test, this caused the Proxy container to be OOM-killed with exit code 137.
This issue requires access to Proxy control-plane/DistServe management endpoints such as `/nodes/add` and `/distserve/connection_warmup`, or an equivalent control-plane position that can register nodes and trigger connection warmup.
### Detail
Analyzed version:
- LMDeploy v0.17.0
- Source snapshot used for analysis: `lmdeploy-main-src/lmdeploy-main`
Relevant source paths:
- `lmdeploy/serve/proxy/proxy.py`
- `lmdeploy/pytorch/disagg/conn/proxy_conn.py`
In `lmdeploy/serve/proxy/proxy.py`, `/distserve/connection_warmup` calls `PDConnectionPool.connect()` for every Prefill x Decode pair:
```python
@app.post('/distserve/connection_warmup', dependencies=[Depends(validate_json_request)])
async def connection_warmup():
await asyncio.gather(*[
node_manager.pd_connection_pool.connect(
PDConnectionMessage(
p_url=p_url,
d_url=d_url,
protocol=node_manager.migration_protocol,
rdma_config=node_manager.rdma_config,
)) for p_url in node_manager.prefill_nodes for d_url in node_manager.decode_nodes
])
return JSONResponse({'SUCCESS': True})
```
In `lmdeploy/pytorch/disagg/conn/proxy_conn.py`, the connection pool has several resource-control gaps:
- `waiting_conn = asyncio.Queue()` has no `maxsize`.
- `connect()` uses `waiting_conn.put_nowait()` for each connection request.
- When a connection pair is already in `Connecting`, `_perform_conn()` creates another `wait_for_conn()` task instead of rejecting, deduplicating, or coalescing the duplicate request.
- `connect()` waits on `await conn_event.wait()` without a request-level timeout.
- `AIOHTTP_TIMEOUT` defaults to `None`, so slow backend HTTP calls may not have a deterministic total timeout.
- `conn_sem` limits only the backend HTTP phase; it does not limit the number of queued warmup requests or `wait_for_conn()` tasks.
As a result, repeated warmup calls against the same finite set of slow P-D pairs can create unbounded pending asyncio tasks even though the connection pool only contains 100 connection keys.
### Reproduction
Test environment:
- LMDeploy v0.17.0
- Proxy-only test container: `lmdeploy0170-pdconn-proxy`
- Proxy memory limit: 4GiB
- Proxy swap limit: 4GiB
- Proxy port: `19200`
- No GPU used by the Proxy test
- 10 fake Prefill HTTP endpoints
- 10 fake Decode HTTP endpoints
- Fake endpoints accepted TCP/HTTP connections but intentionally slept for 3600 seconds to keep connection establishment in `Connecting`
The 4GiB limit was used to make the final availability impact observable in a bounded test run. The underlying issue is unbounded pending task/request accumulation during connection warmup; with a higher memory limit, the same growth pattern would require more warmup requests before causing memory pressure or OOM.
Test parameters:
```text
FAKE_PREFILL_COUNT=10
FAKE_DECODE_COUNT=10
WARMUP_FLOOD_TOTAL=10000
WARMUP_FLOOD_BATCH=100
WARMUP_FLOOD_BATCH_DELAY=0.02
SLOW_NODE_SLEEP_SECONDS=3600
```
Reproduction flow:
1. Start the Proxy with DistServe enabled and a 4GiB memory/swap limit.
2. Start 20 slow fake HTTP endpoints.
3. Register 10 fake Prefill nodes and 10 fake Decode nodes through `/nodes/add`.
4. Repeatedly send requests to `/distserve/connection_warmup`.
5. Observe Proxy connection-pool state, asyncio task count, file descriptors, RSS, and Docker OOM state.
Node registration result:
```text
nodes_count_after_register={"total": 20, "roles": {"2": 10, "3": 10}}
```
Each warmup covered:
```text
10 Prefill x 10 Decode = 100 P-D connection pairs
```
Last successful warmup flood observation before OOM:
```text
attempted=10000
opened=10000
failed=0
held_client_sockets=10000
pool_size=100
pool_status_counts={"Connecting": 100}
waiting_conn_qsize=3300
conn_sem_value=1948
asyncio_tasks_total=345704
asyncio_tasks_pending=345704
PDConnectionPool.connect=169700
PDConnectionPool.connect..wait_for_conn=165900
RequestResponseCycle.run_asgi=10001
PDConnectionPool.connect..conn_worker=100
fd_count=10109
rss_kib=1807900
```
Memory/timeline observations:
```text
2026-09-13T23:51:27+08:00 running=true oom_killed=false exit_code=0 mem=457.3MiB / 4GiB pending_tasks=38505
2026-09-13T23:51:55+08:00 running=true oom_killed=false exit_code=0 mem=2.17GiB / 4GiB pending_tasks=593004
2026-09-13T23:52:21+08:00 running=true oom_killed=false exit_code=0 mem=3.753GiB / 4GiB
2026-09-13T23:52:24+08:00 running=false oom_killed=true exit_code=137 mem=0B / 0B
```
Final Docker state:
```text
name=/lmdeploy0170-pdconn-proxy
running=false
oom_killed=true
exit_code=137
memory=4294967296
memory_swap=4294967296
```
Final test reason:
```text
final_reason=proxy_oom_killed
```
The test used fake slow nodes to keep the PD connection handshake pending. The observed DoS occurred in the Proxy process itself before any real migration was needed.
### Impact
If Proxy management/DistServe control endpoints are reachable by an untrusted or lower-trust actor, that actor can register a limited number of slow or non-completing nodes and repeatedly trigger `/distserve/connection_warmup`.
**Impact**:
- Proxy accumulates unbounded pending `PDConnectionPool.connect()` and `wait_for_conn()` tasks.
- Duplicate warmup requests against the same finite 10P/10D topology are not deduplicated or rejected.
- Proxy file descriptor count and memory usage grow rapidly.
- The Proxy container can be OOM-killed, causing denial of service for the DistServe Proxy.
**Potential mitigations** include bounding `waiting_conn`, coalescing duplicate connection requests per P-D pair, adding request-level timeouts for `conn_event.wait()`, setting deterministic total HTTP timeouts, tracking and cancelling connection tasks, and rate-limiting or authenticating warmup/control-plane endpoints.
### Environment
```Shell
Target version: LMDeploy v0.17.0
Model: Qwen2.5-0.5B-Instruct
Python: 3.12.3 (main, Jul 15 2026, 23:46:41) [GCC 13.3.0]
CUDA available: True
GPU 0: NVIDIA A100 80GB PCIe
GPU 0 Compute Capability: 8.0
CUDA_HOME: /usr/local/cuda
NVCC: Cuda compilation tools, release 13.0, V13.0.88
CUDA Driver Version: 590.48.01
PyTorch: 2.13.0+cu130
sglang: 0.5.19
sglang-kernel: 0.4.6.post1
flashinfer_python: 0.6.18
flashinfer_cubin: 0.6.18
flashinfer_jit_cache: 0.6.18+cu130
triton: 3.7.1
transformers: 5.12.1
numpy: 2.3.5
aiohttp: 3.14.3
fastapi: 0.141.1
huggingface_hub: 1.30.0
interegular: 0.3.3
modelscope: 1.39.1
orjson: 3.12.0
outlines: 0.1.11
packaging: 26.3
psutil: 7.2.2
pydantic: 2.13.5
python-multipart: 0.0.32
pyzmq: 27.2.0
uvicorn: 0.52.4
uvloop: 0.22.1
xgrammar: 0.2.1
openai: 2.6.1
tiktoken: 0.14.0
torchcodec: 0.15.0+cu130
ulimit soft: 1024
```
### Error traceback
```Shell
```
Contributor guide
Assessment
This issue has not been assessed yet.