wehub-resource-sync
134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Base class for platform device specification.
|
|
|
|
Each accelerator sub-package (``platform/cuda``, ``platform/musa``, ...)
|
|
provides a concrete :class:`DeviceSpec` subclass that describes how to
|
|
detect the device and which ops backend to load.
|
|
|
|
The :mod:`~lmcache.v1.platform` module discovers these
|
|
subclasses automatically at import time via ``pkgutil.iter_modules``:
|
|
it imports each sub-package, inspects its module namespace for
|
|
:class:`DeviceSpec` subclasses, instantiates them, and uses the
|
|
resulting objects for device detection and backend selection.
|
|
|
|
No manual registration (e.g. ``DEVICE_SPEC = ...``) is required --
|
|
simply defining the subclass in the sub-package's ``__init__.py`` is
|
|
enough.
|
|
|
|
:class:`DeviceSpec` itself is instantiable and doubles as the fallback
|
|
implementation used when no accelerator sub-package matches the
|
|
detected device: all capabilities default to a safe "no-op / False"
|
|
behaviour, and ``device_type`` / ``torch_module_name`` default to
|
|
``"cpu"``.
|
|
"""
|
|
|
|
# First Party
|
|
from lmcache.v1.platform.base_pin_memory import PinMemoryBackend
|
|
|
|
|
|
# TODO(chunxiaozheng): bind `DeviceIPCWrapper` with `DeviceSpec`?
|
|
class DeviceSpec:
|
|
"""Description of a hardware accelerator backend.
|
|
|
|
Subclasses override the properties / methods below to describe a
|
|
concrete accelerator. Defining a concrete subclass in a platform
|
|
sub-package's ``__init__.py`` is sufficient for auto-discovery.
|
|
|
|
Instantiating :class:`DeviceSpec` directly yields the fallback
|
|
implementation with "no-op / all False" semantics -- this is the
|
|
behaviour used for CPU-only or unknown device types.
|
|
"""
|
|
|
|
# Cached pin-memory backend instance (lazy-initialized).
|
|
_pin_backend_cache: PinMemoryBackend | None = None
|
|
|
|
@property
|
|
def device_type(self) -> str:
|
|
"""Device type string (e.g. ``"cuda"``, ``"musa"``, ``"mlu"``).
|
|
|
|
Defaults to ``"cpu"`` for the fallback implementation.
|
|
"""
|
|
return "cpu"
|
|
|
|
@property
|
|
def torch_module_name(self) -> str:
|
|
"""Attribute name on the ``torch`` package for the device module.
|
|
|
|
For example, ``"cuda"`` corresponds to ``torch.cuda``. Defaults
|
|
to ``"cpu"`` for the fallback implementation.
|
|
"""
|
|
return "cpu"
|
|
|
|
@property
|
|
def ops_module(self) -> str | None:
|
|
"""Fully-qualified module path for the compiled ops backend.
|
|
|
|
Return ``None`` if no custom ops are available (fallback only).
|
|
"""
|
|
return None
|
|
|
|
def is_available(self) -> bool:
|
|
"""Return ``True`` when the device is usable on this system.
|
|
|
|
This method must NOT import from ``lmcache.__init__`` to avoid
|
|
circular dependencies. Use ``import torch`` directly instead.
|
|
The fallback implementation always returns ``False`` so that
|
|
auto-detection never picks it up.
|
|
"""
|
|
return False
|
|
|
|
def is_handle_transfer_available(self) -> bool:
|
|
"""Return ``True`` when the device is usable for handle transfer."""
|
|
# TODO(chunxiaozheng): implement on subclasses
|
|
return True
|
|
|
|
@property
|
|
def pin_memory_backend(self) -> type[PinMemoryBackend] | None:
|
|
"""PinMemoryBackend subclass for this device, or None for default.
|
|
|
|
Subclasses that support host-memory pinning should override this
|
|
property and return the appropriate backend class. Use a lazy
|
|
import inside the property body to avoid heavy imports at class
|
|
definition time.
|
|
"""
|
|
return None
|
|
|
|
def _get_pin_backend(self) -> PinMemoryBackend:
|
|
"""Return the cached pin-memory backend, instantiating on first use."""
|
|
backend = self._pin_backend_cache
|
|
if backend is None:
|
|
backend_cls = self.pin_memory_backend or PinMemoryBackend
|
|
backend = backend_cls()
|
|
self._pin_backend_cache = backend
|
|
return backend
|
|
|
|
def pin_memory(self, ptr: int, size: int, flags: int = 0) -> bool:
|
|
"""Pin a host memory region for DMA access.
|
|
|
|
Args:
|
|
ptr: Raw pointer (data_ptr) to the memory region.
|
|
size: Size in bytes of the region to pin.
|
|
flags: Platform-specific registration flags (e.g.
|
|
``cudaHostRegisterDefault = 0``).
|
|
|
|
Returns:
|
|
True if pinning succeeded, False otherwise.
|
|
"""
|
|
return self._get_pin_backend().pin_memory(ptr, size, flags)
|
|
|
|
def unpin_memory(self, ptr: int) -> bool:
|
|
"""Unpin a previously pinned host memory region.
|
|
|
|
Args:
|
|
ptr: Raw pointer (data_ptr) to the memory region.
|
|
|
|
Returns:
|
|
True if unpinning succeeded, False otherwise.
|
|
"""
|
|
return self._get_pin_backend().unpin_memory(ptr)
|
|
|
|
@property
|
|
def is_pin_supported(self) -> bool:
|
|
"""Whether the current platform supports memory pinning."""
|
|
return self._get_pin_backend().is_pin_supported
|