Skip to content

feat: support manual pipeline layer allocation - #2253

Open
abendrothj wants to merge 2 commits into
exo-explore:mainfrom
abendrothj:feat/1720-manual-layer-allocation
Open

feat: support manual pipeline layer allocation#2253
abendrothj wants to merge 2 commits into
exo-explore:mainfrom
abendrothj:feat/1720-manual-layer-allocation

Conversation

@abendrothj

Copy link
Copy Markdown

Motivation

Memory-proportional automatic allocation cannot express deliberate asymmetric load, such as assigning only a few layers to a laptop while a dedicated larger machine carries the rest.

Changes

  • add optional node_layers to the placement API and PlaceInstance command
  • require the supplied nodes to exactly match a connected pipeline cycle
  • validate positive counts, exact model-layer totals, and per-node memory capacity
  • reject manual allocation for tensor and CFG-parallel modes where layer counts do not map one-to-one to physical nodes
  • generate pipeline shard boundaries from the requested per-node counts
  • preserve automatic proportional placement when node_layers is omitted

Example request field:

{
  "node_layers": {
    "small-node-id": 2,
    "large-node-id": 38
  }
}

Testing

  • uv run pytest src/exo/master/tests/test_placement.py src/exo/master/tests/test_placement_utils.py src/exo/api/tests -m "not slow" (94 passed)
  • repository-wide uv run basedpyright
  • repository-wide uv run ruff check

Fixes #1720


Refiled from #2214, which was closed unintentionally.

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.

Manual layer distribution per device

1 participant