henad_explore/

merge.rs

//! Merges of shard directories into one directory, holding the files one sweep of the whole plan would have written.
//!
//! The inputs are shards of one plan run with one replicate count, each with its own shard index. `runs.csv`
//! and `series.csv` are merged by run id and written together, as a resume replaces them. `summary.csv` is rebuilt,
//! and the manifest lists the merged directories.

use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use std::path::{Path, PathBuf};

use henad_core::explore::plan::Shard;

use crate::output::manifest::{
    BuildRole, Manifest, ManifestError, ManifestMode, ManifestStatus, ManifestTimestamps, RecordedBuild, ResultCounts,
    now_unix_ms, rfc3339,
};
use crate::output::read::{ReadError, RunRecord, RunsCsv, SeriesScan, SeriesSegment, merge_series};
use crate::output::runs_csv::header_line;
use crate::output::{MANIFEST_FILE, OutputDir, OutputError, RUNS_FILE, SERIES_FILE, table_paths};
use crate::progress::{Progress, ProgressEvent};
use crate::sweep::SweepWarning;

/// Maximum number of missing run ids a [`SweepWarning::MissingRuns`] lists.
pub const MAX_LISTED_RUNS: usize = 5;

/// Result of a merge.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MergeReport {
    /// Rows of the merged `runs.csv`, by status.
    pub counts: ResultCounts,
    /// Number of runs in the plan that no input holds.
    pub missing: u64,
    /// Directory that the merged results were written to.
    pub output_dir: PathBuf,
}

/// One shard directory read for a merge.
struct ShardInput {
    dir: PathBuf,
    manifest: Manifest,
    shard: Shard,
    runs: RunsCsv,
    series: SeriesScan,
    /// Ids of the runs `runs.csv` holds.
    run_ids: BTreeSet<u64>,
}

/// Merges the shard directories `shard_dirs` into `output_dir`. The output directory must hold no results.
///
/// Series rows of a run with no `runs.csv` row, and a partial last record, are left out. A plan run that no input
/// holds is reported as a [`ProgressEvent::Warned`] and marks the merged manifest `incomplete`. A resume of the merged
/// directory runs it. Shards whose sessions ran different engine or model builds produce a
/// [`SweepWarning::BuildChanged`] for each build in another shard that matches no build in the lowest shard
/// that records builds. A merge that fails once its manifest is written marks the manifest `failed` when it can.
///
/// # Errors
///
/// Returns [`MergeError`] when there are no inputs, an input cannot be read, the inputs hold different plans,
/// replicate counts, shard counts or columns, two inputs hold the same shard, a run's id does not match its config and
/// replicate, or the merged directory cannot be written.
pub fn merge(
    shard_dirs: &[PathBuf],
    output_dir: &Path,
    progress: &mut dyn Progress,
) -> Result<MergeReport, MergeError> {
    let mut shards = shard_dirs
        .iter()
        .map(|dir| read_input(dir))
        .collect::<Result<Vec<_>, _>>()?;
    shards.sort_by_key(|input| input.shard.index());
    let Some(first) = shards.first() else {
        return Err(MergeError::NoInputs);
    };
    check_shards(&shards)?;
    for warning in build_warnings(&shards) {
        progress.report(&ProgressEvent::Warned(&warning));
    }
    let (runs_header, series_header) = common_headers(&shards)?;
    let plan_runs = first.manifest.plan.runs;
    let (records, counts) = merged_records(&shards, plan_runs, first.manifest.plan.replicates)?;

    let dir = OutputDir::create(output_dir).map_err(MergeError::Output)?;
    let mut manifest = merged_manifest(&shards, shard_dirs);
    dir.write_manifest(&manifest).map_err(MergeError::Output)?;
    if let Err(error) = write_tables(&dir, &shards, &records, &runs_header, &series_header) {
        manifest.fail(now_unix_ms());
        // The merge's own error is the one to report. A manifest that cannot be written says `running`.
        drop(dir.write_manifest(&manifest));
        return Err(MergeError::Output(error));
    }

    let missing = plan_runs - records.len() as u64;
    if missing > 0 {
        let first_missing = (0..plan_runs)
            .filter(|run_id| !records.contains_key(run_id))
            .take(MAX_LISTED_RUNS)
            .collect();
        progress.report(&ProgressEvent::Warned(&SweepWarning::MissingRuns {
            count: missing,
            first: first_missing,
        }));
    }
    let status = if missing == 0 {
        ManifestStatus::Complete
    } else {
        ManifestStatus::Incomplete
    };
    // `merged_manifest` credited every session of every shard, so the counts go in without `Manifest::finish`.
    let finished = now_unix_ms();
    manifest.status = status;
    manifest.results = Some(counts);
    manifest.timestamps.finished_unix_ms = Some(finished);
    manifest.timestamps.finished = Some(rfc3339(finished));
    dir.write_manifest(&manifest).map_err(MergeError::Output)?;
    Ok(MergeReport {
        counts,
        missing,
        output_dir: output_dir.to_owned(),
    })
}

