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//! Field names/sentinel conventions (`-1`/`-1.0` meaning "use the config file's value")
12//! deliberately mirror `flows/ml-pipeline/train.yml`'s `inputs:` and
13//! `mecha10-api/src/jobs/types.rs::CreateJobRequest`, so a job definition maps 1:1 onto either
14//! dispatch path without translation surprises.
15
16use anyhow::{Context, Result};
17use serde::{Deserialize, Serialize};
18use std::path::{Path, PathBuf};
19
20/// Directory (relative to a project's root) job definitions are loaded from.
21pub const JOB_CATALOG_DIR: &str = "configs/rl/jobs";
22
23/// Which Kestra `flows/ml-pipeline/*.yml` flow (and `rl_task_runner.run` mode) a job maps to.
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
25#[serde(rename_all = "kebab-case")]
26pub enum FlowType {
27    Train,
28    Tune,
29    TuneThenTrain,
30    TrainFromTune,
31    ResumeRun,
32}
33
34impl FlowType {
35    /// `rl_task_runner.run`'s mode positional arg (`train` or `tune`) for this flow type -
36    /// `tune`/`tune-then-train` dispatch through `run.py`'s `tune` mode, everything else trains.
37    pub fn run_mode(self) -> &'static str {
38        match self {
39            FlowType::Tune | FlowType::TuneThenTrain => "tune",
40            FlowType::Train | FlowType::TrainFromTune | FlowType::ResumeRun => "train",
41        }
42    }
43}
44
45impl std::fmt::Display for FlowType {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        let s = match self {
48            FlowType::Train => "train",
49            FlowType::Tune => "tune",
50            FlowType::TuneThenTrain => "tune-then-train",
51            FlowType::TrainFromTune => "train-from-tune",
52            FlowType::ResumeRun => "resume-run",
53        };
54        write!(f, "{s}")
55    }
56}
57
58/// A training job definition, loaded from `configs/rl/jobs/<name>.yml`.
59#[derive(Debug, Clone, Deserialize)]
60pub struct JobDefinition {
61    pub name: String,
62    #[serde(default)]
63    pub description: Option<String>,
64    pub flow_type: FlowType,
65    /// rl-task-runner config name, e.g. `local` for `configs/rl/local.yml` (without extension).
66    pub config: String,
67    /// `mecha10_rl.envs` registry name. Empty/omitted = use the config file's `env_name`.
68    #[serde(default)]
69    pub env_name: Option<String>,
70    /// `mecha10_rl.algorithms` registry name. Empty/omitted = use the config file's
71    /// `algorithm_name` for local dispatch - but note remote dispatch requires this to be set
72    /// (see `to_create_job_request`'s doc comment).
73    #[serde(default)]
74    pub algorithm: Option<String>,
75    #[serde(default)]
76    pub timesteps: Option<i64>,
77    #[serde(default)]
78    pub n_steps: Option<i64>,
79    #[serde(default)]
80    pub epochs: Option<i64>,
81    #[serde(default)]
82    pub lr: Option<f64>,
83    #[serde(default)]
84    pub patience: Option<i64>,
85    #[serde(default)]
86    pub min_delta: Option<f64>,
87    /// MLflow run ID to resume from (`resume-run` flow type). See `rl_task_runner.run`'s
88    /// warm-start resolution.
89    #[serde(default)]
90    pub resume_run_id: Option<String>,
91    /// MLflow run ID of a prior Optuna tune (`train-from-tune` flow type). See
92    /// `rl_task_runner.run`'s best-params resolution.
93    #[serde(default)]
94    pub tune_run_id: Option<String>,
95    #[serde(default)]
96    pub n_trials: Option<i64>,
97    #[serde(default)]
98    pub tune_timesteps: Option<i64>,
99    #[serde(default)]
100    pub eval_eps: Option<i64>,
101
102    /// Not part of the YAML - filled in by `load_job_catalog` with the file it was loaded from.
103    #[serde(skip)]
104    pub source_path: PathBuf,
105}
106
107/// Request body for `POST {control_plane_url}/api/model-deployments/jobs` - field-for-field the
108/// same shape as `mecha10-api/src/jobs/types.rs::CreateJobRequest` (and, transitively,
109/// `packages/dashboard`'s `types/ml-training.ts` `CreateJobRequest`/`buildCreateJobRequest`,
110/// PR #278/LAB-1431). Deliberately duplicated rather than shared via a crate dependency - the
111/// CLI and API are different deployables and this struct is the wire contract, not shared code.
112#[derive(Debug, Clone, Serialize)]
113pub struct CreateJobRequest {
114    pub name: String,
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub description: Option<String>,
117    pub algorithm: String,
118    #[serde(skip_serializing_if = "Option::is_none")]
119    pub project_id: Option<String>,
120    #[serde(skip_serializing_if = "Option::is_none")]
121    pub config: Option<String>,
122    #[serde(skip_serializing_if = "Option::is_none")]
123    pub env_name: Option<String>,
124    #[serde(skip_serializing_if = "Option::is_none")]
125    pub timesteps: Option<i64>,
126    #[serde(skip_serializing_if = "Option::is_none")]
127    pub n_steps: Option<i64>,
128    #[serde(skip_serializing_if = "Option::is_none")]
129    pub epochs: Option<i64>,
130    #[serde(skip_serializing_if = "Option::is_none")]
131    pub lr: Option<f64>,
132    #[serde(skip_serializing_if = "Option::is_none")]
133    pub patience: Option<i64>,
134    #[serde(skip_serializing_if = "Option::is_none")]
135    pub min_delta: Option<f64>,
136}
137
138fn non_sentinel_i64(value: Option<i64>) -> Option<i64> {
139    value.filter(|v| *v != -1)
140}
141
142fn non_sentinel_f64(value: Option<f64>) -> Option<f64> {
143    value.filter(|v| *v != -1.0)
144}
145
146impl JobDefinition {
147    /// Env vars for `python3 -m rl_task_runner.run <mode>` - the local dispatch path. Applies
148    /// exactly the same `TIMESTEPS`->`TOTAL_TIMESTEPS`/`EPOCHS`->`N_EPOCHS`/`LR`->`LEARNING_RATE`
149    /// renaming `compose/trainer/docker-compose.yml` applies for a Kestra-dispatched remote run
150    /// (see that file's `environment:` block), so a local run of the same job definition behaves
151    /// identically. `MLFLOW_ENABLED`/`TRACKER_BACKEND=relay`/`RELAY_API_KEY` route metrics
152    /// through mecha10-api's relay (`rl_task_runner.services.relay_tracker.RelayTracker`)
153    /// instead of a direct MLflow connection - no MLflow credentials needed on this machine.
154    pub fn local_env_vars(
155        &self,
156        control_plane_url: &str,
157        relay_api_key: &str,
158        log_level: &str,
159    ) -> Vec<(String, String)> {
160        vec![
161            ("RL_CONFIG_DIR".to_string(), "configs/rl".to_string()),
162            ("CONFIG".to_string(), self.config.clone()),
163            ("ENV_NAME".to_string(), self.env_name.clone().unwrap_or_default()),
164            ("ALGORITHM_NAME".to_string(), self.algorithm.clone().unwrap_or_default()),
165            ("TOTAL_TIMESTEPS".to_string(), self.timesteps.unwrap_or(-1).to_string()),
166            ("N_STEPS".to_string(), self.n_steps.unwrap_or(-1).to_string()),
167            ("N_EPOCHS".to_string(), self.epochs.unwrap_or(-1).to_string()),
168            ("LEARNING_RATE".to_string(), self.lr.unwrap_or(-1.0).to_string()),
169            ("PATIENCE".to_string(), self.patience.unwrap_or(-1).to_string()),
170            ("MIN_DELTA".to_string(), self.min_delta.unwrap_or(-1.0).to_string()),
171            (
172                "RESUME_RUN_ID".to_string(),
173                self.resume_run_id.clone().unwrap_or_default(),
174            ),
175            ("TUNE_RUN_ID".to_string(), self.tune_run_id.clone().unwrap_or_default()),
176            ("N_TRIALS".to_string(), self.n_trials.unwrap_or(-1).to_string()),
177            (
178                "TUNE_TIMESTEPS".to_string(),
179                self.tune_timesteps.unwrap_or(-1).to_string(),
180            ),
181            ("EVAL_EPS".to_string(), self.eval_eps.unwrap_or(-1).to_string()),
182            ("MLFLOW_ENABLED".to_string(), "true".to_string()),
183            ("TRACKER_BACKEND".to_string(), "relay".to_string()),
184            ("MLFLOW_TRACKING_URI".to_string(), control_plane_url.to_string()),
185            ("RELAY_API_KEY".to_string(), relay_api_key.to_string()),
186            ("LOG_LEVEL".to_string(), log_level.to_string()),
187        ]
188    }
189
190    /// `rl_task_runner.run`'s mode positional arg for this job.
191    pub fn run_mode(&self) -> &'static str {
192        self.flow_type.run_mode()
193    }
194
195    /// Whether this job can be dispatched remotely with today's `CreateJobRequest`/`create_job`
196    /// contract - only the plain `train` flow is wired server-side
197    /// (`mecha10-api/src/jobs/handlers.rs::train_flow_inputs` only builds `train.yml` inputs;
198    /// `tune`/`tune-then-train`/`train-from-tune`/`resume-run` aren't dispatchable from
199    /// `create_job` yet). Other flow types are local-only until the API grows dedicated support.
200    pub fn supports_remote(&self) -> bool {
201        matches!(self.flow_type, FlowType::Train)
202    }
203
204    /// Build the remote dispatch request body. Only meaningful when `supports_remote()` is true;
205    /// callers should check that first. `algorithm` is required non-empty by `create_job` (unlike
206    /// local dispatch, where an empty `algorithm_name` falls back to the config file's value) -
207    /// callers should validate before sending.
208    pub fn to_create_job_request(&self, project_id: &str) -> CreateJobRequest {
209        CreateJobRequest {
210            name: self.name.clone(),
211            description: self.description.clone(),
212            algorithm: self.algorithm.clone().unwrap_or_default(),
213            project_id: Some(project_id.to_string()),
214            config: Some(self.config.clone()),
215            env_name: self.env_name.clone(),
216            timesteps: non_sentinel_i64(self.timesteps),
217            n_steps: non_sentinel_i64(self.n_steps),
218            epochs: non_sentinel_i64(self.epochs),
219            lr: non_sentinel_f64(self.lr),
220            patience: non_sentinel_i64(self.patience),
221            min_delta: non_sentinel_f64(self.min_delta),
222        }
223    }
224}
225
226/// Load every `.yml`/`.yaml` job definition under `{project_root}/configs/rl/jobs/`, sorted by
227/// file name for a stable menu order. Returns an empty list (not an error) when the directory
228/// doesn't exist - `mecha10 train` treats that as "no jobs configured yet" and points the user at
229/// `mecha10 init`/the bundled defaults rather than failing.
230pub fn load_job_catalog(project_root: &Path) -> Result<Vec<JobDefinition>> {
231    let jobs_dir = project_root.join(JOB_CATALOG_DIR);
232    if !jobs_dir.exists() {
233        return Ok(Vec::new());
234    }
235
236    let mut paths: Vec<PathBuf> = std::fs::read_dir(&jobs_dir)
237        .with_context(|| format!("failed to read {}", jobs_dir.display()))?
238        .filter_map(|entry| entry.ok())
239        .map(|entry| entry.path())
240        .filter(|path| {
241            path.extension()
242                .and_then(|ext| ext.to_str())
243                .is_some_and(|ext| ext == "yml" || ext == "yaml")
244        })
245        .collect();
246    paths.sort();
247
248    let mut jobs = Vec::with_capacity(paths.len());
249    for path in paths {
250        let content = std::fs::read_to_string(&path).with_context(|| format!("failed to read {}", path.display()))?;
251        let mut job: JobDefinition = serde_yaml::from_str(&content)
252            .with_context(|| format!("failed to parse job definition {}", path.display()))?;
253        job.source_path = path;
254        jobs.push(job);
255    }
256
257    Ok(jobs)
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263    use tempfile::TempDir;
264
265    fn write_job(dir: &Path, file_name: &str, contents: &str) {
266        std::fs::write(dir.join(file_name), contents).unwrap();
267    }
268
269    #[test]
270    fn load_job_catalog_returns_empty_when_directory_missing() {
271        let temp = TempDir::new().unwrap();
272        let jobs = load_job_catalog(temp.path()).unwrap();
273        assert!(jobs.is_empty());
274    }
275
276    #[test]
277    fn load_job_catalog_parses_yaml_files_sorted_by_name() {
278        let temp = TempDir::new().unwrap();
279        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
280        std::fs::create_dir_all(&jobs_dir).unwrap();
281
282        write_job(
283            &jobs_dir,
284            "train.yml",
285            "name: train\nflow_type: train\nconfig: local\nalgorithm: ppo\n",
286        );
287        write_job(
288            &jobs_dir,
289            "hypertune.yml",
290            "name: hypertune\nflow_type: tune\nconfig: local\nn_trials: 10\n",
291        );
292
293        let jobs = load_job_catalog(temp.path()).unwrap();
294
295        assert_eq!(jobs.len(), 2);
296        // Sorted by file name ("hypertune.yml" < "train.yml"), not YAML `name:` field.
297        assert_eq!(jobs[0].name, "hypertune");
298        assert_eq!(jobs[0].flow_type, FlowType::Tune);
299        assert_eq!(jobs[0].run_mode(), "tune");
300        assert_eq!(jobs[1].name, "train");
301        assert_eq!(jobs[1].flow_type, FlowType::Train);
302        assert_eq!(jobs[1].run_mode(), "train");
303    }
304
305    #[test]
306    fn load_job_catalog_ignores_non_yaml_files() {
307        let temp = TempDir::new().unwrap();
308        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
309        std::fs::create_dir_all(&jobs_dir).unwrap();
310
311        write_job(&jobs_dir, "README.md", "not a job");
312        write_job(&jobs_dir, "train.yml", "name: train\nflow_type: train\nconfig: local\n");
313
314        let jobs = load_job_catalog(temp.path()).unwrap();
315        assert_eq!(jobs.len(), 1);
316    }
317
318    #[test]
319    fn flow_type_run_mode_dispatches_tune_family_to_tune_mode() {
320        assert_eq!(FlowType::Train.run_mode(), "train");
321        assert_eq!(FlowType::TrainFromTune.run_mode(), "train");
322        assert_eq!(FlowType::ResumeRun.run_mode(), "train");
323        assert_eq!(FlowType::Tune.run_mode(), "tune");
324        assert_eq!(FlowType::TuneThenTrain.run_mode(), "tune");
325    }
326
327    #[test]
328    fn supports_remote_only_for_plain_train() {
329        let base = JobDefinition {
330            name: "job".to_string(),
331            description: None,
332            flow_type: FlowType::Train,
333            config: "local".to_string(),
334            env_name: None,
335            algorithm: Some("ppo".to_string()),
336            timesteps: None,
337            n_steps: None,
338            epochs: None,
339            lr: None,
340            patience: None,
341            min_delta: None,
342            resume_run_id: None,
343            tune_run_id: None,
344            n_trials: None,
345            tune_timesteps: None,
346            eval_eps: None,
347            source_path: PathBuf::new(),
348        };
349        assert!(base.supports_remote());
350
351        let resume = JobDefinition {
352            flow_type: FlowType::ResumeRun,
353            ..base.clone()
354        };
355        assert!(!resume.supports_remote());
356
357        let tune = JobDefinition {
358            flow_type: FlowType::Tune,
359            ..base
360        };
361        assert!(!tune.supports_remote());
362    }
363
364    #[test]
365    fn to_create_job_request_drops_sentinel_values() {
366        let job = JobDefinition {
367            name: "train".to_string(),
368            description: Some("desc".to_string()),
369            flow_type: FlowType::Train,
370            config: "local".to_string(),
371            env_name: Some("rover".to_string()),
372            algorithm: Some("ppo".to_string()),
373            timesteps: Some(-1),
374            n_steps: Some(2048),
375            epochs: Some(-1),
376            lr: Some(-1.0),
377            patience: Some(50),
378            min_delta: Some(-1.0),
379            resume_run_id: None,
380            tune_run_id: None,
381            n_trials: None,
382            tune_timesteps: None,
383            eval_eps: None,
384            source_path: PathBuf::new(),
385        };
386
387        let req = job.to_create_job_request("my-project");
388
389        assert_eq!(req.name, "train");
390        assert_eq!(req.algorithm, "ppo");
391        assert_eq!(req.project_id.as_deref(), Some("my-project"));
392        assert_eq!(req.timesteps, None);
393        assert_eq!(req.n_steps, Some(2048));
394        assert_eq!(req.epochs, None);
395        assert_eq!(req.lr, None);
396        assert_eq!(req.patience, Some(50));
397        assert_eq!(req.min_delta, None);
398    }
399
400    #[test]
401    fn local_env_vars_applies_docker_compose_renaming_and_relay_settings() {
402        let job = JobDefinition {
403            name: "train".to_string(),
404            description: None,
405            flow_type: FlowType::Train,
406            config: "local".to_string(),
407            env_name: Some("rover".to_string()),
408            algorithm: Some("ppo".to_string()),
409            timesteps: Some(5000),
410            n_steps: None,
411            epochs: None,
412            lr: None,
413            patience: None,
414            min_delta: None,
415            resume_run_id: Some("run-123".to_string()),
416            tune_run_id: None,
417            n_trials: None,
418            tune_timesteps: None,
419            eval_eps: None,
420            source_path: PathBuf::new(),
421        };
422
423        let env = job.local_env_vars("https://control-plane.example.com", "mecha_test_key", "debug");
424        let as_map: std::collections::HashMap<_, _> = env.into_iter().collect();
425
426        assert_eq!(as_map.get("TOTAL_TIMESTEPS").map(String::as_str), Some("5000"));
427        assert_eq!(as_map.get("RESUME_RUN_ID").map(String::as_str), Some("run-123"));
428        assert_eq!(as_map.get("TRACKER_BACKEND").map(String::as_str), Some("relay"));
429        assert_eq!(as_map.get("MLFLOW_ENABLED").map(String::as_str), Some("true"));
430        assert_eq!(
431            as_map.get("MLFLOW_TRACKING_URI").map(String::as_str),
432            Some("https://control-plane.example.com")
433        );
434        assert_eq!(as_map.get("RELAY_API_KEY").map(String::as_str), Some("mecha_test_key"));
435        assert_eq!(as_map.get("LOG_LEVEL").map(String::as_str), Some("debug"));
436    }
437}
mecha10_cli/services/train_jobs.rs

mecha10_cli/services/
train_jobs.rs
