optimization_job_api/
spec.rs1use 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";
11pub const PROTOCOL: &str = "optimization.worker/0.2";
15pub const PROTOCOLS: [&str; 2] = ["optimization.worker/0.1", PROTOCOL];
19pub 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 #[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 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 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}