/// Writes the merged `runs.csv` and `series.csv` into `dir` through [`OutputDir::replace_tables`], then rebuilds
/// `summary.csv`.
fn write_tables(
    dir: &OutputDir,
    shards: &[ShardInput],
    records: &BTreeMap<u64, &RunRecord>,
    runs_header: &str,
    series_header: &str,
) -> Result<(), OutputError> {
    let segments: Vec<SeriesSegment> = shards
        .iter()
        .enumerate()
        .flat_map(|(input_index, input)| {
            input.series.segments.iter().map(move |range| SeriesSegment {
                path: input.series.path.clone(),
                range: range.clone(),
                input_index,
            })
        })
        .collect();
    dir.replace_tables(
        |dest| {
            dest.write_all(runs_header.as_bytes())?;
            records
                .values()
                .try_for_each(|record| dest.write_all(record.text.as_bytes()))
        },
        |dest| {
            merge_series(dest, series_header, &segments, |input_index, run_id| {
                shards[input_index].run_ids.contains(&run_id).then_some(run_id)
            })
        },
    )?;
    dir.write_summary()
}

/// Returns the header lines of `runs.csv` and `series.csv` that every input with those tables shares.
fn common_headers(shards: &[ShardInput]) -> Result<(String, String), MergeError> {
    let runs_header = shards
        .iter()
        .find_map(|input| input.runs.header.as_ref())
        .map(|header| header_line(header));
    let series_header = shards.iter().find_map(|input| input.series.header.clone());
    let (Some(runs_header), Some(series_header)) = (runs_header, series_header) else {
        return Err(MergeError::NoTables);
    };
    for input in shards {
        let differs = |file| MergeError::ColumnsDiffer {
            dir: input.dir.clone(),
            file,
        };
        if input
            .runs
            .header
            .as_ref()
            .is_some_and(|header| header_line(header) != runs_header)
        {
            return Err(differs(RUNS_FILE));
        }
        if input
            .series
            .header
            .as_ref()
            .is_some_and(|header| *header != series_header)
        {
            return Err(differs(SERIES_FILE));
        }
    }
    Ok((runs_header, series_header))
}

/// Returns the records of every input by run id, and their counts by status.
///
/// A record whose run id is not `config_id * replicates + rep` is rejected.
fn merged_records(
    shards: &[ShardInput],
    plan_runs: u64,
    replicates: u64,
) -> Result<(BTreeMap<u64, &RunRecord>, ResultCounts), MergeError> {
    let mut records = BTreeMap::new();
    let mut counts = ResultCounts::default();
    for input in shards {
        for record in &input.runs.records {
            let numbered = record
                .config_id
                .checked_mul(replicates)
                .and_then(|first| first.checked_add(record.rep));
            if record.rep >= replicates || numbered != Some(record.run_id) {
                return Err(MergeError::RunIdDiffers {
                    dir: input.dir.clone(),
                    run_id: record.run_id,
                    config_id: record.config_id,
                    rep: record.rep,
                });
            }
            if !input.shard.contains(record.run_id) || record.run_id >= plan_runs {
                return Err(MergeError::OutsideShard {
                    dir: input.dir.clone(),
                    run_id: record.run_id,
                });
            }
            if records.insert(record.run_id, record).is_some() {
                return Err(MergeError::DuplicateRun {
                    dir: input.dir.clone(),
                    run_id: record.run_id,
                });
            }
            counts.count(record.status);
        }
    }
    Ok((records, counts))
}

