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
5//! selection. `mecha10-templates`'s RL bundle (`create_rl_training_bundle`) ships a standard set
6//! of these into every scaffolded project - a project owner extends the menu simply by adding
7//! more `.yml` files to that directory, no mecha10 core changes needed, same extensibility
8//! philosophy as the `mecha10_rl.envs`/`mecha10_rl.algorithms` registry
9//! (`rl_task_runner.registry`).
10//!
11//! LAB-2970: a job file is now a *complete* training config, not just job identity/overrides
12//! layered onto a separate `configs/rl/local.yml`. The old two-file design (this file's
13//! `timesteps`/`lr`/`epochs` overrides, sentinel-merged at runtime onto `local.yml`'s full
14//! `PPOConfig`) had two names for the same hyperparameter depending on which file you were
15//! reading; that overlap is gone - `JobDefinition` now carries every `PPOConfig` field directly,
16//! using `PPOConfig`'s own field names (`total_timesteps`, `learning_rate`, `n_epochs`, ...), and
17//! `local.yml` no longer exists. `mecha10 train`'s remote dispatch already forwarded the raw
18//! config file's content as `config_yaml`/`CONFIG_YAML_B64` for the trainer container to run
19//! with (see `to_create_job_request`) - that transport is unchanged, it now just reads the job's
20//! own file instead of a separate `local.yml`.
21//!
22//! Field names/sentinel conventions (`-1`/`-1.0` meaning "use this field's built-in default")
23//! deliberately mirror `flows/ml-pipeline/train.yml`'s `inputs:` and
24//! `mecha10-api/src/jobs/types.rs::CreateJobRequest`, so a job definition maps 1:1 onto either
25//! dispatch path without translation surprises.
26
27use anyhow::{bail, Context, Result};
28use serde::{Deserialize, Serialize};
29use std::path::{Path, PathBuf};
30
31/// Directory (relative to a project's root) job definitions are loaded from.
32pub const JOB_CATALOG_DIR: &str = "configs/rl/jobs";
33
34/// Which Kestra `flows/ml-pipeline/*.yml` flow (and `rl_task_runner.run` mode) a job maps to.
35/// Re-exported from `mecha10-job-validation` rather than defined here - that crate is the shared
36/// source of truth for flow_type-keyed validation rules (used by both this CLI and
37/// `mecha10-api`), so this and its `Display`/kebab-case (de)serialization must be the one enum,
38/// not two independently-maintained copies.
39pub use mecha10_job_validation::FlowType;
40
41trait FlowTypeExt {
42    fn run_mode(self) -> &'static str;
43}
44
45impl FlowTypeExt for FlowType {
46    /// `rl_task_runner.run`'s mode positional arg (`train` or `tune`) for this flow type -
47    /// `tune`/`tune-then-train` dispatch through `run.py`'s `tune` mode, everything else trains.
48    fn run_mode(self) -> &'static str {
49        match self {
50            FlowType::Tune | FlowType::TuneThenTrain => "tune",
51            FlowType::Train | FlowType::TrainFromTune | FlowType::ResumeRun => "train",
52        }
53    }
54}
55
56/// A training job definition, loaded from `configs/rl/jobs/<name>.yml`.
57#[derive(Debug, Clone, Deserialize)]
58pub struct JobDefinition {
59    pub name: String,
60    #[serde(default)]
61    pub description: Option<String>,
62    pub flow_type: FlowType,
63    /// `mecha10_rl.envs` registry name. Empty/omitted = resolved from `task_name`'s own
64    /// `task.json` (its `env_name` field - see `resolve_env_name`). A job definition normally
65    /// leaves this unset and relies on `task_name` alone - see `train.yml.template`.
66    #[serde(default)]
67    pub env_name: Option<String>,
68    /// Task-curriculum name (e.g. `find_target`, `task1`) - resolves via `AssetResolver` on the
69    /// Python side (`rl_task_runner.run`'s `TASK_NAME`/`task_name` config key) into the env
70    /// factory's own `task_name` kwarg, and also selects `env_name` when that's left unset (see
71    /// `resolve_env_name`).
72    #[serde(default)]
73    pub task_name: Option<String>,
74    /// `mecha10_rl.algorithms` registry name. Empty/omitted = use this job's own `algorithm_name`
75    /// default for local dispatch - but note remote dispatch requires this to be set (see
76    /// `to_create_job_request`'s doc comment).
77    #[serde(default)]
78    pub algorithm: Option<String>,
79
80    // -- PPO hyperparameters (`rl_task_runner.algorithms.ppo.PPOConfig`, LAB-2970) -------------
81    // Field names mirror `PPOConfig` directly (not CLI-only aliases) now that a job file is the
82    // whole training config, not overrides layered onto a separate `configs/rl/local.yml`.
83    /// SB3 policy class, e.g. `MlpPolicy`/`MultiInputPolicy`. Empty/omitted = `PPOConfig`
84    /// auto-detects from the env's observation space.
85    #[serde(default)]
86    pub policy: Option<String>,
87    #[serde(default)]
88    pub seed: Option<i64>,
89    #[serde(default)]
90    pub environment_map: Option<String>,
91    #[serde(default)]
92    pub total_timesteps: Option<i64>,
93    #[serde(default)]
94    pub learning_rate: Option<f64>,
95    #[serde(default)]
96    pub linear_lr_schedule: Option<bool>,
97    #[serde(default)]
98    pub n_steps: Option<i64>,
99    #[serde(default)]
100    pub batch_size: Option<i64>,
101    #[serde(default)]
102    pub n_epochs: Option<i64>,
103    #[serde(default)]
104    pub gamma: Option<f64>,
105    #[serde(default)]
106    pub gae_lambda: Option<f64>,
107    #[serde(default)]
108    pub clip_range: Option<f64>,
109    #[serde(default)]
110    pub ent_coef: Option<f64>,
111    #[serde(default)]
112    pub vf_coef: Option<f64>,
113    #[serde(default)]
114    pub max_grad_norm: Option<f64>,
115    /// Early-stopping knobs - unrelated to `PPOConfig` proper, but kept alongside the other
116    /// hyperparameters since they're set the same way.
117    #[serde(default)]
118    pub patience: Option<i64>,
119    #[serde(default)]
120    pub min_delta: Option<f64>,
121
122    /// MLflow run ID to resume from (`resume-run` flow type). See `rl_task_runner.run`'s
123    /// warm-start resolution.
124    #[serde(default)]
125    pub resume_run_id: Option<String>,
126    /// MLflow run ID of a prior Optuna tune (`train-from-tune` flow type). See
127    /// `rl_task_runner.run`'s best-params resolution.
128    #[serde(default)]
129    pub tune_run_id: Option<String>,
130    #[serde(default)]
131    pub n_trials: Option<i64>,
132    #[serde(default)]
133    pub tune_timesteps: Option<i64>,
134    #[serde(default)]
135    pub eval_eps: Option<i64>,
136
137    /// Not part of the YAML - filled in by `load_job_catalog` with the file it was loaded from.
138    #[serde(skip)]
139    pub source_path: PathBuf,
140}
141
142/// Request body for `POST {control_plane_url}/api/model-deployments/jobs` - field-for-field the
143/// same shape as `mecha10-api/src/jobs/types.rs::CreateJobRequest` (and, transitively,
144/// `packages/dashboard`'s `types/ml-training.ts` `CreateJobRequest`/`buildCreateJobRequest`,
145/// PR #278/LAB-1431). Deliberately duplicated rather than shared via a crate dependency - the
146/// CLI and API are different deployables and this struct is the wire contract, not shared code.
147#[derive(Debug, Clone, Serialize)]
148pub struct CreateJobRequest {
149    pub name: String,
150    #[serde(skip_serializing_if = "Option::is_none")]
151    pub description: Option<String>,
152    pub algorithm: String,
153    /// FK to the `robots` entity (LAB-2724) - required server-side since `Job` got a real
154    /// `robot_id` column; omitting it fails job creation with a 422 "missing field robot_id".
155    pub robot_id: String,
156    #[serde(skip_serializing_if = "Option::is_none")]
157    pub project_id: Option<String>,
158    /// The project's real name (`mecha10.json`'s `name` field), forwarded to `train.yml`'s
159    /// `project_name` input so `entrypoint.sh` uses it for `mecha10-launcher start
160    /// --project-name` instead of a hardcoded placeholder - mirrors `mecha10 run remote`'s
161    /// `MECHA10_PROJECT_NAME` convention. Distinct from `project_id` above (an MLflow
162    /// experiment-prefix concern, not a robot identity concern).
163    #[serde(skip_serializing_if = "Option::is_none")]
164    pub project_name: Option<String>,
165    #[serde(skip_serializing_if = "Option::is_none")]
166    pub env_name: Option<String>,
167    #[serde(skip_serializing_if = "Option::is_none")]
168    pub task_name: Option<String>,
169    #[serde(skip_serializing_if = "Option::is_none")]
170    pub total_timesteps: Option<i64>,
171    #[serde(skip_serializing_if = "Option::is_none")]
172    pub n_steps: Option<i64>,
173    #[serde(skip_serializing_if = "Option::is_none")]
174    pub n_epochs: Option<i64>,
175    #[serde(skip_serializing_if = "Option::is_none")]
176    pub learning_rate: Option<f64>,
177    #[serde(skip_serializing_if = "Option::is_none")]
178    pub patience: Option<i64>,
179    #[serde(skip_serializing_if = "Option::is_none")]
180    pub min_delta: Option<f64>,
181    /// Raw content of the job's own `configs/rl/jobs/<name>.yml` file, read off local disk right
182    /// here at job-creation time (see `to_create_job_request`) and forwarded by mecha10-api into
183    /// `train.yml`'s `config_yaml_b64` input - the actual training config the trainer container
184    /// runs with. This replaces relying on the trainer container re-pulling that file at runtime
185    /// via mecha10-launcher, which raced the container's own startup and silently no-op'd on
186    /// failure (mecha10-rl's `rl-trainer/entrypoint.sh` doc comment). LAB-2970: the job file
187    /// itself is now the whole training config (no more separate `configs/rl/local.yml`), so this
188    /// is simply the job file's own bytes, unchanged in shape/semantics from before.
189    #[serde(skip_serializing_if = "Option::is_none")]
190    pub config_yaml: Option<String>,
191    /// MLflow run ID to resume from - forwarded to mecha10-api's `CreateJobRequest`, which
192    /// dispatches `resume-run.yml` instead of `train.yml` when this is non-empty (see
193    /// `mecha10-api/src/jobs/handlers.rs::resume_flow_inputs`).
194    #[serde(skip_serializing_if = "Option::is_none")]
195    pub resume_run_id: Option<String>,
196}
197
198/// Read `configs/rl/tasks/<task_name>/task.json`'s `env_name` field, relative to `project_root`.
199/// Returns `None` on any failure (missing file, invalid JSON, missing/non-string field) - callers
200/// treat that the same as "no env_name available from the task", never a hard error: a job
201/// definition can still set `env_name` explicitly instead.
202fn env_name_from_task_json(project_root: &Path, task_name: &str) -> Option<String> {
203    let path = project_root
204        .join(mecha10_cli_core::paths::project::RL_TASKS_DIR)
205        .join(task_name)
206        .join("task.json");
207    let content = std::fs::read_to_string(path).ok()?;
208    let value: serde_json::Value = serde_json::from_str(&content).ok()?;
209    value.get("env_name")?.as_str().map(str::to_string)
210}
211
212/// Drop `value` if it equals `sentinel` - the "unset" placeholder some YAML/env fields use
213/// instead of a real `null` (e.g. `-1` / `-1.0`).
214fn non_sentinel<T: PartialEq>(value: Option<T>, sentinel: T) -> Option<T> {
215    value.filter(|v| *v != sentinel)
216}
217
218fn non_sentinel_i64(value: Option<i64>) -> Option<i64> {
219    non_sentinel(value, -1)
220}
221
222fn non_sentinel_f64(value: Option<f64>) -> Option<f64> {
223    non_sentinel(value, -1.0)
224}
225
226impl JobDefinition {
227    /// Env vars for `python3 -m rl_task_runner.run <mode>` - the local dispatch path. Field names
228    /// are passed through uppercased, unrenamed (LAB-2970 merged the job schema onto `PPOConfig`'s
229    /// own names, so there's no more `TIMESTEPS`->`TOTAL_TIMESTEPS`-style translation to apply).
230    /// `CONFIG`/`RL_CONFIG_DIR` point at this job's own file (`configs/rl/jobs/<name>.yml`) rather
231    /// than a separate `configs/rl/local.yml` - same "which file to load" mechanism, new source,
232    /// same as `to_create_job_request`'s `config_yaml`. `MLFLOW_ENABLED`/`TRACKER_BACKEND=relay`/
233    /// `RELAY_API_KEY` route metrics through mecha10-api's relay
234    /// (`rl_task_runner.services.relay_tracker.RelayTracker`) instead of a direct MLflow
235    /// connection - no MLflow credentials needed on this machine.
236    pub fn local_env_vars(
237        &self,
238        project_root: &Path,
239        control_plane_url: &str,
240        relay_api_key: &str,
241        log_level: &str,
242    ) -> Vec<(String, String)> {
243        let job_stem = self
244            .source_path
245            .file_stem()
246            .and_then(|s| s.to_str())
247            .map(str::to_string)
248            .unwrap_or_else(|| self.name.clone());
249
250        vec![
251            ("RL_CONFIG_DIR".to_string(), "configs/rl/jobs".to_string()),
252            ("CONFIG".to_string(), job_stem),
253            (
254                "ENV_NAME".to_string(),
255                self.resolve_env_name(project_root).unwrap_or_default(),
256            ),
257            ("TASK_NAME".to_string(), self.task_name.clone().unwrap_or_default()),
258            ("ALGORITHM_NAME".to_string(), self.algorithm.clone().unwrap_or_default()),
259            ("POLICY".to_string(), self.policy.clone().unwrap_or_default()),
260            ("SEED".to_string(), self.seed.unwrap_or(-1).to_string()),
261            (
262                "ENVIRONMENT_MAP".to_string(),
263                self.environment_map.clone().unwrap_or_default(),
264            ),
265            (
266                "TOTAL_TIMESTEPS".to_string(),
267                self.total_timesteps.unwrap_or(-1).to_string(),
268            ),
269            (
270                "LEARNING_RATE".to_string(),
271                self.learning_rate.unwrap_or(-1.0).to_string(),
272            ),
273            (
274                "LINEAR_LR_SCHEDULE".to_string(),
275                self.linear_lr_schedule
276                    .map(|v| v.to_string())
277                    .unwrap_or_else(|| "-1".to_string()),
278            ),
279            ("N_STEPS".to_string(), self.n_steps.unwrap_or(-1).to_string()),
280            ("BATCH_SIZE".to_string(), self.batch_size.unwrap_or(-1).to_string()),
281            ("N_EPOCHS".to_string(), self.n_epochs.unwrap_or(-1).to_string()),
282            ("GAMMA".to_string(), self.gamma.unwrap_or(-1.0).to_string()),
283            ("GAE_LAMBDA".to_string(), self.gae_lambda.unwrap_or(-1.0).to_string()),
284            ("CLIP_RANGE".to_string(), self.clip_range.unwrap_or(-1.0).to_string()),
285            ("ENT_COEF".to_string(), self.ent_coef.unwrap_or(-1.0).to_string()),
286            ("VF_COEF".to_string(), self.vf_coef.unwrap_or(-1.0).to_string()),
287            (
288                "MAX_GRAD_NORM".to_string(),
289                self.max_grad_norm.unwrap_or(-1.0).to_string(),
290            ),
291            ("PATIENCE".to_string(), self.patience.unwrap_or(-1).to_string()),
292            ("MIN_DELTA".to_string(), self.min_delta.unwrap_or(-1.0).to_string()),
293            (
294                "RESUME_RUN_ID".to_string(),
295                self.resume_run_id.clone().unwrap_or_default(),
296            ),
297            ("TUNE_RUN_ID".to_string(), self.tune_run_id.clone().unwrap_or_default()),
298            ("N_TRIALS".to_string(), self.n_trials.unwrap_or(-1).to_string()),
299            (
300                "TUNE_TIMESTEPS".to_string(),
301                self.tune_timesteps.unwrap_or(-1).to_string(),
302            ),
303            ("EVAL_EPS".to_string(), self.eval_eps.unwrap_or(-1).to_string()),
304            ("MLFLOW_ENABLED".to_string(), "true".to_string()),
305            ("TRACKER_BACKEND".to_string(), "relay".to_string()),
306            ("MLFLOW_TRACKING_URI".to_string(), control_plane_url.to_string()),
307            ("RELAY_API_KEY".to_string(), relay_api_key.to_string()),
308            ("LOG_LEVEL".to_string(), log_level.to_string()),
309        ]
310    }
311
312    /// `rl_task_runner.run`'s mode positional arg for this job.
313    pub fn run_mode(&self) -> &'static str {
314        self.flow_type.run_mode()
315    }
316
317    /// This job's fields in `mecha10-job-validation`'s shared vocabulary (see
318    /// [`mecha10_job_validation::JobFields`]) - the CLI-specific hyperparameters
319    /// (`n_steps`/`learning_rate`/...) aren't included since that crate has no opinion on them.
320    fn validation_fields(&self) -> mecha10_job_validation::JobFields {
321        mecha10_job_validation::JobFields {
322            task_name: self.task_name.clone(),
323            algorithm: self.algorithm.clone(),
324            total_timesteps: non_sentinel_i64(self.total_timesteps),
325            resume_run_id: self.resume_run_id.clone(),
326            tune_run_id: self.tune_run_id.clone(),
327            n_trials: non_sentinel_i64(self.n_trials),
328            tune_timesteps: non_sentinel_i64(self.tune_timesteps),
329            eval_eps: non_sentinel_i64(self.eval_eps),
330        }
331    }
332
333    /// Validate this job against `mecha10-job-validation`'s flow_type requirement table, plus a
334    /// `task_name`-exists check against the project's own task catalog
335    /// (`configs/rl/tasks/<task_name>/task.json`) - the single validation pass both local and
336    /// remote dispatch should run before doing anything else (see `run_local`/`run_remote` in
337    /// `handlers/train.rs`). Replaces the ad-hoc, dispatch-path-specific checks that used to live
338    /// only in `handlers/train.rs` (e.g. the blank-`algorithm` check for remote dispatch).
339    pub fn validate(&self, project_root: &Path) -> Result<()> {
340        let missing = mecha10_job_validation::validate_job(self.flow_type, &self.validation_fields());
341        if !missing.is_empty() {
342            let details = missing.iter().map(|m| m.to_string()).collect::<Vec<_>>().join("; ");
343            bail!("job '{}' is missing required fields: {details}", self.name);
344        }
345
346        if let Some(task_name) = self.task_name.as_deref() {
347            mecha10_job_validation::task_exists(project_root, task_name)
348                .with_context(|| format!("job '{}' has an invalid task_name", self.name))?;
349        }
350
351        Ok(())
352    }
353
354    /// Resolve which `mecha10_rl.envs` registry name this job targets: an explicit `env_name`
355    /// wins (back-compat for a job definition that still sets it), else it's read from
356    /// `task_name`'s own `task.json` (LAB env-belongs-with-the-task redesign - see
357    /// `train.yml.template`). `None` when neither is available - callers send that through
358    /// unchanged (an empty `ENV_NAME` locally, an omitted `env_name` for remote dispatch); the
359    /// Python side (`rl_task_runner.run`) applies the exact same fallback order independently,
360    /// so this is a CLI-side convenience/consistency resolution, not the only place it happens.
361    pub fn resolve_env_name(&self, project_root: &Path) -> Option<String> {
362        self.env_name.clone().or_else(|| {
363            self.task_name
364                .as_deref()
365                .and_then(|task_name| env_name_from_task_json(project_root, task_name))
366        })
367    }
368
369    /// Whether this job can be dispatched remotely with today's `CreateJobRequest`/`create_job`
370    /// contract. `train` and `resume-run` are wired server-side
371    /// (`mecha10-api/src/jobs/handlers.rs::train_flow_inputs`/`resume_flow_inputs`);
372    /// `tune`/`tune-then-train`/`train-from-tune` aren't dispatchable from `create_job` yet and
373    /// stay local-only until the API grows dedicated support.
374    pub fn supports_remote(&self) -> bool {
375        matches!(self.flow_type, FlowType::Train | FlowType::ResumeRun)
376    }
377
378    /// Build the remote dispatch request body. Only meaningful when `supports_remote()` is true;
379    /// callers should check that first. `algorithm` is required non-empty by `create_job` (unlike
380    /// local dispatch, where an empty `algorithm_name` falls back to the config file's value) -
381    /// callers should validate before sending.
382    pub fn to_create_job_request(
383        &self,
384        project_root: &Path,
385        robot_id: &str,
386        project_name: &str,
387    ) -> Result<CreateJobRequest> {
388        let config_yaml = std::fs::read_to_string(&self.source_path).with_context(|| {
389            format!(
390                "failed to read {} - this job definition's own file must exist before dispatching remotely",
391                self.source_path.display()
392            )
393        })?;
394
395        Ok(CreateJobRequest {
396            name: self.name.clone(),
397            description: self.description.clone(),
398            algorithm: self.algorithm.clone().unwrap_or_default(),
399            robot_id: robot_id.to_string(),
400            project_id: Some(robot_id.to_string()),
401            project_name: Some(project_name.to_string()),
402            // Resolved (not `self.env_name.clone()`) so a job definition that only sets
403            // `task_name` still sends a real `env_name` over the wire too, for back-compat with
404            // any consumer still keying off `env_name` alone.
405            env_name: self.resolve_env_name(project_root),
406            task_name: self.task_name.clone(),
407            total_timesteps: non_sentinel_i64(self.total_timesteps),
408            n_steps: non_sentinel_i64(self.n_steps),
409            n_epochs: non_sentinel_i64(self.n_epochs),
410            learning_rate: non_sentinel_f64(self.learning_rate),
411            patience: non_sentinel_i64(self.patience),
412            min_delta: non_sentinel_f64(self.min_delta),
413            config_yaml: Some(config_yaml),
414            resume_run_id: self
415                .resume_run_id
416                .as_deref()
417                .map(str::trim)
418                .filter(|id| !id.is_empty())
419                .map(str::to_string),
420        })
421    }
422}
423
424/// Load every `.yml`/`.yaml` job definition under `{project_root}/configs/rl/jobs/`, sorted by
425/// file name for a stable menu order. Returns an empty list (not an error) when the directory
426/// doesn't exist - `mecha10 train` treats that as "no jobs configured yet" and points the user at
427/// `mecha10 init`/the bundled defaults rather than failing.
428pub fn load_job_catalog(project_root: &Path) -> Result<Vec<JobDefinition>> {
429    let jobs_dir = project_root.join(JOB_CATALOG_DIR);
430    if !jobs_dir.exists() {
431        return Ok(Vec::new());
432    }
433
434    let mut paths: Vec<PathBuf> = std::fs::read_dir(&jobs_dir)
435        .with_context(|| format!("failed to read {}", jobs_dir.display()))?
436        .filter_map(|entry| entry.ok())
437        .map(|entry| entry.path())
438        .filter(|path| {
439            path.extension()
440                .and_then(|ext| ext.to_str())
441                .is_some_and(|ext| ext == "yml" || ext == "yaml")
442        })
443        .collect();
444    paths.sort();
445
446    let mut jobs = Vec::with_capacity(paths.len());
447    for path in paths {
448        let content = std::fs::read_to_string(&path).with_context(|| format!("failed to read {}", path.display()))?;
449        let mut job: JobDefinition = serde_yaml::from_str(&content)
450            .with_context(|| format!("failed to parse job definition {}", path.display()))?;
451        job.source_path = path;
452        jobs.push(job);
453    }
454
455    Ok(jobs)
456}
457
458#[cfg(test)]
459mod tests {
460    use super::*;
461    use tempfile::TempDir;
462
463    fn write_job(dir: &Path, file_name: &str, contents: &str) {
464        std::fs::write(dir.join(file_name), contents).unwrap();
465    }
466
467    /// A minimal, fully-`None`/default `JobDefinition` for tests to `..base_job()` from, so each
468    /// test only spells out the fields it actually cares about.
469    fn base_job() -> JobDefinition {
470        JobDefinition {
471            name: "job".to_string(),
472            description: None,
473            flow_type: FlowType::Train,
474            env_name: None,
475            task_name: None,
476            algorithm: Some("ppo".to_string()),
477            policy: None,
478            seed: None,
479            environment_map: None,
480            total_timesteps: None,
481            learning_rate: None,
482            linear_lr_schedule: None,
483            n_steps: None,
484            batch_size: None,
485            n_epochs: None,
486            gamma: None,
487            gae_lambda: None,
488            clip_range: None,
489            ent_coef: None,
490            vf_coef: None,
491            max_grad_norm: None,
492            patience: None,
493            min_delta: None,
494            resume_run_id: None,
495            tune_run_id: None,
496            n_trials: None,
497            tune_timesteps: None,
498            eval_eps: None,
499            source_path: PathBuf::new(),
500        }
501    }
502
503    #[test]
504    fn load_job_catalog_returns_empty_when_directory_missing() {
505        let temp = TempDir::new().unwrap();
506        let jobs = load_job_catalog(temp.path()).unwrap();
507        assert!(jobs.is_empty());
508    }
509
510    #[test]
511    fn load_job_catalog_parses_yaml_files_sorted_by_name() {
512        let temp = TempDir::new().unwrap();
513        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
514        std::fs::create_dir_all(&jobs_dir).unwrap();
515
516        write_job(
517            &jobs_dir,
518            "train.yml",
519            "name: train\nflow_type: train\nalgorithm: ppo\ntotal_timesteps: 1000000\n",
520        );
521        write_job(
522            &jobs_dir,
523            "hypertune.yml",
524            "name: hypertune\nflow_type: tune\nalgorithm: ppo\nn_trials: 10\n",
525        );
526
527        let jobs = load_job_catalog(temp.path()).unwrap();
528
529        assert_eq!(jobs.len(), 2);
530        assert_eq!(jobs[0].name, "hypertune");
531        assert_eq!(jobs[0].flow_type, FlowType::Tune);
532        assert_eq!(jobs[0].run_mode(), "tune");
533        assert_eq!(jobs[1].name, "train");
534        assert_eq!(jobs[1].flow_type, FlowType::Train);
535        assert_eq!(jobs[1].run_mode(), "train");
536        assert_eq!(jobs[1].total_timesteps, Some(1000000));
537    }
538
539    #[test]
540    fn load_job_catalog_ignores_non_yaml_files() {
541        let temp = TempDir::new().unwrap();
542        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
543        std::fs::create_dir_all(&jobs_dir).unwrap();
544
545        write_job(&jobs_dir, "README.md", "not a job");
546        write_job(
547            &jobs_dir,
548            "train.yml",
549            "name: train\nflow_type: train\nalgorithm: ppo\n",
550        );
551
552        let jobs = load_job_catalog(temp.path()).unwrap();
553        assert_eq!(jobs.len(), 1);
554    }
555
556    #[test]
557    fn flow_type_run_mode_dispatches_tune_family_to_tune_mode() {
558        assert_eq!(FlowType::Train.run_mode(), "train");
559        assert_eq!(FlowType::TrainFromTune.run_mode(), "train");
560        assert_eq!(FlowType::ResumeRun.run_mode(), "train");
561        assert_eq!(FlowType::Tune.run_mode(), "tune");
562        assert_eq!(FlowType::TuneThenTrain.run_mode(), "tune");
563    }
564
565    #[test]
566    fn supports_remote_for_train_and_resume_run_only() {
567        let base = base_job();
568        assert!(base.supports_remote());
569
570        let resume = JobDefinition {
571            flow_type: FlowType::ResumeRun,
572            ..base.clone()
573        };
574        assert!(resume.supports_remote());
575
576        let tune = JobDefinition {
577            flow_type: FlowType::Tune,
578            ..base
579        };
580        assert!(!tune.supports_remote());
581    }
582
583    #[test]
584    fn to_create_job_request_forwards_resume_run_id() {
585        let temp = TempDir::new().unwrap();
586        let jobs_dir = temp.path().join("configs/rl/jobs");
587        std::fs::create_dir_all(&jobs_dir).unwrap();
588        let source_path = jobs_dir.join("train-from-checkpoint.yml");
589        std::fs::write(&source_path, "name: train-from-checkpoint\nflow_type: resume-run\n").unwrap();
590
591        let job = JobDefinition {
592            name: "train-from-checkpoint".to_string(),
593            flow_type: FlowType::ResumeRun,
594            resume_run_id: Some("  run-abc123  ".to_string()),
595            source_path,
596            ..base_job()
597        };
598
599        let req = job
600            .to_create_job_request(temp.path(), "my-project", "my-robot")
601            .unwrap();
602
603        assert_eq!(req.resume_run_id.as_deref(), Some("run-abc123"));
604    }
605
606    #[test]
607    fn to_create_job_request_omits_empty_resume_run_id() {
608        let temp = TempDir::new().unwrap();
609        let jobs_dir = temp.path().join("configs/rl/jobs");
610        std::fs::create_dir_all(&jobs_dir).unwrap();
611        let source_path = jobs_dir.join("train-from-checkpoint.yml");
612        std::fs::write(&source_path, "name: train-from-checkpoint\nflow_type: resume-run\n").unwrap();
613
614        let job = JobDefinition {
615            name: "train-from-checkpoint".to_string(),
616            flow_type: FlowType::ResumeRun,
617            resume_run_id: Some("   ".to_string()),
618            source_path,
619            ..base_job()
620        };
621
622        let req = job
623            .to_create_job_request(temp.path(), "my-project", "my-robot")
624            .unwrap();
625
626        assert_eq!(req.resume_run_id, None);
627    }
628
629    #[test]
630    fn to_create_job_request_drops_sentinel_values() {
631        let temp = TempDir::new().unwrap();
632        let jobs_dir = temp.path().join("configs/rl/jobs");
633        std::fs::create_dir_all(&jobs_dir).unwrap();
634        let source_path = jobs_dir.join("train.yml");
635        std::fs::write(&source_path, "name: train\nflow_type: train\nenv_name: rover\n").unwrap();
636
637        let job = JobDefinition {
638            name: "train".to_string(),
639            description: Some("desc".to_string()),
640            env_name: Some("rover".to_string()),
641            total_timesteps: Some(-1),
642            n_steps: Some(2048),
643            n_epochs: Some(-1),
644            learning_rate: Some(-1.0),
645            patience: Some(50),
646            min_delta: Some(-1.0),
647            source_path,
648            ..base_job()
649        };
650
651        let req = job
652            .to_create_job_request(temp.path(), "my-project", "my-robot")
653            .unwrap();
654
655        assert_eq!(req.name, "train");
656        assert_eq!(req.algorithm, "ppo");
657        assert_eq!(req.robot_id, "my-project");
658        assert_eq!(req.project_id.as_deref(), Some("my-project"));
659        assert_eq!(req.project_name.as_deref(), Some("my-robot"));
660        assert_eq!(req.total_timesteps, None);
661        assert_eq!(req.n_steps, Some(2048));
662        assert_eq!(req.n_epochs, None);
663        assert_eq!(req.learning_rate, None);
664        assert_eq!(req.patience, Some(50));
665        assert_eq!(req.min_delta, None);
666        assert_eq!(
667            req.config_yaml.as_deref(),
668            Some("name: train\nflow_type: train\nenv_name: rover\n")
669        );
670        assert_eq!(req.task_name, None);
671    }
672
673    #[test]
674    fn to_create_job_request_forwards_task_name() {
675        let temp = TempDir::new().unwrap();
676        let jobs_dir = temp.path().join("configs/rl/jobs");
677        std::fs::create_dir_all(&jobs_dir).unwrap();
678        let source_path = jobs_dir.join("train.yml");
679        std::fs::write(&source_path, "name: train\nflow_type: train\n").unwrap();
680
681        let job = JobDefinition {
682            name: "train".to_string(),
683            task_name: Some("navigate".to_string()),
684            source_path,
685            ..base_job()
686        };
687
688        let req = job
689            .to_create_job_request(temp.path(), "my-project", "my-robot")
690            .unwrap();
691
692        // `task_name` is sent as-is over the wire so remote (Kestra) dispatch has the same
693        // `mecha10_rl.envs` factory argument local dispatch's `TASK_NAME` env var provides -
694        // without this, `GoalConditionedEnv.__init__` fails with a missing required keyword-only
695        // argument.
696        assert_eq!(req.task_name.as_deref(), Some("navigate"));
697    }
698
699    #[test]
700    fn to_create_job_request_errors_when_the_jobs_own_file_is_missing() {
701        let temp = TempDir::new().unwrap();
702        let job = JobDefinition {
703            name: "train".to_string(),
704            source_path: temp.path().join("configs/rl/jobs/train.yml"),
705            ..base_job()
706        };
707
708        let err = job
709            .to_create_job_request(temp.path(), "my-project", "my-robot")
710            .unwrap_err();
711        assert!(err.to_string().contains("configs/rl/jobs/train.yml"));
712    }
713
714    #[test]
715    fn local_env_vars_forwards_ppo_config_native_field_names_and_relay_settings() {
716        let job = JobDefinition {
717            name: "train".to_string(),
718            env_name: Some("rover".to_string()),
719            total_timesteps: Some(5000),
720            learning_rate: Some(0.0003),
721            n_epochs: Some(10),
722            resume_run_id: Some("run-123".to_string()),
723            source_path: PathBuf::from("configs/rl/jobs/train.yml"),
724            ..base_job()
725        };
726
727        let env = job.local_env_vars(
728            &std::path::PathBuf::from("."),
729            "https://control-plane.example.com",
730            "mecha_test_key",
731            "debug",
732        );
733        let as_map: std::collections::HashMap<_, _> = env.into_iter().collect();
734
735        assert_eq!(as_map.get("TOTAL_TIMESTEPS").map(String::as_str), Some("5000"));
736        assert_eq!(as_map.get("LEARNING_RATE").map(String::as_str), Some("0.0003"));
737        assert_eq!(as_map.get("N_EPOCHS").map(String::as_str), Some("10"));
738        assert_eq!(as_map.get("CONFIG").map(String::as_str), Some("train"));
739        assert_eq!(as_map.get("RL_CONFIG_DIR").map(String::as_str), Some("configs/rl/jobs"));
740        assert_eq!(as_map.get("RESUME_RUN_ID").map(String::as_str), Some("run-123"));
741        assert_eq!(as_map.get("TRACKER_BACKEND").map(String::as_str), Some("relay"));
742        assert_eq!(as_map.get("MLFLOW_ENABLED").map(String::as_str), Some("true"));
743        assert_eq!(
744            as_map.get("MLFLOW_TRACKING_URI").map(String::as_str),
745            Some("https://control-plane.example.com")
746        );
747        assert_eq!(as_map.get("RELAY_API_KEY").map(String::as_str), Some("mecha_test_key"));
748        assert_eq!(as_map.get("LOG_LEVEL").map(String::as_str), Some("debug"));
749    }
750}
mecha10_cli/services/train_jobs.rs

mecha10_cli/services/
train_jobs.rs
