From 55fd5522d521616f303692e4ce4ea1bb24718fc6 Mon Sep 17 00:00:00 2001 From: Siju Samuel Date: Sat, 25 Jul 2026 03:00:10 +0000 Subject: [PATCH 1/2] feat: add Intel XPU support with zero-copy SYCL ipc_memory weight update Add Intel XPU as a first-class device backend alongside CUDA and NPU. Device abstraction (device_utils.py): detect xpu, map to the xccl distributed backend, and gate capabilities per backend -- in-place host pinning (cudaHostRegister) and Mooncake device P2P stay CUDA/NPU-only, while cross-process device-tensor IPC is supported on XPU via a native extension. Weight transport (transport.py): the broadcast path shares a device buffer with the colocated inference worker. PyTorch has no XPU tensor IPC (reduce_tensor raises '_share_fd_: only available on CPU'), so a WeightTransport seam selects the CUDA/NPU reduce_tensor path or a native XPU implementation. XPU IPC (xpu_ipc/): implemented on SYCL ext::oneapi::experimental::ipc_memory, which torch's own libsycl (oneAPI >= 2026.0, libsycl.so.9) exports. A SYCL IPC handle is a self-contained, portable byte blob -- verified cross-process on Arc B60 -- so it needs no out-of-band dma-buf fd transfer and no sub-allocation offset handling: the whole handle rides the existing ZMQ channel exactly like CUDA's reduce_tensor tuple. The -fsycl extension coexists with torch's single SYCL runtime in-process. Only the broadcast update method is supported on XPU; P2P is rejected with a clear error (Mooncake has no Level Zero backend for XPU device memory). Robustness: the exported IPC handle is released in an outer try/finally so an early failure (export, ZMQ bind, or first send) can no longer leak it, while the collective barrier stays in the inner finally to avoid deadlocking peers on an asymmetric failure. The P2P store is skipped entirely on backends that do not support device P2P (e.g. XPU) instead of being eagerly initialized and only guarded against ImportError. icpx discovery falls back to PATH (via shutil.which) for oneAPI layouts outside /opt or a sourced setvars.sh that does not export CMPLR_ROOT. Tests: CPU-only coverage for device dispatch, the transport seam, the P2P guard, and the XPU parity paths (portable-handle contract, xccl backend selection, in-place-pin disable, detach-on-early-failure, icpx PATH fallback); hardware-gated tests for the SYCL IPC roundtrip, interior-pointer offset, and the full cross-process broadcast. --- README.md | 22 ++ checkpoint_engine/device_utils.py | 53 ++++- checkpoint_engine/distributed/base.py | 7 +- checkpoint_engine/ps.py | 293 +++++++++++++++---------- checkpoint_engine/transport.py | 124 +++++++++++ checkpoint_engine/worker.py | 44 ++-- checkpoint_engine/xpu_ipc/__init__.py | 163 ++++++++++++++ checkpoint_engine/xpu_ipc/sycl_ipc.cpp | 103 +++++++++ pyproject.toml | 6 +- tests/test_device_manager.py | 174 +++++++++++++++ tests/test_p2p_guard.py | 45 ++++ tests/test_transport.py | 110 ++++++++++ tests/test_xpu_ipc.py | 154 +++++++++++++ tests/test_xpu_parity.py | 225 +++++++++++++++++++ 14 files changed, 1387 insertions(+), 136 deletions(-) create mode 100644 checkpoint_engine/transport.py create mode 100644 checkpoint_engine/xpu_ipc/__init__.py create mode 100644 checkpoint_engine/xpu_ipc/sycl_ipc.cpp create mode 100644 tests/test_device_manager.py create mode 100644 tests/test_p2p_guard.py create mode 100644 tests/test_transport.py create mode 100644 tests/test_xpu_ipc.py create mode 100644 tests/test_xpu_parity.py diff --git a/README.md b/README.md index a656c61..f449b3d 100644 --- a/README.md +++ b/README.md @@ -75,6 +75,28 @@ Use the flexible P2P implementation, notice this will install `mooncake-transfer pip install 'checkpoint-engine[p2p]' ``` +### Intel XPU (build from source) + +Intel XPU is supported for the **broadcast** update path. The cross-process weight handoff uses a native SYCL `ipc_memory` extension (the XPU counterpart of CUDA IPC) that is JIT-compiled at runtime. Install from source since XPU support is not yet in the released package. P2P is not supported on XPU (Mooncake has no Level Zero backend for XPU device memory). + +Requirements: +- Intel XPU build of PyTorch — `torch.xpu.is_available()` returns `True`, `torch>=2.9` (for the device `.uuid` property). See the [PyTorch XPU install guide](https://pytorch.org/docs/stable/notes/get_start_xpu.html); this is not the default PyPI `torch`. +- Intel oneAPI 2026.0+ providing the `icpx` compiler with SYCL `ipc_memory` support. Needed at runtime (first weight update), not at `pip install` time. + +```Bash +git clone https://github.com/MoonshotAI/checkpoint-engine.git +cd checkpoint-engine +pip install -e . # no [p2p] extra on XPU +``` + +Make `icpx` discoverable in the runtime environment, either by sourcing oneAPI (`source /opt/intel/oneapi/setvars.sh`) or by setting `CMPLR_ROOT`. If neither is set, `icpx` is auto-detected under `/opt/intel/oneapi/compiler/*/bin`. The extension then builds automatically on first use; `ParameterServer` also prebuilds it at startup so the one-time compile stays out of the update window. + +Verify the build and the IPC path on the target machine: + +```Bash +pytest tests/test_xpu_ipc.py # hardware-gated; skipped without an Intel GPU + buildable extension +``` + ## Getting Started Prepare an H800 or H20 machine with 8 GPUs with vLLM. Be sure to include [/collective_rpc API endpoint](https://github.com/vllm-project/vllm/commit/f7cf5b512ee41f36613deb2471a44de5f304f70d) commit (available in main branch) since checkpoint-engine will use this endpoint to update weights. vLLM version `v0.10.2` is fully tested and recommended. diff --git a/checkpoint_engine/device_utils.py b/checkpoint_engine/device_utils.py index be5b81f..c866e13 100644 --- a/checkpoint_engine/device_utils.py +++ b/checkpoint_engine/device_utils.py @@ -1,4 +1,5 @@ import ctypes +import gc import os import re import socket @@ -209,9 +210,20 @@ def _is_torch_npu_available(self) -> bool: except ImportError: return False + def _is_torch_xpu_available(self) -> bool: + try: + if hasattr(torch, "xpu") and callable(getattr(torch.xpu, "is_available", None)): + return torch.xpu.is_available() + else: + return False + except Exception: # noqa: BLE001 + return False + def _detect_device_type(self) -> str: if self._is_torch_npu_available(): return "npu" + elif self._is_torch_xpu_available(): + return "xpu" elif torch.cuda.is_available(): return "cuda" else: @@ -222,6 +234,8 @@ def _setup_device_module(self): import torch_npu self.device_module = torch_npu.npu + elif self.device_type == "xpu": + self.device_module = torch.xpu elif self.device_type == "cuda": self.device_module = torch.cuda else: @@ -231,6 +245,8 @@ def _setup_device_module(self): def backend(self) -> str: if self.device_type == "npu": return "hccl" + elif self.device_type == "xpu": + return "xccl" elif self.device_type == "cuda": return "nccl" else: @@ -240,7 +256,7 @@ def backend(self) -> str: def transfer_engine_protocol(self) -> str: if self.device_type == "npu": return "ascend_direct" - elif self.device_type == "cuda": + elif self.device_type in ("cuda", "xpu"): if has_efa_pci(): return "efa" else: @@ -255,3 +271,38 @@ def rdma_device(self, rank: int) -> str: return _get_my_rdma_device(rank, self.device_module.device_count(), _get_rdma_devices()) else: raise TypeError("The current transfer engine protocol is not supported") + + def ipc_collect(self) -> None: + """Reclaim memory held by stale IPC handles where the backend supports it (no-op otherwise).""" + fn = getattr(self.device_module, "ipc_collect", None) + if callable(fn): + fn() + + def supports_inplace_pin(self) -> bool: + """Whether in-place host-memory pinning (cudaHostRegister) is available -- CUDA only.""" + return self.device_type == "cuda" + + def supports_device_ipc(self) -> bool: + """Whether cross-process IPC of *device* tensors works for this backend. + + CUDA/NPU use ``torch.multiprocessing.reductions``; XPU uses the native SYCL + ``ipc_memory`` extension, available when it can be built (oneAPI >= 2026.0). + """ + if self.device_type in ("cuda", "npu"): + return True + if self.device_type == "xpu": + from checkpoint_engine import xpu_ipc + + return xpu_ipc.is_available() + return False + + def supports_device_p2p(self) -> bool: + """Whether P2P (Mooncake) transfer of *device* memory works for this backend (CUDA/NPU only).""" + return self.device_type in ("cuda", "npu") + + def host_empty_cache(self) -> None: + """Release cached pinned host memory (``_host_emptyCache`` on CUDA; else ``gc.collect``).""" + if self.device_type == "cuda": + torch._C._host_emptyCache() + else: + gc.collect() diff --git a/checkpoint_engine/distributed/base.py b/checkpoint_engine/distributed/base.py index 4299394..ea28af7 100644 --- a/checkpoint_engine/distributed/base.py +++ b/checkpoint_engine/distributed/base.py @@ -229,7 +229,12 @@ def use_backend(backend: str | None): "vllm_hccl": ".vllm_hccl.DistributedHccl", } if backend not in mapping: - raise ValueError(f"Unsupported custom backend: {backend}") + # XPU has no custom backend; leave custom_dist unset to use the default xccl TorchBackend. + raise ValueError( + f"Unsupported custom backend: {backend}. " + f"Supported custom backends: {sorted(mapping)}. " + "XPU is not supported here; leave custom_dist unset to use the default xccl backend." + ) module_path, class_name = mapping[backend].rsplit(".", 1) module = importlib.import_module(module_path, "checkpoint_engine.distributed") diff --git a/checkpoint_engine/ps.py b/checkpoint_engine/ps.py index 013a58f..1834110 100644 --- a/checkpoint_engine/ps.py +++ b/checkpoint_engine/ps.py @@ -11,7 +11,6 @@ import torch.distributed import zmq from loguru import logger -from torch.multiprocessing.reductions import reduce_tensor import checkpoint_engine.distributed as dist from checkpoint_engine.data_types import ( @@ -26,6 +25,7 @@ from checkpoint_engine.device_utils import DeviceManager, get_ip, npu_generate_uuid from checkpoint_engine.p2p_store import P2PStore from checkpoint_engine.pin_memory import _ALIGN_SIZE, _register_checkpoint +from checkpoint_engine.transport import build_transport if TYPE_CHECKING: @@ -53,7 +53,14 @@ def _get_physical_gpu_id(device_manager: DeviceManager, device_index: int | None if device_manager.device_type == "npu": return f"NPU-{npu_generate_uuid()}" else: - return f"GPU-{device_manager.device_module.get_device_properties(device_index).uuid!s}" + # CUDA and XPU both expose get_device_properties(idx).uuid. + props = device_manager.device_module.get_device_properties(device_index) + if not hasattr(props, "uuid"): + raise ValueError( + f"{device_manager.device_type} device properties do not expose a 'uuid' " + f"attribute; a newer PyTorch is required (xpu .uuid needs torch>=2.9)" + ) + return f"GPU-{props.uuid!s}" except AssertionError as e: raise ValueError(f"fail to get physical gpu id {device_index}") from e @@ -223,15 +230,39 @@ def __init__( # NPU transfer engine initialization requires prior set_device. device_index = self._local_rank self.device_manager.device_module.set_device(device_index) - try: - self._p2p_store = P2PStore(self.device_manager) - except ImportError as e: - logger.warning(f"[rank{self._rank}] fail to initialize p2p store due to {e}") + # P2P (Mooncake) transfer is only supported on CUDA/NPU; XPU has no Level Zero + # backend for device memory. Skip the store entirely on unsupported backends + # rather than eagerly initializing something that can never be used (and whose + # engine.initialize() may fail with more than just ImportError). + if self.device_manager.supports_device_p2p(): + try: + self._p2p_store = P2PStore(self.device_manager) + except ImportError as e: + logger.warning(f"[rank{self._rank}] fail to initialize p2p store due to {e}") + self._p2p_store = None + else: + logger.info( + f"[rank{self._rank}] p2p store disabled: not supported on device type " + f"'{self.device_manager.device_type}'" + ) self._p2p_store = None self._device_uuid = _get_physical_gpu_id(self.device_manager, device_index) self._rdma_device = None if self._p2p_store is None else self._p2p_store.device + # Build the JIT SYCL IPC extension now, so its multi-second compile is outside + # the first weight-update window. + if self.device_manager.device_type == "xpu": + from checkpoint_engine import xpu_ipc + + if xpu_ipc.prewarm(): + logger.info(f"[rank{self._rank}] XPU SYCL ipc_memory extension prebuilt") + else: + logger.warning( + f"[rank{self._rank}] XPU SYCL ipc_memory extension unavailable at init; " + "weight updates will fail until it can be built" + ) + master_addr = master_addr or os.getenv("MASTER_ADDR") assert master_addr, "master_addr is required" self._store = torch.distributed.TCPStore( @@ -297,7 +328,7 @@ def register_checkpoint( use_inplace_pin_memory: If True (default), allows inplace pin memory for /dev/shm/ safetensors files. This option is ignored when ``use_shared_memory_pool`` is True. """ - if self.device_manager.device_type != "cuda" and use_inplace_pin_memory: + if not self.device_manager.supports_inplace_pin() and use_inplace_pin_memory: logger.warning( f"[rank{self._rank}] Only cuda devices support in-place pin memory, set use_inplace_pin_memory to False" ) @@ -425,14 +456,8 @@ def _unpin(t: torch.Tensor): # we won't delete the memory pool if unpinning fails. del self._memory_pool[checkpoint_name] # see https://github.com/pytorch/pytorch/blob/31d5c675394705f8a6bc767f80ae14bf4f01246b/torch/csrc/cuda/Module.cpp#L2018 - # this works by using torch>=2.5.0 - if self.device_manager.device_type == "cuda": - torch._C._host_emptyCache() - else: - # torch._C._host_emptyCache() is not supported on NPU, so we call gc.collect() to empty host cache. - import gc - - gc.collect() + # this works by using torch>=2.5.0 on cuda; NPU/XPU fall back to gc.collect(). + self.device_manager.host_empty_cache() def gather_metas(self, checkpoint_name: str): """ @@ -730,12 +755,33 @@ def _update_per_bucket( assert len(self._current_global_parameter_metas) != 0, "parameter metas is empty" assert dist.is_initialized(), "process group is not initialized" + # The broadcast shares a device buffer with the colocated worker via cross-process + # device-tensor IPC (CUDA/NPU: torch.multiprocessing; XPU: native SYCL ipc_memory). + # Fail loudly here rather than with an opaque "_share_fd_: only available on CPU" + # deeper in the update. + if not self.device_manager.supports_device_ipc(): + raise RuntimeError( + f"[rank{self._rank}] weight update requires cross-process device-tensor IPC, which " + f"is not available for device type '{self.device_manager.device_type}' in this " + f"environment. On XPU this needs the native SYCL ipc_memory extension " + f"(an oneAPI 'icpx' compiler with SYCL ipc_memory support, i.e. oneAPI >= 2026.0)." + ) + p2p_update = False # if both ranks is None or [], it will use fully broadcast to update to all ranks if not ranks: logger.info(f"[rank{self._rank}] update checkpoint {checkpoint_name}") # if ranks is set, it will use p2p to update to the ranks else: + # The P2P path RDMA-transfers device memory via Mooncake, which has no Level Zero + # backend for XPU. Reject it clearly rather than failing inside p2p_store registration. + if not self.device_manager.supports_device_p2p(): + raise RuntimeError( + f"[rank{self._rank}] P2P weight update (ranks={ranks}) is not supported on " + f"device type '{self.device_manager.device_type}': the Mooncake transfer engine " + f"cannot register {self.device_manager.device_type} device memory. Use the " + f"broadcast update (leave ranks unset) instead." + ) assert self._p2p_store is not None, "p2p store is not initialized" assert ranks, "ranks should be set" @@ -780,113 +826,124 @@ def _update_per_bucket( self._p2p_store.register_named_tensors( {p2p_ipc_buffer_name: buffer if disable_h2d_buffer else h2d_buffer} ) - handle = reduce_tensor(buffer) - - buckets_by_receiver_rank: dict[int, list[H2DBucket]] = defaultdict(list) - max_len = 0 - for receiver_rank, _, bucket in buckets: - buckets_by_receiver_rank[receiver_rank].append(bucket) - if len(buckets_by_receiver_rank[receiver_rank]) > max_len: - max_len = len(buckets_by_receiver_rank[receiver_rank]) - - socket, socket_paths = self._bind_zmq_socket() - req_thread = threading.Thread( - target=req_func, - args=(socket_paths,), - ) - req_thread.start() - socket.send_pyobj(handle) - - gidx = 0 - ret_code = torch.zeros((), device=self.device_manager.device_type, dtype=torch.int64) - buffer_b: torch.Tensor | None = None + transport = build_transport(self.device_manager) + # Outer try guarantees the exported IPC handle is released even when export, + # the socket bind, or the first send below raises -- its finally does nothing + # but detach(). The collective barrier stays in the inner finally so an early + # single-rank failure here cannot deadlock peers that never reach the loop. try: - for i in range(max_len): - if i < len(receiver_rank_buckets) and not disable_h2d_buffer: - self._copy_to_buffer( - checkpoint_name, - receiver_rank_buckets[i][1], - h2d_buffer, - receiver_rank_buckets[i][0] if ranks else None, - ) - for receiver_rank, _buckets in buckets_by_receiver_rank.items(): - if i >= len(_buckets): - continue - bucket = _buckets[i] - alloc, reserved = ( - self.device_manager.device_module.memory_allocated() / 1024 / 1024, - self.device_manager.device_module.memory_reserved() / 1024 / 1024, - ) - self._logger_rank0( - f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} receiver_rank {receiver_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. " - f"Current device allocated {alloc:.2f} MB, " - f"reserved {reserved:.2f} MB." - ) - start = gidx % 2 * bucket_size - buffer_b: torch.Tensor = buffer[start : start + bucket.size] - if receiver_rank == self._rank: - if disable_h2d_buffer: - if p2p_update: - assert bucket == receiver_rank_buckets[i][1] - self._copy_to_buffer( - checkpoint_name, - bucket, - buffer_b, - receiver_rank_buckets[i][0] if p2p_update else None, - ) - else: - buffer_b.data.copy_(h2d_buffer[: bucket.size]) - dist.broadcast(buffer_b, src=receiver_rank, group=ranks_group) - resp = socket.recv() - if resp != b"": - msg = resp.decode("utf-8") - logger.error( - f"[rank{self._rank}] receive error response from rank {receiver_rank} for bucket {gidx} in checkpoint {checkpoint_name}: {msg}" - ) - ret_code.fill_(1) - dist.all_reduce(ret_code, op=torch.distributed.ReduceOp.SUM, group=ranks_group) - self.device_manager.device_module.synchronize() - if ret_code.item() != 0: - # quit early if any rank failed - socket.send_pyobj(RuntimeError("Some workers failed to update weights")) - raise RuntimeError("Failed to update weights due to remote errors") - socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size)) - gidx += 1 - - socket.recv() - device_mem = self.device_manager.device_module.mem_get_info() - logger.info( - f"[rank{self._rank}] weights broadcast done, device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, allocated memory: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, reserved memory: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" + handle = transport.export(buffer) + + buckets_by_receiver_rank: dict[int, list[H2DBucket]] = defaultdict(list) + max_len = 0 + for receiver_rank, _, bucket in buckets: + buckets_by_receiver_rank[receiver_rank].append(bucket) + if len(buckets_by_receiver_rank[receiver_rank]) > max_len: + max_len = len(buckets_by_receiver_rank[receiver_rank]) + + socket, socket_paths = self._bind_zmq_socket() + req_thread = threading.Thread( + target=req_func, + args=(socket_paths,), ) - # Notify worker to release handle - socket.send_pyobj(None) - socket.recv() - # Set to None in correct order (views first, then base tensors) - del buffer_b, h2d_buffer, buffer, handle - self.device_manager.device_module.synchronize() - gc.collect() - self.device_manager.device_module.ipc_collect() - self.device_manager.device_module.empty_cache() - self.device_manager.device_module.synchronize() + req_thread.start() + # The handle is self-contained for every transport, so one ZMQ send completes the handoff. + socket.send_pyobj(handle) - # Log actual memory usage - device_mem = self.device_manager.device_module.mem_get_info() - logger.info( - f"[rank{self._rank}] post-release: device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, " - f"allocated: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, " - f"reserved: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" - ) - # Notify worker to call post_hook - socket.send_pyobj(None) - socket.recv() + gidx = 0 + ret_code = torch.zeros((), device=self.device_manager.device_type, dtype=torch.int64) + buffer_b: torch.Tensor | None = None + try: + for i in range(max_len): + if i < len(receiver_rank_buckets) and not disable_h2d_buffer: + self._copy_to_buffer( + checkpoint_name, + receiver_rank_buckets[i][1], + h2d_buffer, + receiver_rank_buckets[i][0] if ranks else None, + ) + for receiver_rank, _buckets in buckets_by_receiver_rank.items(): + if i >= len(_buckets): + continue + bucket = _buckets[i] + alloc, reserved = ( + self.device_manager.device_module.memory_allocated() / 1024 / 1024, + self.device_manager.device_module.memory_reserved() / 1024 / 1024, + ) + self._logger_rank0( + f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} receiver_rank {receiver_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. " + f"Current device allocated {alloc:.2f} MB, " + f"reserved {reserved:.2f} MB." + ) + start = gidx % 2 * bucket_size + buffer_b: torch.Tensor = buffer[start : start + bucket.size] + if receiver_rank == self._rank: + if disable_h2d_buffer: + if p2p_update: + assert bucket == receiver_rank_buckets[i][1] + self._copy_to_buffer( + checkpoint_name, + bucket, + buffer_b, + receiver_rank_buckets[i][0] if p2p_update else None, + ) + else: + buffer_b.data.copy_(h2d_buffer[: bucket.size]) + dist.broadcast(buffer_b, src=receiver_rank, group=ranks_group) + resp = socket.recv() + if resp != b"": + msg = resp.decode("utf-8") + logger.error( + f"[rank{self._rank}] receive error response from rank {receiver_rank} for bucket {gidx} in checkpoint {checkpoint_name}: {msg}" + ) + ret_code.fill_(1) + dist.all_reduce( + ret_code, op=torch.distributed.ReduceOp.SUM, group=ranks_group + ) + self.device_manager.device_module.synchronize() + if ret_code.item() != 0: + # quit early if any rank failed + socket.send_pyobj(RuntimeError("Some workers failed to update weights")) + raise RuntimeError("Failed to update weights due to remote errors") + socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size)) + gidx += 1 + + socket.recv() + device_mem = self.device_manager.device_module.mem_get_info() + logger.info( + f"[rank{self._rank}] weights broadcast done, device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, allocated memory: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, reserved memory: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" + ) + # Notify worker to release handle + socket.send_pyobj(None) + socket.recv() + # Set to None in correct order (views first, then base tensors) + del buffer_b, h2d_buffer, buffer, handle + self.device_manager.device_module.synchronize() + gc.collect() + self.device_manager.ipc_collect() + self.device_manager.device_module.empty_cache() + self.device_manager.device_module.synchronize() + + # Log actual memory usage + device_mem = self.device_manager.device_module.mem_get_info() + logger.info( + f"[rank{self._rank}] post-release: device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, " + f"allocated: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, " + f"reserved: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" + ) + # Notify worker to call post_hook + socket.send_pyobj(None) + socket.recv() + finally: + req_thread.join() + dist.barrier(group=ranks_group) + socket.close() + if p2p_update: + self._p2p_store.unregister_named_tensors([p2p_ipc_buffer_name]) + + self.device_manager.device_module.empty_cache() finally: - req_thread.join() - dist.barrier(group=ranks_group) - socket.close() - if p2p_update: - self._p2p_store.unregister_named_tensors([p2p_ipc_buffer_name]) - - self.device_manager.device_module.empty_cache() + transport.detach() # we need this CLI entry point for compatibility with former versions diff --git a/checkpoint_engine/transport.py b/checkpoint_engine/transport.py new file mode 100644 index 0000000..0f97cdf --- /dev/null +++ b/checkpoint_engine/transport.py @@ -0,0 +1,124 @@ +"""Pluggable device-buffer handoff between the ParameterServer and the worker. + +The broadcast path shares a device buffer with the colocated worker. CUDA/NPU use +:class:`IpcWeightTransport` (``torch.multiprocessing`` CUDA IPC, wire-format +unchanged); XPU uses :class:`XpuIpcWeightTransport` (native SYCL ``ipc_memory``). +The handle is always a picklable, self-contained value, so the producer's +``export`` -> ZMQ ``send_pyobj`` -> consumer ``attach`` flow is identical for both; +each side calls ``detach`` on cleanup. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any + +import torch +from loguru import logger +from torch.multiprocessing.reductions import reduce_tensor + + +if TYPE_CHECKING: + from collections.abc import Callable + + from checkpoint_engine.device_utils import DeviceManager + + +def _rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) -> torch.Tensor: + func, args = handle + list_args = list(args) + if device_id is not None: + # the key is to change device id to the current device id + # in case two processes have different CUDA_VISIBLE_DEVICES + list_args[6] = device_id + return func(*list_args) + + +class WeightTransport(ABC): + """Hands a device buffer from the producer (ps) to the consumer (worker).""" + + @abstractmethod + def export(self, buffer: torch.Tensor) -> Any: + """Producer: return the picklable handle to send over ZMQ.""" + + @abstractmethod + def attach(self, handle: Any, device_id: int) -> torch.Tensor: + """Consumer: reconstruct the shared device buffer from ``handle``.""" + + def detach(self) -> None: + """Release IPC resources on either side. No-op by default.""" + + +class IpcWeightTransport(WeightTransport): + """CUDA/NPU zero-copy handoff via torch.multiprocessing CUDA IPC (unchanged).""" + + def export(self, buffer: torch.Tensor) -> Any: + return reduce_tensor(buffer) + + def attach(self, handle: Any, device_id: int) -> torch.Tensor: + assert isinstance(handle, tuple), f"expected reduce_tensor tuple, got {type(handle)}" + buffer = _rebuild_ipc(handle, device_id) + assert buffer.dtype == torch.uint8 + return buffer + + +class XpuIpcWeightTransport(WeightTransport): + """Intel XPU zero-copy handoff via native SYCL ``ipc_memory`` (portable byte blob).""" + + kind = "xpu_sycl" # wire tag: identifies this handle format to the consumer + + def __init__(self) -> None: + self._opened_ptr: int | None = None # consumer: the mapping to unmap + self._exported_ptr: int | None = None # producer: the retained exporter handle + + def export(self, buffer: torch.Tensor) -> Any: + from checkpoint_engine import xpu_ipc + + ptr = buffer.data_ptr() + handle_bytes = xpu_ipc.get_handle(ptr) + # Release only in detach(): freeing before the consumer opens can drop the + # fd under the level-zero-v2 UR adapter. + self._exported_ptr = ptr + return { + "kind": self.kind, + "handle_bytes": handle_bytes, + "nbytes": buffer.nbytes, + } + + def attach(self, handle: Any, device_id: int) -> torch.Tensor: + from checkpoint_engine import xpu_ipc + + assert isinstance(handle, dict) and handle.get("kind") == self.kind, ( + f"expected {self.kind} handle dict, got {type(handle)}" + ) + ptr = xpu_ipc.open_handle(handle["handle_bytes"], device_id) + self._opened_ptr = ptr + buffer = xpu_ipc.wrap_tensor(ptr, handle["nbytes"], device_id) + assert buffer.dtype == torch.uint8 + return buffer + + def detach(self) -> None: + # Consumer unmaps its opened pointer; producer releases its exported handle. + # At most one of the two is set on any given instance. + from checkpoint_engine import xpu_ipc + + if self._opened_ptr is not None: + try: + torch.xpu.synchronize() # no in-flight reads before unmapping + xpu_ipc.close_handle(self._opened_ptr) + except Exception as e: # noqa: BLE001 + logger.debug(f"xpu ipc close_handle failed during detach: {e}") + self._opened_ptr = None + if self._exported_ptr is not None: + try: + xpu_ipc.release_handle(self._exported_ptr) + except Exception as e: # noqa: BLE001 + logger.debug(f"xpu ipc release_handle failed during detach: {e}") + self._exported_ptr = None + + +def build_transport(device_manager: DeviceManager) -> WeightTransport: + """Select the weight transport for the current device backend.""" + if device_manager.device_type == "xpu": + return XpuIpcWeightTransport() + return IpcWeightTransport() diff --git a/checkpoint_engine/worker.py b/checkpoint_engine/worker.py index ea170ca..4e773ad 100644 --- a/checkpoint_engine/worker.py +++ b/checkpoint_engine/worker.py @@ -8,20 +8,24 @@ import zmq from checkpoint_engine.device_utils import DeviceManager, npu_generate_uuid +from checkpoint_engine.transport import ( + IpcWeightTransport, + WeightTransport, + XpuIpcWeightTransport, +) _WEIGHTS_TYPE = list[tuple[str, torch.Tensor]] -def _rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) -> torch.Tensor: - func, args = handle - list_args = list(args) - if device_id is not None: - # the key is to change device id to the current device id - # in case two processes have different CUDA_VISIBLE_DEVICES - list_args[6] = device_id - buffer = func(*list_args) - return buffer +def _transport_for_handle(handle: object) -> WeightTransport: + """Pick the consumer-side transport based on the handle wire format. + + CUDA/NPU send a ``reduce_tensor`` tuple; XPU sends a dict tagged with its kind. + """ + if isinstance(handle, dict) and handle.get("kind") == XpuIpcWeightTransport.kind: + return XpuIpcWeightTransport() + return IpcWeightTransport() class FlattenedTensorMetadata(TypedDict): @@ -59,10 +63,11 @@ def update_weights_from_ipc( socket.connect(zmq_handle) buffer: torch.Tensor | None = None device_manager = DeviceManager() + transport: WeightTransport | None = None try: - ipc_handle: tuple[Callable, tuple] = socket.recv_pyobj() - assert isinstance(ipc_handle, tuple) - buffer = _rebuild_ipc(ipc_handle, device_id) + ipc_handle = socket.recv_pyobj() + transport = _transport_for_handle(ipc_handle) + buffer = transport.attach(ipc_handle, device_id) assert buffer.dtype == torch.uint8 socket.send(b"") except Exception as e: @@ -91,10 +96,11 @@ def update_weights_from_ipc( device_manager.device_module.synchronize() released = True buffer = None - del ipc_handle + if transport is not None: + transport.detach() gc.collect() - device_manager.device_module.ipc_collect() + device_manager.ipc_collect() device_manager.device_module.empty_cache() device_manager.device_module.synchronize() socket.send(b"") @@ -119,6 +125,8 @@ def update_weights_from_ipc( finally: socket.close() del buffer + if transport is not None: + transport.detach() gc.collect() device_manager.device_module.empty_cache() @@ -147,6 +155,9 @@ def _device_uuid(self) -> str: return current_platform.get_device_uuid(self.device.index) elif current_platform.device_type == "npu": return f"NPU-{npu_generate_uuid()}" + elif current_platform.device_type == "xpu": + # Must match ps.py::_get_physical_gpu_id ("GPU-") for the ZMQ key to resolve. + return f"GPU-{torch.xpu.get_device_properties(self.device.index).uuid!s}" else: raise ValueError(f"Unsupported device type: {current_platform.device_type}") @@ -170,9 +181,10 @@ def update_weights_from_ipc(self, zmq_handles: dict[str, str]): The device UUID is platform-specific: - For CUDA: UUID from `current_platform.get_device_uuid()` - For NPU: Format "NPU-{generated_uuid}" + - For XPU: Format "GPU-{torch.xpu device uuid}" Raises: - ValueError: If the device type is not supported (not CUDA or NPU). + ValueError: If the device type is not supported (not CUDA, NPU, or XPU). AssertionError: If the device is not properly initialized. Note: @@ -185,6 +197,8 @@ def update_weights_from_ipc(self, zmq_handles: dict[str, str]): # vllm-ascend not init device if current_platform.device_type == "npu" and self.device is None: self.device = torch.device(f"npu:{self.local_rank}") + elif current_platform.device_type == "xpu" and self.device is None: + self.device = torch.device(f"xpu:{self.local_rank}") assert self.device is not None def _load_weights(weights: _WEIGHTS_TYPE): diff --git a/checkpoint_engine/xpu_ipc/__init__.py b/checkpoint_engine/xpu_ipc/__init__.py new file mode 100644 index 0000000..bb1e461 --- /dev/null +++ b/checkpoint_engine/xpu_ipc/__init__.py @@ -0,0 +1,163 @@ +"""Cross-process device-buffer IPC for Intel XPU via SYCL ``ipc_memory``. + +``sycl_ipc.cpp`` wraps ``ipc_memory`` (``get``/``open``/``close``), exported by +torch's own libsycl (oneAPI >= 2026.0); this module JIT-compiles it with +``-fsycl``. The handle is a self-contained portable byte blob (no dma-buf fd, no +offset to carry), so it rides the existing ZMQ channel like CUDA's +``reduce_tensor`` tuple -- see ``XpuIpcWeightTransport``. +""" + +from __future__ import annotations + +import functools +import glob +import os +import shutil +from pathlib import Path +from typing import TYPE_CHECKING + +from loguru import logger + + +if TYPE_CHECKING: + from types import ModuleType + + import torch + + +def _find_sycl_include_dir() -> str | None: + """Locate a directory containing .""" + candidates: list[str] = [] + root = os.getenv("CMPLR_ROOT") + if root: + candidates.append(os.path.join(root, "include")) + # Common oneAPI install layouts (versioned + `latest` symlink). + candidates += sorted(glob.glob("/opt/intel/oneapi/compiler/*/include"), reverse=True) + # Derive from the discovered compiler (/bin/icpx -> /include), which + # covers a PATH-only icpx whose oneAPI root is outside /opt. + icpx = _find_icpx() + if icpx: + candidates.append(os.path.join(os.path.dirname(os.path.dirname(icpx)), "include")) + for inc in candidates: + if os.path.exists( + os.path.join(inc, "sycl", "ext", "oneapi", "experimental", "ipc_memory.hpp") + ): + return inc + return None + + +def _find_icpx() -> str | None: + """Locate the icpx (SYCL) compiler needed for the -fsycl build.""" + candidates: list[str] = [] + root = os.getenv("CMPLR_ROOT") + if root: + candidates.append(os.path.join(root, "bin", "icpx")) + candidates += sorted(glob.glob("/opt/intel/oneapi/compiler/*/bin/icpx"), reverse=True) + for cand in candidates: + if os.path.exists(cand): + return cand + # Fallback to PATH: covers oneAPI layouts outside /opt and a sourced setvars.sh + # that puts icpx on PATH without exporting CMPLR_ROOT. + return shutil.which("icpx") + + +@functools.lru_cache(maxsize=1) +def load_ext() -> ModuleType: + """JIT-compile (``-fsycl``, linking torch's libsycl) and cache the SYCL IPC extension. + + Raises on any failure; callers treat an exception as "XPU IPC unavailable". + """ + icpx = _find_icpx() + if icpx is None: + raise RuntimeError("icpx (oneAPI SYCL compiler) not found; cannot build XPU IPC extension") + icx = os.path.join(os.path.dirname(icpx), "icx") + + from torch.utils.cpp_extension import load + + src = Path(__file__).with_name("sycl_ipc.cpp") + + sycl_include_flags: list[str] = [] + inc = _find_sycl_include_dir() + if inc: + sycl_include_flags = [f"-I{inc}", f"-I{os.path.join(inc, 'sycl')}"] + + # torch.utils.cpp_extension picks the compiler from CC/CXX. Force icx/icpx for the + # -fsycl build: a conda/CI env often exports CXX=g++ (gxx_linux-64), which cannot + # compile -fsycl, and setdefault would keep it. Save/restore the process env. + prev_cc, prev_cxx = os.environ.get("CC"), os.environ.get("CXX") + if prev_cxx and os.path.realpath(prev_cxx) != os.path.realpath(icpx): + logger.debug(f"overriding CXX={prev_cxx!r} with icpx for the SYCL IPC build ({icpx})") + os.environ["CC"], os.environ["CXX"] = icx, icpx + try: + # Do NOT pin -std: torch.utils.cpp_extension injects the standard its ATen + # headers require, and a pin here would override it. + module = load( + name="checkpoint_engine_sycl_ipc", + sources=[str(src)], + extra_cflags=["-fsycl", "-O2", *sycl_include_flags], + extra_ldflags=["-fsycl"], + verbose=False, + ) + finally: + for var, prev in (("CC", prev_cc), ("CXX", prev_cxx)): + if prev is None: + os.environ.pop(var, None) + else: + os.environ[var] = prev + return module + + +# Cache only a *successful* probe so a transient first failure can be retried +# (the build itself is memoised by load_ext()'s lru_cache). +_AVAILABLE: bool = False + + +def is_available() -> bool: + """Whether native XPU SYCL IPC can be built and used here (successes cached, failures retried).""" + global _AVAILABLE + if _AVAILABLE: + return True + try: + import torch + + if not (hasattr(torch, "xpu") and torch.xpu.is_available()): + return False + load_ext() + except Exception as e: # noqa: BLE001 + logger.warning(f"xpu sycl ipc unavailable: {e}") + return False + _AVAILABLE = True + return True + + +def prewarm() -> bool: + """Build the extension ahead of time (outside any weight-update timeout); safe on non-XPU hosts.""" + return is_available() + + +def get_handle(ptr: int) -> bytes: + """Portable IPC handle bytes for a device pointer (interior pointers ok; offset is in the blob).""" + return bytes(load_ext().ipc_get_handle(ptr)) + + +def open_handle(handle_bytes: bytes, device: int) -> int: + """Open another process's handle -> device pointer (offset included); free via :func:`close_handle`.""" + return load_ext().ipc_open_handle(list(handle_bytes), device) + + +def release_handle(ptr: int) -> None: + """Release the exporter handle from :func:`get_handle`; no-op if ``ptr`` was never exported. + + Deferred until the consumer has opened: releasing earlier can free the fd under + the level-zero-v2 UR adapter. + """ + load_ext().ipc_release_handle(ptr) + + +def close_handle(ptr: int) -> None: + load_ext().ipc_close_handle(ptr) + + +def wrap_tensor(ptr: int, nbytes: int, device: int) -> torch.Tensor: + """Wrap an IPC-mapped device pointer as a non-owning torch XPU uint8 tensor.""" + return load_ext().ipc_wrap_tensor(ptr, nbytes, device) diff --git a/checkpoint_engine/xpu_ipc/sycl_ipc.cpp b/checkpoint_engine/xpu_ipc/sycl_ipc.cpp new file mode 100644 index 0000000..f39de46 --- /dev/null +++ b/checkpoint_engine/xpu_ipc/sycl_ipc.cpp @@ -0,0 +1,103 @@ +// Cross-process IPC for Intel XPU tensors via SYCL ipc_memory. The handle is a +// self-contained, portable byte blob: the consumer opens it from the bytes alone +// (no dma-buf fd, no offset to carry -- get() takes torch's interior data_ptr and +// open() restores it). Context and device come from torch (c10::xpu) so the +// mapping lands on torch's own SYCL context. +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace ipc = sycl::ext::oneapi::experimental::ipc_memory; + +namespace { + +std::vector to_bytes(const std::vector& in) { + return {reinterpret_cast(in.data()), + reinterpret_cast(in.data()) + in.size()}; +} + +// Exporter-side handles kept alive until ipc_release_handle(). We must not +// ipc::put() before the consumer opens: under the UR level-zero-v2 adapter +// put_ipc_handle frees the exporter fd and can race the consumer's open. +// ipc::handle is a copyable, non-owning value (freed only via ipc::put), so +// storing it by value needs no manual new/delete. +std::mutex g_handles_mu; +std::unordered_map g_handles; + +} // namespace + +// Portable IPC handle bytes for the allocation backing `ptr` (interior pointers +// are fine -- the offset is in the blob). Handle retained until ipc_release_handle(). +std::vector ipc_get_handle(uintptr_t ptr) { + sycl::context ctx = c10::xpu::get_device_context(); + ipc::handle h = ipc::get(reinterpret_cast(ptr), ctx); + ipc::handle_data_t data = h.data(); // owning copy of the blob, independent of `h` + { + std::lock_guard lk(g_handles_mu); + auto it = g_handles.find(ptr); + if (it != g_handles.end()) { + ipc::put(it->second, ctx); // release stale handle for a reused address + it->second = h; + } else { + g_handles.emplace(ptr, h); + } + } + return {reinterpret_cast(data.data()), + reinterpret_cast(data.data()) + data.size()}; +} + +// Release the exporter handle from ipc_get_handle(ptr); no-op if unregistered. +// Call only after all consumers have opened their mappings (see level-zero-v2 note). +void ipc_release_handle(uintptr_t ptr) { + std::optional h; + { + std::lock_guard lk(g_handles_mu); + auto it = g_handles.find(ptr); + if (it == g_handles.end()) { + return; + } + h = it->second; + g_handles.erase(it); + } + ipc::put(*h, c10::xpu::get_device_context()); +} + +// Open a handle from another process -> mapped device pointer (offset included); +// pass it back to ipc_close_handle to release the mapping. +uintptr_t ipc_open_handle(const std::vector& blob, int64_t device) { + sycl::context ctx = c10::xpu::get_device_context(); + sycl::device dev = c10::xpu::get_raw_device(static_cast(device)); + std::vector data = to_bytes(blob); + void* p = ipc::open(data, ctx, dev); + return reinterpret_cast(p); +} + +// Close an IPC mapping; `ptr` must be from ipc_open_handle. +void ipc_close_handle(uintptr_t ptr) { + ipc::close(reinterpret_cast(ptr), c10::xpu::get_device_context()); +} + +// Wrap an external device pointer as a non-owning torch XPU uint8 tensor (the XPU +// analogue of rebuild_cuda_tensor) so the worker can read weight slices. +torch::Tensor ipc_wrap_tensor(uintptr_t dptr, int64_t nbytes, int64_t device) { + auto opts = torch::TensorOptions().dtype(torch::kUInt8).device( + torch::kXPU, static_cast(device)); + return torch::from_blob(reinterpret_cast(dptr), {nbytes}, [](void*) {}, opts); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ipc_get_handle", &ipc_get_handle, "Get portable IPC handle bytes for a device ptr"); + m.def("ipc_release_handle", &ipc_release_handle, "Release the retained exporter handle for a ptr"); + m.def("ipc_open_handle", &ipc_open_handle, "Open an IPC handle blob -> device ptr"); + m.def("ipc_close_handle", &ipc_close_handle, "Close an IPC mapping"); + m.def("ipc_wrap_tensor", &ipc_wrap_tensor, "Wrap an external XPU ptr as a torch tensor"); +} diff --git a/pyproject.toml b/pyproject.toml index c200382..b5a9f28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,11 @@ requires = ["setuptools", "setuptools-scm>=8"] build-backend = "setuptools.build_meta" [tool.setuptools.packages.find] -include = ["checkpoint_engine"] +include = ["checkpoint_engine", "checkpoint_engine.*"] + +[tool.setuptools.package-data] +# Ship the SYCL ipc_memory source; it is JIT-compiled at runtime on XPU hosts. +"checkpoint_engine.xpu_ipc" = ["*.cpp"] [tool.setuptools_scm] version_file = "checkpoint_engine/_version.py" diff --git a/tests/test_device_manager.py b/tests/test_device_manager.py new file mode 100644 index 0000000..aa2046d --- /dev/null +++ b/tests/test_device_manager.py @@ -0,0 +1,174 @@ +"""Unit tests for DeviceManager multi-accelerator dispatch (cuda / npu / xpu). + +These tests mock the device backends so they run on CPU-only CI (``-m "not gpu"``). A separate +hardware-gated test exercises the real ``torch.xpu`` path when an Intel GPU is present. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from checkpoint_engine.device_utils import DeviceManager + + +def _make_manager(device_type: str, device_module: object) -> DeviceManager: + """Build a DeviceManager without touching real hardware by stubbing detection/setup.""" + dm = DeviceManager.__new__(DeviceManager) + dm.device_type = device_type + dm.device_module = device_module + return dm + + +@pytest.mark.parametrize( + "device_type,expected_backend", + [("cuda", "nccl"), ("npu", "hccl"), ("xpu", "xccl")], +) +def test_backend_mapping(device_type: str, expected_backend: str): + dm = _make_manager(device_type, SimpleNamespace()) + assert dm.backend == expected_backend + + +def test_backend_unsupported(): + dm = _make_manager("tpu", SimpleNamespace()) + with pytest.raises(TypeError): + _ = dm.backend + + +@pytest.mark.parametrize("device_type", ["cuda", "xpu"]) +def test_transfer_engine_protocol_rdma(device_type: str): + dm = _make_manager(device_type, SimpleNamespace()) + with patch("checkpoint_engine.device_utils.has_efa_pci", return_value=False): + assert dm.transfer_engine_protocol == "rdma" + with patch("checkpoint_engine.device_utils.has_efa_pci", return_value=True): + assert dm.transfer_engine_protocol == "efa" + + +def test_transfer_engine_protocol_npu(): + dm = _make_manager("npu", SimpleNamespace()) + assert dm.transfer_engine_protocol == "ascend_direct" + + +def test_ipc_collect_present_is_called(): + called = [] + module = SimpleNamespace(ipc_collect=lambda: called.append(True)) + dm = _make_manager("cuda", module) + dm.ipc_collect() + assert called == [True] + + +def test_ipc_collect_absent_is_noop(): + # torch.xpu has no ipc_collect attribute; ipc_collect() must not raise. + dm = _make_manager("xpu", SimpleNamespace()) + dm.ipc_collect() # no attribute -> silent no-op + + +@pytest.mark.parametrize( + "device_type,expected", + [("cuda", True), ("npu", False), ("xpu", False)], +) +def test_supports_inplace_pin(device_type: str, expected: bool): + dm = _make_manager(device_type, SimpleNamespace()) + assert dm.supports_inplace_pin() is expected + + +@pytest.mark.parametrize("device_type", ["cuda", "npu"]) +def test_supports_device_ipc_true_for_cuda_npu(device_type: str): + dm = _make_manager(device_type, SimpleNamespace()) + assert dm.supports_device_ipc() is True + + +def test_supports_device_ipc_xpu_uses_sycl_extension(): + # On XPU we do not rely on torch reductions (PyTorch has no XPU tensor IPC); + # instead we detect our native SYCL ipc_memory extension via xpu_ipc.is_available(). + dm = _make_manager("xpu", SimpleNamespace()) + with patch("checkpoint_engine.xpu_ipc.is_available", return_value=True): + assert dm.supports_device_ipc() is True + with patch("checkpoint_engine.xpu_ipc.is_available", return_value=False): + assert dm.supports_device_ipc() is False + + +def test_supports_device_ipc_unknown_device(): + dm = _make_manager("tpu", SimpleNamespace()) + assert dm.supports_device_ipc() is False + + +@pytest.mark.parametrize( + "device_type,expected", + [("cuda", True), ("npu", True), ("xpu", False), ("tpu", False)], +) +def test_supports_device_p2p(device_type: str, expected: bool): + # Mooncake has no Level Zero backend, so XPU device-memory P2P is unsupported. + dm = _make_manager(device_type, SimpleNamespace()) + assert dm.supports_device_p2p() is expected + + +def test_host_empty_cache_noncuda_uses_gc(): + dm = _make_manager("xpu", SimpleNamespace()) + with patch("checkpoint_engine.device_utils.gc.collect") as gc_collect: + dm.host_empty_cache() + gc_collect.assert_called_once() + + +def test_host_empty_cache_cuda_uses_torch(): + dm = _make_manager("cuda", SimpleNamespace()) + with patch.object(torch._C, "_host_emptyCache", create=True) as host_empty: + dm.host_empty_cache() + host_empty.assert_called_once() + + +# -------------------------------------------------------------------------------------------- +# Hardware-gated: real torch.xpu behavior on an Intel GPU host. +# -------------------------------------------------------------------------------------------- + +_HAS_XPU = hasattr(torch, "xpu") and torch.xpu.is_available() + + +@pytest.mark.gpu +@pytest.mark.skipif(not _HAS_XPU, reason="requires an Intel XPU device") +def test_real_xpu_device_manager(): + dm = DeviceManager() + assert dm.device_type == "xpu" + assert dm.backend == "xccl" + assert dm.device_module is torch.xpu + # ipc_collect must be a harmless no-op (torch.xpu has no ipc_collect). + dm.ipc_collect() + # XPU has no torch-native device-tensor IPC, but checkpoint-engine ships its own + # native SYCL ipc_memory transport; supports_device_ipc() must agree with whether + # that extension can actually be built/loaded in this environment. + from checkpoint_engine import xpu_ipc + + assert dm.supports_device_ipc() is xpu_ipc.is_available() + # Mooncake has no Level Zero backend, so device-memory P2P stays unsupported on XPU. + assert dm.supports_device_p2p() is False + # cudaHostRegister-style in-place pinning is CUDA-only. + assert dm.supports_inplace_pin() is False + + +@pytest.mark.gpu +@pytest.mark.skipif(not _HAS_XPU, reason="requires an Intel XPU device") +def test_real_xpu_device_ipc_available_when_extension_builds(): + """When the SYCL ipc_memory extension builds on real XPU hardware, the + broadcast path must be reported as supported. This guards the torch>=2.14 + c++20 build regression that silently disabled XPU broadcast.""" + from checkpoint_engine import xpu_ipc + + if not xpu_ipc.is_available(): + pytest.skip("SYCL ipc_memory extension could not be built in this environment") + dm = DeviceManager() + assert dm.supports_device_ipc() is True + + +@pytest.mark.gpu +@pytest.mark.skipif(not _HAS_XPU, reason="requires an Intel XPU device") +def test_real_xpu_physical_uuid_matches_worker_format(): + from checkpoint_engine.ps import _get_physical_gpu_id + + dm = DeviceManager() + uuid0 = _get_physical_gpu_id(dm, 0) + # Same format the vLLM worker derives from torch.xpu.get_device_properties(idx).uuid. + expected = f"GPU-{torch.xpu.get_device_properties(0).uuid!s}" + assert uuid0 == expected + if dm.device_module.device_count() > 1: + assert _get_physical_gpu_id(dm, 1) != uuid0, "per-device uuids must be distinct" diff --git a/tests/test_p2p_guard.py b/tests/test_p2p_guard.py new file mode 100644 index 0000000..0b2b2d8 --- /dev/null +++ b/tests/test_p2p_guard.py @@ -0,0 +1,45 @@ +"""The P2P update path must reject devices whose memory Mooncake cannot register. + +XPU device memory has no Level Zero backend in Mooncake, so a P2P update +(``ranks`` set) must raise a clear error before touching the transfer engine. +CPU-only: we stub the ParameterServer internals up to the guard. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +import checkpoint_engine.distributed as dist +from checkpoint_engine.ps import ParameterServer + + +def _ps_with_device(device_type: str, *, supports_ipc: bool, supports_p2p: bool) -> ParameterServer: + ps = ParameterServer.__new__(ParameterServer) + ps._rank = 0 + ps.device_manager = SimpleNamespace( + device_type=device_type, + supports_device_ipc=lambda: supports_ipc, + supports_device_p2p=lambda: supports_p2p, + ) + # Non-empty metas so the leading assert passes; content is irrelevant (guard fires first). + ps._current_global_parameter_metas = {0: object()} + return ps + + +def test_p2p_update_rejected_on_xpu(): + ps = _ps_with_device("xpu", supports_ipc=True, supports_p2p=False) + with ( + patch.object(dist, "is_initialized", return_value=True), + pytest.raises(RuntimeError, match=r"P2P weight update .* is not supported"), + ): + ps._update_per_bucket("ckpt", req_func=lambda _paths: None, ranks_group=None, ranks=[0]) + + +def test_ipc_unavailable_rejected(): + ps = _ps_with_device("xpu", supports_ipc=False, supports_p2p=False) + with ( + patch.object(dist, "is_initialized", return_value=True), + pytest.raises(RuntimeError, match="cross-process device-tensor IPC"), + ): + ps._update_per_bucket("ckpt", req_func=lambda _paths: None, ranks_group=None, ranks=None) diff --git a/tests/test_transport.py b/tests/test_transport.py new file mode 100644 index 0000000..24b21be --- /dev/null +++ b/tests/test_transport.py @@ -0,0 +1,110 @@ +"""Unit tests for the weight-transport seam (CPU-only, no accelerator required). + +The zero-copy device handoff itself is exercised by a hardware-gated end-to-end +test on real XPU/CUDA. Here we cover the dispatch logic and wire formats. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from checkpoint_engine.transport import ( + IpcWeightTransport, + XpuIpcWeightTransport, + build_transport, +) +from checkpoint_engine.worker import _transport_for_handle + + +def _dm(device_type: str) -> object: + return SimpleNamespace(device_type=device_type) + + +@pytest.mark.parametrize( + "device_type,expected", + [("cuda", IpcWeightTransport), ("npu", IpcWeightTransport), ("xpu", XpuIpcWeightTransport)], +) +def test_build_transport_dispatch(device_type: str, expected: type): + assert isinstance(build_transport(_dm(device_type)), expected) + + +def test_consumer_dispatch_by_handle_shape(): + # CUDA/NPU send a reduce_tensor tuple; XPU sends a tagged dict. + tuple_handle = (lambda *a: None, (1, 2, 3)) + assert isinstance(_transport_for_handle(tuple_handle), IpcWeightTransport) + + xpu_handle = {"kind": XpuIpcWeightTransport.kind, "handle_bytes": b"", "nbytes": 0} + assert isinstance(_transport_for_handle(xpu_handle), XpuIpcWeightTransport) + + # An unrelated dict must not be mistaken for the XPU handle. + assert isinstance(_transport_for_handle({"foo": "bar"}), IpcWeightTransport) + + +def test_ipc_transport_export_uses_reduce_tensor(): + sentinel = ("REDUCED",) + with patch("checkpoint_engine.transport.reduce_tensor", return_value=sentinel) as m: + t = IpcWeightTransport() + out = t.export(SimpleNamespace()) + assert out is sentinel + m.assert_called_once() + + +def test_xpu_export_returns_self_contained_handle(): + # The SYCL handle bytes travel over ZMQ as a picklable dict; no fd, no offset, + # no companion socket. Mock the native extension so this runs on CPU CI. + buffer = SimpleNamespace(data_ptr=lambda: 0xDEAD, nbytes=256) + with patch("checkpoint_engine.xpu_ipc.get_handle", return_value=b"HANDLE") as m: + handle = XpuIpcWeightTransport().export(buffer) + m.assert_called_once_with(0xDEAD) + assert handle == {"kind": "xpu_sycl", "handle_bytes": b"HANDLE", "nbytes": 256} + + +def test_xpu_export_defers_release_until_detach(): + # The exporter handle must be released only in detach() (not export()), against + # the exact pointer exported -- releasing early can free the fd under UR v2. + buffer = SimpleNamespace(data_ptr=lambda: 0xBEEF, nbytes=128) + with ( + patch("checkpoint_engine.xpu_ipc.get_handle", return_value=b"H"), + patch("checkpoint_engine.xpu_ipc.release_handle") as release, + ): + t = XpuIpcWeightTransport() + t.export(buffer) + release.assert_not_called() # not released during export + t.detach() + release.assert_called_once_with(0xBEEF) + # Idempotent: a second detach must not double-release. + t.detach() + release.assert_called_once_with(0xBEEF) + + +def test_xpu_transport_detach_is_safe_when_unused(): + # detach() before any export/attach must not raise (and must not touch the ext). + with ( + patch("checkpoint_engine.xpu_ipc.release_handle") as release, + patch("checkpoint_engine.xpu_ipc.close_handle") as close, + ): + XpuIpcWeightTransport().detach() + release.assert_not_called() + close.assert_not_called() + + +def test_xpu_consumer_detach_closes_opened_mapping(): + # The consumer (attach) side unmaps its opened pointer on detach, and must not + # try to release an exporter handle it never took. + handle = {"kind": "xpu_sycl", "handle_bytes": b"H", "nbytes": 64} + with ( + patch("checkpoint_engine.xpu_ipc.open_handle", return_value=0x7000), + patch("checkpoint_engine.xpu_ipc.wrap_tensor") as wrap, + patch("checkpoint_engine.transport.torch.xpu.synchronize"), + patch("checkpoint_engine.xpu_ipc.close_handle") as close, + patch("checkpoint_engine.xpu_ipc.release_handle") as release, + ): + import torch as _torch + + wrap.return_value = SimpleNamespace(dtype=_torch.uint8) + t = XpuIpcWeightTransport() + t.attach(handle, device_id=0) + t.detach() + close.assert_called_once_with(0x7000) + release.assert_not_called() diff --git a/tests/test_xpu_ipc.py b/tests/test_xpu_ipc.py new file mode 100644 index 0000000..992e818 --- /dev/null +++ b/tests/test_xpu_ipc.py @@ -0,0 +1,154 @@ +"""Hardware-gated tests for the native SYCL ipc_memory transport on Intel XPU. + +These are skipped unless an Intel GPU is present and the SYCL IPC extension can +be built (needs an oneAPI ``icpx`` with SYCL ipc_memory support). They are marked +``gpu`` so CPU-only CI (``-m "not gpu"``) skips them. +""" + +import os + +import pytest +import torch + + +pytestmark = pytest.mark.gpu + + +def _xpu_ipc_available() -> bool: + try: + if not (hasattr(torch, "xpu") and torch.xpu.is_available()): + return False + from checkpoint_engine import xpu_ipc + + return xpu_ipc.is_available() + except Exception: # noqa: BLE001 + return False + + +skip_no_xpu_ipc = pytest.mark.skipif( + not _xpu_ipc_available(), reason="Intel XPU with buildable SYCL ipc_memory extension required" +) + +_N_TENSORS = 8 + + +def _gen_tensors() -> dict[str, torch.Tensor]: + gen = torch.Generator().manual_seed(0) + return { + f"w{i}": (torch.randn(128 + i, 64, generator=gen) * 50).to(torch.bfloat16) + for i in range(_N_TENSORS) + } + + +def _worker_proc( + device_uuid: str, + expected: dict[str, torch.Tensor], + inq: object, + outq: object, +) -> None: + # Module-level so the "spawn" start method can pickle it. + import zmq + + from checkpoint_engine.worker import update_weights_from_ipc + + torch.xpu.set_device(0) + exp = {k: v.to("xpu:0") for k, v in expected.items()} + ctx = zmq.Context() + state = {"n": 0, "ok": True} + + def run(weights: list[tuple[str, torch.Tensor]]) -> None: + for name, w in weights: + if name in exp and not torch.equal(w.to(torch.bfloat16), exp[name]): + state["ok"] = False + elif name in exp: + state["n"] += 1 + + while True: + socket_paths = inq.get() + if socket_paths is None: + break + update_weights_from_ipc( + ctx, + dict(socket_paths)[device_uuid], + device_id=0, + run=run, + post_hook=lambda: torch.xpu.synchronize(), + ) + outq.put((state["n"], state["ok"])) + + +@skip_no_xpu_ipc +def test_sycl_ipc_same_process_roundtrip(): + """get_handle -> open_handle (same process) maps back to the original bytes.""" + from checkpoint_engine import xpu_ipc + + torch.xpu.set_device(0) + t = torch.arange(256, device="xpu:0", dtype=torch.uint8) + torch.xpu.synchronize() + + handle_bytes = xpu_ipc.get_handle(t.data_ptr()) + ptr = xpu_ipc.open_handle(handle_bytes, 0) + wrapped = xpu_ipc.wrap_tensor(ptr, t.numel(), 0) + try: + assert torch.equal(wrapped, t) + finally: + xpu_ipc.close_handle(ptr) + + +@skip_no_xpu_ipc +def test_sycl_ipc_interior_pointer_offset_preserved(): + """A SYCL IPC handle for an interior (sub-allocation) pointer must reopen at + the same offset -- the offset is encoded in the portable handle bytes, not + carried separately.""" + from checkpoint_engine import xpu_ipc + + torch.xpu.set_device(0) + big = torch.zeros(4096, device="xpu:0", dtype=torch.uint8) + view = big[1024:1280] + view.fill_(0xAB) + torch.xpu.synchronize() + + handle_bytes = xpu_ipc.get_handle(view.data_ptr()) + ptr = xpu_ipc.open_handle(handle_bytes, 0) + wrapped = xpu_ipc.wrap_tensor(ptr, view.numel(), 0) + try: + assert torch.equal(wrapped, view) + finally: + xpu_ipc.close_handle(ptr) + + +@skip_no_xpu_ipc +def test_sycl_ipc_cross_process_broadcast(): + """Full ParameterServer broadcast -> colocated worker over the XPU SYCL transport.""" + from torch.multiprocessing import get_context + + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29570") + + from checkpoint_engine.ps import ParameterServer, _get_physical_gpu_id + + tensors = _gen_tensors() + + torch.xpu.set_device(0) + ps = ParameterServer(auto_pg=True) + uuid = _get_physical_gpu_id(ps.device_manager, 0) + + mp = get_context("spawn") + inq, outq = mp.Queue(), mp.Queue() + proc = mp.Process(target=_worker_proc, args=(uuid, tensors, inq, outq)) + proc.start() + try: + ps.register_checkpoint("ckpt", named_tensors=tensors) + ps.init_process_group() + ps.gather_metas("ckpt") + ps.update("ckpt", inq.put) + inq.put(None) + n, ok = outq.get(timeout=60) + assert ok, "received weights did not match originals" + assert n == _N_TENSORS, f"expected {_N_TENSORS} tensors checked, got {n}" + finally: + proc.join(timeout=30) + if proc.is_alive(): + proc.terminate() diff --git a/tests/test_xpu_parity.py b/tests/test_xpu_parity.py new file mode 100644 index 0000000..1b2e7d2 --- /dev/null +++ b/tests/test_xpu_parity.py @@ -0,0 +1,225 @@ +"""CPU-only parity tests for the XPU support paths that diverge from CUDA/NPU. + +These cover device-agnostic logic that would otherwise have no unit coverage: +* The portable-handle contract of ``xpu_ipc.get_handle``/``open_handle`` (bytes + in, bytes out -- no fd, no offset). +* The custom-distributed backend rejection for XPU. +* ``register_checkpoint`` forcing in-place pinning off on non-CUDA devices. + +The zero-copy device handoff itself is exercised by the hardware-gated tests in +``test_xpu_ipc.py``; here we isolate the surrounding logic so it runs on CPU-only +CI (``-m "not gpu"``). +""" + +import os +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +import checkpoint_engine.distributed as dist +from checkpoint_engine import xpu_ipc +from checkpoint_engine.distributed.base import TorchBackend, use_backend + + +# ------------------------------------------------------------------------------ +# get_handle / open_handle: a SYCL ipc_memory handle is self-contained portable +# bytes -- get_handle returns the raw blob and open_handle passes it straight to +# the extension (offset and any fd are encoded inside the bytes, not carried). +# ------------------------------------------------------------------------------ + + +def test_get_handle_returns_raw_portable_bytes(): + fake_ext = MagicMock() + blob = list(b"\xab" * 120) # SYCL handle blob (opaque, self-contained) + fake_ext.ipc_get_handle.return_value = blob + with patch("checkpoint_engine.xpu_ipc.load_ext", return_value=fake_ext): + handle_bytes = xpu_ipc.get_handle(0xDEAD) + fake_ext.ipc_get_handle.assert_called_once_with(0xDEAD) + assert handle_bytes == bytes(blob) + + +def test_open_handle_passes_bytes_and_device_through(): + fake_ext = MagicMock() + fake_ext.ipc_open_handle.return_value = 0x5000 # mapped ptr (offset already applied) + with patch("checkpoint_engine.xpu_ipc.load_ext", return_value=fake_ext): + ptr = xpu_ipc.open_handle(b"\xab" * 120, device=2) + assert ptr == 0x5000 + (blob_arg, device_arg) = fake_ext.ipc_open_handle.call_args.args + assert device_arg == 2 + assert bytes(blob_arg) == b"\xab" * 120 + + +# ------------------------------------------------------------------------------ +# load_ext: the -fsycl build must force CC/CXX to icx/icpx (a conda/CI env often +# exports CXX=g++, which cannot compile -fsycl) and restore the env afterwards. +# ------------------------------------------------------------------------------ + + +def test_load_ext_forces_icpx_over_existing_gpp_and_restores_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("CXX", "/usr/bin/g++") # what conda gxx_linux-64 exports + monkeypatch.delenv("CC", raising=False) + + captured: dict[str, str | None] = {} + + def fake_load(**kwargs: object) -> MagicMock: + # Record what torch.utils.cpp_extension.load would see. + captured["CC"] = os.environ.get("CC") + captured["CXX"] = os.environ.get("CXX") + return MagicMock() + + xpu_ipc.load_ext.cache_clear() + try: + with ( + patch("checkpoint_engine.xpu_ipc._find_icpx", return_value="/opt/oneapi/bin/icpx"), + patch("checkpoint_engine.xpu_ipc._find_sycl_include_dir", return_value=None), + patch("torch.utils.cpp_extension.load", side_effect=fake_load), + ): + xpu_ipc.load_ext() + finally: + xpu_ipc.load_ext.cache_clear() + + # During the build the SYCL compiler must win over the inherited g++. + assert captured["CXX"] == "/opt/oneapi/bin/icpx" + assert captured["CC"] == "/opt/oneapi/bin/icx" + # ...and the caller's environment must be restored afterwards. + assert os.environ["CXX"] == "/usr/bin/g++" + assert "CC" not in os.environ + + +def test_find_icpx_falls_back_to_path(monkeypatch: pytest.MonkeyPatch) -> None: + # A sourced setvars.sh may put icpx on PATH without exporting CMPLR_ROOT, and + # oneAPI need not live under /opt. Without the PATH fallback the compiler is + # reported missing and XPU IPC is wrongly declared unavailable. + monkeypatch.delenv("CMPLR_ROOT", raising=False) + with ( + patch("checkpoint_engine.xpu_ipc.glob.glob", return_value=[]), + patch("checkpoint_engine.xpu_ipc.shutil.which", return_value="/custom/bin/icpx") as which, + ): + assert xpu_ipc._find_icpx() == "/custom/bin/icpx" + which.assert_called_once_with("icpx") + + +def test_is_available_caches_only_success_and_retries_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import checkpoint_engine.xpu_ipc as mod + + monkeypatch.setattr(mod, "_AVAILABLE", False) + calls = {"n": 0} + + def flaky_load() -> MagicMock: + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("transient: compiler env not warm yet") + return MagicMock() + + fake_torch = SimpleNamespace(xpu=SimpleNamespace(is_available=lambda: True)) + with ( + patch.dict("sys.modules", {"torch": fake_torch}), + patch.object(mod, "load_ext", side_effect=flaky_load), + ): + assert mod.is_available() is False # transient failure NOT cached + assert mod.is_available() is True # retried, now succeeds + assert mod.is_available() is True # success cached (no third load_ext) + assert calls["n"] == 2 + + +# ------------------------------------------------------------------------------ +# Custom distributed backend: XPU must reject custom_dist and fall back to the +# native "xccl" TorchBackend (there is no vLLM PyXcclCommunicator to subclass). +# ------------------------------------------------------------------------------ + + +def test_use_backend_rejects_xpu_custom_dist(): + with pytest.raises(ValueError, match="XPU is not supported here"): + use_backend("vllm_xccl") + + +def test_use_backend_none_keeps_default_torch_backend(): + # A falsy backend must leave the (default) TorchBackend in place, which is what + # XPU relies on for xccl. + before = dist.is_initialized # attribute presence sanity + use_backend(None) + from checkpoint_engine.distributed.base import _BACKEND_INSTANCE + + assert isinstance(_BACKEND_INSTANCE, TorchBackend) + assert before is dist.is_initialized + + +# ------------------------------------------------------------------------------ +# register_checkpoint: in-place pinning (cudaHostRegister) is CUDA-only; on XPU +# it must be silently disabled rather than attempted. +# ------------------------------------------------------------------------------ + + +# ------------------------------------------------------------------------------ +# _update_per_bucket: the exported IPC handle is retained until detach(). A +# failure after export but before the broadcast loop (e.g. the ZMQ bind) must +# still release it -- otherwise the exporter handle leaks on every failed update. +# ------------------------------------------------------------------------------ + + +def test_update_per_bucket_detaches_transport_on_early_failure(): + from checkpoint_engine.ps import ParameterServer + + ps = ParameterServer.__new__(ParameterServer) + ps._rank = 0 + ps.device_manager = SimpleNamespace( + device_type="cpu", + supports_device_ipc=lambda: True, + supports_device_p2p=lambda: False, + ) + ps._current_global_parameter_metas = {0: object()} + ps._local_rdma_devices = None + ps._remote_rdma_devices = None + + fake_transport = MagicMock() + fake_transport.export.return_value = {"kind": "fake"} + + with ( + patch.object(dist, "is_initialized", return_value=True), + patch("checkpoint_engine.ps.build_transport", return_value=fake_transport), + patch("checkpoint_engine.ps._gen_h2d_buckets", return_value=[]), + patch.object(ps, "_detect_bucket_size", return_value=(16, False)), + patch.object(ps, "_bind_zmq_socket", side_effect=RuntimeError("bind failed")), + pytest.raises(RuntimeError, match="bind failed"), + ): + ps._update_per_bucket("ckpt", req_func=lambda _paths: None, ranks_group=None, ranks=None) + + # export happened, so the exporter handle is live and must be released even + # though the failure struck before the broadcast loop's own cleanup. + fake_transport.export.assert_called_once() + fake_transport.detach.assert_called_once_with() + + +def test_register_checkpoint_disables_inplace_pin_on_xpu(): + from checkpoint_engine.ps import ParameterServer + + ps = ParameterServer.__new__(ParameterServer) + ps._rank = 0 + ps.device_manager = SimpleNamespace( + device_type="xpu", + supports_inplace_pin=lambda: False, + ) + ps._memory_pool = {} + ps._current_shared_memory_pool_user = "" + ps._p2p_store = None + ps.shared_memory_pool_name = ParameterServer.shared_memory_pool_name + + captured = {} + + def fake_register(*, inplace_pin: bool, **kwargs: object) -> list: + captured["inplace_pin"] = inplace_pin + return [] + + with patch("checkpoint_engine.ps._register_checkpoint", side_effect=fake_register): + ps.register_checkpoint("ckpt", named_tensors={}, use_inplace_pin_memory=True) + # Requested True, but XPU cannot in-place pin -> must be forced False. + assert captured["inplace_pin"] is False + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) From c90ff3fe15f40f07718aebf58e1f179f82bd4215 Mon Sep 17 00:00:00 2001 From: Siju Samuel Date: Thu, 30 Jul 2026 03:51:49 +0000 Subject: [PATCH 2/2] review-rework: address PR #96 review comments Rework of the review comments from @weixiao-huang and @HubertZhang. Drop `from __future__ import annotations` from the two new modules and quote the TYPE_CHECKING-only annotations instead; the rest of the repo already relies on native PEP 604 unions (requires-python >= 3.10). DeviceManager.ipc_collect no longer duck-types via getattr/callable: it dispatches on device_type like every other method in the class and raises TypeError on an unsupported backend. XPU is an explicit no-op (SYCL frees on close_handle; torch.xpu has no ipc_collect). Hoist the IPC-handle lifetime into update(), which opens the handler as a context manager and passes it to _update_per_bucket. The exported handle is still released on every exit path -- including a failure before the broadcast loop's own cleanup -- but the loop keeps upstream's exact try/finally nesting, so the only functional change inside it is reduce_tensor(buffer) -> ipc_handler.export(buffer). Rename the abstraction to match what it does -- it hands over an IPC handle rather than transporting bytes: WeightTransport -> IPCHandler, IpcWeightTransport -> TorchIPCHandler, XpuIpcWeightTransport -> XpuIPCHandler, build_transport -> build_ipc_handler, transport.py -> ipc_handler.py (and tests/test_transport.py -> tests/test_ipc_handler.py). The module and class docstrings now say the object exchanges a handle, not the buffer. The "xpu_sycl" wire tag is unchanged, so producer and consumer stay compatible across builds. Guard the SYCL namespace with __has_include: upstream intel/llvm split the API (functions -> ipc::memory, handle types -> the parent ipc namespace) and deprecated flat ipc_memory, but no oneAPI release ships that layout yet and the extension's feature-test macro is unversioned, so both spellings are supported. Two aliases are needed because qualified lookup cannot reach the parent namespace's types through a nested alias. Verified: the shim compiles clean under -Werror=deprecated-declarations on oneAPI 2026.1. Build the extension with torch's with_sycl=True, which supplies the SYCL include paths and device link. This removes _find_sycl_include_dir entirely and the CC/CXX override block -- confirmed unnecessary, a clean build now succeeds with CXX=g++ exported. _find_icpx stays because torch shells out to a bare "icpx", so the located compiler is prepended to PATH for the build; -O2 is kept since without it the host object is compiled -O0 (6x larger .so). Sort the icpx candidates by numeric version parts so 2026.10 outranks 2026.9 (the previous lexicographic sort got that backwards); the non-numeric `latest` symlink still wins, since it points at the newest install. Reject an icpx that predates SYCL ipc_memory (oneAPI < 2026.0) up front, by probing for the header in the compiler's own tree. Such a compiler builds a device image that torch's newer libsycl cannot load, and the failure is a C++ terminate at dlopen (SIGABRT) that is_available()'s try/except cannot catch -- it killed the whole process, including pytest collection. Now it degrades to is_available() == False with a clear message. Also cover the new ipc_collect TypeError branch and drop the stale test comment that described the removed getattr probe. Make ipc_handler a required positional parameter of _update_per_bucket, placed after req_func, instead of an Optional with an assert; a caller that forgets it now fails at call time rather than deep inside the update. --- README.md | 6 +- checkpoint_engine/device_utils.py | 12 +- .../{transport.py => ipc_handler.py} | 48 ++-- checkpoint_engine/ps.py | 226 +++++++++--------- checkpoint_engine/worker.py | 32 +-- checkpoint_engine/xpu_ipc/__init__.py | 116 +++++---- checkpoint_engine/xpu_ipc/sycl_ipc.cpp | 18 +- tests/test_device_manager.py | 17 +- ...{test_transport.py => test_ipc_handler.py} | 44 ++-- tests/test_p2p_guard.py | 23 +- tests/test_xpu_ipc.py | 4 +- tests/test_xpu_parity.py | 168 ++++++++++--- 12 files changed, 426 insertions(+), 288 deletions(-) rename checkpoint_engine/{transport.py => ipc_handler.py} (73%) rename tests/{test_transport.py => test_ipc_handler.py} (70%) diff --git a/README.md b/README.md index f449b3d..df64369 100644 --- a/README.md +++ b/README.md @@ -80,8 +80,8 @@ pip install 'checkpoint-engine[p2p]' Intel XPU is supported for the **broadcast** update path. The cross-process weight handoff uses a native SYCL `ipc_memory` extension (the XPU counterpart of CUDA IPC) that is JIT-compiled at runtime. Install from source since XPU support is not yet in the released package. P2P is not supported on XPU (Mooncake has no Level Zero backend for XPU device memory). Requirements: -- Intel XPU build of PyTorch — `torch.xpu.is_available()` returns `True`, `torch>=2.9` (for the device `.uuid` property). See the [PyTorch XPU install guide](https://pytorch.org/docs/stable/notes/get_start_xpu.html); this is not the default PyPI `torch`. -- Intel oneAPI 2026.0+ providing the `icpx` compiler with SYCL `ipc_memory` support. Needed at runtime (first weight update), not at `pip install` time. +- Intel XPU build of PyTorch — `torch.xpu.is_available()` returns `True`, `torch>=2.9` (for the device `.uuid` property; the SYCL extension build also needs `torch>=2.7`). See the [PyTorch XPU install guide](https://pytorch.org/docs/stable/notes/get_start_xpu.html); this is not the default PyPI `torch`. +- Intel oneAPI 2026.0+ providing the `icpx` compiler with SYCL IPC memory support. Needed at runtime (first weight update), not at `pip install` time. ```Bash git clone https://github.com/MoonshotAI/checkpoint-engine.git @@ -89,7 +89,7 @@ cd checkpoint-engine pip install -e . # no [p2p] extra on XPU ``` -Make `icpx` discoverable in the runtime environment, either by sourcing oneAPI (`source /opt/intel/oneapi/setvars.sh`) or by setting `CMPLR_ROOT`. If neither is set, `icpx` is auto-detected under `/opt/intel/oneapi/compiler/*/bin`. The extension then builds automatically on first use; `ParameterServer` also prebuilds it at startup so the one-time compile stays out of the update window. +Make `icpx` discoverable in the runtime environment, either by sourcing oneAPI (`source /opt/intel/oneapi/setvars.sh`) or by setting `CMPLR_ROOT`. If neither is set, `icpx` is auto-detected under `/opt/intel/oneapi/compiler/*/bin` and then on `PATH`. The extension then builds automatically on first use; `ParameterServer` also prebuilds it at startup so the one-time compile stays out of the update window. Verify the build and the IPC path on the target machine: diff --git a/checkpoint_engine/device_utils.py b/checkpoint_engine/device_utils.py index c866e13..72d76d4 100644 --- a/checkpoint_engine/device_utils.py +++ b/checkpoint_engine/device_utils.py @@ -273,10 +273,14 @@ def rdma_device(self, rank: int) -> str: raise TypeError("The current transfer engine protocol is not supported") def ipc_collect(self) -> None: - """Reclaim memory held by stale IPC handles where the backend supports it (no-op otherwise).""" - fn = getattr(self.device_module, "ipc_collect", None) - if callable(fn): - fn() + """Reclaim memory held by stale IPC handles.""" + if self.device_type in ("cuda", "npu"): + self.device_module.ipc_collect() + elif self.device_type == "xpu": + # SYCL ipc_memory frees on close_handle; there is no cache to collect. + pass + else: + raise TypeError("The current device type is not supported") def supports_inplace_pin(self) -> bool: """Whether in-place host-memory pinning (cudaHostRegister) is available -- CUDA only.""" diff --git a/checkpoint_engine/transport.py b/checkpoint_engine/ipc_handler.py similarity index 73% rename from checkpoint_engine/transport.py rename to checkpoint_engine/ipc_handler.py index 0f97cdf..4806997 100644 --- a/checkpoint_engine/transport.py +++ b/checkpoint_engine/ipc_handler.py @@ -1,15 +1,15 @@ -"""Pluggable device-buffer handoff between the ParameterServer and the worker. - -The broadcast path shares a device buffer with the colocated worker. CUDA/NPU use -:class:`IpcWeightTransport` (``torch.multiprocessing`` CUDA IPC, wire-format -unchanged); XPU uses :class:`XpuIpcWeightTransport` (native SYCL ``ipc_memory``). -The handle is always a picklable, self-contained value, so the producer's -``export`` -> ZMQ ``send_pyobj`` -> consumer ``attach`` flow is identical for both; -each side calls ``detach`` on cleanup. +"""Pluggable IPC-handle exchange between the ParameterServer and the worker. + +The broadcast path shares a device buffer with the colocated worker. Nothing here +copies or moves the buffer: it only exchanges the IPC handle that lets the worker +map the same device memory. CUDA/NPU use :class:`TorchIPCHandler` +(``torch.multiprocessing`` CUDA IPC, wire-format unchanged); XPU uses +:class:`XpuIPCHandler` (native SYCL ``ipc_memory``). The handle is always a +picklable, self-contained value, so the producer's ``export`` -> ZMQ +``send_pyobj`` -> consumer ``attach`` flow is identical for both; each side calls +``detach`` on cleanup. """ -from __future__ import annotations - from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any @@ -21,10 +21,12 @@ if TYPE_CHECKING: from collections.abc import Callable + from typing_extensions import Self + from checkpoint_engine.device_utils import DeviceManager -def _rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) -> torch.Tensor: +def _rebuild_ipc(handle: tuple["Callable", tuple], device_id: int | None = None) -> torch.Tensor: func, args = handle list_args = list(args) if device_id is not None: @@ -34,8 +36,8 @@ def _rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) - return func(*list_args) -class WeightTransport(ABC): - """Hands a device buffer from the producer (ps) to the consumer (worker).""" +class IPCHandler(ABC): + """Hands an IPC handle for a device buffer from the producer (ps) to the consumer (worker).""" @abstractmethod def export(self, buffer: torch.Tensor) -> Any: @@ -48,8 +50,16 @@ def attach(self, handle: Any, device_id: int) -> torch.Tensor: def detach(self) -> None: """Release IPC resources on either side. No-op by default.""" + # Used as a context manager so the handle is always released, without the + # caller needing its own try/finally. + def __enter__(self) -> "Self": + return self + + def __exit__(self, *exc_info: object) -> None: + self.detach() + -class IpcWeightTransport(WeightTransport): +class TorchIPCHandler(IPCHandler): """CUDA/NPU zero-copy handoff via torch.multiprocessing CUDA IPC (unchanged).""" def export(self, buffer: torch.Tensor) -> Any: @@ -62,7 +72,7 @@ def attach(self, handle: Any, device_id: int) -> torch.Tensor: return buffer -class XpuIpcWeightTransport(WeightTransport): +class XpuIPCHandler(IPCHandler): """Intel XPU zero-copy handoff via native SYCL ``ipc_memory`` (portable byte blob).""" kind = "xpu_sycl" # wire tag: identifies this handle format to the consumer @@ -117,8 +127,8 @@ def detach(self) -> None: self._exported_ptr = None -def build_transport(device_manager: DeviceManager) -> WeightTransport: - """Select the weight transport for the current device backend.""" +def build_ipc_handler(device_manager: "DeviceManager") -> IPCHandler: + """Select the IPC handler for the current device backend.""" if device_manager.device_type == "xpu": - return XpuIpcWeightTransport() - return IpcWeightTransport() + return XpuIPCHandler() + return TorchIPCHandler() diff --git a/checkpoint_engine/ps.py b/checkpoint_engine/ps.py index 1834110..7e50151 100644 --- a/checkpoint_engine/ps.py +++ b/checkpoint_engine/ps.py @@ -23,9 +23,9 @@ ParameterMeta, ) from checkpoint_engine.device_utils import DeviceManager, get_ip, npu_generate_uuid +from checkpoint_engine.ipc_handler import IPCHandler, build_ipc_handler from checkpoint_engine.p2p_store import P2PStore from checkpoint_engine.pin_memory import _ALIGN_SIZE, _register_checkpoint -from checkpoint_engine.transport import build_transport if TYPE_CHECKING: @@ -597,7 +597,10 @@ def update( self.init_process_group(timeout=timeout) # if ranks is None or [], it will use fully broadcast to update to all ranks ranks_group = dist.new_group(ranks) if ranks else None - self._update_per_bucket(checkpoint_name, req_func, ranks_group, ranks) + # `with` releases the exported IPC handle on every exit path, including a + # failure before the broadcast loop's own cleanup starts. + with build_ipc_handler(self.device_manager) as ipc_handler: + self._update_per_bucket(checkpoint_name, req_func, ipc_handler, ranks_group, ranks) self.store_based_barrier() except Exception as e: logger.exception( @@ -749,6 +752,7 @@ def _update_per_bucket( self, checkpoint_name: str, req_func: Callable[[list[tuple[str, str]]], None], + ipc_handler: IPCHandler, ranks_group: dist.DistributedProcessGroup | None, ranks: list[int] | None = None, ): @@ -826,124 +830,114 @@ def _update_per_bucket( self._p2p_store.register_named_tensors( {p2p_ipc_buffer_name: buffer if disable_h2d_buffer else h2d_buffer} ) - transport = build_transport(self.device_manager) - # Outer try guarantees the exported IPC handle is released even when export, - # the socket bind, or the first send below raises -- its finally does nothing - # but detach(). The collective barrier stays in the inner finally so an early - # single-rank failure here cannot deadlock peers that never reach the loop. - try: - handle = transport.export(buffer) - - buckets_by_receiver_rank: dict[int, list[H2DBucket]] = defaultdict(list) - max_len = 0 - for receiver_rank, _, bucket in buckets: - buckets_by_receiver_rank[receiver_rank].append(bucket) - if len(buckets_by_receiver_rank[receiver_rank]) > max_len: - max_len = len(buckets_by_receiver_rank[receiver_rank]) - - socket, socket_paths = self._bind_zmq_socket() - req_thread = threading.Thread( - target=req_func, - args=(socket_paths,), - ) - req_thread.start() - # The handle is self-contained for every transport, so one ZMQ send completes the handoff. - socket.send_pyobj(handle) + handle = ipc_handler.export(buffer) + + buckets_by_receiver_rank: dict[int, list[H2DBucket]] = defaultdict(list) + max_len = 0 + for receiver_rank, _, bucket in buckets: + buckets_by_receiver_rank[receiver_rank].append(bucket) + if len(buckets_by_receiver_rank[receiver_rank]) > max_len: + max_len = len(buckets_by_receiver_rank[receiver_rank]) + + socket, socket_paths = self._bind_zmq_socket() + req_thread = threading.Thread( + target=req_func, + args=(socket_paths,), + ) + req_thread.start() + # The handle is self-contained for every handler, so one ZMQ send completes the handoff. + socket.send_pyobj(handle) - gidx = 0 - ret_code = torch.zeros((), device=self.device_manager.device_type, dtype=torch.int64) - buffer_b: torch.Tensor | None = None - try: - for i in range(max_len): - if i < len(receiver_rank_buckets) and not disable_h2d_buffer: - self._copy_to_buffer( - checkpoint_name, - receiver_rank_buckets[i][1], - h2d_buffer, - receiver_rank_buckets[i][0] if ranks else None, - ) - for receiver_rank, _buckets in buckets_by_receiver_rank.items(): - if i >= len(_buckets): - continue - bucket = _buckets[i] - alloc, reserved = ( - self.device_manager.device_module.memory_allocated() / 1024 / 1024, - self.device_manager.device_module.memory_reserved() / 1024 / 1024, - ) - self._logger_rank0( - f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} receiver_rank {receiver_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. " - f"Current device allocated {alloc:.2f} MB, " - f"reserved {reserved:.2f} MB." - ) - start = gidx % 2 * bucket_size - buffer_b: torch.Tensor = buffer[start : start + bucket.size] - if receiver_rank == self._rank: - if disable_h2d_buffer: - if p2p_update: - assert bucket == receiver_rank_buckets[i][1] - self._copy_to_buffer( - checkpoint_name, - bucket, - buffer_b, - receiver_rank_buckets[i][0] if p2p_update else None, - ) - else: - buffer_b.data.copy_(h2d_buffer[: bucket.size]) - dist.broadcast(buffer_b, src=receiver_rank, group=ranks_group) - resp = socket.recv() - if resp != b"": - msg = resp.decode("utf-8") - logger.error( - f"[rank{self._rank}] receive error response from rank {receiver_rank} for bucket {gidx} in checkpoint {checkpoint_name}: {msg}" + gidx = 0 + ret_code = torch.zeros((), device=self.device_manager.device_type, dtype=torch.int64) + buffer_b: torch.Tensor | None = None + try: + for i in range(max_len): + if i < len(receiver_rank_buckets) and not disable_h2d_buffer: + self._copy_to_buffer( + checkpoint_name, + receiver_rank_buckets[i][1], + h2d_buffer, + receiver_rank_buckets[i][0] if ranks else None, + ) + for receiver_rank, _buckets in buckets_by_receiver_rank.items(): + if i >= len(_buckets): + continue + bucket = _buckets[i] + alloc, reserved = ( + self.device_manager.device_module.memory_allocated() / 1024 / 1024, + self.device_manager.device_module.memory_reserved() / 1024 / 1024, + ) + self._logger_rank0( + f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} receiver_rank {receiver_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. " + f"Current device allocated {alloc:.2f} MB, " + f"reserved {reserved:.2f} MB." + ) + start = gidx % 2 * bucket_size + buffer_b: torch.Tensor = buffer[start : start + bucket.size] + if receiver_rank == self._rank: + if disable_h2d_buffer: + if p2p_update: + assert bucket == receiver_rank_buckets[i][1] + self._copy_to_buffer( + checkpoint_name, + bucket, + buffer_b, + receiver_rank_buckets[i][0] if p2p_update else None, ) - ret_code.fill_(1) - dist.all_reduce( - ret_code, op=torch.distributed.ReduceOp.SUM, group=ranks_group + else: + buffer_b.data.copy_(h2d_buffer[: bucket.size]) + dist.broadcast(buffer_b, src=receiver_rank, group=ranks_group) + resp = socket.recv() + if resp != b"": + msg = resp.decode("utf-8") + logger.error( + f"[rank{self._rank}] receive error response from rank {receiver_rank} for bucket {gidx} in checkpoint {checkpoint_name}: {msg}" ) - self.device_manager.device_module.synchronize() - if ret_code.item() != 0: - # quit early if any rank failed - socket.send_pyobj(RuntimeError("Some workers failed to update weights")) - raise RuntimeError("Failed to update weights due to remote errors") - socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size)) - gidx += 1 - - socket.recv() - device_mem = self.device_manager.device_module.mem_get_info() - logger.info( - f"[rank{self._rank}] weights broadcast done, device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, allocated memory: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, reserved memory: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" - ) - # Notify worker to release handle - socket.send_pyobj(None) - socket.recv() - # Set to None in correct order (views first, then base tensors) - del buffer_b, h2d_buffer, buffer, handle - self.device_manager.device_module.synchronize() - gc.collect() - self.device_manager.ipc_collect() - self.device_manager.device_module.empty_cache() - self.device_manager.device_module.synchronize() - - # Log actual memory usage - device_mem = self.device_manager.device_module.mem_get_info() - logger.info( - f"[rank{self._rank}] post-release: device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, " - f"allocated: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, " - f"reserved: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" - ) - # Notify worker to call post_hook - socket.send_pyobj(None) - socket.recv() - finally: - req_thread.join() - dist.barrier(group=ranks_group) - socket.close() - if p2p_update: - self._p2p_store.unregister_named_tensors([p2p_ipc_buffer_name]) - - self.device_manager.device_module.empty_cache() + ret_code.fill_(1) + dist.all_reduce(ret_code, op=torch.distributed.ReduceOp.SUM, group=ranks_group) + self.device_manager.device_module.synchronize() + if ret_code.item() != 0: + # quit early if any rank failed + socket.send_pyobj(RuntimeError("Some workers failed to update weights")) + raise RuntimeError("Failed to update weights due to remote errors") + socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size)) + gidx += 1 + + socket.recv() + device_mem = self.device_manager.device_module.mem_get_info() + logger.info( + f"[rank{self._rank}] weights broadcast done, device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, allocated memory: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, reserved memory: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" + ) + # Notify worker to release handle + socket.send_pyobj(None) + socket.recv() + # Set to None in correct order (views first, then base tensors) + del buffer_b, h2d_buffer, buffer, handle + self.device_manager.device_module.synchronize() + gc.collect() + self.device_manager.ipc_collect() + self.device_manager.device_module.empty_cache() + self.device_manager.device_module.synchronize() + + # Log actual memory usage + device_mem = self.device_manager.device_module.mem_get_info() + logger.info( + f"[rank{self._rank}] post-release: device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, " + f"allocated: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, " + f"reserved: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" + ) + # Notify worker to call post_hook + socket.send_pyobj(None) + socket.recv() finally: - transport.detach() + req_thread.join() + dist.barrier(group=ranks_group) + socket.close() + if p2p_update: + self._p2p_store.unregister_named_tensors([p2p_ipc_buffer_name]) + + self.device_manager.device_module.empty_cache() # we need this CLI entry point for compatibility with former versions diff --git a/checkpoint_engine/worker.py b/checkpoint_engine/worker.py index 4e773ad..180cc09 100644 --- a/checkpoint_engine/worker.py +++ b/checkpoint_engine/worker.py @@ -8,24 +8,24 @@ import zmq from checkpoint_engine.device_utils import DeviceManager, npu_generate_uuid -from checkpoint_engine.transport import ( - IpcWeightTransport, - WeightTransport, - XpuIpcWeightTransport, +from checkpoint_engine.ipc_handler import ( + IPCHandler, + TorchIPCHandler, + XpuIPCHandler, ) _WEIGHTS_TYPE = list[tuple[str, torch.Tensor]] -def _transport_for_handle(handle: object) -> WeightTransport: - """Pick the consumer-side transport based on the handle wire format. +def _ipc_handler_for_handle(handle: object) -> IPCHandler: + """Pick the consumer-side IPC handler based on the handle wire format. CUDA/NPU send a ``reduce_tensor`` tuple; XPU sends a dict tagged with its kind. """ - if isinstance(handle, dict) and handle.get("kind") == XpuIpcWeightTransport.kind: - return XpuIpcWeightTransport() - return IpcWeightTransport() + if isinstance(handle, dict) and handle.get("kind") == XpuIPCHandler.kind: + return XpuIPCHandler() + return TorchIPCHandler() class FlattenedTensorMetadata(TypedDict): @@ -63,11 +63,11 @@ def update_weights_from_ipc( socket.connect(zmq_handle) buffer: torch.Tensor | None = None device_manager = DeviceManager() - transport: WeightTransport | None = None + ipc_handler: IPCHandler | None = None try: ipc_handle = socket.recv_pyobj() - transport = _transport_for_handle(ipc_handle) - buffer = transport.attach(ipc_handle, device_id) + ipc_handler = _ipc_handler_for_handle(ipc_handle) + buffer = ipc_handler.attach(ipc_handle, device_id) assert buffer.dtype == torch.uint8 socket.send(b"") except Exception as e: @@ -96,8 +96,8 @@ def update_weights_from_ipc( device_manager.device_module.synchronize() released = True buffer = None - if transport is not None: - transport.detach() + if ipc_handler is not None: + ipc_handler.detach() gc.collect() device_manager.ipc_collect() @@ -125,8 +125,8 @@ def update_weights_from_ipc( finally: socket.close() del buffer - if transport is not None: - transport.detach() + if ipc_handler is not None: + ipc_handler.detach() gc.collect() device_manager.device_module.empty_cache() diff --git a/checkpoint_engine/xpu_ipc/__init__.py b/checkpoint_engine/xpu_ipc/__init__.py index bb1e461..bc9c7f0 100644 --- a/checkpoint_engine/xpu_ipc/__init__.py +++ b/checkpoint_engine/xpu_ipc/__init__.py @@ -1,14 +1,12 @@ -"""Cross-process device-buffer IPC for Intel XPU via SYCL ``ipc_memory``. +"""Cross-process device-buffer IPC for Intel XPU via SYCL IPC memory. -``sycl_ipc.cpp`` wraps ``ipc_memory`` (``get``/``open``/``close``), exported by -torch's own libsycl (oneAPI >= 2026.0); this module JIT-compiles it with -``-fsycl``. The handle is a self-contained portable byte blob (no dma-buf fd, no -offset to carry), so it rides the existing ZMQ channel like CUDA's -``reduce_tensor`` tuple -- see ``XpuIpcWeightTransport``. +``sycl_ipc.cpp`` wraps the SYCL IPC memory API (``get``/``open``/``close``), +exported by torch's own libsycl (oneAPI >= 2026.0); this module JIT-compiles it +with ``with_sycl``. The handle is a self-contained portable byte blob (no dma-buf +fd, no offset to carry), so it rides the existing ZMQ channel like CUDA's +``reduce_tensor`` tuple -- see ``XpuIPCHandler``. """ -from __future__ import annotations - import functools import glob import os @@ -25,85 +23,85 @@ import torch -def _find_sycl_include_dir() -> str | None: - """Locate a directory containing .""" - candidates: list[str] = [] - root = os.getenv("CMPLR_ROOT") - if root: - candidates.append(os.path.join(root, "include")) - # Common oneAPI install layouts (versioned + `latest` symlink). - candidates += sorted(glob.glob("/opt/intel/oneapi/compiler/*/include"), reverse=True) - # Derive from the discovered compiler (/bin/icpx -> /include), which - # covers a PATH-only icpx whose oneAPI root is outside /opt. - icpx = _find_icpx() - if icpx: - candidates.append(os.path.join(os.path.dirname(os.path.dirname(icpx)), "include")) - for inc in candidates: - if os.path.exists( - os.path.join(inc, "sycl", "ext", "oneapi", "experimental", "ipc_memory.hpp") - ): - return inc - return None +def _has_ipc_memory(icpx: str) -> bool: + """Whether this icpx ships the SYCL IPC memory header (oneAPI >= 2026.0).""" + root = os.path.dirname(os.path.dirname(icpx)) + header = os.path.join( + root, "include", "sycl", "ext", "oneapi", "experimental", "ipc_memory.hpp" + ) + return os.path.exists(header) + + +def _icpx_version_key(path: str) -> tuple[int, list[int]]: + """Sort key for ``.../compiler//bin/icpx``: numeric parts, so 2026.10 > 2026.9. + + The non-numeric ``latest`` symlink sorts first (it points at the newest install). + """ + version = path.split("/")[-3] + parts = version.split(".") + if not all(p.isdigit() for p in parts): + return (1, []) + return (0, [int(p) for p in parts]) def _find_icpx() -> str | None: - """Locate the icpx (SYCL) compiler needed for the -fsycl build.""" + """Locate an icpx (SYCL) compiler new enough for the SYCL IPC memory build. + + Compilers without the header are skipped: they build a device image that + torch's newer libsycl cannot load, aborting the process on dlopen rather + than raising something we could catch. + """ candidates: list[str] = [] root = os.getenv("CMPLR_ROOT") if root: candidates.append(os.path.join(root, "bin", "icpx")) - candidates += sorted(glob.glob("/opt/intel/oneapi/compiler/*/bin/icpx"), reverse=True) - for cand in candidates: - if os.path.exists(cand): - return cand + candidates += sorted( + glob.glob("/opt/intel/oneapi/compiler/*/bin/icpx"), + key=_icpx_version_key, + reverse=True, + ) # Fallback to PATH: covers oneAPI layouts outside /opt and a sourced setvars.sh # that puts icpx on PATH without exporting CMPLR_ROOT. - return shutil.which("icpx") + which = shutil.which("icpx") + if which: + candidates.append(which) + return next( + (c for c in candidates if os.path.exists(c) and _has_ipc_memory(c)), + None, + ) @functools.lru_cache(maxsize=1) -def load_ext() -> ModuleType: - """JIT-compile (``-fsycl``, linking torch's libsycl) and cache the SYCL IPC extension. +def load_ext() -> "ModuleType": + """JIT-compile (``with_sycl``, linking torch's libsycl) and cache the SYCL IPC extension. Raises on any failure; callers treat an exception as "XPU IPC unavailable". """ icpx = _find_icpx() if icpx is None: - raise RuntimeError("icpx (oneAPI SYCL compiler) not found; cannot build XPU IPC extension") - icx = os.path.join(os.path.dirname(icpx), "icx") + raise RuntimeError( + "no icpx with SYCL ipc_memory support found (needs oneAPI >= 2026.0); " + "cannot build XPU IPC extension" + ) from torch.utils.cpp_extension import load src = Path(__file__).with_name("sycl_ipc.cpp") - sycl_include_flags: list[str] = [] - inc = _find_sycl_include_dir() - if inc: - sycl_include_flags = [f"-I{inc}", f"-I{os.path.join(inc, 'sycl')}"] - - # torch.utils.cpp_extension picks the compiler from CC/CXX. Force icx/icpx for the - # -fsycl build: a conda/CI env often exports CXX=g++ (gxx_linux-64), which cannot - # compile -fsycl, and setdefault would keep it. Save/restore the process env. - prev_cc, prev_cxx = os.environ.get("CC"), os.environ.get("CXX") - if prev_cxx and os.path.realpath(prev_cxx) != os.path.realpath(icpx): - logger.debug(f"overriding CXX={prev_cxx!r} with icpx for the SYCL IPC build ({icpx})") - os.environ["CC"], os.environ["CXX"] = icx, icpx + # with_sycl=True supplies the SYCL include paths and device link, but torch invokes + # a bare "icpx", so it must be on PATH; keep -O2 or the host object is built -O0. + prev_path = os.environ.get("PATH", "") + os.environ["PATH"] = os.path.dirname(icpx) + os.pathsep + prev_path try: - # Do NOT pin -std: torch.utils.cpp_extension injects the standard its ATen - # headers require, and a pin here would override it. module = load( name="checkpoint_engine_sycl_ipc", sources=[str(src)], - extra_cflags=["-fsycl", "-O2", *sycl_include_flags], - extra_ldflags=["-fsycl"], + extra_cflags=["-O2"], + with_sycl=True, verbose=False, ) finally: - for var, prev in (("CC", prev_cc), ("CXX", prev_cxx)): - if prev is None: - os.environ.pop(var, None) - else: - os.environ[var] = prev + os.environ["PATH"] = prev_path return module @@ -158,6 +156,6 @@ def close_handle(ptr: int) -> None: load_ext().ipc_close_handle(ptr) -def wrap_tensor(ptr: int, nbytes: int, device: int) -> torch.Tensor: +def wrap_tensor(ptr: int, nbytes: int, device: int) -> "torch.Tensor": """Wrap an IPC-mapped device pointer as a non-owning torch XPU uint8 tensor.""" return load_ext().ipc_wrap_tensor(ptr, nbytes, device) diff --git a/checkpoint_engine/xpu_ipc/sycl_ipc.cpp b/checkpoint_engine/xpu_ipc/sycl_ipc.cpp index f39de46..ef9b933 100644 --- a/checkpoint_engine/xpu_ipc/sycl_ipc.cpp +++ b/checkpoint_engine/xpu_ipc/sycl_ipc.cpp @@ -16,7 +16,15 @@ #include #include +// Upstream split this API (functions -> ipc::memory, types -> parent ipc) and +// deprecated flat ipc_memory; no oneAPI release ships it yet, so probe for it. +#if __has_include() +namespace ipc = sycl::ext::oneapi::experimental::ipc::memory; +namespace ipc_types = sycl::ext::oneapi::experimental::ipc; +#else namespace ipc = sycl::ext::oneapi::experimental::ipc_memory; +namespace ipc_types = sycl::ext::oneapi::experimental::ipc_memory; +#endif namespace { @@ -28,10 +36,10 @@ std::vector to_bytes(const std::vector& in) { // Exporter-side handles kept alive until ipc_release_handle(). We must not // ipc::put() before the consumer opens: under the UR level-zero-v2 adapter // put_ipc_handle frees the exporter fd and can race the consumer's open. -// ipc::handle is a copyable, non-owning value (freed only via ipc::put), so +// The handle is a copyable, non-owning value (freed only via ipc::put), so // storing it by value needs no manual new/delete. std::mutex g_handles_mu; -std::unordered_map g_handles; +std::unordered_map g_handles; } // namespace @@ -39,8 +47,8 @@ std::unordered_map g_handles; // are fine -- the offset is in the blob). Handle retained until ipc_release_handle(). std::vector ipc_get_handle(uintptr_t ptr) { sycl::context ctx = c10::xpu::get_device_context(); - ipc::handle h = ipc::get(reinterpret_cast(ptr), ctx); - ipc::handle_data_t data = h.data(); // owning copy of the blob, independent of `h` + ipc_types::handle h = ipc::get(reinterpret_cast(ptr), ctx); + ipc_types::handle_data_t data = h.data(); // owning copy of the blob, independent of `h` { std::lock_guard lk(g_handles_mu); auto it = g_handles.find(ptr); @@ -58,7 +66,7 @@ std::vector ipc_get_handle(uintptr_t ptr) { // Release the exporter handle from ipc_get_handle(ptr); no-op if unregistered. // Call only after all consumers have opened their mappings (see level-zero-v2 note). void ipc_release_handle(uintptr_t ptr) { - std::optional h; + std::optional h; { std::lock_guard lk(g_handles_mu); auto it = g_handles.find(ptr); diff --git a/tests/test_device_manager.py b/tests/test_device_manager.py index aa2046d..f7cdb03 100644 --- a/tests/test_device_manager.py +++ b/tests/test_device_manager.py @@ -58,10 +58,19 @@ def test_ipc_collect_present_is_called(): assert called == [True] -def test_ipc_collect_absent_is_noop(): - # torch.xpu has no ipc_collect attribute; ipc_collect() must not raise. +def test_ipc_collect_xpu_is_noop(): + # SYCL frees on close_handle, so XPU has no handle cache to collect (and + # torch.xpu has no ipc_collect); it must be a no-op, not an error. dm = _make_manager("xpu", SimpleNamespace()) - dm.ipc_collect() # no attribute -> silent no-op + dm.ipc_collect() + + +def test_ipc_collect_rejects_unsupported_device(): + # An unsupported backend must fail loudly rather than silently skipping the + # collect, matching backend/transfer_engine_protocol/_setup_device_module. + dm = _make_manager("tpu", SimpleNamespace()) + with pytest.raises(TypeError, match="not supported"): + dm.ipc_collect() @pytest.mark.parametrize( @@ -135,7 +144,7 @@ def test_real_xpu_device_manager(): # ipc_collect must be a harmless no-op (torch.xpu has no ipc_collect). dm.ipc_collect() # XPU has no torch-native device-tensor IPC, but checkpoint-engine ships its own - # native SYCL ipc_memory transport; supports_device_ipc() must agree with whether + # native SYCL IPC memory handler; supports_device_ipc() must agree with whether # that extension can actually be built/loaded in this environment. from checkpoint_engine import xpu_ipc diff --git a/tests/test_transport.py b/tests/test_ipc_handler.py similarity index 70% rename from tests/test_transport.py rename to tests/test_ipc_handler.py index 24b21be..7acbdfd 100644 --- a/tests/test_transport.py +++ b/tests/test_ipc_handler.py @@ -1,4 +1,4 @@ -"""Unit tests for the weight-transport seam (CPU-only, no accelerator required). +"""Unit tests for the IPC-handler seam (CPU-only, no accelerator required). The zero-copy device handoff itself is exercised by a hardware-gated end-to-end test on real XPU/CUDA. Here we cover the dispatch logic and wire formats. @@ -9,12 +9,12 @@ import pytest -from checkpoint_engine.transport import ( - IpcWeightTransport, - XpuIpcWeightTransport, - build_transport, +from checkpoint_engine.ipc_handler import ( + TorchIPCHandler, + XpuIPCHandler, + build_ipc_handler, ) -from checkpoint_engine.worker import _transport_for_handle +from checkpoint_engine.worker import _ipc_handler_for_handle def _dm(device_type: str) -> object: @@ -23,28 +23,28 @@ def _dm(device_type: str) -> object: @pytest.mark.parametrize( "device_type,expected", - [("cuda", IpcWeightTransport), ("npu", IpcWeightTransport), ("xpu", XpuIpcWeightTransport)], + [("cuda", TorchIPCHandler), ("npu", TorchIPCHandler), ("xpu", XpuIPCHandler)], ) -def test_build_transport_dispatch(device_type: str, expected: type): - assert isinstance(build_transport(_dm(device_type)), expected) +def test_build_ipc_handler_dispatch(device_type: str, expected: type): + assert isinstance(build_ipc_handler(_dm(device_type)), expected) def test_consumer_dispatch_by_handle_shape(): # CUDA/NPU send a reduce_tensor tuple; XPU sends a tagged dict. tuple_handle = (lambda *a: None, (1, 2, 3)) - assert isinstance(_transport_for_handle(tuple_handle), IpcWeightTransport) + assert isinstance(_ipc_handler_for_handle(tuple_handle), TorchIPCHandler) - xpu_handle = {"kind": XpuIpcWeightTransport.kind, "handle_bytes": b"", "nbytes": 0} - assert isinstance(_transport_for_handle(xpu_handle), XpuIpcWeightTransport) + xpu_handle = {"kind": XpuIPCHandler.kind, "handle_bytes": b"", "nbytes": 0} + assert isinstance(_ipc_handler_for_handle(xpu_handle), XpuIPCHandler) # An unrelated dict must not be mistaken for the XPU handle. - assert isinstance(_transport_for_handle({"foo": "bar"}), IpcWeightTransport) + assert isinstance(_ipc_handler_for_handle({"foo": "bar"}), TorchIPCHandler) -def test_ipc_transport_export_uses_reduce_tensor(): +def test_torch_handler_export_uses_reduce_tensor(): sentinel = ("REDUCED",) - with patch("checkpoint_engine.transport.reduce_tensor", return_value=sentinel) as m: - t = IpcWeightTransport() + with patch("checkpoint_engine.ipc_handler.reduce_tensor", return_value=sentinel) as m: + t = TorchIPCHandler() out = t.export(SimpleNamespace()) assert out is sentinel m.assert_called_once() @@ -55,7 +55,7 @@ def test_xpu_export_returns_self_contained_handle(): # no companion socket. Mock the native extension so this runs on CPU CI. buffer = SimpleNamespace(data_ptr=lambda: 0xDEAD, nbytes=256) with patch("checkpoint_engine.xpu_ipc.get_handle", return_value=b"HANDLE") as m: - handle = XpuIpcWeightTransport().export(buffer) + handle = XpuIPCHandler().export(buffer) m.assert_called_once_with(0xDEAD) assert handle == {"kind": "xpu_sycl", "handle_bytes": b"HANDLE", "nbytes": 256} @@ -68,7 +68,7 @@ def test_xpu_export_defers_release_until_detach(): patch("checkpoint_engine.xpu_ipc.get_handle", return_value=b"H"), patch("checkpoint_engine.xpu_ipc.release_handle") as release, ): - t = XpuIpcWeightTransport() + t = XpuIPCHandler() t.export(buffer) release.assert_not_called() # not released during export t.detach() @@ -78,13 +78,13 @@ def test_xpu_export_defers_release_until_detach(): release.assert_called_once_with(0xBEEF) -def test_xpu_transport_detach_is_safe_when_unused(): +def test_xpu_handler_detach_is_safe_when_unused(): # detach() before any export/attach must not raise (and must not touch the ext). with ( patch("checkpoint_engine.xpu_ipc.release_handle") as release, patch("checkpoint_engine.xpu_ipc.close_handle") as close, ): - XpuIpcWeightTransport().detach() + XpuIPCHandler().detach() release.assert_not_called() close.assert_not_called() @@ -96,14 +96,14 @@ def test_xpu_consumer_detach_closes_opened_mapping(): with ( patch("checkpoint_engine.xpu_ipc.open_handle", return_value=0x7000), patch("checkpoint_engine.xpu_ipc.wrap_tensor") as wrap, - patch("checkpoint_engine.transport.torch.xpu.synchronize"), + patch("checkpoint_engine.ipc_handler.torch.xpu.synchronize"), patch("checkpoint_engine.xpu_ipc.close_handle") as close, patch("checkpoint_engine.xpu_ipc.release_handle") as release, ): import torch as _torch wrap.return_value = SimpleNamespace(dtype=_torch.uint8) - t = XpuIpcWeightTransport() + t = XpuIPCHandler() t.attach(handle, device_id=0) t.detach() close.assert_called_once_with(0x7000) diff --git a/tests/test_p2p_guard.py b/tests/test_p2p_guard.py index 0b2b2d8..7888bb8 100644 --- a/tests/test_p2p_guard.py +++ b/tests/test_p2p_guard.py @@ -6,7 +6,7 @@ """ from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -29,17 +29,34 @@ def _ps_with_device(device_type: str, *, supports_ipc: bool, supports_p2p: bool) def test_p2p_update_rejected_on_xpu(): ps = _ps_with_device("xpu", supports_ipc=True, supports_p2p=False) + ipc_handler = MagicMock() with ( patch.object(dist, "is_initialized", return_value=True), pytest.raises(RuntimeError, match=r"P2P weight update .* is not supported"), ): - ps._update_per_bucket("ckpt", req_func=lambda _paths: None, ranks_group=None, ranks=[0]) + ps._update_per_bucket( + "ckpt", + req_func=lambda _paths: None, + ipc_handler=ipc_handler, + ranks_group=None, + ranks=[0], + ) + # The guard must fire before any handle is exported. + ipc_handler.export.assert_not_called() def test_ipc_unavailable_rejected(): ps = _ps_with_device("xpu", supports_ipc=False, supports_p2p=False) + ipc_handler = MagicMock() with ( patch.object(dist, "is_initialized", return_value=True), pytest.raises(RuntimeError, match="cross-process device-tensor IPC"), ): - ps._update_per_bucket("ckpt", req_func=lambda _paths: None, ranks_group=None, ranks=None) + ps._update_per_bucket( + "ckpt", + req_func=lambda _paths: None, + ipc_handler=ipc_handler, + ranks_group=None, + ranks=None, + ) + ipc_handler.export.assert_not_called() diff --git a/tests/test_xpu_ipc.py b/tests/test_xpu_ipc.py index 992e818..598c458 100644 --- a/tests/test_xpu_ipc.py +++ b/tests/test_xpu_ipc.py @@ -1,4 +1,4 @@ -"""Hardware-gated tests for the native SYCL ipc_memory transport on Intel XPU. +"""Hardware-gated tests for the native SYCL IPC memory handler on Intel XPU. These are skipped unless an Intel GPU is present and the SYCL IPC extension can be built (needs an oneAPI ``icpx`` with SYCL ipc_memory support). They are marked @@ -119,7 +119,7 @@ def test_sycl_ipc_interior_pointer_offset_preserved(): @skip_no_xpu_ipc def test_sycl_ipc_cross_process_broadcast(): - """Full ParameterServer broadcast -> colocated worker over the XPU SYCL transport.""" + """Full ParameterServer broadcast -> colocated worker over the XPU SYCL handler.""" from torch.multiprocessing import get_context os.environ.setdefault("RANK", "0") diff --git a/tests/test_xpu_parity.py b/tests/test_xpu_parity.py index 1b2e7d2..63d28e6 100644 --- a/tests/test_xpu_parity.py +++ b/tests/test_xpu_parity.py @@ -12,6 +12,7 @@ """ import os +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -51,42 +52,54 @@ def test_open_handle_passes_bytes_and_device_through(): # ------------------------------------------------------------------------------ -# load_ext: the -fsycl build must force CC/CXX to icx/icpx (a conda/CI env often -# exports CXX=g++, which cannot compile -fsycl) and restore the env afterwards. +# load_ext: with_sycl=True makes torch shell out to a bare "icpx", so the located +# compiler must be on PATH for the build and PATH restored afterwards. # ------------------------------------------------------------------------------ -def test_load_ext_forces_icpx_over_existing_gpp_and_restores_env( +def test_load_ext_puts_icpx_on_path_and_restores_it( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("CXX", "/usr/bin/g++") # what conda gxx_linux-64 exports - monkeypatch.delenv("CC", raising=False) + monkeypatch.setenv("PATH", "/usr/bin") - captured: dict[str, str | None] = {} + captured: dict[str, object] = {} def fake_load(**kwargs: object) -> MagicMock: # Record what torch.utils.cpp_extension.load would see. - captured["CC"] = os.environ.get("CC") - captured["CXX"] = os.environ.get("CXX") + captured["PATH"] = os.environ.get("PATH") + captured["kwargs"] = kwargs return MagicMock() xpu_ipc.load_ext.cache_clear() try: with ( patch("checkpoint_engine.xpu_ipc._find_icpx", return_value="/opt/oneapi/bin/icpx"), - patch("checkpoint_engine.xpu_ipc._find_sycl_include_dir", return_value=None), patch("torch.utils.cpp_extension.load", side_effect=fake_load), ): xpu_ipc.load_ext() finally: xpu_ipc.load_ext.cache_clear() - # During the build the SYCL compiler must win over the inherited g++. - assert captured["CXX"] == "/opt/oneapi/bin/icpx" - assert captured["CC"] == "/opt/oneapi/bin/icx" + # torch runs `icpx --version`, so its directory must lead PATH during the build. + assert captured["PATH"] == "/opt/oneapi/bin:/usr/bin" + # The SYCL toolchain comes from with_sycl; -O2 must stay or the host object is -O0. + kwargs = captured["kwargs"] + assert kwargs["with_sycl"] is True + assert kwargs["extra_cflags"] == ["-O2"] # ...and the caller's environment must be restored afterwards. - assert os.environ["CXX"] == "/usr/bin/g++" - assert "CC" not in os.environ + assert os.environ["PATH"] == "/usr/bin" + + +def test_icpx_version_key_orders_numerically() -> None: + # Lexicographic sort would rank 2026.9 above 2026.10; the key must compare + # version parts as numbers so the newest install really wins. + paths = [f"/opt/intel/oneapi/compiler/{v}/bin/icpx" for v in ("2026.9", "2026.10", "2025.3")] + ordered = sorted(paths, key=xpu_ipc._icpx_version_key, reverse=True) + assert [p.split("/")[-3] for p in ordered] == ["2026.10", "2026.9", "2025.3"] + # The non-numeric `latest` symlink points at the newest install, so it wins. + with_latest = [*paths, "/opt/intel/oneapi/compiler/latest/bin/icpx"] + best = max(with_latest, key=xpu_ipc._icpx_version_key) + assert best.split("/")[-3] == "latest" def test_find_icpx_falls_back_to_path(monkeypatch: pytest.MonkeyPatch) -> None: @@ -97,11 +110,41 @@ def test_find_icpx_falls_back_to_path(monkeypatch: pytest.MonkeyPatch) -> None: with ( patch("checkpoint_engine.xpu_ipc.glob.glob", return_value=[]), patch("checkpoint_engine.xpu_ipc.shutil.which", return_value="/custom/bin/icpx") as which, + patch("checkpoint_engine.xpu_ipc.os.path.exists", return_value=True), + patch("checkpoint_engine.xpu_ipc._has_ipc_memory", return_value=True), ): assert xpu_ipc._find_icpx() == "/custom/bin/icpx" which.assert_called_once_with("icpx") +def test_find_icpx_skips_compiler_without_ipc_memory(monkeypatch: pytest.MonkeyPatch) -> None: + # An icpx older than oneAPI 2026.0 builds a device image that torch's libsycl + # cannot load, aborting the process (SIGABRT) on dlopen -- uncatchable from + # Python. Such compilers must be rejected up front, before any build. + monkeypatch.setenv("CMPLR_ROOT", "/opt/intel/oneapi/compiler/2025.3") + with ( + patch("checkpoint_engine.xpu_ipc.glob.glob", return_value=[]), + patch("checkpoint_engine.xpu_ipc.shutil.which", return_value=None), + patch("checkpoint_engine.xpu_ipc.os.path.exists", return_value=True), + patch("checkpoint_engine.xpu_ipc._has_ipc_memory", return_value=False), + ): + assert xpu_ipc._find_icpx() is None + + +def test_has_ipc_memory_probes_the_header(tmp_path: Path) -> None: + # The discriminator is the header's presence in the compiler's own tree + # (/bin/icpx -> /include/sycl/.../ipc_memory.hpp). + icpx = tmp_path / "bin" / "icpx" + icpx.parent.mkdir(parents=True) + icpx.touch() + assert xpu_ipc._has_ipc_memory(str(icpx)) is False + + header = tmp_path / "include" / "sycl" / "ext" / "oneapi" / "experimental" / "ipc_memory.hpp" + header.parent.mkdir(parents=True) + header.touch() + assert xpu_ipc._has_ipc_memory(str(icpx)) is True + + def test_is_available_caches_only_success_and_retries_failure( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -156,43 +199,98 @@ def test_use_backend_none_keeps_default_torch_backend(): # ------------------------------------------------------------------------------ -# _update_per_bucket: the exported IPC handle is retained until detach(). A -# failure after export but before the broadcast loop (e.g. the ZMQ bind) must -# still release it -- otherwise the exporter handle leaks on every failed update. +# update(): the exported IPC handle is retained until detach(). A failure inside +# _update_per_bucket must still release it -- otherwise the exporter handle leaks +# on every failed weight update. # ------------------------------------------------------------------------------ -def test_update_per_bucket_detaches_transport_on_early_failure(): +def test_update_releases_ipc_handle_when_update_fails(): + from checkpoint_engine.ipc_handler import IPCHandler from checkpoint_engine.ps import ParameterServer + # A real IPCHandler (not a MagicMock) so detach() must be reached through the + # context manager's __exit__ rather than by mocked attribute access. + class RecordingHandler(IPCHandler): + def __init__(self) -> None: + self.detached = 0 + + def export(self, buffer: object) -> dict: + return {"kind": "fake"} + + def attach(self, handle: object, device_id: int) -> None: + raise AssertionError("attach is the consumer side; not used here") + + def detach(self) -> None: + self.detached += 1 + ps = ParameterServer.__new__(ParameterServer) ps._rank = 0 + ps._auto_pg = False ps.device_manager = SimpleNamespace( device_type="cpu", - supports_device_ipc=lambda: True, - supports_device_p2p=lambda: False, + device_module=SimpleNamespace( + empty_cache=lambda: None, + memory_allocated=lambda: 0, + memory_reserved=lambda: 0, + ), ) - ps._current_global_parameter_metas = {0: object()} - ps._local_rdma_devices = None - ps._remote_rdma_devices = None + handler = RecordingHandler() - fake_transport = MagicMock() - fake_transport.export.return_value = {"kind": "fake"} + with ( + patch.object(dist, "is_initialized", return_value=True), + patch("checkpoint_engine.ps.build_ipc_handler", return_value=handler), + patch.object( + ps, "_update_per_bucket", side_effect=RuntimeError("update failed") + ) as per_bucket, + pytest.raises(RuntimeError, match="update failed"), + ): + ps.update("ckpt", req_func=lambda _paths: None) + + # The handler is handed to _update_per_bucket and released regardless of outcome. + assert handler in per_bucket.call_args.args + assert handler.detached == 1 + + +def test_update_releases_ipc_handle_on_success(): + from checkpoint_engine.ipc_handler import IPCHandler + from checkpoint_engine.ps import ParameterServer + + class RecordingHandler(IPCHandler): + def __init__(self) -> None: + self.detached = 0 + + def export(self, buffer: object) -> dict: + return {"kind": "fake"} + + def attach(self, handle: object, device_id: int) -> None: + raise AssertionError("attach is the consumer side; not used here") + + def detach(self) -> None: + self.detached += 1 + + ps = ParameterServer.__new__(ParameterServer) + ps._rank = 0 + ps._auto_pg = False + ps.device_manager = SimpleNamespace( + device_type="cpu", + device_module=SimpleNamespace( + empty_cache=lambda: None, + memory_allocated=lambda: 0, + memory_reserved=lambda: 0, + ), + ) + handler = RecordingHandler() with ( patch.object(dist, "is_initialized", return_value=True), - patch("checkpoint_engine.ps.build_transport", return_value=fake_transport), - patch("checkpoint_engine.ps._gen_h2d_buckets", return_value=[]), - patch.object(ps, "_detect_bucket_size", return_value=(16, False)), - patch.object(ps, "_bind_zmq_socket", side_effect=RuntimeError("bind failed")), - pytest.raises(RuntimeError, match="bind failed"), + patch("checkpoint_engine.ps.build_ipc_handler", return_value=handler), + patch.object(ps, "_update_per_bucket"), + patch.object(ps, "store_based_barrier"), ): - ps._update_per_bucket("ckpt", req_func=lambda _paths: None, ranks_group=None, ranks=None) + ps.update("ckpt", req_func=lambda _paths: None) - # export happened, so the exporter handle is live and must be released even - # though the failure struck before the broadcast loop's own cleanup. - fake_transport.export.assert_called_once() - fake_transport.detach.assert_called_once_with() + assert handler.detached == 1 def test_register_checkpoint_disables_inplace_pin_on_xpu():