/// Reads the manifest and tables of the shard directory `dir`.
fn read_input(dir: &Path) -> Result<ShardInput, MergeError> {
    let manifest = Manifest::read(&dir.join(MANIFEST_FILE)).map_err(MergeError::Manifest)?;
    if manifest.mode != ManifestMode::Sweep {
        return Err(MergeError::NotASweep { dir: dir.to_owned() });
    }
    let shard = manifest
        .shard
        .to_shard()
        .ok_or_else(|| MergeError::BadShard { dir: dir.to_owned() })?;
    let (runs_path, series_path) = table_paths(dir);
    let runs = RunsCsv::read(&runs_path).map_err(MergeError::Table)?;
    let run_ids: BTreeSet<u64> = runs.records.iter().map(|record| record.run_id).collect();
    let series = SeriesScan::read(&series_path, |run_id| run_ids.contains(&run_id)).map_err(MergeError::Table)?;
    Ok(ShardInput {
        dir: dir.to_owned(),
        manifest,
        shard,
        runs,
        series,
        run_ids,
    })
}

/// Checks that `shards`, sorted by shard index, are distinct shards of one plan with one replicate count.
fn check_shards(shards: &[ShardInput]) -> Result<(), MergeError> {
    let Some(first) = shards.first() else {
        return Err(MergeError::NoInputs);
    };
    for pair in shards.windows(2) {
        if pair[0].shard == pair[1].shard {
            return Err(MergeError::SameShard {
                shard: pair[0].shard,
                first: pair[0].dir.clone(),
                second: pair[1].dir.clone(),
            });
        }
    }
    for input in shards {
        let (plan, expected) = (&input.manifest.plan, &first.manifest.plan);
        if plan.plan_hash != expected.plan_hash
            || input.manifest.model.schema_hash != first.manifest.model.schema_hash
            || plan.replicates != expected.replicates
            || plan.runs != expected.runs
        {
            return Err(MergeError::PlanDiffers {
                dir: input.dir.clone(),
                first: first.dir.clone(),
            });
        }
        if input.shard.count() != first.shard.count() {
            return Err(MergeError::ShardCountDiffers {
                dir: input.dir.clone(),
                first: first.dir.clone(),
            });
        }
    }
    Ok(())
}

/// Returns a [`SweepWarning::BuildChanged`] for each engine or model build of another shard that matches no build of
/// the reference shard, the lowest shard that records a build for the role.
///
/// The warning sets `recorded` to the reference shard's build that is [treated as](RecordedBuild::reads_as) the other
/// shard's build, or else to the reference shard's first build. Builds of other shards that are treated as the same
/// build produce one warning.
fn build_warnings(shards: &[ShardInput]) -> Vec<SweepWarning> {
    let mut warnings = Vec::new();
    for role in [BuildRole::Engine, BuildRole::Model] {
        let mut listed = shards.iter().map(|input| input.manifest.recorded_builds(role));
        let Some(reference) = listed.by_ref().find(|builds| !builds.is_empty()) else {
            continue;
        };
        let mut others: Vec<RecordedBuild> = Vec::new();
        for build in listed.flatten() {
            if !others.iter().any(|known| known.reads_as(&build)) {
                others.push(build);
            }
        }
        for other in others {
            if reference.iter().any(|build| build.same_build(&other)) {
                continue;
            }
            let recorded = reference
                .iter()
                .find(|build| build.reads_as(&other))
                .unwrap_or(&reference[0]);
            warnings.push(SweepWarning::BuildChanged {
                role,
                recorded: Box::new(recorded.clone()),
                current: Box::new(other),
                between_shards: true,
            });
        }
    }
    warnings
}

