mecha10-job-validation 0.9.0

//! Shared validation for `mecha10 train`'s job catalog (`configs/rl/jobs/*.yml`) and its
//! `task_name` references into the task catalog (`configs/rl/tasks/*/task.json`).
//!
//! This crate is the single source of truth for "what's required/optional/not-applicable per
//! `flow_type`" - before this crate existed, that question was answered three different ways in
//! three different places (the CLI's job struct defaults, the CLI's dispatch-time `if` checks,
//! and mecha10-api's `create_job` handler), which could and did drift out of sync with each
//! other. The CLI (`packages/cli`) depends on this crate and calls
//! [`validate_job`]/[`task_exists`] instead of re-deriving these rules locally; `mecha10-api`'s
//! `create_job` handler (`services/api`) is intended to do the same, but as of now does not
//! depend on this crate and does not call either function.
//!
//! Deliberately Python-free: `rl_task_runner.run` (the Python runtime) does not depend on this
//! crate and does not re-validate - it trusts that a job it's asked to run already passed
//! validation upstream (CLI local dispatch, or mecha10-api's `create_job` for remote dispatch).

use std::path::Path;

/// Which training pipeline flow a job runs as.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)]
#[serde(rename_all = "kebab-case")]
pub enum FlowType {
    Train,
    Tune,
    TuneThenTrain,
    TrainFromTune,
    ResumeRun,
}

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

/// Whether a field is required, optional (sentinel/default-backed), or not applicable for a
/// given `flow_type` - e.g. `resume_run_id` is `Required` for `ResumeRun`, `NotApplicable` for
/// `Train`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FieldRequirement {
    Required,
    Optional,
    NotApplicable,
}

/// The fields this crate has an opinion about. Deliberately not every `PPOConfig` hyperparameter
/// (`learning_rate`, `gamma`, `clip_range`, ...) - those are legitimately optional/sentinel-
/// defaulted for every `flow_type` today and aren't in dispute; this covers job identity
/// (`task_name`/`algorithm`) and the fields that control how long a flow runs or which prior run
/// it targets, which is where "required vs. optional vs. n/a" actually varies by `flow_type`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JobField {
    TaskName,
    Algorithm,
    TotalTimesteps,
    ResumeRunId,
    TuneRunId,
    NTrials,
    TuneTimesteps,
    EvalEps,
}

impl JobField {
    /// The wire/YAML field name a `CreateJobRequest`/job-definition struct is expected to carry
    /// this field under - the single list a contract test in each consumer (the CLI's own
    /// `CreateJobRequest`, `mecha10-api`'s `CreateJobRequest`) checks its own hand-typed struct
    /// against, so a new variant here shows up as a failing test in both instead of silently
    /// missing a field the way `Tune`'s `n_trials`/`tune_timesteps`/`eval_eps` did before either
    /// request struct picked them up.
    pub fn wire_name(self) -> &'static str {
        match self {
            JobField::TaskName => "task_name",
            JobField::Algorithm => "algorithm",
            JobField::TotalTimesteps => "total_timesteps",
            JobField::ResumeRunId => "resume_run_id",
            JobField::TuneRunId => "tune_run_id",
            JobField::NTrials => "n_trials",
            JobField::TuneTimesteps => "tune_timesteps",
            JobField::EvalEps => "eval_eps",
        }
    }

    /// Requiredness of this field for `flow_type`, the single table the CLI reads instead of
    /// re-deriving it; `mecha10-api` is intended to read it too, once it's wired up to depend on
    /// this crate.
    pub fn requirement(self, flow_type: FlowType) -> FieldRequirement {
        use FieldRequirement::*;
        use FlowType::*;
        use JobField::*;

        match (self, flow_type) {
            (TaskName, _) => Required,
            (Algorithm, _) => Required,

            (TotalTimesteps, Train | TuneThenTrain | TrainFromTune | ResumeRun) => Required,
            (TotalTimesteps, Tune) => NotApplicable,

            (ResumeRunId, ResumeRun) => Required,
            (ResumeRunId, _) => NotApplicable,

            (TuneRunId, TrainFromTune) => Required,
            (TuneRunId, _) => NotApplicable,

            (NTrials | TuneTimesteps | EvalEps, Tune | TuneThenTrain) => Required,
            (NTrials | TuneTimesteps | EvalEps, _) => NotApplicable,
        }
    }
}

