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. The CLI (`packages/cli`) depends on this crate and calls
9//! [`validate_job`]/[`task_exists`] instead of re-deriving these rules locally; `mecha10-api`'s
10//! `create_job` handler (`services/api`) is intended to do the same, but as of now does not
11//! depend on this crate and does not call either function.
12//!
13//! Deliberately Python-free: `rl_task_runner.run` (the Python runtime) does not depend on this
14//! crate and does not re-validate - it trusts that a job it's asked to run already passed
15//! validation upstream (CLI local dispatch, or mecha10-api's `create_job` for remote dispatch).
16
17use std::path::Path;
18
19/// Which training pipeline flow a job runs as.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)]
21#[serde(rename_all = "kebab-case")]
22pub enum FlowType {
23    Train,
24    Tune,
25    TuneThenTrain,
26    TrainFromTune,
27    ResumeRun,
28}
29
30impl std::fmt::Display for FlowType {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        let s = match self {
33            FlowType::Train => "train",
34            FlowType::Tune => "tune",
35            FlowType::TuneThenTrain => "tune-then-train",
36            FlowType::TrainFromTune => "train-from-tune",
37            FlowType::ResumeRun => "resume-run",
38        };
39        write!(f, "{s}")
40    }
41}
42
43/// Whether a field is required, optional (sentinel/default-backed), or not applicable for a
44/// given `flow_type` - e.g. `resume_run_id` is `Required` for `ResumeRun`, `NotApplicable` for
45/// `Train`.
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum FieldRequirement {
48    Required,
49    Optional,
50    NotApplicable,
51}
52
53/// The fields this crate has an opinion about. Deliberately not every `PPOConfig` hyperparameter
54/// (`learning_rate`, `gamma`, `clip_range`, ...) - those are legitimately optional/sentinel-
55/// defaulted for every `flow_type` today and aren't in dispute; this covers job identity
56/// (`task_name`/`algorithm`) and the fields that control how long a flow runs or which prior run
57/// it targets, which is where "required vs. optional vs. n/a" actually varies by `flow_type`.
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum JobField {
60    TaskName,
61    Algorithm,
62    TotalTimesteps,
63    ResumeRunId,
64    TuneRunId,
65    NTrials,
66    TuneTimesteps,
67    EvalEps,
68}
69
70impl JobField {
71    /// The wire/YAML field name a `CreateJobRequest`/job-definition struct is expected to carry
72    /// this field under - the single list a contract test in each consumer (the CLI's own
73    /// `CreateJobRequest`, `mecha10-api`'s `CreateJobRequest`) checks its own hand-typed struct
74    /// against, so a new variant here shows up as a failing test in both instead of silently
75    /// missing a field the way `Tune`'s `n_trials`/`tune_timesteps`/`eval_eps` did before either
76    /// request struct picked them up.
77    pub fn wire_name(self) -> &'static str {
78        match self {
79            JobField::TaskName => "task_name",
80            JobField::Algorithm => "algorithm",
81            JobField::TotalTimesteps => "total_timesteps",
82            JobField::ResumeRunId => "resume_run_id",
83            JobField::TuneRunId => "tune_run_id",
84            JobField::NTrials => "n_trials",
85            JobField::TuneTimesteps => "tune_timesteps",
86            JobField::EvalEps => "eval_eps",
87        }
88    }
89
90    /// Requiredness of this field for `flow_type`, the single table the CLI reads instead of
91    /// re-deriving it; `mecha10-api` is intended to read it too, once it's wired up to depend on
92    /// this crate.
93    pub fn requirement(self, flow_type: FlowType) -> FieldRequirement {
94        use FieldRequirement::*;
95        use FlowType::*;
96        use JobField::*;
97
98        match (self, flow_type) {
99            (TaskName, _) => Required,
100            (Algorithm, _) => Required,
101
102            (TotalTimesteps, Train | TuneThenTrain | TrainFromTune | ResumeRun) => Required,
103            (TotalTimesteps, Tune) => NotApplicable,
104
105            (ResumeRunId, ResumeRun) => Required,
106            (ResumeRunId, _) => NotApplicable,
107
108            (TuneRunId, TrainFromTune) => Required,
109            (TuneRunId, _) => NotApplicable,
110
111            (NTrials | TuneTimesteps | EvalEps, Tune | TuneThenTrain) => Required,
112            (NTrials | TuneTimesteps | EvalEps, _) => NotApplicable,
113        }
114    }
115}
116
117/// The subset of a job definition's fields this crate validates. The CLI's `JobDefinition`
118/// (`packages/cli/src/services/train_jobs.rs`) and mecha10-api's `CreateJobRequest`
119/// (`services/api/src/jobs/types.rs`) each convert their own richer type into this one (via a
120/// small `From`/helper on their side) rather than this crate depending on either of them -
121/// this stays the shared validation vocabulary, not the wire/storage shape.
122#[derive(Debug, Clone, Default)]
123pub struct JobFields {
124    pub task_name: Option<String>,
125    pub algorithm: Option<String>,
126    pub total_timesteps: Option<i64>,
127    pub resume_run_id: Option<String>,
128    pub tune_run_id: Option<String>,
129    pub n_trials: Option<i64>,
130    pub tune_timesteps: Option<i64>,
131    pub eval_eps: Option<i64>,
132}
133
134impl JobFields {
135    fn get(&self, field: JobField) -> FieldPresence<'_> {
136        match field {
137            JobField::TaskName => FieldPresence::Str(self.task_name.as_deref()),
138            JobField::Algorithm => FieldPresence::Str(self.algorithm.as_deref()),
139            JobField::TotalTimesteps => FieldPresence::Int(self.total_timesteps),
140            JobField::ResumeRunId => FieldPresence::Str(self.resume_run_id.as_deref()),
141            JobField::TuneRunId => FieldPresence::Str(self.tune_run_id.as_deref()),
142            JobField::NTrials => FieldPresence::Int(self.n_trials),
143            JobField::TuneTimesteps => FieldPresence::Int(self.tune_timesteps),
144            JobField::EvalEps => FieldPresence::Int(self.eval_eps),
145        }
146    }
147}
148
149enum FieldPresence<'a> {
150    Str(Option<&'a str>),
151    Int(Option<i64>),
152}
153
154impl FieldPresence<'_> {
155    /// A string field counts as absent when blank/whitespace-only (matches the CLI's existing
156    /// `.trim().is_empty()` checks for `algorithm`/`resume_run_id`); an int field counts as
157    /// absent when unset - `-1`/`-1.0` sentinels are a *value*, not absence, and are out of
158    /// scope for this presence check (a duration-controlling field left at its sentinel is a
159    /// job-authoring smell, not a validation error).
160    fn is_present(&self) -> bool {
161        match self {
162            FieldPresence::Str(s) => s.is_some_and(|s| !s.trim().is_empty()),
163            FieldPresence::Int(i) => i.is_some(),
164        }
165    }
166}
167
168/// One validation failure: a field that's `Required` for this job's `flow_type` but absent.
169#[derive(Debug, Clone, PartialEq, Eq)]
170pub struct MissingField {
171    pub field: JobField,
172    pub flow_type: FlowType,
173}
174
175impl std::fmt::Display for MissingField {
176    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177        write!(
178            f,
179            "'{:?}' is required for flow_type '{}' but was not set",
180            self.field, self.flow_type
181        )
182    }
183}
184
185const ALL_FIELDS: &[JobField] = &[
186    JobField::TaskName,
187    JobField::Algorithm,
188    JobField::TotalTimesteps,
189    JobField::ResumeRunId,
190    JobField::TuneRunId,
191    JobField::NTrials,
192    JobField::TuneTimesteps,
193    JobField::EvalEps,
194];
195
196/// Every [`JobField`] this crate has an opinion about - the list a consumer's own contract
197/// test (asserting its hand-typed `CreateJobRequest` carries a matching field for each one
198/// that's applicable to a flow type it actually dispatches) iterates over.
199pub fn all_fields() -> &'static [JobField] {
200    ALL_FIELDS
201}
202
203/// Validate `fields` against `flow_type`'s requirement table. Returns every missing required
204/// field (not just the first), so a caller can report the whole list at once.
205pub fn validate_job(flow_type: FlowType, fields: &JobFields) -> Vec<MissingField> {
206    ALL_FIELDS
207        .iter()
208        .filter(|&&field| field.requirement(flow_type) == FieldRequirement::Required)
209        .filter(|&&field| !fields.get(field).is_present())
210        .map(|&field| MissingField { field, flow_type })
211        .collect()
212}
213
214/// Error from [`task_exists`]: `task_name` has no corresponding `task.json` under the task
215/// catalog, or that file isn't valid JSON.
216#[derive(Debug, Clone, PartialEq, Eq)]
217pub enum TaskLookupError {
218    NotFound { task_name: String },
219    InvalidJson { task_name: String, message: String },
220}
221
222impl std::fmt::Display for TaskLookupError {
223    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
224        match self {
225            TaskLookupError::NotFound { task_name } => write!(
226                f,
227                "task '{task_name}' has no configs/rl/tasks/{task_name}/task.json - check the \
228                 task_name against the project's task catalog"
229            ),
230            TaskLookupError::InvalidJson { task_name, message } => {
231                write!(f, "configs/rl/tasks/{task_name}/task.json is not valid JSON: {message}")
232            }
233        }
234    }
235}
236
237impl std::error::Error for TaskLookupError {}
238
239/// `configs/rl/tasks` relative to a project root - kept as a single literal here (rather than a
240/// dependency on `mecha10-cli-core`, to keep this crate's dependency footprint minimal for
241/// mecha10-api's sake) and covered by `RL_TASKS_DIR`'s own equality test in
242/// `mecha10-cli-core::paths::project`, so the two can't silently drift.
243pub const RL_TASKS_DIR: &str = "configs/rl/tasks";
244
245/// Confirm `task_name` has a real `task.json` under `{project_root}/configs/rl/tasks/` and that
246/// it parses as JSON. This is the check that was previously only ever done, best-effort, deep in
247/// `rl_task_runner.run` at actual training time (and even there, on failure it only logged a
248/// warning and continued rather than failing) - moving it here means a bad `task_name` is caught
249/// at CLI/API accept-time instead of surfacing as unexplained training behavior much later.
250pub fn task_exists(project_root: &Path, task_name: &str) -> Result<(), TaskLookupError> {
251    let path = project_root.join(RL_TASKS_DIR).join(task_name).join("task.json");
252    let content = std::fs::read_to_string(&path).map_err(|_| TaskLookupError::NotFound {
253        task_name: task_name.to_string(),
254    })?;
255    serde_json::from_str::<serde_json::Value>(&content).map_err(|e| TaskLookupError::InvalidJson {
256        task_name: task_name.to_string(),
257        message: e.to_string(),
258    })?;
259    Ok(())
260}
261
262/// List every task name under `{project_root}/configs/rl/tasks/` that has a `task.json` - the
263/// catalog `mecha10 train`'s task picker prompts from after a job is selected. A job's own
264/// `configs/rl/jobs/*.yml` no longer pins `task_name` itself (that field is filled in at
265/// dispatch time from this picker instead); this is the list it prompts against. Sorted for a
266/// stable menu order; empty (not an error) when the directory doesn't exist.
267pub fn list_tasks(project_root: &Path) -> std::io::Result<Vec<String>> {
268    let dir = project_root.join(RL_TASKS_DIR);
269    if !dir.exists() {
270        return Ok(Vec::new());
271    }
272
273    let mut names: Vec<String> = std::fs::read_dir(&dir)?
274        .filter_map(|entry| entry.ok())
275        .filter(|entry| entry.path().join("task.json").is_file())
276        .filter_map(|entry| entry.file_name().into_string().ok())
277        .collect();
278    names.sort();
279    Ok(names)
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[test]
287    fn every_job_field_has_a_unique_wire_name() {
288        let names: Vec<&str> = ALL_FIELDS.iter().map(|f| f.wire_name()).collect();
289        let unique: std::collections::HashSet<&str> = names.iter().copied().collect();
290        assert_eq!(names.len(), unique.len());
291    }
292
293    fn fields_with_task_and_algo() -> JobFields {
294        JobFields {
295            task_name: Some("find_target".to_string()),
296            algorithm: Some("ppo".to_string()),
297            ..Default::default()
298        }
299    }
300
301    #[test]
302    fn train_requires_total_timesteps() {
303        let fields = fields_with_task_and_algo();
304        let missing = validate_job(FlowType::Train, &fields);
305        assert_eq!(
306            missing,
307            vec![MissingField {
308                field: JobField::TotalTimesteps,
309                flow_type: FlowType::Train
310            }]
311        );
312    }
313
314    #[test]
315    fn train_passes_with_total_timesteps_set() {
316        let mut fields = fields_with_task_and_algo();
317        fields.total_timesteps = Some(1000);
318        assert!(validate_job(FlowType::Train, &fields).is_empty());
319    }
320
321    #[test]
322    fn tune_does_not_require_total_timesteps() {
323        let mut fields = fields_with_task_and_algo();
324        fields.n_trials = Some(10);
325        fields.tune_timesteps = Some(1000);
326        fields.eval_eps = Some(2);
327        assert!(validate_job(FlowType::Tune, &fields).is_empty());
328    }
329
330    #[test]
331    fn tune_requires_tune_specific_fields() {
332        let fields = fields_with_task_and_algo();
333        let missing = validate_job(FlowType::Tune, &fields);
334        assert_eq!(missing.len(), 3);
335        assert!(missing.iter().all(|m| m.flow_type == FlowType::Tune));
336    }
337
338    #[test]
339    fn resume_run_requires_resume_run_id_and_total_timesteps() {
340        let fields = fields_with_task_and_algo();
341        let missing = validate_job(FlowType::ResumeRun, &fields);
342        assert_eq!(missing.len(), 2);
343        assert!(missing.iter().any(|m| m.field == JobField::ResumeRunId));
344        assert!(missing.iter().any(|m| m.field == JobField::TotalTimesteps));
345    }
346
347    #[test]
348    fn train_from_tune_requires_tune_run_id() {
349        let mut fields = fields_with_task_and_algo();
350        fields.total_timesteps = Some(1000);
351        let missing = validate_job(FlowType::TrainFromTune, &fields);
352        assert_eq!(
353            missing,
354            vec![MissingField {
355                field: JobField::TuneRunId,
356                flow_type: FlowType::TrainFromTune
357            }]
358        );
359    }
360
361    #[test]
362    fn missing_task_name_and_algorithm_reported_for_every_flow_type() {
363        let missing = validate_job(FlowType::Train, &JobFields::default());
364        assert!(missing.iter().any(|m| m.field == JobField::TaskName));
365        assert!(missing.iter().any(|m| m.field == JobField::Algorithm));
366    }
367
368    #[test]
369    fn blank_string_fields_count_as_absent() {
370        let mut fields = fields_with_task_and_algo();
371        fields.resume_run_id = Some("   ".to_string());
372        let missing = validate_job(FlowType::ResumeRun, &fields);
373        assert!(missing.iter().any(|m| m.field == JobField::ResumeRunId));
374    }
375
376    #[test]
377    fn task_exists_finds_bundled_find_target_task() {
378        let dir = tempfile::tempdir().unwrap();
379        let task_dir = dir.path().join("configs/rl/tasks/find_target");
380        std::fs::create_dir_all(&task_dir).unwrap();
381        std::fs::write(task_dir.join("task.json"), r#"{"env_name": "basic_arena"}"#).unwrap();
382        assert!(task_exists(dir.path(), "find_target").is_ok());
383    }
384
385    #[test]
386    fn task_exists_reports_missing_task() {
387        let dir = tempfile::tempdir().unwrap();
388        std::fs::create_dir_all(dir.path().join("configs/rl/tasks")).unwrap();
389        let err = task_exists(dir.path(), "nonexistent").unwrap_err();
390        assert_eq!(
391            err,
392            TaskLookupError::NotFound {
393                task_name: "nonexistent".to_string()
394            }
395        );
396    }
397
398    #[test]
399    fn list_tasks_returns_empty_when_directory_missing() {
400        let dir = tempfile::tempdir().unwrap();
401        assert_eq!(list_tasks(dir.path()).unwrap(), Vec::<String>::new());
402    }
403
404    #[test]
405    fn list_tasks_lists_directories_with_a_task_json_sorted() {
406        let dir = tempfile::tempdir().unwrap();
407        for name in ["task2", "find_target", "task1"] {
408            let task_dir = dir.path().join("configs/rl/tasks").join(name);
409            std::fs::create_dir_all(&task_dir).unwrap();
410            std::fs::write(task_dir.join("task.json"), "{}").unwrap();
411        }
412        std::fs::create_dir_all(dir.path().join("configs/rl/tasks/not-a-task")).unwrap();
413
414        assert_eq!(
415            list_tasks(dir.path()).unwrap(),
416            vec!["find_target".to_string(), "task1".to_string(), "task2".to_string()]
417        );
418    }
419
420    #[test]
421    fn task_exists_reports_invalid_json() {
422        let dir = tempfile::tempdir().unwrap();
423        let task_dir = dir.path().join("configs/rl/tasks/broken");
424        std::fs::create_dir_all(&task_dir).unwrap();
425        std::fs::write(task_dir.join("task.json"), "not json").unwrap();
426        let err = task_exists(dir.path(), "broken").unwrap_err();
427        assert!(matches!(err, TaskLookupError::InvalidJson { .. }));
428    }
429}