349 lines
12 KiB
Python
349 lines
12 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""GPU Memory Service RPC Client.
|
|
|
|
Low-level RPC client stub. The client provides a simple interface for acquiring
|
|
locks and performing allocation operations. The socket connection IS the lock.
|
|
|
|
This module has NO PyTorch dependency.
|
|
|
|
Usage:
|
|
# Writer (acquires RW lock in constructor)
|
|
with GMSRPCClient(socket_path, lock_type=RequestedLockType.RW) as client:
|
|
alloc_id, aligned_size = client.allocate(size=1024*1024)
|
|
fd = client.export(alloc_id)
|
|
# ... write weights using fd ...
|
|
client.commit()
|
|
# Lock released on exit
|
|
|
|
# Reader (acquires RO lock in constructor)
|
|
client = GMSRPCClient(socket_path, lock_type=RequestedLockType.RO)
|
|
if client.committed: # Check if weights are valid
|
|
allocations = client.list_allocations()
|
|
for alloc in allocations:
|
|
fd = client.export(alloc["allocation_id"])
|
|
# ... import and map fd ...
|
|
# Keep connection open during inference!
|
|
# client.close() only when done with inference
|
|
"""
|
|
|
|
import logging
|
|
import socket
|
|
from typing import Dict, List, Optional, Tuple, Type, TypeVar
|
|
|
|
from gpu_memory_service.common.protocol.messages import (
|
|
AllocateRequest,
|
|
AllocateResponse,
|
|
ClearAllRequest,
|
|
ClearAllResponse,
|
|
CommitRequest,
|
|
CommitResponse,
|
|
ErrorResponse,
|
|
ExportRequest,
|
|
FreeRequest,
|
|
FreeResponse,
|
|
GetAllocationRequest,
|
|
GetAllocationResponse,
|
|
GetAllocationStateRequest,
|
|
GetAllocationStateResponse,
|
|
GetLockStateRequest,
|
|
GetLockStateResponse,
|
|
GetStateHashRequest,
|
|
GetStateHashResponse,
|
|
HandshakeRequest,
|
|
HandshakeResponse,
|
|
ListAllocationsRequest,
|
|
ListAllocationsResponse,
|
|
MetadataDeleteRequest,
|
|
MetadataDeleteResponse,
|
|
MetadataGetRequest,
|
|
MetadataGetResponse,
|
|
MetadataListRequest,
|
|
MetadataListResponse,
|
|
MetadataPutRequest,
|
|
MetadataPutResponse,
|
|
)
|
|
from gpu_memory_service.common.protocol.wire import recv_message_sync, send_message_sync
|
|
from gpu_memory_service.common.types import (
|
|
RW_REQUIRED,
|
|
GrantedLockType,
|
|
RequestedLockType,
|
|
)
|
|
|
|
T = TypeVar("T")
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class GMSRPCClient:
|
|
"""GPU Memory Service RPC Client.
|
|
|
|
CRITICAL: Socket connection IS the lock.
|
|
- Constructor blocks until lock is acquired
|
|
- close() releases the lock
|
|
- committed property tells readers if weights are valid
|
|
|
|
For writers (lock_type=RequestedLockType.RW):
|
|
- Use context manager (with statement) for automatic lock release
|
|
- Call commit() after weights are written
|
|
- Call clear_all() before loading new model
|
|
|
|
For readers (lock_type=RequestedLockType.RO):
|
|
- Check committed property after construction
|
|
- Keep connection open during inference lifetime
|
|
- Only call close() when shutting down or allowing weight updates
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
socket_path: str,
|
|
lock_type: RequestedLockType = RequestedLockType.RO,
|
|
timeout_ms: Optional[int] = None,
|
|
):
|
|
"""Connect to Allocation Server and acquire lock.
|
|
|
|
Args:
|
|
socket_path: Path to server's Unix domain socket
|
|
lock_type: Requested lock type (RW, RO, or RW_OR_RO)
|
|
timeout_ms: Timeout in milliseconds for lock acquisition.
|
|
None means wait indefinitely.
|
|
|
|
Raises:
|
|
ConnectionError: If connection fails
|
|
TimeoutError: If timeout_ms expires waiting for lock
|
|
"""
|
|
self.socket_path = socket_path
|
|
self._requested_lock_type = lock_type
|
|
self._socket: Optional[socket.socket] = None
|
|
self._recv_buffer = bytearray()
|
|
self._committed = False
|
|
self._granted_lock_type: Optional[GrantedLockType] = None
|
|
|
|
# Connect and acquire lock
|
|
self._connect(timeout_ms=timeout_ms)
|
|
|
|
def _connect(self, timeout_ms: Optional[int]) -> None:
|
|
"""Connect to server and perform handshake (lock acquisition)."""
|
|
self._socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
try:
|
|
self._socket.connect(self.socket_path)
|
|
except FileNotFoundError:
|
|
self._socket.close()
|
|
self._socket = None
|
|
raise ConnectionError(f"Server not running at {self.socket_path}") from None
|
|
except Exception as e:
|
|
self._socket.close()
|
|
self._socket = None
|
|
raise ConnectionError(f"Failed to connect: {e}") from e
|
|
|
|
# Handshake I/O — clean up socket on any failure
|
|
try:
|
|
request = HandshakeRequest(
|
|
lock_type=self._requested_lock_type,
|
|
timeout_ms=timeout_ms,
|
|
)
|
|
send_message_sync(self._socket, request)
|
|
|
|
# May block waiting for lock
|
|
response, _, self._recv_buffer = recv_message_sync(
|
|
self._socket, self._recv_buffer
|
|
)
|
|
except Exception:
|
|
self._socket.close()
|
|
self._socket = None
|
|
raise
|
|
|
|
if isinstance(response, ErrorResponse):
|
|
self._socket.close()
|
|
self._socket = None
|
|
raise ConnectionError(f"Handshake error: {response.error}")
|
|
|
|
if not isinstance(response, HandshakeResponse):
|
|
self._socket.close()
|
|
self._socket = None
|
|
raise ConnectionError(f"Unexpected response: {type(response)}")
|
|
|
|
if not response.success:
|
|
self._socket.close()
|
|
self._socket = None
|
|
raise TimeoutError("Timeout waiting for lock")
|
|
|
|
self._committed = response.committed
|
|
# Store granted lock type (may differ from requested for rw_or_ro mode)
|
|
if response.granted_lock_type is not None:
|
|
self._granted_lock_type = response.granted_lock_type
|
|
elif self._requested_lock_type == RequestedLockType.RW:
|
|
self._granted_lock_type = GrantedLockType.RW
|
|
else:
|
|
self._granted_lock_type = GrantedLockType.RO
|
|
logger.info(
|
|
f"Connected with {self._requested_lock_type.value} lock (granted={self._granted_lock_type.value}), "
|
|
f"committed={self._committed}"
|
|
)
|
|
|
|
@property
|
|
def committed(self) -> bool:
|
|
"""Check if weights are committed (valid)."""
|
|
return self._committed
|
|
|
|
@property
|
|
def lock_type(self) -> Optional[GrantedLockType]:
|
|
"""Get the lock type actually granted by the server.
|
|
|
|
For rw_or_ro mode, this tells you whether RW or RO was granted.
|
|
"""
|
|
return self._granted_lock_type
|
|
|
|
@property
|
|
def is_connected(self) -> bool:
|
|
"""Check if client is connected."""
|
|
return self._socket is not None
|
|
|
|
def _send_recv(self, request) -> Tuple[object, int]:
|
|
"""Send request and receive response. Returns (response, fd)."""
|
|
if not self._socket:
|
|
raise RuntimeError("Client not connected")
|
|
|
|
send_message_sync(self._socket, request)
|
|
response, fd, self._recv_buffer = recv_message_sync(
|
|
self._socket, self._recv_buffer
|
|
)
|
|
|
|
if isinstance(response, ErrorResponse):
|
|
raise RuntimeError(f"Server error: {response.error}")
|
|
|
|
return response, fd
|
|
|
|
def _call(self, request, response_type: Type[T]) -> T:
|
|
"""Send request, validate response type, return typed response."""
|
|
if type(request) in RW_REQUIRED and self.lock_type != GrantedLockType.RW:
|
|
raise RuntimeError("Operation requires RW connection")
|
|
response, _ = self._send_recv(request)
|
|
if not isinstance(response, response_type):
|
|
raise RuntimeError(f"Unexpected response: {type(response)}")
|
|
return response
|
|
|
|
def get_lock_state(self) -> GetLockStateResponse:
|
|
return self._call(GetLockStateRequest(), GetLockStateResponse)
|
|
|
|
def get_allocation_state(self) -> GetAllocationStateResponse:
|
|
return self._call(GetAllocationStateRequest(), GetAllocationStateResponse)
|
|
|
|
def is_ready(self) -> bool:
|
|
return self.committed
|
|
|
|
def commit(self) -> bool:
|
|
"""Commit weights and release RW lock. Returns True on success."""
|
|
if CommitRequest in RW_REQUIRED and self.lock_type != GrantedLockType.RW:
|
|
raise RuntimeError("Operation requires RW connection")
|
|
|
|
try:
|
|
response, _ = self._send_recv(CommitRequest())
|
|
ok = isinstance(response, CommitResponse) and response.success
|
|
except (ConnectionResetError, BrokenPipeError, OSError) as e:
|
|
# Server closes RW socket as part of commit
|
|
logger.debug(
|
|
f"Commit saw socket error ({type(e).__name__}); verifying via RO connect"
|
|
)
|
|
self.close()
|
|
try:
|
|
ro = GMSRPCClient(
|
|
self.socket_path, lock_type=RequestedLockType.RO, timeout_ms=1000
|
|
)
|
|
try:
|
|
ok = ro.committed
|
|
finally:
|
|
ro.close()
|
|
except TimeoutError:
|
|
ok = False
|
|
|
|
if ok:
|
|
self._committed = True
|
|
self.close()
|
|
logger.info("Committed weights and released RW connection")
|
|
return True
|
|
|
|
return False
|
|
|
|
def allocate(self, size: int, tag: str = "default") -> Tuple[str, int]:
|
|
"""Returns (allocation_id, aligned_size)."""
|
|
r = self._call(AllocateRequest(size=size, tag=tag), AllocateResponse)
|
|
return r.allocation_id, r.aligned_size
|
|
|
|
def export(self, allocation_id: str) -> int:
|
|
"""Export allocation as POSIX FD. Caller must close."""
|
|
_, fd = self._send_recv(ExportRequest(allocation_id=allocation_id))
|
|
if fd < 0:
|
|
raise RuntimeError("No FD received from server")
|
|
return fd
|
|
|
|
def get_allocation(self, allocation_id: str) -> GetAllocationResponse:
|
|
return self._call(
|
|
GetAllocationRequest(allocation_id=allocation_id), GetAllocationResponse
|
|
)
|
|
|
|
def list_allocations(self, tag: Optional[str] = None) -> List[Dict]:
|
|
return self._call(
|
|
ListAllocationsRequest(tag=tag), ListAllocationsResponse
|
|
).allocations
|
|
|
|
def free(self, allocation_id: str) -> bool:
|
|
return self._call(
|
|
FreeRequest(allocation_id=allocation_id), FreeResponse
|
|
).success
|
|
|
|
def clear_all(self) -> int:
|
|
return self._call(ClearAllRequest(), ClearAllResponse).cleared_count
|
|
|
|
def metadata_put(
|
|
self, key: str, allocation_id: str, offset_bytes: int, value: bytes
|
|
) -> bool:
|
|
req = MetadataPutRequest(
|
|
key=key, allocation_id=allocation_id, offset_bytes=offset_bytes, value=value
|
|
)
|
|
return self._call(req, MetadataPutResponse).success
|
|
|
|
def metadata_get(self, key: str) -> Optional[tuple[str, int, bytes]]:
|
|
"""Returns (allocation_id, offset_bytes, value) or None if not found."""
|
|
r = self._call(MetadataGetRequest(key=key), MetadataGetResponse)
|
|
return (r.allocation_id, r.offset_bytes, r.value) if r.found else None
|
|
|
|
def metadata_delete(self, key: str) -> bool:
|
|
return self._call(
|
|
MetadataDeleteRequest(key=key), MetadataDeleteResponse
|
|
).deleted
|
|
|
|
def metadata_list(self, prefix: str = "") -> List[str]:
|
|
return self._call(MetadataListRequest(prefix=prefix), MetadataListResponse).keys
|
|
|
|
def get_memory_layout_hash(self) -> str:
|
|
"""Get state hash (hash of allocations + metadata). Empty if not committed."""
|
|
return self._call(
|
|
GetStateHashRequest(), GetStateHashResponse
|
|
).memory_layout_hash
|
|
|
|
def close(self) -> None:
|
|
"""Close connection and release lock."""
|
|
if self._socket:
|
|
try:
|
|
self._socket.close()
|
|
except Exception:
|
|
pass
|
|
self._socket = None
|
|
lock_str = self.lock_type.value if self.lock_type else "unknown"
|
|
logger.info(f"Closed {lock_str} connection")
|
|
|
|
def __enter__(self) -> "GMSRPCClient":
|
|
"""Context manager entry."""
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
|
"""Context manager exit."""
|
|
self.close()
|
|
|
|
def __del__(self):
|
|
"""Destructor: warn if connection not closed."""
|
|
if self._socket:
|
|
logger.warning("GMSRPCClient not closed properly")
|