diff --git a/src/twinkle/checkpoint_engine/ipc_checkpoint_engine.py b/src/twinkle/checkpoint_engine/ipc_checkpoint_engine.py index b008f15e8..2353e0e4b 100644 --- a/src/twinkle/checkpoint_engine/ipc_checkpoint_engine.py +++ b/src/twinkle/checkpoint_engine/ipc_checkpoint_engine.py @@ -31,7 +31,8 @@ import zmq from typing import Any, AsyncGenerator, Generator -from twinkle import get_logger +from twinkle import Platform, get_logger +from twinkle.utils.framework import Torch from .base import CheckpointEngine logger = get_logger() @@ -49,7 +50,7 @@ class IPCCheckpointEngine(CheckpointEngine): - """Hand weights to a sampler on the same GPU by mapping memory instead of copying it.""" + """Hand weights to a sampler on the same device by mapping memory instead of copying it.""" def __init__(self, bucket_size: int = 512 << 20, **kwargs) -> None: # Smaller default than the NCCL engine's 3 GB: a bigger bucket buys nothing when the transfer @@ -65,21 +66,22 @@ def __init__(self, bucket_size: int = 512 << 20, **kwargs) -> None: self.socket = None self._context = None self._handle = None + self._shm = None # Receiver side: the mapping of the sender's buffer, kept across buckets. Re-mapping per # bucket is what makes device memory appear to grow during a sync. self._mapped: torch.Tensor | None = None + self._mapped_shms = [] self._mapped_signature = None # ── rendezvous ─────────────────────────────────────────────────────── @staticmethod def endpoint() -> str: - """The socket both peers derive independently from the GPU they share. + """The socket both peers derive independently from the device they share. - The device's UUID rather than its index: under Ray each role sees its own GPU as index 0, so - indices collide across ranks while UUIDs do not. + The platform helper obtains the physical device UUID for the current local device. """ - uuid = str(torch.cuda.get_device_properties(torch.cuda.current_device()).uuid) + uuid = str(Platform.get_vllm_device_uuid(Torch.get_current_device())) return f'ipc:///tmp/twinkle-colocate-{uuid}.sock' def prepare(self) -> dict[str, Any]: @@ -165,9 +167,17 @@ def finalize(self): path = self.endpoint().removeprefix('ipc://') if os.path.exists(path): os.unlink(path) + if self._shm is not None: + self.send_buf = None + self._shm.close() + self._shm.unlink() + self._shm = None self.send_buf = None self._handle = None self._mapped = None + for shm in self._mapped_shms: + shm.close() + self._mapped_shms.clear() self._mapped_signature = None self.rank = None @@ -178,9 +188,29 @@ def _ensure_buffer(self, min_size: int) -> None: if self.send_buf is not None and self.send_buf.numel() >= min_size: return size = max(self.bucket_size, min_size) - self.send_buf = torch.empty(size, dtype=torch.uint8, device=torch.cuda.current_device()) + platform = Platform.get_platform() + if platform.device_prefix() == 'npu' and not platform.is_ipc_supported(): + from multiprocessing import shared_memory + + if self._shm is not None: + self.send_buf = None + self._shm.close() + self._shm.unlink() + self._shm = None + self._shm = shared_memory.SharedMemory(create=True, size=size) + self.send_buf = torch.frombuffer(self._shm.buf, dtype=torch.uint8, count=size) + self._handle = {'name': self._shm.name, 'size': size} + return + + self.send_buf = torch.empty( + size, + dtype=torch.uint8, + device=f'{platform.device_prefix()}:{Torch.get_current_device()}', + ) # One handle per buffer, reused for every bucket: the buffer is refilled, not reallocated, so # the mapping stays valid and the receiver can keep it. + if platform.device_prefix() == 'npu': + import torch_npu # noqa: F401 from torch.multiprocessing.reductions import reduce_tensor self._handle = reduce_tensor(self.send_buf) @@ -223,7 +253,7 @@ async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, def _flush(self, bucket_meta: list[dict], is_last: bool) -> None: """Publish the filled part of the buffer and wait until the receiver is done with it.""" # The copies above are non_blocking; without this the receiver could map bytes not yet written. - torch.cuda.synchronize() + Torch.synchronize() self.socket.send(pickle.dumps({'handle': self._handle, 'bucket_meta': bucket_meta, 'is_last': is_last})) # The receiver copies out of this buffer, so it must say so before we overwrite it. self.socket.recv() @@ -245,7 +275,7 @@ async def receive_weights(self) -> AsyncGenerator[tuple[str, torch.Tensor], None yield meta['name'], buffer[start:start + nbytes].view(meta['dtype']).view(meta['shape']) # Consumers copy with non_blocking=True, so the acknowledgement has to wait for the copies # and not merely for the loop above. - torch.cuda.synchronize() + Torch.synchronize() self.socket.send(b'ack') if message['is_last']: break @@ -259,13 +289,25 @@ def _map(self, handle) -> torch.Tensor: signature = self._handle_signature(handle) if self._mapped is not None and signature == self._mapped_signature: return self._mapped - from torch.multiprocessing.reductions import rebuild_cuda_tensor + if isinstance(handle, dict): + from multiprocessing import shared_memory + + mapped_shm = shared_memory.SharedMemory(name=handle['name']) + self._mapped_shms.append(mapped_shm) + self._mapped = torch.frombuffer( + mapped_shm.buf, + dtype=torch.uint8, + count=handle['size'], + ) + self._mapped_signature = signature + return self._mapped + func, args = handle args = list(args) - # Both peers see the shared GPU as their own device 0, but be explicit rather than trust the - # index the sender happened to record. - args[6] = torch.cuda.current_device() - self._mapped = func(*args) if callable(func) else rebuild_cuda_tensor(*args) + if Platform.device_prefix() == 'npu': + import torch_npu # noqa: F401 + args[6] = Torch.get_current_device() + self._mapped = func(*args) self._mapped_signature = signature return self._mapped @@ -276,6 +318,8 @@ def _handle_signature(handle) -> tuple: Locally implemented rather than shared with the sampler's worker extension, which has the same helper: the sampler imports this package, so importing it back would be circular. """ + if isinstance(handle, dict): + return tuple(handle.items()) _, args = handle return tuple((type(v).__name__, bytes(v) if isinstance(v, (bytes, bytearray)) else v) for v in args if isinstance(v, (bytes, bytearray, int, float, bool, str)) or v is None) diff --git a/src/twinkle/checkpoint_engine/manager.py b/src/twinkle/checkpoint_engine/manager.py index 9f7a24666..cdf30a0d8 100644 --- a/src/twinkle/checkpoint_engine/manager.py +++ b/src/twinkle/checkpoint_engine/manager.py @@ -127,8 +127,6 @@ def decide_backend_engine( platform_name = Platform.get_platform(platform).__name__ if mode == 'colocate': - if platform_name != 'GPU': - raise NotImplementedError("mode='colocate' currently requires the GPU platform.") from twinkle.checkpoint_engine import IPCCheckpointEngine return IPCCheckpointEngine if mode != 'standalone': @@ -166,8 +164,8 @@ def sync_weights(self, merge_and_sync=True): self._sync_weights_naive(merge_and_sync) return - is_master = [True] + [False] * (self.model.device_mesh.world_size - 1) - model_metadata = self.model.prepare_checkpoint_engine(is_master) + model_metadata = self.model.prepare_checkpoint_engine([True] + + [False] * (self.model.device_mesh.world_size - 1)) self.sampler.prepare_checkpoint_engine(False) model_kwargs, sampler_kwargs = self.backend_cls.build_topology( self.model.device_mesh.world_size, diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index b1e1790de..dc9825429 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -51,7 +51,7 @@ class VLLMEngine(BaseSamplerEngine): This engine uses vLLM v1's AsyncLLM and supports: - Tinker-compatible sample() API with logprobs - Multi-tenant LoRA adapters for client-server mode - - Weight synchronization via load_weights (colocated) or CUDA IPC + - Weight synchronization via load_weights (colocated) or device IPC - Sleep/wake_up for GPU memory management in colocated training Deployment scenarios: @@ -109,10 +109,10 @@ def __init__( # ``list_loras()`` per request. self._synced_lora_request: Optional[Any] = None - # Long-lived CUDA IPC bucket reused across all update_weights() + # Long-lived device IPC bucket reused across all update_weights() # calls. Allocating a new IPC buffer (and hence a new IPC handle) - # per sync forces every worker to create a new CUDA IPC mapping via - # ``rebuild_cuda_tensor`` because PyTorch's ``shared_cache`` cannot + # per sync forces every worker to create a new device IPC mapping via + # the reducer callable because PyTorch's ``shared_cache`` cannot # hit on unseen storage handles. The driver reclaims those mappings # lazily, which is the root cause of the slow GPU memory drift we # observed under frequent LoRA syncs. By pinning a single buffer @@ -578,14 +578,14 @@ async def update_weights( bucket_size_mb: int = 2048, **kwargs, ) -> None: - """Update model weights via ZMQ + CUDA IPC to worker extension. + """Update model weights via ZMQ + device IPC to worker extension. Accepts **either** a ``dict[str, Tensor]`` (legacy) **or** an async generator / sync generator of ``(name, tensor)`` pairs (streaming). The streaming path avoids accumulating a full model copy on GPU: tensors are consumed one-by-one from the generator, copied into a - GPU IPC bucket, and flushed to the vLLM worker subprocess when the + device IPC bucket, and flushed to the vLLM worker subprocess when the bucket is full. Args: @@ -621,15 +621,19 @@ async def _sync_iter(): weight_aiter = _sync_iter() - # Peek first tensor to detect device (GPU → IPC, CPU → SHM). + # Peek first tensor to detect device (supported accelerator → IPC, CPU → SHM). try: first_name, first_tensor = await weight_aiter.__anext__() except StopAsyncIteration: logger.warning('update_weights called with empty weights') return - use_gpu_ipc = first_tensor.is_cuda - use_shm = not use_gpu_ipc + use_device_ipc = first_tensor.is_cuda + if first_tensor.device.type == 'npu': + from twinkle.utils.platforms import NPU + + use_device_ipc = NPU.is_ipc_supported() + use_shm = not use_device_ipc # Use a per-sync unique IPC endpoint to avoid cross-actor collisions # when multiple sampler actors share the same device UUID. @@ -650,13 +654,16 @@ async def _sync_iter(): buffer = None shm = None - if use_gpu_ipc: + if use_device_ipc: + if first_tensor.device.type == 'npu': + # torch_npu registers the NPU reducer used by reduce_tensor. + import torch_npu # noqa: F401 from torch.multiprocessing.reductions import reduce_tensor # Reuse a long-lived IPC bucket whenever the requested size # fits. The handle is produced once and shipped to every # subsequent sync so each worker's ``shared_cache`` stays warm - # and no new CUDA IPC mapping is created per sync. + # and no new device IPC mapping is created per sync. need_realloc = ( self._ipc_buffer is None or self._ipc_buffer_size < bucket_size or self._ipc_buffer.device != first_tensor.device) @@ -714,7 +721,7 @@ def _zmq_send_recv(payload, where: str): )) # Send IPC/SHM handle, wait for worker ready (non-blocking) - handle_payload = ipc_handle if use_gpu_ipc else {'name': shm_name, 'size': bucket_size} + handle_payload = ipc_handle if use_device_ipc else {'name': shm_name, 'size': bucket_size} await loop.run_in_executor(None, _zmq_send_recv, handle_payload, 'handle handshake') # Stream weights into buckets and send to worker @@ -821,7 +828,7 @@ async def _flush_bucket(is_last: bool) -> None: elapsed = time.time() - start_time mode = 'LoRA' if base_sync_done and peft_config else 'base' logger.info(f'Updated {n_weights} {mode} weights via ' - f"{'IPC' if use_gpu_ipc else 'SHM'} in {elapsed:.2f}s") + f"{'IPC' if use_device_ipc else 'SHM'} in {elapsed:.2f}s") async def shutdown(self) -> None: """Shutdown the vLLM engine and release all resources. diff --git a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py index 2ed5f5227..0a4ec3e56 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py @@ -45,26 +45,20 @@ def set_death_signal(): def _rebuild_ipc(handle, device_id: Optional[int] = None) -> torch.Tensor: - """Rebuild CUDA tensor from IPC handle.""" - from torch.multiprocessing.reductions import rebuild_cuda_tensor - + """Rebuild an accelerator tensor from an IPC reducer handle.""" func, args = handle list_args = list(args) if device_id is not None: list_args[6] = device_id - - if callable(func): - return func(*list_args) - else: - return rebuild_cuda_tensor(*list_args) + return func(*list_args) def _ipc_handle_signature(handle) -> Optional[tuple]: - """Derive a stable signature for a CUDA IPC handle. + """Derive a stable signature for an accelerator IPC handle. ``reduce_tensor`` returns ``(func, args)`` where ``args`` contains the - CUDA IPC storage handle bytes, storage size, ref-counter handle, etc. - Two handles are equivalent (i.e. map the same CUDA memory region) when + IPC storage handle bytes, storage size, ref-counter handle, etc. + Two handles are equivalent (i.e. map the same device memory region) when these inner fields match. We hash only the parts that are picklable and comparable to avoid accidental mismatches due to local objects. """ @@ -127,12 +121,12 @@ def update_weights_from_ipc( use_shm: bool = False, zmq_handle: Optional[str] = None, ) -> None: - """Receive and load weights via ZMQ + CUDA IPC/SHM. + """Receive and load weights via ZMQ + device IPC/SHM. Called via ``collective_rpc("update_weights_from_ipc", ...)`` from :meth:`VLLMEngine.update_weights`. The VLLMEngine sends weights - in buckets over a ZMQ REQ/REP channel backed by CUDA IPC (GPU - tensors) or shared memory (CPU tensors). + in buckets over a ZMQ REQ/REP channel backed by device IPC + (accelerator tensors) or shared memory (CPU tensors). For TP > 1, only TP rank 0 communicates with the VLLMEngine over ZMQ. It broadcasts the IPC handle and bucket metadata to other @@ -142,7 +136,7 @@ def update_weights_from_ipc( Args: peft_config: If provided with base_sync_done, loads as LoRA. base_sync_done: If True and peft_config, replaces existing LoRA. - use_shm: If True, use shared memory instead of CUDA IPC. + use_shm: If True, use shared memory instead of device IPC. zmq_handle: Optional ZMQ IPC endpoint. If None, uses _get_zmq_handle(). """ import torch.distributed as dist @@ -196,6 +190,10 @@ def _broadcast_obj(obj): # ── Step 2: Receive and broadcast IPC/SHM handle ── buffer, shm = None, None + if not use_shm and self.device.type == 'npu': + # Register the NPU reducer before recv_pyobj() unpickles its callable. + import torch_npu # noqa: F401 + if is_driver: try: comm_metadata = socket.recv_pyobj() @@ -210,9 +208,9 @@ def _broadcast_obj(obj): if not use_shm: handle = comm_metadata # All TP ranks rebuild the IPC buffer from the same handle. - # CUDA IPC allows any process on the same node to map the memory. + # Device IPC allows any process on the same node to map the memory. # Reuse a cached buffer across syncs when the sender reuses the - # same IPC handle: this avoids creating a fresh CUDA IPC mapping + # same IPC handle: this avoids creating a fresh device IPC mapping # per sync, which the driver releases lazily and is the root # cause of the apparent GPU memory growth under frequent syncs. handle_signature = _ipc_handle_signature(handle) diff --git a/src/twinkle/utils/platforms/npu.py b/src/twinkle/utils/platforms/npu.py index de15707f6..fb76f1293 100644 --- a/src/twinkle/utils/platforms/npu.py +++ b/src/twinkle/utils/platforms/npu.py @@ -1,11 +1,14 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import hashlib import os +import platform import re import socket import subprocess from typing import Optional +from packaging import version + from .base import Platform # ref: https://www.hiascend.com/document/detail/zh/canncommercial/83RC1/maintenref/envvar/envref_07_0144.html @@ -16,7 +19,6 @@ # NPU-side socket port pool used by HCCL for device communication channels. _HCCL_NPU_SOCKET_PORT_RANGE_ENV = 'HCCL_NPU_SOCKET_PORT_RANGE' - def _derive_hccl_socket_env_defaults(master_port: int) -> dict: """Derive deterministic default HCCL socket env values from master_port.""" # Keep values stable per job and spread jobs across non-overlapping ranges. @@ -99,6 +101,61 @@ def ensure_npu_backend() -> None: class NPU(Platform): + @staticmethod + def is_ipc_supported() -> bool: + """Return whether HDK and CANN meet the NPU device-IPC requirement.""" + try: + result = subprocess.run( + ['npu-smi', 'info', '-t', 'board', '-i', '1'], capture_output=True, text=True, check=True) + except subprocess.CalledProcessError: + visible_devices = (os.environ.get('ASCEND_VISIBLE_DEVICES') + or os.environ.get('ASCEND_RT_VISIBLE_DEVICES')) + if not visible_devices: + raise + device_id = int(visible_devices.split(',')[0]) + try: + result = subprocess.run( + ['npu-smi', 'info', '-t', 'board', '-i', str(device_id)], + capture_output=True, + text=True, + check=True, + ) + except subprocess.CalledProcessError: + result = subprocess.run( + ['npu-smi', 'info', '-t', 'board', '-i', str(device_id // 2)], + capture_output=True, + text=True, + check=True, + ) + + software_version = next( + (line.split(':', 1)[1].strip().lower() + for line in result.stdout.splitlines() if 'Software Version' in line), + None, + ) + if software_version is None: + raise RuntimeError('Could not find Software Version in npu-smi output') + + ascend_home = os.environ.get('ASCEND_HOME_PATH', '/usr/local/Ascend/ascend-toolkit/latest') + info_file = os.path.join(ascend_home, f'{platform.machine()}-linux', 'ascend_toolkit_install.info') + with open(info_file) as info: + cann_version = next( + (line.split('=', 1)[1].strip().lower() for line in info if line.startswith('version=')), + None, + ) + if cann_version is None: + raise RuntimeError('Could not find version in CANN toolkit info file') + + pattern = r'(\d+\.\d+(?=\.t))|(\d+\.\d+(?:\.(?:rc\d+|\d+))?)' + software_match = re.match(pattern, software_version) + cann_match = re.match(pattern, cann_version) + if software_match is None or cann_match is None: + raise RuntimeError(f'Invalid NPU versions: HDK={software_version}, CANN={cann_version}') + software_base = software_match.group(1) or software_match.group(2) + cann_base = cann_match.group(1) or cann_match.group(2) + return (version.parse(software_base) >= version.parse('25.3.rc1') + and version.parse(cann_base) >= version.parse('8.3.rc1')) + @staticmethod def visible_device_env(): # Ascend runtime uses ASCEND_RT_VISIBLE_DEVICES. diff --git a/tests/checkpoint_engine/test_ipc_checkpoint_engine_device_neutral.py b/tests/checkpoint_engine/test_ipc_checkpoint_engine_device_neutral.py new file mode 100644 index 000000000..396cee822 --- /dev/null +++ b/tests/checkpoint_engine/test_ipc_checkpoint_engine_device_neutral.py @@ -0,0 +1,149 @@ +"""CPU-only coverage for the device-neutral parts of the IPC checkpoint engine.""" + +from unittest.mock import Mock + +import pytest +import torch + +import twinkle.checkpoint_engine.ipc_checkpoint_engine as ipc_module +from twinkle.checkpoint_engine import IPCCheckpointEngine + + +def test_endpoint_uses_platform_uuid_without_touching_cuda(monkeypatch): + """Endpoint construction must not query the current CUDA device.""" + monkeypatch.setattr(torch.cuda, 'current_device', Mock(side_effect=AssertionError('CUDA initialized'))) + monkeypatch.setattr(torch.cuda, 'get_device_properties', Mock(side_effect=AssertionError('CUDA initialized'))) + monkeypatch.setattr(ipc_module.Torch, 'get_current_device', lambda: 0) + monkeypatch.setattr( + ipc_module.Platform, + 'get_vllm_device_uuid', + staticmethod(lambda device_id=0, platform=None: f'uuid-{device_id}'), + ) + + assert IPCCheckpointEngine.endpoint() == 'ipc:///tmp/twinkle-colocate-uuid-0.sock' + + +def test_map_uses_reducer_callable_and_receiver_device(monkeypatch): + calls = [] + + def rebuild(*args): + calls.append(args) + return 'mapped' + + monkeypatch.setattr(ipc_module.Platform, 'device_prefix', staticmethod(lambda: 'cuda')) + monkeypatch.setattr(ipc_module.Torch, 'get_current_device', lambda: 3) + sender_args = [None, None, None, None, None, None, 17] + + engine = IPCCheckpointEngine() + assert engine._map((rebuild, sender_args)) == 'mapped' + assert calls[0][6] == 3 + assert sender_args[6] == 17 + + +def test_vllm_worker_uses_reducer_callable_and_receiver_device(): + from twinkle.sampler.vllm_sampler.vllm_worker_extension import _rebuild_ipc + + calls = [] + + def rebuild(*args): + calls.append(args) + return 'mapped' + + sender_args = [None, None, None, None, None, None, 17] + + assert _rebuild_ipc((rebuild, sender_args), device_id=3) == 'mapped' + assert calls[0][6] == 3 + assert sender_args[6] == 17 + + +def test_buffer_and_flush_use_framework_device_and_sync_helpers(monkeypatch): + class CPU: + @staticmethod + def device_prefix(): + return 'cpu' + + monkeypatch.setattr(ipc_module.Platform, 'get_platform', staticmethod(lambda platform=None: CPU)) + monkeypatch.setattr(ipc_module.Torch, 'get_current_device', lambda: 0) + + from torch.multiprocessing import reductions + + monkeypatch.setattr(reductions, 'reduce_tensor', lambda tensor: (None, [None])) + engine = IPCCheckpointEngine(bucket_size=8) + engine._ensure_buffer(8) + assert engine.send_buf.device.type == 'cpu' + + class Socket: + def send(self, payload): + self.payload = payload + + def recv(self): + return b'ack' + + engine.socket = Socket() + synchronize = Mock() + monkeypatch.setattr(ipc_module.Torch, 'synchronize', synchronize) + engine._flush([], is_last=True) + synchronize.assert_called_once_with() + + +def test_npu_without_device_ipc_uses_shared_memory_and_closes_it(monkeypatch): + class NPU: + @staticmethod + def device_prefix(): + return 'npu' + + @staticmethod + def is_ipc_supported(): + return False + + monkeypatch.setattr(ipc_module.Platform, 'get_platform', staticmethod(lambda platform=None: NPU)) + monkeypatch.setattr(ipc_module.Torch, 'get_current_device', lambda: 0) + + sender = IPCCheckpointEngine(bucket_size=8) + sender._ensure_buffer(8) + handle = sender._handle + assert set(handle) == {'name', 'size'} + assert sender.send_buf.device.type == 'cpu' + + receiver = IPCCheckpointEngine() + mapped = receiver._map(handle) + assert mapped.device.type == 'cpu' + assert mapped.numel() == 8 + + sender.finalize() + del mapped + receiver.finalize() + + from multiprocessing import shared_memory + + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=handle['name']) + + +def test_shared_memory_growth_keeps_the_previous_receiver_mapping_alive(monkeypatch): + class NPU: + @staticmethod + def device_prefix(): + return 'npu' + + @staticmethod + def is_ipc_supported(): + return False + + monkeypatch.setattr(ipc_module.Platform, 'get_platform', staticmethod(lambda platform=None: NPU)) + + sender = IPCCheckpointEngine(bucket_size=8) + receiver = IPCCheckpointEngine() + sender._ensure_buffer(8) + first = receiver._map(sender._handle) + first[0] = 7 + + sender._ensure_buffer(16) + second = receiver._map(sender._handle) + + assert first[0].item() == 7 + assert second.numel() == 16 + + del first, second + sender.finalize() + receiver.finalize() diff --git a/tests/utils/test_npu_ipc_support.py b/tests/utils/test_npu_ipc_support.py new file mode 100644 index 000000000..f61d1594a --- /dev/null +++ b/tests/utils/test_npu_ipc_support.py @@ -0,0 +1,39 @@ +from types import SimpleNamespace + +import pytest + +from twinkle.utils.platforms import npu + + +@pytest.mark.parametrize( + ('software_version', 'cann_version', 'expected'), + [ + ('25.3.rc1.2', '8.3.rc1', True), + ('25.5.t3.b001', '8.3.0', True), + ('25.2.0', '8.3.rc1', False), + ('25.3.rc1', '8.2.0', False), + ], +) +def test_npu_ipc_version_gate(monkeypatch, tmp_path, software_version, cann_version, expected): + monkeypatch.setattr( + npu.subprocess, + 'run', + lambda *args, **kwargs: SimpleNamespace(stdout=f'Software Version : {software_version}\n'), + ) + monkeypatch.setattr(npu.platform, 'machine', lambda: 'x86_64') + cann_dir = tmp_path / 'x86_64-linux' + cann_dir.mkdir() + (cann_dir / 'ascend_toolkit_install.info').write_text(f'version={cann_version}\n') + monkeypatch.setenv('ASCEND_HOME_PATH', str(tmp_path)) + + assert npu.NPU.is_ipc_supported() is expected + + +def test_npu_ipc_detection_errors_are_not_hidden(monkeypatch): + def fail(*args, **kwargs): + raise RuntimeError('npu-smi failed') + + monkeypatch.setattr(npu.subprocess, 'run', fail) + + with pytest.raises(RuntimeError, match='npu-smi failed'): + npu.NPU.is_ipc_supported()