nodedb_vector/collection/
budget.rs

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

nodedb_vector/collection/
budget.rs
