Stabilize disaggregated decode burst handoff
Burst decode can sit in disaggregation prealloc/transfer queues while health probes still need an immediate scheduler-alive response, and EAGLE prebuilt metadata must outlive the synthetic prebuilt step until the first real decode consume. The transfer path also needs to send final aux metadata even when there are no new KV pages in the final chunk. This keeps decode metadata slot ownership narrow instead of cloning EAGLE tensors, adds fail-safe cleanup for prebuilt exceptions/finished prebuilt requests, fixes final empty chunk metadata transfer, and records the investigation ledger for future debugging. Constraint: bs=1 historically worked, so decode compute/sampling behavior must not be changed without direct evidence. Rejected: Clone EAGLE metadata tensors on transfer commit | avoids lifetime issues but adds hot-path CPU/GPU memory traffic. Rejected: Treat prefill AbortReq logs as root cause | current evidence shows they can be downstream of decode/router aborts. Confidence: medium Scope-risk: moderate Directive: Do not move EAGLE metadata slot release earlier than first real decode result processing without proving the H2D/spec_info consume point changed. Tested: g0034 docker py_compile for touched scheduler/disagg/test files; PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/mem_cache/test_req_to_token_pool.py test/registered/unit/managers/test_scheduler_health_check.py -> 24 passed Not-tested: Fresh end-to-end burst traffic after restart; decode node_rank=1 persistent log capture.
This commit is contained in:
laoyao0822
@@ -258,6 +258,112 @@ spec_accept_length=2.0
|
||||
|
||||
已看函数：
|
||||
|
||||
## 5. 2026-06-05 SSH 恢复后的远端证据补充
|
||||
|
||||
### 5.1 最新 cjy decode 进程已不在，且 decode0 日志无 Python fatal
|
||||
|
||||
检查时间：2026-06-05 17:10 CST 左右。
|
||||
|
||||
远端 `g0035/g0036/g0033` 上当前活跃的 `sglang.launch_server` 是其他 workload：
|
||||
|
||||
```text
|
||||
python3 -u -m sglang.launch_server ... --port 30000 ... --disaggregation-mode prefill ...
|
||||
```
|
||||
|
||||
不是本轮 cjy decode 进程。共享日志目录最新 cjy 相关文件仍是：
|
||||
|
||||
```text
|
||||
/mnt/beegfs/cjy/log/decode0_20260605_003305.log
|
||||
/mnt/beegfs/cjy/log/sglang_cp_hicache_20260605_003303.log
|
||||
```
|
||||
|
||||
`decode0_20260605_003305.log` 中：
|
||||
|
||||
- decode server args 明确是 `disaggregation_mode='decode'`；
|
||||
- `enable_nsa_prefill_context_parallel=False`；
|
||||
- `enable_nsa_prefill_cp_shared_kv=False`；
|
||||
- `enable_cp_shared_kv_prefill_bs_gt1=False`；
|
||||
- `enable_hierarchical_cache=False`；
|
||||
- `kv_cache_dtype='fp8_e4m3'`；
|
||||
- `speculative_algorithm='EAGLE'`；
|
||||
- `nsa_decode_backend='flashmla_kv'`。
|
||||
|
||||
所以 decode 端不直接启用 CP shared-KV/bs>1；它仍会消费 prefill transfer 过来的 KV/metadata。
|
||||
|
||||
同一日志没有出现：
|
||||
|
||||
```text
|
||||
Traceback / Scheduler hit / SIGQUIT / Health check failed / FAIL_FAST / Metadata corruption / Decode transfer failed
|
||||
```
|
||||
|
||||
文件最后停在 00:39:54 多条 `POST /v1/chat/completions HTTP/1.1" 200 OK`。目前不能从该文件确认 decode 是 Python 异常退出。
|
||||
|
||||
### 5.2 router stdout/stderr 没有持久化到 cjy 日志
|
||||
|
||||
`g0033` router PID `5292`、`761528` 以及 `g0035` router PID `1136943` 的 stdout/stderr 都指向 `/dev/pts/0`：
|
||||
|
||||
```text
|
||||
/proc/<pid>/fd/1 -> /dev/pts/0
|
||||
/proc/<pid>/fd/2 -> /dev/pts/0
|
||||
```
|
||||
|
||||
`/mnt/beegfs/cjy` 和 `/tmp` 下没有找到对应 router/gateway/smg 持久日志。因此 00:39:54 后 router 侧是否先触发 abort/断连，在当前日志面上没有证据。
|
||||
|
||||
### 5.3 prefill 00:40:04 的 AbortReq 更像后果，不是当前第一故障点
|
||||
|
||||
`sglang_cp_hicache_20260605_003303.log` 在 00:40:04 出现大量：
|
||||
|
||||
```text
|
||||
Prefill bootstrap failed ... KVTransferError(...): Aborted by AbortReq
|
||||
[OUTPUT_ZERO_DEBUG] ... finish_reason={'type': 'abort', ...}
|
||||
```
|
||||
|
||||
这些 abort 覆盖多个 rid、多个 bootstrap room，发生在 decode0 日志停止写入之后约 10 秒。当前证据更像 decode/router/client 侧断开后 prefill transfer 被 AbortReq 取消；不能把它作为 prefill 计算/CP shared-KV 直接崩溃的根因。
|
||||
|
||||
### 5.4 fake warmup 的 EAGLE metadata 是单独问题，不应再当作真实 transfer 乱码证据
|
||||
|
||||
已读代码：
|
||||
|
||||
- `python/sglang/srt/disaggregation/decode.py::_is_fake_transfer()`
|
||||
- `python/sglang/srt/disaggregation/fake/conn.py::FakeKVReceiver.init()`
|
||||
- `python/sglang/srt/disaggregation/utils.py::MetadataBuffers`
|
||||
|
||||
事实：
|
||||
|
||||
- fake warmup 使用 `FAKE_BOOTSTRAP_HOST = "2.2.2.2"`；
|
||||
- decode `_commit_transfer_to_req()` 遇到 fake transfer 时跳过 `bootstrap_room` 校验；
|
||||
- `FakeKVReceiver.init()` 只设置 `has_init=True`，不会写 `MetadataBuffers`；
|
||||
- EAGLE decode warmup 仍会从 metadata buffer 读取 `output_id/topk/hidden`。
|
||||
|
||||
因此 decode warmup 的：
|
||||
|
||||
```text
|
||||
! of the of the of the
|
||||
!!!!!!!??!??!
|
||||
```
|
||||
|
||||
可以由未初始化/默认 metadata 解释。它是 warmup 路径 bug 或诊断噪声，但不能直接证明真实 Mooncake transfer 或 prefill CP 写入损坏。
|
||||
|
||||
### 5.5 当前更值得验证的真实请求风险：decode EAGLE metadata slot 生命周期
|
||||
|
||||
已读代码：
|
||||
|
||||
- `python/sglang/srt/disaggregation/decode.py::_commit_transfer_to_req()`
|
||||
- `python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py::process_prebuilt()`
|
||||
- `python/sglang/srt/managers/scheduler_output_processor_mixin.py::_free_decode_metadata_index_if_held()`
|
||||
|
||||
当前逻辑：
|
||||
|
||||
1. `_commit_transfer_to_req()` 从 `MetadataBuffers.get_buf(idx)` 取 views；
|
||||
2. EAGLE 时把 `output_topk_p/output_topk_index/output_hidden_states` view 直接挂到 `req`；
|
||||
3. `process_prebuilt()` 用 `torch.as_tensor(..., device=cuda)` 和 `torch.stack(...).to(cuda)` 消费这些 CPU/pinned views；
|
||||
4. `get_new_prebuilt_batch()` 的 `finally` 随后调用 `_free_decode_metadata_index_if_held()` 释放 slot。
|
||||
|
||||
待验证风险：
|
||||
|
||||
- 如果 CPU->GPU copy 在当前 stream 上是异步排队，而 metadata slot 在 copy 真正消费前被释放并被后续 transfer 覆盖，bs>1/突发请求会更容易 corrupt EAGLE topk/hidden，表现为 accept len 低和输出异常。
|
||||
- `e9c341afe` 明确把原来的 clone mitigation 改成 slot ownership，commit 自身也标注 `Not-tested: Fresh ETE runtime validation...`。这是当前优先验证假设，但还没有日志直接证明。
|
||||
|
||||
- `is_current_only_extend_batch()`
|
||||
- `should_reuse_current_extend_kv()`
|
||||
- `current_extend_kv_rows_for_reuse()`
|
||||
@@ -1496,3 +1602,182 @@ Verification:
|
||||
Remaining runtime validation:
|
||||
|
||||
- Need user restart/run real traffic to confirm output text and EAGLE accept length recover. This fix addresses a concrete coordinate corruption in the bs>1-era path that also affected bs=1; it does not prove all true bs>1 cases are complete.
|
||||
|
||||
## 2026-06-05 multi-request hang after topk-offset fix
|
||||
|
||||
User reports: sending several requests at once can hang. Treat this as a new failure mode after the output-quality fix; do not assume it is the same topk coordinate bug unless logs prove it.
|
||||
|
||||
Facts inherited from the previous inspection:
|
||||
|
||||
- Latest observed prefill log still showed `#new-seq: 1` only; the bs>1 scheduler gate was enabled but the visible trace had not actually formed a multi-request prefill batch yet.
|
||||
- Prefill can remain with `#inflight-req: 1` while repeating 64-token health/warmup batches, which points to a disaggregated KV transfer lifecycle stuck after prefill compute, not necessarily active GPU compute.
|
||||
- Decode configuration in the inspected run had `tp_size=16`, `dp_size=16`, `attn_cp_size=1`, `enable_dp_attention=True`, so `attn_tp_cpu_group` is effectively size 1. Therefore decode-side `poll_and_all_reduce(..., attn_tp_cpu_group)` is not the direct collective hang cause for that configuration.
|
||||
- Prefill showed Mooncake bootstrap failures/timeouts for some rids and decode showed detokenizer watchdog pending rids. This narrows the next audit to prefill `send_kv_chunk` / inflight polling and decode prealloc/transfer metadata readiness.
|
||||
|
||||
High-signal suspects to validate next:
|
||||
|
||||
1. `PrefillBootstrapQueue.send_kv_chunk()` can return without sending when `page_indices` is empty. If the caller has already appended the req to `disagg_prefill_inflight_queue`, this creates a request that can never complete transfer polling.
|
||||
2. `DecodeTransferQueue._commit_transfer_to_req()` treats `metadata.bootstrap_room == 0` after transfer success as "not ready yet". If metadata never arrives, decode can retain a transferred request indefinitely instead of failing with an explicit corruption/ordering error.
|
||||
3. EAGLE metadata buffer slots are intentionally held until `process_prebuilt()` consumes them; if a multi-request path stores references beyond that point, slot reuse can still corrupt later decode state. This is lower priority until logs show transferred requests reaching prebuilt decode.
|
||||
|
||||
Next rule: add bounded fail-fast/diagnostic evidence at the disagg transfer state-machine boundaries before changing scheduling policy or disabling bs>1.
|
||||
|
||||
### 2026-06-05 remote deploy mismatch found
|
||||
|
||||
Remote active tree on `g0034:/mnt/beegfs/cjy/sglang-dev` differs from local HEAD for `python/sglang/srt/disaggregation/mooncake/conn.py`.
|
||||
|
||||
Diff summary:
|
||||
|
||||
- Local HEAD imports and calls `validate_transfer_page_count_or_raise(...)` before Mooncake KV transfer.
|
||||
- Remote still has the old workaround that silently truncates `kv_chunk.prefill_kv_indices` to `len(chunked_dst_kv_indice)` when destination pages are fewer.
|
||||
|
||||
Implication:
|
||||
|
||||
- The active run can still silently drop CP shared-KV pages during transfer. This invalidates any conclusion drawn from that run about the latest local fail-fast transfer contract.
|
||||
- This is directly relevant to multi-request hangs/corruption because a mismatched page count can make prefill/decode disagree about transferred KV pages while the transfer layer proceeds instead of failing early.
|
||||
|
||||
Next action before further ETE conclusions: sync `mooncake/conn.py` (and any other changed transfer files if hashes differ) to remote and restart before testing. Do not keep debugging latest-run behavior as if the fail-fast transfer validation were active.
|
||||
|
||||
### 2026-06-05 fix: final empty KV chunk must still send aux metadata
|
||||
|
||||
Code audit found a real disaggregated-transfer lifecycle bug independent of the remote sync mismatch:
|
||||
|
||||
- `send_kv_chunk()` wrote final metadata via `disagg_metadata_buffers.set_buf(req)` and then returned early when `page_indices` was empty.
|
||||
- This is valid for non-final chunk sends because partial tail pages should be delayed.
|
||||
- It is invalid for the final chunk: if previous chunk sends already transferred all origin KV pages, the final step may have zero new KV pages but still must call `req.disagg_kv_sender.send(empty_pages, state_indices)` so Mooncake queues the final aux metadata transfer and emits terminal status.
|
||||
- Without that final send, decode can wait for metadata/status while prefill has already put the request into the inflight transfer lifecycle.
|
||||
|
||||
Fix:
|
||||
|
||||
- Non-final empty chunks still skip.
|
||||
- Final empty chunks now log `[CP_SHARED_KV_TRANSFER][final_empty_chunk]` and call sender.send with an empty int32 page array so aux metadata/status can complete.
|
||||
|
||||
Verification:
|
||||
|
||||
- Added `test_prefill_final_empty_kv_chunk_still_sends_aux_metadata`.
|
||||
- Remote `g0034` container passed targeted test and full file:
|
||||
`PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_req_to_token_pool.py` → `7 passed`.
|
||||
- Remote `py_compile` passed for `prefill.py`, `mooncake/conn.py`, and the test file.
|
||||
|
||||
Runtime note:
|
||||
|
||||
- Code has been synced to `g0034:/mnt/beegfs/cjy/sglang-dev`, but the currently running prefill process was started before this sync. It must be restarted before ETE can validate this fix.
|
||||
|
||||
### 2026-06-05 decode death after burst: detokenizer watchdog, not prefill fail-fast
|
||||
|
||||
User reports decode died after the transfer lifecycle fix was synced. Current evidence from the latest inspected run:
|
||||
|
||||
- Prefill process on `g0034` stayed alive.
|
||||
- Decode process on `g0035/g0036` was gone after the failure.
|
||||
- Latest decode log `decode0_20260605_001131.log` showed:
|
||||
- a first long request completed with HTTP 200 and normal decode batches;
|
||||
- then a burst of multiple dispatched chat requests;
|
||||
- then detokenizer health failure: `Server couldn't get a response from detokenizer for last 20 seconds`, followed by `30 requests still pending`;
|
||||
- shutdown hit `AttributeError: 'ReqState' object has no attribute 'response_sent_to_client_ts'`.
|
||||
- Latest prefill log around the same time mainly showed `Prefill bootstrap failed ... Aborted by AbortReq`, which is consistent with decode aborting pending requests during watchdog shutdown. No prior CP shared-KV page-count fail-fast was seen in the inspected grep.
|
||||
|
||||
Implication:
|
||||
|
||||
- The `response_sent_to_client_ts` exception is a shutdown-cleanup bug, but likely not the primary death cause.
|
||||
- The primary failure is currently narrowed to decode-side pending request / detokenizer heartbeat / disaggregated transfer lifecycle under burst load.
|
||||
- Next audit should inspect `TokenizerManager` pending request lifecycle and decode prealloc/transfer queues before changing CP scheduling or transfer contracts again.
|
||||
|
||||
### 2026-06-05 root cause: decode health probe treats disagg queue wait as idle
|
||||
|
||||
Code audit found a concrete decode-side watchdog false-positive:
|
||||
|
||||
- `Scheduler.is_fully_idle(for_health_check=True)` intentionally ignored disaggregation prealloc/transfer queues.
|
||||
- Therefore, when decode had no running batch but many requests waiting in decode prealloc/transfer, `/health` believed the scheduler was idle and injected a `HEALTH_CHECK_*` generation request instead of returning a scheduler-alive signal.
|
||||
- That health request could sit behind real disaggregated KV work and produce no detokenizer output within the 20s health timeout.
|
||||
- The HTTP server then marked the server unhealthy and triggered graceful shutdown, which caused the observed prefill-side `Aborted by AbortReq` bootstrap failures.
|
||||
- RED tests on remote confirm current behavior:
|
||||
- decode prealloc queue still reports idle for health;
|
||||
- busy health check is deferred instead of immediately sending `HealthCheckOutput`.
|
||||
|
||||
Fix direction:
|
||||
|
||||
- Treat disaggregated prealloc/transfer/bootstrap/inflight queues as busy for health checks.
|
||||
- When a health check arrives while busy, send `HealthCheckOutput` immediately from `process_input_requests`; do not defer it until `process_batch_result`, because a disagg queue wait may have no batch result to piggyback on.
|
||||
- Keep the shutdown `ReqState.response_sent_to_client_ts` issue as secondary cleanup; it is not the root cause of decode death.
|
||||
|
||||
Implemented:
|
||||
|
||||
- `process_input_requests()` now sends `HealthCheckOutput` immediately for busy health probes.
|
||||
- `is_fully_idle(for_health_check=True)` now counts disaggregated prefill bootstrap/inflight queues and decode prealloc/transfer queues as busy.
|
||||
- `_finish_all_pending_requests_on_shutdown()` now checks `ReqState.time_stats.response_sent_to_client_time` instead of assuming a nonexistent `response_sent_to_client_ts` field, so a future watchdog shutdown can finish pending requests cleanly rather than crashing during cleanup.
|
||||
|
||||
Verification:
|
||||
|
||||
- Remote RED before fix: health tests failed because decode prealloc queue reported idle and busy health check did not send an immediate signal.
|
||||
- Remote GREEN after fix:
|
||||
`PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_scheduler_health_check.py` → `3 passed`.
|
||||
- Remote `py_compile` passed for `scheduler.py`, `tokenizer_manager.py`, and the new test.
|
||||
- Hashes match local and `g0034:/mnt/beegfs/cjy/sglang-dev` for the synced files.
|
||||
|
||||
## 2026-06-05 decode 再次退出：当前证据
|
||||
|
||||
本轮用户反馈 decode 进程再次死亡后，检查了 g0035/g0036/g0034 最新日志：
|
||||
|
||||
- g0035/g0036 上没有存活的 `sglang.launch_server` / `python -m sglang` decode 进程。
|
||||
- 最新 decode 日志为 `/mnt/beegfs/cjy/log/decode0_20260605_003305.log`，文件停在 00:39:54，最后是大量 `POST /v1/chat/completions HTTP/1.1" 200 OK`，没有 `Traceback`、`Health check failed`、`SIGQUIT`、`RuntimeError`、`NCCL` 或 `CUDA` fatal 行。
|
||||
- `dmesg` 与 `docker events` 在 00:38-00:41 窗口没有显示容器 OOM / kill / Xid 证据。
|
||||
- g0036 没有对应的最新 `decode1_20260605_*` 文件；如果 node_rank=1 的 stdout 没有 tee 到文件，则当前固定日志无法覆盖远端第二节点的首个 fatal。
|
||||
- 最新 prefill 日志 `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260605_003303.log` 在 00:40:04 大量出现 `Prefill bootstrap failed ... KVTransferError(...): Aborted by AbortReq.` 和 `[OUTPUT_ZERO_DEBUG]`。这说明 prefill 侧收到 abort；在没有更早 prefill traceback 的情况下，它更像 decode/router/client 侧中断后的结果，而不是 prefill 首因。
|
||||
- 之前用户贴出的 decode 启动失败 `ImportError: cannot import name 'validate_transfer_page_count_or_raise'` 是代码同步不一致的明确证据；当前 g0035 容器内 `/sgl-workspace/sglang-tai` 的 `disaggregation/utils.py` 与本地 hash 已一致，并包含该符号。
|
||||
|
||||
当前结论：这次不是上一轮已修的 health-check false-positive（至少最新 decode log 没有 health failure 证据）。目前缺失的是 decode node_rank=1 / 子进程首个 fatal 行；不能再基于 prefill `AbortReq` 反推 prefill 是根因。
|
||||
|
||||
下一步需要补强的代码侧证据：
|
||||
|
||||
1. decode/prealloc/transfer 队列 abort 时应打印 request/room/queue/receiver 状态，避免只在 prefill 看到 `Aborted by AbortReq`。
|
||||
2. 多节点 decode 启动必须保证 node_rank=1 也有文件日志，或在 rank0 汇总子节点失败原因。
|
||||
3. 如果再次无 traceback 退出，需要优先找 node_rank=1 stdout / shell 退出码，而不是继续在 prefill HiCache 写入路径重复排查。
|
||||
|
||||
## 2026-06-05 SSH 恢复后的复核
|
||||
|
||||
SSH 恢复后重新检查远端状态：
|
||||
|
||||
- g0034/g0033/g0035/g0036 当前存活的是另一套 `--port 30000 --disaggregation-mode prefill` 进程，不是本轮 cjy 的 `50001` prefill/decode 进程。
|
||||
- cjy 最新日志仍是：
|
||||
- prefill: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260605_003303.log`
|
||||
- decode: `/mnt/beegfs/cjy/log/decode0_20260605_003305.log`
|
||||
- `decode0_20260605_003305.log` 结尾仍停在 00:39:54 的 HTTP 200，没有 fatal traceback。
|
||||
- `sglang_cp_hicache_20260605_003303.log` 00:40:04 的大量 `KVTransferError(...): Aborted by AbortReq` 是 prefill 侧收到 abort 后的表现；该日志本身仍不能证明 prefill 是首因。
|
||||
- `decode0_20260605_001131.log` 的旧失败是 health watchdog + `ReqState.response_sent_to_client_ts` 清理异常，已经由本地 `scheduler.py` / `tokenizer_manager.py` 修复覆盖；不能把这个旧日志继续当作当前版本的新根因。
|
||||
|
||||
新增诊断：
|
||||
|
||||
- `Scheduler.abort_request()` 在 disaggregation 模式下增加 `[DISAGG_ABORT_TRACE]` warning 汇总，打印 abort rid/abort_all、waiting/running、prefill bootstrap/inflight、decode prealloc/transfer 队列长度和 abort_message。
|
||||
- 目的：下一次复现时能在 decode/prefill scheduler 自己的日志中看到 abort 进入点和队列状态，不再只能从 prefill Mooncake 的 `Aborted by AbortReq` 反推。
|
||||
- 该日志只在 disaggregation 模式且存在 disagg 队列/abort_all 时触发，不在普通每 tick 热路径触发。
|
||||
|
||||
## 2026-06-05 EAGLE metadata slot 生命周期复核
|
||||
|
||||
代码审计发现一个比继续追 prefill `AbortReq` 更直接的本地风险点：decode 端 EAGLE 初始 draft metadata 的 pinned CPU buffer slot 被过早释放。
|
||||
|
||||
路径：
|
||||
|
||||
1. `DecodeTransferQueue._commit_transfer_to_req()` 从 `metadata_buffers.get_buf(idx)` 取出 `output_topk_p` / `output_topk_index` / `hidden_states_tensor` 等 CPU buffer view，挂到 `req` 上。
|
||||
2. `SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch()` 调用 `new_batch.process_prebuilt(...)`，把这些 metadata 纳入 prebuilt decode batch。
|
||||
3. 旧逻辑在 `finally` 中立刻 `_free_decode_metadata_index_if_held(req)`。
|
||||
4. 但 `process_prebuilt()` 不是最终消费点；真正的 EAGLE draft state 还要等后续一次真实 decode forward/H2D copy 才算安全消费完成。
|
||||
5. burst 请求下，slot 早释放后可被下一批 transfer 复用并覆盖，导致 decode 初始 draft state 错乱，表现可对应 accept len 降低、输出乱码或后续请求卡死。
|
||||
|
||||
修复策略：
|
||||
|
||||
- 不 clone 大 tensor；保持原来的零拷贝 CPU metadata view 模型。
|
||||
- `get_new_prebuilt_batch()` 成功路径不再释放 metadata slot。
|
||||
- `process_prebuilt()` 抛异常时立即释放，避免错误路径泄漏。
|
||||
- synthetic prebuilt 步骤内已经结束的请求在 `process_batch_result_prebuilt()` 释放。
|
||||
- 普通请求在第一轮真实 `process_batch_result_decode()` 中等待 `result.copy_done.synchronize()` 后释放，保证初始 draft metadata 至少存活到第一轮 decode consume 完成。
|
||||
|
||||
风险/边界：
|
||||
|
||||
- metadata slot 会比旧逻辑多持有一个 decode iteration；这是正确生命周期成本，不是 clone 成本。需要 allocator 容量覆盖 burst prebuilt 请求数。
|
||||
- 当前修复只证明 slot 生命周期不再早释放；如果下一次仍出现 burst hang，需要结合新增 `[DISAGG_ABORT_TRACE]` 与 `SGLANG_EAGLE_ACCEPT_DEBUG=1` 对比 `metadata_set -> metadata_get -> prebuilt_state`。
|
||||
|
||||
验证：
|
||||
|
||||
- 本地 `py_compile` 通过：`decode.py`、`scheduler_output_processor_mixin.py`、`test_decode_queue_compaction.py`。
|
||||
- 远端 g0034 容器 targeted tests 通过：
|
||||
`PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/managers/test_scheduler_health_check.py` → `17 passed`。
|
||||
|
||||
@@ -1482,9 +1482,10 @@ class SchedulerDisaggregationDecodeMixin:
|
||||
new_batch.prepare_for_prebuilt()
|
||||
try:
|
||||
new_batch.process_prebuilt(self.server_args, self.future_map)
|
||||
finally:
|
||||
except BaseException:
|
||||
for req in can_run_list:
|
||||
self._free_decode_metadata_index_if_held(req)
|
||||
raise
|
||||
|
||||
return new_batch
|
||||
|
||||
|
||||
@@ -946,11 +946,18 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
page_size,
|
||||
)
|
||||
|
||||
if len(page_indices) == 0:
|
||||
if len(page_indices) == 0 and not last_chunk:
|
||||
logger.info(
|
||||
f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"
|
||||
f"Skip sending non-final kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"
|
||||
)
|
||||
return
|
||||
if len(page_indices) == 0:
|
||||
logger.warning(
|
||||
"[CP_SHARED_KV_TRANSFER][final_empty_chunk] "
|
||||
"sending final aux metadata without new KV pages: "
|
||||
f"{req.rid=} {req.bootstrap_room=} {start_idx=} {end_idx=} "
|
||||
f"origin_len={len(req.origin_input_ids)} fill_len={len(req.fill_ids)}"
|
||||
)
|
||||
prefill_queue = getattr(self, "disagg_prefill_bootstrap_queue", None)
|
||||
has_draft_pool = (
|
||||
getattr(prefill_queue, "draft_token_to_kv_pool", None) is not None
|
||||
|
||||
@@ -1567,8 +1567,15 @@ class Scheduler(
|
||||
if is_health_check_generate_req(recv_req) and not self.is_fully_idle(
|
||||
for_health_check=True
|
||||
):
|
||||
self.return_health_check_ipcs.append(
|
||||
getattr(recv_req, "http_worker_ipc", None)
|
||||
# Send the liveness signal immediately. Deferring it until
|
||||
# process_batch_result can false-fail disaggregated decode when
|
||||
# requests are waiting in prealloc/transfer queues and no decode
|
||||
# batch is currently producing output.
|
||||
self.send_to_tokenizer.send_output(
|
||||
HealthCheckOutput(
|
||||
http_worker_ipc=getattr(recv_req, "http_worker_ipc", None)
|
||||
),
|
||||
recv_req,
|
||||
)
|
||||
continue
|
||||
|
||||
@@ -2902,11 +2909,11 @@ class Scheduler(
|
||||
return ClearHiCacheReqOutput(success=if_success)
|
||||
|
||||
def is_fully_idle(self, for_health_check=False) -> bool:
|
||||
# Health check piggybacks on running requests in process_output.
|
||||
# Only running_batch + waiting_queue guarantee active GPU processing;
|
||||
# disagg queues (bootstrap/prealloc/transfer) may have items without
|
||||
# any request actually running on GPU — e.g. stuck handshake, full
|
||||
# KV cache, or stalled transfer — so they can't carry health info.
|
||||
# Health checks need a scheduler-liveness answer. Disaggregated queues
|
||||
# may wait for prefill/transfer without a local decode batch result, so
|
||||
# they must still count as busy; otherwise /health injects a generation
|
||||
# request that can sit behind the same queue and trigger a false
|
||||
# detokenizer-hang shutdown.
|
||||
# Batch running status
|
||||
idle = (
|
||||
self.running_batch.is_empty()
|
||||
@@ -2921,17 +2928,18 @@ class Scheduler(
|
||||
# Waiting queues: waiting + bootstrapping + preallocation + kv transfer (decode)
|
||||
idle &= len(self.waiting_queue) == 0
|
||||
|
||||
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
||||
idle &= len(self.disagg_prefill_inflight_queue) == 0
|
||||
idle &= len(self.disagg_prefill_bootstrap_queue.queue) == 0
|
||||
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE:
|
||||
idle &= len(self.disagg_decode_prealloc_queue.queue) == 0
|
||||
idle &= len(self.disagg_decode_transfer_queue.queue) == 0
|
||||
|
||||
if not for_health_check:
|
||||
# Grammar queue and prefill inflight queue may not produce batch
|
||||
# results instantly, but they still indicate the server is not idle.
|
||||
idle &= len(self.grammar_manager.grammar_queue) == 0
|
||||
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
||||
idle &= len(self.disagg_prefill_inflight_queue) == 0
|
||||
idle &= len(self.disagg_prefill_bootstrap_queue.queue) == 0
|
||||
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE:
|
||||
idle &= len(self.disagg_decode_prealloc_queue.queue) == 0
|
||||
idle &= len(self.disagg_decode_transfer_queue.queue) == 0
|
||||
|
||||
# HiCache: in-flight async ops (GPU↔Host↔L3) must drain before
|
||||
# destructive operations like attach/detach/flush_cache.
|
||||
@@ -3189,6 +3197,60 @@ class Scheduler(
|
||||
return RpcReqOutput(success, "" if not exec else str(exec))
|
||||
|
||||
def abort_request(self, recv_req: AbortReq):
|
||||
if self.disaggregation_mode != DisaggregationMode.NULL:
|
||||
running_reqs = (
|
||||
len(self.running_batch.reqs)
|
||||
if self.running_batch is not None and self.running_batch.reqs is not None
|
||||
else 0
|
||||
)
|
||||
prefill_bootstrap_reqs = (
|
||||
len(self.disagg_prefill_bootstrap_queue.queue)
|
||||
if self.disaggregation_mode == DisaggregationMode.PREFILL
|
||||
and getattr(self, "disagg_prefill_bootstrap_queue", None) is not None
|
||||
else 0
|
||||
)
|
||||
prefill_inflight_reqs = (
|
||||
len(self.disagg_prefill_inflight_queue)
|
||||
if self.disaggregation_mode == DisaggregationMode.PREFILL
|
||||
and getattr(self, "disagg_prefill_inflight_queue", None) is not None
|
||||
else 0
|
||||
)
|
||||
decode_prealloc_reqs = (
|
||||
len(self.disagg_decode_prealloc_queue.queue)
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE
|
||||
and getattr(self, "disagg_decode_prealloc_queue", None) is not None
|
||||
else 0
|
||||
)
|
||||
decode_transfer_reqs = (
|
||||
len(self.disagg_decode_transfer_queue.queue)
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE
|
||||
and getattr(self, "disagg_decode_transfer_queue", None) is not None
|
||||
else 0
|
||||
)
|
||||
if (
|
||||
recv_req.abort_all
|
||||
or prefill_bootstrap_reqs
|
||||
or prefill_inflight_reqs
|
||||
or decode_prealloc_reqs
|
||||
or decode_transfer_reqs
|
||||
):
|
||||
logger.warning(
|
||||
"[DISAGG_ABORT_TRACE] scheduler abort_request rid=%s abort_all=%s "
|
||||
"mode=%s waiting=%s running=%s prefill_bootstrap=%s "
|
||||
"prefill_inflight=%s decode_prealloc=%s decode_transfer=%s "
|
||||
"abort_message=%s",
|
||||
recv_req.rid,
|
||||
recv_req.abort_all,
|
||||
self.disaggregation_mode,
|
||||
len(self.waiting_queue),
|
||||
running_reqs,
|
||||
prefill_bootstrap_reqs,
|
||||
prefill_inflight_reqs,
|
||||
decode_prealloc_reqs,
|
||||
decode_transfer_reqs,
|
||||
recv_req.abort_message,
|
||||
)
|
||||
|
||||
# todo hisparse, release resources for abort requests in hisparse coordinator
|
||||
# Delete requests in the waiting queue
|
||||
to_del = []
|
||||
|
||||
@@ -64,9 +64,11 @@ class SchedulerOutputProcessorMixin:
|
||||
allocator.free(idx)
|
||||
req.metadata_buffer_index = -1
|
||||
# The EAGLE handoff tensors are views into the reusable metadata slot.
|
||||
# process_prebuilt consumes them into the batch before this helper is
|
||||
# called; keeping the views afterwards makes the request observe a
|
||||
# future transfer that reuses the same slot.
|
||||
# They are safe to release only after either the request finishes in the
|
||||
# synthetic prebuilt step or the first real decode forward consumes the
|
||||
# spec_info built from those views. Releasing immediately after
|
||||
# process_prebuilt lets a burst transfer overwrite pinned CPU metadata
|
||||
# before pending H2D copies have been consumed.
|
||||
req.output_topk_p = None
|
||||
req.output_topk_index = None
|
||||
req.hidden_states_tensor = None
|
||||
@@ -184,6 +186,7 @@ class SchedulerOutputProcessorMixin:
|
||||
if req.finished():
|
||||
req.time_stats.set_quick_finish_time()
|
||||
release_kv_cache(req, self.tree_cache)
|
||||
self._free_decode_metadata_index_if_held(req)
|
||||
|
||||
# Note: Logprobs should be handled on the prefill engine.
|
||||
self.stream_output(batch.reqs, batch.return_logprob)
|
||||
@@ -459,6 +462,13 @@ class SchedulerOutputProcessorMixin:
|
||||
if result.copy_done is not None:
|
||||
result.copy_done.synchronize()
|
||||
|
||||
# Requests entering decode through disaggregated prebuilt EAGLE hold
|
||||
# reusable metadata slots until the first real decode forward consumes
|
||||
# the initial draft state. Free here, not in get_new_prebuilt_batch(),
|
||||
# so burst transfers cannot overwrite the source views too early.
|
||||
for req in batch.reqs:
|
||||
self._free_decode_metadata_index_if_held(req)
|
||||
|
||||
logits_output, next_token_ids, can_run_cuda_graph = (
|
||||
result.logits_output,
|
||||
result.next_token_ids,
|
||||
|
||||
@@ -2187,7 +2187,14 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
||||
continue
|
||||
|
||||
is_stream = getattr(state.obj, "stream", False)
|
||||
already_sent_data = state.response_sent_to_client_ts > 0
|
||||
time_stats = getattr(state, "time_stats", None)
|
||||
already_sent_data = bool(
|
||||
getattr(state, "response_sent_to_client_ts", 0) > 0
|
||||
or (
|
||||
time_stats is not None
|
||||
and getattr(time_stats, "response_sent_to_client_time", 0) > 0
|
||||
)
|
||||
)
|
||||
|
||||
if is_stream and already_sent_data:
|
||||
# Streaming request that already sent partial data to client.
|
||||
|
||||
@@ -14,7 +14,10 @@ from sglang.srt.disaggregation.decode import (
|
||||
DecodeTransferQueue,
|
||||
SchedulerDisaggregationDecodeMixin,
|
||||
)
|
||||
from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator
|
||||
from sglang.srt.disaggregation.utils import (
|
||||
DisaggregationMode,
|
||||
ReqToMetadataIdxAllocator,
|
||||
)
|
||||
from sglang.srt.managers.scheduler_output_processor_mixin import (
|
||||
SchedulerOutputProcessorMixin,
|
||||
)
|
||||
@@ -41,6 +44,13 @@ class FakeReq:
|
||||
self.cached_tokens = 0
|
||||
self.init_next_round_calls = []
|
||||
self.sampling_params = SimpleNamespace(max_new_tokens=0)
|
||||
self.return_hidden_states = False
|
||||
self.grammar = None
|
||||
self.token_ids_logprob = None
|
||||
self.top_logprobs_num = 0
|
||||
self.multimodal_inputs = None
|
||||
self.mamba_ping_pong_track_buffer = None
|
||||
self.to_finish = None
|
||||
|
||||
class FakeTimeStats:
|
||||
def __init__(self):
|
||||
@@ -64,6 +74,12 @@ class FakeReq:
|
||||
def set_quick_finish_time(self):
|
||||
return None
|
||||
|
||||
def set_last_decode_finish_time(self):
|
||||
return None
|
||||
|
||||
def set_completion_time(self):
|
||||
return None
|
||||
|
||||
self.time_stats = FakeTimeStats()
|
||||
|
||||
def add_latency(self, stage):
|
||||
@@ -75,7 +91,7 @@ class FakeReq:
|
||||
def init_next_round_input(self, tree_cache):
|
||||
self.init_next_round_calls.append(tree_cache)
|
||||
|
||||
def check_finished(self):
|
||||
def check_finished(self, *args):
|
||||
return None
|
||||
|
||||
def finished(self):
|
||||
@@ -146,6 +162,18 @@ class FakeAllocator(ReqToMetadataIdxAllocator):
|
||||
self.freed.append(free_index)
|
||||
|
||||
|
||||
class FakeTokenToKVAllocator:
|
||||
def __init__(self):
|
||||
self.begin_calls = 0
|
||||
self.end_calls = 0
|
||||
|
||||
def free_group_begin(self):
|
||||
self.begin_calls += 1
|
||||
|
||||
def free_group_end(self):
|
||||
self.end_calls += 1
|
||||
|
||||
|
||||
class TestDecodeQueueCompaction(CustomTestCase):
|
||||
def test_decode_transfer_queue_compacts_in_one_pass(self):
|
||||
streamed = []
|
||||
@@ -692,12 +720,127 @@ class TestDecodeQueueCompaction(CustomTestCase):
|
||||
captured["batch"].processed,
|
||||
[(scheduler.server_args, scheduler.future_map)],
|
||||
)
|
||||
self.assertEqual(allocator.freed, [20, 21, 22])
|
||||
# EAGLE metadata slots are intentionally held past process_prebuilt().
|
||||
# The initial draft state is consumed by the first real decode forward,
|
||||
# so releasing here would let burst transfers overwrite the pinned CPU
|
||||
# source views before the GPU copy/consume is complete.
|
||||
self.assertEqual(allocator.freed, [])
|
||||
for req in captured["reqs"]:
|
||||
self.assertEqual(req.metadata_buffer_index, -1)
|
||||
self.assertIsNone(getattr(req, "output_topk_p", None))
|
||||
self.assertIsNone(getattr(req, "output_topk_index", None))
|
||||
self.assertIsNone(getattr(req, "hidden_states_tensor", None))
|
||||
self.assertGreaterEqual(req.metadata_buffer_index, 20)
|
||||
|
||||
def test_get_new_prebuilt_batch_frees_metadata_on_prebuilt_error(self):
|
||||
allocator = FakeAllocator()
|
||||
scheduler = cast(Any, SimpleNamespace())
|
||||
scheduler.grammar_manager = SimpleNamespace(
|
||||
has_waiting_grammars=lambda: False,
|
||||
get_ready_grammar_requests=lambda: [],
|
||||
)
|
||||
scheduler._add_request_to_queue = lambda req: None
|
||||
req0 = FakeReq("req-0", 0)
|
||||
scheduler.waiting_queue = [req0]
|
||||
scheduler.running_batch = SimpleNamespace(batch_size=lambda: 0)
|
||||
scheduler.req_to_token_pool = SimpleNamespace(size=8)
|
||||
scheduler.max_running_requests = 8
|
||||
scheduler.tree_cache = object()
|
||||
scheduler.token_to_kv_pool_allocator = object()
|
||||
scheduler.model_config = object()
|
||||
scheduler.enable_overlap = False
|
||||
scheduler.spec_algorithm = object()
|
||||
scheduler.server_args = object()
|
||||
scheduler.future_map = object()
|
||||
scheduler.req_to_metadata_buffer_idx_allocator = allocator
|
||||
scheduler._free_decode_metadata_index_if_held = (
|
||||
SchedulerOutputProcessorMixin._free_decode_metadata_index_if_held.__get__(
|
||||
scheduler
|
||||
)
|
||||
)
|
||||
req0.metadata_buffer_index = 42
|
||||
|
||||
class FailingBatch(FakeBatch):
|
||||
def process_prebuilt(self, server_args, future_map):
|
||||
raise RuntimeError("prebuilt failed")
|
||||
|
||||
with patch(
|
||||
"sglang.srt.disaggregation.decode.ScheduleBatch.init_new",
|
||||
lambda reqs, *args, **kwargs: FailingBatch(),
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, "prebuilt failed"):
|
||||
SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch(scheduler)
|
||||
|
||||
self.assertEqual(allocator.freed, [42])
|
||||
self.assertEqual(req0.metadata_buffer_index, -1)
|
||||
|
||||
def test_process_batch_result_prebuilt_frees_finished_metadata(self):
|
||||
allocator = FakeAllocator()
|
||||
scheduler = SchedulerOutputProcessorMixin.__new__(SchedulerOutputProcessorMixin)
|
||||
scheduler.disaggregation_mode = DisaggregationMode.DECODE
|
||||
scheduler.req_to_metadata_buffer_idx_allocator = allocator
|
||||
scheduler.tree_cache = object()
|
||||
scheduler.stream_output = lambda reqs, return_logprob: None
|
||||
|
||||
req = FakeReq("finished", 0)
|
||||
req.metadata_buffer_index = 43
|
||||
req.output_topk_p = torch.ones((1,), dtype=torch.float32)
|
||||
req.output_topk_index = torch.ones((1,), dtype=torch.int64)
|
||||
req.hidden_states_tensor = torch.ones((4,), dtype=torch.float32)
|
||||
req.finished_reason = FINISH_ABORT("done")
|
||||
batch = SimpleNamespace(reqs=[req], return_logprob=False)
|
||||
|
||||
with patch(
|
||||
"sglang.srt.managers.scheduler_output_processor_mixin.release_kv_cache",
|
||||
lambda *args, **kwargs: None,
|
||||
):
|
||||
scheduler.process_batch_result_prebuilt(batch)
|
||||
|
||||
self.assertEqual(allocator.freed, [43])
|
||||
self.assertEqual(req.metadata_buffer_index, -1)
|
||||
|
||||
def test_process_batch_result_decode_releases_prebuilt_metadata_after_consume(self):
|
||||
allocator = FakeAllocator()
|
||||
token_allocator = FakeTokenToKVAllocator()
|
||||
scheduler = SchedulerOutputProcessorMixin.__new__(SchedulerOutputProcessorMixin)
|
||||
scheduler.req_to_metadata_buffer_idx_allocator = allocator
|
||||
scheduler.server_args = SimpleNamespace(
|
||||
disaggregation_decode_enable_offload_kvcache=False
|
||||
)
|
||||
scheduler.enable_hisparse = False
|
||||
scheduler.enable_overlap = False
|
||||
scheduler.enable_metrics = False
|
||||
scheduler.token_to_kv_pool_allocator = token_allocator
|
||||
scheduler.tree_cache = object()
|
||||
scheduler.forward_ct_decode = 0
|
||||
scheduler.num_generated_tokens = 0
|
||||
scheduler.stream_output = lambda reqs, return_logprob: None
|
||||
scheduler.report_decode_stats = lambda *args, **kwargs: None
|
||||
scheduler.update_spec_metrics = lambda *args, **kwargs: None
|
||||
scheduler._maybe_log_eagle_accept_debug = lambda *args, **kwargs: None
|
||||
|
||||
req = FakeReq("decode", 0)
|
||||
req.metadata_buffer_index = 44
|
||||
req.output_topk_p = torch.ones((1,), dtype=torch.float32)
|
||||
req.output_topk_index = torch.ones((1,), dtype=torch.int64)
|
||||
req.hidden_states_tensor = torch.ones((4,), dtype=torch.float32)
|
||||
batch = SimpleNamespace(
|
||||
reqs=[req],
|
||||
return_logprob=False,
|
||||
spec_algorithm=SimpleNamespace(is_none=lambda: True),
|
||||
is_spec_v2=False,
|
||||
)
|
||||
result = SimpleNamespace(
|
||||
copy_done=None,
|
||||
logits_output=SimpleNamespace(hidden_states=None, customized_info=None),
|
||||
next_token_ids=torch.tensor([5], dtype=torch.int64),
|
||||
can_run_cuda_graph=True,
|
||||
num_accepted_tokens=1,
|
||||
)
|
||||
|
||||
scheduler.process_batch_result_decode(batch, result)
|
||||
|
||||
self.assertEqual(allocator.freed, [44])
|
||||
self.assertEqual(req.metadata_buffer_index, -1)
|
||||
self.assertEqual(req.output_ids, [5])
|
||||
self.assertEqual(token_allocator.begin_calls, 1)
|
||||
self.assertEqual(token_allocator.end_calls, 1)
|
||||
|
||||
def test_get_new_prebuilt_batch_keeps_waiting_queue_when_no_capacity(self):
|
||||
scheduler = cast(Any, SimpleNamespace())
|
||||
|
||||
88
test/registered/unit/managers/test_scheduler_health_check.py
Normal file
88
test/registered/unit/managers/test_scheduler_health_check.py
Normal file
@@ -0,0 +1,88 @@
|
||||
import unittest
|
||||
from collections import deque
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.managers.io_struct import HealthCheckOutput
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
from sglang.srt.managers.tokenizer_manager import TokenizerManager
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=2, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class TestSchedulerHealthCheck(CustomTestCase):
|
||||
def _idle_scheduler(self, *, disaggregation_mode=DisaggregationMode.NULL):
|
||||
return SimpleNamespace(
|
||||
running_batch=SimpleNamespace(is_empty=lambda: True, reqs=[]),
|
||||
chunked_req=None,
|
||||
dllm_manager=SimpleNamespace(any_staging_reqs=lambda: False),
|
||||
last_batch=None,
|
||||
cur_batch=None,
|
||||
enable_overlap=False,
|
||||
result_queue=deque(),
|
||||
pp_size=1,
|
||||
waiting_queue=[],
|
||||
grammar_manager=SimpleNamespace(grammar_queue=[]),
|
||||
disaggregation_mode=disaggregation_mode,
|
||||
enable_hierarchical_cache=False,
|
||||
)
|
||||
|
||||
def test_decode_prealloc_queue_is_busy_for_health_check(self):
|
||||
scheduler = self._idle_scheduler(disaggregation_mode=DisaggregationMode.DECODE)
|
||||
scheduler.disagg_decode_prealloc_queue = SimpleNamespace(queue=[object()])
|
||||
scheduler.disagg_decode_transfer_queue = SimpleNamespace(queue=[])
|
||||
scheduler.server_args = SimpleNamespace(
|
||||
disaggregation_decode_enable_offload_kvcache=False
|
||||
)
|
||||
|
||||
self.assertFalse(Scheduler.is_fully_idle(scheduler, for_health_check=True))
|
||||
|
||||
def test_busy_health_check_sends_immediate_signal(self):
|
||||
sent = []
|
||||
|
||||
scheduler = SimpleNamespace(
|
||||
session_controller=SimpleNamespace(maybe_reap=MagicMock()),
|
||||
is_fully_idle=lambda for_health_check=False: False,
|
||||
send_to_tokenizer=SimpleNamespace(
|
||||
send_output=lambda output, recv_obj=None: sent.append((output, recv_obj))
|
||||
),
|
||||
return_health_check_ipcs=deque(),
|
||||
_request_dispatcher=MagicMock(),
|
||||
)
|
||||
req = SimpleNamespace(rid="HEALTH_CHECK_unit", http_worker_ipc="ipc://unit")
|
||||
|
||||
Scheduler.process_input_requests(scheduler, [req])
|
||||
|
||||
scheduler._request_dispatcher.assert_not_called()
|
||||
self.assertEqual(len(sent), 1)
|
||||
self.assertIsInstance(sent[0][0], HealthCheckOutput)
|
||||
self.assertEqual(sent[0][0].http_worker_ipc, "ipc://unit")
|
||||
self.assertIs(sent[0][1], req)
|
||||
|
||||
def test_shutdown_pending_request_without_response_ts_aborts_cleanly(self):
|
||||
event = SimpleNamespace(set=MagicMock())
|
||||
state = SimpleNamespace(
|
||||
finished=False,
|
||||
obj=SimpleNamespace(stream=True),
|
||||
text="",
|
||||
output_ids=[],
|
||||
out_list=[],
|
||||
event=event,
|
||||
)
|
||||
tokenizer_manager = SimpleNamespace(rid_to_state={"rid": state})
|
||||
|
||||
TokenizerManager._finish_all_pending_requests_on_shutdown(tokenizer_manager)
|
||||
|
||||
self.assertTrue(state.finished)
|
||||
self.assertEqual(len(state.out_list), 1)
|
||||
self.assertEqual(
|
||||
state.out_list[0]["meta_info"]["finish_reason"]["type"], "abort"
|
||||
)
|
||||
event.set.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -13,6 +13,7 @@ from sglang.srt.disaggregation.decode import (
|
||||
_kv_locs_to_page_indices_cpu,
|
||||
)
|
||||
from sglang.srt.disaggregation.prefill import (
|
||||
SchedulerDisaggregationPrefillMixin,
|
||||
_kv_locs_to_page_indices_cpu as _prefill_kv_locs_to_page_indices_cpu,
|
||||
)
|
||||
from sglang.srt.managers.schedule_batch import Req
|
||||
@@ -200,6 +201,61 @@ class TestDecodePreallocQueue(CustomTestCase):
|
||||
self.assertEqual(page_indices.dtype.name, "int32")
|
||||
self.assertEqual(page_indices.tolist(), [32, 33, 34])
|
||||
|
||||
|
||||
def test_prefill_final_empty_kv_chunk_still_sends_aux_metadata(self):
|
||||
class FakeAllocator:
|
||||
page_size = 4
|
||||
|
||||
def get_kvcache(self):
|
||||
return object()
|
||||
|
||||
class FakeMetadataBuffers:
|
||||
def __init__(self):
|
||||
self.requests = []
|
||||
|
||||
def set_buf(self, req):
|
||||
self.requests.append(req)
|
||||
|
||||
class FakeSender:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def send(self, page_indices, state_indices):
|
||||
self.calls.append((page_indices, state_indices))
|
||||
|
||||
fake_scheduler = SimpleNamespace(
|
||||
token_to_kv_pool_allocator=FakeAllocator(),
|
||||
req_to_token_pool=SimpleNamespace(
|
||||
req_to_token=torch.arange(100, 112, dtype=torch.int64).reshape(1, -1)
|
||||
),
|
||||
disagg_metadata_buffers=FakeMetadataBuffers(),
|
||||
disagg_prefill_bootstrap_queue=SimpleNamespace(draft_token_to_kv_pool=None),
|
||||
)
|
||||
sender = FakeSender()
|
||||
req = SimpleNamespace(
|
||||
rid="r-final-empty",
|
||||
bootstrap_room=123,
|
||||
start_send_idx=8,
|
||||
fill_ids=list(range(8)),
|
||||
origin_input_ids=list(range(8)),
|
||||
req_pool_idx=0,
|
||||
disagg_kv_sender=sender,
|
||||
prefix_indices=[],
|
||||
host_hit_length=0,
|
||||
)
|
||||
|
||||
SchedulerDisaggregationPrefillMixin.send_kv_chunk(
|
||||
fake_scheduler, req, last_chunk=True
|
||||
)
|
||||
|
||||
self.assertEqual(req.start_send_idx, 8)
|
||||
self.assertEqual(fake_scheduler.disagg_metadata_buffers.requests, [req])
|
||||
self.assertEqual(len(sender.calls), 1)
|
||||
page_indices, state_indices = sender.calls[0]
|
||||
self.assertEqual(page_indices.dtype.name, "int32")
|
||||
self.assertEqual(page_indices.tolist(), [])
|
||||
self.assertIsNone(state_indices)
|
||||
|
||||
def test_prefill_kv_locs_to_page_indices_cpu_copies_only_page_starts(self):
|
||||
kv_locs = torch.tensor(
|
||||
[
|
||||
|
||||
Reference in New Issue
Block a user