Repository navigation
[REFACT][M9][1/N]: Type PD control-plane protocol and transfer phases - #131
YouHengfei wants to merge 2 commits into
Conversation
|
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. 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. |
Purpose
Refs #32 and #111 (M9 track A).
Types the PD control-plane protocol without changing the wire format or behavior:
pd/protocol.py:CommandKind(coordinator→worker),EventKind(worker→coordinator),PendingPhase, andProfileActionasstrenums. Members compare equal to their string values; messages over the Pipe remain the same plain dicts.Pending.phasebecomes a typedPendingPhaseinstead of a raw string (finding F6 indocs/pd_walkthrough.md, currently in [Doc][M9] PD and KV transfer walkthrough #64).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
tests/test_pd.py(mocked processes/pipes).pytest tests/test_pd.py tests/test_profiling.py(CPU).pytest tests/test_pd.py --run-gpuwith real multi-process NIXL on 2× H100.Test Result
tests/test_pd.py+tests/test_profiling.py(CPU): 38 passed, 15 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.