diff --git a/lib/gpu_memory_service/GMS_MULTI_DEVICE.md b/lib/gpu_memory_service/GMS_MULTI_DEVICE.md new file mode 100644 index 000000000000..fee099381a20 --- /dev/null +++ b/lib/gpu_memory_service/GMS_MULTI_DEVICE.md @@ -0,0 +1,185 @@ +# GMS Multiple Device Enablement — Design Proposal + +## 1. Overview + +GPU Memory Service (GMS) manages cross-process virtual memory on accelerators for weight /KV-cache sharing in Dynamo inference clusters. + +This design introduces a **device-agnostic VMM abstraction layer** so that Intel XPU can plug in alongside the existing CUDA path without touching the server, client, or snapshot logic. + +--- + +## 2. Goals + +| # | Goal | Status | +|---|------|--------| +| G1 | Define a vendor-neutral `VMMDevice` ABC covering all device operations GMS needs | ✅ Phase 1 | +| G2 | Wrap existing CUDA helpers in a `CudaVMM` class that inherits the ABC | ✅ Phase 1 | +| G3 | Process-global VMM singleton via `init_vmm()` / `get_vmm()`, CLI `--device-type` | ✅ Phase 1 | +| G4 | Implement `XpuVMM` | ⬜ Phase 2 | +| G5 | Add XPU torch mempool dispatch via `torch.xpu` APIs | ⬜ Phase 2 | +| G6 | Enable snapshot save/load for XPU | ⬜ Phase 2 | + +--- + +## 3. Architecture (Phase 1 — Delivered) + +``` +┌──────────────────────────────────────────────────────┐ +│ CLI / Supervisor │ +│ server.py ──→ --device-type {cuda,xpu} │ +└────────────────────────┬─────────────────────────────┘ + │ VMMDeviceType enum + ▼ +┌──────────────────────────────────────────────────────┐ +│ common/vmm/__init__.py | +│ Singleton instance: | +│ init_vmm(device_type) │ +│ get_vmm() → VMMDevice │ +│ │ +│ ┌───────────────┐ ┌───────────────┐ │ +│ │ CudaVMM │ │ XpuVMM │ │ +│ │ (cuda_utils) │ │ (xpu_utils) │ │ +│ └───────────────┘ └───────────────┘ │ +└──────────────────────────────────────────────────────┘ + │ + ┌──────────────┼──────────────┐ + ▼ ▼ ▼ + GMSRPCServer GMS GMSAllocationManager + (server/rpc.py) (server/gms.py) (server/allocations.py) + │ + ▼ + GMSClientMemoryManager → GMSStorageClient (snapshot) + (client/memory_manager.py) (snapshot/storage_client.py) +``` + +### 3.1 New Module: `common/vmm/` + +| File | Purpose | +|------|---------| +| `__init__.py` | `VMMDeviceType` enum, `init_vmm()` singleton initializer, `get_vmm()` accessor, `get_vmm_device_type()` | +| `device.py` | `VMMDevice` — `abc.ABC` with 24 `@abstractmethod` vendor-neutral methods | +| `cuda_utils.py` | Module-level CUDA helpers (unchanged logic) + `CudaVMM(VMMDevice)` class | + +- `init_vmm(device_type)`: explicit process-level backend selection, e.g. + CLI `--device-type`. On a mixed node (e.g., CUDA and other device coexist), + --device-type device_name overrides the auto-detection because it calls + init_vmm(VMMDeviceType.DeviceName) at process startup — before any get_vmm() lazy path fires. +- `get_vmm()`: returns the singleton and may lazily initialize the default/autodetected + backend for public client/integration compatibility via init_vmm(detected_device_type). +- `vmm.ensure_initialized()`: initializes the selected backend driver/runtime, e.g. CUDA `cuInit` + +### 3.2 `VMMDevice` ABC Surface (24 methods) + +``` +Category Method +───────────────────────── ───────────────────────────────────── +Driver lifecycle ensure_initialized() + synchronize() +Discovery / sizing list_devices() → list[int] + get_allocation_granularity(device) → int +Physical memory create_tolerate_oom(size, device) → (bool, int) + release(handle) +Shareable handles export_to_shareable_handle(handle) → int (FD) + import_shareable_handle_close_fd(fd) → int +VA space + mapping address_reserve(size, granularity) → int + address_free(va, size) + map(va, size, handle) + unmap(va, size) + set_access(va, size, device, access) +Monitoring / sizing device_memory_info(device) → (free_bytes, total_bytes) +Pointer validation validate_pointer(va) +Runtime helpers runtime_check_result(result, name) + runtime_set_device(device) + host_register(ptr, size) + host_unregister(ptr) +Stream management stream_create_nonblocking() → opaque + stream_destroy(stream) + stream_synchronize(stream) +Async copy memcpy_h2d_async(dst, src, size, stream) + memcpy_d2h_async(dst, src, size, stream) +``` + +### 3.3 Singleton Lifecycle + +``` +Process startup (cli/server.py, cli/runner.py, snapshot loader/saver): + init_vmm(VMMDeviceType.from_str(args.device_type)) + +Any module needing VMM: + vmm = get_vmm() # returns cached singleton + vmm.runtime_set_device(device) + vmm.list_devices() + ... +``` + +The singleton is immutable after initialization. Conflicting re-initialization +raises `RuntimeError`. Thread-safe via `threading.Lock`. + +### 3.4 CLI Argument Propagation + +``` +gms-server (supervisor) --device-type cuda|xpu + └─→ init_vmm(device_type) + └─→ gms-server (per-device) --device-type (forwarded to child process) + └─→ init_vmm(device_type) + └─→ get_vmm() used throughout server, allocations, client +``` + +No constructor threading — every module calls `get_vmm()` directly. + +## 4. Files Changed (Phase 1) + +| File | Change | +|------|--------| +| `cli/server.py` | Parse `--device-type`, call `init_vmm()`, use `get_vmm().list_devices()` | +| `cli/runner.py` | Call `init_vmm(config.device_type)` at startup | +| `cli/snapshot/loader.py` | Call `init_vmm(device_type)`, `vmm = get_vmm()` in helpers | +| `cli/snapshot/saver.py` | Call `init_vmm(device_type)`, `vmm = get_vmm()` in helpers | +| `client/memory_manager.py` | `self._vmm = get_vmm()`, property `device_type` via `get_vmm_device_type()` | +| `client/torch/allocator.py` | CUDA-only guards via `get_vmm_device_type()` | +| `common/vmm/__init__.py` | **NEW** — enum + singleton | +| `common/vmm/device.py` | **NEW** — VMMDevice ABC | +| `common/vmm/cuda_utils.py` | **NEW** (moved from `common/cuda_utils.py`) — CUDA helpers + CudaVMM class | +| `common/utils.py` | Add `align_to_granularity()` utility | +| `server/allocations.py` | `self._vmm = get_vmm()` | +| `server/gms.py` | No longer forwards device params — singleton | +| `server/rpc.py` | No longer forwards device params — singleton | +| `snapshot/storage_client.py` | No longer accepts device_type — uses singleton | +| `snapshot/disk.py` | Fix import path: `common.vmm` | +| `snapshot/backends/nixl_staging.py` | Fix import path: `common.vmm` | +| `snapshot/backends/pinned_host.py` | Fix import path: `common.vmm` | +| `tests/report_pytest_markers.py` | Update stub module list | + +--- + +## 5. Phase 2 — XPU Implementation + +### 5.1 `XpuVMM(VMMDevice)` + +Implement the methods on the OneAPI/Torch runtime. + +### 5.2 Torch Allocator Dispatch + +The PyTorch front door (`client/torch/`) routes torch tensor allocations through GMS via a pluggable allocator inside a `gms_use_mem_pool(tag)` context; the dispatch key is the existing `get_vmm_device_type()`. + +- **`extensions/allocator.cpp` + `setup.py` — no change.** `my_malloc(ssize_t, int, void* stream)` forwards the opaque stream/queue to Python *without dereferencing it*, so it is ABI-compatible with `XPUPluggableAllocator`'s + `void* alloc_fn(size_t, int, sycl::queue*)`. The **same** `_allocator_ext.so` + and `"my_malloc"` / `"my_free"` symbols serve both backends. +- **`client/torch/allocator.py` — only change.** Swap the four CUDA-hardcoded + spots for device dispatch (`torch.cuda.*`, `torch.xpu.*`): + + | Spot | CUDA | XPU | + |------|------|-----| + | `_ensure_callbacks_initialized` | `torch.cuda.CUDAPluggableAllocator` | `torch.xpu.memory.XPUPluggableAllocator` | + | `_create_mem_pool` | `torch.cuda.memory.MemPool` | `torch.xpu.memory.MemPool` | + | `gms_use_mem_pool` | `torch.cuda.use_mem_pool` | `torch.xpu.memory.use_mem_pool` | + | `prune_allocations` | `torch.cuda.synchronize` | `torch.xpu.synchronize` | + + The torch APIs are symmetric, so a small accessor returning the active device's + `(PluggableAllocator, MemPool, use_mem_pool, synchronize)` is sufficient. + +--- + +## 6. Backward Compatibility + +CUDA path remains unchanged. All existing `--device-type cuda` deployments work identically. The default value is `cuda` everywhere. diff --git a/lib/gpu_memory_service/cli/args.py b/lib/gpu_memory_service/cli/args.py index e543832b1450..c7f23bf741af 100644 --- a/lib/gpu_memory_service/cli/args.py +++ b/lib/gpu_memory_service/cli/args.py @@ -9,6 +9,7 @@ from typing import Optional from gpu_memory_service.common.utils import GMS_TAGS, get_socket_path +from gpu_memory_service.common.vmm import VMMDeviceType logger = logging.getLogger(__name__) @@ -23,6 +24,7 @@ class Config: alloc_retry_interval: float alloc_retry_timeout: Optional[float] verbose: bool + device_type: VMMDeviceType def parse_args(argv: Optional[list[str]] = None) -> list[Config]: @@ -71,6 +73,13 @@ def parse_args(argv: Optional[list[str]] = None) -> list[Config]: help="Max seconds to wait for allocation retries before failing (default: 60.0). " "Pass an explicit large value if you need essentially-unbounded retry.", ) + parser.add_argument( + "--device-type", + type=str, + default=VMMDeviceType.CUDA.value, + choices=[d.value for d in VMMDeviceType], + help="VMM device type (vendor driver) to use (default: cuda).", + ) args = parser.parse_args(argv) @@ -94,6 +103,7 @@ def parse_args(argv: Optional[list[str]] = None) -> list[Config]: alloc_retry_interval=args.alloc_retry_interval, alloc_retry_timeout=args.alloc_retry_timeout, verbose=args.verbose, + device_type=VMMDeviceType.from_str(args.device_type), ) for tag in tags ] diff --git a/lib/gpu_memory_service/cli/runner.py b/lib/gpu_memory_service/cli/runner.py index 5353c57a1e2d..be71f44bbf7b 100644 --- a/lib/gpu_memory_service/cli/runner.py +++ b/lib/gpu_memory_service/cli/runner.py @@ -22,6 +22,7 @@ from collections.abc import Sequence import uvloop +from gpu_memory_service.common.vmm import init_vmm from gpu_memory_service.server.rpc import GMSRPCServer from .args import Config, parse_args @@ -53,11 +54,13 @@ async def serve_configs(configs: Sequence[Config]) -> None: logging.getLogger().setLevel(logging.DEBUG) logging.getLogger("gpu_memory_service").setLevel(logging.DEBUG) + init_vmm(configs[0].device_type) servers = [] for config in configs: logger.info("Starting GPU Memory Service Server for device %d", config.device) logger.info("GMS tag: %s", config.tag) logger.info("Socket path: %s", config.socket_path) + logger.info("VMM device type: %s", config.device_type.value) logger.info( "Allocation retry config: interval=%ss timeout=%s", config.alloc_retry_interval, diff --git a/lib/gpu_memory_service/cli/server.py b/lib/gpu_memory_service/cli/server.py index 2f0993c1ddec..4e180f781068 100644 --- a/lib/gpu_memory_service/cli/server.py +++ b/lib/gpu_memory_service/cli/server.py @@ -11,13 +11,14 @@ from __future__ import annotations +import argparse import logging import signal import subprocess import sys import time -from gpu_memory_service.common.cuda_utils import list_devices +from gpu_memory_service.common.vmm import VMMDeviceType, get_vmm, init_vmm logging.basicConfig( level=logging.INFO, @@ -26,9 +27,17 @@ logger = logging.getLogger(__name__) -def _child_command(device: int) -> list[str]: +def _child_command(device: int, device_type: str) -> list[str]: """Command for one child process serving every production tag on one GPU.""" - return [sys.executable, "-m", "gpu_memory_service", "--device", str(device)] + return [ + sys.executable, + "-m", + "gpu_memory_service", + "--device", + str(device), + "--device-type", + device_type, + ] def _terminate_all(processes: list[subprocess.Popen]) -> None: @@ -50,10 +59,31 @@ def _supervise(processes: list[subprocess.Popen]) -> int: def main() -> None: + parser = argparse.ArgumentParser( + description="GPU Memory Service supervisor (one server per (device, tag))." + ) + parser.add_argument( + "--device-type", + type=str, + default=VMMDeviceType.CUDA.value, + choices=[d.value for d in VMMDeviceType], + help="VMM device type forwarded to server (default: cuda).", + ) + args = parser.parse_args() + + init_vmm(VMMDeviceType.from_str(args.device_type)) + vmm = get_vmm() + vmm.ensure_initialized() + devices = vmm.list_devices() processes = [] - for device in list_devices(): - proc = subprocess.Popen(_child_command(device)) - logger.info("Started GMS device=%d pid=%d", device, proc.pid) + for device in devices: + proc = subprocess.Popen(_child_command(device, args.device_type)) + logger.info( + "Started GMS device=%d device_type=%s pid=%d", + device, + args.device_type, + proc.pid, + ) processes.append(proc) def terminate(*_args) -> None: diff --git a/lib/gpu_memory_service/cli/snapshot/loader.py b/lib/gpu_memory_service/cli/snapshot/loader.py index 9df77e38f1e7..e0d3e805a6f4 100644 --- a/lib/gpu_memory_service/cli/snapshot/loader.py +++ b/lib/gpu_memory_service/cli/snapshot/loader.py @@ -18,8 +18,8 @@ from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -from gpu_memory_service.common import cuda_utils from gpu_memory_service.common.utils import get_socket_path +from gpu_memory_service.common.vmm import VMMDeviceType, get_vmm, init_vmm from gpu_memory_service.snapshot.backends.sharded_ssd import parse_sharded_ssd_roots from gpu_memory_service.snapshot.storage_client import GMSStorageClient from gpu_memory_service.snapshot.transfer import TransferBackendKind @@ -52,7 +52,9 @@ def _load_device( # GMSStorageClient still publishes the restored layout from this thread. # Ensure the loader's main per-device thread has a current CUDA context for # the final synchronize/unmap/commit path. - cuda_utils.cuda_runtime_set_device(device) + vmm = get_vmm() + vmm.ensure_initialized() + vmm.runtime_set_device(device) client = GMSStorageClient( socket_path=get_socket_path(device), device=device, @@ -105,11 +107,22 @@ def _build_parser() -> argparse.ArgumentParser: default=2, help="Number of independent sharded-ssd restore queues per SSD root.", ) + parser.add_argument( + "--device-type", + type=str, + default=VMMDeviceType.CUDA.value, + choices=[d.value for d in VMMDeviceType], + help="VMM device type (default: cuda).", + ) return parser -def _list_checkpoint_devices(checkpoint_dir: str | None) -> list[int]: - devices = cuda_utils.list_devices() +def _list_checkpoint_devices( + checkpoint_dir: str | None, +) -> list[int]: + vmm = get_vmm() + vmm.ensure_initialized() + devices = vmm.list_devices() if not checkpoint_dir: return devices @@ -155,6 +168,8 @@ def main(argv: list[str] | None = None) -> None: parser.error("--sharded-ssd-queues-per-root must be a positive integer") checkpoint_dir = args.checkpoint_dir max_workers = args.max_workers + device_type = VMMDeviceType.from_str(args.device_type) + init_vmm(device_type) transfer_backend = args.transfer_backend sharded_ssd_roots = parse_sharded_ssd_roots(args.sharded_ssd_roots) sharded_ssd_queues_per_root = args.sharded_ssd_queues_per_root diff --git a/lib/gpu_memory_service/cli/snapshot/saver.py b/lib/gpu_memory_service/cli/snapshot/saver.py index 6c445f80bc65..d8db65a3c6a8 100644 --- a/lib/gpu_memory_service/cli/snapshot/saver.py +++ b/lib/gpu_memory_service/cli/snapshot/saver.py @@ -16,8 +16,8 @@ import time from concurrent.futures import ThreadPoolExecutor, as_completed -from gpu_memory_service.common import cuda_utils from gpu_memory_service.common.utils import get_socket_path +from gpu_memory_service.common.vmm import VMMDeviceType, get_vmm, init_vmm from gpu_memory_service.snapshot.backends.sharded_ssd import ( device_sharded_ssd_roots, parse_sharded_ssd_roots, @@ -55,9 +55,11 @@ def _save_device( ",".join(shard_roots) or "-", ) t0 = time.monotonic() - # This runs on a ThreadPoolExecutor thread; bind its CUDA device before + # This runs on a ThreadPoolExecutor thread; bind its device before # any device work, mirroring the loader's _load_device. - cuda_utils.cuda_runtime_set_device(device) + vmm = get_vmm() + vmm.ensure_initialized() + vmm.runtime_set_device(device) GMSStorageClient( output_dir, socket_path=get_socket_path(device), @@ -107,6 +109,13 @@ def _build_parser() -> argparse.ArgumentParser: default="", help="Comma-separated SSD roots for sharded prototype saves.", ) + parser.add_argument( + "--device-type", + type=str, + default=VMMDeviceType.CUDA.value, + choices=[d.value for d in VMMDeviceType], + help="VMM device type (default: cuda).", + ) return parser @@ -122,7 +131,11 @@ def main(argv: list[str] | None = None) -> None: shard_size_bytes = args.shard_size_bytes sharded_ssd_roots = parse_sharded_ssd_roots(args.sharded_ssd_roots) - devices = cuda_utils.list_devices() + device_type = VMMDeviceType.from_str(args.device_type) + init_vmm(device_type) + vmm = get_vmm() + vmm.ensure_initialized() + devices = vmm.list_devices() logger.info( "Starting GMS save for %d devices lock_timeout_ms=%d sharded_ssd_roots=%s", len(devices), diff --git a/lib/gpu_memory_service/client/memory_manager.py b/lib/gpu_memory_service/client/memory_manager.py index db8763e94bd8..c2c44675f91c 100644 --- a/lib/gpu_memory_service/client/memory_manager.py +++ b/lib/gpu_memory_service/client/memory_manager.py @@ -35,23 +35,10 @@ from typing import Dict, List, Optional from gpu_memory_service.client.session import _GMSClientSession -from gpu_memory_service.common.cuda_utils import ( - align_to_granularity, - cuda_ensure_initialized, - cuda_synchronize, - cuda_validate_pointer, - cumem_address_free, - cumem_address_reserve, - cumem_create_tolerate_oom, - cumem_get_allocation_granularity, - cumem_import_from_shareable_handle_close_fd, - cumem_map, - cumem_release, - cumem_set_access, - cumem_unmap, -) from gpu_memory_service.common.locks import GrantedLockType, RequestedLockType from gpu_memory_service.common.protocol.messages import GetAllocationResponse +from gpu_memory_service.common.utils import align_to_granularity +from gpu_memory_service.common.vmm import VMMDeviceType, get_vmm, get_vmm_device_type logger = logging.getLogger(__name__) @@ -171,6 +158,7 @@ def __init__( self.device = device self.tag = tag self.scratch_size = scratch_size + self._vmm = get_vmm() self._client: Optional[_GMSClientSession] = None @@ -191,11 +179,15 @@ def __init__( self._va_preserved = False self._last_memory_layout_hash: str = "" - cuda_ensure_initialized() - self.granularity = cumem_get_allocation_granularity(device) + self._vmm.ensure_initialized() + self.granularity = self._vmm.get_allocation_granularity(device) # ==================== Properties ==================== + @property + def device_type(self) -> VMMDeviceType: + return get_vmm_device_type() + @property def granted_lock_type(self) -> Optional[GrantedLockType]: return self._granted_lock_type @@ -335,7 +327,7 @@ def commit(self) -> bool: self._require_rw() # Publish barrier: all writer-side GPU work must be visible before commit. - cuda_synchronize() + self._vmm.synchronize() for mapping in list(self._mappings.values()): if mapping.handle != 0: @@ -376,7 +368,7 @@ def metadata_delete(self, key: str) -> bool: def reserve_va(self, size: int) -> int: """Reserve virtual address space (cuMemAddressReserve). No tracking.""" aligned_size = align_to_granularity(size, self.granularity) - return cumem_address_reserve(aligned_size, self.granularity) + return self._vmm.address_reserve(aligned_size, self.granularity) def map_va( self, @@ -393,9 +385,9 @@ def map_va( """ assert self._granted_lock_type is not None aligned_size = align_to_granularity(size, self.granularity) - handle = cumem_import_from_shareable_handle_close_fd(fd) - cumem_map(va, aligned_size, handle) - cumem_set_access(va, aligned_size, self.device, self._granted_lock_type) + handle = self._vmm.import_shareable_handle_close_fd(fd) + self._vmm.map(va, aligned_size, handle) + self._vmm.set_access(va, aligned_size, self.device, self._granted_lock_type) self._track_mapping( LocalMapping( allocation_id=allocation_id, @@ -418,8 +410,8 @@ def unmap_va(self, va: int) -> None: mapping = self._mappings.get(va) if mapping is None or mapping.handle == 0: return - cumem_unmap(va, mapping.aligned_size) - cumem_release(mapping.handle) + self._vmm.unmap(va, mapping.aligned_size) + self._vmm.release(mapping.handle) self._mappings[va] = mapping.with_handle(0) def free_va(self, va: int) -> None: @@ -435,7 +427,7 @@ def free_va(self, va: int) -> None: mapping = self._mappings.get(va) if mapping is None: return - cumem_address_free(va, mapping.va_reserved_size) + self._vmm.address_free(va, mapping.va_reserved_size) self._mappings.pop(va, None) self._inverse_mapping.pop(mapping.allocation_id, None) @@ -507,7 +499,7 @@ def unmap_all_vas(self) -> None: """Synchronize + unmap all VAs (real mappings AND scratch mappings). Preserves VA reservations for remap. """ - cuda_synchronize() + self._vmm.synchronize() unmapped_count = 0 total_bytes = 0 @@ -523,8 +515,8 @@ def unmap_all_vas(self) -> None: for base_va, scratch in self._scratch_mappings.items(): if scratch.scratch_handle == 0: continue - cumem_unmap(base_va, scratch.va_reserved_size) - cumem_release(scratch.scratch_handle) + self._vmm.unmap(base_va, scratch.va_reserved_size) + self._vmm.release(scratch.scratch_handle) scratch.scratch_handle = 0 unmapped_count += 1 total_bytes += scratch.va_reserved_size @@ -586,14 +578,13 @@ def remap_all_vas(self) -> None: ) if str(alloc_info.tag) != mapping.tag: raise StaleMemoryLayoutError( - f"Layout rank {rank} tag changed: " - f"{mapping.tag} vs {alloc_info.tag}" + f"Layout rank {rank} tag changed: {mapping.tag} vs {alloc_info.tag}" ) fd = self.export_handle(alloc_info.allocation_id) - handle = cumem_import_from_shareable_handle_close_fd(fd) - cumem_map(va, mapping.aligned_size, handle) - cumem_set_access( + handle = self._vmm.import_shareable_handle_close_fd(fd) + self._vmm.map(va, mapping.aligned_size, handle) + self._vmm.set_access( va, mapping.aligned_size, self.device, self._granted_lock_type ) remapped_vas.append(va) @@ -609,9 +600,9 @@ def remap_all_vas(self) -> None: total_bytes += mapping.aligned_size if remapped_vas: - cuda_synchronize() + self._vmm.synchronize() for va in remapped_vas: - cuda_validate_pointer(va) + self._vmm.validate_pointer(va) self._va_preserved = False self._unmapped = False @@ -703,17 +694,20 @@ def create_scratch_mapping(self, size: int, tag: str = "kv_cache") -> int: aligned_size = align_to_granularity(size, self.granularity) va_reserved_size = align_to_granularity(size, self.scratch_size) - ok, scratch_handle = cumem_create_tolerate_oom(self.scratch_size, self.device) + ok, scratch_handle = self._vmm.create_tolerate_oom( + self.scratch_size, self.device + ) if not ok: raise RuntimeError( - "cuMemCreate failed to allocate the scratch chunk " - f"({self.scratch_size // (1 << 20)} MiB) on device {self.device}" + f"VMM physical memory allocation failed " + f"({self.scratch_size // (1 << 20)} MiB) on " + f"{self.device_type.value} device {self.device}" ) - va = cumem_address_reserve(va_reserved_size, self.scratch_size) + va = self._vmm.address_reserve(va_reserved_size, self.scratch_size) for offset in range(0, va_reserved_size, self.scratch_size): - cumem_map(va + offset, self.scratch_size, scratch_handle) - cumem_set_access(va, va_reserved_size, self.device, GrantedLockType.RW) + self._vmm.map(va + offset, self.scratch_size, scratch_handle) + self._vmm.set_access(va, va_reserved_size, self.device, GrantedLockType.RW) self._scratch_mappings[va] = _ScratchMapping( size=size, @@ -797,11 +791,11 @@ def destroy_scratch_mapping(self, base_va: int) -> bool: if scratch is None: return False - cuda_synchronize() + self._vmm.synchronize() if scratch.scratch_handle: - cumem_unmap(base_va, scratch.va_reserved_size) - cumem_release(scratch.scratch_handle) - cumem_address_free(base_va, scratch.va_reserved_size) + self._vmm.unmap(base_va, scratch.va_reserved_size) + self._vmm.release(scratch.scratch_handle) + self._vmm.address_free(base_va, scratch.va_reserved_size) return True # ==================== Lifecycle ==================== @@ -812,10 +806,10 @@ def close(self, *, best_effort: bool = False) -> None: synchronize + unmap all + free all VAs + abort. Args: - best_effort: If True, skip cuda_synchronize and swallow + best_effort: If True, skip self._vmm.synchronize() and swallow errors during cleanup. Used after checkpoint where cuda-checkpoint may have torn down the device context - (cuda_synchronize calls os._exit via fail()). + (self._vmm.synchronize() calls os._exit via fail()). """ if best_effort: try: @@ -826,7 +820,7 @@ def close(self, *, best_effort: bool = False) -> None: self._inverse_mapping.clear() self._scratch_mappings.clear() else: - cuda_synchronize() + self._vmm.synchronize() for base_va in list(self._scratch_mappings.keys()): self.destroy_scratch_mapping(base_va) for va in list(self._mappings.keys()): diff --git a/lib/gpu_memory_service/client/torch/allocator.py b/lib/gpu_memory_service/client/torch/allocator.py index 98a834429a99..1346d14c3f66 100644 --- a/lib/gpu_memory_service/client/torch/allocator.py +++ b/lib/gpu_memory_service/client/torch/allocator.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Iterator, Optional from gpu_memory_service.common.locks import GrantedLockType, RequestedLockType +from gpu_memory_service.common.vmm import VMMDeviceType, get_vmm_device_type if TYPE_CHECKING: import torch @@ -80,6 +81,11 @@ def _gms_free(ptr: int, size: int, device: int, stream: int) -> None: def _ensure_callbacks_initialized() -> None: global _callbacks_initialized, _pluggable_alloc + if get_vmm_device_type() != VMMDeviceType.CUDA: + raise NotImplementedError( + f"GMS torch mempool integration is CUDA-only; device_type={get_vmm_device_type().value} " + ) + from gpu_memory_service.client.torch.extensions import _allocator_ext as cumem from torch.cuda import CUDAPluggableAllocator @@ -92,6 +98,11 @@ def _ensure_callbacks_initialized() -> None: def _create_mem_pool() -> "MemPool": + if get_vmm_device_type() != VMMDeviceType.CUDA: + raise NotImplementedError( + f"GMS torch mempool integration is CUDA-only; device_type={get_vmm_device_type().value} " + ) + from torch.cuda.memory import MemPool assert _pluggable_alloc is not None @@ -284,9 +295,9 @@ def prune_allocations( return if synchronize: - import torch + from gpu_memory_service.integrations.common.utils import torch_device - torch.cuda.synchronize(manager.device) + torch_device().synchronize(manager.device) keep = {str(allocation_id) for allocation_id in referenced_allocation_ids} @@ -328,6 +339,11 @@ def gms_use_mem_pool(tag: str, device: "torch.device | int") -> Iterator[None]: if state.mem_pool is None: raise RuntimeError(f"GMS allocator tag={tag} does not have a mempool") + if get_vmm_device_type() != VMMDeviceType.CUDA: + raise NotImplementedError( + f"gms_use_mem_pool is CUDA-only; device_type={get_vmm_device_type().value} " + ) + token = _active_tag.set(tag) try: with torch.cuda.use_mem_pool(state.mem_pool, device=device): diff --git a/lib/gpu_memory_service/common/utils.py b/lib/gpu_memory_service/common/utils.py index 8f589ac726a9..1e04b31aab8b 100644 --- a/lib/gpu_memory_service/common/utils.py +++ b/lib/gpu_memory_service/common/utils.py @@ -75,3 +75,16 @@ def get_socket_path(device: int, tag: str = "weights") -> str: _uuid_cache[device] = uuid socket_dir = os.environ.get("GMS_SOCKET_DIR") or tempfile.gettempdir() return os.path.join(socket_dir, f"gms_{uuid}_{tag}.sock") + + +def align_to_granularity(size: int, granularity: int) -> int: + """Align size up to VMM granularity. + + Args: + size: Size in bytes + granularity: Allocation granularity + + Returns: + Aligned size + """ + return ((size + granularity - 1) // granularity) * granularity diff --git a/lib/gpu_memory_service/common/vmm/__init__.py b/lib/gpu_memory_service/common/vmm/__init__.py new file mode 100644 index 000000000000..50458dc493d3 --- /dev/null +++ b/lib/gpu_memory_service/common/vmm/__init__.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPU Memory Service — VMM device abstraction. + +GMS depends on a per-vendor virtual-memory-management surface +(allocate physical memory, export/import shareable handles, reserve and +map virtual addresses, ...). + +The VMM instance is process-global singleton and immutable once initialized. +- `init_vmm(device_type)`: explicit process-level backend selection, e.g. + CLI `--device-type`. On a mixed node (e.g., CUDA and other device coexist), + --device-type device_name overrides the auto-detection because it calls + init_vmm(VMMDeviceType.DeviceName) at process startup — before any get_vmm() + lazy path fires. +- `get_vmm()`: returns the singleton and may lazily initialize the default/autodetected + backend for public client/integration compatibility via init_vmm(detected_device_type). +- `vmm.ensure_initialized()`: initializes the selected backend driver/runtime, + e.g. CUDA `cuInit` + +""" + +from __future__ import annotations + +import threading +from enum import Enum + +from .device import VMMDevice + + +class VMMDeviceType(str, Enum): + """Identify which vendor's VMM driver a GMS instance should use.""" + + CUDA = "cuda" + XPU = "xpu" + + @classmethod + def from_str(cls, value: str) -> "VMMDeviceType": + try: + return cls(value.lower()) + except ValueError as exc: + valid = ", ".join(b.value for b in cls) + raise ValueError( + f"Unknown VMM device type {value!r}; expected one of: {valid}" + ) from exc + + +# --------------------------------------------------------------------------- +# Process-global singleton +# --------------------------------------------------------------------------- + +_lock = threading.Lock() +_vmm_instance: VMMDevice | None = None +_vmm_device_type: VMMDeviceType | None = None + + +def init_vmm(device_type: VMMDeviceType) -> None: + """Initialize the process-global VMM singleton. Idempotent for same kind.""" + global _vmm_instance, _vmm_device_type + with _lock: + if _vmm_instance is not None: + if _vmm_device_type != device_type: + raise RuntimeError( + f"VMM already initialized as {_vmm_device_type!r}; " + f"cannot reinitialize as {device_type!r}" + ) + return + inst = _create_vmm(device_type) + _vmm_device_type = device_type + _vmm_instance = inst + + +def _detect_device_type() -> VMMDeviceType: + """Auto-detect the available accelerator device type at runtime. + + Priority: CUDA > XPU > fallback to CUDA (will fail at actual device use). + """ + try: + import torch + + if torch.cuda.is_available(): + return VMMDeviceType.CUDA + if hasattr(torch, "xpu") and torch.xpu.is_available(): + return VMMDeviceType.XPU + except Exception: + pass + return VMMDeviceType.CUDA + + +def get_vmm() -> VMMDevice: + """Return the process-global VMM singleton. + + If ``init_vmm()`` has not been called, lazily auto-detects the device + type via ``_detect_device_type()`` and initializes. This preserves + backward compatibility for integration paths (vLLM, SGLang, TRTLLM, + gms-storage-client) that construct helpers without explicitly + bootstrapping the VMM singleton. + + Explicit ``init_vmm(device_type)`` calls (e.g. from CLI ``--device-type``) + still take priority since they run before any ``get_vmm()`` call. + """ + inst = _vmm_instance + if inst is None: + init_vmm(_detect_device_type()) + inst = _vmm_instance + return inst # type: ignore[return-value] + + +def get_vmm_device_type() -> VMMDeviceType: + """Return the active device type. + + Lazily auto-detects if ``init_vmm()`` has not been called. + """ + kind = _vmm_device_type + if kind is None: + init_vmm(_detect_device_type()) + kind = _vmm_device_type + return kind # type: ignore[return-value] + + +def _create_vmm(device_type: VMMDeviceType) -> VMMDevice: + """Construct the appropriate VMMDevice implementation.""" + if device_type is VMMDeviceType.CUDA: + from .cuda_utils import CudaVMM + + return CudaVMM() + + if device_type is VMMDeviceType.XPU: + raise NotImplementedError("'xpu' VMM backend is not implemented yet") + + raise ValueError(f"Unhandled VMM device type: {device_type!r}") + + +# --------------------------------------------------------------------------- +# Test support +# --------------------------------------------------------------------------- + + +def _reset_vmm_singleton() -> None: + """Reset the singleton for test isolation. NOT for production use.""" + global _vmm_instance, _vmm_device_type + with _lock: + _vmm_instance = None + _vmm_device_type = None + + +__all__ = [ + "VMMDevice", + "VMMDeviceType", + "get_vmm", + "get_vmm_device_type", + "init_vmm", +] diff --git a/lib/gpu_memory_service/common/cuda_utils.py b/lib/gpu_memory_service/common/vmm/cuda_utils.py similarity index 71% rename from lib/gpu_memory_service/common/cuda_utils.py rename to lib/gpu_memory_service/common/vmm/cuda_utils.py index 2abee0d6ef3e..7afae531e073 100644 --- a/lib/gpu_memory_service/common/cuda_utils.py +++ b/lib/gpu_memory_service/common/vmm/cuda_utils.py @@ -9,6 +9,7 @@ from gpu_memory_service.common.locks import GrantedLockType from gpu_memory_service.common.utils import fail +from gpu_memory_service.common.vmm.device import VMMDevice try: from cuda.bindings import driver as cuda @@ -37,7 +38,7 @@ def __getattr__(self, name): cuda_runtime = _MissingCudaRuntime() -def list_devices() -> list[int]: +def list_cuda_devices() -> list[int]: """Return list of CUDA device indices visible to this process via NVML.""" import pynvml @@ -51,7 +52,7 @@ def list_devices() -> list[int]: return list(range(count)) -def device_memory_info(device: int) -> tuple[int, int]: +def cuda_device_memory_info(device: int) -> tuple[int, int]: """Return ``(free_bytes, total_bytes)`` for a CUDA device via NVML.""" import pynvml @@ -131,19 +132,6 @@ def cumem_export_to_shareable_handle(handle: int) -> int: return int(fd) -def align_to_granularity(size: int, granularity: int) -> int: - """Align size up to VMM granularity. - - Args: - size: Size in bytes - granularity: Allocation granularity - - Returns: - Aligned size - """ - return ((size + granularity - 1) // granularity) * granularity - - def cumem_import_from_shareable_handle_close_fd(fd: int) -> int: try: result, handle = cuda.cuMemImportFromShareableHandle( @@ -306,3 +294,111 @@ def cuda_memcpy_d2h_async( ), "cudaMemcpyAsync", ) + + +class CudaVMM(VMMDevice): + """``VMMDevice`` Protocol implementation backed by the CUDA driver API. + + Methods delegate to the module-level helper. + + """ + + # ----- driver lifecycle ------------------------------------------------- + + def ensure_initialized(self) -> None: + cuda_ensure_initialized() + + def synchronize(self) -> None: + cuda_synchronize() + + # ----- discovery / sizing ----------------------------------------------- + + def list_devices(self) -> list[int]: + return list_cuda_devices() + + def device_memory_info(self, device: int) -> tuple[int, int]: + return cuda_device_memory_info(device) + + def get_allocation_granularity(self, device: int) -> int: + return cumem_get_allocation_granularity(device) + + # ----- physical memory -------------------------------------------------- + + def create_tolerate_oom(self, size: int, device: int) -> tuple[bool, int]: + return cumem_create_tolerate_oom(size, device) + + def release(self, handle: int) -> None: + cumem_release(handle) + + # ----- shareable-handle export / import --------------------------------- + + def export_to_shareable_handle(self, handle: int) -> int: + return cumem_export_to_shareable_handle(handle) + + def import_shareable_handle_close_fd(self, fd: int) -> int: + return cumem_import_from_shareable_handle_close_fd(fd) + + # ----- virtual address space + mapping ---------------------------------- + + def address_reserve(self, size: int, granularity: int) -> int: + return cumem_address_reserve(size, granularity) + + def address_free(self, va: int, size: int) -> None: + cumem_address_free(va, size) + + def map(self, va: int, size: int, handle: int) -> None: + cumem_map(va, size, handle) + + def unmap(self, va: int, size: int) -> None: + cumem_unmap(va, size) + + def set_access( + self, va: int, size: int, device: int, access: GrantedLockType + ) -> None: + cumem_set_access(va, size, device, access) + + # ----- pointer validation ----------------------------------------------- + + def validate_pointer(self, va: int) -> None: + cuda_validate_pointer(va) + + # ----- runtime helpers -------------------------------------------------- + + def runtime_check_result(self, result, name: str) -> None: + cuda_runtime_check_result(result, name) + + def runtime_set_device(self, device: int) -> None: + cuda_runtime_set_device(device) + + def host_register(self, ptr: int, size: int) -> None: + cuda_host_register(ptr, size) + + def host_unregister(self, ptr: int) -> None: + cuda_host_unregister(ptr) + + def stream_create_nonblocking(self): + return cuda_stream_create_nonblocking() + + def stream_destroy(self, stream) -> None: + cuda_stream_destroy(stream) + + def stream_synchronize(self, stream) -> None: + cuda_stream_synchronize(stream) + + def memcpy_h2d_async( + self, + dst_ptr: int, + src_ptr: int, + size: int, + stream, + ) -> None: + cuda_memcpy_h2d_async(dst_ptr, src_ptr, size, stream) + + def memcpy_d2h_async( + self, + dst_ptr: int, + src_ptr: int, + size: int, + stream, + ) -> None: + cuda_memcpy_d2h_async(dst_ptr, src_ptr, size, stream) diff --git a/lib/gpu_memory_service/common/vmm/device.py b/lib/gpu_memory_service/common/vmm/device.py new file mode 100644 index 000000000000..f78e59a4cac3 --- /dev/null +++ b/lib/gpu_memory_service/common/vmm/device.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Abstract VMM device base class. + +Defines the per-vendor virtual-memory-management surface that GMS depends on. +Method names are vendor-neutral verbs (``map``, ``address_reserve``, +``create_tolerate_oom``, etc.). + +All derived classes must implement every abstract method. + +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from gpu_memory_service.common.locks import GrantedLockType + + +class VMMDevice(ABC): + """Per-vendor virtual-memory-management device contract. + + A device instance is obtained via ``get_vmm()`` + and used by GMS to allocate physical memory, export/import shareable + handles for cross-process sharing, reserve and map virtual addresses. + + All subclasses must implement every abstract method; instantiation will + raise ``TypeError`` if any are missing. + + """ + + # ----- driver lifecycle ------------------------------------------------- + + @abstractmethod + def ensure_initialized(self) -> None: + """Initialize the underlying driver. Idempotent.""" + + @abstractmethod + def synchronize(self) -> None: + """Block until all in-flight device work for the current context completes.""" + + # ----- discovery / sizing ----------------------------------------------- + + @abstractmethod + def list_devices(self) -> list[int]: + """Return device indices visible to this process.""" + + @abstractmethod + def device_memory_info(self, device: int) -> tuple[int, int]: + """Return ``(free_bytes, total_bytes)`` for ``device``.""" + + @abstractmethod + def get_allocation_granularity(self, device: int) -> int: + """Minimum allocation granularity in bytes for ``device``.""" + + # ----- physical memory -------------------------------------------------- + + @abstractmethod + def create_tolerate_oom(self, size: int, device: int) -> tuple[bool, int]: + """Allocate physical memory of exactly ``size`` bytes on ``device``. + + Returns ``(allocated, handle)``. ``allocated`` is ``False`` and + ``handle`` is ``0`` on OOM; on any other error the implementation + raises. + """ + + @abstractmethod + def release(self, handle: int) -> None: + """Release a physical memory handle returned by ``create_tolerate_oom``.""" + + # ----- shareable-handle export / import --------------------------------- + + @abstractmethod + def export_to_shareable_handle(self, handle: int) -> int: + """Return a POSIX FD that can be passed cross-process via SCM_RIGHTS.""" + + @abstractmethod + def import_shareable_handle_close_fd(self, fd: int) -> int: + """Import a shareable FD into a local physical-memory handle. + + The FD is closed on success or failure (matches the existing CUDA + helper's contract). + """ + + # ----- virtual address space + mapping ---------------------------------- + + @abstractmethod + def address_reserve(self, size: int, granularity: int) -> int: + """Reserve a contiguous VA range. Returns the base VA.""" + + @abstractmethod + def address_free(self, va: int, size: int) -> None: + """Release a VA reservation.""" + + @abstractmethod + def map(self, va: int, size: int, handle: int) -> None: + """Bind the VA range to a physical handle.""" + + @abstractmethod + def unmap(self, va: int, size: int) -> None: + """Unbind the VA range. The reservation itself is preserved.""" + + @abstractmethod + def set_access( + self, va: int, size: int, device: int, access: GrantedLockType + ) -> None: + """Set device-side access permissions for a mapped VA range.""" + + # ----- pointer validation ----------------------------------------------- + + @abstractmethod + def validate_pointer(self, va: int) -> None: + """Best-effort check that ``va`` refers to a valid device allocation.""" + + # ----- runtime helpers -------------------------------------------------- + + @abstractmethod + def runtime_check_result(self, result, name: str) -> None: + """Check a device-runtime return code; raise on failure.""" + + @abstractmethod + def runtime_set_device(self, device: int) -> None: + """Set the active device for the current thread.""" + + @abstractmethod + def host_register(self, ptr: int, size: int) -> None: + """Pin host memory for DMA access.""" + + @abstractmethod + def host_unregister(self, ptr: int) -> None: + """Unpin previously registered host memory.""" + + @abstractmethod + def stream_create_nonblocking(self): + """Create a non-blocking execution stream. Returns an opaque handle.""" + + @abstractmethod + def stream_destroy(self, stream) -> None: + """Destroy an execution stream.""" + + @abstractmethod + def stream_synchronize(self, stream) -> None: + """Block until all work on ``stream`` completes.""" + + @abstractmethod + def memcpy_h2d_async( + self, + dst_ptr: int, + src_ptr: int, + size: int, + stream, + ) -> None: + """Async host-to-device copy on ``stream``.""" + + @abstractmethod + def memcpy_d2h_async( + self, + dst_ptr: int, + src_ptr: int, + size: int, + stream, + ) -> None: + """Async device-to-host copy on ``stream``.""" diff --git a/lib/gpu_memory_service/integrations/common/utils.py b/lib/gpu_memory_service/integrations/common/utils.py index 1fa24d98c105..66a0b5986e5a 100644 --- a/lib/gpu_memory_service/integrations/common/utils.py +++ b/lib/gpu_memory_service/integrations/common/utils.py @@ -23,6 +23,18 @@ logger = logging.getLogger(__name__) +def torch_device(): + """Return the torch device module (torch.cuda or torch.xpu) for the active VMM device.""" + from gpu_memory_service.common.vmm import VMMDeviceType, get_vmm_device_type + + device_type = get_vmm_device_type() + if device_type == VMMDeviceType.CUDA: + return torch.cuda + if device_type == VMMDeviceType.XPU: + return torch.xpu + raise RuntimeError(f"Unsupported VMM device type: {device_type!r}") + + @dataclass(frozen=True) class GMSCommittedMemoryStats: committed_bytes: int diff --git a/lib/gpu_memory_service/integrations/vllm/worker.py b/lib/gpu_memory_service/integrations/vllm/worker.py index c526fad48ca8..2194089bc5fa 100644 --- a/lib/gpu_memory_service/integrations/vllm/worker.py +++ b/lib/gpu_memory_service/integrations/vllm/worker.py @@ -38,6 +38,7 @@ from gpu_memory_service.integrations.common.utils import ( get_gms_lock_mode, get_gms_ro_connect_timeout_ms, + torch_device, ) from gpu_memory_service.integrations.vllm.model_loader import ( abort_pending_gms_write, @@ -197,10 +198,10 @@ def _determine_available_memory_before_gms_publish(self) -> int: # adding below. has_pending_write = has_pending_gms_write() - torch.cuda.reset_peak_memory_stats() + torch_device().reset_peak_memory_stats() self.model_runner.profile_run() - torch.cuda.synchronize() - torch_peak = torch.cuda.max_memory_allocated() + torch_device().synchronize() + torch_peak = torch_device().max_memory_allocated() cudagraph_memory_estimate = 0 if ( @@ -335,7 +336,7 @@ def sleep(self, level: int = 1) -> None: reservations. Wake reconnects and rebuilds via the standard prepare_scratch_for_reallocation → reallocate → remap pipeline. """ - free_bytes_before = torch.cuda.mem_get_info()[0] + free_bytes_before = torch_device().mem_get_info()[0] # Pause MX serving before GMS unmap mx_ctx = get_mx_load_context() @@ -350,9 +351,9 @@ def sleep(self, level: int = 1) -> None: manager.abort() gc.collect() - torch.cuda.empty_cache() + torch_device().empty_cache() - free_bytes_after, total = torch.cuda.mem_get_info() + free_bytes_after, total = torch_device().mem_get_info() freed_bytes = free_bytes_after - free_bytes_before used_bytes = total - free_bytes_after logger.info( diff --git a/lib/gpu_memory_service/server/allocations.py b/lib/gpu_memory_service/server/allocations.py index 17175fb9fcb2..df63cd9af2be 100644 --- a/lib/gpu_memory_service/server/allocations.py +++ b/lib/gpu_memory_service/server/allocations.py @@ -1,7 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Server-side CUDA allocation store.""" +"""Server-side VMM allocation store. + +Uses the ``VMMDevice`` abstraction from ``common.vmm`` so that future +non-CUDA device types can plug in without touching this file. + +""" from __future__ import annotations @@ -13,15 +18,8 @@ from typing import Callable, Optional from uuid import uuid4 -from gpu_memory_service.common.cuda_utils import ( - align_to_granularity, - cuda_ensure_initialized, - cumem_create_tolerate_oom, - cumem_export_to_shareable_handle, - cumem_get_allocation_granularity, - cumem_release, - device_memory_info, -) +from gpu_memory_service.common.utils import align_to_granularity +from gpu_memory_service.common.vmm import get_vmm logger = logging.getLogger(__name__) @@ -62,15 +60,16 @@ def __init__( ) self._device = device + self._vmm = get_vmm() self._allocations: dict[str, AllocationInfo] = {} self._next_layout_slot = 0 - cuda_ensure_initialized() - self._granularity = cumem_get_allocation_granularity(device) + self._vmm.ensure_initialized() + self._granularity = self._vmm.get_allocation_granularity(device) self._allocation_retry_interval = allocation_retry_interval self._allocation_retry_timeout = allocation_retry_timeout logger.info( - "GMSAllocationManager initialized: device=%d, granularity=%d, " - "alloc_retry_interval=%.3f, alloc_retry_timeout=%s", + "GMSAllocationManager initialized: device=%d, " + "granularity=%d, alloc_retry_interval=%.3f, alloc_retry_timeout=%s", device, self._granularity, self._allocation_retry_interval, @@ -108,7 +107,9 @@ async def allocate( "RW client disconnected during allocation retry" ) - allocated, handle = cumem_create_tolerate_oom(aligned_size, self._device) + allocated, handle = self._vmm.create_tolerate_oom( + aligned_size, self._device + ) if allocated: break @@ -130,7 +131,7 @@ async def allocate( # rather than silent. free_b, total_b = -1, -1 try: - free_b, total_b = device_memory_info(self._device) + free_b, total_b = self._vmm.device_memory_info(self._device) except Exception: logger.debug( "NVML memory info failed for device %d", @@ -150,7 +151,7 @@ async def allocate( ) await asyncio.sleep(self._allocation_retry_interval) - export_fd = int(cumem_export_to_shareable_handle(int(handle))) + export_fd = int(self._vmm.export_to_shareable_handle(int(handle))) info = AllocationInfo( allocation_id=str(uuid4()), size=size, @@ -182,7 +183,7 @@ def free_allocation(self, allocation_id: str) -> bool: if info is None: return False os.close(info.export_fd) - cumem_release(info.handle) + self._vmm.release(info.handle) self._allocations.pop(allocation_id, None) logger.debug("Freed allocation: %s", allocation_id) return True @@ -192,7 +193,7 @@ def clear_all(self) -> int: for allocation_id in allocation_ids: info = self._allocations[allocation_id] os.close(info.export_fd) - cumem_release(info.handle) + self._vmm.release(info.handle) self._allocations.pop(allocation_id, None) if allocation_ids: logger.info("Cleared %d allocations", len(allocation_ids)) diff --git a/lib/gpu_memory_service/server/gms.py b/lib/gpu_memory_service/server/gms.py index c6cfad765ea3..2ec427496294 100644 --- a/lib/gpu_memory_service/server/gms.py +++ b/lib/gpu_memory_service/server/gms.py @@ -79,7 +79,10 @@ def __init__( self._events: deque[GMSRuntimeEvent] = deque(maxlen=self._MAX_EVENTS) self._metadata: dict[str, MetadataEntry] = {} self._memory_layout_hash = "" - logger.info("GMS initialized: device=%d", device) + logger.info( + "GMS initialized: device=%d", + device, + ) @property def state(self) -> ServerState: diff --git a/lib/gpu_memory_service/server/rpc.py b/lib/gpu_memory_service/server/rpc.py index 8044d81c6a0f..22ff21eeb418 100644 --- a/lib/gpu_memory_service/server/rpc.py +++ b/lib/gpu_memory_service/server/rpc.py @@ -72,7 +72,10 @@ def __init__( allocation_retry_timeout=allocation_retry_timeout, ) self._server: Optional[asyncio.Server] = None - logger.info("GMSRPCServer initialized: device=%d", device) + logger.info( + "GMSRPCServer initialized: device=%d", + device, + ) def _prepare_socket_path(self) -> None: if not os.path.exists(self.socket_path): diff --git a/lib/gpu_memory_service/setup.py b/lib/gpu_memory_service/setup.py index 93fec96bcdb7..50f8ea2c8403 100644 --- a/lib/gpu_memory_service/setup.py +++ b/lib/gpu_memory_service/setup.py @@ -68,6 +68,7 @@ def _create_ext_modules(): "gpu_memory_service.cli.snapshot", "gpu_memory_service.common", "gpu_memory_service.common.protocol", + "gpu_memory_service.common.vmm", "gpu_memory_service.server", "gpu_memory_service.client", "gpu_memory_service.client.torch", @@ -88,6 +89,7 @@ def _create_ext_modules(): "gpu_memory_service.cli.snapshot": "cli/snapshot", "gpu_memory_service.common": "common", "gpu_memory_service.common.protocol": "common/protocol", + "gpu_memory_service.common.vmm": "common/vmm", "gpu_memory_service.server": "server", "gpu_memory_service.client": "client", "gpu_memory_service.client.torch": "client/torch", diff --git a/lib/gpu_memory_service/snapshot/backends/nixl_staging.py b/lib/gpu_memory_service/snapshot/backends/nixl_staging.py index b1cfb029ff28..75ec159729bd 100644 --- a/lib/gpu_memory_service/snapshot/backends/nixl_staging.py +++ b/lib/gpu_memory_service/snapshot/backends/nixl_staging.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from typing import Callable, List, Mapping, Optional, Sequence -from gpu_memory_service.common import cuda_utils +from gpu_memory_service.common.vmm import get_vmm from gpu_memory_service.snapshot.backends.nixl_common import ( DRAM_MEM_TYPE, FILE_MEM_TYPE, @@ -293,7 +293,7 @@ def _prepare_group( ) if self._cancel_event.is_set(): raise CancelledError(f"{self._backend_name} cancelled") - cuda_utils.cuda_runtime_set_device(self._device) + get_vmm().runtime_set_device(self._device) agent = create_nixl_agent( api, agent_name=agent_name, @@ -302,7 +302,7 @@ def _prepare_group( ) if self._cancel_event.is_set(): raise CancelledError(f"{self._backend_name} cancelled") - slots = make_pinned_copy_slots(_PINNED_COPY_BUFFERS_PER_WORKER) + slots = make_pinned_copy_slots(get_vmm(), _PINNED_COPY_BUFFERS_PER_WORKER) prep_elapsed_s = time.monotonic() - prep_t0 logger.info( "%s prepared %s=%s files=%d prep_elapsed=%.3fs " @@ -335,7 +335,7 @@ def _restore_prepared_group( prepared: _PreparedNixlGroup, targets: Mapping[str, GMSTransferTarget], ) -> None: - cuda_utils.cuda_runtime_set_device(self._device) + get_vmm().runtime_set_device(self._device) group_t0 = time.monotonic() group_bytes = 0 try: @@ -394,7 +394,7 @@ def restore_file_groups_with_nixl_staging( next_slot = 0 try: if owned_slots: - slots = make_pinned_copy_slots(buffers_per_worker) + slots = make_pinned_copy_slots(get_vmm(), buffers_per_worker) for file_path, sources in file_groups: fd = open_direct_read_fd(file_path, logger=logger, require_direct=True) try: diff --git a/lib/gpu_memory_service/snapshot/backends/pinned_host.py b/lib/gpu_memory_service/snapshot/backends/pinned_host.py index 6d502f76105e..7cc9d9516a0e 100644 --- a/lib/gpu_memory_service/snapshot/backends/pinned_host.py +++ b/lib/gpu_memory_service/snapshot/backends/pinned_host.py @@ -10,7 +10,7 @@ import os from typing import Any, List, Sequence, Tuple -from gpu_memory_service.common import cuda_utils +from gpu_memory_service.common.vmm import VMMDevice PINNED_COPY_CHUNK_SIZE = 64 * 1024 * 1024 @@ -42,9 +42,15 @@ def _free_aligned_buffer(view: memoryview, ptr: int) -> None: class PinnedCopySlot: - """One reusable pinned host buffer and CUDA stream.""" + """One reusable pinned host buffer and copy stream. - def __init__(self, size: int = PINNED_COPY_CHUNK_SIZE) -> None: + Args: + vmm: VMMDevice instance for device operations. + size: Size of the pinned buffer in bytes. + """ + + def __init__(self, vmm: VMMDevice, size: int = PINNED_COPY_CHUNK_SIZE) -> None: + self._vmm = vmm size = int(size) self.view, self._raw, self.ptr = _allocate_aligned_buffer(size) self.stream = None @@ -52,29 +58,29 @@ def __init__(self, size: int = PINNED_COPY_CHUNK_SIZE) -> None: self._registered = False self._closed = False try: - self.stream = cuda_utils.cuda_stream_create_nonblocking() - cuda_utils.cuda_host_register(self.ptr, size) + self.stream = self._vmm.stream_create_nonblocking() + self._vmm.host_register(self.ptr, size) self._registered = True except Exception: try: if self.stream is not None: - cuda_utils.cuda_stream_destroy(self.stream) + self._vmm.stream_destroy(self.stream) finally: _free_aligned_buffer(self.view, self.ptr) raise def copy_to_device_async(self, dst_ptr: int, size: int) -> None: - cuda_utils.cuda_memcpy_h2d_async(dst_ptr, self.ptr, size, self.stream) + self._vmm.memcpy_h2d_async(dst_ptr, self.ptr, size, self.stream) self.busy = True def copy_from_device_async(self, src_ptr: int, size: int) -> None: - cuda_utils.cuda_memcpy_d2h_async(self.ptr, src_ptr, size, self.stream) + self._vmm.memcpy_d2h_async(self.ptr, src_ptr, size, self.stream) self.busy = True def wait(self) -> None: if not self.busy: return - cuda_utils.cuda_stream_synchronize(self.stream) + self._vmm.stream_synchronize(self.stream) self.busy = False def close(self) -> None: @@ -87,7 +93,7 @@ def close(self) -> None: error = exc try: if self._registered: - cuda_utils.cuda_host_unregister(self.ptr) + self._vmm.host_unregister(self.ptr) self._registered = False except Exception as exc: # noqa: BLE001 if error is None: @@ -98,13 +104,13 @@ def close(self) -> None: ) try: if self.stream is not None: - cuda_utils.cuda_stream_destroy(self.stream) + self._vmm.stream_destroy(self.stream) self.stream = None except Exception as exc: # noqa: BLE001 if error is None: error = exc else: - _LOGGER.warning("failed to destroy CUDA copy stream", exc_info=True) + _LOGGER.warning("failed to destroy copy stream", exc_info=True) try: _free_aligned_buffer(self.view, self.ptr) self._closed = True @@ -117,11 +123,11 @@ def close(self) -> None: raise error -def make_pinned_copy_slots(count: int) -> List[PinnedCopySlot]: +def make_pinned_copy_slots(vmm: VMMDevice, count: int) -> List[PinnedCopySlot]: slots: List[PinnedCopySlot] = [] try: for _ in range(count): - slots.append(PinnedCopySlot()) + slots.append(PinnedCopySlot(vmm)) except Exception: for slot in slots: try: diff --git a/lib/gpu_memory_service/snapshot/disk.py b/lib/gpu_memory_service/snapshot/disk.py index 0d11bdbd75bd..3ab6ebbd83c3 100644 --- a/lib/gpu_memory_service/snapshot/disk.py +++ b/lib/gpu_memory_service/snapshot/disk.py @@ -8,8 +8,8 @@ import os from typing import Any, Dict, Optional, Sequence, Tuple -from gpu_memory_service.common import cuda_utils from gpu_memory_service.common.protocol.messages import GetAllocationResponse +from gpu_memory_service.common.vmm import get_vmm from gpu_memory_service.snapshot.backends.pinned_host import ( PINNED_COPY_CHUNK_SIZE, close_pinned_copy_slots, @@ -65,9 +65,10 @@ def __init__( raise ValueError("buffers must be positive") if chunk_size <= 0: raise ValueError("chunk_size must be positive") + self._vmm = get_vmm() if device is not None: - cuda_utils.cuda_runtime_set_device(device) - self._slots = make_pinned_copy_slots(buffers) + self._vmm.runtime_set_device(device) + self._slots = make_pinned_copy_slots(self._vmm, buffers) self._slot_index = 0 self._closed = False try: diff --git a/lib/gpu_memory_service/tests/_deps.py b/lib/gpu_memory_service/tests/_deps.py index d2ccd599c3f7..09639adb3910 100644 --- a/lib/gpu_memory_service/tests/_deps.py +++ b/lib/gpu_memory_service/tests/_deps.py @@ -34,4 +34,7 @@ def _check_gms_usable() -> bool: if HAS_TORCH: import torch - HAS_CUDA = torch.cuda.is_available() + try: + HAS_CUDA = torch.cuda.is_available() + except Exception: + HAS_CUDA = False diff --git a/lib/gpu_memory_service/tests/_fake_vmm.py b/lib/gpu_memory_service/tests/_fake_vmm.py new file mode 100644 index 000000000000..bccbe2c31932 --- /dev/null +++ b/lib/gpu_memory_service/tests/_fake_vmm.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared VMMDevice mock for unit tests. + +Provides a device-agnostic ``FakeVMM(VMMDevice)`` that stubs all abstract +methods with in-memory counters and ``os.pipe()`` for FD simulation. +Import this in any test that needs to monkeypatch the VMM singleton. +""" + +from __future__ import annotations + +import itertools +import os + +from gpu_memory_service.common.vmm import VMMDevice + + +class FakeVMM(VMMDevice): + """Device-agnostic VMMDevice mock for unit tests. + + Works regardless of whether the real backend is CUDA or XPU — all + VMMDevice methods are stubbed with in-memory counters and os.pipe() + for FD export/import simulation. + """ + + def __init__(self, devices: list[int] | None = None): + self._handles = itertools.count(1000) + self._vas = itertools.count(0x100000, 0x10000) + self._devices = devices if devices is not None else [0] + self.calls: list[tuple] = [] + + def ensure_initialized(self): + pass + + def synchronize(self): + pass + + def list_devices(self): + return self._devices + + def device_memory_info(self, device): + return (8 * 1024**3, 16 * 1024**3) + + def get_allocation_granularity(self, device): + return 4096 + + def create_tolerate_oom(self, size, device): + return (True, next(self._handles)) + + def release(self, handle): + pass + + def export_to_shareable_handle(self, handle): + read_fd, write_fd = os.pipe() + os.close(write_fd) + return read_fd + + def import_shareable_handle_close_fd(self, fd): + os.close(fd) + return next(self._handles) + + def address_reserve(self, size, granularity): + return next(self._vas) + + def address_free(self, va, size): + pass + + def map(self, va, size, handle): + pass + + def unmap(self, va, size): + pass + + def set_access(self, va, size, device, access): + pass + + def validate_pointer(self, va): + pass + + def runtime_check_result(self, result, name): + pass + + def runtime_set_device(self, device): + self.calls.append(("set_device", device)) + + def host_register(self, ptr, size): + pass + + def host_unregister(self, ptr): + pass + + def stream_create_nonblocking(self): + return "fake_stream" + + def stream_destroy(self, stream): + pass + + def stream_synchronize(self, stream): + pass + + def memcpy_h2d_async(self, dst_ptr, src_ptr, size, stream): + pass + + def memcpy_d2h_async(self, dst_ptr, src_ptr, size, stream): + pass diff --git a/lib/gpu_memory_service/tests/test_cli_server.py b/lib/gpu_memory_service/tests/test_cli_server.py index cbbb13c82193..3768bb71cff9 100644 --- a/lib/gpu_memory_service/tests/test_cli_server.py +++ b/lib/gpu_memory_service/tests/test_cli_server.py @@ -30,12 +30,14 @@ def test_child_command_launches_default_multi_tag_runner(): - assert server._child_command(3) == [ + assert server._child_command(3, "cuda") == [ sys.executable, "-m", "gpu_memory_service", "--device", "3", + "--device-type", + "cuda", ] diff --git a/lib/gpu_memory_service/tests/test_runtime_flows.py b/lib/gpu_memory_service/tests/test_runtime_flows.py index 6977aaa24ab9..0a1a61e78d1c 100644 --- a/lib/gpu_memory_service/tests/test_runtime_flows.py +++ b/lib/gpu_memory_service/tests/test_runtime_flows.py @@ -11,7 +11,6 @@ from __future__ import annotations import asyncio -import itertools import os import signal import socket @@ -33,7 +32,8 @@ if HAS_PYNVML: import pynvml -from gpu_memory_service.client import memory_manager as client_memory_manager +import gpu_memory_service.common.vmm as _vmm_module +from _fake_vmm import FakeVMM from gpu_memory_service.client.memory_manager import ( GMSClientMemoryManager, StaleMemoryLayoutError, @@ -47,7 +47,7 @@ GetRuntimeStateRequest, GetRuntimeStateResponse, ) -from gpu_memory_service.server import allocations as server_allocations +from gpu_memory_service.common.vmm import VMMDeviceType from gpu_memory_service.server.allocations import GMSAllocationManager from gpu_memory_service.server.fsm import ServerState from gpu_memory_service.server.rpc import GMSRPCServer @@ -229,81 +229,12 @@ async def _disconnect_rw_session(self) -> None: @pytest.fixture def running_gms(monkeypatch, tmp_path): - server_handles = itertools.count(1000) - client_handles = itertools.count(10000) - next_va = itertools.count(0x100000, 0x10000) - - monkeypatch.setattr(server_allocations, "cuda_ensure_initialized", lambda: None) - monkeypatch.setattr( - server_allocations, - "cumem_get_allocation_granularity", - lambda device: 4096, - ) - monkeypatch.setattr( - server_allocations, - "cumem_create_tolerate_oom", - lambda size, device: (True, next(server_handles)), - ) - monkeypatch.setattr(server_allocations, "cumem_release", lambda handle: None) - - def export_fd(handle: int) -> int: - read_fd, write_fd = os.pipe() - os.close(write_fd) - return read_fd - - monkeypatch.setattr( - server_allocations, "cumem_export_to_shareable_handle", export_fd - ) + fake_vmm = FakeVMM() - monkeypatch.setattr(client_memory_manager, "cuda_ensure_initialized", lambda: None) - monkeypatch.setattr( - client_memory_manager, - "cuda_set_current_device", - lambda device: None, - raising=False, - ) - monkeypatch.setattr( - client_memory_manager, - "cumem_get_allocation_granularity", - lambda device: 4096, - ) - monkeypatch.setattr( - client_memory_manager, - "cumem_create_tolerate_oom", - lambda size, device: (True, next(client_handles)), - ) - monkeypatch.setattr(client_memory_manager, "cuda_synchronize", lambda: None) - monkeypatch.setattr( - client_memory_manager, - "cumem_address_reserve", - lambda size, granularity: next(next_va), - ) - monkeypatch.setattr( - client_memory_manager, - "cumem_address_free", - lambda va, size: None, - ) - monkeypatch.setattr( - client_memory_manager, "cumem_map", lambda va, size, handle: None - ) - monkeypatch.setattr( - client_memory_manager, - "cumem_set_access", - lambda va, size, device, mode: None, - ) - monkeypatch.setattr(client_memory_manager, "cumem_unmap", lambda va, size: None) - monkeypatch.setattr(client_memory_manager, "cumem_release", lambda handle: None) - monkeypatch.setattr(client_memory_manager, "cuda_validate_pointer", lambda va: True) - - def import_fd(fd: int) -> int: - os.close(fd) - return next(client_handles) - - monkeypatch.setattr( - client_memory_manager, - "cumem_import_from_shareable_handle_close_fd", - import_fd, - ) + # Inject fake VMM into the process-global singleton so that + # GMSAllocationManager and GMSClientMemoryManager both use it. + monkeypatch.setattr(_vmm_module, "_vmm_instance", fake_vmm) + monkeypatch.setattr(_vmm_module, "_vmm_device_type", VMMDeviceType.CUDA) socket_path = str(tmp_path / "gms.sock") server = GMSRPCServer(socket_path, device=0, allocation_retry_interval=0.01) @@ -986,29 +917,20 @@ def test_scratch_reallocation_keeps_committed_allocation_on_cuda_granularity( async def test_allocation_manager_caches_exported_fd(monkeypatch): export_calls = 0 - monkeypatch.setattr(server_allocations, "cuda_ensure_initialized", lambda: None) - monkeypatch.setattr( - server_allocations, - "cumem_get_allocation_granularity", - lambda device: 4096, - ) - monkeypatch.setattr( - server_allocations, - "cumem_create_tolerate_oom", - lambda size, device: (True, 4242), - ) - monkeypatch.setattr(server_allocations, "cumem_release", lambda handle: None) + class _CountingVMM(FakeVMM): + def create_tolerate_oom(self, size, device): + return (True, 4242) - def export_fd(handle: int) -> int: - nonlocal export_calls - export_calls += 1 - read_fd, write_fd = os.pipe() - os.close(write_fd) - return read_fd + def export_to_shareable_handle(self, handle): + nonlocal export_calls + export_calls += 1 + read_fd, write_fd = os.pipe() + os.close(write_fd) + return read_fd - monkeypatch.setattr( - server_allocations, "cumem_export_to_shareable_handle", export_fd - ) + fake_vmm = _CountingVMM() + monkeypatch.setattr(_vmm_module, "_vmm_instance", fake_vmm) + monkeypatch.setattr(_vmm_module, "_vmm_device_type", VMMDeviceType.CUDA) allocations = GMSAllocationManager(device=0) info = await allocations.allocate(size=4096, tag="weights") diff --git a/lib/gpu_memory_service/tests/test_snapshot_loader.py b/lib/gpu_memory_service/tests/test_snapshot_loader.py index 1fbea8744160..41b3d89c3690 100644 --- a/lib/gpu_memory_service/tests/test_snapshot_loader.py +++ b/lib/gpu_memory_service/tests/test_snapshot_loader.py @@ -4,6 +4,7 @@ """Unit tests for the GMS snapshot loader CLI.""" import pytest +from _fake_vmm import FakeVMM try: from gpu_memory_service.cli.snapshot import loader @@ -30,7 +31,7 @@ def test_list_checkpoint_devices_requires_exact_visible_device_match( (tmp_path / "device-0-copy").mkdir() (tmp_path / "not-a-device").mkdir() (tmp_path / "device-1").write_text("not a directory", encoding="utf-8") - monkeypatch.setattr(loader.cuda_utils, "list_devices", lambda: [0, 2]) + monkeypatch.setattr(loader, "get_vmm", lambda: FakeVMM(devices=[0, 2])) assert loader._list_checkpoint_devices(str(tmp_path)) == [0, 2] @@ -53,7 +54,7 @@ def test_list_checkpoint_devices_rejects_mismatched_checkpoints( ): for dirname in checkpoint_dirs: (tmp_path / dirname).mkdir() - monkeypatch.setattr(loader.cuda_utils, "list_devices", lambda: visible_devices) + monkeypatch.setattr(loader, "get_vmm", lambda: FakeVMM(devices=visible_devices)) with pytest.raises(RuntimeError, match=expected): loader._list_checkpoint_devices(str(tmp_path)) @@ -61,6 +62,8 @@ def test_list_checkpoint_devices_rejects_mismatched_checkpoints( def test_load_device_sets_cuda_context_before_storage_client(monkeypatch): calls = [] + fake_vmm = FakeVMM(devices=[3]) + fake_vmm.calls = calls # share the calls list class FakeStorageClient: def __init__(self, **kwargs): @@ -80,11 +83,7 @@ def load_to_gms(self, input_dir, *, max_workers, clear_existing): monkeypatch.setattr(loader, "get_socket_path", lambda device: f"/tmp/gms-{device}") monkeypatch.setattr(loader, "GMSStorageClient", FakeStorageClient) - monkeypatch.setattr( - loader.cuda_utils, - "cuda_runtime_set_device", - lambda device: calls.append(("set_device", device)), - ) + monkeypatch.setattr(loader, "get_vmm", lambda: fake_vmm) loader._load_device( "/checkpoints/run/versions/1", diff --git a/lib/gpu_memory_service/tests/test_snapshot_nixl_staging.py b/lib/gpu_memory_service/tests/test_snapshot_nixl_staging.py index 32ee45c6c078..5374b5598836 100644 --- a/lib/gpu_memory_service/tests/test_snapshot_nixl_staging.py +++ b/lib/gpu_memory_service/tests/test_snapshot_nixl_staging.py @@ -23,6 +23,8 @@ allow_module_level=True, ) +from _fake_vmm import FakeVMM + pytestmark = [ pytest.mark.pre_merge, pytest.mark.unit, @@ -110,13 +112,10 @@ def fake_load_nixl_api(): assert allow_finish.wait(timeout=1.0) return FakeApi() - monkeypatch.setattr( - nixl_staging.cuda_utils, - "cuda_runtime_set_device", - lambda _device: None, - ) + fake_vmm = FakeVMM() + monkeypatch.setattr(nixl_staging, "get_vmm", lambda: fake_vmm) monkeypatch.setattr(nixl_staging, "load_nixl_api", fake_load_nixl_api) - monkeypatch.setattr(nixl_staging, "make_pinned_copy_slots", lambda _count: []) + monkeypatch.setattr(nixl_staging, "make_pinned_copy_slots", lambda _vmm, _count: []) session = _NixlPosixStagingTransferSession( backend_name="test-backend", diff --git a/lib/gpu_memory_service/tests/test_snapshot_saver.py b/lib/gpu_memory_service/tests/test_snapshot_saver.py index f1e77f21f312..bc2962780c7b 100644 --- a/lib/gpu_memory_service/tests/test_snapshot_saver.py +++ b/lib/gpu_memory_service/tests/test_snapshot_saver.py @@ -24,6 +24,13 @@ def test_save_device_sets_cuda_context_before_storage_client(monkeypatch): calls = [] + class FakeVMM: + def ensure_initialized(self): + calls.append(("ensure_initialized",)) + + def runtime_set_device(self, device): + calls.append(("set_device", device)) + class FakeStorageClient: def __init__(self, output_dir, **kwargs): calls.append(("init", output_dir, kwargs)) @@ -33,11 +40,7 @@ def save(self, *, max_workers): monkeypatch.setattr(saver, "get_socket_path", lambda device: f"/tmp/gms-{device}") monkeypatch.setattr(saver, "GMSStorageClient", FakeStorageClient) - monkeypatch.setattr( - saver.cuda_utils, - "cuda_runtime_set_device", - lambda device: calls.append(("set_device", device)), - ) + monkeypatch.setattr(saver, "get_vmm", lambda: FakeVMM()) saver._save_device( "/checkpoints/run/versions/1", @@ -48,9 +51,10 @@ def save(self, *, max_workers): [], ) - assert calls[0] == ("set_device", 3) - assert calls[1][0] == "init" - assert calls[1][1] == "/checkpoints/run/versions/1/device-3" - assert calls[1][2]["socket_path"] == "/tmp/gms-3" - assert calls[1][2]["device"] == 3 - assert calls[2] == ("save", {"max_workers": 8}) + assert calls[0] == ("ensure_initialized",) + assert calls[1] == ("set_device", 3) + assert calls[2][0] == "init" + assert calls[2][1] == "/checkpoints/run/versions/1/device-3" + assert calls[2][2]["socket_path"] == "/tmp/gms-3" + assert calls[2][2]["device"] == 3 + assert calls[3] == ("save", {"max_workers": 8}) diff --git a/lib/gpu_memory_service/tests/test_torch_integration.py b/lib/gpu_memory_service/tests/test_torch_integration.py index f6d3e3d918d0..7762d9dd4815 100644 --- a/lib/gpu_memory_service/tests/test_torch_integration.py +++ b/lib/gpu_memory_service/tests/test_torch_integration.py @@ -40,6 +40,7 @@ ) from gpu_memory_service.client.torch.tensor import _tensor_from_pointer from gpu_memory_service.common.locks import RequestedLockType +from gpu_memory_service.common.vmm import _reset_vmm_singleton from gpu_memory_service.server.rpc import GMSRPCServer pytestmark = [ @@ -136,6 +137,7 @@ def cancel() -> None: raise thread_error if os.path.exists(socket_path): os.unlink(socket_path) + _reset_vmm_singleton() def _make_gms_tensor( @@ -401,3 +403,32 @@ def test_materialized_module_from_gms_matches_plain_module_forward(running_gms): ) reader.close() + + +def test_integration_helper_without_explicit_init_vmm(tmp_path): + """Ensure GMSClientMemoryManager works without pre-seeding the VMM singleton. + + Integration paths (vLLM, SGLang, TRTLLM, gms-storage-client) construct + GMSClientMemoryManager without calling init_vmm() first. The lazy + auto-detection in get_vmm() must initialize transparently based on + available hardware. + """ + # Reset singleton to simulate a fresh process that never called init_vmm() + _reset_vmm_singleton() + + from gpu_memory_service.common.vmm import ( + _detect_device_type, + get_vmm, + get_vmm_device_type, + ) + + # get_vmm() should lazily auto-detect and initialize without raising + vmm = get_vmm() + assert vmm is not None + + # device type should match what auto-detection would pick + expected = _detect_device_type() + assert get_vmm_device_type() == expected + + # Clean up + _reset_vmm_singleton() diff --git a/tests/report_pytest_markers.py b/tests/report_pytest_markers.py index d8eac4708ee5..6a84f77de131 100755 --- a/tests/report_pytest_markers.py +++ b/tests/report_pytest_markers.py @@ -146,7 +146,9 @@ "gpu_memory_service.client.torch.tensor", "gpu_memory_service.common", "gpu_memory_service.common.locks", - "gpu_memory_service.common.cuda_utils", + "gpu_memory_service.common.vmm", + "gpu_memory_service.common.vmm.device", + "gpu_memory_service.common.vmm.cuda_utils", "gpu_memory_service.common.protocol", "gpu_memory_service.common.protocol.messages", "gpu_memory_service.common.protocol.wire",