taktwerk 0.6.0

//! `taktwerk campaign`: plan the runs of a campaign file, print the plan, and with `--write` run
//! them, each a `taktwerk run` subprocess with its own scenario and record directory; resume
//! by skipping runs already done with the same spec.
//!
//! ```text
//! <out>/campaign.json            the plan: file, hashes, every run's values, seed and spec hash
//! <out>/runs.csv                 one row per run: values, seed, status, stop, cycles, wall time
//! <out>/runs/<name>/run.toml     the run's effective scenario; `taktwerk run --scenario` reruns it
//! <out>/runs/<name>/status.json  status, exit code, stop, wall time, sha256 of every record file
//! <out>/runs/<name>/stderr.log   the run's log
//! <out>/runs/<name>/<project>/<recorder>/<run-id>/   the run's records
//! ```

mod plan;
mod sample;

use std::collections::{BTreeMap, VecDeque};
use std::fmt::Write as _;
use std::path::{Path, PathBuf};
use std::process::{Command, ExitCode, Stdio};
use std::sync::{Arc, Mutex};
use std::time::Instant;

use anyhow::{Context as _, bail};
use serde::{Deserialize, Serialize};
use taktwerk_core::project::Project;
use taktwerk_core::scenario::ScenarioFile;
use taktwerk_record::util;

use crate::setup::model_path;
pub use plan::CampaignFile;

/// Format of `campaign.json` and `status.json`.
const FORMAT: u32 = 1;

/// Reads scenario files from disk.
struct Disk;

impl plan::Files for Disk {
    fn scenario(&self, path: &Path) -> Result<ScenarioFile, String> {
        let text = std::fs::read_to_string(path)
            .map_err(|e| format!("scenario {}: {e}", path.display()))?;
        toml::from_str(&text).map_err(|e| format!("scenario {}: {e}", path.display()))
    }
}

/// `campaign.json`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Manifest {
    /// [`FORMAT`].
    pub format: u32,
    /// taktwerk version that planned it.
    pub taktwerk: String,
    /// The campaign file.
    pub campaign: String,
    /// `sha256:<hex>` of it.
    pub campaign_hash: String,
    /// The project, absolute.
    pub project: String,
    /// `sha256:<hex>` of it.
    pub project_hash: String,
    /// `--speed` of every run.
    pub speed: String,
    /// The campaign's seed.
    pub seed: u64,
    /// Every run, in plan order.
    pub runs: Vec<RunEntry>,
}

/// One run of `campaign.json`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RunEntry {
    /// Name and directory under `runs/`.
    pub name: String,
    /// Case.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub case: Option<String>,
    /// Grid point.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub grid: Option<usize>,
    /// Sample.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub sample: Option<usize>,
    /// Replicate.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub replicate: Option<usize>,
    /// Seed.
    pub seed: u64,
    /// The values its axes set.
    pub values: BTreeMap<String, toml::Value>,
    /// `sha256:<hex>` over the taktwerk version, the project and its model packages and the run's
    /// scenario: a run whose spec hash changed is run again.
    pub spec_hash: String,
}

/// `status.json` of a run.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Status {
    /// [`FORMAT`].
    pub format: u32,
    /// The run.
    pub name: String,
    /// Its spec hash when it ran.
    pub spec_hash: String,
    /// `ok`: it ended on a limit or condition; `failed`: anything else.
    pub status: String,
    /// The process's exit code; absent when it was killed.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub exit_code: Option<i32>,
    /// Why it stopped, as `events.csv` words it.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub stop: Option<String>,
    /// Ticks it ran.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cycles: Option<u64>,
    /// Wall time, seconds.
    pub wall_s: f64,
    /// Path under the run directory → `sha256:<hex>`, every file but this one and the log.
    pub files: BTreeMap<String, String>,
}

