feat: on-demand profiling endpoints (/start_profile, /stop_profile, /profile_status) - #3
Open
sufubao wants to merge 18 commits into
Open
feat: on-demand profiling endpoints (/start_profile, /stop_profile, /profile_status)#3sufubao wants to merge 18 commits into
sufubao wants to merge 18 commits into
Conversation
Implement ProfileControlReq, StartProfileCmd, and StopProfileCmd dataclasses for on-demand profiling control. ProfileControlReq is sent from HTTP server to router and converted to worker commands via to_worker_cmd() method. All classes are pickle-serializable for shared-memory transport. - ProfileControlReq: httpserver -> router control message with defaults - StartProfileCmd: worker command for profile initiation with full config - StopProfileCmd: worker command for profile termination
Implements WorkerProfilerManager: IDLE→ARMED→RUNNING→FLUSHING→IDLE state machine with on_cmd/on_step_boundary interface. Failures report to ProfileStatusBoard and reset to IDLE for recovery. Includes 8 TDD unit tests covering auto-stop, start_step delay, manual stop, cancel while armed, duplicate-start guard, factory/export error paths, and idle fast-path forward counting.
- Import StartProfileCmd/StopProfileCmd and WorkerProfilerManager in base_backend.py - Instantiate WorkerProfilerManager after shm buffers are set up, before infer threads start - Dispatch StartProfileCmd/StopProfileCmd to profiler_manager.on_cmd() in _read_reqs_buffer_and_init_reqs - Call profiler_manager.on_step_boundary() as first statement in is_prefill() and is_decode() branches of ChunkedPrefillBackend.infer_loop and DPChunkedPrefillBackend.infer_loop; is_pass() branches unchanged
…ing disabled - 501 /start_profile in multinode tensor-parallel mode (nnodes>1, dp=1): locally injected cmds would desync the nodes' lockstep shm cmd buffers and hang NCCL collectives; router also drops such cmds defensively - create ProfileStatusBoard / WorkerProfilerManager only under --enable_profiling (NullProfilerManager no-op otherwise), so disabled servers allocate no /dev/shm segments
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
torch.profilercapture to a running server, gated by a new--enable_profilingflag:POST /start_profile(withnum_stepsauto-stop /start_steparming),POST /stop_profile, andGET /profile_status(per-rank state via a shared-memory status board).ProfileControlReq), router → all worker ranks via theShmObjsIOBuffercommand broadcast (same path asAbortedReqCmd), and a per-rankWorkerProfilerManagerstate machine that starts/stops capture at forward-step boundaries. One gzipped chrome trace per rank ({prefix}-{id}-TP-{t}-DP-{d}.trace.json.gz), viewable in perfetto./start_profileinput is fully validated (output_dir sandbox, prefix whitelist, positive-int steps); multinode-TP mode is gated off (local cmd injection would desync the nodes' lockstep cmd buffers and hang NCCL); servers without the flag allocate no profiler state at all.Design doc:
docs/profile_server/design.md· Usage:docs/profile_server/usage.mdTest Plan
unit_tests/server/core/objs/test_profile_cmd.py,test_profile_status_board.py,unit_tests/server/router/test_worker_profiler_manager.py) — state machine incl. cross-thread deferred stop, atomic-export failure cleanup, shm boardnum_stepsforwards → idle; two valid chrome traces (13.9k events, 3.4k GPU kernels); manual-stop path; all six 400/501 validation paths; flag-off 501 gate; default (flag-off) server serves normally