Struct TransferManager

Source

pub struct TransferManager { /* private fields */ }

Expand description

Public entry point for layout and transfer management.

TransferManager combines layout registration/metadata management with transfer execution capabilities, providing a unified API for:

Registering local layouts and obtaining handles
Exporting/importing layout metadata for remote workers
Executing transfers between layouts using handles
Managing CUDA, NIXL, and other execution resources

Implementations§

Source §

impl TransferManager

Source

pub fn builder() -> TransferConfigBuilder

Create a new TransferManager builder.

The builder configures the worker ID, NIXL agent, CUDA device, and other execution parameters before creating the manager.

§Example

let manager = TransferManager::builder()
    .worker_id(0)  // NIXL agent name defaults to "worker-0"
    .nixl_backend("ucx")  // Optional: defaults to UCX from env
    .cuda_device_id(0)
    .build()?;

// Or with custom agent name:
let manager = TransferManager::builder()
    .worker_id(0)
    .nixl_agent_name("custom-agent")
    .build()?;

Source

pub fn register_layout(&self, layout: PhysicalLayout) -> Result<LayoutHandle>

Register a local physical layout and return a unique handle.

This registers the layout with the embedded memory manager, assigning it a unique handle that can be used for handle-based transfers.

§Arguments

layout - Physical layout to register

§Returns

Unique handle for the registered layout

§Errors

Returns an error if layout IDs are exhausted (u16::MAX reached)

Source

pub fn export_metadata(&self) -> Result<SerializedLayout>

Export layout metadata for transmission to remote workers.

This exports all registered local layouts along with NIXL metadata needed for remote memory registration.

§Returns

Packed metadata ready for transmission to remote workers

Source

pub fn import_metadata( &self, metadata: SerializedLayout, ) -> Result<Vec<LayoutHandle>>

Import remote layout metadata.

This loads NIXL metadata and reconstructs physical layouts from a remote worker’s exported metadata.

§Arguments

metadata - Packed metadata from remote worker

§Returns

Vector of handles for the imported remote layouts

§Errors

Returns an error if the remote worker was already loaded or if metadata loading/reconstruction fails

Source

pub fn build_logical_descriptor( &self, handle: LayoutHandle, logical_type: LogicalLayoutHandle, ) -> Result<LogicalLayoutDescriptor>

Build a logical layout descriptor for a specific handle.

This creates a descriptor that includes the logical layout type (G1, G2, G3, G4) for use in RDMA metadata exchange. The caller must provide the logical type mapping since only the caller (e.g., DirectWorker) knows which handle corresponds to which logical tier.

§Arguments

handle - Handle to the local layout
logical_type - The logical tier (G1, G2, G3, G4) this handle represents

§Returns

A LogicalLayoutDescriptor ready for serialization

§Errors

Returns an error if the handle is not found or serialization fails

Source

pub fn get_nixl_metadata(&self) -> Result<Vec<u8>>

Get the NIXL metadata for this worker.

Returns the raw NIXL metadata bytes needed for remote registration.

Source

pub fn worker_address(&self) -> WorkerAddress

Get the worker address for this manager.

Source

pub fn nixl_agent(&self) -> &NixlAgent

Get a reference to the NIXL agent.

This is useful for building layouts that need to register memory with the same agent that the TransferManager uses.

Source

pub fn get_layout_config(&self, handle: LayoutHandle) -> Result<LayoutConfig>

Get the layout configuration for a registered layout.

Returns a clone of the layout’s configuration, which includes dimensions like num_blocks, num_layers, page_size, etc.

§Arguments

handle - Handle to a registered layout (local or remote)

§Returns

A clone of the layout’s configuration

§Errors

Returns an error if the handle is not found

Source

pub fn execute_transfer( &self, src_handle: LayoutHandle, src_blocks: &[BlockId], dst_handle: LayoutHandle, dst_blocks: &[BlockId], options: TransferOptions, ) -> Result<TransferCompleteNotification>

Transfer complete blocks between layouts using handles.

This function copies entire blocks (all layers and outer dimensions) between the source and destination layouts identified by their handles. The transfer strategy (memcpy, CUDA, NIXL) is automatically selected based on storage locations.

The lock on the registry is held only briefly during layout lookup, then released before executing the actual transfer.