Skip to main content

sim_lib_compute_auto/
store.rs

1//! Table-backed measured profile persistence.
2
3use sim_kernel::{Cx, Error, Expr, NumberLiteral, Result, Symbol, Value};
4use sim_lib_compute_model::ModeledComputeProfile;
5
6use crate::{
7    ComputeDeviceIdentity, ComputeProfileLimits, ComputeProfileProvenance, ComputeProfileSamples,
8    ComputeThermalPowerContext, MeasuredComputeProfile,
9};
10
11const PROFILE_SCHEMA: &str = "sim.compute.auto.profile.v1";
12const DEFAULT_MAX_KEY_BYTES: usize = 96;
13const DEFAULT_MAX_PROFILE_BYTES: usize = 8192;
14
15/// Policy bounding profile persistence.
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct ProfileStorePolicy {
18    /// Maximum UTF-8 bytes accepted in a profile key.
19    pub max_key_bytes: usize,
20    /// Maximum estimated bytes accepted for one profile record.
21    pub max_profile_bytes: usize,
22}
23
24impl Default for ProfileStorePolicy {
25    fn default() -> Self {
26        Self {
27            max_key_bytes: DEFAULT_MAX_KEY_BYTES,
28            max_profile_bytes: DEFAULT_MAX_PROFILE_BYTES,
29        }
30    }
31}
32
33/// Table-backed measured profile store.
34pub struct ProfileStore {
35    table: Value,
36    policy: ProfileStorePolicy,
37}
38
39impl ProfileStore {
40    /// Builds a store from a caller-supplied Table or Dir value.
41    pub fn new(table: Value, policy: ProfileStorePolicy) -> Result<Self> {
42        if table.object().as_table_impl().is_none() {
43            return Err(Error::TypeMismatch {
44                expected: "Table or Dir",
45                found: "non-table",
46            });
47        }
48        Ok(Self { table, policy })
49    }
50
51    /// Stores a profile with one bounded, atomic table update.
52    pub fn save(&self, cx: &mut Cx, key: Symbol, profile: &MeasuredComputeProfile) -> Result<()> {
53        self.validate_key(&key)?;
54        if profile.estimated_profile_bytes() > self.policy.max_profile_bytes {
55            return Err(Error::Eval(
56                "compute profile exceeds byte policy".to_owned(),
57            ));
58        }
59        let value = profile_to_value(cx, profile)?;
60        self.table
61            .object()
62            .as_table_impl()
63            .expect("profile store table")
64            .set(cx, key, value)
65    }
66
67    /// Loads a checked profile from a bounded key.
68    pub fn load(&self, cx: &mut Cx, key: Symbol) -> Result<Option<MeasuredComputeProfile>> {
69        self.validate_key(&key)?;
70        let table = self
71            .table
72            .object()
73            .as_table_impl()
74            .expect("profile store table");
75        if !table.has(cx, key.clone())? {
76            return Ok(None);
77        }
78        let value = table.get(cx, key)?;
79        let profile = profile_from_value(cx, &value)?;
80        if profile.estimated_profile_bytes() > self.policy.max_profile_bytes {
81            return Err(Error::Eval(
82                "stored compute profile exceeds byte policy".to_owned(),
83            ));
84        }
85        Ok(Some(profile))
86    }
87
88    /// Lists bounded profile keys.
89    pub fn keys(&self, cx: &mut Cx) -> Result<Vec<Symbol>> {
90        let keys = self
91            .table
92            .object()
93            .as_table_impl()
94            .expect("profile store table")
95            .keys(cx)?;
96        for key in &keys {
97            self.validate_key(key)?;
98        }
99        Ok(keys)
100    }
101
102    fn validate_key(&self, key: &Symbol) -> Result<()> {
103        if key.to_string().len() > self.policy.max_key_bytes {
104            return Err(Error::Eval(
105                "compute profile key exceeds byte policy".to_owned(),
106            ));
107        }
108        Ok(())
109    }
110}
111
112fn profile_to_value(cx: &mut Cx, profile: &MeasuredComputeProfile) -> Result<Value> {
113    let entries = vec![
114        (sym("schema"), string(cx, PROFILE_SCHEMA)?),
115        (sym("adapter"), string(cx, &profile.identity.adapter)?),
116        (sym("driver"), string(cx, &profile.identity.driver)?),
117        (sym("backend"), string(cx, &profile.identity.backend)?),
118        (
119            sym("max-resident-bytes"),
120            number(cx, profile.limits.max_resident_bytes)?,
121        ),
122        (
123            sym("max-storage-binding-bytes"),
124            number(cx, profile.limits.max_storage_binding_bytes)?,
125        ),
126        (
127            sym("max-queue-depth"),
128            number(cx, profile.limits.max_queue_depth as u64)?,
129        ),
130        (
131            sym("max-queue-bytes"),
132            number(cx, profile.limits.max_queue_bytes)?,
133        ),
134        (
135            sym("submission-deadline-ticks"),
136            number(cx, profile.limits.submission_deadline_ticks)?,
137        ),
138        (
139            sym("upload-bytes-per-tick"),
140            number_list(cx, &profile.samples.upload_bytes_per_tick)?,
141        ),
142        (
143            sym("download-bytes-per-tick"),
144            number_list(cx, &profile.samples.download_bytes_per_tick)?,
145        ),
146        (
147            sym("launch-ticks"),
148            number_list(cx, &profile.samples.launch_ticks)?,
149        ),
150        (
151            sym("element-elements-per-tick"),
152            number_list(cx, &profile.samples.element_elements_per_tick)?,
153        ),
154        (
155            sym("reduction-elements-per-tick"),
156            number_list(cx, &profile.samples.reduction_elements_per_tick)?,
157        ),
158        (
159            sym("matmul-ops-per-tick"),
160            number_list(cx, &profile.samples.matmul_ops_per_tick)?,
161        ),
162        (sym("thermal"), string(cx, &profile.context.thermal)?),
163        (sym("power"), string(cx, &profile.context.power)?),
164        (sym("tile-bytes"), number(cx, profile.tile_bytes)?),
165        (
166            sym("allocation-bytes"),
167            number_list(cx, &profile.allocation_bytes)?,
168        ),
169        (sym("producer"), string(cx, &profile.provenance.producer)?),
170        (
171            sym("measured-at-tick"),
172            number(cx, profile.provenance.measured_at_tick)?,
173        ),
174        (
175            sym("stale-after-ticks"),
176            number(cx, profile.provenance.stale_after_ticks)?,
177        ),
178        (
179            sym("modeled-provider"),
180            string(cx, &profile.modeled.provider)?,
181        ),
182    ];
183    cx.factory().table(entries)
184}
185
186fn profile_from_value(cx: &mut Cx, value: &Value) -> Result<MeasuredComputeProfile> {
187    let entries = value
188        .object()
189        .as_table_impl()
190        .ok_or(Error::TypeMismatch {
191            expected: "profile table",
192            found: "non-table",
193        })?
194        .entries(cx)?;
195    let get = |name: &str| -> Result<Value> {
196        entries
197            .iter()
198            .find(|(key, _)| *key == sym(name))
199            .map(|(_, value)| value.clone())
200            .ok_or_else(|| Error::Eval(format!("compute profile missing {name}")))
201    };
202    if string_value(cx, &get("schema")?)? != PROFILE_SCHEMA {
203        return Err(Error::Eval("unsupported compute profile schema".to_owned()));
204    }
205    let max_queue_depth = number_value(cx, &get("max-queue-depth")?)?;
206    Ok(MeasuredComputeProfile {
207        identity: ComputeDeviceIdentity {
208            adapter: string_value(cx, &get("adapter")?)?,
209            driver: string_value(cx, &get("driver")?)?,
210            backend: string_value(cx, &get("backend")?)?,
211        },
212        limits: ComputeProfileLimits {
213            max_resident_bytes: number_value(cx, &get("max-resident-bytes")?)?,
214            max_storage_binding_bytes: number_value(cx, &get("max-storage-binding-bytes")?)?,
215            max_queue_depth: usize::try_from(max_queue_depth)
216                .map_err(|_| Error::Eval("max-queue-depth exceeds usize".to_owned()))?,
217            max_queue_bytes: number_value(cx, &get("max-queue-bytes")?)?,
218            submission_deadline_ticks: number_value(cx, &get("submission-deadline-ticks")?)?,
219        },
220        samples: ComputeProfileSamples {
221            upload_bytes_per_tick: number_vec(cx, &get("upload-bytes-per-tick")?)?,
222            download_bytes_per_tick: number_vec(cx, &get("download-bytes-per-tick")?)?,
223            launch_ticks: number_vec(cx, &get("launch-ticks")?)?,
224            element_elements_per_tick: number_vec(cx, &get("element-elements-per-tick")?)?,
225            reduction_elements_per_tick: number_vec(cx, &get("reduction-elements-per-tick")?)?,
226            matmul_ops_per_tick: number_vec(cx, &get("matmul-ops-per-tick")?)?,
227        },
228        context: ComputeThermalPowerContext {
229            thermal: string_value(cx, &get("thermal")?)?,
230            power: string_value(cx, &get("power")?)?,
231        },
232        tile_bytes: number_value(cx, &get("tile-bytes")?)?,
233        allocation_bytes: number_vec(cx, &get("allocation-bytes")?)?,
234        provenance: ComputeProfileProvenance {
235            producer: string_value(cx, &get("producer")?)?,
236            measured_at_tick: number_value(cx, &get("measured-at-tick")?)?,
237            stale_after_ticks: number_value(cx, &get("stale-after-ticks")?)?,
238        },
239        modeled: ModeledComputeProfile {
240            provider: string_value(cx, &get("modeled-provider")?)?,
241            max_queue_depth: usize::try_from(max_queue_depth)
242                .map_err(|_| Error::Eval("max-queue-depth exceeds usize".to_owned()))?,
243            max_queue_bytes: number_value(cx, &get("max-queue-bytes")?)?,
244            max_resident_bytes: number_value(cx, &get("max-resident-bytes")?)?,
245            segment_tile_bytes: number_value(cx, &get("tile-bytes")?)?,
246            max_storage_binding_bytes: number_value(cx, &get("max-storage-binding-bytes")?)?,
247            submission_deadline_ticks: number_value(cx, &get("submission-deadline-ticks")?)?,
248            fault: None,
249            auto_flush_batches: false,
250        },
251    })
252}
253
254fn number(cx: &mut Cx, value: u64) -> Result<Value> {
255    cx.factory()
256        .number_literal(Symbol::qualified("numbers", "u64"), value.to_string())
257}
258
259fn string(cx: &mut Cx, value: &str) -> Result<Value> {
260    cx.factory().string(value.to_owned())
261}
262
263fn number_list(cx: &mut Cx, values: &[u64]) -> Result<Value> {
264    let values = values
265        .iter()
266        .map(|value| number(cx, *value))
267        .collect::<Result<Vec<_>>>()?;
268    cx.factory().list(values)
269}
270
271fn string_value(cx: &mut Cx, value: &Value) -> Result<String> {
272    match value.object().as_expr(cx)? {
273        Expr::String(value) => Ok(value),
274        _ => Err(Error::TypeMismatch {
275            expected: "string",
276            found: "non-string",
277        }),
278    }
279}
280
281fn number_value(cx: &mut Cx, value: &Value) -> Result<u64> {
282    match value.object().as_expr(cx)? {
283        Expr::Number(NumberLiteral { canonical, .. }) => canonical
284            .parse::<u64>()
285            .map_err(|_| Error::Eval("invalid u64 profile value".to_owned())),
286        _ => Err(Error::TypeMismatch {
287            expected: "number",
288            found: "non-number",
289        }),
290    }
291}
292
293fn number_vec(cx: &mut Cx, value: &Value) -> Result<Vec<u64>> {
294    match value.object().as_expr(cx)? {
295        Expr::List(values) => values
296            .iter()
297            .map(|expr| match expr {
298                Expr::Number(NumberLiteral { canonical, .. }) => canonical
299                    .parse::<u64>()
300                    .map_err(|_| Error::Eval("invalid u64 profile sample".to_owned())),
301                _ => Err(Error::TypeMismatch {
302                    expected: "number list",
303                    found: "non-number list",
304                }),
305            })
306            .collect(),
307        _ => Err(Error::TypeMismatch {
308            expected: "list",
309            found: "non-list",
310        }),
311    }
312}
313
314fn sym(name: &str) -> Symbol {
315    Symbol::qualified("compute-profile", name)
316}