nodedb_vector/collection/
budget.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! RAM budget enforcement and mmap spillover for `VectorCollection`.
4
5use nodedb_mem::ScopedMemory;
6
7use crate::collection::tier::StorageTier;
8use crate::hnsw::HnswIndex;
9use crate::mmap_segment::MmapVectorSegment;
10
11use super::lifecycle::VectorCollection;
12
13impl VectorCollection {
14    /// Set the data directory for mmap segment files.
15    pub fn set_data_dir(&mut self, dir: std::path::PathBuf) {
16        self.data_dir = Some(dir);
17    }
18
19    /// Set the RAM budget for vector data (FP32 in sealed segments).
20    pub fn set_ram_budget(&mut self, bytes: usize) {
21        self.ram_budget_bytes = bytes;
22    }
23
24    /// Estimate current RAM usage for vector data.
25    pub fn ram_usage_bytes(&self) -> usize {
26        let bytes_per_vector = self.dim * std::mem::size_of::<f32>();
27        let growing = self.growing.len() * bytes_per_vector;
28        let building: usize = self
29            .building
30            .iter()
31            .map(|b| b.flat.len() * bytes_per_vector)
32            .sum();
33        let sealed_ram: usize = self
34            .sealed
35            .iter()
36            .filter(|s| s.tier == StorageTier::L0Ram)
37            .map(|s| s.index.len() * bytes_per_vector)
38            .sum();
39        let ivf = self.ivf.as_ref().map_or(0, |ivf| ivf.memory_bytes());
40        growing + building + sealed_ram + ivf
41    }
42
43    /// Whether the RAM budget is exceeded.
44    pub fn is_budget_exceeded(&self) -> bool {
45        self.ram_budget_bytes > 0 && self.ram_usage_bytes() >= self.ram_budget_bytes
46    }
47
48    /// Number of segments that fell back to mmap.
49    pub fn mmap_fallback_count(&self) -> u32 {
50        self.mmap_fallback_count
51    }
52
53    /// Number of currently active mmap segments.
54    pub fn mmap_segment_count(&self) -> u32 {
55        self.mmap_segment_count
56    }
57
58    /// Determine storage tier and optionally create mmap segment for a completed build.
59    ///
60    /// `base_id` is the global vector ID offset for the first vector in `index`.
61    /// Surrogate IDs are looked up from `self.surrogate_map` for rows
62    /// `[base_id, base_id + N)` and written into the segment's surrogate block.
63    pub(crate) fn resolve_tier_for_build(
64        &mut self,
65        segment_id: u32,
66        base_id: u32,
67        index: &HnswIndex,
68        memory: &ScopedMemory,
69    ) -> (StorageTier, Option<MmapVectorSegment>) {
70        if !self.is_budget_exceeded() {
71            return (StorageTier::L0Ram, None);
72        }
73
74        let Some(dir) = &self.data_dir else {
75            return (StorageTier::L0Ram, None);
76        };
77
78        let seg_name = format!("seg-{segment_id}.vseg");
79        let seg_path = dir.join(&seg_name);
80        let count = index.len();
81
82        let refs: Vec<Vec<f32>> = (0..count)
83            .filter_map(|i| index.get_vector(i as u32).map(|v| v.to_vec()))
84            .collect();
85        let ref_slices: Vec<&[f32]> = refs.iter().map(|v| v.as_slice()).collect();
86
87        // Build the parallel surrogate ID array (u64 per row).
88        let surrogate_ids: Vec<u64> = (0..count as u32)
89            .map(|local_id| {
90                let global_id = base_id + local_id;
91                self.surrogate_map
92                    .get(&global_id)
93                    .map(|s| s.as_u32() as u64)
94                    .unwrap_or(0)
95            })
96            .collect();
97
98        match MmapVectorSegment::create_with_surrogates(
99            dir,
100            &seg_name,
101            self.dim,
102            &ref_slices,
103            &surrogate_ids,
104            memory,
105        ) {
106            Ok(mmap) => {
107                self.mmap_fallback_count += 1;
108                self.mmap_segment_count += 1;
109                tracing::info!(
110                    segment_id,
111                    vectors = count,
112                    path = %seg_path.display(),
113                    "vector segment spilled to mmap (L1 NVMe)"
114                );
115                (StorageTier::L1Nvme, Some(mmap))
116            }
117            Err(e) => {
118                tracing::warn!(
119                    segment_id,
120                    error = %e,
121                    "mmap fallback failed, keeping vectors in RAM"
122                );
123                (StorageTier::L0Ram, None)
124            }
125        }
126    }
127}
nodedb_vector/collection/budget.rs

nodedb_vector/collection/
budget.rs
