From f363288c432ad9d3b9458a6499e2067456416e98 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Wed, 8 Jul 2026 14:47:26 +0800 Subject: [PATCH] feat: add offpolicy collector infer device --- conf/offpolicy/config.yaml | 1 + scripts/train_offpolicy.py | 4 + .../algos/torch/flash_sac/double_buffer.py | 3 + .../torch/offpolicy/double_buffer_runner.py | 6 ++ .../algos/torch/offpolicy/multi_gpu_runner.py | 18 +++- src/unilab/algos/torch/offpolicy/runner.py | 16 ++- src/unilab/algos/torch/offpolicy/worker.py | 33 ++++-- src/unilab/utils/device.py | 102 ++++++++++++++++++ .../test_offpolicy_double_buffer_runner.py | 55 ++++++++++ tests/algos/test_offpolicy_runner_unit.py | 22 ++++ tests/utils/test_device.py | 94 ++++++++++++++++ 11 files changed, 339 insertions(+), 15 deletions(-) create mode 100644 tests/utils/test_device.py diff --git a/conf/offpolicy/config.yaml b/conf/offpolicy/config.yaml index 74984f226..73ef04728 100644 --- a/conf/offpolicy/config.yaml +++ b/conf/offpolicy/config.yaml @@ -6,6 +6,7 @@ defaults: training: task_name: G1WalkFlat device: null + collector_infer_device: cpu logger: tensorboard wandb_project: unilab wandb_entity: null diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index daf7d98de..be9c143c7 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -178,6 +178,7 @@ def build_runner(algo_name: str, cfg: DictConfig): num_gpus = int(getattr(cfg.training, "num_gpus", 1)) multi_gpu_sync_mode = str(getattr(cfg.training, "multi_gpu_sync_mode", "local_sgd")) multi_gpu_sync_interval = int(getattr(cfg.training, "multi_gpu_sync_interval", 1)) + collector_infer_device = str(getattr(cfg.training, "collector_infer_device", "cpu") or "cpu") sync_collection = not bool(cfg.training.no_sync_collection) @@ -318,6 +319,7 @@ def build_runner(algo_name: str, cfg: DictConfig): trace_cuda_events=cfg.training.trace_cuda_events, seed=cfg.algo.seed, nan_guard_cfg=_nan_guard_cfg, + collector_infer_device=collector_infer_device, ) _learner = _learner_cls(device=_device, **_learner_kwargs) @@ -348,6 +350,7 @@ def build_runner(algo_name: str, cfg: DictConfig): verbose_metrics=verbose_metrics, seed=cfg.algo.seed, nan_guard_cfg=_nan_guard_cfg, + collector_infer_device=collector_infer_device, ) if algo_name == "td3": @@ -433,6 +436,7 @@ def build_runner(algo_name: str, cfg: DictConfig): verbose_metrics=verbose_metrics, actor_kwargs=_actor_kwargs, nan_guard_cfg=_nan_guard_cfg, + collector_infer_device=collector_infer_device, ) if algo_name == "flashsac": diff --git a/src/unilab/algos/torch/flash_sac/double_buffer.py b/src/unilab/algos/torch/flash_sac/double_buffer.py index 178906698..698e10424 100644 --- a/src/unilab/algos/torch/flash_sac/double_buffer.py +++ b/src/unilab/algos/torch/flash_sac/double_buffer.py @@ -55,6 +55,7 @@ def build_flashsac_double_buffer_runner( ensure_registries() apply_training_seed(cfg.algo.seed, torch_runtime=True, cuda=True) device = cfg.training.device or get_default_device() + collector_infer_device = str(getattr(cfg.training, "collector_infer_device", "cpu") or "cpu") _validate_flashsac_double_buffer_runtime( cfg, device=device, @@ -152,6 +153,7 @@ def build_flashsac_double_buffer_runner( trace_thread_time=cfg.training.trace_thread_time, trace_cuda_events=cfg.training.trace_cuda_events, nan_guard_cfg=nan_guard_cfg, + collector_infer_device=collector_infer_device, ) learner = FlashSACLearner(device=device, **learner_kwargs) @@ -183,4 +185,5 @@ def build_flashsac_double_buffer_runner( replay_prefetch_mode=replay_prefetch_mode, verbose_metrics=verbose_metrics, nan_guard_cfg=nan_guard_cfg, + collector_infer_device=collector_infer_device, ) diff --git a/src/unilab/algos/torch/offpolicy/double_buffer_runner.py b/src/unilab/algos/torch/offpolicy/double_buffer_runner.py index 73194c1dd..c51ddf9c0 100644 --- a/src/unilab/algos/torch/offpolicy/double_buffer_runner.py +++ b/src/unilab/algos/torch/offpolicy/double_buffer_runner.py @@ -282,6 +282,10 @@ def learn( f"{self.replay_transfer_backend.get('backend')} " f"({self.replay_transfer_backend.get('device_family')})" ) + logger.log_status( + "Collector infer device: " + f"{self.collector_infer_device_raw} -> {self.collector_infer_device}" + ) logger.log_status("Replay learner lightweight: fixed (log_interval=1)") if self.verbose_metrics: logger.log_status("Verbose metrics: enabled (field-level pack CSV)") @@ -332,6 +336,8 @@ def learn( "obs_dim": self.obs_dim, "action_dim": self.action_dim, "actor_kwargs": self.actor_kwargs, + "collector_infer_device": self.collector_infer_device, + "collector_infer_device_raw": self.collector_infer_device_raw, "seed": derive_worker_seed(self.seed, worker_index=0), "trace_enabled": self.trace_enabled, "trace_thread_time": self.trace_thread_time, diff --git a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py index 377b4f035..c23dfc6f7 100644 --- a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py +++ b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py @@ -2,7 +2,7 @@ Architecture: Main process → creates ReplayBuffer (host-only), WeightSync, queues - → spawns Collector subprocess (CPU, env simulation) + → spawns Collector subprocess (CPU env I/O, configurable inference device) → spawns N Learner workers via mp.spawn (one per GPU) Learner rank i → samples packed CPU replay rows to its rank device through a rank-local H2D pipeline, then either averages gradients @@ -276,6 +276,11 @@ def _learner_worker( f"{sync_mode} (interval={sync_interval} iteration" f"{'s' if sync_interval != 1 else ''})" ) + logger.log_status( + "Collector infer device: " + f"{runner_kwargs.get('collector_infer_device_raw', 'cpu')} -> " + f"{runner_kwargs.get('collector_infer_device', 'cpu')}" + ) if sync_mode == "local_sgd": logger.log_status( "Local-SGD optimizer state: rank-local; parameters averaged at sync boundary" @@ -571,8 +576,9 @@ def _learner_worker( class MultiGPUOffPolicyRunner(OffPolicyRunner): """Multi-GPU off-policy runner. - Keeps a single Collector on CPU and spawns *num_gpus* Learner workers via - ``torch.multiprocessing.spawn``. Each worker processes an independent + Keeps a single Collector process and spawns *num_gpus* Learner workers via + ``torch.multiprocessing.spawn``. Env I/O remains CPU/numpy while collector + actor inference can use a configured device. Each worker processes an independent mini-batch from the same shared ReplayBuffer through a rank-local H2D pipeline. SAC defaults to local-SGD: ranks apply local updates and average parameters at runner-controlled synchronization boundaries. Strict per-update @@ -733,7 +739,7 @@ def _learn_multi_gpu( for _ in range(self.num_gpus) ] - # --- Start Collector (CPU, single process, unchanged) --- + # --- Start Collector (single process, device-configurable inference) --- weight_param_shapes = {k: v.shape for k, v in self.learner.actor.state_dict().items()} collector_kwargs = { "env_name": self.env_name, @@ -758,6 +764,8 @@ def _learn_multi_gpu( "obs_dim": self.obs_dim, "action_dim": self.action_dim, "actor_kwargs": self.actor_kwargs, + "collector_infer_device": self.collector_infer_device, + "collector_infer_device_raw": self.collector_infer_device_raw, "seed": derive_worker_seed(self.seed, worker_index=0), "collector_pack_request_queue": collector_pack_request_queues, "collector_pack_ready_queue": collector_pack_ready_queues, @@ -798,6 +806,8 @@ def _learn_multi_gpu( "algo_type": self.algo_type, "obs_normalization": self.obs_normalization, "shared_obs_normalizer_stats": shared_obs_normalizer_stats, + "collector_infer_device": self.collector_infer_device, + "collector_infer_device_raw": self.collector_infer_device_raw, } try: diff --git a/src/unilab/algos/torch/offpolicy/runner.py b/src/unilab/algos/torch/offpolicy/runner.py index e9bd7a5fa..6e12a4c72 100644 --- a/src/unilab/algos/torch/offpolicy/runner.py +++ b/src/unilab/algos/torch/offpolicy/runner.py @@ -17,7 +17,7 @@ from unilab.ipc.replay_buffer import ReplayBuffer from unilab.logging import OffPolicyLogger, TraceRecorder from unilab.training.seed import apply_training_seed, derive_worker_seed -from unilab.utils.device import get_default_device +from unilab.utils.device import get_default_device, resolve_torch_device_alias from unilab.utils.nan_guard import NanGuardCfg @@ -176,13 +176,19 @@ def __init__( trace_thread_time: bool = False, trace_cuda_events: bool = True, nan_guard_cfg: NanGuardCfg | None = None, + collector_infer_device: str | None = "cpu", ): + self.collector_infer_device_raw = str(collector_infer_device or "cpu") + self.collector_infer_device = resolve_torch_device_alias( + self.collector_infer_device_raw, + default="cpu", + ) super().__init__( env_name=env_name, env_cfg_overrides={}, rl_cfg={}, device=device, - collector_device="cpu", + collector_device=self.collector_infer_device, num_envs=num_envs, sim_backend=sim_backend, ) @@ -332,6 +338,8 @@ def learn( "obs_dim": self.obs_dim, "action_dim": self.action_dim, "actor_kwargs": self.actor_kwargs, + "collector_infer_device": self.collector_infer_device, + "collector_infer_device_raw": self.collector_infer_device_raw, "seed": derive_worker_seed(self.seed, worker_index=0), "trace_enabled": self.trace_enabled, "trace_thread_time": self.trace_thread_time, @@ -362,6 +370,10 @@ def learn( logger.set_collection_sync(self.sync_collection, self.env_steps_per_sync) if hasattr(self.learner, "use_symmetry") and self.learner.use_symmetry: logger.log_status("Symmetry augmentation: enabled") + logger.log_status( + "Collector infer device: " + f"{self.collector_infer_device_raw} -> {self.collector_infer_device}" + ) self._active_logger = logger logger.start() diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index e64b9d699..eb853bc61 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -490,6 +490,8 @@ def off_policy_collector_fn( collector_pack_ready_queue=None, collector_pack_shared_slots=None, nan_guard_cfg=None, + collector_infer_device: str = "cpu", + collector_infer_device_raw: str | None = None, **kwargs, ): """Entry point for the off-policy collector subprocess. @@ -529,6 +531,8 @@ def off_policy_collector_fn( collector_pack_ready_queue=collector_pack_ready_queue, collector_pack_shared_slots=collector_pack_shared_slots, nan_guard_cfg=nan_guard_cfg, + collector_infer_device=collector_infer_device, + collector_infer_device_raw=collector_infer_device_raw, ) @@ -563,6 +567,8 @@ def _run_collector( collector_pack_ready_queue, collector_pack_shared_slots, nan_guard_cfg=None, + collector_infer_device: str = "cpu", + collector_infer_device_raw: str | None = None, ): del learning_starts from unilab.base import registry @@ -601,7 +607,10 @@ def _run_collector( weight_sync.trace_recorder = trace_recorder weight_sync.trace_thread_time = trace_thread_time - # Build actor (always on CPU for env interaction) + collector_infer_device = str(collector_infer_device or "cpu") + collector_infer_device_raw = str(collector_infer_device_raw or collector_infer_device) + + # Build actor on the resolved collector inference device. Env I/O remains numpy. obs_dim, action_dim = resolve_collector_actor_dims( env, obs_dim=obs_dim, @@ -613,7 +622,7 @@ def _run_collector( action_dim, actor_hidden_dim, use_layer_norm, - "cpu", + collector_infer_device, num_envs, **(actor_kwargs or {}), ) @@ -635,8 +644,8 @@ def _run_collector( from collections import defaultdict ep_reward_components = defaultdict(list) - timing_accum_ms = defaultdict(float) - timing_counts = defaultdict(int) + timing_accum_ms: defaultdict[str, float] = defaultdict(float) + timing_counts: defaultdict[str, int] = defaultdict(int) done_count_window = 0 timeout_count_window = 0 terminated_count_window = 0 @@ -711,8 +720,8 @@ def _run_collector( # Select action with torch.no_grad(): _t_infer_ns = _time.perf_counter_ns() - obs_torch = torch.from_numpy(obs_np_input) - dones_torch = torch.from_numpy(prev_dones_np) + obs_torch = torch.from_numpy(obs_np_input).to(collector_infer_device) + dones_torch = torch.from_numpy(prev_dones_np).to(collector_infer_device) priv_info_np = resolve_offpolicy_actor_priv_info( algo_type=algo_type, obs_np=obs_np, @@ -720,7 +729,9 @@ def _run_collector( info=info_dict, ) priv_info_torch = ( - torch.from_numpy(priv_info_np) if priv_info_np is not None else None + torch.from_numpy(priv_info_np).to(collector_infer_device) + if priv_info_np is not None + else None ) actions_torch = sample_offpolicy_actions( actor=actor, @@ -729,13 +740,17 @@ def _run_collector( prev_dones_torch=dones_torch, priv_info_torch=priv_info_torch, ) - actions_np = actions_torch.numpy() + actions_np = actions_torch.detach().cpu().numpy() if trace_recorder: trace_recorder.add_slice( - "collector/actor_infer_cpu", + "collector/actor_infer", category="collector", start_ns=_t_infer_ns, end_ns=_time.perf_counter_ns(), + args={ + "collector_infer_device_raw": collector_infer_device_raw, + "collector_infer_device": collector_infer_device, + }, ) phase_start_ns = _record_phase_ms(cycle_timing_ms, "action_select_ms", phase_start_ns) diff --git a/src/unilab/utils/device.py b/src/unilab/utils/device.py index c2c9fbab5..f5a4cdcee 100644 --- a/src/unilab/utils/device.py +++ b/src/unilab/utils/device.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import Callable, cast + import torch @@ -18,3 +20,103 @@ def get_default_device() -> str: if torch.backends.mps.is_available(): return "mps" return "cpu" + + +def _device_count(device_type: str) -> int | None: + if device_type == "cuda": + return int(torch.cuda.device_count()) + if device_type == "xpu": + xpu = getattr(torch, "xpu", None) + device_count = getattr(xpu, "device_count", None) + if callable(device_count): + return int(cast(Callable[[], int], device_count)()) + return None + + +def _mps_available() -> bool: + mps = getattr(torch.backends, "mps", None) + is_available = getattr(mps, "is_available", None) + return bool(callable(is_available) and is_available()) + + +def _parse_device_alias(value: str) -> tuple[str, int | None]: + raw = value.strip().lower() + if not raw: + raise ValueError("Device alias must not be empty") + if ":" not in raw: + return raw, None + base, index_text = raw.split(":", 1) + if not index_text: + raise ValueError(f"Device alias {value!r} has an empty index") + try: + index = int(index_text) + except ValueError as exc: + raise ValueError(f"Device alias {value!r} has a non-integer index") from exc + if index < 0: + raise ValueError(f"Device alias {value!r} has a negative index") + return base, index + + +def _resolve_indexed_device(device_type: str, index: int | None, original: str) -> str: + count = _device_count(device_type) + if count is not None and index is not None and index >= count: + raise ValueError( + f"Requested device {original!r} resolves to {device_type}:{index}, " + f"but only {count} {device_type} device(s) are available" + ) + return device_type if index is None else f"{device_type}:{index}" + + +def _resolve_mps_alias(index: int | None, original: str) -> str: + if not _mps_available(): + raise ValueError(f"Requested device {original!r} requires MPS, but MPS is unavailable") + if index not in (None, 0): + raise ValueError( + f"Requested device {original!r} cannot be mapped to MPS; only index 0 is valid" + ) + return "mps" + + +def resolve_torch_device_alias(device: str | None, *, default: str = "cpu") -> str: + """Resolve a cross-platform torch device alias to a concrete device string. + + ``gpu`` is an abstract accelerator alias. ``cuda`` is also accepted on + macOS/MPS for config portability and maps to ``mps`` when CUDA is absent. + The function validates the resolved device and never silently falls back to + CPU for unavailable accelerators. + """ + original = default if device is None else str(device) + base, index = _parse_device_alias(original) + + if base == "cpu": + if index is not None: + raise ValueError(f"CPU device {original!r} must not include an index") + return "cpu" + + if base == "mps": + return _resolve_mps_alias(index, original) + + if base == "xpu": + if not _xpu_available(): + raise ValueError(f"Requested device {original!r} requires XPU, but XPU is unavailable") + return _resolve_indexed_device("xpu", index, original) + + if base == "cuda": + if torch.cuda.is_available(): + return _resolve_indexed_device("cuda", index, original) + if _mps_available(): + return _resolve_mps_alias(index, original) + raise ValueError(f"Requested device {original!r} requires CUDA, but CUDA is unavailable") + + if base == "gpu": + if torch.cuda.is_available(): + return _resolve_indexed_device("cuda", index, original) + if _xpu_available(): + return _resolve_indexed_device("xpu", index, original) + if _mps_available(): + return _resolve_mps_alias(index, original) + raise ValueError( + f"Requested device {original!r} requires an accelerator, but none is available" + ) + + raise ValueError(f"Unsupported device alias {original!r}; expected cpu, gpu, cuda, mps, or xpu") diff --git a/tests/algos/test_offpolicy_double_buffer_runner.py b/tests/algos/test_offpolicy_double_buffer_runner.py index ad8b538e0..b4e9301b9 100644 --- a/tests/algos/test_offpolicy_double_buffer_runner.py +++ b/tests/algos/test_offpolicy_double_buffer_runner.py @@ -58,6 +58,11 @@ def test_default_replay_prefetch_mode_is_one_tick(): assert cfg.training.replay_prefetch_mode == "one_tick" +def test_default_collector_infer_device_is_cpu(): + cfg = _offpolicy_cfg() + assert cfg.training.collector_infer_device == "cpu" + + def test_b_path_internal_knobs_are_not_configured(): cfg = _offpolicy_cfg() assert "replay_pack_layout" not in cfg.training @@ -315,10 +320,60 @@ def __init__(self, *args, **kwargs): assert isinstance(runner, _FakeRunner) assert runner.kwargs["device"] == device + assert runner.kwargs["collector_infer_device"] == "cpu" assert runner.kwargs["replay_prefetch_mode"] == "one_tick" assert runner.kwargs["learner"].kwargs["use_compile"] is True +def test_sac_collector_infer_device_override_is_passed_to_runner( + monkeypatch: pytest.MonkeyPatch, +): + import gymnasium as gym + + mod = _offpolicy() + cfg = _offpolicy_cfg( + [ + "algo=sac", + "training.device=cpu", + "training.collector_infer_device=gpu:0", + "algo.use_symmetry=false", + ] + ) + + class _FakeEnv: + obs_groups_spec = {"obs": 4, "critic": 6} + action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) + + def close(self): + pass + + class _FakeLearner: + class actor: + @staticmethod + def state_dict(): + return {"w": MagicMock(shape=(4,))} + + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + class _FakeRunner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + monkeypatch.setattr(mod, "ensure_registries", lambda: None) + monkeypatch.setattr(mod, "create_env", lambda *args, **kwargs: _FakeEnv()) + + import unilab.algos.torch.fast_sac.learner as learner_mod + import unilab.algos.torch.offpolicy.double_buffer_runner as db_mod + + monkeypatch.setattr(learner_mod, "FastSACLearner", _FakeLearner) + monkeypatch.setattr(db_mod, "DoubleBufferOffPolicyRunner", _FakeRunner) + + runner = mod.build_runner("sac", cfg) + + assert runner.kwargs["collector_infer_device"] == "gpu:0" + + def test_sac_compile_override_is_passed_to_learner(monkeypatch: pytest.MonkeyPatch): import gymnasium as gym diff --git a/tests/algos/test_offpolicy_runner_unit.py b/tests/algos/test_offpolicy_runner_unit.py index 952fd653c..188b70f31 100644 --- a/tests/algos/test_offpolicy_runner_unit.py +++ b/tests/algos/test_offpolicy_runner_unit.py @@ -393,6 +393,7 @@ def _make_runner( policy_frequency: int = 1, trace_enabled: bool = False, trace_output_dir: str | None = None, + collector_infer_device: str | None = "cpu", ) -> OffPolicyRunner: monkeypatch.setattr(runner_module, "ReplayBuffer", _FakeReplayBuffer) monkeypatch.setattr(runner_module, "SharedWeightSync", _FakeWeightSync) @@ -416,12 +417,33 @@ def _make_runner( device="cpu", trace_enabled=trace_enabled, trace_output_dir=trace_output_dir, + collector_infer_device=collector_infer_device, ) monkeypatch.setattr(runner, "_start_collector", lambda *args, **kwargs: None) runner._collector_process = _FakeProcess() return runner +def test_offpolicy_runner_resolves_collector_infer_device( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + runner_module, + "resolve_torch_device_alias", + lambda device, *, default="cpu": "mps" if device == "gpu" else "cpu", + ) + + runner = _make_runner( + monkeypatch, + sync_collection=True, + collector_infer_device="gpu", + ) + + assert runner.collector_infer_device_raw == "gpu" + assert runner.collector_infer_device == "mps" + assert runner.collector_device == "mps" + + def test_offpolicy_runner_sync_waits_for_train_start_threshold( monkeypatch: pytest.MonkeyPatch, tmp_path ) -> None: diff --git a/tests/utils/test_device.py b/tests/utils/test_device.py new file mode 100644 index 000000000..db3967a74 --- /dev/null +++ b/tests/utils/test_device.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +import unilab.utils.device as device_mod + + +def _patch_devices( + monkeypatch: pytest.MonkeyPatch, + *, + cuda: bool = False, + cuda_count: int = 0, + xpu: bool = False, + xpu_count: int = 0, + mps: bool = False, +) -> None: + monkeypatch.setattr(device_mod.torch.cuda, "is_available", lambda: cuda) + monkeypatch.setattr(device_mod.torch.cuda, "device_count", lambda: cuda_count) + monkeypatch.setattr( + device_mod.torch, + "xpu", + SimpleNamespace( + is_available=lambda: xpu, + device_count=lambda: xpu_count, + ), + raising=False, + ) + monkeypatch.setattr(device_mod.torch.backends.mps, "is_available", lambda: mps) + + +def test_resolve_torch_device_alias_defaults_to_cpu(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_devices(monkeypatch) + + assert device_mod.resolve_torch_device_alias(None) == "cpu" + assert device_mod.resolve_torch_device_alias("cpu") == "cpu" + + +def test_resolve_torch_device_alias_gpu_prefers_cuda( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_devices(monkeypatch, cuda=True, cuda_count=2, mps=True) + + assert device_mod.resolve_torch_device_alias("gpu") == "cuda" + assert device_mod.resolve_torch_device_alias("gpu:1") == "cuda:1" + + +def test_resolve_torch_device_alias_gpu_uses_xpu_before_mps( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_devices(monkeypatch, xpu=True, xpu_count=2, mps=True) + + assert device_mod.resolve_torch_device_alias("gpu") == "xpu" + assert device_mod.resolve_torch_device_alias("gpu:1") == "xpu:1" + + +@pytest.mark.parametrize("alias", ["gpu", "gpu:0", "cuda", "cuda:0"]) +def test_resolve_torch_device_alias_macos_compat_maps_gpu_and_cuda_to_mps( + monkeypatch: pytest.MonkeyPatch, + alias: str, +) -> None: + _patch_devices(monkeypatch, mps=True) + + assert device_mod.resolve_torch_device_alias(alias) == "mps" + + +@pytest.mark.parametrize("alias", ["gpu:1", "cuda:1", "mps:1"]) +def test_resolve_torch_device_alias_mps_rejects_nonzero_index( + monkeypatch: pytest.MonkeyPatch, + alias: str, +) -> None: + _patch_devices(monkeypatch, mps=True) + + with pytest.raises(ValueError, match="MPS"): + device_mod.resolve_torch_device_alias(alias) + + +def test_resolve_torch_device_alias_rejects_missing_cuda_index( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_devices(monkeypatch, cuda=True, cuda_count=1) + + with pytest.raises(ValueError, match="only 1 cuda device"): + device_mod.resolve_torch_device_alias("gpu:1") + + +def test_resolve_torch_device_alias_rejects_unavailable_accelerator( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_devices(monkeypatch) + + with pytest.raises(ValueError, match="none is available"): + device_mod.resolve_torch_device_alias("gpu")