spark_storage/
config.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// Configuration shape for `--high-speed-swap`. All fields are required
4// (PCND); validation runs at startup and a `HighSpeedSwap` orchestrator
5// cannot be constructed with a partial / inconsistent config. The shape
6// matches the locked CLI flag set in the plan.
7
8use anyhow::{Result, bail};
9use serde::Deserialize;
10use std::path::PathBuf;
11
12#[derive(Clone, Debug, Deserialize)]
13pub struct HighSpeedSwapConfig {
14    /// `--high-speed-swap-dir`: directory where per-layer KV files live.
15    pub dir: PathBuf,
16    /// `--high-speed-swap-bytes`: total disk budget; the layout fails fast
17    /// if the budget can't fit `num_layers × bytes_per_layer`.
18    pub bytes: u64,
19    /// `--high-speed-swap-resident-blocks`: HBM scratch slot count. The
20    /// scratch pool allocates exactly this many slots × per-slot bytes.
21    pub resident_blocks: u32,
22    /// `--high-speed-swap-rank`: predictor low-rank dimension.
23    pub rank: u32,
24    /// `--high-speed-swap-qd`: io_uring submission queue depth. Phase-3
25    /// shows QD=8 reaches 3.4 GB/s on this hardware (random 64 KiB).
26    pub qd: u32,
27    /// `--high-speed-swap-graph`: capture the per-layer body in a CUDA
28    /// graph and replay (Phase 4).
29    pub graph: bool,
30    /// Predictor seed; when omitted in CLI we'd error rather than default.
31    pub projection_seed: u64,
32}
33
34impl HighSpeedSwapConfig {
35    /// Validate cross-field invariants. Returns `Ok` only if the config is
36    /// internally consistent and the directory is plausibly usable.
37    pub fn validate(&self) -> Result<()> {
38        if self.bytes == 0 {
39            bail!("--high-speed-swap-bytes must be > 0");
40        }
41        if self.resident_blocks == 0 {
42            bail!("--high-speed-swap-resident-blocks must be > 0");
43        }
44        if self.rank == 0 || self.rank > 128 {
45            bail!(
46                "--high-speed-swap-rank must be in 1..=128, got {}",
47                self.rank
48            );
49        }
50        if self.qd == 0 || self.qd > 64 {
51            bail!("--high-speed-swap-qd must be in 1..=64, got {}", self.qd);
52        }
53        // The scratch pool must hold *at least one full tile*; tile_capacity
54        // is taken to equal resident_blocks (single-tile fast path). Any
55        // smaller would degrade to streaming-only-per-block.
56        Ok(())
57    }
58
59    /// Validate, then ensure the directory exists or can be created.
60    pub fn validate_and_prepare(&self) -> Result<()> {
61        self.validate()?;
62        std::fs::create_dir_all(&self.dir)
63            .map_err(|e| anyhow::anyhow!("create {}: {e}", self.dir.display()))?;
64        // Cross-mount check vs --swap-space-gb is performed at the
65        // spark-server layer where we know both paths; out of scope here.
66        Ok(())
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    fn cfg() -> HighSpeedSwapConfig {
75        HighSpeedSwapConfig {
76            dir: PathBuf::from("/tmp/atlas-hss-cfg"),
77            bytes: 64 << 30,
78            resident_blocks: 8192,
79            rank: 32,
80            qd: 8,
81            graph: true,
82            projection_seed: 0xCAFE_F00D,
83        }
84    }
85
86    #[test]
87    fn happy_path() {
88        cfg().validate().unwrap();
89    }
90
91    #[test]
92    fn rejects_zero_bytes() {
93        let mut c = cfg();
94        c.bytes = 0;
95        assert!(c.validate().is_err());
96    }
97
98    #[test]
99    fn rejects_zero_resident_blocks() {
100        let mut c = cfg();
101        c.resident_blocks = 0;
102        assert!(c.validate().is_err());
103    }
104
105    #[test]
106    fn rejects_out_of_range_rank() {
107        let mut c = cfg();
108        c.rank = 0;
109        assert!(c.validate().is_err());
110        c.rank = 129;
111        assert!(c.validate().is_err());
112    }
113
114    #[test]
115    fn rejects_out_of_range_qd() {
116        let mut c = cfg();
117        c.qd = 0;
118        assert!(c.validate().is_err());
119        c.qd = 65;
120        assert!(c.validate().is_err());
121    }
122}