[Distributed] Implement Multi-GPU Session Load Balancer

[DsThakurRawat]

Description (Medium)

The server currently only utilizes device_id: 0.

Task

• Create a MultiGpuDispatcher that can instantiate N ONNX Runtime sessions across N GPUs.
• Implement a round-robin routing policy to distribute /act requests across available compute resources.

[huynna12]

Is it still available? I'm happy to work on this problem.

[DsThakurRawat]

Yes, still available, and it's yours. Thanks for picking it up.

Fair warning that the repo has moved around quite a bit since this issue was written, so here's what's worth knowing before you start.

The package got renamed

• It was reflex, it's tether now (as of v0.12.0). All the real code is in src/tether/. There's still a src/reflex/ directory but it's just a compat shim that re-exports and prints a deprecation warning, so don't put anything there.
• To get set up the same way CI does: uv pip install -e ".[dev,onnx,serve]" then pytest tests/.

Where the hardcoded device_id: 0 actually lives

• It's in src/tether/runtime/ort_providers.py, around line 320 for the TensorRT provider options and line 343 for the CUDA ones.
• That function is really the only place you need to touch to get a device index flowing through, since basically every session in the codebase gets built from the plan it returns. It already accepts a requested_providers override, so widening OrtProviderPlan to carry a device index is a lot cleaner than having the dispatcher assemble providers on its own.

One thing in the issue body is misleading, sorry about that

• The task list says round-robin the /act requests. Please don't do that literally, it'll break correctness.
• Serving here is stateful per episode. Action chunk buffers, temporal_vla_cache, the RTC reset that fires when episode_id changes (server.py around line 3493), and the A2C2 hook all assume that consecutive /act calls from the same episode land on the same session.
• So the round-robin belongs at the point where a new episode gets assigned to a GPU, and after that the episode should stick to it. That's what the "Session" in the issue title is getting at.

There's already a router doing almost exactly this

• Have a look at runtime/policy_router.py and runtime/two_policy_dispatcher.py. They already do episode-sticky routing (SHA-256 on episode_id, first request picks the slot, everything after that follows it) plus draining a slot when it keeps crashing. Worth modelling MultiGpuDispatcher on that rather than starting fresh.
• Also worth copying the shape: TwoPolicyDispatcher is pure composition, no I/O of its own, doesn't touch FastAPI. Keeping MultiGpuDispatcher the same way is what'll make it testable without a GPU, which matters (see below).
• When episode_id is missing it falls back to session_id and logs a one-time degraded routing warning. Matching that behaviour would be good.

Two device affinity gotchas that'll bite you

• The TensorRT engine and timing caches all currently sit under a single path in the export dir (ort_providers.py around line 285). TRT engines are built per GPU architecture, so on a box with mixed GPUs a shared cache will happily hand one GPU an engine compiled for a different one. Needs to be split per device.
• CUDA graphs (runtime/cuda_graphs.py) get captured on one specific device and can't be moved. So any CudaGraphWrapper has to be created per device, after the session is already pinned.

Testing, and the CI situation

• CI is Linux only and has no GPU at all, so none of the multi-GPU paths can run there. If you build it against an injectable session factory you can at least unit test the routing logic with fakes, which is the only way this gets meaningful coverage.
• Anything that genuinely needs hardware goes behind the existing hardware marker (RUN_HARDWARE_TESTS=1).
• Also a heads up on the quality ratchet: it diffs ruff and mypy findings between the base commit and your branch and fails if you add any new ones. It also won't let a PR edit the workflows or the [tool.ruff] / [tool.mypy] sections of pyproject.toml, so silencing things that way won't get through. Pinned versions are ruff==0.15.10 and mypy==1.20.1.

Process stuff

• Branch off main, one concern per PR, and add a regression test.
• If you end up adding any Prometheus metrics, please make sure they actually get incremented and that the label values match the documented enums. We've got [R3] Registered Prometheus metrics never incremented (dead series) #279 and [R4] Safety/SLO metric label values don't match documented enums #280 open from metrics that were registered and then never fired.

If you want to sketch the design in this thread before writing much code, happy to look it over. And just ask if any of the above is unclear.