/// Returns the manifest of the merged directory while it is written, based on the manifest of the lowest shard.
///
/// The merge holds the whole plan, lists every session of every shard with the builds it ran, and keeps the
/// execution and runtime of the lowest shard. A shard's session that records no engine build receives the build
/// that its shard's manifest implies. The last session of a shard left `running` or `failed` is credited with the rows
/// the shard holds beyond those it found, as a resume credits it. The engine block records the build that performs
/// the merge, and `replays_exactly` is true only when it is true in every shard.
fn merged_manifest(shards: &[ShardInput], shard_dirs: &[PathBuf]) -> Manifest {
    let mut manifest = shards[0].manifest.clone();
    manifest.status = ManifestStatus::Running;
    manifest.shard = Shard::WHOLE.into();
    manifest.engine = RecordedBuild::engine();
    manifest.model.replays_exactly = shards.iter().all(|input| input.manifest.model.replays_exactly);
    manifest.sessions = shards
        .iter()
        .flat_map(|input| {
            let mut recorded = input.manifest.clone();
            recorded.record_session_engines();
            // A session whose process ended before `Manifest::finish` never had its runs counted.
            if matches!(recorded.status, ManifestStatus::Running | ManifestStatus::Failed)
                && let Some(last) = recorded.sessions.last_mut()
            {
                last.ran = (input.runs.records.len() as u64).saturating_sub(last.skipped);
            }
            recorded.sessions
        })
        .collect();
    let started = shards
        .iter()
        .map(|input| input.manifest.timestamps.started_unix_ms)
        .min()
        .unwrap_or_default();
    manifest.timestamps = ManifestTimestamps::started_at(started);
    manifest.results = None;
    manifest.merged_shards = Some(shard_dirs.iter().map(|dir| dir.display().to_string()).collect());
    manifest
}

/// Shard directories that cannot be merged.
#[derive(Debug)]
pub enum MergeError {
    /// A merge with no input directories.
    NoInputs,
    /// A manifest that cannot be read.
    Manifest(ManifestError),
    /// A table that cannot be read back.
    Table(ReadError),
    /// A manifest in `dir` whose shard no plan can have.
    BadShard {
        /// Path of the input directory.
        dir: PathBuf,
    },
    /// Inputs with no `runs.csv` or `series.csv` header between them.
    NoTables,
    /// Input `dir` holds a different plan, model schema or replicate count than input `first`.
    PlanDiffers {
        /// Path of the input directory that differs.
        dir: PathBuf,
        /// Path of the input directory with the lowest shard index.
        first: PathBuf,
    },
    /// Input `dir` splits the plan into a different number of shards than input `first`.
    ShardCountDiffers {
        /// Path of the input directory that differs.
        dir: PathBuf,
        /// Path of the input directory with the lowest shard index.
        first: PathBuf,
    },
    /// Inputs `first` and `second` both hold `shard`.
    SameShard {
        /// Shard both inputs record.
        shard: Shard,
        /// Path of the input directory given earlier.
        first: PathBuf,
        /// Path of the input directory given later.
        second: PathBuf,
    },
    /// `file` of input `dir` has different columns from the other inputs.
    ColumnsDiffer {
        /// Path of the input directory that differs.
        dir: PathBuf,
        /// File name of the table, `runs.csv` or `series.csv`.
        file: &'static str,
    },
    /// Run `run_id` of input `dir`, written as replicate `rep` of config `config_id` under a different id than the
    /// one that the manifest's replicate count implies.
    RunIdDiffers {
        /// Path of the input directory.
        dir: PathBuf,
        /// Run id as written in `runs.csv`.
        run_id: u64,
        /// Config id as written in `runs.csv`.
        config_id: u64,
        /// Replicate index as written in `runs.csv`, counting from 0.
        rep: u64,
    },
    /// Run `run_id` of input `dir`, outside its shard or past the plan's runs.
    OutsideShard {
        /// Path of the input directory.
        dir: PathBuf,
        /// Run id as written in `runs.csv`.
        run_id: u64,
    },
    /// Run `run_id`, listed twice in `runs.csv` of input `dir`.
    DuplicateRun {
        /// Path of the input directory.
        dir: PathBuf,
        /// Run id as written in `runs.csv`.
        run_id: u64,
    },
    /// The merged directory cannot be written.
    Output(OutputError),
    /// Input `dir` holds a search. A search runs whole and has no shards.
    NotASweep {
        /// Path of the input directory.
        dir: PathBuf,
    },
}

