use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
#[cfg(not(target_arch = "wasm32"))]
use std::fs::OpenOptions;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use web_time::Instant;
use henad_core::action::Schedule;
use henad_core::explore::factor::FactorSpec;
use henad_core::explore::fingerprint::{run_key, search_hash};
use henad_core::explore::measure::MeasurePlan;
use henad_core::explore::outcome::{PlannedRun, RunOutcome};
use henad_core::explore::plan::{Config, MAX_RUNS, ModelSchema, Plan, PlanError, Shard};
use henad_core::explore::replay::Replay;
use henad_core::explore::search::genome::{Genome, SearchSpace, SearchSpaceError};
use henad_core::explore::search::pse::{
ArchiveEntry, PatternAxis, PatternCell, PatternPlacement, PatternSpaceSettings,
};
use henad_core::explore::search::{
Candidate, CandidateOrigin, Evaluation, GenerationSummary, Objective, RankingEntry, SearchAlgorithm, SearchReport,
SearchSpec, SearchSpecError, Searcher,
};
use henad_core::explore::seed::search_seed;
use henad_core::explore::spec::SweepSpec;
use crate::exec::{BatchEnd, RunRequest};
#[cfg(not(target_arch = "wasm32"))]
use crate::exec::{Executor, RunSink};
use crate::output::manifest::{
Manifest, ManifestAxisRanges, ManifestMode, ManifestPlan, ManifestSearch, ResultCounts, now_unix_ms,
};
use crate::output::memory::{SweepFiles, memory_writer};
use crate::output::read::{ReadError, RunsCsv, SeriesScan, parse_one};
use crate::output::resume::{ResumeError, check_model};
use crate::output::runs_csv::{ID_COLUMNS, OUTCOME_COLUMNS, column_names};
use crate::output::search_tables::{ConfigColumns, SearchTablesWriter, write_archive, write_ranking};
use crate::output::{
ARCHIVE_FILE, BATCHES_FILE, BEST_FILE, EVALUATIONS_FILE, GENERATIONS_FILE, MANIFEST_FILE, OutputDir, OutputError,
OutputWriter, RUNS_FILE, SERIES_FILE, runs_csv, series_csv, table_paths,
};
use crate::probe::{ProbeReport, TimedProbe, check_capacity};
#[cfg(not(target_arch = "wasm32"))]
use crate::progress::ProgressMeter;
use crate::progress::{Progress, ProgressEvent};
use crate::sweep::{
ExploreError, ManifestParts, SweepEnd, SweepInputs, SweepOutline, SweepRecord, SweepReport, SweepWarning,
build_warnings, finish_manifest, hex, running_manifest, sized_layout,
};
#[derive(Debug, Clone)]
pub struct SearchPlan {
base: Arc<Plan>,
search: SearchSpec,
space: SearchSpace,
search_hash: u64,
run_count: u64,
}
impl SearchPlan {
pub fn new(spec: &SweepSpec, schema: &ModelSchema<'_>) -> Result<Self, SearchPlanError> {
let search = spec.search.as_ref().ok_or(SearchPlanError::NotASearch)?;
if !spec.blocks.is_empty() {
return Err(SearchPlanError::Blocks);
}
search.check().map_err(SearchPlanError::Settings)?;
let base = spec.plan(schema).map_err(SearchPlanError::Plan)?;
let space = SearchSpace::resolve(&search.space, schema.params, &spec.actions, &spec.fixed)
.map_err(SearchPlanError::Space)?;
if (search.batch_size as u64).saturating_mul(base.replicates()) > MAX_RUNS {
return Err(SearchPlanError::BatchTooLarge {
batch_size: search.batch_size,
replicates: base.replicates(),
});
}
let run_count = search
.max_evaluations
.checked_mul(base.replicates())
.ok_or(SearchPlanError::TooManyRuns)?;
Ok(Self {
search_hash: search_hash(base.plan_hash(), base.replicates(), search),
base: Arc::new(base),
search: search.clone(),
space,
run_count,
})
}
pub fn base(&self) -> &Arc<Plan> {
&self.base
}
pub fn search(&self) -> &SearchSpec {
&self.search
}
pub fn space(&self) -> &SearchSpace {
&self.space
}
pub fn search_hash(&self) -> u64 {
self.search_hash
}
pub fn run_count(&self) -> u64 {
self.run_count
}
pub fn replicates(&self) -> u64 {
self.base.replicates()
}
pub fn config(&self, genome: &Genome) -> Config {
let base = self.base.config(0).expect("a plan with no blocks has one config");
self.space.decode(genome, base)
}
pub fn run(&self, candidate: &Candidate, index: u64) -> PlannedRun {
let replicates = self.replicates();
let replicate = candidate.replicate_offset + index;
let seeds = self.base.seed_settings();
let seed_config_id = candidate.origin.reevaluated_id().unwrap_or(candidate.id);
PlannedRun {
run_id: candidate.id * replicates + index,
config_id: candidate.id,
rep: replicate,
seed: seeds.scheme.seed(seeds.root, seed_config_id, replicate),
}
}
pub fn run_key(&self, run: &PlannedRun, config: &Config) -> u64 {
run_key(
self.base.results_fingerprint(),
&config.params,
&config.action_ticks,
run.seed,
)
}
pub fn replay(&self, run: &PlannedRun, config: &Config) -> Replay {
let settings = self.base.run_settings();
Replay {
model: self.base.model().to_owned(),
params: config.params.clone(),
seed: run.seed,
schedule: self.base.schedule(config),
ticks: settings.warmup + settings.steps,
label: format!(
"Search run {}: candidate {}, replicate {}",
run.run_id, run.config_id, run.rep
),
}
}
pub fn watched_reducers(&self, measure: &MeasurePlan) -> Result<Vec<usize>, SearchPlanError> {
let names = measure.reducers().names();
self.search
.watched_columns()
.into_iter()
.map(|column| {
names
.iter()
.position(|name| name == column)
.ok_or_else(|| SearchPlanError::UnknownColumn {
column: column.to_owned(),
known: names.to_vec(),
})
})
.collect()
}
pub fn outline(&self) -> SearchOutline {
SearchOutline {
algorithm: self.search.algorithm.as_str(),
max_evaluations: self.search.max_evaluations,
batch_size: self.search.batch_size,
objective: self.search.objective.clone(),
watched_columns: self.search.watched_columns().into_iter().map(str::to_owned).collect(),
pattern_axes: match &self.search.algorithm {
SearchAlgorithm::PatternSpaceExploration(settings) => {
Some([settings.x_axis.clone(), settings.y_axis.clone()])
}
SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => None,
},
space: self.search.space.clone(),
search_seed: search_seed(self.base.seed_settings().root),
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum SearchPlanError {
NotASearch,
Blocks,
Plan(PlanError),
Settings(SearchSpecError),
Space(SearchSpaceError),
UnknownColumn {
column: String,
known: Vec<String>,
},
BatchTooLarge {
batch_size: usize,
replicates: u64,
},
TooManyRuns,
Sharded,
RetryFailed,
}
impl fmt::Display for SearchPlanError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::NotASearch => f.write_str("the spec has no [search] table"),
Self::Blocks => f.write_str("a search selects its own configs and takes no blocks"),
Self::Plan(_) => f.write_str("cannot plan the search's fixed values and actions"),
Self::Settings(_) => f.write_str("search settings"),
Self::Space(_) => f.write_str("search space"),
Self::UnknownColumn { column, known } if known.is_empty() => {
write!(f, "search watches '{column}', and the runs have no reducer columns")
}
Self::UnknownColumn { column, known } => write!(
f,
"no reducer writes '{column}', the column the search watches. The reducer columns are {}",
known.join(", ")
),
Self::BatchTooLarge { batch_size, replicates } => write!(
f,
"a batch of {batch_size} candidates at {replicates} replicates each has more than {MAX_RUNS} runs"
),
Self::TooManyRuns => f.write_str("search budget has more runs than a 64-bit integer can hold"),
Self::Sharded => f.write_str("a search cannot run as a shard"),
Self::RetryFailed => {
f.write_str("a resumed search replays its failed runs as they were, and cannot run them again")
}
}
}
}
impl std::error::Error for SearchPlanError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Plan(error) => Some(error),
Self::Settings(error) => Some(error),
Self::Space(error) => Some(error),
Self::NotASearch
| Self::Blocks
| Self::UnknownColumn { .. }
| Self::BatchTooLarge { .. }
| Self::TooManyRuns
| Self::Sharded
| Self::RetryFailed => None,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct SearchOutline {
pub algorithm: &'static str,
pub max_evaluations: u64,
pub batch_size: usize,
pub objective: Option<Objective>,
pub watched_columns: Vec<String>,
pub pattern_axes: Option<[PatternAxis; 2]>,
pub space: Vec<FactorSpec>,
pub search_seed: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct SearchUpdate {
pub batch: u64,
pub evaluations: u64,
pub runs: u64,
pub evaluated: Vec<EvaluatedCandidate>,
pub best: Option<RankingEntry>,
pub generations: Vec<GenerationSummary>,
pub landed_entries: Vec<ArchiveEntry>,
pub filled_cells: u64,
pub pattern_settings: Option<PatternSpaceSettings>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct EvaluatedCandidate {
pub candidate_id: u64,
pub batch: u64,
pub origin: CandidateOrigin,
pub replicate_offset: u64,
pub replicates: u64,
pub config: Config,
pub failed_count: u64,
pub reading: EvaluationReading,
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum EvaluationReading {
Objective {
objective: f64,
pooled_objective: f64,
pooled_replicates: u64,
},
Placement {
placement: Option<PatternPlacement>,
new_cell: bool,
},
}
pub(crate) struct SearchSession {
plan: Arc<SearchPlan>,
searcher: Box<dyn Searcher + Send>,
watched_reducers: Vec<usize>,
recorded_runs: Arc<[RecordedRun]>,
configs: BTreeMap<u64, Config>,
batch_count: u64,
evaluations: u64,
runs: u64,
reported_generations: usize,
reported_ranges: bool,
}
pub(crate) struct AskedBatch {
index: u64,
candidates: Vec<Candidate>,
configs: Vec<Config>,
schedules: Vec<Schedule>,
runs: Vec<BatchRun>,
recorded_values: Vec<Vec<Option<f64>>>,
}
struct BatchRun {
run: PlannedRun,
run_key: u64,
candidate_index: usize,
}
impl AskedBatch {
#[cfg(not(target_arch = "wasm32"))]
pub(crate) fn requests(&self) -> Vec<RunRequest<'_>> {
self.runs[self.recorded_values.len()..]
.iter()
.map(|batch_run| self.request_of(batch_run))
.collect()
}
#[cfg(any(target_arch = "wasm32", test))]
pub(crate) fn request(&self, position: usize) -> Option<RunRequest<'_>> {
let batch_run = self.runs.get(self.recorded_values.len() + position)?;
Some(self.request_of(batch_run))
}
fn request_of(&self, batch_run: &BatchRun) -> RunRequest<'_> {
RunRequest {
run: batch_run.run,
run_key: batch_run.run_key,
params: &self.configs[batch_run.candidate_index].params,
schedule: self.schedules[batch_run.candidate_index].clone(),
}
}
pub(crate) fn config_of(&self, run: &PlannedRun) -> &Config {
let index = self
.candidates
.binary_search_by_key(&run.config_id, |candidate| candidate.id)
.expect("a run of the batch belongs to one of its candidates");
&self.configs[index]
}
}
impl SearchSession {
pub(crate) fn new(
plan: Arc<SearchPlan>,
watched_reducers: Vec<usize>,
recorded_runs: Arc<[RecordedRun]>,
) -> Result<Self, SearchPlanError> {
let root = plan.base.seed_settings().root;
let searcher = plan
.search
.searcher(&plan.space, root)
.map_err(SearchPlanError::Settings)?;
Ok(Self {
plan,
searcher,
watched_reducers,
recorded_runs,
configs: BTreeMap::new(),
batch_count: 0,
evaluations: 0,
runs: 0,
reported_generations: 0,
reported_ranges: false,
})
}
pub(crate) fn watched_reducers(&self) -> &[usize] {
&self.watched_reducers
}
pub(crate) fn ask(&mut self) -> Result<Option<AskedBatch>, ResumeError> {
if self.searcher.is_done() {
return Ok(None);
}
let candidates = self.searcher.ask(self.plan.search.batch_size);
if candidates.is_empty() {
return Ok(None);
}
let replicates = self.plan.replicates();
let mut batch = AskedBatch {
index: self.batch_count,
configs: Vec::with_capacity(candidates.len()),
schedules: Vec::with_capacity(candidates.len()),
runs: Vec::with_capacity(candidates.len() * replicates as usize),
recorded_values: Vec::new(),
candidates: Vec::new(),
};
for (candidate_index, candidate) in candidates.iter().enumerate() {
let config = self.plan.config(&candidate.genome);
for index in 0..replicates {
let run = self.plan.run(candidate, index);
batch.runs.push(BatchRun {
run,
run_key: self.plan.run_key(&run, &config),
candidate_index,
});
}
if candidate.origin.reevaluated_id().is_none() {
self.configs.insert(candidate.id, config.clone());
}
batch.schedules.push(self.plan.base.schedule(&config));
batch.configs.push(config);
}
for batch_run in &batch.runs {
let Some(recorded) = usize::try_from(batch_run.run.run_id)
.ok()
.and_then(|run_id| self.recorded_runs.get(run_id))
else {
break;
};
if recorded.run_key != batch_run.run_key {
return Err(ResumeError::SearchRunChanged {
run_id: batch_run.run.run_id,
});
}
batch.recorded_values.push(recorded.values.clone());
}
batch.candidates = candidates;
Ok(Some(batch))
}
pub(crate) fn replay_recorded_runs(&mut self) -> Result<(), ResumeError> {
while let Some(batch) = self.ask()? {
if batch.recorded_values.len() < batch.runs.len() {
break;
}
self.tell(batch, Vec::new());
}
Ok(())
}
pub(crate) fn tell(&mut self, batch: AskedBatch, values: Vec<Vec<Option<f64>>>) -> SearchUpdate {
let AskedBatch {
index,
candidates,
configs,
runs,
mut recorded_values,
..
} = batch;
recorded_values.extend(values);
let values = recorded_values;
assert_eq!(values.len(), runs.len(), "every run of the batch has its values");
let replicates = self.plan.replicates() as usize;
let evaluations: Vec<Evaluation> = candidates
.iter()
.zip(values.chunks(replicates.max(1)))
.map(|(candidate, outputs)| Evaluation {
candidate_id: candidate.id,
outputs: outputs.to_vec(),
})
.collect();
self.searcher.tell(&evaluations);
let searcher = &self.searcher;
let pattern_settings = match &self.plan.search.algorithm {
SearchAlgorithm::PatternSpaceExploration(settings) => Some(searcher.pattern_settings().unwrap_or(settings)),
SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => None,
};
let mut touched = BTreeSet::new();
let evaluated: Vec<EvaluatedCandidate> = candidates
.into_iter()
.zip(configs)
.zip(&evaluations)
.map(|((candidate, config), evaluation)| {
let failed_count = evaluation
.outputs
.iter()
.filter(|row| row.iter().any(|value| !value.is_some_and(f64::is_finite)))
.count() as u64;
let reading = if let Some(settings) = pattern_settings {
let placement = settings.place(evaluation);
let cell = placement.and_then(|placement| placement.cell);
if let Some(cell) = cell {
touched.insert(cell);
}
let new_cell = cell.is_some_and(|cell| {
searcher
.archive_entry(cell)
.is_some_and(|entry| entry.candidate_id == candidate.id)
});
EvaluationReading::Placement { placement, new_cell }
} else {
let objective = self
.plan
.search
.objective
.as_ref()
.expect("a scored search has an objective");
let pooled_id = candidate.origin.reevaluated_id().unwrap_or(candidate.id);
let pooled = searcher.ranking_entry(pooled_id);
EvaluationReading::Objective {
objective: objective.score(evaluation),
pooled_objective: pooled.as_ref().map_or(objective.goal.worst(), |entry| entry.objective),
pooled_replicates: pooled.as_ref().map_or(0, |entry| entry.replicate_count),
}
};
EvaluatedCandidate {
candidate_id: candidate.id,
batch: index,
origin: candidate.origin,
replicate_offset: candidate.replicate_offset,
replicates: replicates as u64,
config,
failed_count,
reading,
}
})
.collect();
self.batch_count += 1;
self.evaluations += evaluated.len() as u64;
self.runs += runs.len() as u64;
let generations = searcher.generations();
let new_generations = generations[self.reported_generations.min(generations.len())..].to_vec();
self.reported_generations = generations.len();
let (best, filled_cells) = (searcher.best(), searcher.filled_cells());
let pattern_settings = searcher.pattern_settings().cloned();
SearchUpdate {
batch: index,
evaluations: self.evaluations,
runs: self.runs,
evaluated,
best,
generations: new_generations,
landed_entries: self.landed_entries(touched),
filled_cells,
pattern_settings,
}
}
fn landed_entries(&mut self, touched: BTreeSet<PatternCell>) -> Vec<ArchiveEntry> {
if !self.reported_ranges && self.searcher.pattern_settings().is_some() {
self.reported_ranges = true;
return self.searcher.report().archive;
}
touched
.into_iter()
.filter_map(|cell| self.searcher.archive_entry(cell))
.collect()
}
pub(crate) fn report(&self) -> SearchReport {
self.searcher.report()
}
pub(crate) fn manifest_search(&self) -> ManifestSearch {
let mut record = manifest_search(&self.plan);
record.evaluations = self.evaluations;
record.batch_count = self.batch_count;
match &self.plan.search.algorithm {
SearchAlgorithm::PatternSpaceExploration(_) => {
record.filled_cells = Some(self.searcher.filled_cells());
record.axis_ranges = self.searcher.pattern_settings().and_then(ManifestAxisRanges::of);
}
SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => {
let best = self.searcher.best();
record.best_candidate_id = best.as_ref().map(|entry| entry.candidate_id);
record.best_objective = best
.map(|entry| entry.objective)
.filter(|objective| objective.is_finite());
}
}
record
}
pub(crate) fn write_closing_table<W: Write>(
&self,
report: &SearchReport,
dest: W,
columns: &ConfigColumns,
) -> io::Result<W> {
match &self.plan.search.algorithm {
SearchAlgorithm::PatternSpaceExploration(settings) => {
let settings = self.searcher.pattern_settings().unwrap_or(settings);
write_archive(dest, settings, &report.archive, &self.configs, columns)
}
SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => {
write_ranking(dest, &report.ranking, &self.configs, columns)
}
}
}
}
fn manifest_search(plan: &SearchPlan) -> ManifestSearch {
let pattern_settings = match &plan.search.algorithm {
SearchAlgorithm::PatternSpaceExploration(settings) => Some(settings),
SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => None,
};
ManifestSearch {
algorithm: plan.search.algorithm.as_str().to_owned(),
search_hash: hex(plan.search_hash),
max_evaluations: plan.search.max_evaluations,
batch_size: plan.search.batch_size,
search_seed: search_seed(plan.base.seed_settings().root),
watched_columns: plan.search.watched_columns().into_iter().map(str::to_owned).collect(),
evaluations: 0,
batch_count: 0,
best_candidate_id: None,
best_objective: None,
filled_cells: pattern_settings.map(|_| 0),
axis_ranges: pattern_settings.and_then(ManifestAxisRanges::of),
}
}
fn closing_table_file(plan: &SearchPlan) -> &'static str {
match plan.search.algorithm {
SearchAlgorithm::PatternSpaceExploration(_) => ARCHIVE_FILE,
SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => BEST_FILE,
}
}
fn writes_generations(plan: &SearchPlan) -> bool {
matches!(plan.search.algorithm, SearchAlgorithm::Genetic(_))
}
pub(crate) fn watched_values(outcome: &RunOutcome, watched_reducers: &[usize]) -> Vec<Option<f64>> {
watched_reducers
.iter()
.map(|&position| {
if outcome.status.is_failure() {
None
} else {
outcome
.reducers
.get(position)
.copied()
.flatten()
.filter(|value| value.is_finite())
}
})
.collect()
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct RecordedRun {
run_key: u64,
values: Vec<Option<f64>>,
}
#[derive(Debug)]
#[cfg_attr(target_arch = "wasm32", expect(dead_code, reason = "a browser resumes no folder"))]
pub(crate) struct RecordedSearch {
pub(crate) recorded: Manifest,
runs: Arc<[RecordedRun]>,
counts: ResultCounts,
runs_bytes: u64,
series_bytes: u64,
fresh_tables: bool,
}
impl RecordedSearch {
pub(crate) fn read(
path: &Path,
plan: &SearchPlan,
runs_header: &str,
series_header: &str,
watched_reducers: &[usize],
config_column_count: usize,
) -> Result<Self, ResumeError> {
let recorded = Manifest::read(&path.join(MANIFEST_FILE)).map_err(ResumeError::Manifest)?;
if recorded.mode != ManifestMode::Search {
return Err(ResumeError::ModeChanged {
recorded: recorded.mode,
current: ManifestMode::Search,
});
}
check_model(&recorded, plan.base.model())?;
let current_schema = hex(plan.base.schema_hash());
if recorded.model.schema_hash != current_schema {
return Err(ResumeError::SchemaChanged {
recorded: recorded.model.schema_hash.clone(),
current: current_schema,
});
}
let current_search = hex(plan.search_hash);
let recorded_search = recorded
.search
.as_ref()
.map_or_else(String::new, |search| search.search_hash.clone());
if recorded_search != current_search {
return Err(ResumeError::SearchChanged {
recorded: recorded_search,
current: current_search,
});
}
let (runs_path, series_path) = table_paths(path);
let table = RunsCsv::read(&runs_path).map_err(ResumeError::Table)?;
if let Some(header) = &table.header
&& runs_csv::header_line(header) != runs_header
{
return Err(ResumeError::ColumnsChanged { file: RUNS_FILE });
}
let reducers_start = ID_COLUMNS.len() + config_column_count + OUTCOME_COLUMNS.len();
let mut runs = Vec::with_capacity(table.records.len());
let mut counts = ResultCounts::default();
for (position, record) in table.records.iter().enumerate() {
if position as u64 >= plan.run_count() {
return Err(ResumeError::UnknownRun { run_id: record.run_id });
}
if record.run_id != position as u64 {
return Err(ResumeError::SearchRunChanged { run_id: record.run_id });
}
let fields =
parse_one(&record.text, 0..record.text.len(), &runs_path, position + 2).map_err(ResumeError::Table)?;
let values = watched_reducers
.iter()
.map(|&reducer| {
let column = reducers_start + reducer;
let field = fields.get(column).map_or("", String::as_str);
if record.status.is_failure() || field.is_empty() {
return Ok(None);
}
field.parse::<f64>().ok().map(Some).ok_or_else(|| {
ResumeError::Table(ReadError::BadField {
path: runs_path.clone(),
record_number: position + 2,
column: table
.header
.as_ref()
.and_then(|header| header.get(column))
.cloned()
.unwrap_or_default(),
text: field.to_owned(),
})
})
})
.collect::<Result<_, _>>()?;
counts.count(record.status);
runs.push(RecordedRun {
run_key: record.run_key,
values,
});
}
let recorded_runs = runs.len() as u64;
let series = SeriesScan::read(&series_path, |run_id| run_id < recorded_runs).map_err(ResumeError::Table)?;
if series.header.as_deref().is_some_and(|header| header != series_header) {
return Err(ResumeError::ColumnsChanged { file: SERIES_FILE });
}
if series.kept_after_dropped || series.segments.len() > 1 {
return Err(ResumeError::SeriesOutOfOrder);
}
let fresh_tables = table.header.is_none() || series.header.is_none();
if fresh_tables {
runs.clear();
counts = ResultCounts::default();
}
Ok(Self {
recorded,
runs: runs.into(),
counts,
runs_bytes: table.complete_bytes,
series_bytes: series.first_dropped_offset.unwrap_or(series.complete_bytes),
fresh_tables,
})
}
#[cfg(not(target_arch = "wasm32"))]
fn repair(&self, dir: &OutputDir) -> Result<(), OutputError> {
for (file, length) in [(RUNS_FILE, self.runs_bytes), (SERIES_FILE, self.series_bytes)] {
let path = dir.path().join(file);
let repaired = if self.fresh_tables {
match std::fs::remove_file(&path) {
Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
removed => removed,
}
} else {
OpenOptions::new()
.write(true)
.open(&path)
.and_then(|table| table.set_len(length))
};
repaired.map_err(|source| OutputError::Write { path, source })?;
}
Ok(())
}
}
#[cfg(not(target_arch = "wasm32"))]
pub(crate) fn run_search_into_directory(
inputs: &SweepInputs<'_>,
plan: Option<Arc<SearchPlan>>,
output_dir: &Path,
progress: &mut dyn Progress,
) -> Result<SweepRecord, ExploreError> {
let preparation = SearchPreparation::new(inputs, plan, None)?;
preparation.announce(inputs, progress);
let record = preparation.write_directory(inputs, output_dir, progress)?;
progress.report(&ProgressEvent::Ended(&record.report));
Ok(record)
}
#[cfg(not(target_arch = "wasm32"))]
pub(crate) fn run_search_in_memory(
inputs: &SweepInputs<'_>,
plan: Option<Arc<SearchPlan>>,
progress: &mut dyn Progress,
) -> Result<SweepRecord, ExploreError> {
let preparation = SearchPreparation::new(inputs, plan, None)?;
preparation.announce(inputs, progress);
let (mut output, manifest) = preparation.memory_output(inputs)?;
let executor = preparation.executor(inputs)?;
let end = preparation.run_all(&executor, &mut output, progress)?;
let record = preparation.finish_in_memory(output, manifest, end)?;
progress.report(&ProgressEvent::Ended(&record.report));
Ok(record)
}
pub(crate) struct SearchPreparation {
plan: Arc<SearchPlan>,
probe: ProbeReport,
measure: Arc<MeasurePlan>,
watched_reducers: Vec<usize>,
resumed: Option<RecordedSearch>,
#[cfg_attr(
target_arch = "wasm32",
expect(dead_code, reason = "a search in a browser writes no directory")
)]
locked: Option<OutputDir>,
outline: SweepOutline,
started: Instant,
started_unix_ms: u64,
}
pub(crate) struct SearchWriters<W: Write> {
pub(crate) session: SearchSession,
pub(crate) writer: OutputWriter<W>,
pub(crate) tables: SearchTablesWriter<W>,
pub(crate) columns: ConfigColumns,
}
impl SearchPreparation {
pub(crate) fn new(
inputs: &SweepInputs<'_>,
plan: Option<Arc<SearchPlan>>,
probe: Option<TimedProbe>,
) -> Result<Self, ExploreError> {
let (entry, gpu, options) = (inputs.entry, inputs.gpu, inputs.options);
if options.shard != Shard::WHOLE {
return Err(SearchPlanError::Sharded.into());
}
if options.retry_failed {
return Err(SearchPlanError::RetryFailed.into());
}
let (started, started_unix_ms, probe) = match probe {
Some(timed) => (timed.started, timed.started_unix_ms, Some(timed.report)),
None => (Instant::now(), now_unix_ms(), None),
};
let plan = match plan {
Some(plan) => plan,
None => Arc::new(SearchPlan::new(inputs.spec, &entry.schema())?),
};
let resume_dir = inputs
.folder
.filter(|output_dir| options.resume && OutputDir::holds_results(output_dir));
let locked = match (resume_dir, inputs.folder) {
(None, Some(output_dir)) => {
OutputDir::check_free(output_dir)?;
None
}
(Some(output_dir), _) if !inputs.dry_run => Some(OutputDir::open(output_dir)?),
_ => None,
};
let base = plan.base();
if let Some(ctx) = gpu {
check_capacity(entry, base, &ctx.device.limits())?;
}
let probe = match probe {
Some(probe) => probe,
None => ProbeReport::for_plan(entry, gpu, base)?,
};
let measure = MeasurePlan::new(base.run_settings(), base.measure_settings(), probe.columns.clone())?;
let watched_reducers = plan.watched_reducers(&measure)?;
let resumed = resume_dir
.map(|output_dir| {
let params = entry.param_descriptors();
let columns = column_names(params, base.actions(), measure.reducers().names());
RecordedSearch::read(
output_dir,
&plan,
&runs_csv::header_line(&columns),
&series_csv::header_line(measure.columns()),
&watched_reducers,
params.len() + base.actions().len(),
)
})
.transpose()?;
if let Some(recorded) = resumed.as_ref().filter(|recorded| !recorded.runs.is_empty()) {
SearchSession::new(Arc::clone(&plan), watched_reducers.clone(), Arc::clone(&recorded.runs))?
.replay_recorded_runs()?;
}
let skipped = resumed.as_ref().map_or(0, |recorded| recorded.runs.len() as u64);
let pending = plan.run_count().saturating_sub(skipped);
let batch_runs = (plan.search.batch_size as u64).saturating_mul(plan.replicates());
let (layout, projected_bytes) = sized_layout(entry, gpu, options, &probe, pending.min(batch_runs))?;
let outline = SweepOutline {
model: entry.id().to_owned(),
backend: entry.metadata().backend,
configs: None,
replicates: plan.replicates(),
runs: plan.run_count(),
blocks: Vec::new(),
shard: Shard::WHOLE,
skipped,
pending,
layout,
projected_bytes,
series_rows: pending.saturating_mul(measure.series_row_count()),
stat_columns: (0..measure.columns().len())
.map(|column| measure.columns().name(column).to_owned())
.collect(),
reducer_columns: measure.reducers().names().to_vec(),
dry_run: inputs.dry_run,
search: Some(plan.outline()),
};
Ok(Self {
plan,
probe,
measure: Arc::new(measure),
watched_reducers,
resumed,
locked,
outline,
started,
started_unix_ms,
})
}
pub(crate) fn announce(&self, inputs: &SweepInputs<'_>, progress: &mut dyn Progress) {
progress.report(&ProgressEvent::Planned(&self.outline));
let mut warnings: Vec<SweepWarning> = self
.plan
.base
.warnings()
.iter()
.cloned()
.map(SweepWarning::Plan)
.collect();
if let Some(resumed) = &self.resumed {
warnings.extend(build_warnings(&resumed.recorded, inputs.provenance, inputs.entry));
}
for warning in &warnings {
progress.report(&ProgressEvent::Warned(warning));
}
}
#[cfg(any(target_arch = "wasm32", test))]
pub(crate) fn measure(&self) -> &Arc<MeasurePlan> {
&self.measure
}
#[cfg(any(target_arch = "wasm32", test))]
pub(crate) fn plan(&self) -> &Arc<SearchPlan> {
&self.plan
}
fn manifest(&self, inputs: &SweepInputs<'_>) -> Result<Manifest, OutputError> {
running_manifest(
inputs,
&ManifestParts {
mode: ManifestMode::Search,
plan: &self.plan.base,
probe: &self.probe,
outline: &self.outline,
manifest_plan: ManifestPlan {
plan_hash: hex(self.plan.base.plan_hash()),
results_fingerprint: hex(self.plan.base.results_fingerprint()),
configs: self.outline.configs,
replicates: self.outline.replicates,
runs: self.outline.runs,
blocks: Vec::new(),
},
started_unix_ms: self.started_unix_ms,
recorded: self.resumed.as_ref().map(|resumed| &resumed.recorded),
search: Some(manifest_search(&self.plan)),
},
)
}
fn session(&self) -> Result<SearchSession, ExploreError> {
let recorded = self
.resumed
.as_ref()
.map(|resumed| Arc::clone(&resumed.runs))
.unwrap_or_default();
Ok(SearchSession::new(
Arc::clone(&self.plan),
self.watched_reducers.clone(),
recorded,
)?)
}
pub(crate) fn memory_output(
&self,
inputs: &SweepInputs<'_>,
) -> Result<(SearchWriters<Vec<u8>>, Manifest), ExploreError> {
let params = inputs.entry.param_descriptors();
let writer = memory_writer(&self.plan.base, params, &self.measure)?;
let generations = writes_generations(&self.plan).then(Vec::new);
let tables = SearchTablesWriter::new(
Vec::new(),
Vec::new(),
generations,
params,
self.plan.base.actions(),
self.is_pattern_search(),
)
.map_err(|source| table_error(Path::new(EVALUATIONS_FILE), source))?;
let output = SearchWriters {
session: self.session()?,
writer,
tables,
columns: ConfigColumns::new(params, self.plan.base.actions()),
};
Ok((output, self.manifest(inputs)?))
}
fn is_pattern_search(&self) -> bool {
matches!(self.plan.search.algorithm, SearchAlgorithm::PatternSpaceExploration(_))
}
#[cfg(not(target_arch = "wasm32"))]
fn write_directory(
&self,
inputs: &SweepInputs<'_>,
output_dir: &Path,
progress: &mut dyn Progress,
) -> Result<SweepRecord, ExploreError> {
let created;
let dir = if let Some(dir) = &self.locked {
dir
} else {
created = OutputDir::create(output_dir)?;
&created
};
let mut manifest = self.manifest(inputs)?;
dir.write_manifest(&manifest)?;
let mut standing = None;
let (end, counts, session, search_report) = match self.run_into(dir, inputs, progress, &mut standing) {
Ok(finished) => finished,
Err(error) => {
manifest.fail(now_unix_ms());
if let Some(standing) = standing {
manifest.search = Some(standing);
}
drop(dir.write_manifest(&manifest));
return Err(error);
}
};
let end = finish_manifest(&mut manifest, end, counts);
manifest.search = Some(session.manifest_search());
dir.write_manifest(&manifest)?;
Ok(SweepRecord {
report: self.report(end, counts, Some(output_dir.to_owned())),
manifest,
files: None,
search: Some(search_report),
})
}
#[cfg(not(target_arch = "wasm32"))]
fn run_into(
&self,
dir: &OutputDir,
inputs: &SweepInputs<'_>,
progress: &mut dyn Progress,
standing: &mut Option<ManifestSearch>,
) -> Result<(BatchEnd, ResultCounts, SearchSession, SearchReport), ExploreError> {
let params = inputs.entry.param_descriptors();
let executor = self.executor(inputs)?;
let fresh_tables = match &self.resumed {
Some(resumed) => {
resumed.repair(dir)?;
resumed.fresh_tables
}
None => true,
};
let writer = if fresh_tables {
dir.open_writer(&self.plan.base, params, &self.measure)?
} else {
dir.append_writer(&self.plan.base, params)?
};
let generations = if writes_generations(&self.plan) {
Some(dir.create_table(GENERATIONS_FILE)?)
} else {
None
};
let tables = SearchTablesWriter::new(
dir.create_table(EVALUATIONS_FILE)?,
dir.create_table(BATCHES_FILE)?,
generations,
params,
self.plan.base.actions(),
self.is_pattern_search(),
)
.map_err(|source| table_error(&dir.path().join(EVALUATIONS_FILE), source))?;
let mut output = SearchWriters {
session: self.session()?,
writer,
tables,
columns: ConfigColumns::new(params, self.plan.base.actions()),
};
let end = self.run_all(&executor, &mut output, progress);
*standing = Some(output.session.manifest_search());
let end = end?;
let SearchWriters {
session,
writer,
tables,
columns,
} = output;
let written = writer.counts();
writer.finish().map_err(|source| OutputError::Write {
path: dir.path().to_owned(),
source,
})?;
tables
.into_inner()
.map_err(|source| table_error(&dir.path().join(EVALUATIONS_FILE), source))?;
let closing_file = closing_table_file(&self.plan);
let search_report = session.report();
session
.write_closing_table(&search_report, dir.create_table(closing_file)?, &columns)
.map_err(|source| table_error(&dir.path().join(closing_file), source))?;
dir.write_summary()?;
let kept = self.resumed.as_ref().map(|resumed| resumed.counts).unwrap_or_default();
Ok((end, kept + written, session, search_report))
}
#[cfg(not(target_arch = "wasm32"))]
fn executor<'a>(&self, inputs: &SweepInputs<'a>) -> Result<Executor<'a>, ExploreError> {
let options = inputs.options;
let executor = Executor::new(
inputs.entry,
inputs.gpu,
Arc::clone(&self.measure),
self.outline.layout,
options.control.clone(),
)?;
Ok(executor
.with_timeout(self.plan.base.run_settings().timeout)
.with_active_runs(options.active_runs.clone())
.with_gpu_memory_budget(options.gpu_memory_budget))
}
#[cfg(not(target_arch = "wasm32"))]
fn run_all<W: Write>(
&self,
executor: &Executor<'_>,
output: &mut SearchWriters<W>,
progress: &mut dyn Progress,
) -> Result<BatchEnd, ExploreError> {
let mut meter = ProgressMeter::new(self.outline.pending);
loop {
if executor.control().is_aborted() {
return Ok(BatchEnd::Aborted);
}
let Some(batch) = output.session.ask()? else {
return Ok(BatchEnd::Complete);
};
let requests = batch.requests();
let mut sink = SearchSink {
batch: &batch,
writer: &mut output.writer,
watched_reducers: output.session.watched_reducers(),
values: Vec::with_capacity(requests.len()),
progress: &mut *progress,
meter: &mut meter,
};
let end = if requests.is_empty() {
BatchEnd::Complete
} else {
executor.run_batch(&requests, &mut sink)?
};
let values = sink.values;
drop(requests);
if end != BatchEnd::Complete {
return Ok(end);
}
let update = output.session.tell(batch, values);
output
.tables
.write_batch(&update)
.map_err(|source| table_error(Path::new(EVALUATIONS_FILE), source))?;
progress.report(&ProgressEvent::SearchBatchTold(&update));
}
}
pub(crate) fn finish_in_memory(
&self,
output: SearchWriters<Vec<u8>>,
mut manifest: Manifest,
end: BatchEnd,
) -> Result<SweepRecord, ExploreError> {
let SearchWriters {
session,
writer,
tables,
columns,
} = output;
let counts = writer.counts();
let end = finish_manifest(&mut manifest, end, counts);
manifest.search = Some(session.manifest_search());
let mut files = SweepFiles::assemble(writer, &manifest)?;
let (evaluations, batches, generations) = tables
.into_inner()
.map_err(|source| table_error(Path::new(EVALUATIONS_FILE), source))?;
let closing_file = closing_table_file(&self.plan);
let search_report = session.report();
let closing = session
.write_closing_table(&search_report, Vec::new(), &columns)
.map_err(|source| table_error(Path::new(closing_file), source))?;
files.search_tables.push((EVALUATIONS_FILE, evaluations));
files.search_tables.push((BATCHES_FILE, batches));
if let Some(generations) = generations {
files.search_tables.push((GENERATIONS_FILE, generations));
}
files.search_tables.push((closing_file, closing));
Ok(SweepRecord {
report: self.report(end, counts, None),
manifest,
files: Some(files),
search: Some(search_report),
})
}
pub(crate) fn report(&self, end: SweepEnd, counts: ResultCounts, output_dir: Option<PathBuf>) -> SweepReport {
SweepReport {
outline: self.outline.clone(),
end,
counts,
elapsed: self.started.elapsed(),
output_dir,
}
}
}
fn table_error(path: &Path, source: io::Error) -> ExploreError {
ExploreError::Output(OutputError::Write {
path: path.to_owned(),
source,
})
}
#[cfg(not(target_arch = "wasm32"))]
struct SearchSink<'s, W: Write> {
batch: &'s AskedBatch,
writer: &'s mut OutputWriter<W>,
watched_reducers: &'s [usize],
values: Vec<Vec<Option<f64>>>,
progress: &'s mut dyn Progress,
meter: &'s mut ProgressMeter,
}
#[cfg(not(target_arch = "wasm32"))]
impl<W: Write> RunSink for SearchSink<'_, W> {
fn commit(&mut self, outcome: RunOutcome) -> io::Result<()> {
self.writer
.write_config_run(&outcome, self.batch.config_of(&outcome.run))?;
self.values.push(watched_values(&outcome, self.watched_reducers));
self.progress.report(&ProgressEvent::RunCommitted(&outcome));
Ok(())
}
fn finished(&mut self, outcome: &RunOutcome) {
if let Some(update) = self.meter.record_finished_run(outcome.status) {
self.progress.report(&ProgressEvent::Progressed(update));
}
}
}