mecha10-cli 0.6.3

//! Job catalog for `mecha10 train` (LAB-1776).
//!
//! A "job" is a YAML file under a project's `configs/rl/jobs/` directory that maps to one of
//! the Kestra `flows/ml-pipeline/*.yml` flow types plus an env/algorithm/hyperparameter
//! selection. `mecha10-templates`'s RL bundle (`create_rl_training_bundle`) ships a standard set
//! of these into every scaffolded project - a project owner extends the menu simply by adding
//! more `.yml` files to that directory, no mecha10 core changes needed, same extensibility
//! philosophy as the `mecha10_rl.envs`/`mecha10_rl.algorithms` registry
//! (`rl_task_runner.registry`).
//!
//! Field names/sentinel conventions (`-1`/`-1.0` meaning "use the config file's value")
//! deliberately mirror `flows/ml-pipeline/train.yml`'s `inputs:` and
//! `mecha10-api/src/jobs/types.rs::CreateJobRequest`, so a job definition maps 1:1 onto either
//! dispatch path without translation surprises.

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};

/// Directory (relative to a project's root) job definitions are loaded from.
pub const JOB_CATALOG_DIR: &str = "configs/rl/jobs";

/// Which Kestra `flows/ml-pipeline/*.yml` flow (and `rl_task_runner.run` mode) a job maps to.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum FlowType {
    Train,
    Tune,
    TuneThenTrain,
    TrainFromTune,
    ResumeRun,
}

impl FlowType {
    /// `rl_task_runner.run`'s mode positional arg (`train` or `tune`) for this flow type -
    /// `tune`/`tune-then-train` dispatch through `run.py`'s `tune` mode, everything else trains.
    pub fn run_mode(self) -> &'static str {
        match self {
            FlowType::Tune | FlowType::TuneThenTrain => "tune",
            FlowType::Train | FlowType::TrainFromTune | FlowType::ResumeRun => "train",
        }
    }
}

impl std::fmt::Display for FlowType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let s = match self {
            FlowType::Train => "train",
            FlowType::Tune => "tune",
            FlowType::TuneThenTrain => "tune-then-train",
            FlowType::TrainFromTune => "train-from-tune",
            FlowType::ResumeRun => "resume-run",
        };
        write!(f, "{s}")
    }
}

/// A training job definition, loaded from `configs/rl/jobs/<name>.yml`.
#[derive(Debug, Clone, Deserialize)]
pub struct JobDefinition {
    pub name: String,
    #[serde(default)]
    pub description: Option<String>,
    pub flow_type: FlowType,
    /// rl-task-runner config name, e.g. `local` for `configs/rl/local.yml` (without extension).
    pub config: String,
    /// `mecha10_rl.envs` registry name. Empty/omitted = use the config file's `env_name`.
    #[serde(default)]
    pub env_name: Option<String>,
    /// `mecha10_rl.algorithms` registry name. Empty/omitted = use the config file's
    /// `algorithm_name` for local dispatch - but note remote dispatch requires this to be set
    /// (see `to_create_job_request`'s doc comment).
    #[serde(default)]
    pub algorithm: Option<String>,
    #[serde(default)]
    pub timesteps: Option<i64>,
    #[serde(default)]
    pub n_steps: Option<i64>,
    #[serde(default)]
    pub epochs: Option<i64>,
    #[serde(default)]
    pub lr: Option<f64>,
    #[serde(default)]
    pub patience: Option<i64>,
    #[serde(default)]
    pub min_delta: Option<f64>,
    /// MLflow run ID to resume from (`resume-run` flow type). See `rl_task_runner.run`'s
    /// warm-start resolution.
    #[serde(default)]
    pub resume_run_id: Option<String>,
    /// MLflow run ID of a prior Optuna tune (`train-from-tune` flow type). See
    /// `rl_task_runner.run`'s best-params resolution.
    #[serde(default)]
    pub tune_run_id: Option<String>,
    #[serde(default)]
    pub n_trials: Option<i64>,
    #[serde(default)]
    pub tune_timesteps: Option<i64>,
    #[serde(default)]
    pub eval_eps: Option<i64>,

    /// Not part of the YAML - filled in by `load_job_catalog` with the file it was loaded from.
    #[serde(skip)]
    pub source_path: PathBuf,
}