/// The subset of a job definition's fields this crate validates. The CLI's `JobDefinition`
/// (`packages/cli/src/services/train_jobs.rs`) and mecha10-api's `CreateJobRequest`
/// (`services/api/src/jobs/types.rs`) each convert their own richer type into this one (via a
/// small `From`/helper on their side) rather than this crate depending on either of them -
/// this stays the shared validation vocabulary, not the wire/storage shape.
#[derive(Debug, Clone, Default)]
pub struct JobFields {
    pub task_name: Option<String>,
    pub algorithm: Option<String>,
    pub total_timesteps: Option<i64>,
    pub resume_run_id: Option<String>,
    pub tune_run_id: Option<String>,
    pub n_trials: Option<i64>,
    pub tune_timesteps: Option<i64>,
    pub eval_eps: Option<i64>,
}

impl JobFields {
    fn get(&self, field: JobField) -> FieldPresence<'_> {
        match field {
            JobField::TaskName => FieldPresence::Str(self.task_name.as_deref()),
            JobField::Algorithm => FieldPresence::Str(self.algorithm.as_deref()),
            JobField::TotalTimesteps => FieldPresence::Int(self.total_timesteps),
            JobField::ResumeRunId => FieldPresence::Str(self.resume_run_id.as_deref()),
            JobField::TuneRunId => FieldPresence::Str(self.tune_run_id.as_deref()),
            JobField::NTrials => FieldPresence::Int(self.n_trials),
            JobField::TuneTimesteps => FieldPresence::Int(self.tune_timesteps),
            JobField::EvalEps => FieldPresence::Int(self.eval_eps),
        }
    }
}

enum FieldPresence<'a> {
    Str(Option<&'a str>),
    Int(Option<i64>),
}

impl FieldPresence<'_> {
    /// A string field counts as absent when blank/whitespace-only (matches the CLI's existing
    /// `.trim().is_empty()` checks for `algorithm`/`resume_run_id`); an int field counts as
    /// absent when unset - `-1`/`-1.0` sentinels are a *value*, not absence, and are out of
    /// scope for this presence check (a duration-controlling field left at its sentinel is a
    /// job-authoring smell, not a validation error).
    fn is_present(&self) -> bool {
        match self {
            FieldPresence::Str(s) => s.is_some_and(|s| !s.trim().is_empty()),
            FieldPresence::Int(i) => i.is_some(),
        }
    }
}

/// One validation failure: a field that's `Required` for this job's `flow_type` but absent.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MissingField {
    pub field: JobField,
    pub flow_type: FlowType,
}

impl std::fmt::Display for MissingField {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "'{:?}' is required for flow_type '{}' but was not set",
            self.field, self.flow_type
        )
    }
}

const ALL_FIELDS: &[JobField] = &[
    JobField::TaskName,
    JobField::Algorithm,
    JobField::TotalTimesteps,
    JobField::ResumeRunId,
    JobField::TuneRunId,
    JobField::NTrials,
    JobField::TuneTimesteps,
    JobField::EvalEps,
];

/// Every [`JobField`] this crate has an opinion about - the list a consumer's own contract
/// test (asserting its hand-typed `CreateJobRequest` carries a matching field for each one
/// that's applicable to a flow type it actually dispatches) iterates over.
pub fn all_fields() -> &'static [JobField] {
    ALL_FIELDS
}

/// Validate `fields` against `flow_type`'s requirement table. Returns every missing required
/// field (not just the first), so a caller can report the whole list at once.
pub fn validate_job(flow_type: FlowType, fields: &JobFields) -> Vec<MissingField> {
    ALL_FIELDS
        .iter()
        .filter(|&&field| field.requirement(flow_type) == FieldRequirement::Required)
        .filter(|&&field| !fields.get(field).is_present())
        .map(|&field| MissingField { field, flow_type })
        .collect()
}

/// Error from [`task_exists`]: `task_name` has no corresponding `task.json` under the task
/// catalog, or that file isn't valid JSON.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TaskLookupError {
    NotFound { task_name: String },
    InvalidJson { task_name: String, message: String },
}