/// How `taktwerk campaign` was asked to run.
#[derive(Debug, Clone, Default)]
pub struct Options {
    /// Run, rather than print the plan.
    pub write: bool,
    /// Start over when the plan changed, and run every run again.
    pub force: bool,
    /// Runs at a time, instead of the file's `jobs`.
    pub jobs: Option<usize>,
}

/// Plan `file`, print the plan; with `write`, run what is not done.
///
/// # Errors
/// The file, its project or a scenario cannot be read; the plan is bad; the output directory
/// holds another plan and `force` is not set; a run cannot be started.
pub fn campaign(file: &Path, options: &Options) -> anyhow::Result<ExitCode> {
    let file = std::path::absolute(file)?;
    let text =
        std::fs::read_to_string(&file).with_context(|| format!("read {}", file.display()))?;
    let spec: CampaignFile =
        toml::from_str(&text).with_context(|| format!("campaign {}", file.display()))?;
    let dir = file
        .parent()
        .unwrap_or_else(|| Path::new("/"))
        .to_path_buf();
    crate::cli::parse_speed(&spec.speed).map_err(|e| anyhow::anyhow!("speed: {e}"))?;
    let runs = plan::plan(&spec, &dir, &Disk).map_err(anyhow::Error::msg)?;
    let project_file = dir.join(&spec.project);
    let project_text = std::fs::read(&project_file)
        .with_context(|| format!("read the project {}", project_file.display()))?;
    let project = Project::load(&project_file)?;
    let mut base = String::new();
    let _ = write!(
        base,
        "{} {}",
        env!("CARGO_PKG_VERSION"),
        util::sha256(&project_text)
    );
    for (id, m) in &project.models {
        let path = model_path(&project_file, &m.path);
        let hash = util::sha256_package(&path)
            .with_context(|| format!("hash model `{id}` at {}", path.display()))?;
        let _ = write!(base, " {id}={hash}");
    }
    let stem = file
        .file_stem()
        .map_or_else(|| "campaign".into(), |s| s.to_string_lossy().into_owned());
    let out = dir.join(spec.out.clone().unwrap_or_else(|| PathBuf::from(stem)));
    let mut entries = Vec::with_capacity(runs.len());
    let mut tomls = Vec::with_capacity(runs.len());
    for r in &runs {
        let run_toml = toml::to_string(&r.scenario).context("a run's scenario")?;
        entries.push(RunEntry {
            name: r.name.clone(),
            case: r.case.clone(),
            grid: r.grid,
            sample: r.sample,
            replicate: r.replicate,
            seed: r.seed,
            values: r.values.iter().cloned().collect(),
            spec_hash: util::sha256(format!("{base}\n{run_toml}").as_bytes()),
        });
        tomls.push(run_toml);
    }
    let manifest = Manifest {
        format: FORMAT,
        taktwerk: env!("CARGO_PKG_VERSION").to_owned(),
        campaign: file.display().to_string(),
        campaign_hash: util::sha256(text.as_bytes()),
        project: project_file.display().to_string(),
        project_hash: util::sha256(&project_text),
        speed: spec.speed.clone(),
        seed: spec.seed,
        runs: entries,
    };
    let done: Vec<bool> = manifest
        .runs
        .iter()
        .map(|e| {
            read_status(&out, &e.name)
                .is_some_and(|s| s.status == "ok" && s.spec_hash == e.spec_hash)
        })
        .collect();
    let jobs = options.jobs.unwrap_or(spec.jobs).max(1);
    say(&plan_table(&manifest, &done, &out, jobs));
    if !options.write {
        say("dry run: nothing was run; --write runs the campaign\n");
        return Ok(ExitCode::SUCCESS);
    }

    let manifest_path = out.join("campaign.json");
    if let Ok(old) = std::fs::read_to_string(&manifest_path)
        && let Ok(old) = serde_json::from_str::<Manifest>(&old)
        && old
            .runs
            .iter()
            .map(|r| &r.name)
            .ne(manifest.runs.iter().map(|r| &r.name))
        && !options.force
    {
        bail!(
            "{} holds another plan (other runs); --force starts over",
            out.display()
        );
    }
    if options.force {
        let runs_dir = out.join("runs");
        if runs_dir.exists() {
            std::fs::remove_dir_all(&runs_dir)
                .with_context(|| format!("remove {}", runs_dir.display()))?;
        }
    }
    std::fs::create_dir_all(out.join("runs"))?;
    write_atomic(&manifest_path, &serde_json::to_string_pretty(&manifest)?)?;

    let pending: VecDeque<usize> = (0..runs.len())
        .filter(|&i| options.force || !done[i])
        .collect();
    let total = pending.len();
    let queue = Arc::new(Mutex::new(pending));
    let finished = Arc::new(Mutex::new(0usize));
    let exe = std::env::current_exe().context("the taktwerk binary")?;
    let job = Job {
        exe,
        project: project_file,
        speed: spec.speed.clone(),
        out: out.clone(),
    };
    std::thread::scope(|scope| {
        for _ in 0..jobs.min(total.max(1)) {
            let queue = Arc::clone(&queue);
            let finished = Arc::clone(&finished);
            let (job, manifest, tomls) = (&job, &manifest, &tomls);
            scope.spawn(move || {
                loop {
                    let next = queue.lock().map(|mut q| q.pop_front()).ok().flatten();
                    let Some(i) = next else { break };
                    let entry = &manifest.runs[i];
                    let line = match job.run(entry, &tomls[i]) {
                        Ok(s) => format!(
                            "{:<7} {:<24} {:>8.1} s  {}",
                            s.status,
                            entry.name,
                            s.wall_s,
                            s.stop.as_deref().unwrap_or("-")
                        ),
                        Err(e) => format!("error   {:<24} {e:#}", entry.name),
                    };
                    let n = finished.lock().map(|mut n| {
                        *n += 1;
                        *n
                    });
                    say(&format!("[{}/{total}] {line}\n", n.unwrap_or(0)));
                }
            });
        }
    });

    let statuses: Vec<Option<Status>> = manifest
        .runs
        .iter()
        .map(|e| read_status(&out, &e.name))
        .collect();
    write_atomic(&out.join("runs.csv"), &runs_csv(&manifest, &statuses))?;
    let ok = manifest
        .runs
        .iter()
        .zip(&statuses)
        .filter(|(e, s)| {
            s.as_ref()
                .is_some_and(|s| s.status == "ok" && s.spec_hash == e.spec_hash)
        })
        .count();
    say(&format!(
        "{ok} of {} runs ok; {}\n",
        manifest.runs.len(),
        out.join("runs.csv").display()
    ));
    Ok(if ok == manifest.runs.len() {
        ExitCode::SUCCESS
    } else {
        ExitCode::FAILURE
    })
}

