mecha10-templates 0.6.2

//! RL training bundle template generation (LAB-1743, Phase 3 of the robot-7 merge, LAB-1688)
//!
//! Packages the Phase 2 (LAB-1732/1733/1734) `rover`/`ppo` entry points
//! (`packages/rl-task-runner`, `packages/rl-trainer`) into the "standard bundle" a new `rover`
//! template project gets for free:
//!
//! - A default sentinel-override local/dev config for the `rover`/`ppo` combination
//!   (`configs/rl/local.yml`), consumed by `rl_task_runner.local_config` (LAB-1734).
//! - The Python package dependencies that provide the bundled env/algorithm, appended to
//!   `requirements.txt` (both packages register their entry points in their own `pyproject.toml`
//!   already - this bundle only needs to declare the dependency, per LAB-1732's design).
//! - The real Kestra flow templates (`kestra/flows/ml-pipeline/*.yml`) that dispatch a
//!   training/tune run against the bundled `rover`/`ppo` combination (`train`, `tune`,
//!   `tune-then-train`, `train-from-tune`, `resume-run`), plus a `kestra/README.md` explaining
//!   how to use them. These are embedded copies of the repo's own `flows/ml-pipeline/*.yml`
//!   (LAB-1745, Phase 5 of the robot-7 merge, LAB-1688) - parameterized via Kestra `inputs:`
//!   (`env_name`/`algorithm_name`/`config`/...) and `secret()` calls at flow-execution time, not
//!   Rust-side string substitution, so the same flow YAML works for any project (LAB-1765).
//!
//! Follows the same embedded-template-file pattern as the rest of this crate (see
//! `MetaTemplates`/`ProjectTemplateService`'s `MODEL_JSON_TEMPLATE`/`ENVIRONMENT_JSON_TEMPLATE`)
//! rather than inventing a second scaffolding mechanism - static files embedded at compile time
//! via `include_str!()`, written into the new project's directory tree.

use crate::paths;
use anyhow::{Context, Result};
use std::path::Path;
use tokio::io::AsyncWriteExt;

/// Default sentinel-override config for the bundled `rover`/`ppo` combination.
const RL_LOCAL_CONFIG: &str = include_str!("../../templates/rl/local.yml.template");

/// RL bundle's Python dependencies, appended to the project's `requirements.txt`.
const RL_REQUIREMENTS: &str = include_str!("../../templates/rl/requirements-rl.txt.template");

/// README describing the bundled Kestra flows and how they're parameterized.
const KESTRA_README: &str = include_str!("../../templates/rl/kestra-README.md.template");

/// The five generic ml-pipeline flows, embedded verbatim from the repo's own
/// `flows/ml-pipeline/*.yml` (LAB-1745) - not modified, only relocated into this crate's
/// templates directory so they can be embedded at compile time and bundled into new projects.
const KESTRA_FLOW_TRAIN: &str = include_str!("../../templates/rl/kestra-flows/train.yml.template");
const KESTRA_FLOW_TUNE: &str = include_str!("../../templates/rl/kestra-flows/tune.yml.template");
const KESTRA_FLOW_TUNE_THEN_TRAIN: &str = include_str!("../../templates/rl/kestra-flows/tune-then-train.yml.template");
const KESTRA_FLOW_TRAIN_FROM_TUNE: &str = include_str!("../../templates/rl/kestra-flows/train-from-tune.yml.template");
const KESTRA_FLOW_RESUME_RUN: &str = include_str!("../../templates/rl/kestra-flows/resume-run.yml.template");

/// The standard job catalog `mecha10 train`'s menu ships with (LAB-1776): plain train,
/// train-from-checkpoint (resume-run), and hypertune (tune). A project owner extends this by
/// adding more `*.yml` files under `configs/rl/jobs/` directly - `mecha10 train` scans that
/// whole directory, no mecha10 core changes needed (same extensibility model as the
/// `mecha10_rl.envs`/`mecha10_rl.algorithms` registry).
const JOB_TRAIN: &str = include_str!("../../templates/rl/jobs/train.yml.template");
const JOB_TRAIN_FROM_CHECKPOINT: &str = include_str!("../../templates/rl/jobs/train-from-checkpoint.yml.template");
const JOB_HYPERTUNE: &str = include_str!("../../templates/rl/jobs/hypertune.yml.template");