impl fmt::Display for MergeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::NoInputs => f.write_str("a merge needs at least one shard directory"),
            Self::Manifest(_) => f.write_str("cannot read the manifest of a shard"),
            Self::Table(_) => f.write_str("cannot read the tables of a shard"),
            Self::BadShard { dir } => write!(f, "'{}' records an invalid shard", dir.display()),
            Self::NoTables => f.write_str("no shard holds a runs.csv and a series.csv with a header"),
            Self::PlanDiffers { dir, first } => write!(
                f,
                "'{}' has a different plan, model schema or replicate count from '{}'",
                dir.display(),
                first.display()
            ),
            Self::ShardCountDiffers { dir, first } => write!(
                f,
                "'{}' has a different shard count from '{}'",
                dir.display(),
                first.display()
            ),
            Self::SameShard { shard, first, second } => write!(
                f,
                "'{}' and '{}' both hold shard {shard}",
                first.display(),
                second.display()
            ),
            Self::ColumnsDiffer { dir, file } => {
                write!(
                    f,
                    "{file} of '{}' has different columns from the other shards",
                    dir.display()
                )
            }
            Self::RunIdDiffers {
                dir,
                run_id,
                config_id,
                rep,
            } => write!(
                f,
                "'{}' holds replicate {rep} of config {config_id} as run {run_id}, and its manifest gives a \
                 different run id. Resume that shard before merging it",
                dir.display()
            ),
            Self::OutsideShard { dir, run_id } => {
                write!(f, "'{}' holds run {run_id}, outside its shard", dir.display())
            }
            Self::DuplicateRun { dir, run_id } => {
                write!(f, "'{}' lists run {run_id} twice", dir.display())
            }
            Self::Output(_) => f.write_str("cannot write the merged results"),
            Self::NotASweep { dir } => write!(f, "'{}' holds a search. A search has no shards to merge", dir.display()),
        }
    }
}

