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            total_timesteps: non_sentinel_i64(self.total_timesteps),
214            resume_run_id: resume_run_id.map(str::to_string),
215            tune_run_id: None,
216            n_trials: non_sentinel_i64(self.n_trials),
217            tune_timesteps: non_sentinel_i64(self.tune_timesteps),
218            eval_eps: non_sentinel_i64(self.eval_eps),
219        }
220    }
221
222    /// Validate this job against `mecha10-job-validation`'s flow_type requirement table, plus a
223    /// `task_name`-exists check against the project's own task catalog.
224    pub fn validate(&self, project_root: &Path, resume_run_id: Option<&str>) -> Result<()> {
225        let missing = mecha10_job_validation::validate_job(self.flow_type, &self.validation_fields(resume_run_id));
226        if !missing.is_empty() {
227            let details = missing.iter().map(|m| m.to_string()).collect::<Vec<_>>().join("; ");
228            bail!("job '{}' is missing required fields: {details}", self.name);
229        }
230
231        let mut missing_hyperparams = Vec::new();
232        if non_sentinel_i64(self.total_timesteps).is_none() {
233            missing_hyperparams.push("total_timesteps");
234        }
235        if self.flow_type != FlowType::ResumeRun && non_sentinel_i64(self.n_steps).is_none() {
236            missing_hyperparams.push("n_steps");
237        }
238        if !missing_hyperparams.is_empty() {
239            bail!(
240                "job '{}' is missing required field(s): {}",
241                self.name,
242                missing_hyperparams.join(", ")
243            );
244        }
245
246        if let Some(task_name) = self.task_name.as_deref() {
247            mecha10_job_validation::task_exists(project_root, task_name)
248                .with_context(|| format!("job '{}' has an invalid task_name", self.name))?;
249        }
250
251        Ok(())
252    }
253
254    /// Resolve which `mecha10_rl.envs` registry name this job targets: an explicit `env_name`
255    /// wins, else it's read from `task_name`'s own `task.json`. `None` when neither is available.
256    pub fn resolve_env_name(&self, project_root: &Path) -> Option<String> {
257        self.env_name.clone().or_else(|| {
258            self.task_name
259                .as_deref()
260                .and_then(|task_name| env_name_from_task_json(project_root, task_name))
261        })
262    }
263
264    /// Whether this job can be dispatched remotely with today's `CreateJobRequest`/`create_job`
265    /// contract. `tune-then-train`/`train-from-tune` aren't dispatchable yet and stay local-only.
266    pub fn supports_remote(&self) -> bool {
267        matches!(self.flow_type, FlowType::Train | FlowType::ResumeRun | FlowType::Tune)
268    }
269
270    /// Build the remote dispatch request body. Only meaningful when `supports_remote()` is true;
271    /// callers should check that first.
272    pub fn to_create_job_request(
273        &self,
274        project_root: &Path,
275        robot_id: &str,
276        project_name: &str,
277        resume_run_id: Option<&str>,
278    ) -> Result<CreateJobRequest> {
279        let config_yaml = std::fs::read_to_string(&self.source_path).with_context(|| {
280            format!(
281                "failed to read {} - this job definition's own file must exist before dispatching remotely",
282                self.source_path.display()
283            )
284        })?;
285
286        Ok(CreateJobRequest {
287            name: self.name.clone(),
288            description: self.description.clone(),
289            algorithm: self.algorithm.clone().unwrap_or_default(),
290            robot_id: robot_id.to_string(),
291            project_id: Some(robot_id.to_string()),
292            project_name: Some(project_name.to_string()),
293            env_name: self.resolve_env_name(project_root),
294            task_name: self.task_name.clone(),
295            total_timesteps: non_sentinel_i64(self.total_timesteps),
296            n_steps: non_sentinel_i64(self.n_steps),
297            n_epochs: non_sentinel_i64(self.n_epochs),
298            learning_rate: non_sentinel_f64(self.learning_rate),
299            patience: non_sentinel_i64(self.patience),
300            min_delta: non_sentinel_f64(self.min_delta),
301            n_envs: non_sentinel_i64(self.n_envs),
302            config_yaml: Some(config_yaml),
303            resume_run_id: resume_run_id
304                .map(str::trim)
305                .filter(|id| !id.is_empty())
306                .map(str::to_string),
307            n_trials: non_sentinel_i64(self.n_trials),
308            tune_timesteps: non_sentinel_i64(self.tune_timesteps),
309            eval_eps: non_sentinel_i64(self.eval_eps),
310            search_space: self.search_space.clone(),
311        })
312    }
313}
314
315/// Load every `.yml`/`.yaml` job definition under `{project_root}/configs/rl/jobs/`, sorted by
316/// file name for a stable menu order. Returns an empty list (not an error) when the directory
317/// doesn't exist - `mecha10 train` treats that as "no jobs configured yet" and points the user at
318/// `mecha10 init`/the bundled defaults rather than failing.
319pub fn load_job_catalog(project_root: &Path) -> Result<Vec<JobDefinition>> {
320    let jobs_dir = project_root.join(JOB_CATALOG_DIR);
321    if !jobs_dir.exists() {
322        return Ok(Vec::new());
323    }
324
325    let mut paths: Vec<PathBuf> = std::fs::read_dir(&jobs_dir)
326        .with_context(|| format!("failed to read {}", jobs_dir.display()))?
327        .filter_map(|entry| entry.ok())
328        .map(|entry| entry.path())
329        .filter(|path| {
330            path.extension()
331                .and_then(|ext| ext.to_str())
332                .is_some_and(|ext| ext == "yml" || ext == "yaml")
333        })
334        .collect();
335    paths.sort();
336
337    let mut jobs = Vec::with_capacity(paths.len());
338    for path in paths {
339        let content = std::fs::read_to_string(&path).with_context(|| format!("failed to read {}", path.display()))?;
340        let mut job: JobDefinition = serde_yaml::from_str(&content)
341            .with_context(|| format!("failed to parse job definition {}", path.display()))?;
342        job.source_path = path;
343        jobs.push(job);
344    }
345
346    Ok(jobs)
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352    use tempfile::TempDir;
353
354    fn write_job(dir: &Path, file_name: &str, contents: &str) {
355        std::fs::write(dir.join(file_name), contents).unwrap();
356    }
357
358    fn write_task_json(project_root: &Path, task_name: &str, contents: &str) {
359        let task_dir = project_root
360            .join(mecha10_cli_core::paths::project::RL_TASKS_DIR)
361            .join(task_name);
362        std::fs::create_dir_all(&task_dir).unwrap();
363        std::fs::write(task_dir.join("task.json"), contents).unwrap();
364    }
365
366    #[test]
367    fn resolve_env_name_prefers_an_explicit_env_name_over_the_tasks_own() {
368        let dir = TempDir::new().unwrap();
369        write_task_json(dir.path(), "find_target", r#"{"env_name": "basic_arena"}"#);
370
371        let job = JobDefinition {
372            env_name: Some("rover_arena".to_string()),
373            task_name: Some("find_target".to_string()),
374            ..base_job()
375        };
376        assert_eq!(job.resolve_env_name(dir.path()), Some("rover_arena".to_string()));
377    }
378
379    #[test]
380    fn resolve_env_name_falls_back_to_the_tasks_own_task_json() {
381        let dir = TempDir::new().unwrap();
382        write_task_json(
383            dir.path(),
384            "find_target",
385            r#"{"env_name": "basic_arena", "model": "humanoid"}"#,
386        );
387
388        let job = JobDefinition {
389            task_name: Some("find_target".to_string()),
390            ..base_job()
391        };
392        assert_eq!(job.resolve_env_name(dir.path()), Some("basic_arena".to_string()));
393    }
394
395    #[test]
396    fn resolve_env_name_is_none_when_task_has_no_env_name_field() {
397        let dir = TempDir::new().unwrap();
398        write_task_json(dir.path(), "find_target", r#"{"model": "humanoid"}"#);
399
400        let job = JobDefinition {
401            task_name: Some("find_target".to_string()),
402            ..base_job()
403        };
404        assert_eq!(job.resolve_env_name(dir.path()), None);
405    }
406
407    #[test]
408    fn resolve_env_name_is_none_without_an_env_name_or_a_task_name() {
409        let dir = TempDir::new().unwrap();
410        let job = base_job();
411        assert_eq!(job.resolve_env_name(dir.path()), None);
412    }
413
414    /// A minimal, fully-`None`/default `JobDefinition` for tests to `..base_job()` from, so each
415    /// test only spells out the fields it actually cares about.
416    fn base_job() -> JobDefinition {
417        JobDefinition {
418            name: "job".to_string(),
419            description: None,
420            flow_type: FlowType::Train,
421            env_name: None,
422            task_name: None,
423            algorithm: Some("ppo".to_string()),
424            policy: None,
425            seed: None,
426            environment_map: None,
427            total_timesteps: None,
428            learning_rate: None,
429            linear_lr_schedule: None,
430            n_steps: None,
431            batch_size: None,
432            n_epochs: None,
433            gamma: None,
434            gae_lambda: None,
435            clip_range: None,
436            ent_coef: None,
437            vf_coef: None,
438            max_grad_norm: None,
439            patience: None,
440            min_delta: None,
441            n_envs: None,
442            n_trials: None,
443            tune_timesteps: None,
444            eval_eps: None,
445            search_space: None,
446            source_path: PathBuf::new(),
447        }
448    }
449
450    #[test]
451    fn load_job_catalog_returns_empty_when_directory_missing() {
452        let temp = TempDir::new().unwrap();
453        let jobs = load_job_catalog(temp.path()).unwrap();
454        assert!(jobs.is_empty());
455    }
456
457    #[test]
458    fn load_job_catalog_parses_yaml_files_sorted_by_name() {
459        let temp = TempDir::new().unwrap();
460        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
461        std::fs::create_dir_all(&jobs_dir).unwrap();
462
463        write_job(
464            &jobs_dir,
465            "train.yml",
466            "name: train\nflow_type: train\nalgorithm: ppo\ntotal_timesteps: 1000000\n",
467        );
468        write_job(
469            &jobs_dir,
470            "hypertune.yml",
471            "name: hypertune\nflow_type: tune\nalgorithm: ppo\nn_trials: 10\n",
472        );
473
474        let jobs = load_job_catalog(temp.path()).unwrap();
475
476        assert_eq!(jobs.len(), 2);
477        assert_eq!(jobs[0].name, "hypertune");
478        assert_eq!(jobs[0].flow_type, FlowType::Tune);
479        assert_eq!(jobs[1].name, "train");
480        assert_eq!(jobs[1].flow_type, FlowType::Train);
481        assert_eq!(jobs[1].total_timesteps, Some(1000000));
482    }
483
484    #[test]
485    fn load_job_catalog_ignores_non_yaml_files() {
486        let temp = TempDir::new().unwrap();
487        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
488        std::fs::create_dir_all(&jobs_dir).unwrap();
489
490        write_job(&jobs_dir, "README.md", "not a job");
491        write_job(
492            &jobs_dir,
493            "train.yml",
494            "name: train\nflow_type: train\nalgorithm: ppo\n",
495        );
496
497        let jobs = load_job_catalog(temp.path()).unwrap();
498        assert_eq!(jobs.len(), 1);
499    }
500
501    #[test]
502    fn supports_remote_for_train_resume_run_and_tune() {
503        let base = base_job();
504        assert!(base.supports_remote());
505
506        let resume = JobDefinition {
507            flow_type: FlowType::ResumeRun,
508            ..base.clone()
509        };
510        assert!(resume.supports_remote());
511
512        let tune = JobDefinition {
513            flow_type: FlowType::Tune,
514            ..base.clone()
515        };
516        assert!(tune.supports_remote());
517
518        let tune_then_train = JobDefinition {
519            flow_type: FlowType::TuneThenTrain,
520            ..base.clone()
521        };
522        assert!(!tune_then_train.supports_remote());
523
524        let train_from_tune = JobDefinition {
525            flow_type: FlowType::TrainFromTune,
526            ..base
527        };
528        assert!(!train_from_tune.supports_remote());
529    }
530
531    #[test]
532    fn to_create_job_request_forwards_resume_run_id() {
533        let temp = TempDir::new().unwrap();
534        let jobs_dir = temp.path().join("configs/rl/jobs");
535        std::fs::create_dir_all(&jobs_dir).unwrap();
536        let source_path = jobs_dir.join("train-from-checkpoint.yml");
537        std::fs::write(&source_path, "name: train-from-checkpoint\nflow_type: resume-run\n").unwrap();
538
539        let job = JobDefinition {
540            name: "train-from-checkpoint".to_string(),
541            flow_type: FlowType::ResumeRun,
542            source_path,
543            ..base_job()
544        };
545
546        let req = job
547            .to_create_job_request(temp.path(), "my-project", "my-robot", Some("  run-abc123  "))
548            .unwrap();
549
550        assert_eq!(req.resume_run_id.as_deref(), Some("run-abc123"));
551    }
552
553    #[test]
554    fn to_create_job_request_omits_empty_resume_run_id() {
555        let temp = TempDir::new().unwrap();
556        let jobs_dir = temp.path().join("configs/rl/jobs");
557        std::fs::create_dir_all(&jobs_dir).unwrap();
558        let source_path = jobs_dir.join("train-from-checkpoint.yml");
559        std::fs::write(&source_path, "name: train-from-checkpoint\nflow_type: resume-run\n").unwrap();
560
561        let job = JobDefinition {
562            name: "train-from-checkpoint".to_string(),
563            flow_type: FlowType::ResumeRun,
564            source_path,
565            ..base_job()
566        };
567
568        let req = job
569            .to_create_job_request(temp.path(), "my-project", "my-robot", Some("   "))
570            .unwrap();
571
572        assert_eq!(req.resume_run_id, None);
573    }
574
575    /// Contract test against `mecha10-job-validation`'s `JobField` list: every field applicable
576    /// to a remotely-dispatchable flow type must have a matching wire key on this struct.
577    #[test]
578    fn create_job_request_has_a_field_for_every_dispatchable_job_field() {
579        let dispatchable = [FlowType::Train, FlowType::ResumeRun, FlowType::Tune];
580
581        let req = CreateJobRequest {
582            name: "job".to_string(),
583            description: None,
584            algorithm: "ppo".to_string(),
585            robot_id: "robot-1".to_string(),
586            project_id: None,
587            project_name: None,
588            env_name: None,
589            task_name: Some("find_target".to_string()),
590            total_timesteps: Some(1_000_000),
591            n_steps: None,
592            n_epochs: None,
593            learning_rate: None,
594            patience: None,
595            min_delta: None,
596            n_envs: None,
597            config_yaml: None,
598            resume_run_id: Some("run-123".to_string()),
599            n_trials: Some(10),
600            tune_timesteps: Some(500_000),
601            eval_eps: Some(5),
602            search_space: None,
603        };
604        let json = serde_json::to_value(&req).expect("CreateJobRequest serializes");
605
606        for field in mecha10_job_validation::all_fields() {
607            let applicable = dispatchable
608                .iter()
609                .any(|&ft| field.requirement(ft) != mecha10_job_validation::FieldRequirement::NotApplicable);
610            if applicable {
611                assert!(
612                    json.get(field.wire_name()).is_some(),
613                    "CreateJobRequest has no '{}' field for {field:?}, which is applicable to a \
614                     dispatchable flow type",
615                    field.wire_name()
616                );
617            }
618        }
619    }
620
621    #[test]
622    fn to_create_job_request_forwards_tune_fields() {
623        let temp = TempDir::new().unwrap();
624        let jobs_dir = temp.path().join("configs/rl/jobs");
625        std::fs::create_dir_all(&jobs_dir).unwrap();
626        let source_path = jobs_dir.join("hypertune.yml");
627        std::fs::write(&source_path, "name: hypertune\nflow_type: tune\n").unwrap();
628
629        let job = JobDefinition {
630            name: "hypertune".to_string(),
631            flow_type: FlowType::Tune,
632            n_trials: Some(20),
633            tune_timesteps: Some(250_000),
634            eval_eps: Some(3),
635            source_path,
636            ..base_job()
637        };
638
639        let req = job
640            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
641            .unwrap();
642
643        assert_eq!(req.n_trials, Some(20));
644        assert_eq!(req.tune_timesteps, Some(250_000));
645        assert_eq!(req.eval_eps, Some(3));
646    }
647
648    #[test]
649    fn to_create_job_request_forwards_n_envs() {
650        let temp = TempDir::new().unwrap();
651        let jobs_dir = temp.path().join("configs/rl/jobs");
652        std::fs::create_dir_all(&jobs_dir).unwrap();
653        let source_path = jobs_dir.join("train.yml");
654        std::fs::write(&source_path, "name: train\nflow_type: train\nn_envs: 8\n").unwrap();
655
656        let job = JobDefinition {
657            name: "train".to_string(),
658            n_envs: Some(8),
659            source_path,
660            ..base_job()
661        };
662
663        let req = job
664            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
665            .unwrap();
666
667        assert_eq!(req.n_envs, Some(8));
668    }
669
670    #[test]
671    fn to_create_job_request_forwards_search_space() {
672        let temp = TempDir::new().unwrap();
673        let jobs_dir = temp.path().join("configs/rl/jobs");
674        std::fs::create_dir_all(&jobs_dir).unwrap();
675        let source_path = jobs_dir.join("hypertune.yml");
676        let yaml = "name: hypertune\nflow_type: tune\nsearch_space:\n  - name: learning_rate\n    \
677                    type: float\n    low: 1e-5\n    high: 1e-3\n    log: true\n  - name: n_steps\n    \
678                    type: categorical\n    choices: [1024, 2048, 4096]\n  - name: n_epochs\n    \
679                    type: int\n    low: 5\n    high: 20\n";
680        std::fs::write(&source_path, yaml).unwrap();
681        let job: JobDefinition = serde_yaml::from_str(yaml).unwrap();
682        let job = JobDefinition { source_path, ..job };
683
684        let req = job
685            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
686            .unwrap();
687
688        let search_space = req.search_space.expect("search_space forwarded");
689        assert_eq!(search_space.len(), 3);
690        assert_eq!(search_space[0].name, "learning_rate");
691        assert_eq!(search_space[0].type_, SearchSpaceType::Float);
692        assert_eq!(search_space[0].low, Some(1e-5));
693        assert_eq!(search_space[0].high, Some(1e-3));
694        assert_eq!(search_space[0].log, Some(true));
695        assert_eq!(search_space[1].name, "n_steps");
696        assert_eq!(search_space[1].type_, SearchSpaceType::Categorical);
697        assert_eq!(
698            search_space[1].choices,
699            Some(vec![
700                serde_json::json!(1024),
701                serde_json::json!(2048),
702                serde_json::json!(4096)
703            ])
704        );
705        assert_eq!(search_space[2].name, "n_epochs");
706        assert_eq!(search_space[2].type_, SearchSpaceType::Int);
707        assert_eq!(search_space[2].low, Some(5.0));
708        assert_eq!(search_space[2].high, Some(20.0));
709    }
710
711    #[test]
712    fn to_create_job_request_omits_search_space_when_unset() {
713        let temp = TempDir::new().unwrap();
714        let jobs_dir = temp.path().join("configs/rl/jobs");
715        std::fs::create_dir_all(&jobs_dir).unwrap();
716        let source_path = jobs_dir.join("train.yml");
717        std::fs::write(&source_path, "name: train\nflow_type: train\n").unwrap();
718
719        let job = JobDefinition {
720            name: "train".to_string(),
721            source_path,
722            ..base_job()
723        };
724
725        let req = job
726            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
727            .unwrap();
728
729        assert_eq!(req.search_space, None);
730    }
731
732    #[test]
733    fn to_create_job_request_drops_sentinel_values() {
734        let temp = TempDir::new().unwrap();
735        let jobs_dir = temp.path().join("configs/rl/jobs");
736        std::fs::create_dir_all(&jobs_dir).unwrap();
737        let source_path = jobs_dir.join("train.yml");
738        std::fs::write(&source_path, "name: train\nflow_type: train\nenv_name: rover\n").unwrap();
739
740        let job = JobDefinition {
741            name: "train".to_string(),
742            description: Some("desc".to_string()),
743            env_name: Some("rover".to_string()),
744            total_timesteps: Some(-1),
745            n_steps: Some(2048),
746            n_epochs: Some(-1),
747            learning_rate: Some(-1.0),
748            patience: Some(50),
749            min_delta: Some(-1.0),
750            n_envs: Some(-1),
751            source_path,
752            ..base_job()
753        };
754
755        let req = job
756            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
757            .unwrap();
758
759        assert_eq!(req.name, "train");
760        assert_eq!(req.algorithm, "ppo");
761        assert_eq!(req.robot_id, "my-project");
762        assert_eq!(req.project_id.as_deref(), Some("my-project"));
763        assert_eq!(req.project_name.as_deref(), Some("my-robot"));
764        assert_eq!(req.total_timesteps, None);
765        assert_eq!(req.n_steps, Some(2048));
766        assert_eq!(req.n_epochs, None);
767        assert_eq!(req.learning_rate, None);
768        assert_eq!(req.patience, Some(50));
769        assert_eq!(req.min_delta, None);
770        assert_eq!(req.n_envs, None);
771        assert_eq!(
772            req.config_yaml.as_deref(),
773            Some("name: train\nflow_type: train\nenv_name: rover\n")
774        );
775        assert_eq!(req.task_name, None);
776    }
777
778    #[test]
779    fn to_create_job_request_forwards_task_name() {
780        let temp = TempDir::new().unwrap();
781        let jobs_dir = temp.path().join("configs/rl/jobs");
782        std::fs::create_dir_all(&jobs_dir).unwrap();
783        let source_path = jobs_dir.join("train.yml");
784        std::fs::write(&source_path, "name: train\nflow_type: train\n").unwrap();
785
786        let job = JobDefinition {
787            name: "train".to_string(),
788            task_name: Some("navigate".to_string()),
789            source_path,
790            ..base_job()
791        };
792
793        let req = job
794            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
795            .unwrap();
796
797        assert_eq!(req.task_name.as_deref(), Some("navigate"));
798    }
799
800    #[test]
801    fn to_create_job_request_errors_when_the_jobs_own_file_is_missing() {
802        let temp = TempDir::new().unwrap();
803        let job = JobDefinition {
804            name: "train".to_string(),
805            source_path: temp.path().join("configs/rl/jobs/train.yml"),
806            ..base_job()
807        };
808
809        let err = job
810            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
811            .unwrap_err();
812        assert!(err.to_string().contains("configs/rl/jobs/train.yml"));
813    }
814}
mecha10_cli/services/train_jobs.rs

mecha10_cli/services/
train_jobs.rs