/// What every run of a campaign shares.
struct Job {
    exe: PathBuf,
    project: PathBuf,
    speed: String,
    out: PathBuf,
}

impl Job {
    /// Run `entry` with scenario `run_toml` in a fresh directory and write its `status.json`.
    fn run(&self, entry: &RunEntry, run_toml: &str) -> anyhow::Result<Status> {
        let dir = self.out.join("runs").join(&entry.name);
        if dir.exists() {
            std::fs::remove_dir_all(&dir).with_context(|| format!("clear {}", dir.display()))?;
        }
        std::fs::create_dir_all(&dir)?;
        let scenario = dir.join("run.toml");
        std::fs::write(&scenario, run_toml)?;
        let log = std::fs::File::create(dir.join("stderr.log"))?;
        let started = Instant::now();
        let output = Command::new(&self.exe)
            .arg("run")
            .arg(&self.project)
            .args(["--speed", &self.speed, "--no-connectors", "--scenario"])
            .arg(&scenario)
            .arg("--record-dir")
            .arg(&dir)
            .stdin(Stdio::null())
            .stderr(log)
            .output()
            .with_context(|| format!("start {}", self.exe.display()))?;
        let wall_s = started.elapsed().as_secs_f64();
        let stdout = String::from_utf8_lossy(&output.stdout);
        let (cycles, stop) = summary_of(&stdout);
        let ok = output.status.success()
            && stop
                .as_deref()
                .is_some_and(|s| s != "requested" && s != "failed");
        let mut files = BTreeMap::new();
        hash_files(&dir, &dir, &mut files)?;
        let status = Status {
            format: FORMAT,
            name: entry.name.clone(),
            spec_hash: entry.spec_hash.clone(),
            status: if ok { "ok" } else { "failed" }.to_owned(),
            exit_code: output.status.code(),
            stop,
            cycles,
            wall_s,
            files,
        };
        write_atomic(
            &dir.join("status.json"),
            &serde_json::to_string_pretty(&status)?,
        )?;
        Ok(status)
    }
}

