Skip to content

feat(weight-transfer): support router-discovered NIXL autoscaling#3119

Draft
S1ro1 wants to merge 6 commits into
mainfrom
exp/router-autoscaling-poc
Draft

feat(weight-transfer): support router-discovered NIXL autoscaling#3119
S1ro1 wants to merge 6 commits into
mainfrom
exp/router-autoscaling-poc

Conversation

@S1ro1

@S1ro1 S1ro1 commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Summary

PoC support for adding inference replicas to a running NIXL + ModelExpress RL job through the inference router.

  • adds router_url to the shared trainer/orchestrator NIXL configuration
  • moves inference-worker discovery and /init_broadcaster setup to the trainer
  • discovers backend admin URLs through the router's /workers endpoint
  • assigns stable rank ranges to newly discovered inference servers
  • updates ModelExpress participant counts as the inference pool grows
  • keeps orchestrator pause/update/resume traffic on the router endpoint

Validated topology

Validated on the H200 E2E Kubernetes cluster with:

  • two trainer nodes, 16 trainer ranks, EP=8
  • two initial TP=8 inference replicas
  • one additional TP=8 inference replica started about 90 seconds later
  • Laguna-XS.2, 100-step math run, NIXL + ModelExpress

The third replica registered as unavailable for rollout traffic, joined the next weight update as inference ranks 16–23, became routable after that update, and the job continued with 24 participating inference ranks.

Stack

Current PoC limitations

  • only homogeneous inference servers are supported; ranks per server are fixed from the initial pool
  • scale-up was validated; arbitrary shrink/grow churn and stable-URL pod replacement still need lifecycle work
  • rank assignments and membership are process-local and are not restored after trainer/router restart
  • a worker failure during an in-flight synchronization remains fail-stop

Checks

  • ruff format --check on changed Python files
  • ruff check on changed Python files
  • end-to-end staggered 2×TP8 → 3×TP8 H200 run

@S1ro1
S1ro1 changed the base branch from fix/nixl-replay-strided-views to main July 23, 2026 05:43
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.

1 participant