zsfm_core/
input.rs

1use anyhow::Context;
2
3/// Maximum stdin JSON payload accepted by `infer` commands (64 MiB).
4pub const MAX_STDIN_BYTES: u64 = 64 << 20;
5/// Maximum forecast horizon steps accepted by `infer` commands.
6pub const MAX_HORIZON: usize = 4096;
7/// Maximum total context values accepted per request (flat count).
8pub const MAX_CONTEXT_VALUES: usize = 1_000_000;
9
10/// Parse the `context` field of a forecast request into `[batch][variate][time]`.
11///
12/// Accepted shapes:
13/// - `[t0, t1, ...]` → batch=1, n_var=1 (flat single series)
14/// - `[[t0, t1, ...], ...]` → batch=N, n_var=1 (batch of univariate)
15/// - `[[[v0_t0, ...], [v1_t0, ...]], ...]` → batch=N, n_var=M (batch of multivariate)
16pub fn parse_mv_contexts(val: serde_json::Value) -> anyhow::Result<Vec<Vec<Vec<f32>>>> {
17    let parsed = parse_mv_contexts_inner(val)?;
18    validate_contexts(&parsed)?;
19    Ok(parsed)
20}
21
22fn parse_mv_contexts_inner(val: serde_json::Value) -> anyhow::Result<Vec<Vec<Vec<f32>>>> {
23    match val {
24        serde_json::Value::Array(arr) if arr.is_empty() => {
25            anyhow::bail!("context must be a non-empty array")
26        }
27        serde_json::Value::Array(arr) => {
28            let first = arr.first().context("context must be a non-empty array")?;
29            if !first.is_array() {
30                // Flat: [t0, t1, ...] → batch=1, n_var=1
31                let ctx = serde_json::from_value::<Vec<f32>>(serde_json::Value::Array(arr))
32                    .context("context must be a JSON array of numbers")?;
33                anyhow::ensure!(!ctx.is_empty(), "context series must not be empty");
34                Ok(vec![vec![ctx]])
35            } else if first
36                .as_array()
37                .and_then(|a| a.first())
38                .map(|v| v.is_array())
39                .unwrap_or(false)
40            {
41                // 3D: [batch][variate][time]
42                let out: Vec<Vec<Vec<f32>>> = arr
43                    .into_iter()
44                    .enumerate()
45                    .map(|(i, batch_item)| {
46                        let variates = serde_json::from_value::<Vec<Vec<f32>>>(batch_item)
47                            .with_context(|| {
48                                format!("context[{i}] must be an array of variate arrays")
49                            })?;
50                        anyhow::ensure!(!variates.is_empty(), "context[{i}] must not be empty");
51                        for (vi, series) in variates.iter().enumerate() {
52                            anyhow::ensure!(
53                                !series.is_empty(),
54                                "context[{i}][{vi}] must not be empty"
55                            );
56                        }
57                        Ok(variates)
58                    })
59                    .collect::<anyhow::Result<_>>()?;
60                anyhow::ensure!(!out.is_empty(), "context must be a non-empty array");
61                Ok(out)
62            } else {
63                // 2D: [batch][time] → each series is n_var=1
64                let out: Vec<Vec<Vec<f32>>> = arr
65                    .into_iter()
66                    .enumerate()
67                    .map(|(i, v)| {
68                        let series = serde_json::from_value::<Vec<f32>>(v)
69                            .with_context(|| format!("context[{i}] must be an array of numbers"))?;
70                        anyhow::ensure!(!series.is_empty(), "context[{i}] must not be empty");
71                        Ok(vec![series])
72                    })
73                    .collect::<anyhow::Result<_>>()?;
74                anyhow::ensure!(!out.is_empty(), "context must be a non-empty array");
75                Ok(out)
76            }
77        }
78        _ => anyhow::bail!("context must be a JSON array"),
79    }
80}
81
82/// Reject empty, oversized, or non-finite contexts before inference.
83pub fn validate_contexts(contexts: &[Vec<Vec<f32>>]) -> anyhow::Result<()> {
84    anyhow::ensure!(!contexts.is_empty(), "context must be a non-empty array");
85    let mut total: usize = 0;
86    for (bi, batch) in contexts.iter().enumerate() {
87        anyhow::ensure!(!batch.is_empty(), "context[{bi}] must not be empty");
88        for (vi, series) in batch.iter().enumerate() {
89            anyhow::ensure!(!series.is_empty(), "context[{bi}][{vi}] must not be empty");
90            total = total.saturating_add(series.len());
91            anyhow::ensure!(
92                total <= MAX_CONTEXT_VALUES,
93                "context too large — got {total} values, max is {MAX_CONTEXT_VALUES}"
94            );
95            for &v in series {
96                anyhow::ensure!(
97                    v.is_finite(),
98                    "context[{bi}][{vi}] must contain only finite numbers"
99                );
100            }
101        }
102    }
103    Ok(())
104}
105
106/// Parse and bound the `horizon` field of a forecast request.
107pub fn parse_horizon(val: &serde_json::Value) -> anyhow::Result<usize> {
108    let horizon = val
109        .get("horizon")
110        .and_then(|v| v.as_u64())
111        .context("horizon must be a positive integer")? as usize;
112    validate_horizon(horizon)?;
113    Ok(horizon)
114}
115
116/// Enforce `1 <= horizon <= MAX_HORIZON`.
117pub fn validate_horizon(horizon: usize) -> anyhow::Result<()> {
118    anyhow::ensure!(horizon >= 1, "horizon must be a positive integer");
119    anyhow::ensure!(
120        horizon <= MAX_HORIZON,
121        "horizon too large — got {horizon}, max is {MAX_HORIZON}"
122    );
123    Ok(())
124}
125
126/// Parse a 2D numeric matrix for tabular `infer` requests.
127pub fn parse_matrix(val: &serde_json::Value) -> anyhow::Result<Vec<Vec<f32>>> {
128    let mat: Vec<Vec<f32>> =
129        serde_json::from_value(val.clone()).context("expected a 2D JSON array")?;
130    anyhow::ensure!(!mat.is_empty(), "matrix must be a non-empty array");
131    let cols = mat[0].len();
132    anyhow::ensure!(cols > 0, "matrix rows must not be empty");
133    anyhow::ensure!(
134        mat.len() * cols <= MAX_CONTEXT_VALUES,
135        "matrix too large — max is {MAX_CONTEXT_VALUES} values"
136    );
137    for (ri, row) in mat.iter().enumerate() {
138        anyhow::ensure!(row.len() == cols, "row {ri} has inconsistent width");
139        for &v in row {
140            anyhow::ensure!(
141                v.is_finite(),
142                "matrix row {ri} must contain only finite numbers"
143            );
144        }
145    }
146    Ok(mat)
147}
148
149/// Numerically stable softmax over host `f32` logits.
150pub fn softmax(logits: &[f32]) -> Vec<f32> {
151    if logits.is_empty() {
152        return Vec::new();
153    }
154    let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
155    // `max` is finite here: callers validate finiteness, empty case handled above.
156    let exps: Vec<f32> = logits.iter().map(|&v| (v - max).exp()).collect();
157    let sum: f32 = exps.iter().sum();
158    if sum == 0.0 || !sum.is_finite() {
159        return vec![1.0 / logits.len() as f32; logits.len()];
160    }
161    exps.into_iter().map(|v| v / sum).collect()
162}
163
164/// Read stdin with a byte cap so a huge pipe cannot OOM before JSON parsing.
165pub fn read_stdin_limited() -> anyhow::Result<String> {
166    use std::io::Read;
167    let mut buf = String::new();
168    std::io::stdin()
169        .take(MAX_STDIN_BYTES + 1)
170        .read_to_string(&mut buf)
171        .context("read stdin")?;
172    anyhow::ensure!(
173        (buf.len() as u64) <= MAX_STDIN_BYTES,
174        "stdin input too large — max is {MAX_STDIN_BYTES} bytes"
175    );
176    Ok(buf)
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182
183    #[test]
184    fn rejects_empty_contexts() {
185        assert!(parse_mv_contexts(serde_json::json!([])).is_err());
186        assert!(parse_mv_contexts(serde_json::json!([[]])).is_err());
187        assert!(parse_mv_contexts(serde_json::json!([[[]]])).is_err());
188    }
189
190    #[test]
191    fn rejects_non_finite_context() {
192        let nan = f64::NAN;
193        // JSON cannot encode NaN, so inject via f32 bit pattern through Value.
194        let v = serde_json::json!([1.0, 2.0]);
195        let mut ctx = parse_mv_contexts(v).unwrap();
196        ctx[0][0][0] = f32::from_bits(0x7fc0_0000);
197        assert!(validate_contexts(&ctx).is_err());
198        let _ = nan;
199    }
200
201    #[test]
202    fn horizon_bounds() {
203        assert!(validate_horizon(0).is_err());
204        assert!(validate_horizon(1).is_ok());
205        assert!(validate_horizon(MAX_HORIZON).is_ok());
206        assert!(validate_horizon(MAX_HORIZON + 1).is_err());
207    }
208
209    #[test]
210    fn matrix_validation() {
211        assert!(parse_matrix(&serde_json::json!([])).is_err());
212        assert!(parse_matrix(&serde_json::json!([[1.0], [2.0, 3.0]])).is_err());
213        assert!(parse_matrix(&serde_json::json!([[1.0, 2.0]])).is_ok());
214        assert_eq!(softmax(&[]).len(), 0);
215        let p = softmax(&[1.0, 2.0]);
216        assert!((p.iter().sum::<f32>() - 1.0).abs() < 1e-5);
217    }
218}
zsfm_core/input.rs

zsfm_core/
input.rs