/// Ticks and stop reason from `taktwerk run`'s summary line, `cycles N ... stop <reason>`.
fn summary_of(stdout: &str) -> (Option<u64>, Option<String>) {
    let Some(line) = stdout.lines().rev().find(|l| l.starts_with("cycles ")) else {
        return (None, None);
    };
    let cycles = line.split_whitespace().nth(1).and_then(|n| n.parse().ok());
    let stop = line.split_once("  stop ").map(|(_, r)| r.trim().to_owned());
    (cycles, stop)
}

/// `sha256` of every file under `dir` but `status.json` and `stderr.log`, by path relative to
/// `root`.
fn hash_files(root: &Path, dir: &Path, out: &mut BTreeMap<String, String>) -> anyhow::Result<()> {
    let mut entries: Vec<_> = std::fs::read_dir(dir)?.collect::<Result<_, _>>()?;
    entries.sort_by_key(std::fs::DirEntry::file_name);
    for e in entries {
        let path = e.path();
        if e.file_type()?.is_dir() {
            hash_files(root, &path, out)?;
            continue;
        }
        let rel = path
            .strip_prefix(root)
            .unwrap_or(&path)
            .to_string_lossy()
            .into_owned();
        if rel == "status.json" || rel == "stderr.log" {
            continue;
        }
        out.insert(rel, util::sha256(&std::fs::read(&path)?));
    }
    Ok(())
}

fn read_status(out: &Path, name: &str) -> Option<Status> {
    let text = std::fs::read_to_string(out.join("runs").join(name).join("status.json")).ok()?;
    serde_json::from_str(&text).ok()
}

/// Print `text` to stdout; a closed stdout (`| head`) is no error, the campaign goes on.
fn say(text: &str) {
    use std::io::Write as _;
    let mut out = std::io::stdout().lock();
    let _ = out.write_all(text.as_bytes()).and_then(|()| out.flush());
}

/// Write `text` to `path` through a temporary file, so a reader never sees half of it.
fn write_atomic(path: &Path, text: &str) -> anyhow::Result<()> {
    let tmp = path.with_extension("tmp");
    std::fs::write(&tmp, text).with_context(|| format!("write {}", tmp.display()))?;
    std::fs::rename(&tmp, path).with_context(|| format!("write {}", path.display()))?;
    Ok(())
}