/// Directory the default RL config is written to, relative to the project root.
const RL_CONFIG_DIR: &str = "configs/rl";

/// Path to the default RL config, relative to the project root.
const RL_LOCAL_CONFIG_PATH: &str = "configs/rl/local.yml";

/// Directory the standard job catalog is written to, relative to the project root.
const RL_JOBS_DIR: &str = "configs/rl/jobs";

/// The bundled job definitions as `(file name, embedded content)` pairs, written under
/// `RL_JOBS_DIR`.
const RL_JOBS: [(&str, &str); 3] = [
    ("train.yml", JOB_TRAIN),
    ("train-from-checkpoint.yml", JOB_TRAIN_FROM_CHECKPOINT),
    ("hypertune.yml", JOB_HYPERTUNE),
];

/// Directory the Kestra bundle (README + flows) is written to, relative to the project root.
const KESTRA_DIR: &str = "kestra";

/// Path to the Kestra bundle's README, relative to the project root.
const KESTRA_README_PATH: &str = "kestra/README.md";

/// Directory the Kestra flow YAML files are written to, relative to the project root - mirrors
/// the source repo's own `flows/ml-pipeline/` layout, nested under the bundle's `kestra/` dir.
const KESTRA_FLOWS_DIR: &str = "kestra/flows/ml-pipeline";

/// The bundled flows as `(file name, embedded content)` pairs, written under
/// `KESTRA_FLOWS_DIR`.
const KESTRA_FLOWS: [(&str, &str); 5] = [
    ("train.yml", KESTRA_FLOW_TRAIN),
    ("tune.yml", KESTRA_FLOW_TUNE),
    ("tune-then-train.yml", KESTRA_FLOW_TUNE_THEN_TRAIN),
    ("train-from-tune.yml", KESTRA_FLOW_TRAIN_FROM_TUNE),
    ("resume-run.yml", KESTRA_FLOW_RESUME_RUN),
];

/// RL training bundle template generator
pub struct RlBundleTemplates;

impl RlBundleTemplates {
    /// Create a new RlBundleTemplates instance
    pub fn new() -> Self {
        Self
    }

    /// Write the default `rover`/`ppo` sentinel-override config to `configs/rl/local.yml`
    pub async fn create_rl_config(&self, path: &Path) -> Result<()> {
        let dest_dir = path.join(RL_CONFIG_DIR);
        tokio::fs::create_dir_all(&dest_dir)
            .await
            .with_context(|| format!("Failed to create {}", RL_CONFIG_DIR))?;

        tokio::fs::write(path.join(RL_LOCAL_CONFIG_PATH), RL_LOCAL_CONFIG)
            .await
            .context("Failed to write configs/rl/local.yml")?;

        Ok(())
    }

    /// Write the standard job catalog (`train`, `train-from-checkpoint`, `hypertune`) to
    /// `configs/rl/jobs/` - the menu `mecha10 train` lists by default for every scaffolded
    /// project (LAB-1776). A project owner adds more job definitions by dropping additional
    /// `*.yml` files in this same directory.
    pub async fn create_rl_jobs(&self, path: &Path) -> Result<()> {
        let jobs_dir = path.join(RL_JOBS_DIR);
        tokio::fs::create_dir_all(&jobs_dir)
            .await
            .with_context(|| format!("Failed to create {}", RL_JOBS_DIR))?;

        for (file_name, content) in RL_JOBS {
            tokio::fs::write(jobs_dir.join(file_name), content)
                .await
                .with_context(|| format!("Failed to write {}/{}", RL_JOBS_DIR, file_name))?;
        }

        Ok(())
    }

