mecha10-templates 0.8.2

//! RL training bundle template generation.
//!
//! Packages the standard `configs/rl/jobs/` job catalog and the RL bundle's Python

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

const RL_JOBS_DIR: &str = "configs/rl/jobs";

/// `(destination file name, source file name under `<assets_root>/rl/jobs/`)` for the standard
/// job catalog.
const RL_JOBS: [(&str, &str); 3] = [
    ("train.yml", "train.yml.template"),
    ("train-from-checkpoint.yml", "train-from-checkpoint.yml.template"),
    ("hypertune.yml", "hypertune.yml.template"),
];

pub struct RlBundleTemplates;

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

    /// Write the standard job catalog (`train`, `train-from-checkpoint`, `hypertune`) to
    /// `configs/rl/jobs/`, the menu `mecha10 train` lists by default, reading each job template
    /// from the freshly-fetched `mecha10-assets` extraction at `assets_root`.
    pub async fn create_rl_jobs(&self, path: &Path, assets_root: &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, template_name) in RL_JOBS {
            let template_path = assets_root.join("rl/jobs").join(template_name);
            let content = tokio::fs::read_to_string(&template_path)
                .await
                .with_context(|| format!("Failed to read {:?} from downloaded mecha10-assets", template_path))?;

            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`, reading
    /// them from the freshly-fetched `mecha10-assets` extraction at `assets_root`. Assumes
    /// `requirements.txt` already exists (i.e. `create_requirements_txt` has run).
    pub async fn append_rl_requirements(&self, path: &Path, assets_root: &Path) -> Result<()> {
        let requirements_path = path.join(paths::meta::REQUIREMENTS_TXT);

        let template_path = assets_root.join("rl/requirements-rl.txt.template");
        let rl_requirements = tokio::fs::read_to_string(&template_path)
            .await
            .with_context(|| format!("Failed to read {:?} from downloaded mecha10-assets", template_path))?;

        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(())
    }
}

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

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

    const FIXTURE_TRAIN: &str = r#"name: train
description: Train PPO from scratch
flow_type: train

algorithm: ppo

seed: 0
environment_map: ""

policy: ""
total_timesteps: 1000
learning_rate: 0.0003
linear_lr_schedule: true
n_steps: 256
batch_size: 64
n_epochs: 10
gamma: 0.99
gae_lambda: 0.95
clip_range: 0.2
ent_coef: 0.01
vf_coef: 0.5
max_grad_norm: 0.5

n_envs: 1

patience: -1
min_delta: -1.0
"#;

    const FIXTURE_TRAIN_FROM_CHECKPOINT: &str = r#"name: train-from-checkpoint
description: Resume training from a prior run's checkpoint
flow_type: resume-run

algorithm: ppo

total_timesteps: 1000
patience: -1
min_delta: -1.0

n_envs: 1
"#;

    const FIXTURE_HYPERTUNE: &str = r#"name: hypertune
description: Optuna hyperparameter search over learning rate / rollout steps / epochs
flow_type: tune

algorithm: ppo

n_trials: 2
tune_timesteps: 1000
eval_eps: 1

total_timesteps: 10000
n_steps: 2048

n_envs: 1

search_space:
  - name: learning_rate
    type: float
    low: 0.00001
    high: 0.001
    log: true
"#;

    const FIXTURE_RL_REQUIREMENTS: &str = "\nmecha10-rl-task-runner @ git+ssh://git@github.com/mecha-industries/mecha10-rl.git@main#subdirectory=rl-task-runner\nmecha10-rl @ git+ssh://git@github.com/mecha-industries/mecha10-rl.git@main#subdirectory=rl-trainer\n";

    /// Stage a fake extracted `mecha10-assets` root containing the standard RL job templates
    /// and RL requirements template, mirroring the real `mecha10-assets` layout
    /// (`rl/jobs/*.yml.template`, `rl/requirements-rl.txt.template`).
    fn stage_assets_root() -> TempDir {
        let assets_root = TempDir::new().unwrap();

        let jobs_dir = assets_root.path().join("rl/jobs");
        std::fs::create_dir_all(&jobs_dir).unwrap();
        std::fs::write(jobs_dir.join("train.yml.template"), FIXTURE_TRAIN).unwrap();
        std::fs::write(
            jobs_dir.join("train-from-checkpoint.yml.template"),
            FIXTURE_TRAIN_FROM_CHECKPOINT,
        )
        .unwrap();
        std::fs::write(jobs_dir.join("hypertune.yml.template"), FIXTURE_HYPERTUNE).unwrap();

        std::fs::write(
            assets_root.path().join("rl/requirements-rl.txt.template"),
            FIXTURE_RL_REQUIREMENTS,
        )
        .unwrap();

        assets_root
    }

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

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

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

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

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

        for file_name in ["train.yml", "train-from-checkpoint.yml", "hypertune.yml"] {
            let content = tokio::fs::read_to_string(temp.path().join("configs/rl/jobs").join(file_name))
                .await
                .unwrap();
            assert!(content.contains("n_envs: 1"), "{file_name} missing n_envs: 1");
        }
    }

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

        bundle.create_rl_jobs(temp.path(), assets_root.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("algorithm: ppo"), "{file_name} missing algorithm: ppo");
        }
    }

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

        bundle.create_rl_jobs(temp.path(), assets_root.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 assets_root = stage_assets_root();
        let bundle = RlBundleTemplates::new();

        bundle.create_rl_jobs(temp.path(), assets_root.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 assets_root = stage_assets_root();
        let bundle = RlBundleTemplates::new();

        bundle.create_rl_jobs(temp.path(), assets_root.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_create_rl_jobs_do_not_pin_a_task_name() {
        let temp = TempDir::new().unwrap();
        let assets_root = stage_assets_root();
        let bundle = RlBundleTemplates::new();

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

        for file in ["train.yml", "train-from-checkpoint.yml", "hypertune.yml"] {
            let content = tokio::fs::read_to_string(temp.path().join("configs/rl/jobs").join(file))
                .await
                .unwrap();
            assert!(!content.contains("task_name"), "{file} should not pin a task_name");
        }
    }

    #[tokio::test]
    async fn test_append_rl_requirements_preserves_existing_content() {
        let temp = TempDir::new().unwrap();
        let assets_root = stage_assets_root();
        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(), assets_root.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 assets_root = stage_assets_root();
        let bundle = RlBundleTemplates::new();

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

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