impl std::error::Error for MergeError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Manifest(error) => Some(error),
            Self::Table(error) => Some(error),
            Self::Output(error) => Some(error),
            Self::NoInputs
            | Self::BadShard { .. }
            | Self::NoTables
            | Self::PlanDiffers { .. }
            | Self::ShardCountDiffers { .. }
            | Self::SameShard { .. }
            | Self::ColumnsDiffer { .. }
            | Self::RunIdDiffers { .. }
            | Self::OutsideShard { .. }
            | Self::DuplicateRun { .. }
            | Self::NotASweep { .. } => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::fs;
    use std::path::{Path, PathBuf};

    use henad_core::explore::plan::Shard;
    use henad_core::explore::spec::SweepSpec;

    use super::{MergeError, merge};
    use crate::exec::Concurrency;
    use crate::output::RUNS_FILE;
    use crate::output::manifest::{BuildRole, RecordedBuild};
    use crate::progress::NoProgress;
    use crate::sweep::SweepOptions;
    use crate::tests::support::{
        Recorder, ScratchDir, entry, manifest, other_engine, provenance, rewrite_manifest, sweep, sweep_options,
        sweep_with,
    };

    /// Returns Game of Life on an 8 by 8 grid for 2 steps, with 4 replicates.
    fn life_spec() -> SweepSpec {
        let mut spec = SweepSpec::new("game_of_life");
        spec.fixed = vec![
            ("grid_width".to_owned(), "8".to_owned()),
            ("grid_height".to_owned(), "8".to_owned()),
        ];
        spec.run.steps = 2;
        spec.run.replicates = 4;
        spec
    }

    /// Runs shard `index` of 2 of [`life_spec`] into `dir` under `engine`.
    fn shard_under(dir: &Path, index: u64, engine: RecordedBuild) {
        let options = SweepOptions {
            provenance: provenance().with_engine(engine),
            shard: Shard::new(index, 2).expect("a valid shard"),
            ..sweep_options(false)
        };
        let life = entry("game_of_life", None);
        sweep_with(&life, None, &life_spec(), dir, &options, &mut Recorder::default()).expect("the shard runs");
    }

    /// Returns the directories of 2 shards under `scratch`.
    fn shard_dirs(scratch: &ScratchDir) -> Vec<PathBuf> {
        (0..2)
            .map(|index| scratch.path().join(format!("shard-{index}")))
            .collect()
    }

    #[test]
    fn a_build_the_lowest_shard_ran_is_never_warned_about() {
        let scratch = ScratchDir::new("merge-reference-builds");
        let dirs = shard_dirs(&scratch);
        shard_under(&dirs[0], 0, RecordedBuild::engine());
        shard_under(&dirs[1], 1, RecordedBuild::engine());
        // Shard 0 was resumed under another build, which wrote one of its runs.
        rewrite_manifest(&dirs[0], |recorded| {
            let mut resumed = recorded.sessions[0].clone();
            (resumed.skipped, resumed.ran, resumed.engine) = (1, 1, Some(other_engine()));
            recorded.sessions[0].ran = 1;
            recorded.sessions.push(resumed);
        });
        let mut progress = Recorder::default();
        merge(&dirs, &scratch.path().join("merged"), &mut progress).expect("the shards merge");
        assert_eq!(progress.warnings, [], "shard 1 ran a build shard 0 ran");

        let reversed = ScratchDir::new("merge-reference-other");
        let dirs = shard_dirs(&reversed);
        shard_under(&dirs[0], 0, RecordedBuild::engine());
        shard_under(&dirs[1], 1, other_engine());
        let mut progress = Recorder::default();
        merge(&dirs, &reversed.path().join("merged"), &mut progress).expect("the shards merge");
        assert_eq!(progress.warnings.len(), 1, "{:?}", progress.warnings);
    }

    #[test]
    fn a_merge_credits_a_shard_session_that_never_ended() {
        let scratch = ScratchDir::new("merge-credits");
        let dirs = shard_dirs(&scratch);
        shard_under(&dirs[0], 0, RecordedBuild::engine());
        shard_under(&dirs[1], 1, RecordedBuild::engine());
        rewrite_manifest(&dirs[1], |recorded| {
            recorded.status = crate::output::manifest::ManifestStatus::Running;
            recorded.results = None;
            recorded.sessions[0].ran = 0;
            recorded.model.replays_exactly = false;
        });
        let merged = scratch.path().join("merged");
        merge(&dirs, &merged, &mut NoProgress).expect("the shards merge");
        let recorded = manifest(&merged);
        let ran: Vec<u64> = recorded.sessions.iter().map(|session| session.ran).collect();
        assert_eq!(ran, [2, 2], "the killed shard's session wrote both of its runs");
        assert_eq!(recorded.recorded_builds(BuildRole::Engine), [RecordedBuild::engine()]);
        assert!(
            !recorded.model.replays_exactly,
            "a model replays exactly only when every shard's does"
        );
    }

    #[test]
    fn a_shard_listing_a_run_twice_is_refused() {
        let scratch = ScratchDir::new("merge-duplicate");
        let shard_dir = scratch.path().join("shard");
        let mut spec = SweepSpec::new("game_of_life");
        spec.fixed = vec![
            ("grid_width".to_owned(), "8".to_owned()),
            ("grid_height".to_owned(), "8".to_owned()),
        ];
        spec.run.steps = 2;
        spec.run.replicates = 2;
        sweep(&entry("game_of_life", None), None, &spec, &shard_dir, Concurrency::Auto);
        let runs_path = shard_dir.join(RUNS_FILE);
        let runs = fs::read_to_string(&runs_path).expect("runs.csv is written");
        let last_row = runs.lines().last().expect("runs.csv holds a run");
        fs::write(&runs_path, format!("{runs}{last_row}\n")).expect("runs.csv is written again");

        let error = merge(
            std::slice::from_ref(&shard_dir),
            &scratch.path().join("merged"),
            &mut NoProgress,
        )
        .expect_err("a run listed twice");
        assert!(matches!(error, MergeError::DuplicateRun { run_id: 1, .. }), "{error:?}");
        assert_eq!(
            error.to_string(),
            format!("'{}' lists run 1 twice", shard_dir.display())
        );
    }
}