impl std::fmt::Display for TaskLookupError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TaskLookupError::NotFound { task_name } => write!(
                f,
                "task '{task_name}' has no configs/rl/tasks/{task_name}/task.json - check the \
                 task_name against the project's task catalog"
            ),
            TaskLookupError::InvalidJson { task_name, message } => {
                write!(f, "configs/rl/tasks/{task_name}/task.json is not valid JSON: {message}")
            }
        }
    }
}

impl std::error::Error for TaskLookupError {}

/// `configs/rl/tasks` relative to a project root - kept as a single literal here (rather than a
/// dependency on `mecha10-cli-core`, to keep this crate's dependency footprint minimal for
/// mecha10-api's sake) and covered by `RL_TASKS_DIR`'s own equality test in
/// `mecha10-cli-core::paths::project`, so the two can't silently drift.
pub const RL_TASKS_DIR: &str = "configs/rl/tasks";

/// Confirm `task_name` has a real `task.json` under `{project_root}/configs/rl/tasks/` and that
/// it parses as JSON. This is the check that was previously only ever done, best-effort, deep in
/// `rl_task_runner.run` at actual training time (and even there, on failure it only logged a
/// warning and continued rather than failing) - moving it here means a bad `task_name` is caught
/// at CLI/API accept-time instead of surfacing as unexplained training behavior much later.
pub fn task_exists(project_root: &Path, task_name: &str) -> Result<(), TaskLookupError> {
    let path = project_root.join(RL_TASKS_DIR).join(task_name).join("task.json");
    let content = std::fs::read_to_string(&path).map_err(|_| TaskLookupError::NotFound {
        task_name: task_name.to_string(),
    })?;
    serde_json::from_str::<serde_json::Value>(&content).map_err(|e| TaskLookupError::InvalidJson {
        task_name: task_name.to_string(),
        message: e.to_string(),
    })?;
    Ok(())
}

