Skip to content

feat(aorta): multi-node disaggregated launch via single cluster.json (AIMVT-173) - #171

Open
speriaswamy-amd wants to merge 8 commits into
mainfrom
surya/aorta-multinode-disaggregated
Open

feat(aorta): multi-node disaggregated launch via single cluster.json (AIMVT-173)#171
speriaswamy-amd wants to merge 8 commits into
mainfrom
surya/aorta-multinode-disaggregated

Conversation

@speriaswamy-amd

@speriaswamy-amd speriaswamy-amd commented May 14, 2026

Copy link
Copy Markdown
Contributor

Summary

Implements AIMVT-173: run the CVS Aorta benchmark across N nodes from a single cluster.json, mirroring the disaggregated launch pattern used by the existing PyTorch xDiT and SGLang multi-node test suites.

The AortaRunner now orchestrates torchrun on every node in parallel (rendezvousing on the head), and consolidates per-node torch_profiler trees into <aorta_path>/combined_traces/node_<rank>/ so the host parser sees a single unified set. Single-node behavior is unchanged: multi_node.master_launch_mode='auto' keeps the legacy experiment_script path for 1-node clusters, and pre-existing yamls without a multi_node: block still validate (Pydantic supplies sensible defaults).

What's in the diff

  • cvs/runners/aorta.py — AortaMultiNodeConfig dataclass; _resolve_launch_mode, _pick_master_port, _build_torchrun_command, _run_single_node, _collect_multi_node_traces + local/remote copy helpers; refactored run()
  • cvs/parsers/schemas.py — AortaMultiNodeConfigFile Pydantic schema + train_script existence check
  • cvs/input/config_file/aorta/aorta_benchmark.yaml — new multi_node: block with inline docs
  • docs/reference/configuration-files/aorta.rst — new "Multi-node disaggregated launch" section + parameter table
  • cvs/runners/unittests/test_aorta_multinode.py — 24 unit tests (launch-mode resolution, port selection, command construction, env merging, trace-tree copy, schema validation, single --override group invariant)
  • cvs/tests/benchmark/test_aorta.py — wire multi_node block through the runner-config fixture

Validation

End-to-end cvs run test_aorta against a real 2-node cluster (g17u19 head + f16u13 worker, 16xMI300X total) — 5/5 pytest cases pass in 148s, traces collected from both nodes, host parser produced metrics for all 16 ranks. Four runtime bugs surfaced and were fixed during this validation:

  1. Container ran as jenkins UID and couldn't open /dev/kfd despite --privileged → now passes user="root" and group_add=["video","render"]
  2. UnboundLocalError: trace_mtime in the freshest-trace selector → initialised upfront
  3. Head-node traces never collected when orchestrator is a separate login host → falls back to SSH/rsync
  4. --override key=val --override key=val … collapsed to last group only (aorta train.py uses argparse(nargs="*")) → packed behind a single --override

Test plan

  • ruff check . --exclude .venv — clean
  • ruff format --check — clean
  • python -m unittest discover -s cvs — 288/288 pass (existing 264 + 24 new)
  • cvs run test_aorta on real 2-node cluster — 5/5 pass
  • Backward compat: yaml without multi_node: block still validates and runs single-node
  • CI on PR

Made with Cursor

speriaswamy-amd and others added 2 commits May 14, 2026 14:01
@speriaswamy-amd @cursoragent
feat(aorta): multi-node disaggregated launch via single cluster.json …
7dacef3 
…(AIMVT-173)

Run the CVS Aorta pipeline across N nodes from one cluster.json by orchestrating
torchrun on every node in parallel, rendezvous-ing on the head, and consolidating
per-node torch_profiler trees into <aorta_path>/combined_traces/node_<rank>/ for
the host parser. Single-node behavior is unchanged: master_launch_mode='auto'
keeps the legacy script path for 1-node clusters.

Notable runtime fixes shaken out by the 2-node validation on g17u19+f16u13:
- launch container as root (+render group) so /dev/kfd is accessible
- pull head-node traces over SSH when orchestrator != head physical host
- pack all training_overrides behind a single --override (aorta argparse uses
  nargs="*" and silently drops earlier groups otherwise)
- initialise trace_mtime before the freshest-trace comparison

Adds AortaMultiNodeConfig + Pydantic schema, refactors AortaRunner.run(),
documents the new block in docs/reference/configuration-files/aorta.rst.

