Hongsheng Liuandhsliuustc0106
c608910d20
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com> Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
64 lines
1.8 KiB
Python
64 lines
1.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Platform-neutral lifecycle for registering existing host mappings."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from typing import Protocol
|
|
|
|
import torch
|
|
|
|
from vllm_omni.host_weight_runtime import MappedHostRegion
|
|
|
|
|
|
class HostRegistrationError(RuntimeError):
|
|
"""Raised when existing host ranges cannot be registered safely."""
|
|
|
|
|
|
class HostRegistrationBudgetError(HostRegistrationError):
|
|
"""Raised before mutation when the complete mapping exceeds its budget."""
|
|
|
|
|
|
class HostRegistrationCleanupError(HostRegistrationError):
|
|
"""Raised when registered ranges cannot be released safely."""
|
|
|
|
def __init__(self, message: str, *, active_registration: HostRegistration | None = None) -> None:
|
|
super().__init__(message)
|
|
self.active_registration = active_registration
|
|
|
|
|
|
class HostRegistration(Protocol):
|
|
"""Transport-owned registration for one process-local host mapping."""
|
|
|
|
@property
|
|
def total_bytes(self) -> int: ...
|
|
|
|
@property
|
|
def region_count(self) -> int: ...
|
|
|
|
def close(self) -> tuple[str, ...]: ...
|
|
|
|
|
|
def register_host_mappings(
|
|
regions: Sequence[MappedHostRegion],
|
|
*,
|
|
device: torch.device,
|
|
max_bytes: int | None,
|
|
) -> HostRegistration:
|
|
"""Register immutable ranges with the active platform implementation."""
|
|
if device.type == "cuda":
|
|
from .cuda_host_registration import CudaHostRegistration
|
|
|
|
return CudaHostRegistration.create(regions, max_bytes=max_bytes)
|
|
raise HostRegistrationError(f"host-mapping registration is not supported on {device.type}")
|
|
|
|
|
|
__all__ = [
|
|
"HostRegistration",
|
|
"HostRegistrationBudgetError",
|
|
"HostRegistrationCleanupError",
|
|
"HostRegistrationError",
|
|
"register_host_mappings",
|
|
]
|