From 22913ab7a3626ccb7d5dbe7c91a08d3b98353889 Mon Sep 17 00:00:00 2001 From: Kotaro Yoshimoto Date: Mon, 7 Sep 2026 18:27:33 +0900 Subject: [PATCH] fix(closed_loop): bind ONNX providers to the rank's CUDA device The CLI selects cuda:{local_rank}, but ONNX Runtime reads no device from torch and defaults to GPU 0, so every rank of a distributed ONNX evaluation piles onto the first visible GPU. Pass the requested index to the CUDA and TensorRT provider options; an unindexed "cuda" follows torch's current device. CPU-only sessions and the TensorRT cache options are unchanged. The provider tests grew a third fake session to cover this, so fold all three onto one fake and one helper, and assert whole option dicts per provider rather than a single key, which also catches stray options. --- scenario_generation/simulate.py | 7 ++ .../tests/test_onnx_provider.py | 111 ++++++++++-------- 2 files changed, 69 insertions(+), 49 deletions(-) diff --git a/scenario_generation/simulate.py b/scenario_generation/simulate.py index fe863ad15..c33f7a01e 100644 --- a/scenario_generation/simulate.py +++ b/scenario_generation/simulate.py @@ -120,6 +120,13 @@ def __init__( else {} for p in providers ] + target = torch.device(device) + if target.type == "cuda" and any(p in _ACCELERATED for p in providers): + # ORT defaults to GPU 0 independently of torch.cuda.set_device() in each rank. + device_id = target.index if target.index is not None else torch.cuda.current_device() + for provider, option in zip(providers, options): + if provider in _ACCELERATED: + option["device_id"] = device_id self.session = ort.InferenceSession( str(onnx_path), providers=providers, provider_options=options ) diff --git a/scenario_generation/tests/test_onnx_provider.py b/scenario_generation/tests/test_onnx_provider.py index 29e6f8112..d12402311 100644 --- a/scenario_generation/tests/test_onnx_provider.py +++ b/scenario_generation/tests/test_onnx_provider.py @@ -9,13 +9,45 @@ from scenario_generation.simulate import CPU_EP, CUDA_EP, TENSORRT_EP, _require_accelerator +# Distinct from every index the cases ask for, so a device_id off the wrong source shows up. +_CURRENT_DEVICE = 2 +_TRT_CACHE = "/engines" +_TRT_CACHE_OPTIONS = { + "trt_engine_cache_enable": True, + "trt_engine_cache_path": _TRT_CACHE, + "trt_timing_cache_enable": True, +} + class _Session: - def __init__(self, active): - self._active = active + """Stands in for ``ort.InferenceSession``, reporting the providers it was handed as active.""" + + def __init__(self, providers, provider_options=None): + self.providers = providers + self.provider_options = provider_options def get_providers(self): - return self._active + return self.providers + + def get_inputs(self): + return [] + + def get_outputs(self): + return [] + + +def _open_session(monkeypatch, device, providers=None, **kwargs): + """Build an ``_OnnxModel`` against a fake onnxruntime and hand back the session it opened.""" + import scenario_generation.simulate as simulate + + def _factory(path, providers, provider_options): + return _Session(providers, provider_options) + + monkeypatch.setattr(simulate.torch.cuda, "current_device", lambda: _CURRENT_DEVICE) + monkeypatch.setitem( + sys.modules, "onnxruntime", types.SimpleNamespace(InferenceSession=_factory) + ) + return simulate._OnnxModel("m.onnx", device, providers, **kwargs).session def test_a_gpu_request_that_landed_on_cpu_raises(): @@ -46,52 +78,33 @@ def test_asking_for_cpu_is_not_a_failure(): def test_the_default_does_not_reach_for_tensorrt(monkeypatch): """TensorRT partitions the graph up front and refuses ops it cannot build, so defaulting to it turns a model it dislikes into a session that never opens. It has to be asked for.""" - seen = {} - - class _Session: - def __init__(self, path, providers=None, provider_options=None): - seen["providers"] = providers - - def get_providers(self): - return [CUDA_EP, CPU_EP] - - def get_inputs(self): - return [] - - def get_outputs(self): - return [] - - import scenario_generation.simulate as simulate - - monkeypatch.setitem( - sys.modules, "onnxruntime", types.SimpleNamespace(InferenceSession=_Session) - ) - simulate._OnnxModel("m.onnx", "cuda") - - assert seen["providers"] == [CUDA_EP, CPU_EP] + assert _open_session(monkeypatch, "cuda").providers == [CUDA_EP, CPU_EP] def test_asking_for_cpu_by_device_does_not_request_a_gpu_provider(monkeypatch): - seen = {} - - class _Session: - def __init__(self, path, providers=None, provider_options=None): - seen["providers"] = providers - - def get_providers(self): - return [CPU_EP] - - def get_inputs(self): - return [] - - def get_outputs(self): - return [] - - import scenario_generation.simulate as simulate - - monkeypatch.setitem( - sys.modules, "onnxruntime", types.SimpleNamespace(InferenceSession=_Session) - ) - simulate._OnnxModel("m.onnx", "cpu") - - assert seen["providers"] == [CPU_EP] + assert _open_session(monkeypatch, "cpu").providers == [CPU_EP] + + +@pytest.mark.parametrize( + "device,providers,expected", + [ + # 0 is a real GPU, not "unset": it must not read as absent and fall back to the default. + ("cuda:0", [CUDA_EP, CPU_EP], {CUDA_EP: {"device_id": 0}, CPU_EP: {}}), + ("cuda", [CUDA_EP, CPU_EP], {CUDA_EP: {"device_id": _CURRENT_DEVICE}, CPU_EP: {}}), + ( + "cuda:1", + [TENSORRT_EP, CUDA_EP, CPU_EP], + { + TENSORRT_EP: {"device_id": 1, **_TRT_CACHE_OPTIONS}, + CUDA_EP: {"device_id": 1}, + CPU_EP: {}, + }, + ), + ], +) +def test_gpu_providers_follow_requested_device(monkeypatch, device, providers, expected): + """ORT ignores torch.cuda.set_device() and defaults to GPU 0, so every rank of a distributed + run piles onto the first visible GPU unless the provider itself carries the index.""" + session = _open_session(monkeypatch, device, providers, engine_cache_dir=_TRT_CACHE) + + assert dict(zip(session.providers, session.provider_options)) == expected