/// Request body for `POST {control_plane_url}/api/model-deployments/jobs` - field-for-field the
/// same shape as `mecha10-api/src/jobs/types.rs::CreateJobRequest` (and, transitively,
/// `packages/dashboard`'s `types/ml-training.ts` `CreateJobRequest`/`buildCreateJobRequest`,
/// PR #278/LAB-1431). Deliberately duplicated rather than shared via a crate dependency - the
/// CLI and API are different deployables and this struct is the wire contract, not shared code.
#[derive(Debug, Clone, Serialize)]
pub struct CreateJobRequest {
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    pub algorithm: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub project_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub config: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub env_name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub timesteps: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub n_steps: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub epochs: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub lr: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub patience: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub min_delta: Option<f64>,
}

fn non_sentinel_i64(value: Option<i64>) -> Option<i64> {
    value.filter(|v| *v != -1)
}

fn non_sentinel_f64(value: Option<f64>) -> Option<f64> {
    value.filter(|v| *v != -1.0)
}

impl JobDefinition {
    /// Env vars for `python3 -m rl_task_runner.run <mode>` - the local dispatch path. Applies
    /// exactly the same `TIMESTEPS`->`TOTAL_TIMESTEPS`/`EPOCHS`->`N_EPOCHS`/`LR`->`LEARNING_RATE`
    /// renaming `compose/trainer/docker-compose.yml` applies for a Kestra-dispatched remote run
    /// (see that file's `environment:` block), so a local run of the same job definition behaves
    /// identically. `MLFLOW_ENABLED`/`TRACKER_BACKEND=relay`/`RELAY_API_KEY` route metrics
    /// through mecha10-api's relay (`rl_task_runner.services.relay_tracker.RelayTracker`)
    /// instead of a direct MLflow connection - no MLflow credentials needed on this machine.
    pub fn local_env_vars(
        &self,
        control_plane_url: &str,
        relay_api_key: &str,
        log_level: &str,
    ) -> Vec<(String, String)> {
        vec![
            ("RL_CONFIG_DIR".to_string(), "configs/rl".to_string()),
            ("CONFIG".to_string(), self.config.clone()),
            ("ENV_NAME".to_string(), self.env_name.clone().unwrap_or_default()),
            ("ALGORITHM_NAME".to_string(), self.algorithm.clone().unwrap_or_default()),
            ("TOTAL_TIMESTEPS".to_string(), self.timesteps.unwrap_or(-1).to_string()),
            ("N_STEPS".to_string(), self.n_steps.unwrap_or(-1).to_string()),
            ("N_EPOCHS".to_string(), self.epochs.unwrap_or(-1).to_string()),
            ("LEARNING_RATE".to_string(), self.lr.unwrap_or(-1.0).to_string()),
            ("PATIENCE".to_string(), self.patience.unwrap_or(-1).to_string()),
            ("MIN_DELTA".to_string(), self.min_delta.unwrap_or(-1.0).to_string()),
            (
                "RESUME_RUN_ID".to_string(),
                self.resume_run_id.clone().unwrap_or_default(),
            ),
            ("TUNE_RUN_ID".to_string(), self.tune_run_id.clone().unwrap_or_default()),
            ("N_TRIALS".to_string(), self.n_trials.unwrap_or(-1).to_string()),
            (
                "TUNE_TIMESTEPS".to_string(),
                self.tune_timesteps.unwrap_or(-1).to_string(),
            ),
            ("EVAL_EPS".to_string(), self.eval_eps.unwrap_or(-1).to_string()),
            ("MLFLOW_ENABLED".to_string(), "true".to_string()),
            ("TRACKER_BACKEND".to_string(), "relay".to_string()),
            ("MLFLOW_TRACKING_URI".to_string(), control_plane_url.to_string()),
            ("RELAY_API_KEY".to_string(), relay_api_key.to_string()),
            ("LOG_LEVEL".to_string(), log_level.to_string()),
        ]
    }

