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}