mecha10_cli/services/
train_jobs.rs

1//! Job catalog for `mecha10 train` (LAB-1776).
2//!
3//! A "job" is a YAML file under a project's `configs/rl/jobs/` directory that maps to one of
4//! the Kestra `flows/ml-pipeline/*.yml` flow types plus an env/algorithm/hyperparameter selection.
5
6use anyhow::{bail, Context, Result};
7use serde::{Deserialize, Serialize};
8use std::path::{Path, PathBuf};
9
10pub const JOB_CATALOG_DIR: &str = "configs/rl/jobs";
11
12pub use mecha10_job_validation::FlowType;
13
14/// One dimension of a hyperparameter search space for a tuning job.
15#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, schemars::JsonSchema)]
16pub struct SearchSpaceEntry {
17    pub name: String,
18    #[serde(rename = "type")]
19    pub type_: SearchSpaceType,
20    #[serde(default)]
21    pub low: Option<f64>,
22    #[serde(default)]
23    pub high: Option<f64>,
24    #[serde(default)]
25    pub log: Option<bool>,
26    #[serde(default)]
27    pub choices: Option<Vec<serde_json::Value>>,
28}
29
30/// `search_space` entry kind - `float`/`int` take `low`/`high` (and optional `log`),
31/// `categorical` takes `choices`.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize, schemars::JsonSchema)]
33#[serde(rename_all = "lowercase")]
34pub enum SearchSpaceType {
35    Float,
36    Int,
37    Categorical,
38}
39
40/// A training job definition, loaded from `configs/rl/jobs/<name>.yml`.
41#[derive(Debug, Clone, Deserialize, schemars::JsonSchema)]
42#[allow(dead_code)]
43pub struct JobDefinition {
44    pub name: String,
45    #[serde(default)]
46    pub description: Option<String>,
47    pub flow_type: FlowType,
48    /// Name of the simulation environment to train in. If omitted, it's resolved from the
49    /// selected task.
50    #[serde(default)]
51    #[schemars(example = Some("warehouse_nav".to_string()))]
52    pub env_name: Option<String>,
53    /// Name of the task/curriculum to train on, selected from the project's task catalog.
54    #[serde(default)]
55    #[schemars(example = Some("patrol_curriculum".to_string()))]
56    pub task_name: Option<String>,
57    /// Name of the reinforcement learning algorithm to use for this job.
58    #[serde(default)]
59    #[schemars(example = Some("ppo".to_string()))]
60    pub algorithm: Option<String>,
61
62    /// Name of the policy network architecture to use. If omitted, it's auto-detected from the
63    /// environment.
64    #[serde(default)]
65    #[schemars(example = Some("MlpPolicy".to_string()))]
66    #[allow(dead_code, reason = "schema/deserialize-only, see doc comment above")]
67    pub policy: Option<String>,
68    #[serde(default)]
69    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
70    pub seed: Option<i64>,
71    #[serde(default)]
72    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
73    pub environment_map: Option<String>,
74    #[serde(default)]
75    pub total_timesteps: Option<i64>,
76    #[serde(default)]
77    pub learning_rate: Option<f64>,
78    #[serde(default)]
79    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
80    pub linear_lr_schedule: Option<bool>,
81    #[serde(default)]
82    pub n_steps: Option<i64>,
83    #[serde(default)]
84    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
85    pub batch_size: Option<i64>,
86    #[serde(default)]
87    pub n_epochs: Option<i64>,
88    #[serde(default)]
89    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
90    pub gamma: Option<f64>,
91    #[serde(default)]
92    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
93    pub gae_lambda: Option<f64>,
94    #[serde(default)]
95    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
96    pub clip_range: Option<f64>,
97    #[serde(default)]
98    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
99    pub ent_coef: Option<f64>,
100    #[serde(default)]
101    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
102    pub vf_coef: Option<f64>,
103    #[serde(default)]
104    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
105    pub max_grad_norm: Option<f64>,
106    /// Number of evaluations with no improvement to wait before stopping training early.
107    #[serde(default)]
108    #[schemars(example = Some(10))]
109    pub patience: Option<i64>,
110    #[serde(default)]
111    pub min_delta: Option<f64>,
112    /// Number of parallel environments to run during training.
113    #[serde(default)]
114    #[schemars(example = Some(8))]
115    pub n_envs: Option<i64>,
116
117    #[serde(default)]
118    pub n_trials: Option<i64>,
119    #[serde(default)]
120    pub tune_timesteps: Option<i64>,
121    #[serde(default)]
122    pub eval_eps: Option<i64>,
123    /// Hyperparameter search space to explore for a tuning job. If omitted, a default search
124    /// space is used.
125    #[serde(default)]
126    pub search_space: Option<Vec<SearchSpaceEntry>>,
127
128    #[serde(skip)]
129    #[schemars(skip)]
130    pub source_path: PathBuf,
131}
132
133#[derive(Debug, Clone, Serialize)]
134pub struct CreateJobRequest {
135    pub name: String,
136    #[serde(skip_serializing_if = "Option::is_none")]
137    pub description: Option<String>,
138    pub algorithm: String,
139    pub robot_id: String,
140    #[serde(skip_serializing_if = "Option::is_none")]
141    pub project_id: Option<String>,
142    #[serde(skip_serializing_if = "Option::is_none")]
143    pub project_name: Option<String>,
144    #[serde(skip_serializing_if = "Option::is_none")]
145    pub env_name: Option<String>,
146    #[serde(skip_serializing_if = "Option::is_none")]
147    pub task_name: Option<String>,
148    #[serde(skip_serializing_if = "Option::is_none")]
149    pub total_timesteps: Option<i64>,
150    #[serde(skip_serializing_if = "Option::is_none")]
151    pub n_steps: Option<i64>,
152    #[serde(skip_serializing_if = "Option::is_none")]
153    pub n_epochs: Option<i64>,
154    #[serde(skip_serializing_if = "Option::is_none")]
155    pub learning_rate: Option<f64>,
156    #[serde(skip_serializing_if = "Option::is_none")]
157    pub patience: Option<i64>,
158    #[serde(skip_serializing_if = "Option::is_none")]
159    pub min_delta: Option<f64>,
160    #[serde(skip_serializing_if = "Option::is_none")]
161    pub n_envs: Option<i64>,
162    #[serde(skip_serializing_if = "Option::is_none")]
163    pub config_yaml: Option<String>,
164    #[serde(skip_serializing_if = "Option::is_none")]
165    pub resume_run_id: Option<String>,
166    #[serde(skip_serializing_if = "Option::is_none")]
167    pub n_trials: Option<i64>,
168    #[serde(skip_serializing_if = "Option::is_none")]
169    pub tune_timesteps: Option<i64>,
170    #[serde(skip_serializing_if = "Option::is_none")]
171    pub eval_eps: Option<i64>,
172    #[serde(skip_serializing_if = "Option::is_none")]
173    pub search_space: Option<Vec<SearchSpaceEntry>>,
174}
175
176/// Read `configs/rl/tasks/<task_name>/task.json`'s `env_name` field, relative to `project_root`.
177/// Returns `None` on any failure (missing file, invalid JSON, missing/non-string field).
178fn env_name_from_task_json(project_root: &Path, task_name: &str) -> Option<String> {
179    task_json_field(project_root, task_name, "env_name")
180}
181
182fn task_json_field(project_root: &Path, task_name: &str, field: &str) -> Option<String> {
183    let path = project_root
184        .join(mecha10_cli_core::paths::project::RL_TASKS_DIR)
185        .join(task_name)
186        .join("task.json");
187    let content = std::fs::read_to_string(path).ok()?;
188    let value: serde_json::Value = serde_json::from_str(&content).ok()?;
189    value.get(field)?.as_str().map(str::to_string)
190}
191
192/// Drop `value` if it equals `sentinel` - the "unset" placeholder some YAML/env fields use
193/// instead of a real `null` (e.g. `-1` / `-1.0`).
194fn non_sentinel<T: PartialEq>(value: Option<T>, sentinel: T) -> Option<T> {
195    value.filter(|v| *v != sentinel)
196}
197
198fn non_sentinel_i64(value: Option<i64>) -> Option<i64> {
199    non_sentinel(value, -1)
200}
201
202fn non_sentinel_f64(value: Option<f64>) -> Option<f64> {
203    non_sentinel(value, -1.0)
204}
205
206impl JobDefinition {
207    /// This job's fields in `mecha10-job-validation`'s shared vocabulary (see
208    /// [`mecha10_job_validation::JobFields`]).
209    fn validation_fields(&self, resume_run_id: Option<&str>) -> mecha10_job_validation::JobFields {
210        mecha10_job_validation::JobFields {
211            task_name: self.task_name.clone(),
212            algorithm: self.algorithm.clone(),
213            resume_run_id: resume_run_id.map(str::to_string),
214            tune_run_id: None,
215            n_trials: non_sentinel_i64(self.n_trials),
216            tune_timesteps: non_sentinel_i64(self.tune_timesteps),
217            eval_eps: non_sentinel_i64(self.eval_eps),
218        }
219    }
220
221    /// This job's per-algorithm hyperparameter fields in `mecha10-job-validation`'s shared
222    /// vocabulary (see [`mecha10_job_validation::AlgorithmFields`]) - the by-algorithm
223    /// counterpart to [`Self::validation_fields`].
224    fn validation_fields_for_algorithm(&self) -> mecha10_job_validation::AlgorithmFields {
225        mecha10_job_validation::AlgorithmFields {
226            total_timesteps: non_sentinel_i64(self.total_timesteps),
227            learning_rate: non_sentinel_f64(self.learning_rate),
228            n_steps: non_sentinel_i64(self.n_steps),
229            n_epochs: non_sentinel_i64(self.n_epochs),
230            patience: non_sentinel_i64(self.patience),
231            min_delta: non_sentinel_f64(self.min_delta),
232        }
233    }
234
235    /// Validate this job against `mecha10-job-validation`'s flow_type requirement table and its
236    /// by-algorithm hyperparameter requirement table, plus a `task_name`-exists check against
237    /// the project's own task catalog. Algorithm-awareness (e.g. `ppo` needing `n_steps`, `act`
238    /// needing none of these fields) lives entirely in `mecha10-job-validation`'s
239    /// `AlgorithmField` table - this method only composes the two generic validators' results,
240    /// the same way `compose_node_schema` composes a base schema with a node's own custom
241    /// schema rather than validating them separately.
242    pub fn validate(&self, project_root: &Path, resume_run_id: Option<&str>) -> Result<()> {
243        let missing_basic =
244            mecha10_job_validation::validate_job(self.flow_type, &self.validation_fields(resume_run_id));
245        let missing_algorithm = mecha10_job_validation::validate_algorithm_fields(
246            self.algorithm.as_deref().unwrap_or_default(),
247            self.flow_type,
248            &self.validation_fields_for_algorithm(),
249        );
250
251        if !missing_basic.is_empty() || !missing_algorithm.is_empty() {
252            let details = missing_basic
253                .iter()
254                .map(|m| m.to_string())
255                .chain(missing_algorithm.iter().map(|m| m.to_string()))
256                .collect::<Vec<_>>()
257                .join("; ");
258            bail!("job '{}' is missing required fields: {details}", self.name);
259        }
260
261        if let Some(task_name) = self.task_name.as_deref() {
262            mecha10_job_validation::task_exists(project_root, task_name)
263                .with_context(|| format!("job '{}' has an invalid task_name", self.name))?;
264        }
265
266        Ok(())
267    }
268
269    /// Resolve which `mecha10_rl.envs` registry name this job targets: an explicit `env_name`
270    /// wins, else it's read from `task_name`'s own `task.json`. `None` when neither is available.
271    pub fn resolve_env_name(&self, project_root: &Path) -> Option<String> {
272        self.env_name.clone().or_else(|| {
273            self.task_name
274                .as_deref()
275                .and_then(|task_name| env_name_from_task_json(project_root, task_name))
276        })
277    }
278
279    /// Whether this job can be dispatched remotely with today's `CreateJobRequest`/`create_job`
280    /// contract. `tune-then-train`/`train-from-tune` aren't dispatchable yet and stay local-only.
281    pub fn supports_remote(&self) -> bool {
282        matches!(self.flow_type, FlowType::Train | FlowType::ResumeRun | FlowType::Tune)
283    }
284
285    /// Build the remote dispatch request body. Only meaningful when `supports_remote()` is true;
286    /// callers should check that first.
287    pub fn to_create_job_request(
288        &self,
289        project_root: &Path,
290        robot_id: &str,
291        project_name: &str,
292        resume_run_id: Option<&str>,
293    ) -> Result<CreateJobRequest> {
294        let config_yaml = std::fs::read_to_string(&self.source_path).with_context(|| {
295            format!(
296                "failed to read {} - this job definition's own file must exist before dispatching remotely",
297                self.source_path.display()
298            )
299        })?;
300
301        Ok(CreateJobRequest {
302            name: self.name.clone(),
303            description: self.description.clone(),
304            algorithm: self.algorithm.clone().unwrap_or_default(),
305            robot_id: robot_id.to_string(),
306            project_id: Some(robot_id.to_string()),
307            project_name: Some(project_name.to_string()),
308            env_name: self.resolve_env_name(project_root),
309            task_name: self.task_name.clone(),
310            total_timesteps: non_sentinel_i64(self.total_timesteps),
311            n_steps: non_sentinel_i64(self.n_steps),
312            n_epochs: non_sentinel_i64(self.n_epochs),
313            learning_rate: non_sentinel_f64(self.learning_rate),
314            patience: non_sentinel_i64(self.patience),
315            min_delta: non_sentinel_f64(self.min_delta),
316            n_envs: non_sentinel_i64(self.n_envs),
317            config_yaml: Some(config_yaml),
318            resume_run_id: resume_run_id
319                .map(str::trim)
320                .filter(|id| !id.is_empty())
321                .map(str::to_string),
322            n_trials: non_sentinel_i64(self.n_trials),
323            tune_timesteps: non_sentinel_i64(self.tune_timesteps),
324            eval_eps: non_sentinel_i64(self.eval_eps),
325            search_space: self.search_space.clone(),
326        })
327    }
328}
329
330/// Load every `.yml`/`.yaml` job definition under `{project_root}/configs/rl/jobs/`, sorted by
331/// file name for a stable menu order. Returns an empty list (not an error) when the directory
332/// doesn't exist - `mecha10 train` treats that as "no jobs configured yet" and points the user at
333/// `mecha10 init`/the bundled defaults rather than failing.
334pub fn load_job_catalog(project_root: &Path) -> Result<Vec<JobDefinition>> {
335    let jobs_dir = project_root.join(JOB_CATALOG_DIR);
336    if !jobs_dir.exists() {
337        return Ok(Vec::new());
338    }
339
340    let mut paths: Vec<PathBuf> = std::fs::read_dir(&jobs_dir)
341        .with_context(|| format!("failed to read {}", jobs_dir.display()))?
342        .filter_map(|entry| entry.ok())
343        .map(|entry| entry.path())
344        .filter(|path| {
345            path.extension()
346                .and_then(|ext| ext.to_str())
347                .is_some_and(|ext| ext == "yml" || ext == "yaml")
348        })
349        .collect();
350    paths.sort();
351
352    let mut jobs = Vec::with_capacity(paths.len());
353    for path in paths {
354        let content = std::fs::read_to_string(&path).with_context(|| format!("failed to read {}", path.display()))?;
355        let mut job: JobDefinition = serde_yaml::from_str(&content)
356            .with_context(|| format!("failed to parse job definition {}", path.display()))?;
357        job.source_path = path;
358        jobs.push(job);
359    }
360
361    Ok(jobs)
362}
363
364#[cfg(test)]
365mod tests {
366    use super::*;
367    use tempfile::TempDir;
368
369    fn write_job(dir: &Path, file_name: &str, contents: &str) {
370        std::fs::write(dir.join(file_name), contents).unwrap();
371    }
372
373    fn write_task_json(project_root: &Path, task_name: &str, contents: &str) {
374        let task_dir = project_root
375            .join(mecha10_cli_core::paths::project::RL_TASKS_DIR)
376            .join(task_name);
377        std::fs::create_dir_all(&task_dir).unwrap();
378        std::fs::write(task_dir.join("task.json"), contents).unwrap();
379    }
380
381    #[test]
382    fn resolve_env_name_prefers_an_explicit_env_name_over_the_tasks_own() {
383        let dir = TempDir::new().unwrap();
384        write_task_json(dir.path(), "find_target", r#"{"env_name": "basic_arena"}"#);
385
386        let job = JobDefinition {
387            env_name: Some("rover_arena".to_string()),
388            task_name: Some("find_target".to_string()),
389            ..base_job()
390        };
391        assert_eq!(job.resolve_env_name(dir.path()), Some("rover_arena".to_string()));
392    }
393
394    #[test]
395    fn resolve_env_name_falls_back_to_the_tasks_own_task_json() {
396        let dir = TempDir::new().unwrap();
397        write_task_json(
398            dir.path(),
399            "find_target",
400            r#"{"env_name": "basic_arena", "model": "humanoid"}"#,
401        );
402
403        let job = JobDefinition {
404            task_name: Some("find_target".to_string()),
405            ..base_job()
406        };
407        assert_eq!(job.resolve_env_name(dir.path()), Some("basic_arena".to_string()));
408    }
409
410    #[test]
411    fn resolve_env_name_is_none_when_task_has_no_env_name_field() {
412        let dir = TempDir::new().unwrap();
413        write_task_json(dir.path(), "find_target", r#"{"model": "humanoid"}"#);
414
415        let job = JobDefinition {
416            task_name: Some("find_target".to_string()),
417            ..base_job()
418        };
419        assert_eq!(job.resolve_env_name(dir.path()), None);
420    }
421
422    #[test]
423    fn resolve_env_name_is_none_without_an_env_name_or_a_task_name() {
424        let dir = TempDir::new().unwrap();
425        let job = base_job();
426        assert_eq!(job.resolve_env_name(dir.path()), None);
427    }
428
429    /// A minimal, fully-`None`/default `JobDefinition` for tests to `..base_job()` from, so each
430    /// test only spells out the fields it actually cares about.
431    fn base_job() -> JobDefinition {
432        JobDefinition {
433            name: "job".to_string(),
434            description: None,
435            flow_type: FlowType::Train,
436            env_name: None,
437            task_name: None,
438            algorithm: Some("ppo".to_string()),
439            policy: None,
440            seed: None,
441            environment_map: None,
442            total_timesteps: None,
443            learning_rate: None,
444            linear_lr_schedule: None,
445            n_steps: None,
446            batch_size: None,
447            n_epochs: None,
448            gamma: None,
449            gae_lambda: None,
450            clip_range: None,
451            ent_coef: None,
452            vf_coef: None,
453            max_grad_norm: None,
454            patience: None,
455            min_delta: None,
456            n_envs: None,
457            n_trials: None,
458            tune_timesteps: None,
459            eval_eps: None,
460            search_space: None,
461            source_path: PathBuf::new(),
462        }
463    }
464
465    /// A `ppo` job with every one of its required hyperparameters set (per
466    /// `mecha10_job_validation::AlgorithmField::requirement`) except `total_timesteps`/
467    /// `n_steps`, which individual tests below null out to isolate the field they're checking.
468    fn ppo_job_with_hyperparams_set() -> JobDefinition {
469        JobDefinition {
470            task_name: Some("find_target".to_string()),
471            algorithm: Some("ppo".to_string()),
472            total_timesteps: Some(1000),
473            learning_rate: Some(0.0003),
474            n_steps: Some(2048),
475            n_epochs: Some(10),
476            patience: Some(10),
477            min_delta: Some(0.01),
478            ..base_job()
479        }
480    }
481
482    #[test]
483    fn validate_does_not_require_any_hyperparameter_fields_for_a_non_ppo_algorithm() {
484        let temp = TempDir::new().unwrap();
485        write_task_json(temp.path(), "find_target", r#"{"env_name": "basic_arena"}"#);
486
487        let job = JobDefinition {
488            task_name: Some("find_target".to_string()),
489            algorithm: Some("act".to_string()),
490            total_timesteps: Some(1000),
491            n_steps: None,
492            ..base_job()
493        };
494
495        assert!(job.validate(temp.path(), None).is_ok());
496    }
497
498    /// Regression test for the bug this crate's base/by-algorithm split exists to prevent from
499    /// recurring: `act` has no `total_timesteps` concept at all (it reads `total_steps`
500    /// instead, an untyped field today), so a `Train`-flow `act` job that omits
501    /// `total_timesteps` entirely must still pass `validate()` - it used to fail, because the
502    /// old hand-rolled check required `total_timesteps` unconditionally for every algorithm.
503    #[test]
504    fn validate_does_not_require_total_timesteps_for_an_act_train_job() {
505        let temp = TempDir::new().unwrap();
506        write_task_json(temp.path(), "find_target", r#"{"env_name": "basic_arena"}"#);
507
508        let job = JobDefinition {
509            flow_type: FlowType::Train,
510            task_name: Some("find_target".to_string()),
511            algorithm: Some("act".to_string()),
512            total_timesteps: None,
513            n_steps: None,
514            ..base_job()
515        };
516
517        assert!(job.validate(temp.path(), None).is_ok());
518    }
519
520    #[test]
521    fn validate_still_requires_n_steps_for_ppo() {
522        let temp = TempDir::new().unwrap();
523        write_task_json(temp.path(), "find_target", r#"{"env_name": "basic_arena"}"#);
524
525        let job = JobDefinition {
526            n_steps: None,
527            ..ppo_job_with_hyperparams_set()
528        };
529
530        let err = job.validate(temp.path(), None).unwrap_err();
531        assert!(
532            err.to_string().contains("n_steps"),
533            "expected error to mention n_steps, got: {err}"
534        );
535    }
536
537    #[test]
538    fn validate_passes_for_ppo_with_every_required_hyperparameter_set() {
539        let temp = TempDir::new().unwrap();
540        write_task_json(temp.path(), "find_target", r#"{"env_name": "basic_arena"}"#);
541
542        let job = ppo_job_with_hyperparams_set();
543
544        assert!(job.validate(temp.path(), None).is_ok());
545    }
546
547    /// A `ppo` `resume-run` job restores its rollout length (`n_steps`) from the checkpoint it
548    /// resumes, so `n_steps` specifically is *not* required for that flow - matching the real,
549    /// currently-shipping `train-from-checkpoint.yml` job template (`algorithm: ppo`,
550    /// `flow_type: resume-run`, no `n_steps`). This carve-out is encoded in
551    /// `mecha10_job_validation::AlgorithmField::requirement`, which takes `flow_type` for exactly
552    /// this reason; it predates this crate as a hand-rolled `if` in the CLI
553    /// (`flow_type != ResumeRun && algorithm == "ppo" && n_steps missing`).
554    #[test]
555    fn validate_does_not_require_n_steps_for_a_ppo_resume_run() {
556        let temp = TempDir::new().unwrap();
557        write_task_json(temp.path(), "find_target", r#"{"env_name": "basic_arena"}"#);
558
559        let without_n_steps = JobDefinition {
560            flow_type: FlowType::ResumeRun,
561            n_steps: None,
562            ..ppo_job_with_hyperparams_set()
563        };
564        assert!(
565            without_n_steps.validate(temp.path(), Some("run-123")).is_ok(),
566            "a ppo resume-run job without n_steps must still validate"
567        );
568    }
569
570    /// The `n_steps` carve-out above is `n_steps`-specific: every *other* `ppo` hyperparameter is
571    /// still required for a `resume-run`. Dropping `total_timesteps` must still fail validation.
572    #[test]
573    fn validate_still_requires_total_timesteps_for_a_ppo_resume_run() {
574        let temp = TempDir::new().unwrap();
575        write_task_json(temp.path(), "find_target", r#"{"env_name": "basic_arena"}"#);
576
577        let missing_total_timesteps = JobDefinition {
578            flow_type: FlowType::ResumeRun,
579            n_steps: None,
580            total_timesteps: None,
581            ..ppo_job_with_hyperparams_set()
582        };
583        let err = missing_total_timesteps
584            .validate(temp.path(), Some("run-123"))
585            .unwrap_err();
586        assert!(
587            err.to_string().contains("total_timesteps"),
588            "expected error to mention total_timesteps, got: {err}"
589        );
590        assert!(
591            !err.to_string().contains("n_steps"),
592            "n_steps should remain exempt for a ppo resume-run, got: {err}"
593        );
594    }
595
596    #[test]
597    fn load_job_catalog_returns_empty_when_directory_missing() {
598        let temp = TempDir::new().unwrap();
599        let jobs = load_job_catalog(temp.path()).unwrap();
600        assert!(jobs.is_empty());
601    }
602
603    #[test]
604    fn load_job_catalog_parses_yaml_files_sorted_by_name() {
605        let temp = TempDir::new().unwrap();
606        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
607        std::fs::create_dir_all(&jobs_dir).unwrap();
608
609        write_job(
610            &jobs_dir,
611            "train.yml",
612            "name: train\nflow_type: train\nalgorithm: ppo\ntotal_timesteps: 1000000\n",
613        );
614        write_job(
615            &jobs_dir,
616            "hypertune.yml",
617            "name: hypertune\nflow_type: tune\nalgorithm: ppo\nn_trials: 10\n",
618        );
619
620        let jobs = load_job_catalog(temp.path()).unwrap();
621
622        assert_eq!(jobs.len(), 2);
623        assert_eq!(jobs[0].name, "hypertune");
624        assert_eq!(jobs[0].flow_type, FlowType::Tune);
625        assert_eq!(jobs[1].name, "train");
626        assert_eq!(jobs[1].flow_type, FlowType::Train);
627        assert_eq!(jobs[1].total_timesteps, Some(1000000));
628    }
629
630    #[test]
631    fn load_job_catalog_ignores_non_yaml_files() {
632        let temp = TempDir::new().unwrap();
633        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
634        std::fs::create_dir_all(&jobs_dir).unwrap();
635
636        write_job(&jobs_dir, "README.md", "not a job");
637        write_job(
638            &jobs_dir,
639            "train.yml",
640            "name: train\nflow_type: train\nalgorithm: ppo\n",
641        );
642
643        let jobs = load_job_catalog(temp.path()).unwrap();
644        assert_eq!(jobs.len(), 1);
645    }
646
647    #[test]
648    fn supports_remote_for_train_resume_run_and_tune() {
649        let base = base_job();
650        assert!(base.supports_remote());
651
652        let resume = JobDefinition {
653            flow_type: FlowType::ResumeRun,
654            ..base.clone()
655        };
656        assert!(resume.supports_remote());
657
658        let tune = JobDefinition {
659            flow_type: FlowType::Tune,
660            ..base.clone()
661        };
662        assert!(tune.supports_remote());
663
664        let tune_then_train = JobDefinition {
665            flow_type: FlowType::TuneThenTrain,
666            ..base.clone()
667        };
668        assert!(!tune_then_train.supports_remote());
669
670        let train_from_tune = JobDefinition {
671            flow_type: FlowType::TrainFromTune,
672            ..base
673        };
674        assert!(!train_from_tune.supports_remote());
675    }
676
677    #[test]
678    fn to_create_job_request_forwards_resume_run_id() {
679        let temp = TempDir::new().unwrap();
680        let jobs_dir = temp.path().join("configs/rl/jobs");
681        std::fs::create_dir_all(&jobs_dir).unwrap();
682        let source_path = jobs_dir.join("train-from-checkpoint.yml");
683        std::fs::write(&source_path, "name: train-from-checkpoint\nflow_type: resume-run\n").unwrap();
684
685        let job = JobDefinition {
686            name: "train-from-checkpoint".to_string(),
687            flow_type: FlowType::ResumeRun,
688            source_path,
689            ..base_job()
690        };
691
692        let req = job
693            .to_create_job_request(temp.path(), "my-project", "my-robot", Some("  run-abc123  "))
694            .unwrap();
695
696        assert_eq!(req.resume_run_id.as_deref(), Some("run-abc123"));
697    }
698
699    #[test]
700    fn to_create_job_request_omits_empty_resume_run_id() {
701        let temp = TempDir::new().unwrap();
702        let jobs_dir = temp.path().join("configs/rl/jobs");
703        std::fs::create_dir_all(&jobs_dir).unwrap();
704        let source_path = jobs_dir.join("train-from-checkpoint.yml");
705        std::fs::write(&source_path, "name: train-from-checkpoint\nflow_type: resume-run\n").unwrap();
706
707        let job = JobDefinition {
708            name: "train-from-checkpoint".to_string(),
709            flow_type: FlowType::ResumeRun,
710            source_path,
711            ..base_job()
712        };
713
714        let req = job
715            .to_create_job_request(temp.path(), "my-project", "my-robot", Some("   "))
716            .unwrap();
717
718        assert_eq!(req.resume_run_id, None);
719    }
720
721    /// Contract test against `mecha10-job-validation`'s `JobField` list: every field applicable
722    /// to a remotely-dispatchable flow type must have a matching wire key on this struct.
723    #[test]
724    fn create_job_request_has_a_field_for_every_dispatchable_job_field() {
725        let dispatchable = [FlowType::Train, FlowType::ResumeRun, FlowType::Tune];
726
727        let req = CreateJobRequest {
728            name: "job".to_string(),
729            description: None,
730            algorithm: "ppo".to_string(),
731            robot_id: "robot-1".to_string(),
732            project_id: None,
733            project_name: None,
734            env_name: None,
735            task_name: Some("find_target".to_string()),
736            total_timesteps: Some(1_000_000),
737            n_steps: None,
738            n_epochs: None,
739            learning_rate: None,
740            patience: None,
741            min_delta: None,
742            n_envs: None,
743            config_yaml: None,
744            resume_run_id: Some("run-123".to_string()),
745            n_trials: Some(10),
746            tune_timesteps: Some(500_000),
747            eval_eps: Some(5),
748            search_space: None,
749        };
750        let json = serde_json::to_value(&req).expect("CreateJobRequest serializes");
751
752        for field in mecha10_job_validation::all_fields() {
753            let applicable = dispatchable
754                .iter()
755                .any(|&ft| field.requirement(ft) != mecha10_job_validation::FieldRequirement::NotApplicable);
756            if applicable {
757                assert!(
758                    json.get(field.wire_name()).is_some(),
759                    "CreateJobRequest has no '{}' field for {field:?}, which is applicable to a \
760                     dispatchable flow type",
761                    field.wire_name()
762                );
763            }
764        }
765    }
766
767    #[test]
768    fn to_create_job_request_forwards_tune_fields() {
769        let temp = TempDir::new().unwrap();
770        let jobs_dir = temp.path().join("configs/rl/jobs");
771        std::fs::create_dir_all(&jobs_dir).unwrap();
772        let source_path = jobs_dir.join("hypertune.yml");
773        std::fs::write(&source_path, "name: hypertune\nflow_type: tune\n").unwrap();
774
775        let job = JobDefinition {
776            name: "hypertune".to_string(),
777            flow_type: FlowType::Tune,
778            n_trials: Some(20),
779            tune_timesteps: Some(250_000),
780            eval_eps: Some(3),
781            source_path,
782            ..base_job()
783        };
784
785        let req = job
786            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
787            .unwrap();
788
789        assert_eq!(req.n_trials, Some(20));
790        assert_eq!(req.tune_timesteps, Some(250_000));
791        assert_eq!(req.eval_eps, Some(3));
792    }
793
794    #[test]
795    fn to_create_job_request_forwards_n_envs() {
796        let temp = TempDir::new().unwrap();
797        let jobs_dir = temp.path().join("configs/rl/jobs");
798        std::fs::create_dir_all(&jobs_dir).unwrap();
799        let source_path = jobs_dir.join("train.yml");
800        std::fs::write(&source_path, "name: train\nflow_type: train\nn_envs: 8\n").unwrap();
801
802        let job = JobDefinition {
803            name: "train".to_string(),
804            n_envs: Some(8),
805            source_path,
806            ..base_job()
807        };
808
809        let req = job
810            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
811            .unwrap();
812
813        assert_eq!(req.n_envs, Some(8));
814    }
815
816    #[test]
817    fn to_create_job_request_forwards_search_space() {
818        let temp = TempDir::new().unwrap();
819        let jobs_dir = temp.path().join("configs/rl/jobs");
820        std::fs::create_dir_all(&jobs_dir).unwrap();
821        let source_path = jobs_dir.join("hypertune.yml");
822        let yaml = "name: hypertune\nflow_type: tune\nsearch_space:\n  - name: learning_rate\n    \
823                    type: float\n    low: 1e-5\n    high: 1e-3\n    log: true\n  - name: n_steps\n    \
824                    type: categorical\n    choices: [1024, 2048, 4096]\n  - name: n_epochs\n    \
825                    type: int\n    low: 5\n    high: 20\n";
826        std::fs::write(&source_path, yaml).unwrap();
827        let job: JobDefinition = serde_yaml::from_str(yaml).unwrap();
828        let job = JobDefinition { source_path, ..job };
829
830        let req = job
831            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
832            .unwrap();
833
834        let search_space = req.search_space.expect("search_space forwarded");
835        assert_eq!(search_space.len(), 3);
836        assert_eq!(search_space[0].name, "learning_rate");
837        assert_eq!(search_space[0].type_, SearchSpaceType::Float);
838        assert_eq!(search_space[0].low, Some(1e-5));
839        assert_eq!(search_space[0].high, Some(1e-3));
840        assert_eq!(search_space[0].log, Some(true));
841        assert_eq!(search_space[1].name, "n_steps");
842        assert_eq!(search_space[1].type_, SearchSpaceType::Categorical);
843        assert_eq!(
844            search_space[1].choices,
845            Some(vec![
846                serde_json::json!(1024),
847                serde_json::json!(2048),
848                serde_json::json!(4096)
849            ])
850        );
851        assert_eq!(search_space[2].name, "n_epochs");
852        assert_eq!(search_space[2].type_, SearchSpaceType::Int);
853        assert_eq!(search_space[2].low, Some(5.0));
854        assert_eq!(search_space[2].high, Some(20.0));
855    }
856
857    #[test]
858    fn to_create_job_request_omits_search_space_when_unset() {
859        let temp = TempDir::new().unwrap();
860        let jobs_dir = temp.path().join("configs/rl/jobs");
861        std::fs::create_dir_all(&jobs_dir).unwrap();
862        let source_path = jobs_dir.join("train.yml");
863        std::fs::write(&source_path, "name: train\nflow_type: train\n").unwrap();
864
865        let job = JobDefinition {
866            name: "train".to_string(),
867            source_path,
868            ..base_job()
869        };
870
871        let req = job
872            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
873            .unwrap();
874
875        assert_eq!(req.search_space, None);
876    }
877
878    #[test]
879    fn to_create_job_request_drops_sentinel_values() {
880        let temp = TempDir::new().unwrap();
881        let jobs_dir = temp.path().join("configs/rl/jobs");
882        std::fs::create_dir_all(&jobs_dir).unwrap();
883        let source_path = jobs_dir.join("train.yml");
884        std::fs::write(&source_path, "name: train\nflow_type: train\nenv_name: rover\n").unwrap();
885
886        let job = JobDefinition {
887            name: "train".to_string(),
888            description: Some("desc".to_string()),
889            env_name: Some("rover".to_string()),
890            total_timesteps: Some(-1),
891            n_steps: Some(2048),
892            n_epochs: Some(-1),
893            learning_rate: Some(-1.0),
894            patience: Some(50),
895            min_delta: Some(-1.0),
896            n_envs: Some(-1),
897            source_path,
898            ..base_job()
899        };
900
901        let req = job
902            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
903            .unwrap();
904
905        assert_eq!(req.name, "train");
906        assert_eq!(req.algorithm, "ppo");
907        assert_eq!(req.robot_id, "my-project");
908        assert_eq!(req.project_id.as_deref(), Some("my-project"));
909        assert_eq!(req.project_name.as_deref(), Some("my-robot"));
910        assert_eq!(req.total_timesteps, None);
911        assert_eq!(req.n_steps, Some(2048));
912        assert_eq!(req.n_epochs, None);
913        assert_eq!(req.learning_rate, None);
914        assert_eq!(req.patience, Some(50));
915        assert_eq!(req.min_delta, None);
916        assert_eq!(req.n_envs, None);
917        assert_eq!(
918            req.config_yaml.as_deref(),
919            Some("name: train\nflow_type: train\nenv_name: rover\n")
920        );
921        assert_eq!(req.task_name, None);
922    }
923
924    #[test]
925    fn to_create_job_request_forwards_task_name() {
926        let temp = TempDir::new().unwrap();
927        let jobs_dir = temp.path().join("configs/rl/jobs");
928        std::fs::create_dir_all(&jobs_dir).unwrap();
929        let source_path = jobs_dir.join("train.yml");
930        std::fs::write(&source_path, "name: train\nflow_type: train\n").unwrap();
931
932        let job = JobDefinition {
933            name: "train".to_string(),
934            task_name: Some("navigate".to_string()),
935            source_path,
936            ..base_job()
937        };
938
939        let req = job
940            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
941            .unwrap();
942
943        assert_eq!(req.task_name.as_deref(), Some("navigate"));
944    }
945
946    #[test]
947    fn to_create_job_request_errors_when_the_jobs_own_file_is_missing() {
948        let temp = TempDir::new().unwrap();
949        let job = JobDefinition {
950            name: "train".to_string(),
951            source_path: temp.path().join("configs/rl/jobs/train.yml"),
952            ..base_job()
953        };
954
955        let err = job
956            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
957            .unwrap_err();
958        assert!(err.to_string().contains("configs/rl/jobs/train.yml"));
959    }
960}
mecha10_cli/services/train_jobs.rs

mecha10_cli/services/
train_jobs.rs
