wehub-resource-sync
739 lines
31 KiB
Python
739 lines
31 KiB
Python
import functools
|
|
import glob
|
|
import logging
|
|
import os
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from collections import OrderedDict
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
|
|
|
import ray
|
|
from ray._private.ray_constants import (
|
|
NIXL_REMOTE_AGENT_CACHE_MAXSIZE,
|
|
)
|
|
from ray.experimental.rdt.nixl_memory_pool import MemoryPoolManager
|
|
from ray.experimental.rdt.tensor_transport_manager import (
|
|
CommunicatorMetadata,
|
|
FetchRequest,
|
|
TensorTransportManager,
|
|
TensorTransportMetadata,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def _is_efa_available() -> bool:
|
|
"""Detect whether AWS EFA (Elastic Fabric Adapter) devices are present.
|
|
|
|
A bare host exposes ``efa*`` netdevs, but inside a container/Kubernetes pod
|
|
netdevs are network-namespaced away and only the rdma-verbs devices under
|
|
``/sys/class/infiniband`` are mounted in. Those verbs devices are not
|
|
EFA-specific -- ordinary InfiniBand/RoCE NICs appear there too -- so we
|
|
confirm each one is bound to the kernel ``efa`` driver before treating it as
|
|
EFA. Without that check, non-AWS RDMA nodes would wrongly auto-select the
|
|
LIBFABRIC backend instead of UCX.
|
|
"""
|
|
if glob.glob("/sys/class/net/efa*"):
|
|
return True
|
|
for ib_dev in glob.glob("/sys/class/infiniband/*"):
|
|
# A stale or broken sysfs entry shouldn't abort the scan; skip it and
|
|
# keep looking (defaulting to UCX if nothing resolves to the efa driver).
|
|
try:
|
|
driver = os.path.realpath(os.path.join(ib_dev, "device", "driver"))
|
|
except OSError:
|
|
continue
|
|
if os.path.basename(driver) == "efa":
|
|
return True
|
|
return False
|
|
|
|
|
|
def _nixl_transport_available_in_process() -> bool:
|
|
"""Returns whether the NIXL tensor transport can be initialized in this process.
|
|
|
|
Returns:
|
|
True if the NIXL agent initializes successfully, False on any failure
|
|
(e.g. nixl not installed, LIBFABRIC/EFA probe failure, or other backend
|
|
init errors).
|
|
"""
|
|
try:
|
|
from ray.experimental.rdt.util import get_tensor_transport_manager
|
|
|
|
get_tensor_transport_manager("NIXL").get_nixl_agent()
|
|
return True
|
|
except Exception:
|
|
logger.debug("NIXL tensor transport unavailable on actor.", exc_info=True)
|
|
return False
|
|
|
|
|
|
@dataclass
|
|
class NixlCommunicatorMetadata(CommunicatorMetadata):
|
|
"""Metadata for the NIXL communicator."""
|
|
|
|
|
|
@dataclass
|
|
class NixlTransportMetadata(TensorTransportMetadata):
|
|
"""Metadata for tensors stored in the GPU object store for NIXL transport.
|
|
|
|
Args:
|
|
nixl_serialized_descs: Serialized tensor descriptors for NIXL transport.
|
|
nixl_agent_meta: The additional metadata of the remote NIXL agent.
|
|
nixl_agent_name: The name of the NIXL agent.
|
|
nixl_agent_meta_version: The version of the NIXL agent metadata.
|
|
"""
|
|
|
|
nixl_serialized_descs: Optional[bytes] = None
|
|
nixl_agent_meta: Optional[bytes] = None
|
|
nixl_agent_name: Optional[str] = None
|
|
nixl_agent_meta_version: Optional[int] = 0
|
|
|
|
__eq__ = object.__eq__
|
|
__hash__ = object.__hash__
|
|
|
|
|
|
@dataclass
|
|
class TensorDesc:
|
|
# nixlRegDList handle, or None for pool-managed tensors (pool memory is
|
|
# registered once at pool creation, so individual tensors don't need their
|
|
# own NIXL registration).
|
|
reg_desc: Any
|
|
# tracks the number of NIXL metadata containing the tensor.
|
|
metadata_count: int
|
|
|
|
|
|
@dataclass
|
|
class NixlFetchRequest(FetchRequest):
|
|
"""NIXL-specific FetchRequest carrying the async transfer state.
|
|
|
|
Returned by fetch_multiple_tensors and consumed by wait_fetch_complete.
|
|
|
|
Args:
|
|
obj_id: Inherited. The object ID for the transfer, used for abort checks and cleanup.
|
|
tensors: Inherited. Pre-allocated output tensors (populated before the transfer starts).
|
|
xfer_handle: NIXL transfer request handle.
|
|
nixl_agent: Reference to the NIXL agent.
|
|
remote_name: Name of the remote NIXL agent.
|
|
remove_tensor_descs: Whether to remove tensor descriptors from the cache during cleanup.
|
|
"""
|
|
|
|
xfer_handle: Any = None
|
|
nixl_agent: Any = None
|
|
remote_name: Optional[str] = None
|
|
remove_tensor_descs: bool = False
|
|
transport: Any = None
|
|
|
|
def __del__(self):
|
|
if self.transport is not None:
|
|
self.transport._cleanup_transfer(
|
|
self.obj_id,
|
|
self.tensors,
|
|
self.xfer_handle,
|
|
self.remote_name,
|
|
self.remove_tensor_descs,
|
|
)
|
|
|
|
|
|
class NixlTensorTransport(TensorTransportManager):
|
|
def __init__(self):
|
|
# This is lazily initialized because it requires NIXL to actually be installed and we want to allow an owner that is just coordinating to not need to have NIXL installed.
|
|
self._nixl_agent = None
|
|
self._aborted_transfer_obj_ids = set()
|
|
self._aborted_transfer_obj_ids_lock = threading.Lock()
|
|
# Mapping from tensor storage data pointer to the NIXL descriptor and reference count.
|
|
# Unlike _managed_meta_nixl, we only deregister tensors when ALL metadata containing the tensor is freed.
|
|
# For pool-managed tensors, reg_desc is None and the pool block is returned instead of deregistering.
|
|
self._tensor_desc_cache: Dict[int, TensorDesc] = {}
|
|
# Mapping from object ID to the NIXL managed meta.
|
|
# The lifetime of _managed_meta_nixl is tied to the object ref and freed when the ref goes out of scope.
|
|
self._managed_meta_nixl: Dict[str, Any] = {}
|
|
# Lock protecting _tensor_desc_cache and _managed_meta_nixl since they can be
|
|
# accessed from the main task execution thread or the _ray_system thread.
|
|
self._cache_lock = threading.RLock()
|
|
# LRU cache of remote agent names. When full, the least
|
|
# recently used remote agent is evicted and remove_remote_agent is called.
|
|
self._remote_agents: OrderedDict = OrderedDict()
|
|
# Increment the version whenever memory is deregistered.
|
|
self._nixl_agent_meta_version = 0
|
|
self._memory_pool: Optional[MemoryPoolManager] = None
|
|
# The NIXL backend the agent was actually created with ("UCX" or "LIBFABRIC").
|
|
self._backend: Optional[str] = None
|
|
|
|
def tensor_transport_backend(self) -> str:
|
|
return "NIXL"
|
|
|
|
@staticmethod
|
|
def is_one_sided() -> bool:
|
|
return True
|
|
|
|
@staticmethod
|
|
def can_abort_transport() -> bool:
|
|
return True
|
|
|
|
def register_nixl_memory(self, tensor: "torch.Tensor") -> None:
|
|
"""Registers the tensor's memory with NIXL and bumps the reference count so the memory region is never deregistered."""
|
|
self._add_tensor_descs([tensor])
|
|
|
|
def register_nixl_memory_pool(self, size: int, device: "torch.device") -> None:
|
|
"""Pre-allocates a memory pool and registers it with NIXL.
|
|
|
|
Args:
|
|
size: Size of the memory pool in bytes.
|
|
device: Device to allocate the pool on (cpu or cuda).
|
|
|
|
Raises:
|
|
ValueError: If a memory pool is already registered.
|
|
"""
|
|
if self._memory_pool is not None:
|
|
raise ValueError(
|
|
"A memory pool is already registered. "
|
|
"Only one memory pool is supported."
|
|
)
|
|
nixl_agent = self.get_nixl_agent()
|
|
pool = MemoryPoolManager(pool_size=size, device=device)
|
|
nixl_agent.register_memory(pool.get_pool_tensor())
|
|
self._memory_pool = pool
|
|
|
|
def deregister_nixl_memory(self, tensor: "torch.Tensor") -> None:
|
|
"""Decrements the reference count for the tensor's NIXL memory registration.
|
|
If the count reaches 0, the memory is deregistered from NIXL.
|
|
"""
|
|
self._remove_tensor_descs([tensor])
|
|
|
|
def select_backend(self) -> str:
|
|
"""Returns the NIXL backend to attempt.
|
|
|
|
Prefers LIBFABRIC when EFA devices are present and UCX everywhere else.
|
|
LIBFABRIC requires GPUDirect (GDR) for CUDA registration; if it isn't
|
|
available, ``_add_tensor_descs`` surfaces a clear error at registration
|
|
time with backend-specific troubleshooting guidance.
|
|
"""
|
|
return "LIBFABRIC" if _is_efa_available() else "UCX"
|
|
|
|
def _make_nixl_agent(self, backend: str):
|
|
"""Creates a NIXL agent configured with the given backend."""
|
|
from nixl._api import nixl_agent, nixl_agent_config
|
|
|
|
agent_config = nixl_agent_config(backends=[backend])
|
|
ctx = ray.get_runtime_context()
|
|
actor_id = ctx.get_actor_id()
|
|
if actor_id is None:
|
|
# If the actor id is None, it means the current process is a driver.
|
|
import uuid
|
|
|
|
actor_id = f"RAY-DRIVER-{uuid.uuid4()}"
|
|
return nixl_agent(actor_id, agent_config)
|
|
|
|
def get_nixl_agent(self):
|
|
"""Returns the NIXL agent, building it once on first use."""
|
|
if self._nixl_agent is None:
|
|
self._nixl_agent = self._init_nixl_agent()
|
|
return self._nixl_agent
|
|
|
|
def _init_nixl_agent(self):
|
|
"""Builds the NIXL agent for the selected backend."""
|
|
backend = self.select_backend()
|
|
agent = self._make_nixl_agent(backend)
|
|
self._backend = backend
|
|
logger.info("Using NIXL backend: %s", backend)
|
|
return agent
|
|
|
|
def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool:
|
|
# TODO(dayshah): This is called on a .remote RDT call, so it's quite expensive.
|
|
def __ray_actor_has_tensor_transport__(
|
|
self: "ray.actor.ActorHandle",
|
|
) -> bool:
|
|
return _nixl_transport_available_in_process()
|
|
|
|
return ray.get(
|
|
actor.__ray_call__.options(concurrency_group="_ray_system").remote(
|
|
__ray_actor_has_tensor_transport__
|
|
)
|
|
)
|
|
|
|
def extract_tensor_transport_metadata(
|
|
self,
|
|
obj_id: str,
|
|
rdt_object: List["torch.Tensor"],
|
|
) -> NixlTransportMetadata:
|
|
import torch
|
|
|
|
with self._cache_lock:
|
|
device = None
|
|
tensor_meta = []
|
|
|
|
if rdt_object:
|
|
# We assume all tensors in one RDT object have the same device type,
|
|
# but we don't assume they're all on the same device.
|
|
devices = set()
|
|
device = rdt_object[0].device
|
|
for t in rdt_object:
|
|
if t.device.type != device.type:
|
|
raise ValueError(
|
|
"All tensors in an RDT object must have the same device type."
|
|
)
|
|
if not t.is_contiguous():
|
|
raise ValueError(
|
|
"All tensors in an RDT object must be contiguous."
|
|
)
|
|
tensor_meta.append((t.shape, t.dtype))
|
|
devices.add(t.device)
|
|
if device.type == "cuda":
|
|
# We have to synchronize before memory registration to assure the
|
|
# object has been created because nixl doesn't guarantee it will.
|
|
for dev in devices:
|
|
torch.cuda.synchronize(dev)
|
|
|
|
nixl_agent = self.get_nixl_agent()
|
|
# Use the pool only when every tensor lives on the exact same
|
|
# device as the pool, AND no tensor already has an existing
|
|
# NIXL registration (via register_nixl_memory).
|
|
pool_eligible = (
|
|
self._memory_pool is not None
|
|
and all(
|
|
t.device == self._memory_pool.get_pool_tensor().device
|
|
for t in rdt_object
|
|
)
|
|
and not any(self._tensor_memory_registered(t) for t in rdt_object)
|
|
)
|
|
if pool_eligible:
|
|
xfer_descs = self._allocate_pool_xfer_descs(rdt_object)
|
|
else:
|
|
self._add_tensor_descs(rdt_object)
|
|
xfer_descs = nixl_agent.get_xfer_descs(rdt_object)
|
|
|
|
serialized_descs = nixl_agent.get_serialized_descs(xfer_descs)
|
|
agent_meta = nixl_agent.get_agent_metadata()
|
|
agent_name = nixl_agent.name
|
|
agent_meta_version = self._nixl_agent_meta_version
|
|
else:
|
|
serialized_descs, agent_meta = None, None
|
|
agent_name, agent_meta_version = None, None
|
|
|
|
ret = NixlTransportMetadata(
|
|
tensor_meta=tensor_meta,
|
|
tensor_device=device.type if device else None,
|
|
nixl_serialized_descs=serialized_descs,
|
|
nixl_agent_meta=agent_meta,
|
|
nixl_agent_name=agent_name,
|
|
nixl_agent_meta_version=agent_meta_version,
|
|
)
|
|
self._put_meta(obj_id, ret)
|
|
return ret
|
|
|
|
def get_communicator_metadata(
|
|
self,
|
|
src_actor: "ray.actor.ActorHandle",
|
|
dst_actor: "ray.actor.ActorHandle",
|
|
backend: Optional[str] = None,
|
|
) -> NixlCommunicatorMetadata:
|
|
return NixlCommunicatorMetadata()
|
|
|
|
def fetch_multiple_tensors(
|
|
self,
|
|
obj_id: str,
|
|
tensor_transport_metadata: TensorTransportMetadata,
|
|
communicator_metadata: CommunicatorMetadata,
|
|
target_buffers: Optional[List["torch.Tensor"]] = None,
|
|
) -> NixlFetchRequest:
|
|
"""Initiates an async transfer for multiple tensors.
|
|
|
|
This triggers the transfer but does not wait for completion.
|
|
Call wait_fetch_complete(fetch_request) to wait for the transfer to
|
|
finish and retrieve the tensors.
|
|
|
|
Args:
|
|
obj_id: The object ID for the transfer.
|
|
tensor_transport_metadata: Metadata for the tensor transport.
|
|
communicator_metadata: Metadata for the communicator.
|
|
target_buffers: Optional pre-allocated buffers to receive tensors into.
|
|
|
|
Returns:
|
|
A NixlFetchRequest carrying the async transfer state.
|
|
"""
|
|
from ray.experimental.rdt.util import (
|
|
create_empty_tensors_from_metadata,
|
|
)
|
|
|
|
tensors = target_buffers or create_empty_tensors_from_metadata(
|
|
tensor_transport_metadata
|
|
)
|
|
|
|
assert isinstance(tensor_transport_metadata, NixlTransportMetadata)
|
|
assert isinstance(communicator_metadata, NixlCommunicatorMetadata)
|
|
|
|
nixl_serialized_descs = tensor_transport_metadata.nixl_serialized_descs
|
|
remote_nixl_agent_meta = tensor_transport_metadata.nixl_agent_meta
|
|
|
|
with self._aborted_transfer_obj_ids_lock:
|
|
if obj_id in self._aborted_transfer_obj_ids:
|
|
self._aborted_transfer_obj_ids.remove(obj_id)
|
|
raise RuntimeError(f"NIXL transfer aborted for object id: {obj_id}")
|
|
|
|
remote_name = None
|
|
xfer_handle = None
|
|
added_tensor_descs = False
|
|
|
|
assert tensors
|
|
|
|
try:
|
|
nixl_agent = self.get_nixl_agent()
|
|
remote_xfer_descs = nixl_agent.deserialize_descs(nixl_serialized_descs)
|
|
# This creates a placeholder for the tensor in the tensor_desc_cache even though it doesn't have an object ref for caching purposes.
|
|
self._add_tensor_descs(tensors)
|
|
added_tensor_descs = True
|
|
local_xfer_descs = nixl_agent.get_xfer_descs(tensors)
|
|
|
|
remote_name = tensor_transport_metadata.nixl_agent_name
|
|
remote_agent_meta_version = (
|
|
tensor_transport_metadata.nixl_agent_meta_version
|
|
)
|
|
|
|
# Nixl agent reuse is enabled.
|
|
if NIXL_REMOTE_AGENT_CACHE_MAXSIZE > 0:
|
|
if remote_name in self._remote_agents:
|
|
# If the remote agent metadata version is different from the cached one,
|
|
# it means there was memory deregistered. We need to remove the remote agent
|
|
# before adding it, because `nixlRemoteSection` currently does not support
|
|
# updating descriptor list in such a case (there is potential memory overlap).
|
|
if remote_agent_meta_version != self._remote_agents[remote_name]:
|
|
nixl_agent.remove_remote_agent(remote_name)
|
|
self._remote_agents.move_to_end(remote_name)
|
|
elif len(self._remote_agents) >= NIXL_REMOTE_AGENT_CACHE_MAXSIZE:
|
|
evicted_agent_name, _ = self._remote_agents.popitem(last=False)
|
|
nixl_agent.remove_remote_agent(evicted_agent_name)
|
|
|
|
self._remote_agents[remote_name] = remote_agent_meta_version
|
|
|
|
nixl_agent.add_remote_agent(remote_nixl_agent_meta)
|
|
|
|
xfer_handle = nixl_agent.initialize_xfer(
|
|
"READ",
|
|
local_xfer_descs,
|
|
remote_xfer_descs,
|
|
remote_name,
|
|
b"UUID",
|
|
)
|
|
|
|
state = nixl_agent.transfer(xfer_handle)
|
|
if state == "ERR":
|
|
raise RuntimeError("NIXL transfer got to Error state.")
|
|
|
|
return NixlFetchRequest(
|
|
tensors=tensors,
|
|
obj_id=obj_id,
|
|
xfer_handle=xfer_handle,
|
|
nixl_agent=nixl_agent,
|
|
remote_name=remote_name,
|
|
remove_tensor_descs=added_tensor_descs,
|
|
transport=self,
|
|
)
|
|
except Exception:
|
|
self._cleanup_transfer(
|
|
obj_id, tensors, xfer_handle, remote_name, added_tensor_descs
|
|
)
|
|
# TODO(swang): There is a circular import error because ray.util
|
|
# currently depends on ray.experimental.internal_kv.
|
|
from ray.exceptions import RayDirectTransportError
|
|
|
|
raise RayDirectTransportError(
|
|
f"The NIXL transfer failed for object id: {obj_id}. The source actor may have died during the transfer. "
|
|
f"The exception thrown from nixl transfer was:\n {traceback.format_exc()}"
|
|
) from None
|
|
|
|
def wait_fetch_complete(
|
|
self, fetch_request: FetchRequest, timeout: float = -1
|
|
) -> List["torch.Tensor"]:
|
|
"""Waits for a previously initiated fetch to complete and returns the tensors.
|
|
|
|
Args:
|
|
fetch_request: The NixlFetchRequest returned by fetch_multiple_tensors.
|
|
timeout: Maximum time in seconds to wait. -1 means wait indefinitely.
|
|
0 means return immediately if not ready.
|
|
|
|
Returns:
|
|
List of tensors that were transferred.
|
|
|
|
Raises:
|
|
RayDirectTransportError: If the transfer failed.
|
|
TimeoutError: If the timeout is exceeded.
|
|
"""
|
|
assert isinstance(fetch_request, NixlFetchRequest)
|
|
obj_id = fetch_request.obj_id
|
|
|
|
if not fetch_request.tensors:
|
|
return fetch_request.tensors
|
|
|
|
try:
|
|
# Check the state of the transfer continuously.
|
|
deadline = None if timeout < 0 else time.monotonic() + timeout
|
|
while True:
|
|
state = self.get_nixl_agent().check_xfer_state(
|
|
fetch_request.xfer_handle
|
|
)
|
|
if state == "ERR":
|
|
raise RuntimeError("NIXL transfer got to Error state.")
|
|
if state == "PROC":
|
|
if deadline is not None and time.monotonic() >= deadline:
|
|
raise TimeoutError(
|
|
f"NIXL transfer timed out after {timeout}s for object id: {obj_id}"
|
|
)
|
|
with self._aborted_transfer_obj_ids_lock:
|
|
if obj_id in self._aborted_transfer_obj_ids:
|
|
self._aborted_transfer_obj_ids.remove(obj_id)
|
|
raise RuntimeError(
|
|
f"NIXL transfer aborted for object id: {obj_id}"
|
|
)
|
|
time.sleep(0.001) # Avoid busy waiting
|
|
elif state == "DONE":
|
|
break
|
|
|
|
return fetch_request.tensors
|
|
except TimeoutError:
|
|
raise
|
|
except Exception:
|
|
from ray.exceptions import RayDirectTransportError
|
|
|
|
raise RayDirectTransportError(
|
|
f"The NIXL transfer failed for object id: {obj_id}. The source actor may have died during the transfer. "
|
|
f"The exception thrown from nixl transfer was:\n {traceback.format_exc()}"
|
|
) from None
|
|
|
|
def _cleanup_transfer(
|
|
self,
|
|
obj_id: str,
|
|
tensors: List["torch.Tensor"],
|
|
xfer_handle: Any,
|
|
remote_name: Optional[str],
|
|
remove_tensor_descs: bool,
|
|
) -> None:
|
|
"""Cleans up resources after a transfer completes or fails."""
|
|
# We could raise errors or NIXL could raise errors like NIXL_ERR_REMOTE_DISCONNECT,
|
|
# so doing best effort cleanup.
|
|
nixl_agent = self._nixl_agent
|
|
if nixl_agent is None:
|
|
return
|
|
# We could raise errors or NIXL could raise errors like NIXL_ERR_REMOTE_DISCONNECT,
|
|
# so doing best effort cleanup.
|
|
with self._aborted_transfer_obj_ids_lock:
|
|
self._aborted_transfer_obj_ids.discard(obj_id)
|
|
if xfer_handle:
|
|
nixl_agent.release_xfer_handle(xfer_handle)
|
|
if NIXL_REMOTE_AGENT_CACHE_MAXSIZE == 0 and remote_name:
|
|
nixl_agent.remove_remote_agent(remote_name)
|
|
if remove_tensor_descs:
|
|
self._remove_tensor_descs(tensors)
|
|
|
|
def recv_multiple_tensors(
|
|
self,
|
|
obj_id: str,
|
|
tensor_transport_metadata: TensorTransportMetadata,
|
|
communicator_metadata: CommunicatorMetadata,
|
|
target_buffers: Optional[List["torch.Tensor"]] = None,
|
|
) -> List["torch.Tensor"]:
|
|
"""Receives multiple tensors synchronously."""
|
|
fetch_request = self.fetch_multiple_tensors(
|
|
obj_id, tensor_transport_metadata, communicator_metadata, target_buffers
|
|
)
|
|
return self.wait_fetch_complete(fetch_request)
|
|
|
|
def send_multiple_tensors(
|
|
self,
|
|
tensors: List["torch.Tensor"],
|
|
tensor_transport_metadata: TensorTransportMetadata,
|
|
communicator_metadata: CommunicatorMetadata,
|
|
):
|
|
raise NotImplementedError(
|
|
"NIXL transport does not support send_multiple_tensors, since it is a one-sided transport."
|
|
)
|
|
|
|
def garbage_collect(
|
|
self,
|
|
obj_id: str,
|
|
tensor_transport_meta: TensorTransportMetadata,
|
|
tensors: List["torch.Tensor"],
|
|
):
|
|
with self._cache_lock:
|
|
assert isinstance(tensor_transport_meta, NixlTransportMetadata)
|
|
if obj_id not in self._managed_meta_nixl:
|
|
return
|
|
self._managed_meta_nixl.pop(obj_id, None)
|
|
self._remove_tensor_descs(tensors)
|
|
|
|
def abort_transport(
|
|
self,
|
|
obj_id: str,
|
|
communicator_metadata: CommunicatorMetadata,
|
|
):
|
|
with self._aborted_transfer_obj_ids_lock:
|
|
self._aborted_transfer_obj_ids.add(obj_id)
|
|
|
|
def _get_num_managed_meta_nixl(self) -> int:
|
|
with self._cache_lock:
|
|
return len(self._managed_meta_nixl)
|
|
|
|
def _get_meta(self, object_id: str) -> Optional[NixlTransportMetadata]:
|
|
"""
|
|
Get the NIXL transport metadata for the given object ID if it exists
|
|
"""
|
|
with self._cache_lock:
|
|
if object_id in self._managed_meta_nixl:
|
|
return self._managed_meta_nixl[object_id]
|
|
return None
|
|
|
|
def _put_meta(self, object_id: str, meta: NixlTransportMetadata):
|
|
"""
|
|
Store the NIXL transport metadata for the given object ID
|
|
"""
|
|
with self._cache_lock:
|
|
self._managed_meta_nixl[object_id] = meta
|
|
|
|
def _remove_tensor_descs(self, tensors: List["torch.Tensor"]):
|
|
"""
|
|
Decrements the reference count for each tensor. If the count reaches 0,
|
|
traditionally-registered memory is deregistered from NIXL, while
|
|
pool-managed blocks (reg_desc is None) are returned to the pool.
|
|
"""
|
|
with self._cache_lock:
|
|
pool_return_tensors: List["torch.Tensor"] = []
|
|
for tensor in tensors:
|
|
key = tensor.untyped_storage().data_ptr()
|
|
if key not in self._tensor_desc_cache:
|
|
continue
|
|
tensor_desc = self._tensor_desc_cache[key]
|
|
tensor_desc.metadata_count -= 1
|
|
if tensor_desc.metadata_count == 0:
|
|
self._tensor_desc_cache.pop(key)
|
|
if tensor_desc.reg_desc is not None:
|
|
# Traditional path: deregister NIXL memory.
|
|
self.get_nixl_agent().deregister_memory(tensor_desc.reg_desc)
|
|
self._nixl_agent_meta_version += 1
|
|
else:
|
|
# Pool path: return block to pool.
|
|
pool_return_tensors.append(tensor)
|
|
if pool_return_tensors and self._memory_pool is not None:
|
|
self._memory_pool.free_tensors(pool_return_tensors)
|
|
|
|
def _add_tensor_descs(self, tensors: List["torch.Tensor"]):
|
|
"""
|
|
If this is the first time the tensor is being registered, we register the
|
|
full underlying pytorch storage object with NIXL. Otherwise, we increment the reference count.
|
|
"""
|
|
with self._cache_lock:
|
|
for tensor in tensors:
|
|
key = tensor.untyped_storage().data_ptr()
|
|
if key in self._tensor_desc_cache:
|
|
self._tensor_desc_cache[key].metadata_count += 1
|
|
continue
|
|
mem_type = "cuda" if tensor.is_cuda else "cpu"
|
|
# the GPU ID of the device the tensor is on.
|
|
# NOTE: we clip this to 0 since the GPU ID is not used for
|
|
# CPU tensors, and get_device returns -1 for CPU tensors.
|
|
# This triggers an error in nixl since it expects an unsigned.
|
|
gpu_id = max(tensor.get_device(), 0)
|
|
# Registering the full underlying pytorch storage object by
|
|
# constructing a memory region with the data pointer, size,
|
|
# GPU ID, and meta info. Doing the equivalent of what nixl
|
|
# does for pytorch tensors internally:
|
|
# https://github.com/ai-dynamo/nixl/blob/dd23ef01bd366aef89fa552f2b042f89a0b45fcb/src/api/python/_api.py#L1034
|
|
try:
|
|
reg_desc = self.get_nixl_agent().register_memory(
|
|
[
|
|
(
|
|
tensor.untyped_storage().data_ptr(),
|
|
tensor.untyped_storage().nbytes(),
|
|
gpu_id,
|
|
"",
|
|
)
|
|
],
|
|
mem_type=mem_type,
|
|
)
|
|
except Exception as e:
|
|
# TODO(xyuzh): Remove the warning after nixl surfaces the error message
|
|
if self._backend == "LIBFABRIC":
|
|
troubleshooting = (
|
|
"See https://github.com/ai-dynamo/nixl/blob/main/src/plugins/libfabric/README.md "
|
|
"for LIBFABRIC troubleshooting. "
|
|
"Set FI_LOG_LEVEL=Debug for libfabric diagnostics."
|
|
)
|
|
else:
|
|
troubleshooting = (
|
|
"See https://docs.ray.io/en/latest/ray-core/direct-transport/direct-transport.html "
|
|
"for NIXL/UCX configuration. "
|
|
"Set UCX_LOG_LEVEL=debug for UCX diagnostics."
|
|
)
|
|
vmm_hint = ""
|
|
alloc_conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "").lower()
|
|
if mem_type == "cuda" and "expandable_segments:true" in alloc_conf:
|
|
vmm_hint = (
|
|
" PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is set; "
|
|
"CUDA VMM memory can't be RDMA-registered — allocate "
|
|
"transferred tensors without expandable_segments."
|
|
)
|
|
raise RuntimeError(
|
|
f"Failed to register {mem_type} memory with NIXL "
|
|
f"(backend={self._backend}, "
|
|
f"size={tensor.untyped_storage().nbytes()} bytes, "
|
|
f"gpu_id={gpu_id}).{vmm_hint} {troubleshooting}"
|
|
) from e
|
|
self._tensor_desc_cache[key] = TensorDesc(reg_desc, 1)
|
|
|
|
def _tensor_memory_registered(self, t: "torch.Tensor") -> bool:
|
|
"""Check if the tensor's memory has been registered with NIXL."""
|
|
entry = self._tensor_desc_cache.get(t.untyped_storage().data_ptr())
|
|
return entry is not None and entry.reg_desc is not None
|
|
|
|
def _add_pool_tensor_descs(self, tensors: List["torch.Tensor"]):
|
|
"""Add pool-managed tensor entries to the unified _tensor_desc_cache.
|
|
|
|
Pool-managed tensors use reg_desc=None since pool memory is registered
|
|
once at pool creation. The metadata_count tracks reference counting
|
|
just like traditional tensors.
|
|
|
|
Note: Entries are keyed by the source tensor's storage ``data_ptr()``.
|
|
If PyTorch frees and reallocates that storage address before GC runs,
|
|
a stale cache entry could map to an unrelated tensor. This is the same
|
|
constraint as the traditional (non-pool) path and is mitigated by the
|
|
fact that pool blocks hold a reference to pool memory, not the source
|
|
storage.
|
|
"""
|
|
with self._cache_lock:
|
|
for tensor in tensors:
|
|
key = tensor.untyped_storage().data_ptr()
|
|
if key in self._tensor_desc_cache:
|
|
self._tensor_desc_cache[key].metadata_count += 1
|
|
else:
|
|
self._tensor_desc_cache[key] = TensorDesc(
|
|
reg_desc=None, metadata_count=1
|
|
)
|
|
|
|
def _allocate_pool_xfer_descs(self, tensors: List["torch.Tensor"]) -> Any:
|
|
"""Allocate pool memory for tensors and return NIXL transfer descriptors.
|
|
|
|
Handles rollback of newly allocated pool blocks if get_xfer_descs
|
|
fails, without disturbing cached blocks from prior calls.
|
|
"""
|
|
pool = self._memory_pool
|
|
# Remember which storages already have a pool block (cache hits)
|
|
# so we don't free them on rollback.
|
|
pre_existing = {
|
|
t.untyped_storage().data_ptr() for t in tensors if pool.has_block(t)
|
|
}
|
|
pool_tensor_views = pool.allocate_for_tensors(tensors)
|
|
try:
|
|
xfer_descs = self._nixl_agent.get_xfer_descs(pool_tensor_views)
|
|
except Exception:
|
|
# Only free newly allocated blocks, not cache hits.
|
|
new_tensors = [
|
|
t for t in tensors if t.untyped_storage().data_ptr() not in pre_existing
|
|
]
|
|
if new_tensors:
|
|
pool.free_tensors(new_tensors)
|
|
raise
|
|
self._add_pool_tensor_descs(tensors)
|
|
return xfer_descs
|