    /// `rl_task_runner.run`'s mode positional arg for this job.
    pub fn run_mode(&self) -> &'static str {
        self.flow_type.run_mode()
    }

    /// Whether this job can be dispatched remotely with today's `CreateJobRequest`/`create_job`
    /// contract - only the plain `train` flow is wired server-side
    /// (`mecha10-api/src/jobs/handlers.rs::train_flow_inputs` only builds `train.yml` inputs;
    /// `tune`/`tune-then-train`/`train-from-tune`/`resume-run` aren't dispatchable from
    /// `create_job` yet). Other flow types are local-only until the API grows dedicated support.
    pub fn supports_remote(&self) -> bool {
        matches!(self.flow_type, FlowType::Train)
    }

    /// Build the remote dispatch request body. Only meaningful when `supports_remote()` is true;
    /// callers should check that first. `algorithm` is required non-empty by `create_job` (unlike
    /// local dispatch, where an empty `algorithm_name` falls back to the config file's value) -
    /// callers should validate before sending.
    pub fn to_create_job_request(&self, project_id: &str) -> CreateJobRequest {
        CreateJobRequest {
            name: self.name.clone(),
            description: self.description.clone(),
            algorithm: self.algorithm.clone().unwrap_or_default(),
            project_id: Some(project_id.to_string()),
            config: Some(self.config.clone()),
            env_name: self.env_name.clone(),
            timesteps: non_sentinel_i64(self.timesteps),
            n_steps: non_sentinel_i64(self.n_steps),
            epochs: non_sentinel_i64(self.epochs),
            lr: non_sentinel_f64(self.lr),
            patience: non_sentinel_i64(self.patience),
            min_delta: non_sentinel_f64(self.min_delta),
        }
    }
}

