register_nixl_memory_pool#

ray.experimental.register_nixl_memory_pool(size: int, device: torch.device) → None[source]#

Pre-allocates a memory pool and registers it with NIXL.

This enables pool-based memory management for NIXL transfers, which can improve performance by avoiding repeated memory registration/deregistration. The pool is registered once with NIXL and individual tensors are copied into it on ray.put.

Only the tensors passed to ray.put are copied (by their own byte size), not their full underlying storage. Contiguous tensors from a single ray.put are packed into as few pool blocks as the free list allows. Pool blocks are freed when the ObjectRef goes out of scope.

Each tensor is placed on a multiple of its own element size, so tensors sharing a dtype pack with no padding between them.

If the pool has insufficient space for an allocation, NixlOutOfMemoryError is raised.

Parameters:

  • size (int) – Size of the memory pool in bytes.

  • device (torch.device) – Device to allocate the pool on (e.g., torch.device("cpu") or torch.device("cuda")).

Example

import torch
import ray
from ray.experimental import register_nixl_memory_pool

@ray.remote(num_gpus=1, enable_tensor_transport=True)
class Trainer:
    def __init__(self):
        # Pre-allocate a 1GB GPU memory pool for NIXL transfers
        register_nixl_memory_pool(1024 * 1024 * 1024, torch.device("cuda"))

    def get_weight_ref(self):
        weight = torch.randn(1000, 1000, device="cuda")
        return ray.put(weight, _tensor_transport="nixl")

PublicAPI (alpha): This API is in alpha and may change before becoming stable.