fix(closed_loop): bind ONNX providers to the rank's CUDA device

[HansRobo]

Problem

ONNX Runtime takes no device from torch. The CLI selects cuda:{local_rank}, but every rank's ORT session opens on GPU 0.

Live CUDA contexts per pid, 4 ranks on 4 H100s, same node and same onnx, two arms differing only by the 7 lines below:

Without the fix

GPU  pid      mem
0    ...510   754 MiB   rank 0: torch + its own ORT session
0    ...511   658 MiB   rank 1's ORT session
0    ...512   658 MiB   rank 2's ORT session
0    ...513   658 MiB   rank 3's ORT session
1    ...511   618 MiB   rank 1: torch only
2    ...512   618 MiB   rank 2: torch only
3    ...513   618 MiB   rank 3: torch only

With the fix

GPU  pid      mem
0    ...324   754 MiB   rank 0
1    ...325   754 MiB   rank 1
2    ...326   754 MiB   rank 2
3    ...327   754 MiB   rank 3

Every non-zero rank held a second context on GPU 0 — its ORT session — while its own GPU carried only the torch context. GPU 0 served all four ranks while GPUs 1-3 sat in a collective-wait kernel: 0% memory bandwidth, ~120 W against a 700 W TDP.

Effect on closed-loop lap

run_all_groups_closed_loop.py over 11 routes x 2 object modes, 4 ranks on one node, ONNX planner. Both arms ran back to back in one allocation on the same tree; only scenario_generation/simulate.py differed.

                          baseline    with fix    ratio
wall                        8902 s      5263 s     1.69x
  startup                   1120 s       629 s     (page cache, not the change)
  route 1 -> route 22       7779 s      4631 s     1.68x

Per-route speedup tracks route length, from 0.92x on a 20-second route to 2.84x on the longest (6604 steps). Short routes are dominated by setup and I/O; long ones by inference, which is what serialised onto GPU 0. It is 1.7x rather than 4x because GPU 0 was not saturated — inference is roughly half the critical path, the rest is CPU-side simulation.

Results are unchanged: the per-mode groups.json of both arms are byte-identical. The cross-mode aggregate differs in the last 1-2 ULP (mean_route_completion ...171 vs ...173), which is summation order, not behaviour.

Fix

Pass the requested index to the CUDA and TensorRT provider options; an unindexed cuda follows torch.cuda.current_device(). Seven lines. CPU-only sessions and the TensorRT cache options are unchanged.

Notes

• Tests add three cases: cuda:0 (index 0 is a real GPU, not "unset"), unindexed cuda, and cuda:1 with TensorRT (nonzero index, cache options preserved). The deletions are the three fake ORT sessions already in test_onnx_provider.py folded onto one helper.
• With device="cuda" (no index) on a host without CUDA, torch.cuda.current_device() now raises before the session opens, replacing _require_accelerator's actionable message with a raw torch error. Left as-is: the eval CLI always passes cuda:{local_rank}, which never reaches that path.

Reference: https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#device_id

[go-sakayori]

LGTM

[kosuke55]

Thanks!