kvbm_physical/manager/
metadata.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Serialization types for exporting/importing layout metadata with NIXL integration.
5
6use super::handle::LayoutHandle;
7use crate::layout::LayoutDescriptor;
8use anyhow::Result;
9use bincode::{Decode, Encode};
10use serde::{Deserialize, Serialize};
11
12use kvbm_common::LogicalLayoutHandle;
13
14/// Worker identification combining worker_id and NIXL agent name.
15#[derive(Debug, Clone, Encode, Decode, PartialEq, Eq)]
16pub struct WorkerAddress {
17    /// Unique identifier for this worker
18    pub worker_id: u64,
19    /// NIXL agent name on this worker
20    pub nixl_agent_name: String,
21}
22
23impl WorkerAddress {
24    /// Create a new worker address.
25    pub fn new(worker_id: u64, nixl_agent_name: String) -> Self {
26        Self {
27            worker_id,
28            nixl_agent_name,
29        }
30    }
31}
32
33/// Layout descriptor with its assigned handle and logical type for RDMA metadata exchange.
34///
35/// This includes the logical layout type (G1, G2, G3, G4) so that remote instances
36/// know which physical handle corresponds to which tier.
37#[derive(Debug, Clone, Encode, Decode)]
38pub struct LogicalLayoutDescriptor {
39    /// Unique handle for this layout
40    pub handle: LayoutHandle,
41    /// The logical layout type (G1, G2, G3, G4)
42    #[bincode(with_serde)]
43    pub logical_type: LogicalLayoutHandle,
44    /// Serialized layout data (uses Serde, bridged via bincode)
45    #[bincode(with_serde)]
46    pub layout: LayoutDescriptor,
47}
48
49impl LogicalLayoutDescriptor {
50    /// Create a new layout descriptor with handle and logical type.
51    pub fn new(
52        handle: LayoutHandle,
53        logical_type: LogicalLayoutHandle,
54        layout: LayoutDescriptor,
55    ) -> Self {
56        Self {
57            handle,
58            logical_type,
59            layout,
60        }
61    }
62
63    /// Create a layout descriptor with G2 as the default logical type.
64    ///
65    /// This is provided for backwards compatibility with code that doesn't
66    /// track logical types. G2 is used as the default since it's the most
67    /// common tier for RDMA transfers (GPU memory for KV cache).
68    ///
69    /// For proper RDMA transfers between instances, use `new()` with the
70    /// correct logical type from the Worker's registered handles.
71    pub fn new_with_default_type(handle: LayoutHandle, layout: LayoutDescriptor) -> Self {
72        Self {
73            handle,
74            logical_type: LogicalLayoutHandle::G2,
75            layout,
76        }
77    }
78}
79
80/// Type alias for backwards compatibility.
81pub type LocalLayoutDescriptor = LogicalLayoutDescriptor;
82
83/// The set of [`LogicalLayoutDescriptor`] that are RDMA enabled. This object packages the detail
84/// about the layouts and the NIXL RDMA metadata required to reconstruct the layouts and access
85/// the memory via NIXL RDMA.
86#[derive(Debug, Encode, Decode)]
87pub struct RdmaLayoutDescriptors {
88    /// Worker identification
89    pub worker_address: WorkerAddress,
90    /// Exported NIXL metadata from nixl_sys::Agent::get_local_md()
91    pub nixl_metadata: Vec<u8>,
92    /// Serialized layouts (handle + logical type + layout data)
93    pub layouts: Vec<LogicalLayoutDescriptor>,
94}
95
96/// Managed memory metadata package for export/import.
97///
98/// This is the wire format for transmitting layout metadata between workers.
99/// It contains everything needed to reconstruct remote layouts and load their
100/// NIXL registration data.
101#[derive(Clone, Serialize, Deserialize, Encode, Decode)]
102#[serde(transparent)]
103pub struct SerializedLayout(Vec<u8>);
104
105impl SerializedLayout {
106    /// Pack metadata into a serialized form.
107    ///
108    /// # Arguments
109    /// * `worker_address` - Worker identification
110    /// * `nixl_metadata` - NIXL metadata blob from get_local_md()
111    /// * `layouts` - Vector of layouts with handles and logical types to export
112    ///
113    /// # Returns
114    /// Packed metadata ready for transmission
115    pub fn pack(
116        worker_address: WorkerAddress,
117        nixl_metadata: Vec<u8>,
118        layouts: Vec<LogicalLayoutDescriptor>,
119    ) -> Result<Self> {
120        let inner = RdmaLayoutDescriptors {
121            worker_address,
122            nixl_metadata,
123            layouts,
124        };
125        let bytes = bincode::encode_to_vec(&inner, bincode::config::standard())
126            .map_err(|e| anyhow::anyhow!("failed to encode managed memory metadata: {}", e))?;
127        Ok(Self(bytes))
128    }
129
130    /// Unpack metadata from serialized form.
131    ///
132    /// # Returns
133    /// Unpacked metadata structure
134    pub fn unpack(&self) -> Result<RdmaLayoutDescriptors> {
135        let (inner, _) = bincode::decode_from_slice(&self.0, bincode::config::standard())
136            .map_err(|e| anyhow::anyhow!("failed to decode managed memory metadata: {}", e))?;
137        Ok(inner)
138    }
139
140    /// Get the raw bytes.
141    pub fn as_bytes(&self) -> &[u8] {
142        &self.0
143    }
144
145    /// Create from raw bytes.
146    pub fn from_bytes(bytes: Vec<u8>) -> Self {
147        Self(bytes)
148    }
149
150    /// Get the size in bytes.
151    pub fn len(&self) -> usize {
152        self.0.len()
153    }
154
155    /// Check if empty.
156    pub fn is_empty(&self) -> bool {
157        self.0.is_empty()
158    }
159}
160
161impl std::fmt::Debug for SerializedLayout {
162    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
163        f.debug_struct("SerializedLayout")
164            .field("size_bytes", &self.len())
165            .finish()
166    }
167}
168
169#[cfg(all(test, feature = "testing-kvbm"))]
170mod tests {
171    use super::*;
172    use crate::layout::{
173        BlockFormat, FullyContiguousDetails, KvBlockLayout, LayoutConfig, LayoutDescriptor,
174        LayoutTypeDetails, NixlMetadata,
175    };
176    use dynamo_memory::{MemoryRegion, StorageKind, nixl};
177    use kvbm_common::LogicalLayoutHandle;
178
179    fn make_test_serialized_layout() -> LayoutDescriptor {
180        let config = LayoutConfig::builder()
181            .num_blocks(2)
182            .num_layers(2)
183            .outer_dim(2)
184            .page_size(4)
185            .inner_dim(8)
186            .dtype_width_bytes(2)
187            .build()
188            .unwrap();
189
190        LayoutDescriptor {
191            version: 1,
192            layout_config: config,
193            location: StorageKind::System,
194            nixl_metadata: NixlMetadata::new("test".to_string(), nixl::MemType::Dram, 0),
195            memory_descriptors: vec![MemoryRegion {
196                addr: 0x1000,
197                size: 4096,
198            }],
199            layout_type_details: LayoutTypeDetails::FullyContiguous(FullyContiguousDetails {
200                block_format: BlockFormat::Operational,
201                kv_block_layout: KvBlockLayout::OperationalNHD,
202            }),
203        }
204    }
205
206    #[test]
207    fn test_worker_address() {
208        let addr = WorkerAddress::new(42, "test_agent".to_string());
209        assert_eq!(addr.worker_id, 42);
210        assert_eq!(addr.nixl_agent_name, "test_agent");
211    }
212
213    #[test]
214    fn test_serialized_layout_with_handle() {
215        let handle = LayoutHandle::new(1, 2);
216        let layout = make_test_serialized_layout();
217        let with_handle = LogicalLayoutDescriptor::new(handle, LogicalLayoutHandle::G2, layout);
218
219        assert_eq!(with_handle.handle, handle);
220        assert_eq!(with_handle.logical_type, LogicalLayoutHandle::G2);
221    }
222
223    #[test]
224    fn test_metadata_pack_unpack() {
225        let worker_address = WorkerAddress::new(100, "worker_100".to_string());
226        let nixl_metadata = vec![1, 2, 3, 4, 5];
227        let layouts = vec![LogicalLayoutDescriptor::new(
228            LayoutHandle::new(100, 1),
229            LogicalLayoutHandle::G2,
230            make_test_serialized_layout(),
231        )];
232
233        let packed =
234            SerializedLayout::pack(worker_address.clone(), nixl_metadata.clone(), layouts).unwrap();
235
236        assert!(!packed.is_empty());
237
238        let unpacked = packed.unpack().unwrap();
239
240        assert_eq!(unpacked.worker_address, worker_address);
241        assert_eq!(unpacked.nixl_metadata, nixl_metadata);
242        assert_eq!(unpacked.layouts.len(), 1);
243        assert_eq!(unpacked.layouts[0].handle.worker_id(), 100);
244        assert_eq!(unpacked.layouts[0].handle.layout_id(), 1);
245        assert_eq!(unpacked.layouts[0].logical_type, LogicalLayoutHandle::G2);
246    }
247
248    #[test]
249    fn test_metadata_multiple_layouts() {
250        let worker_address = WorkerAddress::new(200, "worker_200".to_string());
251        let nixl_metadata = vec![10, 20, 30];
252        let layouts = vec![
253            LogicalLayoutDescriptor::new(
254                LayoutHandle::new(200, 1),
255                LogicalLayoutHandle::G1,
256                make_test_serialized_layout(),
257            ),
258            LogicalLayoutDescriptor::new(
259                LayoutHandle::new(200, 2),
260                LogicalLayoutHandle::G2,
261                make_test_serialized_layout(),
262            ),
263            LogicalLayoutDescriptor::new(
264                LayoutHandle::new(200, 3),
265                LogicalLayoutHandle::G3,
266                make_test_serialized_layout(),
267            ),
268        ];
269
270        let packed =
271            SerializedLayout::pack(worker_address, nixl_metadata, layouts.clone()).unwrap();
272        let unpacked = packed.unpack().unwrap();
273
274        assert_eq!(unpacked.layouts.len(), 3);
275        let expected_logical_types = [
276            LogicalLayoutHandle::G1,
277            LogicalLayoutHandle::G2,
278            LogicalLayoutHandle::G3,
279        ];
280        for (i, layout) in unpacked.layouts.iter().enumerate() {
281            assert_eq!(layout.handle.worker_id(), 200);
282            assert_eq!(layout.handle.layout_id(), (i + 1) as u16);
283            assert_eq!(layout.logical_type, expected_logical_types[i]);
284        }
285    }
286
287    #[test]
288    fn test_metadata_from_bytes() {
289        let worker_address = WorkerAddress::new(42, "test".to_string());
290        let nixl_metadata = vec![1, 2, 3];
291        let layouts = vec![];
292
293        let packed = SerializedLayout::pack(worker_address, nixl_metadata, layouts).unwrap();
294        let bytes = packed.as_bytes().to_vec();
295
296        let restored = SerializedLayout::from_bytes(bytes);
297        let unpacked = restored.unpack().unwrap();
298
299        assert_eq!(unpacked.worker_address.worker_id, 42);
300    }
301}
kvbm_physical/manager/metadata.rs

kvbm_physical/manager/
metadata.rs