    /// Append the RL bundle's Python dependencies to the project's `requirements.txt`
    ///
    /// Assumes `requirements.txt` already exists (i.e. `create_requirements_txt` has already
    /// run) - appends rather than overwrites so this stays a single `pip install -r
    /// requirements.txt` for the whole project.
    pub async fn append_rl_requirements(&self, path: &Path) -> Result<()> {
        let requirements_path = path.join(paths::meta::REQUIREMENTS_TXT);

        let mut file = tokio::fs::OpenOptions::new()
            .append(true)
            .create(true)
            .open(&requirements_path)
            .await
            .context("Failed to open requirements.txt for appending RL bundle dependencies")?;

        file.write_all(RL_REQUIREMENTS.as_bytes())
            .await
            .context("Failed to append RL bundle dependencies to requirements.txt")?;

        Ok(())
    }

    /// Write the `kestra/` bundle: the usage README plus the five real ml-pipeline flow
    /// templates (`train`, `tune`, `tune-then-train`, `train-from-tune`, `resume-run`), under
    /// `kestra/flows/ml-pipeline/`.
    pub async fn create_kestra_flows(&self, path: &Path) -> Result<()> {
        let dest_dir = path.join(KESTRA_DIR);
        tokio::fs::create_dir_all(&dest_dir)
            .await
            .with_context(|| format!("Failed to create {}", KESTRA_DIR))?;

        tokio::fs::write(path.join(KESTRA_README_PATH), KESTRA_README)
            .await
            .context("Failed to write kestra/README.md")?;

        let flows_dir = path.join(KESTRA_FLOWS_DIR);
        tokio::fs::create_dir_all(&flows_dir)
            .await
            .with_context(|| format!("Failed to create {}", KESTRA_FLOWS_DIR))?;

        for (file_name, content) in KESTRA_FLOWS {
            tokio::fs::write(flows_dir.join(file_name), content)
                .await
                .with_context(|| format!("Failed to write {}/{}", KESTRA_FLOWS_DIR, file_name))?;
        }

        Ok(())
    }
}

impl Default for RlBundleTemplates {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[tokio::test]
    async fn test_create_rl_config_writes_expected_defaults() {
        let temp = TempDir::new().unwrap();
        let bundle = RlBundleTemplates::new();

        bundle.create_rl_config(temp.path()).await.unwrap();

        let content = tokio::fs::read_to_string(temp.path().join("configs/rl/local.yml"))
            .await
            .unwrap();
        assert!(content.contains("env: rover"));
        assert!(content.contains("algorithm: ppo"));
    }

    #[tokio::test]
    async fn test_create_rl_jobs_writes_standard_catalog() {
        let temp = TempDir::new().unwrap();
        let bundle = RlBundleTemplates::new();

        bundle.create_rl_jobs(temp.path()).await.unwrap();

        for file_name in ["train.yml", "train-from-checkpoint.yml", "hypertune.yml"] {
            let job_path = temp.path().join("configs/rl/jobs").join(file_name);
            let content = tokio::fs::read_to_string(&job_path)
                .await
                .unwrap_or_else(|e| panic!("failed to read {}: {}", job_path.display(), e));
            assert!(content.contains("name:"), "{file_name} missing name:");
            assert!(content.contains("flow_type:"), "{file_name} missing flow_type:");
            assert!(content.contains("config: local"), "{file_name} missing config: local");
        }
    }

    #[tokio::test]
    async fn test_create_rl_jobs_train_maps_to_the_train_flow() {
        let temp = TempDir::new().unwrap();
        let bundle = RlBundleTemplates::new();

        bundle.create_rl_jobs(temp.path()).await.unwrap();

        let content = tokio::fs::read_to_string(temp.path().join("configs/rl/jobs/train.yml"))
            .await
            .unwrap();
        assert!(content.contains("flow_type: train"));
        assert!(content.contains("algorithm: ppo"));
    }

    #[tokio::test]
    async fn test_create_rl_jobs_train_from_checkpoint_maps_to_resume_run_flow() {
        let temp = TempDir::new().unwrap();
        let bundle = RlBundleTemplates::new();

        bundle.create_rl_jobs(temp.path()).await.unwrap();

        let content = tokio::fs::read_to_string(temp.path().join("configs/rl/jobs/train-from-checkpoint.yml"))
            .await
            .unwrap();
        assert!(content.contains("flow_type: resume-run"));
        assert!(content.contains("resume_run_id"));
    }

