235 lines
8.0 KiB
Python
235 lines
8.0 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Request handlers for GPU Memory Service."""
|
|
|
|
import hashlib
|
|
import logging
|
|
from dataclasses import dataclass
|
|
|
|
from gpu_memory_service.common.protocol.messages import (
|
|
AllocateRequest,
|
|
AllocateResponse,
|
|
ClearAllResponse,
|
|
FreeRequest,
|
|
FreeResponse,
|
|
GetAllocationRequest,
|
|
GetAllocationResponse,
|
|
GetAllocationStateResponse,
|
|
GetLockStateResponse,
|
|
GetStateHashResponse,
|
|
ListAllocationsRequest,
|
|
ListAllocationsResponse,
|
|
MetadataDeleteRequest,
|
|
MetadataDeleteResponse,
|
|
MetadataGetRequest,
|
|
MetadataGetResponse,
|
|
MetadataListRequest,
|
|
MetadataListResponse,
|
|
MetadataPutRequest,
|
|
MetadataPutResponse,
|
|
)
|
|
from gpu_memory_service.common.types import derive_state
|
|
|
|
from .memory_manager import AllocationNotFoundError, GMSServerMemoryManager
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MetadataEntry:
|
|
allocation_id: str
|
|
offset_bytes: int
|
|
value: bytes
|
|
|
|
|
|
class RequestHandler:
|
|
"""Handles allocation and metadata requests."""
|
|
|
|
def __init__(self, device: int = 0):
|
|
self._memory_manager = GMSServerMemoryManager(device)
|
|
self._metadata: dict[str, MetadataEntry] = {}
|
|
self._memory_layout_hash: str = (
|
|
"" # Hash of allocations + metadata, computed on commit
|
|
)
|
|
logger.info(f"RequestHandler initialized: device={device}")
|
|
|
|
@property
|
|
def granularity(self) -> int:
|
|
return self._memory_manager.granularity
|
|
|
|
def on_rw_abort(self) -> None:
|
|
"""Called when RW connection closes without commit."""
|
|
logger.warning("RW aborted; clearing allocations and metadata")
|
|
self._memory_manager.clear_all()
|
|
self._metadata.clear()
|
|
self._memory_layout_hash = ""
|
|
|
|
def on_commit(self) -> None:
|
|
"""Called when RW connection commits. Computes state hash."""
|
|
self._memory_layout_hash = self._compute_memory_layout_hash()
|
|
logger.info(f"Committed with state hash: {self._memory_layout_hash[:16]}...")
|
|
|
|
def _compute_memory_layout_hash(self) -> str:
|
|
"""Compute hash of current allocations + metadata."""
|
|
h = hashlib.sha256()
|
|
# Hash allocations (sorted by ID for determinism)
|
|
for info in sorted(
|
|
self._memory_manager.list_allocations(), key=lambda x: x.allocation_id
|
|
):
|
|
h.update(
|
|
f"{info.allocation_id}:{info.size}:{info.aligned_size}:{info.tag}".encode()
|
|
)
|
|
# Hash metadata (sorted by key for determinism)
|
|
for key in sorted(self._metadata.keys()):
|
|
entry = self._metadata[key]
|
|
h.update(f"{key}:{entry.allocation_id}:{entry.offset_bytes}:".encode())
|
|
h.update(entry.value)
|
|
return h.hexdigest()
|
|
|
|
def on_shutdown(self) -> None:
|
|
"""Called on server shutdown."""
|
|
if self._memory_manager.allocation_count > 0:
|
|
count = self._memory_manager.clear_all()
|
|
self._metadata.clear()
|
|
logger.info(f"Released {count} GPU allocations during shutdown")
|
|
|
|
# ==================== State Queries ====================
|
|
|
|
def handle_get_lock_state(
|
|
self,
|
|
has_rw: bool,
|
|
ro_count: int,
|
|
waiting_writers: int,
|
|
committed: bool,
|
|
) -> GetLockStateResponse:
|
|
"""Get lock/session state."""
|
|
state = derive_state(has_rw, ro_count, committed)
|
|
return GetLockStateResponse(
|
|
state=state.value,
|
|
has_rw_session=has_rw,
|
|
ro_session_count=ro_count,
|
|
waiting_writers=waiting_writers,
|
|
committed=committed,
|
|
is_ready=committed and not has_rw,
|
|
)
|
|
|
|
def handle_get_allocation_state(self) -> GetAllocationStateResponse:
|
|
"""Get allocation state."""
|
|
return GetAllocationStateResponse(
|
|
allocation_count=self._memory_manager.allocation_count,
|
|
total_bytes=self._memory_manager.total_bytes,
|
|
)
|
|
|
|
# ==================== Allocation Operations ====================
|
|
|
|
def handle_allocate(self, req: AllocateRequest) -> AllocateResponse:
|
|
"""Create physical memory allocation.
|
|
|
|
Requires RW connection (enforced by server).
|
|
"""
|
|
info = self._memory_manager.allocate(req.size, req.tag)
|
|
return AllocateResponse(
|
|
allocation_id=info.allocation_id,
|
|
size=info.size,
|
|
aligned_size=info.aligned_size,
|
|
)
|
|
|
|
def handle_export(self, allocation_id: str) -> tuple[GetAllocationResponse, int]:
|
|
"""Export allocation as POSIX FD.
|
|
|
|
Returns (response, fd). Caller must close fd after sending.
|
|
"""
|
|
fd = self._memory_manager.export_fd(allocation_id)
|
|
info = self._memory_manager.get_allocation(allocation_id)
|
|
response = GetAllocationResponse(
|
|
allocation_id=info.allocation_id,
|
|
size=info.size,
|
|
aligned_size=info.aligned_size,
|
|
tag=info.tag,
|
|
)
|
|
return response, fd
|
|
|
|
def handle_get_allocation(self, req: GetAllocationRequest) -> GetAllocationResponse:
|
|
"""Get allocation info."""
|
|
try:
|
|
info = self._memory_manager.get_allocation(req.allocation_id)
|
|
return GetAllocationResponse(
|
|
allocation_id=info.allocation_id,
|
|
size=info.size,
|
|
aligned_size=info.aligned_size,
|
|
tag=info.tag,
|
|
)
|
|
except AllocationNotFoundError:
|
|
raise ValueError(f"Unknown allocation: {req.allocation_id}") from None
|
|
|
|
def handle_list_allocations(
|
|
self, req: ListAllocationsRequest
|
|
) -> ListAllocationsResponse:
|
|
"""List all allocations."""
|
|
allocations = self._memory_manager.list_allocations(req.tag)
|
|
result = [
|
|
{
|
|
"allocation_id": info.allocation_id,
|
|
"size": info.size,
|
|
"aligned_size": info.aligned_size,
|
|
"tag": info.tag,
|
|
}
|
|
for info in allocations
|
|
]
|
|
return ListAllocationsResponse(allocations=result)
|
|
|
|
def handle_free(self, req: FreeRequest) -> FreeResponse:
|
|
"""Free single allocation.
|
|
|
|
Requires RW connection (enforced by server).
|
|
"""
|
|
success = self._memory_manager.free(req.allocation_id)
|
|
return FreeResponse(success=success)
|
|
|
|
def handle_clear_all(self) -> ClearAllResponse:
|
|
"""Clear all allocations and metadata.
|
|
|
|
Requires RW connection (enforced by server).
|
|
"""
|
|
count = self._memory_manager.clear_all()
|
|
self._metadata.clear()
|
|
return ClearAllResponse(cleared_count=count)
|
|
|
|
# ==================== Metadata Operations ====================
|
|
|
|
def handle_metadata_put(self, req: MetadataPutRequest) -> MetadataPutResponse:
|
|
self._metadata[req.key] = MetadataEntry(
|
|
req.allocation_id, req.offset_bytes, req.value
|
|
)
|
|
return MetadataPutResponse(success=True)
|
|
|
|
def handle_metadata_get(self, req: MetadataGetRequest) -> MetadataGetResponse:
|
|
entry = self._metadata.get(req.key)
|
|
if entry is None:
|
|
return MetadataGetResponse(found=False)
|
|
return MetadataGetResponse(
|
|
found=True,
|
|
allocation_id=entry.allocation_id,
|
|
offset_bytes=entry.offset_bytes,
|
|
value=entry.value,
|
|
)
|
|
|
|
def handle_metadata_delete(
|
|
self, req: MetadataDeleteRequest
|
|
) -> MetadataDeleteResponse:
|
|
return MetadataDeleteResponse(
|
|
deleted=self._metadata.pop(req.key, None) is not None
|
|
)
|
|
|
|
def handle_metadata_list(self, req: MetadataListRequest) -> MetadataListResponse:
|
|
keys = (
|
|
[k for k in self._metadata if k.startswith(req.prefix)]
|
|
if req.prefix
|
|
else list(self._metadata)
|
|
)
|
|
return MetadataListResponse(keys=sorted(keys))
|
|
|
|
def handle_get_memory_layout_hash(self) -> GetStateHashResponse:
|
|
return GetStateHashResponse(memory_layout_hash=self._memory_layout_hash)
|