Skip to content

feat: place pipeline shards by measured bandwidth and latency - #2297

Draft
abendrothj wants to merge 2 commits into
exo-explore:mainfrom
abendrothj:feat/957-complete-placement-objective
Draft

abendrothj wants to merge 2 commits into
exo-explore:mainfrom
abendrothj:feat/957-complete-placement-objective

Conversation

@abendrothj

Copy link
Copy Markdown

Problem

Pipeline placement currently chooses the smallest feasible cycle and divides layers by available RAM. That can put too much work on a low-bandwidth node, ignore a faster larger cycle, or select a ring ordering with slower links.

Changes

  • Profile live memory bandwidth on each node and carry it into placement.
  • Allocate layers in proportion to bandwidth while treating available RAM as a hard capacity.
  • Estimate each candidate's token time from its actual capacity-adjusted layer allocation and the measured RTT of each ring hop.
  • Compare fully profiled cycles across device counts, with leaf status, download progress, and RAM used only as tie-breakers.
  • Preserve the existing smallest-cycle and download-aware behavior until profiling has produced a complete candidate.
  • Split placement filtering, selection, and instance construction into focused helpers.

This replaces #1088 and incorporates the review feedback there.

Stack

Depends on #2254 for measured per-link latency. Until #2254 merges, GitHub will include its commits in this PR's diff because both head branches are in a fork. The diff will narrow to the placement changes automatically after #2254 lands.

Tests

  • Added end-to-end coverage for measured-latency selection, bandwidth taking priority over download progress, comparison across cycle sizes, and capacity-adjusted scoring.
  • uv run basedpyright
  • uv run ruff check
  • uv run ruff format --check on changed Python files
  • uv run pytest --ignore=src/exo/download/tests --ignore=rust/exo_rs/tests/test_python.py — 425 passed, 3 skipped, 190 deselected

nix fmt was unavailable locally because Nix is not installed.

Fixes #957

@abendrothj
abendrothj force-pushed the feat/957-complete-placement-objective branch from 064a96a to b1a12ef Compare September 8, 2026 02:33
@abendrothj
abendrothj marked this pull request as draft September 16, 2026 20:17
@abendrothj

Copy link
Copy Markdown
Author

Moving to draft: this branch still carries #2254's commit, so the diff here double-counts that work until #2254 lands. It also rewrites the placement selection and layer-allocation paths that #2252 (accelerator tiebreaker) and #2253 (manual node_layers) change, and the current revision does not preserve either behaviour. I will rebase onto whatever lands first and mark this ready once the diff is placement-only.

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.

[MEDIUM] Better placement algorithm for pipeline parallelism using memory bandwidth + latency

1 participant