/// List every task name under `{project_root}/configs/rl/tasks/` that has a `task.json` - the
/// catalog `mecha10 train`'s task picker prompts from after a job is selected. A job's own
/// `configs/rl/jobs/*.yml` no longer pins `task_name` itself (that field is filled in at
/// dispatch time from this picker instead); this is the list it prompts against. Sorted for a
/// stable menu order; empty (not an error) when the directory doesn't exist.
pub fn list_tasks(project_root: &Path) -> std::io::Result<Vec<String>> {
    let dir = project_root.join(RL_TASKS_DIR);
    if !dir.exists() {
        return Ok(Vec::new());
    }

    let mut names: Vec<String> = std::fs::read_dir(&dir)?
        .filter_map(|entry| entry.ok())
        .filter(|entry| entry.path().join("task.json").is_file())
        .filter_map(|entry| entry.file_name().into_string().ok())
        .collect();
    names.sort();
    Ok(names)
}

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

    #[test]
    fn every_job_field_has_a_unique_wire_name() {
        let names: Vec<&str> = ALL_FIELDS.iter().map(|f| f.wire_name()).collect();
        let unique: std::collections::HashSet<&str> = names.iter().copied().collect();
        assert_eq!(names.len(), unique.len());
    }

    fn fields_with_task_and_algo() -> JobFields {
        JobFields {
            task_name: Some("find_target".to_string()),
            algorithm: Some("ppo".to_string()),
            ..Default::default()
        }
    }

    #[test]
    fn train_requires_total_timesteps() {
        let fields = fields_with_task_and_algo();
        let missing = validate_job(FlowType::Train, &fields);
        assert_eq!(
            missing,
            vec![MissingField {
                field: JobField::TotalTimesteps,
                flow_type: FlowType::Train
            }]
        );
    }

    #[test]
    fn train_passes_with_total_timesteps_set() {
        let mut fields = fields_with_task_and_algo();
        fields.total_timesteps = Some(1000);
        assert!(validate_job(FlowType::Train, &fields).is_empty());
    }

    #[test]
    fn tune_does_not_require_total_timesteps() {
        let mut fields = fields_with_task_and_algo();
        fields.n_trials = Some(10);
        fields.tune_timesteps = Some(1000);
        fields.eval_eps = Some(2);
        assert!(validate_job(FlowType::Tune, &fields).is_empty());
    }

    #[test]
    fn tune_requires_tune_specific_fields() {
        let fields = fields_with_task_and_algo();
        let missing = validate_job(FlowType::Tune, &fields);
        assert_eq!(missing.len(), 3);
        assert!(missing.iter().all(|m| m.flow_type == FlowType::Tune));
    }

    #[test]
    fn resume_run_requires_resume_run_id_and_total_timesteps() {
        let fields = fields_with_task_and_algo();
        let missing = validate_job(FlowType::ResumeRun, &fields);
        assert_eq!(missing.len(), 2);
        assert!(missing.iter().any(|m| m.field == JobField::ResumeRunId));
        assert!(missing.iter().any(|m| m.field == JobField::TotalTimesteps));
    }

    #[test]
    fn train_from_tune_requires_tune_run_id() {
        let mut fields = fields_with_task_and_algo();
        fields.total_timesteps = Some(1000);
        let missing = validate_job(FlowType::TrainFromTune, &fields);
        assert_eq!(
            missing,
            vec![MissingField {
                field: JobField::TuneRunId,
                flow_type: FlowType::TrainFromTune
            }]
        );
    }

    #[test]
    fn missing_task_name_and_algorithm_reported_for_every_flow_type() {
        let missing = validate_job(FlowType::Train, &JobFields::default());
        assert!(missing.iter().any(|m| m.field == JobField::TaskName));
        assert!(missing.iter().any(|m| m.field == JobField::Algorithm));
    }

    #[test]
    fn blank_string_fields_count_as_absent() {
        let mut fields = fields_with_task_and_algo();
        fields.resume_run_id = Some("   ".to_string());
        let missing = validate_job(FlowType::ResumeRun, &fields);
        assert!(missing.iter().any(|m| m.field == JobField::ResumeRunId));
    }

    #[test]
    fn task_exists_finds_bundled_find_target_task() {
        let dir = tempfile::tempdir().unwrap();
        let task_dir = dir.path().join("configs/rl/tasks/find_target");
        std::fs::create_dir_all(&task_dir).unwrap();
        std::fs::write(task_dir.join("task.json"), r#"{"env_name": "basic_arena"}"#).unwrap();
        assert!(task_exists(dir.path(), "find_target").is_ok());
    }

    #[test]
    fn task_exists_reports_missing_task() {
        let dir = tempfile::tempdir().unwrap();
        std::fs::create_dir_all(dir.path().join("configs/rl/tasks")).unwrap();
        let err = task_exists(dir.path(), "nonexistent").unwrap_err();
        assert_eq!(
            err,
            TaskLookupError::NotFound {
                task_name: "nonexistent".to_string()
            }
        );
    }

    #[test]
    fn list_tasks_returns_empty_when_directory_missing() {
        let dir = tempfile::tempdir().unwrap();
        assert_eq!(list_tasks(dir.path()).unwrap(), Vec::<String>::new());
    }

    #[test]
    fn list_tasks_lists_directories_with_a_task_json_sorted() {
        let dir = tempfile::tempdir().unwrap();
        for name in ["task2", "find_target", "task1"] {
            let task_dir = dir.path().join("configs/rl/tasks").join(name);
            std::fs::create_dir_all(&task_dir).unwrap();
            std::fs::write(task_dir.join("task.json"), "{}").unwrap();
        }
        std::fs::create_dir_all(dir.path().join("configs/rl/tasks/not-a-task")).unwrap();

        assert_eq!(
            list_tasks(dir.path()).unwrap(),
            vec!["find_target".to_string(), "task1".to_string(), "task2".to_string()]
        );
    }

    #[test]
    fn task_exists_reports_invalid_json() {
        let dir = tempfile::tempdir().unwrap();
        let task_dir = dir.path().join("configs/rl/tasks/broken");
        std::fs::create_dir_all(&task_dir).unwrap();
        std::fs::write(task_dir.join("task.json"), "not json").unwrap();
        let err = task_exists(dir.path(), "broken").unwrap_err();
        assert!(matches!(err, TaskLookupError::InvalidJson { .. }));
    }
}