metrale_storage/
expert_pack_fs.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The expert store's layer files. `ExpertFileWriter` creates them
4//! and places records; `ExpertFileReader` reads records back with positional
5//! reads.
6//!
7//! Owner: metrale-storage experts.
8//! Invariants:
9//! - A record is written and read at `layout.file_offset(key)` in the file of
10//!   `key.layer`.
11
12use super::*;
13use std::fs::{File, OpenOptions};
14// 2026-09-25: Positional reads and writes on unix and Windows alike.
15use crate::tier::pio;
16use std::path::{Path, PathBuf};
17
18/// 2026-09-25: Creates one file per MoE layer, sized to `bytes_per_layer`,
19/// writes records at their offsets, and writes `manifest.json` in `finish`.
20/// The files are opened without O_DIRECT.
21pub struct ExpertFileWriter {
22    dir: PathBuf,
23    index: ExpertIndex,
24    spec: ExpertRecordSpec,
25    layout: ExpertLayout,
26    files: Vec<File>,
27}
28
29impl ExpertFileWriter {
30    pub fn create(dir: &Path, index: ExpertIndex) -> Result<Self> {
31        std::fs::create_dir_all(dir).with_context(|| format!("mkdir {}", dir.display()))?;
32        let spec = index.spec();
33        let layout = index.layout();
34        let mut files = Vec::with_capacity(index.num_moe_layers as usize);
35        for l in 0..index.num_moe_layers {
36            let p = dir.join(index.file_name(l));
37            let f = OpenOptions::new()
38                .read(true)
39                .write(true)
40                .create(true)
41                .truncate(true)
42                .open(&p)
43                .with_context(|| format!("create {}", p.display()))?;
44            f.set_len(layout.bytes_per_layer())
45                .with_context(|| format!("set_len {}", p.display()))?;
46            files.push(f);
47        }
48        Ok(Self {
49            dir: dir.to_path_buf(),
50            index,
51            spec,
52            layout,
53            files,
54        })
55    }
56
57    pub fn spec(&self) -> &ExpertRecordSpec {
58        &self.spec
59    }
60
61    /// 2026-09-25: Pack and write one record; fails when the layer or expert is
62    /// out of range.
63    pub fn write_record(
64        &self,
65        key: ExpertKey,
66        header: &ExpertRecordHeader,
67        projs: &[ProjData; 3],
68    ) -> Result<()> {
69        if key.layer >= self.index.num_moe_layers {
70            bail!("layer {} out of range", key.layer);
71        }
72        if key.expert >= self.index.num_experts {
73            bail!("expert {} out of range", key.expert);
74        }
75        let rec = pack_record(&self.spec, self.layout.record_stride, header, projs)?;
76        let off = self.layout.file_offset(key);
77        pio::write_all_at(&self.files[key.layer as usize], &rec, off)
78            .with_context(|| format!("write record {:?} at {off}", key))?;
79        Ok(())
80    }
81
82    /// 2026-09-25: fsync every layer file, then write `manifest.json`.
83    pub fn finish(self) -> Result<()> {
84        for f in &self.files {
85            f.sync_all().context("fsync layer file")?;
86        }
87        let p = self.dir.join(ExpertIndex::MANIFEST_NAME);
88        let json = serde_json::to_string_pretty(&self.index)?;
89        std::fs::write(&p, json).with_context(|| format!("write {}", p.display()))?;
90        Ok(())
91    }
92}
93
94/// 2026-09-25: Reads records back with positional reads (`tier::pio`).
95/// `expert_tier.rs` and `expert_peer.rs` open stores through it.
96pub struct ExpertFileReader {
97    index: ExpertIndex,
98    spec: ExpertRecordSpec,
99    layout: ExpertLayout,
100    files: Vec<File>,
101}
102
103impl ExpertFileReader {
104    pub fn open(dir: &Path) -> Result<Self> {
105        let mp = dir.join(ExpertIndex::MANIFEST_NAME);
106        let json =
107            std::fs::read_to_string(&mp).with_context(|| format!("read {}", mp.display()))?;
108        let index: ExpertIndex =
109            serde_json::from_str(&json).with_context(|| format!("parse {}", mp.display()))?;
110        if index.version != ExpertRecordHeader::VERSION {
111            bail!(
112                "manifest version {} != supported {}",
113                index.version,
114                ExpertRecordHeader::VERSION
115            );
116        }
117        let spec = index.spec();
118        let layout = index.layout();
119        let mut files = Vec::with_capacity(index.num_moe_layers as usize);
120        for l in 0..index.num_moe_layers {
121            let p = dir.join(index.file_name(l));
122            files.push(File::open(&p).with_context(|| format!("open {}", p.display()))?);
123        }
124        Ok(Self {
125            index,
126            spec,
127            layout,
128            files,
129        })
130    }
131
132    pub fn index(&self) -> &ExpertIndex {
133        &self.index
134    }
135    pub fn spec(&self) -> &ExpertRecordSpec {
136        &self.spec
137    }
138
139    /// 2026-09-25: Read one record's `record_stride` bytes into a new buffer.
140    /// An out-of-range layer is an error rather than an index panic.
141    pub fn read_record_raw(&self, key: ExpertKey) -> Result<Vec<u8>> {
142        if key.layer as usize >= self.files.len() {
143            bail!(
144                "ExpertFileReader: layer {} out of range ({} layer files)",
145                key.layer,
146                self.files.len()
147            );
148        }
149        let mut buf = vec![0u8; self.layout.record_stride as usize];
150        let off = self.layout.file_offset(key);
151        pio::read_exact_at(&self.files[key.layer as usize], &mut buf, off)
152            .with_context(|| format!("read record {:?} at {off}", key))?;
153        Ok(buf)
154    }
155}
156
157#[cfg(test)]
158mod fs_tests {
159    use super::*;
160
161    fn tmpdir(tag: &str) -> PathBuf {
162        let p = std::env::temp_dir().join(format!(
163            "metrale-xpr-{}-{}-{}",
164            tag,
165            std::process::id(),
166            std::time::SystemTime::now()
167                .duration_since(std::time::UNIX_EPOCH)
168                .map(|d| d.as_nanos())
169                .unwrap_or(0)
170        ));
171        std::fs::create_dir_all(&p).unwrap();
172        p
173    }
174
175    fn synth_index() -> ExpertIndex {
176        ExpertIndex::new(64, 128, 16, 256, 4096, vec![0, 1], 3)
177    }
178
179    fn synth_projs(spec: &ExpertRecordSpec, seed: u8) -> [Vec<(Vec<u8>, Vec<u8>)>; 1] {
180        let mut out = Vec::new();
181        for p in Proj::ALL {
182            let pb = spec.proj_bytes(p);
183            let packed: Vec<u8> = (0..pb.packed_bytes)
184                .map(|i| (i as u8).wrapping_add(seed).wrapping_add(p as u8))
185                .collect();
186            let scale: Vec<u8> = (0..pb.scale_bytes)
187                .map(|i| (i as u8).wrapping_mul(3).wrapping_add(seed))
188                .collect();
189            out.push((packed, scale));
190        }
191        [out]
192    }
193
194    #[test]
195    fn write_then_read_round_trips_bit_identical() {
196        let dir = tmpdir("rt");
197        let index = synth_index();
198        let spec = index.spec();
199
200        let mut expected = std::collections::HashMap::new();
201        {
202            let w = ExpertFileWriter::create(&dir, index.clone()).unwrap();
203            for layer in 0..index.num_moe_layers {
204                for expert in 0..index.num_experts {
205                    let seed = (layer as u8) << 4 | expert as u8;
206                    let raw = synth_projs(&spec, seed);
207                    let projs = [
208                        ProjData {
209                            packed: &raw[0][0].0,
210                            scale: &raw[0][0].1,
211                        },
212                        ProjData {
213                            packed: &raw[0][1].0,
214                            scale: &raw[0][1].1,
215                        },
216                        ProjData {
217                            packed: &raw[0][2].0,
218                            scale: &raw[0][2].1,
219                        },
220                    ];
221                    let header = ExpertRecordHeader {
222                        layer,
223                        expert,
224                        inter: index.inter as u32,
225                        hidden: index.hidden as u32,
226                        group_size: index.group_size as u32,
227                        scale2: [seed as f32, seed as f32 + 0.5, seed as f32 + 1.0],
228                        input_scale: [Some(1.0), None, Some(2.0)],
229                    };
230                    w.write_record(ExpertKey::new(layer, expert), &header, &projs)
231                        .unwrap();
232                    expected.insert((layer, expert), (raw, header));
233                }
234            }
235            w.finish().unwrap();
236        }
237
238        let r = ExpertFileReader::open(&dir).unwrap();
239        assert_eq!(r.index(), &index);
240        for layer in 0..index.num_moe_layers {
241            for expert in 0..index.num_experts {
242                let key = ExpertKey::new(layer, expert);
243                let buf = r.read_record_raw(key).unwrap();
244                let (hdr, views) = unpack_record(r.spec(), &buf).unwrap();
245                let (raw, exp_hdr) = &expected[&(layer, expert)];
246                assert_eq!(&hdr, exp_hdr, "header {:?}", key);
247                for p in Proj::ALL {
248                    assert_eq!(
249                        views[p as usize].packed,
250                        &raw[0][p as usize].0[..],
251                        "packed {:?} {:?}",
252                        key,
253                        p
254                    );
255                    assert_eq!(
256                        views[p as usize].scale,
257                        &raw[0][p as usize].1[..],
258                        "scale {:?} {:?}",
259                        key,
260                        p
261                    );
262                }
263            }
264        }
265        std::fs::remove_dir_all(&dir).ok();
266    }
267
268    #[test]
269    fn wrong_projection_length_errors() {
270        let index = synth_index();
271        let spec = index.spec();
272        let header = ExpertRecordHeader {
273            layer: 0,
274            expert: 0,
275            inter: index.inter as u32,
276            hidden: index.hidden as u32,
277            group_size: index.group_size as u32,
278            scale2: [1.0; 3],
279            input_scale: [Some(1.0); 3],
280        };
281        let bad = vec![0u8; 8];
282        let ok_scale = vec![0u8; spec.proj_bytes(Proj::Gate).scale_bytes as usize];
283        let projs = [
284            ProjData {
285                packed: &bad,
286                scale: &ok_scale,
287            },
288            ProjData {
289                packed: &bad,
290                scale: &ok_scale,
291            },
292            ProjData {
293                packed: &bad,
294                scale: &ok_scale,
295            },
296        ];
297        let err = pack_record(&spec, index.record_stride, &header, &projs);
298        assert!(err.is_err(), "short packed buffer must error");
299    }
300
301    #[test]
302    fn read_record_raw_rejects_out_of_range_layer() {
303        let dir = tmpdir("oob");
304        let index = synth_index();
305        ExpertFileWriter::create(&dir, index)
306            .unwrap()
307            .finish()
308            .unwrap();
309        let r = ExpertFileReader::open(&dir).unwrap();
310        assert!(r.read_record_raw(ExpertKey::new(0, 0)).is_ok());
311        assert!(r.read_record_raw(ExpertKey::new(99, 0)).is_err());
312        std::fs::remove_dir_all(&dir).ok();
313    }
314
315    #[test]
316    fn manifest_geometry_round_trips_through_json() {
317        let index = synth_index();
318        let json = serde_json::to_string(&index).unwrap();
319        let back: ExpertIndex = serde_json::from_str(&json).unwrap();
320        assert_eq!(index, back);
321        assert_eq!(index.layout().record_stride, back.layout().record_stride);
322        assert_eq!(index.spec().raw_bytes(), back.spec().raw_bytes());
323    }
324}
metrale_storage/expert_pack_fs.rs

metrale_storage/
expert_pack_fs.rs
