Skip to main content

optimization_job_api/
spec.rs

1use crate::{Result, coded, legacy, storage};
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4use std::{
5    collections::BTreeSet,
6    fs,
7    path::{Component, Path},
8};
9
10pub const SCHEMA: &str = "optimization.job/0.1";
11/// The worker protocol of this build. It adds the result contract 0.2
12/// (evidence, refutations, objective bounds and the infeasibility and
13/// certified claims) to the process protocol of 0.1.
14pub const PROTOCOL: &str = "optimization.worker/0.2";
15/// Worker protocols the runner accepts. Every 0.1 result record is a valid
16/// 0.2 record, so 0.1 workers keep working; a 0.1 runner rejects 0.2 workers
17/// at registration instead of failing on their journals.
18pub const PROTOCOLS: [&str; 2] = ["optimization.worker/0.1", PROTOCOL];
19/// True if the runner accepts a worker that describes itself with `protocol`.
20pub fn protocol_supported(protocol: &Value) -> bool {
21    protocol.as_str().is_some_and(|p| PROTOCOLS.contains(&p))
22}
23pub const MAX_INPUT_BYTES: u64 = 64 * 1024 * 1024;
24
25#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
26#[serde(deny_unknown_fields)]
27pub struct Resources {
28    pub cpu_slots: usize,
29    pub outer_workers: usize,
30    pub inner_threads: usize,
31}
32#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
33#[serde(deny_unknown_fields)]
34pub struct Limits {
35    pub wall_time_seconds: f64,
36    pub grace_seconds: f64,
37}
38#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
39#[serde(deny_unknown_fields)]
40pub struct Input {
41    pub name: String,
42    pub file: String,
43    pub sha256: String,
44}
45#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
46#[serde(deny_unknown_fields)]
47pub struct JobSpec {
48    pub schema_version: String,
49    pub request_key: String,
50    pub worker: String,
51    pub payload_schema: String,
52    pub payload: Value,
53    pub resources: Resources,
54    pub limits: Limits,
55    #[serde(default)]
56    pub inputs: Vec<Input>,
57    /// Filled by the worker during validation, never inferred by the runner.
58    #[serde(default, skip_serializing_if = "Option::is_none")]
59    pub task_count: Option<usize>,
60}
61pub fn safe_name(s: &str) -> bool {
62    !s.is_empty()
63        && s.len() <= 80
64        && s.bytes()
65            .all(|c| c.is_ascii_alphanumeric() || b"-_.".contains(&c))
66        && s != "."
67        && s != ".."
68}
69impl JobSpec {
70    pub fn validate(&self) -> Result<()> {
71        if self.schema_version != SCHEMA
72            || self.request_key.is_empty()
73            || self.request_key.len() > 200
74            || !safe_name(&self.worker)
75            || self.payload_schema.is_empty()
76            || !self.payload.is_object()
77        {
78            return Err("invalid job envelope, schema, request_key or worker".into());
79        }
80        let r = &self.resources;
81        if r.cpu_slots == 0
82            || r.cpu_slots > 1024
83            || r.outer_workers == 0
84            || r.inner_threads == 0
85            || r.outer_workers
86                .checked_mul(r.inner_threads)
87                .is_none_or(|n| n > r.cpu_slots)
88        {
89            return Err(
90                "resources require 1 <= outer_workers * inner_threads <= cpu_slots <= 1024".into(),
91            );
92        }
93        if !self.limits.wall_time_seconds.is_finite()
94            || !(0.0..=31_536_000.0).contains(&self.limits.wall_time_seconds)
95            || self.limits.wall_time_seconds == 0.0
96            || !self.limits.grace_seconds.is_finite()
97            || !(0.0..=86400.0).contains(&self.limits.grace_seconds)
98        {
99            return Err("invalid wall time or grace".into());
100        }
101        if self.task_count.is_some_and(|n| n == 0 || n > 1_000_000) || self.inputs.len() > 16 {
102            return Err("invalid task/input count".into());
103        }
104        let mut names = BTreeSet::new();
105        for i in &self.inputs {
106            if !safe_name(&i.name)
107                || !names.insert(&i.name)
108                || i.sha256.len() != 64
109                || !i
110                    .sha256
111                    .bytes()
112                    .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
113                || i.file.is_empty()
114                || Path::new(&i.file)
115                    .components()
116                    .any(|c| !matches!(c, Component::Normal(_) | Component::CurDir))
117                || !Path::new(&i.file)
118                    .components()
119                    .any(|c| matches!(c, Component::Normal(_)))
120            {
121                return Err("input requires unique safe name, relative contained path and lowercase SHA-256".into());
122            }
123        }
124        Ok(())
125    }
126    pub fn digest(&self) -> Result<String> {
127        let mut v = serde_json::to_value(self)?;
128        v.as_object_mut().unwrap().remove("request_key");
129        Ok(storage::digest(&serde_json::to_vec(&v)?))
130    }
131    pub fn input_bytes(&self, base: &Path, name: &str) -> Result<Vec<u8>> {
132        let input = self
133            .inputs
134            .iter()
135            .find(|i| i.name == name)
136            .ok_or("missing model input")?;
137        checked_input(base, input)
138    }
139}
140pub fn checked_input(base: &Path, input: &Input) -> Result<Vec<u8>> {
141    let base = fs::canonicalize(base)?;
142    let path = fs::canonicalize(base.join(&input.file))?;
143    if !path.starts_with(&base) || !path.is_file() || fs::metadata(&path)?.len() > MAX_INPUT_BYTES {
144        return Err("input escapes directory or exceeds 64 MiB".into());
145    }
146    // The bounded read also covers a file that grows after the metadata check.
147    use std::io::Read;
148    let mut bytes = Vec::new();
149    fs::File::open(path)?
150        .take(MAX_INPUT_BYTES + 1)
151        .read_to_end(&mut bytes)?;
152    if bytes.len() as u64 > MAX_INPUT_BYTES || storage::digest(&bytes) != input.sha256 {
153        return Err("input SHA-256 mismatch or oversized input".into());
154    }
155    Ok(bytes)
156}
157pub fn stage_inputs(s: &mut JobSpec, from: &Path, to: &Path) -> Result<()> {
158    fs::create_dir_all(to.join("inputs"))?;
159    for i in &mut s.inputs {
160        let bytes = checked_input(from, i)?;
161        i.file = format!("inputs/{}", i.name);
162        let path = to.join(&i.file);
163        fs::write(&path, bytes)?;
164        fs::File::open(path)?.sync_all()?;
165    }
166    fs::File::open(to.join("inputs"))?.sync_all()?;
167    Ok(())
168}
169#[derive(Clone, Debug, Serialize)]
170#[serde(untagged)]
171pub enum Request {
172    Legacy(legacy::JobSpec),
173    Job(JobSpec),
174}
175impl Request {
176    /// Parses and validates a submitted job; every failure is `invalid_job`.
177    pub fn parse(bytes: &[u8]) -> Result<Self> {
178        Self::parse_unchecked(bytes).map_err(|e| coded("invalid_job", e.to_string()))
179    }
180    fn parse_unchecked(bytes: &[u8]) -> Result<Self> {
181        let v: Value = serde_json::from_slice(bytes)?;
182        match v["schema_version"].as_str() {
183            Some("fcmaes.job/0.1") => Ok(Self::Legacy(
184                serde_json::from_value::<legacy::JobSpec>(v)?.resolve()?,
185            )),
186            Some(SCHEMA) => {
187                let s: JobSpec = serde_json::from_value(v)?;
188                s.validate()?;
189                Ok(Self::Job(s))
190            }
191            _ => Err("unsupported job schema".into()),
192        }
193    }
194    pub fn request_key(&self) -> &str {
195        match self {
196            Self::Legacy(s) => &s.request_key,
197            Self::Job(s) => &s.request_key,
198        }
199    }
200    pub fn worker(&self) -> &str {
201        match self {
202            Self::Legacy(_) => "native",
203            Self::Job(s) => &s.worker,
204        }
205    }
206    pub fn slots(&self) -> usize {
207        match self {
208            Self::Legacy(s) => s.execution.workers,
209            Self::Job(s) => s.resources.cpu_slots,
210        }
211    }
212    pub fn wall(&self) -> f64 {
213        match self {
214            Self::Legacy(s) => s.budget.wall_time_seconds,
215            Self::Job(s) => s.limits.wall_time_seconds,
216        }
217    }
218    pub fn grace(&self) -> f64 {
219        match self {
220            Self::Legacy(s) => s.budget.grace_seconds,
221            Self::Job(s) => s.limits.grace_seconds,
222        }
223    }
224    pub fn digest(&self) -> Result<String> {
225        match self {
226            Self::Legacy(s) => s.digest(),
227            Self::Job(s) => s.digest(),
228        }
229    }
230    pub fn save(&self, path: &Path) -> Result<()> {
231        Ok(storage::write_atomic_json(path, self)?)
232    }
233}