    #[tokio::test]
    async fn test_create_rl_jobs_hypertune_maps_to_tune_flow() {
        let temp = TempDir::new().unwrap();
        let bundle = RlBundleTemplates::new();

        bundle.create_rl_jobs(temp.path()).await.unwrap();

        let content = tokio::fs::read_to_string(temp.path().join("configs/rl/jobs/hypertune.yml"))
            .await
            .unwrap();
        assert!(content.contains("flow_type: tune"));
        assert!(content.contains("n_trials"));
    }

    #[tokio::test]
    async fn test_append_rl_requirements_preserves_existing_content() {
        let temp = TempDir::new().unwrap();
        let bundle = RlBundleTemplates::new();

        tokio::fs::write(temp.path().join("requirements.txt"), "onnx>=1.15.0\n")
            .await
            .unwrap();

        bundle.append_rl_requirements(temp.path()).await.unwrap();

        let content = tokio::fs::read_to_string(temp.path().join("requirements.txt"))
            .await
            .unwrap();
        assert!(content.contains("onnx>=1.15.0"));
        assert!(content.contains("mecha10-rl-task-runner"));
        assert!(content.contains("mecha10-rl "));
    }

    #[tokio::test]
    async fn test_append_rl_requirements_creates_file_if_missing() {
        let temp = TempDir::new().unwrap();
        let bundle = RlBundleTemplates::new();

        bundle.append_rl_requirements(temp.path()).await.unwrap();

        let content = tokio::fs::read_to_string(temp.path().join("requirements.txt"))
            .await
            .unwrap();
        assert!(content.contains("mecha10-rl-task-runner"));
    }

    #[tokio::test]
    async fn test_create_kestra_flows_writes_readme() {
        let temp = TempDir::new().unwrap();
        let bundle = RlBundleTemplates::new();

        bundle.create_kestra_flows(temp.path()).await.unwrap();

        let content = tokio::fs::read_to_string(temp.path().join("kestra/README.md"))
            .await
            .unwrap();
        // The placeholder text is gone; the README now documents the bundled flows.
        assert!(!content.contains("not yet part of this bundle"));
        assert!(content.contains("train"));
        assert!(content.contains("env_name"));
        assert!(content.contains("algorithm_name"));
    }

    #[tokio::test]
    async fn test_create_kestra_flows_writes_all_five_flows_generically_parameterized() {
        let temp = TempDir::new().unwrap();
        let bundle = RlBundleTemplates::new();

        bundle.create_kestra_flows(temp.path()).await.unwrap();

        for file_name in [
            "train.yml",
            "tune.yml",
            "tune-then-train.yml",
            "train-from-tune.yml",
            "resume-run.yml",
        ] {
            let flow_path = temp.path().join("kestra/flows/ml-pipeline").join(file_name);
            let content = tokio::fs::read_to_string(&flow_path)
                .await
                .unwrap_or_else(|e| panic!("failed to read {}: {}", flow_path.display(), e));

            // Generically parameterized via Kestra inputs/secrets, not hardcoded to one project.
            assert!(content.contains("id: env_name"), "{file_name} missing env_name input");
            assert!(
                content.contains("id: algorithm_name"),
                "{file_name} missing algorithm_name input"
            );
            assert!(content.contains("secret("), "{file_name} missing secret() usage");
        }
    }

    #[tokio::test]
    async fn test_create_kestra_flows_train_matches_source_flow() {
        let temp = TempDir::new().unwrap();
        let bundle = RlBundleTemplates::new();

        bundle.create_kestra_flows(temp.path()).await.unwrap();

        let content = tokio::fs::read_to_string(temp.path().join("kestra/flows/ml-pipeline/train.yml"))
            .await
            .unwrap();
        assert!(content.contains("id: train"));
        assert!(content.contains("namespace: mecha10.ml-pipeline"));
    }
}