/// Load every `.yml`/`.yaml` job definition under `{project_root}/configs/rl/jobs/`, sorted by
/// file name for a stable menu order. Returns an empty list (not an error) when the directory
/// doesn't exist - `mecha10 train` treats that as "no jobs configured yet" and points the user at
/// `mecha10 init`/the bundled defaults rather than failing.
pub fn load_job_catalog(project_root: &Path) -> Result<Vec<JobDefinition>> {
    let jobs_dir = project_root.join(JOB_CATALOG_DIR);
    if !jobs_dir.exists() {
        return Ok(Vec::new());
    }

    let mut paths: Vec<PathBuf> = std::fs::read_dir(&jobs_dir)
        .with_context(|| format!("failed to read {}", jobs_dir.display()))?
        .filter_map(|entry| entry.ok())
        .map(|entry| entry.path())
        .filter(|path| {
            path.extension()
                .and_then(|ext| ext.to_str())
                .is_some_and(|ext| ext == "yml" || ext == "yaml")
        })
        .collect();
    paths.sort();

    let mut jobs = Vec::with_capacity(paths.len());
    for path in paths {
        let content = std::fs::read_to_string(&path).with_context(|| format!("failed to read {}", path.display()))?;
        let mut job: JobDefinition = serde_yaml::from_str(&content)
            .with_context(|| format!("failed to parse job definition {}", path.display()))?;
        job.source_path = path;
        jobs.push(job);
    }

    Ok(jobs)
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    fn write_job(dir: &Path, file_name: &str, contents: &str) {
        std::fs::write(dir.join(file_name), contents).unwrap();
    }

    #[test]
    fn load_job_catalog_returns_empty_when_directory_missing() {
        let temp = TempDir::new().unwrap();
        let jobs = load_job_catalog(temp.path()).unwrap();
        assert!(jobs.is_empty());
    }

    #[test]
    fn load_job_catalog_parses_yaml_files_sorted_by_name() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
        std::fs::create_dir_all(&jobs_dir).unwrap();

        write_job(
            &jobs_dir,
            "train.yml",
            "name: train\nflow_type: train\nconfig: local\nalgorithm: ppo\n",
        );
        write_job(
            &jobs_dir,
            "hypertune.yml",
            "name: hypertune\nflow_type: tune\nconfig: local\nn_trials: 10\n",
        );

        let jobs = load_job_catalog(temp.path()).unwrap();

        assert_eq!(jobs.len(), 2);
        // Sorted by file name ("hypertune.yml" < "train.yml"), not YAML `name:` field.
        assert_eq!(jobs[0].name, "hypertune");
        assert_eq!(jobs[0].flow_type, FlowType::Tune);
        assert_eq!(jobs[0].run_mode(), "tune");
        assert_eq!(jobs[1].name, "train");
        assert_eq!(jobs[1].flow_type, FlowType::Train);
        assert_eq!(jobs[1].run_mode(), "train");
    }

    #[test]
    fn load_job_catalog_ignores_non_yaml_files() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
        std::fs::create_dir_all(&jobs_dir).unwrap();

        write_job(&jobs_dir, "README.md", "not a job");
        write_job(&jobs_dir, "train.yml", "name: train\nflow_type: train\nconfig: local\n");

        let jobs = load_job_catalog(temp.path()).unwrap();
        assert_eq!(jobs.len(), 1);
    }

    #[test]
    fn flow_type_run_mode_dispatches_tune_family_to_tune_mode() {
        assert_eq!(FlowType::Train.run_mode(), "train");
        assert_eq!(FlowType::TrainFromTune.run_mode(), "train");
        assert_eq!(FlowType::ResumeRun.run_mode(), "train");
        assert_eq!(FlowType::Tune.run_mode(), "tune");
        assert_eq!(FlowType::TuneThenTrain.run_mode(), "tune");
    }

    #[test]
    fn supports_remote_only_for_plain_train() {
        let base = JobDefinition {
            name: "job".to_string(),
            description: None,
            flow_type: FlowType::Train,
            config: "local".to_string(),
            env_name: None,
            algorithm: Some("ppo".to_string()),
            timesteps: None,
            n_steps: None,
            epochs: None,
            lr: None,
            patience: None,
            min_delta: None,
            resume_run_id: None,
            tune_run_id: None,
            n_trials: None,
            tune_timesteps: None,
            eval_eps: None,
            source_path: PathBuf::new(),
        };
        assert!(base.supports_remote());

        let resume = JobDefinition {
            flow_type: FlowType::ResumeRun,
            ..base.clone()
        };
        assert!(!resume.supports_remote());

        let tune = JobDefinition {
            flow_type: FlowType::Tune,
            ..base
        };
        assert!(!tune.supports_remote());
    }

    #[test]
    fn to_create_job_request_drops_sentinel_values() {
        let job = JobDefinition {
            name: "train".to_string(),
            description: Some("desc".to_string()),
            flow_type: FlowType::Train,
            config: "local".to_string(),
            env_name: Some("rover".to_string()),
            algorithm: Some("ppo".to_string()),
            timesteps: Some(-1),
            n_steps: Some(2048),
            epochs: Some(-1),
            lr: Some(-1.0),
            patience: Some(50),
            min_delta: Some(-1.0),
            resume_run_id: None,
            tune_run_id: None,
            n_trials: None,
            tune_timesteps: None,
            eval_eps: None,
            source_path: PathBuf::new(),
        };

        let req = job.to_create_job_request("my-project");

        assert_eq!(req.name, "train");
        assert_eq!(req.algorithm, "ppo");
        assert_eq!(req.project_id.as_deref(), Some("my-project"));
        assert_eq!(req.timesteps, None);
        assert_eq!(req.n_steps, Some(2048));
        assert_eq!(req.epochs, None);
        assert_eq!(req.lr, None);
        assert_eq!(req.patience, Some(50));
        assert_eq!(req.min_delta, None);
    }

    #[test]
    fn local_env_vars_applies_docker_compose_renaming_and_relay_settings() {
        let job = JobDefinition {
            name: "train".to_string(),
            description: None,
            flow_type: FlowType::Train,
            config: "local".to_string(),
            env_name: Some("rover".to_string()),
            algorithm: Some("ppo".to_string()),
            timesteps: Some(5000),
            n_steps: None,
            epochs: None,
            lr: None,
            patience: None,
            min_delta: None,
            resume_run_id: Some("run-123".to_string()),
            tune_run_id: None,
            n_trials: None,
            tune_timesteps: None,
            eval_eps: None,
            source_path: PathBuf::new(),
        };

        let env = job.local_env_vars("https://control-plane.example.com", "mecha_test_key", "debug");
        let as_map: std::collections::HashMap<_, _> = env.into_iter().collect();

        assert_eq!(as_map.get("TOTAL_TIMESTEPS").map(String::as_str), Some("5000"));
        assert_eq!(as_map.get("RESUME_RUN_ID").map(String::as_str), Some("run-123"));
        assert_eq!(as_map.get("TRACKER_BACKEND").map(String::as_str), Some("relay"));
        assert_eq!(as_map.get("MLFLOW_ENABLED").map(String::as_str), Some("true"));
        assert_eq!(
            as_map.get("MLFLOW_TRACKING_URI").map(String::as_str),
            Some("https://control-plane.example.com")
        );
        assert_eq!(as_map.get("RELAY_API_KEY").map(String::as_str), Some("mecha_test_key"));
        assert_eq!(as_map.get("LOG_LEVEL").map(String::as_str), Some("debug"));
    }
}