kvbm_engine/worker/
mod.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4mod coordinated;
5#[doc = include_str!("../../docs/worker-group.md")]
6pub mod group;
7mod physical;
8mod protocol;
9pub mod velo;
10
11pub use coordinated::CoordinatedWorker;
12pub use physical::{PhysicalWorker, PhysicalWorkerBuilder};
13
14/// Compatibility alias for [`PhysicalWorker`].
15pub use physical::PhysicalWorker as DirectWorker;
16
17use anyhow::Result;
18use std::{pin::Pin, sync::Arc};
19
20use crate::object::ObjectBlockOps;
21pub use crate::{BlockId, InstanceId, SequenceHash};
22pub use kvbm_common::LogicalLayoutHandle;
23pub use kvbm_physical::{
24    manager::{LayoutHandle, SerializedLayout},
25    transfer::TransferCompleteNotification,
26};
27
28pub use velo::{VeloWorkerClient, VeloWorkerService, VeloWorkerServiceBuilder};
29
30/// Boxed future for serialized layout responses - allows both typed_unary and raw unary results
31pub type SerializedResponseAwaiter = Pin<Box<dyn Future<Output = Result<SerializedLayout>> + Send>>;
32/// Boxed future for import metadata responses
33pub type ImportMetadataResponseAwaiter =
34    Pin<Box<dyn Future<Output = Result<Vec<LayoutHandle>>> + Send>>;
35
36pub use protocol::*;
37
38pub trait WorkerTransfers: Send + Sync {
39    /// Execute a local transfer between two logical layouts.
40    ///
41    /// # Arguments
42    /// * `src` - The source layout handle
43    /// * `dst` - The destination layout handle
44    /// * `src_block_ids` - The source block IDs
45    /// * `dst_block_ids` - The destination block IDs
46    /// * `options` - Transfer options (layer range, bounce buffers, etc.)
47    ///
48    /// # Returns
49    /// A future that completes when the transfer is complete
50    fn execute_local_transfer(
51        &self,
52        src: LogicalLayoutHandle,
53        dst: LogicalLayoutHandle,
54        src_block_ids: Arc<[BlockId]>,
55        dst_block_ids: Arc<[BlockId]>,
56        options: kvbm_physical::transfer::TransferOptions,
57    ) -> Result<TransferCompleteNotification>;
58
59    /// Execute a remote transfer from a remote layout to a local logical layout.
60    ///
61    /// This represents a NIXL transfer.
62    ///
63    /// # Arguments
64    /// * `src` - Remote sources can take several forms, see [`RemoteDescriptor`]
65    /// * `dst` - The destination layout handle
66    /// * `dst_block_ids` - The destination block IDs
67    /// * `options` - Transfer options (layer range, bounce buffers, etc.)
68    ///
69    /// # Returns
70    /// A future that completes when the transfer is complete
71    fn execute_remote_onboard(
72        &self,
73        src: RemoteDescriptor,
74        dst: LogicalLayoutHandle,
75        dst_block_ids: Arc<[BlockId]>,
76        options: kvbm_physical::transfer::TransferOptions,
77    ) -> Result<TransferCompleteNotification>;
78
79    /// Execute a remote offload from a local logical layout to a remote descriptor.
80    ///
81    /// This represents a NIXL offload.
82    ///
83    /// # Arguments
84    /// * `src` - The source layout handle
85    /// * `dst` - The destination remote descriptor
86    /// * `src_block_ids` - The source block IDs
87    /// * `options` - Transfer options (layer range, bounce buffers, etc.)
88    ///
89    /// # Returns
90    /// A future that completes when the offload is complete
91    fn execute_remote_offload(
92        &self,
93        src: LogicalLayoutHandle,
94        src_block_ids: Arc<[BlockId]>,
95        dst: RemoteDescriptor,
96        options: kvbm_physical::transfer::TransferOptions,
97    ) -> Result<TransferCompleteNotification>;
98
99    /// Connect to a remote instance by importing its metadata and storing handle mappings.
100    ///
101    /// This method stores the handle mappings internally for later use by
102    /// `execute_remote_onboard_for_instance`. The metadata is also imported into
103    /// the underlying transfer manager so NIXL knows about the remote.
104    ///
105    /// # Arguments
106    /// * `instance_id` - The unique identifier of the remote instance
107    /// * `metadata` - Serialized layout metadata from the remote instance.
108    ///   For DirectWorker, expects exactly 1 element.
109    ///   For ReplicatedWorker, expects one element per worker (in rank order).
110    ///
111    /// # Returns
112    /// A response that completes when the metadata has been imported and mappings stored.
113    fn connect_remote(
114        &self,
115        instance_id: InstanceId,
116        metadata: Vec<SerializedLayout>,
117    ) -> Result<ConnectRemoteResponse>;
118
119    /// Check if remote metadata has been imported for an instance.
120    ///
121    /// Returns true if `connect_remote` has been successfully called for this instance.
122    fn has_remote_metadata(&self, instance_id: InstanceId) -> bool;
123
124    /// Execute a remote onboard transfer using stored handle mapping.
125    ///
126    /// This method looks up the remote handle from the stored mapping
127    /// (established via `connect_remote`) and executes the transfer.
128    ///
129    /// # Arguments
130    /// * `instance_id` - The remote instance to pull from
131    /// * `remote_logical_type` - The logical layout type on the remote (e.g., G2)
132    /// * `src_block_ids` - Block IDs on the remote to pull
133    /// * `dst` - Local destination logical layout
134    /// * `dst_block_ids` - Local destination block IDs
135    /// * `options` - Transfer options
136    ///
137    /// # Errors
138    /// Returns error if remote metadata hasn't been imported for this instance.
139    fn execute_remote_onboard_for_instance(
140        &self,
141        instance_id: InstanceId,
142        remote_logical_type: LogicalLayoutHandle,
143        src_block_ids: Vec<BlockId>,
144        dst: LogicalLayoutHandle,
145        dst_block_ids: Arc<[BlockId]>,
146        options: kvbm_physical::transfer::TransferOptions,
147    ) -> Result<TransferCompleteNotification>;
148}
149
150pub trait Worker: WorkerTransfers + ObjectBlockOps + Send + Sync {
151    /// Get the G1 layout handle for this worker (if configured).
152    ///
153    /// Returns None if no G1 layout has been registered with this worker.
154    fn g1_handle(&self) -> Option<LayoutHandle>;
155
156    /// Get the G2 layout handle for this worker (if configured).
157    ///
158    /// Returns None if no G2 layout has been registered with this worker.
159    fn g2_handle(&self) -> Option<LayoutHandle>;
160
161    /// Get the G3 layout handle for this worker (if configured).
162    ///
163    /// Returns None if no G3 layout has been registered with this worker.
164    fn g3_handle(&self) -> Option<LayoutHandle>;
165
166    /// Export the local metadata for this worker.
167    ///
168    /// # Returns
169    /// A [`kvbm_physical::manager::SerializedLayout`] containing the local metadata
170    fn export_metadata(&self) -> Result<SerializedLayoutResponse>;
171
172    /// Import the remote metadata for this worker.
173    ///
174    /// # Arguments
175    /// * `metadata` - A [`kvbm_physical::manager::SerializedLayout`] containing the remote metadata
176    ///
177    /// # Returns
178    /// A vector of [`kvbm_physical::manager::LayoutHandle`] for the imported remote layouts
179    fn import_metadata(&self, metadata: SerializedLayout) -> Result<ImportMetadataResponse>;
180}
kvbm_engine/worker/mod.rs

kvbm_engine/worker/
mod.rs
