Skip to content

[REFACT][M9][1/N]: Type PD control-plane protocol and transfer phases - #131

Closed
YouHengfei wants to merge 2 commits into
InferMatrixLab:mainfrom
YouHengfei:refactor/m9-pd-protocol
Closed

YouHengfei wants to merge 2 commits into
InferMatrixLab:mainfrom
YouHengfei:refactor/m9-pd-protocol

Conversation

@YouHengfei

@YouHengfei YouHengfei commented Oct 17, 2026 •

Copy link
Copy Markdown
Contributor

Purpose

Refs #32 and #111 (M9 track A).

Types the PD control-plane protocol without changing the wire format or behavior:

  • New pd/protocol.py: CommandKind (coordinator→worker), EventKind (worker→coordinator), PendingPhase, and ProfileAction as str enums. Members compare equal to their string values; messages over the Pipe remain the same plain dicts.
  • Pending.phase becomes a typed PendingPhase instead of a raw string (finding F6 in docs/pd_walkthrough.md, currently in [Doc][M9] PD and KV transfer walkthrough #64).
  • Session-scoped profiling commands (profile_*) are dispatched through an explicit command→action map instead of string prefix manipulation, keeping them separate from transfer-scoped messages as [RFC] PyTorch profiling interface with coordinated PD capture and per-rank artifacts #71 requires.

Organized as two commits: the first adds CPU characterization tests that pin the current protocol behavior (state machine phases, credit accounting, cancel/cancel_safe ordering, command field sets, profile-reply isolation); the second performs the typing. The tests pass unmodified before and after.

Test Plan

  • New CPU coordinator tests in tests/test_pd.py (mocked processes/pipes).
  • pytest tests/test_pd.py tests/test_profiling.py (CPU).
  • Full CPU suite.
  • pytest tests/test_pd.py --run-gpu with real multi-process NIXL on 2× H100.

Test Result

  • tests/test_pd.py + tests/test_profiling.py (CPU): 38 passed, 15 skipped (GPU).
  • Full CPU suite: 564 passed, 170 skipped.
  • GPU (CUDA_VISIBLE_DEVICES=2,1, UCX_TLS=tcp,sm,self,cuda_copy,cuda_ipc): 32 passed, 1 skipped (the 4-GPU multi-worker test). Covers Ouro + Nanbeige handoff, last_exited/shared/CUDA Graph combinations, cancel/reuse at every protocol boundary, prefix-reuse transfer reduction, and EOS after handoff.
  • ruff check / ruff format --check: clean.

@YouHengfei
YouHengfei marked this pull request as draft October 17, 2026 06:03
@YouHengfei

Copy link
Copy Markdown
Contributor Author

This PR is intentionally split into two commits:

First commit adds CPU characterization tests that pin the current protocol behavior (phase transitions, credit accounting, cancel/cancel_safe ordering, command field sets) against the unmodified code.
Second commit performs the typing. The tests pass unchanged on both sides, which is the behavior-preservation evidence.

Scope note: This PR types message kinds and transfer phases only. Payloads remain plain dicts on purpose — the Pipe wire format is unchanged, and converting payloads to dataclasses would touch every send site and bloat the diff. Typed payloads (e.g. reserve/commit/ack records in pd/protocol.py) are a natural follow-up once this lands, after the M3/M4/M5 contract discussions in #111 settle the field ownership.
GPU coverage: The multi-process NIXL suite passes on 2× H100 (32 passed, 1 skipped for the 4-GPU case), so the str-enum messages survive pickle/Pipe round trips on the real path.

@YouHengfei
YouHengfei marked this pull request as ready for review October 17, 2026 14:41
@YouHengfei YouHengfei changed the title [REFACT][M9] Type PD control-plane protocol and transfer phases [REFACT][M9][1/N]: Type PD control-plane protocol and transfer phases Oct 17, 2026
@YouHengfei
YouHengfei marked this pull request as draft October 21, 2026 09:39
@YouHengfei YouHengfei closed this Oct 23, 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