/// The plan as a table, with what is done.
fn plan_table(m: &Manifest, done: &[bool], out: &Path, jobs: usize) -> String {
    let keys = value_keys(m);
    let mut s = format!(
        "campaign {}\nproject  {}\nout      {}\nruns     {} ({} done), {jobs} at a time, speed {}\n\n",
        m.campaign,
        m.project,
        out.display(),
        m.runs.len(),
        done.iter().filter(|d| **d).count(),
        m.speed
    );
    let mut table: Vec<Vec<String>> = vec![
        ["run", "seed"]
            .into_iter()
            .map(str::to_owned)
            .chain(keys.iter().cloned())
            .chain(["done".to_owned()])
            .collect(),
    ];
    for (r, d) in m.runs.iter().zip(done) {
        let mut row = vec![r.name.clone(), r.seed.to_string()];
        row.extend(keys.iter().map(|k| {
            r.values
                .get(k)
                .map_or_else(|| "-".to_owned(), ToString::to_string)
        }));
        row.push(if *d { "yes" } else { "no" }.to_owned());
        table.push(row);
    }
    let widths: Vec<usize> = (0..table[0].len())
        .map(|c| table.iter().map(|r| r[c].len()).max().unwrap_or(0))
        .collect();
    for row in &table {
        let cells: Vec<String> = row
            .iter()
            .zip(&widths)
            .map(|(cell, w)| format!("{cell:<w$}"))
            .collect();
        let _ = writeln!(s, "{}", cells.join("  ").trim_end());
    }
    s
}

/// Every key some run's axes set, in order of first appearance.
fn value_keys(m: &Manifest) -> Vec<String> {
    let mut keys: Vec<String> = Vec::new();
    for r in &m.runs {
        for k in r.values.keys() {
            if !keys.contains(k) {
                keys.push(k.clone());
            }
        }
    }
    keys
}

/// `runs.csv`: one row per run in plan order.
fn runs_csv(m: &Manifest, statuses: &[Option<Status>]) -> String {
    let keys = value_keys(m);
    let mut s = String::from("name,case,grid,sample,replicate,seed,status,stop,cycles,wall_s");
    for k in &keys {
        s.push(',');
        s.push_str(&csv_field(k));
    }
    s.push('\n');
    let opt = |v: Option<usize>| v.map(|v| v.to_string()).unwrap_or_default();
    for (r, st) in m.runs.iter().zip(statuses) {
        let current = st.as_ref().filter(|s| s.spec_hash == r.spec_hash);
        let _ = write!(
            s,
            "{},{},{},{},{},{},{},{},{},{}",
            csv_field(&r.name),
            csv_field(r.case.as_deref().unwrap_or("")),
            opt(r.grid),
            opt(r.sample),
            opt(r.replicate),
            r.seed,
            current.map_or("pending", |s| s.status.as_str()),
            csv_field(current.and_then(|s| s.stop.as_deref()).unwrap_or("")),
            current
                .and_then(|s| s.cycles)
                .map(|c| c.to_string())
                .unwrap_or_default(),
            current.map(|s| s.wall_s.to_string()).unwrap_or_default(),
        );
        for k in &keys {
            s.push(',');
            if let Some(v) = r.values.get(k) {
                let text = match v {
                    toml::Value::String(t) => t.clone(),
                    other => other.to_string(),
                };
                s.push_str(&csv_field(&text));
            }
        }
        s.push('\n');
    }
    s
}

/// A CSV field, quoted when it holds a comma, quote or line break.
fn csv_field(text: &str) -> String {
    if text.contains([',', '"', '\n', '\r']) {
        format!("\"{}\"", text.replace('"', "\"\""))
    } else {
        text.to_owned()
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, reason = "tests")]
mod tests {
    use super::*;

    #[test]
    fn the_summary_line_gives_cycles_and_stop() {
        let (c, s) = summary_of(
            "cycles 600  published 600  overruns 0  stale 0  cycle mean 0.1 ms  max 1 ms  stop condition a.x > 1\n",
        );
        assert_eq!((c, s.as_deref()), (Some(600), Some("condition a.x > 1")));
        assert_eq!(summary_of("nothing"), (None, None));
    }

    #[test]
    fn csv_fields_are_quoted_when_needed() {
        assert_eq!(csv_field("a"), "a");
        assert_eq!(csv_field("[1.0, 2.0]"), "\"[1.0, 2.0]\"");
        assert_eq!(csv_field("a\"b"), "\"a\"\"b\"");
    }
}