Co-authored-by: Cursor <cursoragent@cursor.com>
@speriaswamy-amd @cursoragent
test(aorta): unit tests for multi-node launch helpers (AIMVT-173)
ea1c7ba 
24 unittest cases covering the new launch-mode resolution, master-port picking,
torchrun command construction, base-env merging, combined_traces helper,
local trace-tree copy, train_script existence check, and the Pydantic
AortaMultiNodeConfigFile schema. Also pins the "single --override group"
invariant in two places to prevent the argparse(nargs="*") regression.

Co-authored-by: Cursor <cursoragent@cursor.com>
@speriaswamy-amd speriaswamy-amd mentioned this pull request May 14, 2026
Closed
4 tasks
speriaswamy-amd and others added 6 commits July 10, 2026 00:15
@speriaswamy-amd
Merge remote-tracking branch 'origin/main' into surya/aorta-multinode…
ccdc140 
…-disaggregated
@speriaswamy-amd
Merge origin/main into surya/aorta-multinode-disaggregated
1a1a330 
Picks up a ruff-format fix (b9ce6c5) on main that landed after this
branch's last sync, which was causing CI's fmt-check to fail on files
this branch doesn't touch (docker_lib.py, jax_training_lib.py).
@speriaswamy-amd @claude
fix(aorta): pick master_port on head node, enforce launch timeout, us…
beeefe0 
…e vpc_ip for rendezvous

Live 2-node testing on Alola surfaced three real bugs in the
disaggregated multi-node torchrun launch path:

- master_port was picked via a local socket bind on the orchestrator,
  not the head node running rendezvous, contradicting the schema's own
  docs. Now queried over SSH on the head node.
- setup()/run() had no timeout on their ThreadPoolExecutor waits, so a
  stalled node (confirmed live via an NCCL collective stall) hung the
  whole run indefinitely. Now bounded by timeout_seconds via
  concurrent.futures.wait(), surfacing RunStatus.TIMEOUT instead of
  hanging.
- master_addr defaulted to the head node's mgmt/SSH address instead of
  its VPC/RDMA-fabric address, matching the existing node_dict vs.
  vpc_ip distinction already used by rccl_perf.py. Now resolved via
  node_vpc_ips with a safe fallback for single-network clusters.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@speriaswamy-amd @claude
fix(aorta): collect traces from surviving nodes on partial multi-node…
3151935 
… failure

run()'s disaggregated branch returned FAILED/TIMEOUT immediately when any
node failed, before ever collecting profiler traces. One flaky/late node
in an otherwise-good multi-hour run therefore discarded every surviving
node's traces, forcing a full rerun or a manual TraceLensParser salvage.

Trace collection and downstream artifact discovery now run unconditionally;
the run is still correctly reported as FAILED/TIMEOUT with the failing
nodes named in error_message, but artifacts (torch_traces, TraceLens/GEMM
analysis) are populated from whatever the surviving nodes produced.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@speriaswamy-amd @claude
fix(aorta): tear down partial setup, use daemon threads, pass SSH key…
f5c3a61 
… for port pick

BaseRunner.execute() only ran teardown() after a fully successful setup(),
leaking containers already launched on nodes that succeeded before a later
node failed or timed out. teardown() now runs whenever setup() was attempted.

setup()/run() replace ThreadPoolExecutor with a daemon-thread-based bounded
execution helper: CPython's atexit hook joins every ThreadPoolExecutor worker
thread regardless of shutdown(wait=False), so a genuinely stuck node could
still hang the whole process at interpreter exit. Daemon threads are exempt.

_pick_free_port_on() now passes -i self.config.pkey when configured, matching
the existing _copy_remote_torch_profilers() pattern - previously any cluster
with a non-default SSH key would fail to pick a master port.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@speriaswamy-amd
fix(aorta): quote SSH port-pick command, close setup-teardown contain…
9bd621e 
…er race

_pick_free_port_on() passed the remote python snippet as separate argv
elements, so ssh's remote shell re-split it on the snippet's semicolons;
build it as one shlex-quoted string, matching _copy_remote_torch_profilers().

setup() now shares a cancel Event with _setup_single_node(): once the
overall deadline passes, a straggling node that finishes launching its
container after being abandoned tears it down itself instead of
registering it into self._containers after teardown()'s one-time snapshot
has already run and cleared it.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@speriaswamy-amd