Skip to main content

mecha10_job_validation/
lib.rs

1//! Shared validation for `mecha10 train`'s job catalog (`configs/rl/jobs/*.yml`) and its
2//! `task_name` references into the task catalog (`configs/rl/tasks/*/task.json`).
3//!
4//! This crate is the single source of truth for "what's required/optional/not-applicable per
5//! `flow_type`" - before this crate existed, that question was answered three different ways in
6//! three different places (the CLI's job struct defaults, the CLI's dispatch-time `if` checks,
7//! and mecha10-api's `create_job` handler), which could and did drift out of sync with each
8//! other. Both the CLI (`packages/cli`) and `mecha10-api` (`services/api`) depend on this crate
9//! and call [`validate_job`]/[`task_exists`] instead of re-deriving these rules locally.
10//!
11//! Deliberately Python-free: `rl_task_runner.run` (the Python runtime) does not depend on this
12//! crate and does not re-validate - it trusts that a job it's asked to run already passed
13//! validation upstream (CLI local dispatch, or mecha10-api's `create_job` for remote dispatch).
14
15use std::path::Path;
16
17/// Which Kestra `flows/ml-pipeline/*.yml` flow (and `rl_task_runner.run` mode) a job maps to.
18/// Mirrors `packages/cli/src/services/train_jobs.rs::FlowType` - that's the deserialization
19/// target for job YAML, re-exported from here so both it and `mecha10-api` share one enum
20/// instead of two independently-maintained copies.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
22#[serde(rename_all = "kebab-case")]
23pub enum FlowType {
24    Train,
25    Tune,
26    TuneThenTrain,
27    TrainFromTune,
28    ResumeRun,
29}
30
31impl std::fmt::Display for FlowType {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        let s = match self {
34            FlowType::Train => "train",
35            FlowType::Tune => "tune",
36            FlowType::TuneThenTrain => "tune-then-train",
37            FlowType::TrainFromTune => "train-from-tune",
38            FlowType::ResumeRun => "resume-run",
39        };
40        write!(f, "{s}")
41    }
42}
43
44/// Whether a field is required, optional (sentinel/default-backed), or not applicable for a
45/// given `flow_type` - e.g. `resume_run_id` is `Required` for `ResumeRun`, `NotApplicable` for
46/// `Train`.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum FieldRequirement {
49    Required,
50    Optional,
51    NotApplicable,
52}
53
54/// The fields this crate has an opinion about. Deliberately not every `PPOConfig` hyperparameter
55/// (`learning_rate`, `gamma`, `clip_range`, ...) - those are legitimately optional/sentinel-
56/// defaulted for every `flow_type` today and aren't in dispute; this covers job identity
57/// (`task_name`/`algorithm`) and the fields that control how long a flow runs or which prior run
58/// it targets, which is where "required vs. optional vs. n/a" actually varies by `flow_type`.
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60pub enum JobField {
61    TaskName,
62    Algorithm,
63    TotalTimesteps,
64    ResumeRunId,
65    TuneRunId,
66    NTrials,
67    TuneTimesteps,
68    EvalEps,
69}
70
71impl JobField {
72    /// Requiredness of this field for `flow_type`, the single table both the CLI and
73    /// mecha10-api should read instead of re-deriving it.
74    pub fn requirement(self, flow_type: FlowType) -> FieldRequirement {
75        use FieldRequirement::*;
76        use FlowType::*;
77        use JobField::*;
78
79        match (self, flow_type) {
80            // task_name/algorithm identify what to train and are required for every flow type -
81            // `run.py` hard-requires both (`_resolve_task_name`/`_resolve_env_algorithm`), and
82            // `env_name_from_task_json`'s task-catalog lookup only works with a real task_name.
83            (TaskName, _) => Required,
84            (Algorithm, _) => Required,
85
86            // total_timesteps controls how long a *training* run lasts - required wherever a
87            // train phase happens, n/a for a pure tune.
88            (TotalTimesteps, Train | TuneThenTrain | TrainFromTune | ResumeRun) => Required,
89            (TotalTimesteps, Tune) => NotApplicable,
90
91            // resume_run_id only means anything for resume-run.
92            (ResumeRunId, ResumeRun) => Required,
93            (ResumeRunId, _) => NotApplicable,
94
95            // tune_run_id only means anything for train-from-tune.
96            (TuneRunId, TrainFromTune) => Required,
97            (TuneRunId, _) => NotApplicable,
98
99            // n_trials/tune_timesteps/eval_eps control how long/thorough a *tune* phase is -
100            // required wherever a tune phase happens, n/a otherwise.
101            (NTrials | TuneTimesteps | EvalEps, Tune | TuneThenTrain) => Required,
102            (NTrials | TuneTimesteps | EvalEps, _) => NotApplicable,
103        }
104    }
105}
106
107/// The subset of a job definition's fields this crate validates. The CLI's `JobDefinition`
108/// (`packages/cli/src/services/train_jobs.rs`) and mecha10-api's `CreateJobRequest`
109/// (`services/api/src/jobs/types.rs`) each convert their own richer type into this one (via a
110/// small `From`/helper on their side) rather than this crate depending on either of them -
111/// this stays the shared validation vocabulary, not the wire/storage shape.
112#[derive(Debug, Clone, Default)]
113pub struct JobFields {
114    pub task_name: Option<String>,
115    pub algorithm: Option<String>,
116    pub total_timesteps: Option<i64>,
117    pub resume_run_id: Option<String>,
118    pub tune_run_id: Option<String>,
119    pub n_trials: Option<i64>,
120    pub tune_timesteps: Option<i64>,
121    pub eval_eps: Option<i64>,
122}
123
124impl JobFields {
125    fn get(&self, field: JobField) -> FieldPresence<'_> {
126        match field {
127            JobField::TaskName => FieldPresence::Str(self.task_name.as_deref()),
128            JobField::Algorithm => FieldPresence::Str(self.algorithm.as_deref()),
129            JobField::TotalTimesteps => FieldPresence::Int(self.total_timesteps),
130            JobField::ResumeRunId => FieldPresence::Str(self.resume_run_id.as_deref()),
131            JobField::TuneRunId => FieldPresence::Str(self.tune_run_id.as_deref()),
132            JobField::NTrials => FieldPresence::Int(self.n_trials),
133            JobField::TuneTimesteps => FieldPresence::Int(self.tune_timesteps),
134            JobField::EvalEps => FieldPresence::Int(self.eval_eps),
135        }
136    }
137}
138
139enum FieldPresence<'a> {
140    Str(Option<&'a str>),
141    Int(Option<i64>),
142}
143
144impl FieldPresence<'_> {
145    /// A string field counts as absent when blank/whitespace-only (matches the CLI's existing
146    /// `.trim().is_empty()` checks for `algorithm`/`resume_run_id`); an int field counts as
147    /// absent when unset - `-1`/`-1.0` sentinels are a *value*, not absence, and are out of
148    /// scope for this presence check (a duration-controlling field left at its sentinel is a
149    /// job-authoring smell, not a validation error).
150    fn is_present(&self) -> bool {
151        match self {
152            FieldPresence::Str(s) => s.is_some_and(|s| !s.trim().is_empty()),
153            FieldPresence::Int(i) => i.is_some(),
154        }
155    }
156}
157
158/// One validation failure: a field that's `Required` for this job's `flow_type` but absent.
159#[derive(Debug, Clone, PartialEq, Eq)]
160pub struct MissingField {
161    pub field: JobField,
162    pub flow_type: FlowType,
163}
164
165impl std::fmt::Display for MissingField {
166    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167        write!(
168            f,
169            "'{:?}' is required for flow_type '{}' but was not set",
170            self.field, self.flow_type
171        )
172    }
173}
174
175const ALL_FIELDS: &[JobField] = &[
176    JobField::TaskName,
177    JobField::Algorithm,
178    JobField::TotalTimesteps,
179    JobField::ResumeRunId,
180    JobField::TuneRunId,
181    JobField::NTrials,
182    JobField::TuneTimesteps,
183    JobField::EvalEps,
184];
185
186/// Validate `fields` against `flow_type`'s requirement table. Returns every missing required
187/// field (not just the first), so a caller can report the whole list at once.
188pub fn validate_job(flow_type: FlowType, fields: &JobFields) -> Vec<MissingField> {
189    ALL_FIELDS
190        .iter()
191        .filter(|&&field| field.requirement(flow_type) == FieldRequirement::Required)
192        .filter(|&&field| !fields.get(field).is_present())
193        .map(|&field| MissingField { field, flow_type })
194        .collect()
195}
196
197/// Error from [`task_exists`]: `task_name` has no corresponding `task.json` under the task
198/// catalog, or that file isn't valid JSON.
199#[derive(Debug, Clone, PartialEq, Eq)]
200pub enum TaskLookupError {
201    NotFound { task_name: String },
202    InvalidJson { task_name: String, message: String },
203}
204
205impl std::fmt::Display for TaskLookupError {
206    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
207        match self {
208            TaskLookupError::NotFound { task_name } => write!(
209                f,
210                "task '{task_name}' has no configs/rl/tasks/{task_name}/task.json - check the \
211                 task_name against the project's task catalog"
212            ),
213            TaskLookupError::InvalidJson { task_name, message } => {
214                write!(f, "configs/rl/tasks/{task_name}/task.json is not valid JSON: {message}")
215            }
216        }
217    }
218}
219
220impl std::error::Error for TaskLookupError {}
221
222/// `configs/rl/tasks` relative to a project root - kept as a single literal here (rather than a
223/// dependency on `mecha10-cli-core`, to keep this crate's dependency footprint minimal for
224/// mecha10-api's sake) and covered by `RL_TASKS_DIR`'s own equality test in
225/// `mecha10-cli-core::paths::project`, so the two can't silently drift.
226pub const RL_TASKS_DIR: &str = "configs/rl/tasks";
227
228/// Confirm `task_name` has a real `task.json` under `{project_root}/configs/rl/tasks/` and that
229/// it parses as JSON. This is the check that was previously only ever done, best-effort, deep in
230/// `rl_task_runner.run` at actual training time (and even there, on failure it only logged a
231/// warning and continued rather than failing) - moving it here means a bad `task_name` is caught
232/// at CLI/API accept-time instead of surfacing as unexplained training behavior much later.
233pub fn task_exists(project_root: &Path, task_name: &str) -> Result<(), TaskLookupError> {
234    let path = project_root.join(RL_TASKS_DIR).join(task_name).join("task.json");
235    let content = std::fs::read_to_string(&path).map_err(|_| TaskLookupError::NotFound {
236        task_name: task_name.to_string(),
237    })?;
238    serde_json::from_str::<serde_json::Value>(&content).map_err(|e| TaskLookupError::InvalidJson {
239        task_name: task_name.to_string(),
240        message: e.to_string(),
241    })?;
242    Ok(())
243}
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248
249    fn fields_with_task_and_algo() -> JobFields {
250        JobFields {
251            task_name: Some("find_target".to_string()),
252            algorithm: Some("ppo".to_string()),
253            ..Default::default()
254        }
255    }
256
257    #[test]
258    fn train_requires_total_timesteps() {
259        let fields = fields_with_task_and_algo();
260        let missing = validate_job(FlowType::Train, &fields);
261        assert_eq!(
262            missing,
263            vec![MissingField {
264                field: JobField::TotalTimesteps,
265                flow_type: FlowType::Train
266            }]
267        );
268    }
269
270    #[test]
271    fn train_passes_with_total_timesteps_set() {
272        let mut fields = fields_with_task_and_algo();
273        fields.total_timesteps = Some(1000);
274        assert!(validate_job(FlowType::Train, &fields).is_empty());
275    }
276
277    #[test]
278    fn tune_does_not_require_total_timesteps() {
279        let mut fields = fields_with_task_and_algo();
280        fields.n_trials = Some(10);
281        fields.tune_timesteps = Some(1000);
282        fields.eval_eps = Some(2);
283        assert!(validate_job(FlowType::Tune, &fields).is_empty());
284    }
285
286    #[test]
287    fn tune_requires_tune_specific_fields() {
288        let fields = fields_with_task_and_algo();
289        let missing = validate_job(FlowType::Tune, &fields);
290        assert_eq!(missing.len(), 3);
291        assert!(missing.iter().all(|m| m.flow_type == FlowType::Tune));
292    }
293
294    #[test]
295    fn resume_run_requires_resume_run_id_and_total_timesteps() {
296        let fields = fields_with_task_and_algo();
297        let missing = validate_job(FlowType::ResumeRun, &fields);
298        assert_eq!(missing.len(), 2);
299        assert!(missing.iter().any(|m| m.field == JobField::ResumeRunId));
300        assert!(missing.iter().any(|m| m.field == JobField::TotalTimesteps));
301    }
302
303    #[test]
304    fn train_from_tune_requires_tune_run_id() {
305        let mut fields = fields_with_task_and_algo();
306        fields.total_timesteps = Some(1000);
307        let missing = validate_job(FlowType::TrainFromTune, &fields);
308        assert_eq!(
309            missing,
310            vec![MissingField {
311                field: JobField::TuneRunId,
312                flow_type: FlowType::TrainFromTune
313            }]
314        );
315    }
316
317    #[test]
318    fn missing_task_name_and_algorithm_reported_for_every_flow_type() {
319        let missing = validate_job(FlowType::Train, &JobFields::default());
320        assert!(missing.iter().any(|m| m.field == JobField::TaskName));
321        assert!(missing.iter().any(|m| m.field == JobField::Algorithm));
322    }
323
324    #[test]
325    fn blank_string_fields_count_as_absent() {
326        let mut fields = fields_with_task_and_algo();
327        fields.resume_run_id = Some("   ".to_string());
328        let missing = validate_job(FlowType::ResumeRun, &fields);
329        assert!(missing.iter().any(|m| m.field == JobField::ResumeRunId));
330    }
331
332    #[test]
333    fn task_exists_finds_bundled_find_target_task() {
334        let dir = tempfile::tempdir().unwrap();
335        let task_dir = dir.path().join("configs/rl/tasks/find_target");
336        std::fs::create_dir_all(&task_dir).unwrap();
337        std::fs::write(task_dir.join("task.json"), r#"{"env_name": "basic_arena"}"#).unwrap();
338        assert!(task_exists(dir.path(), "find_target").is_ok());
339    }
340
341    #[test]
342    fn task_exists_reports_missing_task() {
343        let dir = tempfile::tempdir().unwrap();
344        std::fs::create_dir_all(dir.path().join("configs/rl/tasks")).unwrap();
345        let err = task_exists(dir.path(), "nonexistent").unwrap_err();
346        assert_eq!(
347            err,
348            TaskLookupError::NotFound {
349                task_name: "nonexistent".to_string()
350            }
351        );
352    }
353
354    #[test]
355    fn task_exists_reports_invalid_json() {
356        let dir = tempfile::tempdir().unwrap();
357        let task_dir = dir.path().join("configs/rl/tasks/broken");
358        std::fs::create_dir_all(&task_dir).unwrap();
359        std::fs::write(task_dir.join("task.json"), "not json").unwrap();
360        let err = task_exists(dir.path(), "broken").unwrap_err();
361        assert!(matches!(err, TaskLookupError::InvalidJson { .. }));
362    }
363}