1use std::collections::{BTreeMap, BTreeSet};
15use std::fmt;
16#[cfg(not(target_arch = "wasm32"))]
17use std::fs::OpenOptions;
18use std::io::{self, Write};
19use std::path::{Path, PathBuf};
20use std::sync::Arc;
21
22use web_time::Instant;
23
24use henad_core::action::Schedule;
25use henad_core::explore::factor::FactorSpec;
26use henad_core::explore::fingerprint::{run_key, search_hash};
27use henad_core::explore::measure::MeasurePlan;
28use henad_core::explore::outcome::{PlannedRun, RunOutcome};
29use henad_core::explore::plan::{Config, MAX_RUNS, ModelSchema, Plan, PlanError, Shard};
30use henad_core::explore::replay::Replay;
31use henad_core::explore::search::genome::{Genome, SearchSpace, SearchSpaceError};
32use henad_core::explore::search::pse::{
33 ArchiveEntry, PatternAxis, PatternCell, PatternPlacement, PatternSpaceSettings,
34};
35use henad_core::explore::search::{
36 Candidate, CandidateOrigin, Evaluation, GenerationSummary, Objective, RankingEntry, SearchAlgorithm, SearchReport,
37 SearchSpec, SearchSpecError, Searcher,
38};
39use henad_core::explore::seed::search_seed;
40use henad_core::explore::spec::SweepSpec;
41
42use crate::exec::{BatchEnd, RunRequest};
43#[cfg(not(target_arch = "wasm32"))]
44use crate::exec::{Executor, RunSink};
45use crate::output::manifest::{
46 Manifest, ManifestAxisRanges, ManifestMode, ManifestPlan, ManifestSearch, ResultCounts, now_unix_ms,
47};
48use crate::output::memory::{SweepFiles, memory_writer};
49use crate::output::read::{ReadError, RunsCsv, SeriesScan, parse_one};
50use crate::output::resume::{ResumeError, check_model};
51use crate::output::runs_csv::{ID_COLUMNS, OUTCOME_COLUMNS, column_names};
52use crate::output::search_tables::{ConfigColumns, SearchTablesWriter, write_archive, write_ranking};
53use crate::output::{
54 ARCHIVE_FILE, BATCHES_FILE, BEST_FILE, EVALUATIONS_FILE, GENERATIONS_FILE, MANIFEST_FILE, OutputDir, OutputError,
55 OutputWriter, RUNS_FILE, SERIES_FILE, runs_csv, series_csv, table_paths,
56};
57use crate::probe::{ProbeReport, TimedProbe, check_capacity};
58#[cfg(not(target_arch = "wasm32"))]
59use crate::progress::ProgressMeter;
60use crate::progress::{Progress, ProgressEvent};
61use crate::sweep::{
62 ExploreError, ManifestParts, SweepEnd, SweepInputs, SweepOutline, SweepRecord, SweepReport, SweepWarning,
63 build_warnings, finish_manifest, hex, running_manifest, sized_layout,
64};
65
66#[derive(Debug, Clone)]
68pub struct SearchPlan {
69 base: Arc<Plan>,
71 search: SearchSpec,
72 space: SearchSpace,
73 search_hash: u64,
74 run_count: u64,
75}
76
77impl SearchPlan {
78 pub fn new(spec: &SweepSpec, schema: &ModelSchema<'_>) -> Result<Self, SearchPlanError> {
86 let search = spec.search.as_ref().ok_or(SearchPlanError::NotASearch)?;
87 if !spec.blocks.is_empty() {
88 return Err(SearchPlanError::Blocks);
89 }
90 search.check().map_err(SearchPlanError::Settings)?;
91 let base = spec.plan(schema).map_err(SearchPlanError::Plan)?;
92 let space = SearchSpace::resolve(&search.space, schema.params, &spec.actions, &spec.fixed)
93 .map_err(SearchPlanError::Space)?;
94 if (search.batch_size as u64).saturating_mul(base.replicates()) > MAX_RUNS {
95 return Err(SearchPlanError::BatchTooLarge {
96 batch_size: search.batch_size,
97 replicates: base.replicates(),
98 });
99 }
100 let run_count = search
101 .max_evaluations
102 .checked_mul(base.replicates())
103 .ok_or(SearchPlanError::TooManyRuns)?;
104 Ok(Self {
105 search_hash: search_hash(base.plan_hash(), base.replicates(), search),
106 base: Arc::new(base),
107 search: search.clone(),
108 space,
109 run_count,
110 })
111 }
112
113 pub fn base(&self) -> &Arc<Plan> {
115 &self.base
116 }
117
118 pub fn search(&self) -> &SearchSpec {
120 &self.search
121 }
122
123 pub fn space(&self) -> &SearchSpace {
125 &self.space
126 }
127
128 pub fn search_hash(&self) -> u64 {
130 self.search_hash
131 }
132
133 pub fn run_count(&self) -> u64 {
135 self.run_count
136 }
137
138 pub fn replicates(&self) -> u64 {
140 self.base.replicates()
141 }
142
143 pub fn config(&self, genome: &Genome) -> Config {
149 let base = self.base.config(0).expect("a plan with no blocks has one config");
150 self.space.decode(genome, base)
151 }
152
153 pub fn run(&self, candidate: &Candidate, index: u64) -> PlannedRun {
157 let replicates = self.replicates();
158 let replicate = candidate.replicate_offset + index;
159 let seeds = self.base.seed_settings();
160 let seed_config_id = candidate.origin.reevaluated_id().unwrap_or(candidate.id);
161 PlannedRun {
162 run_id: candidate.id * replicates + index,
163 config_id: candidate.id,
164 rep: replicate,
165 seed: seeds.scheme.seed(seeds.root, seed_config_id, replicate),
166 }
167 }
168
169 pub fn run_key(&self, run: &PlannedRun, config: &Config) -> u64 {
171 run_key(
172 self.base.results_fingerprint(),
173 &config.params,
174 &config.action_ticks,
175 run.seed,
176 )
177 }
178
179 pub fn replay(&self, run: &PlannedRun, config: &Config) -> Replay {
181 let settings = self.base.run_settings();
182 Replay {
183 model: self.base.model().to_owned(),
184 params: config.params.clone(),
185 seed: run.seed,
186 schedule: self.base.schedule(config),
187 ticks: settings.warmup + settings.steps,
188 label: format!(
189 "Search run {}: candidate {}, replicate {}",
190 run.run_id, run.config_id, run.rep
191 ),
192 }
193 }
194
195 pub fn watched_reducers(&self, measure: &MeasurePlan) -> Result<Vec<usize>, SearchPlanError> {
201 let names = measure.reducers().names();
202 self.search
203 .watched_columns()
204 .into_iter()
205 .map(|column| {
206 names
207 .iter()
208 .position(|name| name == column)
209 .ok_or_else(|| SearchPlanError::UnknownColumn {
210 column: column.to_owned(),
211 known: names.to_vec(),
212 })
213 })
214 .collect()
215 }
216
217 pub fn outline(&self) -> SearchOutline {
219 SearchOutline {
220 algorithm: self.search.algorithm.as_str(),
221 max_evaluations: self.search.max_evaluations,
222 batch_size: self.search.batch_size,
223 objective: self.search.objective.clone(),
224 watched_columns: self.search.watched_columns().into_iter().map(str::to_owned).collect(),
225 pattern_axes: match &self.search.algorithm {
226 SearchAlgorithm::PatternSpaceExploration(settings) => {
227 Some([settings.x_axis.clone(), settings.y_axis.clone()])
228 }
229 SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => None,
230 },
231 space: self.search.space.clone(),
232 search_seed: search_seed(self.base.seed_settings().root),
233 }
234 }
235}
236
237#[derive(Debug, Clone, PartialEq)]
239pub enum SearchPlanError {
240 NotASearch,
242 Blocks,
244 Plan(PlanError),
246 Settings(SearchSpecError),
248 Space(SearchSpaceError),
250 UnknownColumn {
252 column: String,
254 known: Vec<String>,
256 },
257 BatchTooLarge {
259 batch_size: usize,
261 replicates: u64,
263 },
264 TooManyRuns,
266 Sharded,
268 RetryFailed,
270}
271
272impl fmt::Display for SearchPlanError {
273 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
274 match self {
275 Self::NotASearch => f.write_str("the spec has no [search] table"),
276 Self::Blocks => f.write_str("a search selects its own configs and takes no blocks"),
277 Self::Plan(_) => f.write_str("cannot plan the search's fixed values and actions"),
278 Self::Settings(_) => f.write_str("search settings"),
279 Self::Space(_) => f.write_str("search space"),
280 Self::UnknownColumn { column, known } if known.is_empty() => {
281 write!(f, "search watches '{column}', and the runs have no reducer columns")
282 }
283 Self::UnknownColumn { column, known } => write!(
284 f,
285 "no reducer writes '{column}', the column the search watches. The reducer columns are {}",
286 known.join(", ")
287 ),
288 Self::BatchTooLarge { batch_size, replicates } => write!(
289 f,
290 "a batch of {batch_size} candidates at {replicates} replicates each has more than {MAX_RUNS} runs"
291 ),
292 Self::TooManyRuns => f.write_str("search budget has more runs than a 64-bit integer can hold"),
293 Self::Sharded => f.write_str("a search cannot run as a shard"),
294 Self::RetryFailed => {
295 f.write_str("a resumed search replays its failed runs as they were, and cannot run them again")
296 }
297 }
298 }
299}
300
301impl std::error::Error for SearchPlanError {
302 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
303 match self {
304 Self::Plan(error) => Some(error),
305 Self::Settings(error) => Some(error),
306 Self::Space(error) => Some(error),
307 Self::NotASearch
308 | Self::Blocks
309 | Self::UnknownColumn { .. }
310 | Self::BatchTooLarge { .. }
311 | Self::TooManyRuns
312 | Self::Sharded
313 | Self::RetryFailed => None,
314 }
315 }
316}
317
318#[derive(Debug, Clone, PartialEq)]
320pub struct SearchOutline {
321 pub algorithm: &'static str,
323 pub max_evaluations: u64,
325 pub batch_size: usize,
327 pub objective: Option<Objective>,
329 pub watched_columns: Vec<String>,
331 pub pattern_axes: Option<[PatternAxis; 2]>,
333 pub space: Vec<FactorSpec>,
335 pub search_seed: u64,
337}
338
339#[derive(Debug, Clone, PartialEq)]
341pub struct SearchUpdate {
342 pub batch: u64,
344 pub evaluations: u64,
346 pub runs: u64,
348 pub evaluated: Vec<EvaluatedCandidate>,
350 pub best: Option<RankingEntry>,
352 pub generations: Vec<GenerationSummary>,
354 pub landed_entries: Vec<ArchiveEntry>,
357 pub filled_cells: u64,
359 pub pattern_settings: Option<PatternSpaceSettings>,
362}
363
364#[derive(Debug, Clone, PartialEq)]
366pub struct EvaluatedCandidate {
367 pub candidate_id: u64,
369 pub batch: u64,
371 pub origin: CandidateOrigin,
373 pub replicate_offset: u64,
375 pub replicates: u64,
377 pub config: Config,
379 pub failed_count: u64,
381 pub reading: EvaluationReading,
383}
384
385#[derive(Debug, Clone, Copy, PartialEq)]
387pub enum EvaluationReading {
388 Objective {
390 objective: f64,
392 pooled_objective: f64,
395 pooled_replicates: u64,
397 },
398 Placement {
400 placement: Option<PatternPlacement>,
403 new_cell: bool,
405 },
406}
407
408pub(crate) struct SearchSession {
410 plan: Arc<SearchPlan>,
411 searcher: Box<dyn Searcher + Send>,
412 watched_reducers: Vec<usize>,
414 recorded_runs: Arc<[RecordedRun]>,
416 configs: BTreeMap<u64, Config>,
418 batch_count: u64,
419 evaluations: u64,
420 runs: u64,
421 reported_generations: usize,
423 reported_ranges: bool,
425}
426
427pub(crate) struct AskedBatch {
429 index: u64,
430 candidates: Vec<Candidate>,
431 configs: Vec<Config>,
433 schedules: Vec<Schedule>,
435 runs: Vec<BatchRun>,
436 recorded_values: Vec<Vec<Option<f64>>>,
438}
439
440struct BatchRun {
442 run: PlannedRun,
443 run_key: u64,
444 candidate_index: usize,
446}
447
448impl AskedBatch {
449 #[cfg(not(target_arch = "wasm32"))]
451 pub(crate) fn requests(&self) -> Vec<RunRequest<'_>> {
452 self.runs[self.recorded_values.len()..]
453 .iter()
454 .map(|batch_run| self.request_of(batch_run))
455 .collect()
456 }
457
458 #[cfg(any(target_arch = "wasm32", test))]
460 pub(crate) fn request(&self, position: usize) -> Option<RunRequest<'_>> {
461 let batch_run = self.runs.get(self.recorded_values.len() + position)?;
462 Some(self.request_of(batch_run))
463 }
464
465 fn request_of(&self, batch_run: &BatchRun) -> RunRequest<'_> {
466 RunRequest {
467 run: batch_run.run,
468 run_key: batch_run.run_key,
469 params: &self.configs[batch_run.candidate_index].params,
470 schedule: self.schedules[batch_run.candidate_index].clone(),
471 }
472 }
473
474 pub(crate) fn config_of(&self, run: &PlannedRun) -> &Config {
480 let index = self
481 .candidates
482 .binary_search_by_key(&run.config_id, |candidate| candidate.id)
483 .expect("a run of the batch belongs to one of its candidates");
484 &self.configs[index]
485 }
486}
487
488impl SearchSession {
489 pub(crate) fn new(
496 plan: Arc<SearchPlan>,
497 watched_reducers: Vec<usize>,
498 recorded_runs: Arc<[RecordedRun]>,
499 ) -> Result<Self, SearchPlanError> {
500 let root = plan.base.seed_settings().root;
501 let searcher = plan
502 .search
503 .searcher(&plan.space, root)
504 .map_err(SearchPlanError::Settings)?;
505 Ok(Self {
506 plan,
507 searcher,
508 watched_reducers,
509 recorded_runs,
510 configs: BTreeMap::new(),
511 batch_count: 0,
512 evaluations: 0,
513 runs: 0,
514 reported_generations: 0,
515 reported_ranges: false,
516 })
517 }
518
519 pub(crate) fn watched_reducers(&self) -> &[usize] {
521 &self.watched_reducers
522 }
523
524 pub(crate) fn ask(&mut self) -> Result<Option<AskedBatch>, ResumeError> {
532 if self.searcher.is_done() {
533 return Ok(None);
534 }
535 let candidates = self.searcher.ask(self.plan.search.batch_size);
536 if candidates.is_empty() {
537 return Ok(None);
538 }
539 let replicates = self.plan.replicates();
540 let mut batch = AskedBatch {
541 index: self.batch_count,
542 configs: Vec::with_capacity(candidates.len()),
543 schedules: Vec::with_capacity(candidates.len()),
544 runs: Vec::with_capacity(candidates.len() * replicates as usize),
545 recorded_values: Vec::new(),
546 candidates: Vec::new(),
547 };
548 for (candidate_index, candidate) in candidates.iter().enumerate() {
549 let config = self.plan.config(&candidate.genome);
550 for index in 0..replicates {
551 let run = self.plan.run(candidate, index);
552 batch.runs.push(BatchRun {
553 run,
554 run_key: self.plan.run_key(&run, &config),
555 candidate_index,
556 });
557 }
558 if candidate.origin.reevaluated_id().is_none() {
559 self.configs.insert(candidate.id, config.clone());
560 }
561 batch.schedules.push(self.plan.base.schedule(&config));
562 batch.configs.push(config);
563 }
564 for batch_run in &batch.runs {
565 let Some(recorded) = usize::try_from(batch_run.run.run_id)
566 .ok()
567 .and_then(|run_id| self.recorded_runs.get(run_id))
568 else {
569 break;
570 };
571 if recorded.run_key != batch_run.run_key {
572 return Err(ResumeError::SearchRunChanged {
573 run_id: batch_run.run.run_id,
574 });
575 }
576 batch.recorded_values.push(recorded.values.clone());
577 }
578 batch.candidates = candidates;
579 Ok(Some(batch))
580 }
581
582 pub(crate) fn replay_recorded_runs(&mut self) -> Result<(), ResumeError> {
590 while let Some(batch) = self.ask()? {
591 if batch.recorded_values.len() < batch.runs.len() {
592 break;
593 }
594 self.tell(batch, Vec::new());
595 }
596 Ok(())
597 }
598
599 pub(crate) fn tell(&mut self, batch: AskedBatch, values: Vec<Vec<Option<f64>>>) -> SearchUpdate {
608 let AskedBatch {
609 index,
610 candidates,
611 configs,
612 runs,
613 mut recorded_values,
614 ..
615 } = batch;
616 recorded_values.extend(values);
617 let values = recorded_values;
618 assert_eq!(values.len(), runs.len(), "every run of the batch has its values");
619 let replicates = self.plan.replicates() as usize;
620 let evaluations: Vec<Evaluation> = candidates
621 .iter()
622 .zip(values.chunks(replicates.max(1)))
623 .map(|(candidate, outputs)| Evaluation {
624 candidate_id: candidate.id,
625 outputs: outputs.to_vec(),
626 })
627 .collect();
628 self.searcher.tell(&evaluations);
629 let searcher = &self.searcher;
630 let pattern_settings = match &self.plan.search.algorithm {
631 SearchAlgorithm::PatternSpaceExploration(settings) => Some(searcher.pattern_settings().unwrap_or(settings)),
632 SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => None,
633 };
634 let mut touched = BTreeSet::new();
635 let evaluated: Vec<EvaluatedCandidate> = candidates
636 .into_iter()
637 .zip(configs)
638 .zip(&evaluations)
639 .map(|((candidate, config), evaluation)| {
640 let failed_count = evaluation
641 .outputs
642 .iter()
643 .filter(|row| row.iter().any(|value| !value.is_some_and(f64::is_finite)))
644 .count() as u64;
645 let reading = if let Some(settings) = pattern_settings {
646 let placement = settings.place(evaluation);
647 let cell = placement.and_then(|placement| placement.cell);
648 if let Some(cell) = cell {
649 touched.insert(cell);
650 }
651 let new_cell = cell.is_some_and(|cell| {
652 searcher
653 .archive_entry(cell)
654 .is_some_and(|entry| entry.candidate_id == candidate.id)
655 });
656 EvaluationReading::Placement { placement, new_cell }
657 } else {
658 let objective = self
659 .plan
660 .search
661 .objective
662 .as_ref()
663 .expect("a scored search has an objective");
664 let pooled_id = candidate.origin.reevaluated_id().unwrap_or(candidate.id);
665 let pooled = searcher.ranking_entry(pooled_id);
666 EvaluationReading::Objective {
667 objective: objective.score(evaluation),
668 pooled_objective: pooled.as_ref().map_or(objective.goal.worst(), |entry| entry.objective),
669 pooled_replicates: pooled.as_ref().map_or(0, |entry| entry.replicate_count),
670 }
671 };
672 EvaluatedCandidate {
673 candidate_id: candidate.id,
674 batch: index,
675 origin: candidate.origin,
676 replicate_offset: candidate.replicate_offset,
677 replicates: replicates as u64,
678 config,
679 failed_count,
680 reading,
681 }
682 })
683 .collect();
684 self.batch_count += 1;
685 self.evaluations += evaluated.len() as u64;
686 self.runs += runs.len() as u64;
687 let generations = searcher.generations();
688 let new_generations = generations[self.reported_generations.min(generations.len())..].to_vec();
689 self.reported_generations = generations.len();
690 let (best, filled_cells) = (searcher.best(), searcher.filled_cells());
691 let pattern_settings = searcher.pattern_settings().cloned();
692 SearchUpdate {
693 batch: index,
694 evaluations: self.evaluations,
695 runs: self.runs,
696 evaluated,
697 best,
698 generations: new_generations,
699 landed_entries: self.landed_entries(touched),
700 filled_cells,
701 pattern_settings,
702 }
703 }
704
705 fn landed_entries(&mut self, touched: BTreeSet<PatternCell>) -> Vec<ArchiveEntry> {
710 if !self.reported_ranges && self.searcher.pattern_settings().is_some() {
711 self.reported_ranges = true;
712 return self.searcher.report().archive;
713 }
714 touched
715 .into_iter()
716 .filter_map(|cell| self.searcher.archive_entry(cell))
717 .collect()
718 }
719
720 pub(crate) fn report(&self) -> SearchReport {
722 self.searcher.report()
723 }
724
725 pub(crate) fn manifest_search(&self) -> ManifestSearch {
727 let mut record = manifest_search(&self.plan);
728 record.evaluations = self.evaluations;
729 record.batch_count = self.batch_count;
730 match &self.plan.search.algorithm {
731 SearchAlgorithm::PatternSpaceExploration(_) => {
732 record.filled_cells = Some(self.searcher.filled_cells());
733 record.axis_ranges = self.searcher.pattern_settings().and_then(ManifestAxisRanges::of);
734 }
735 SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => {
736 let best = self.searcher.best();
737 record.best_candidate_id = best.as_ref().map(|entry| entry.candidate_id);
738 record.best_objective = best
739 .map(|entry| entry.objective)
740 .filter(|objective| objective.is_finite());
741 }
742 }
743 record
744 }
745
746 pub(crate) fn write_closing_table<W: Write>(
752 &self,
753 report: &SearchReport,
754 dest: W,
755 columns: &ConfigColumns,
756 ) -> io::Result<W> {
757 match &self.plan.search.algorithm {
758 SearchAlgorithm::PatternSpaceExploration(settings) => {
759 let settings = self.searcher.pattern_settings().unwrap_or(settings);
760 write_archive(dest, settings, &report.archive, &self.configs, columns)
761 }
762 SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => {
763 write_ranking(dest, &report.ranking, &self.configs, columns)
764 }
765 }
766 }
767}
768
769fn manifest_search(plan: &SearchPlan) -> ManifestSearch {
771 let pattern_settings = match &plan.search.algorithm {
772 SearchAlgorithm::PatternSpaceExploration(settings) => Some(settings),
773 SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => None,
774 };
775 ManifestSearch {
776 algorithm: plan.search.algorithm.as_str().to_owned(),
777 search_hash: hex(plan.search_hash),
778 max_evaluations: plan.search.max_evaluations,
779 batch_size: plan.search.batch_size,
780 search_seed: search_seed(plan.base.seed_settings().root),
781 watched_columns: plan.search.watched_columns().into_iter().map(str::to_owned).collect(),
782 evaluations: 0,
783 batch_count: 0,
784 best_candidate_id: None,
785 best_objective: None,
786 filled_cells: pattern_settings.map(|_| 0),
787 axis_ranges: pattern_settings.and_then(ManifestAxisRanges::of),
788 }
789}
790
791fn closing_table_file(plan: &SearchPlan) -> &'static str {
793 match plan.search.algorithm {
794 SearchAlgorithm::PatternSpaceExploration(_) => ARCHIVE_FILE,
795 SearchAlgorithm::Random | SearchAlgorithm::HillClimb(_) | SearchAlgorithm::Genetic(_) => BEST_FILE,
796 }
797}
798
799fn writes_generations(plan: &SearchPlan) -> bool {
801 matches!(plan.search.algorithm, SearchAlgorithm::Genetic(_))
802}
803
804pub(crate) fn watched_values(outcome: &RunOutcome, watched_reducers: &[usize]) -> Vec<Option<f64>> {
809 watched_reducers
810 .iter()
811 .map(|&position| {
812 if outcome.status.is_failure() {
813 None
814 } else {
815 outcome
816 .reducers
817 .get(position)
818 .copied()
819 .flatten()
820 .filter(|value| value.is_finite())
821 }
822 })
823 .collect()
824}
825
826#[derive(Debug, Clone, PartialEq)]
828pub(crate) struct RecordedRun {
829 run_key: u64,
830 values: Vec<Option<f64>>,
832}
833
834#[derive(Debug)]
836#[cfg_attr(target_arch = "wasm32", expect(dead_code, reason = "a browser resumes no folder"))]
837pub(crate) struct RecordedSearch {
838 pub(crate) recorded: Manifest,
840 runs: Arc<[RecordedRun]>,
842 counts: ResultCounts,
843 runs_bytes: u64,
845 series_bytes: u64,
848 fresh_tables: bool,
850}
851
852impl RecordedSearch {
853 pub(crate) fn read(
865 path: &Path,
866 plan: &SearchPlan,
867 runs_header: &str,
868 series_header: &str,
869 watched_reducers: &[usize],
870 config_column_count: usize,
871 ) -> Result<Self, ResumeError> {
872 let recorded = Manifest::read(&path.join(MANIFEST_FILE)).map_err(ResumeError::Manifest)?;
873 if recorded.mode != ManifestMode::Search {
874 return Err(ResumeError::ModeChanged {
875 recorded: recorded.mode,
876 current: ManifestMode::Search,
877 });
878 }
879 check_model(&recorded, plan.base.model())?;
880 let current_schema = hex(plan.base.schema_hash());
881 if recorded.model.schema_hash != current_schema {
882 return Err(ResumeError::SchemaChanged {
883 recorded: recorded.model.schema_hash.clone(),
884 current: current_schema,
885 });
886 }
887 let current_search = hex(plan.search_hash);
888 let recorded_search = recorded
889 .search
890 .as_ref()
891 .map_or_else(String::new, |search| search.search_hash.clone());
892 if recorded_search != current_search {
893 return Err(ResumeError::SearchChanged {
894 recorded: recorded_search,
895 current: current_search,
896 });
897 }
898 let (runs_path, series_path) = table_paths(path);
899 let table = RunsCsv::read(&runs_path).map_err(ResumeError::Table)?;
900 if let Some(header) = &table.header
901 && runs_csv::header_line(header) != runs_header
902 {
903 return Err(ResumeError::ColumnsChanged { file: RUNS_FILE });
904 }
905 let reducers_start = ID_COLUMNS.len() + config_column_count + OUTCOME_COLUMNS.len();
906 let mut runs = Vec::with_capacity(table.records.len());
907 let mut counts = ResultCounts::default();
908 for (position, record) in table.records.iter().enumerate() {
909 if position as u64 >= plan.run_count() {
910 return Err(ResumeError::UnknownRun { run_id: record.run_id });
911 }
912 if record.run_id != position as u64 {
913 return Err(ResumeError::SearchRunChanged { run_id: record.run_id });
914 }
915 let fields =
916 parse_one(&record.text, 0..record.text.len(), &runs_path, position + 2).map_err(ResumeError::Table)?;
917 let values = watched_reducers
918 .iter()
919 .map(|&reducer| {
920 let column = reducers_start + reducer;
921 let field = fields.get(column).map_or("", String::as_str);
922 if record.status.is_failure() || field.is_empty() {
923 return Ok(None);
924 }
925 field.parse::<f64>().ok().map(Some).ok_or_else(|| {
926 ResumeError::Table(ReadError::BadField {
927 path: runs_path.clone(),
928 record_number: position + 2,
929 column: table
930 .header
931 .as_ref()
932 .and_then(|header| header.get(column))
933 .cloned()
934 .unwrap_or_default(),
935 text: field.to_owned(),
936 })
937 })
938 })
939 .collect::<Result<_, _>>()?;
940 counts.count(record.status);
941 runs.push(RecordedRun {
942 run_key: record.run_key,
943 values,
944 });
945 }
946 let recorded_runs = runs.len() as u64;
947 let series = SeriesScan::read(&series_path, |run_id| run_id < recorded_runs).map_err(ResumeError::Table)?;
948 if series.header.as_deref().is_some_and(|header| header != series_header) {
949 return Err(ResumeError::ColumnsChanged { file: SERIES_FILE });
950 }
951 if series.kept_after_dropped || series.segments.len() > 1 {
952 return Err(ResumeError::SeriesOutOfOrder);
953 }
954 let fresh_tables = table.header.is_none() || series.header.is_none();
955 if fresh_tables {
956 runs.clear();
957 counts = ResultCounts::default();
958 }
959 Ok(Self {
960 recorded,
961 runs: runs.into(),
962 counts,
963 runs_bytes: table.complete_bytes,
964 series_bytes: series.first_dropped_offset.unwrap_or(series.complete_bytes),
965 fresh_tables,
966 })
967 }
968
969 #[cfg(not(target_arch = "wasm32"))]
971 fn repair(&self, dir: &OutputDir) -> Result<(), OutputError> {
972 for (file, length) in [(RUNS_FILE, self.runs_bytes), (SERIES_FILE, self.series_bytes)] {
973 let path = dir.path().join(file);
974 let repaired = if self.fresh_tables {
975 match std::fs::remove_file(&path) {
976 Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
977 removed => removed,
978 }
979 } else {
980 OpenOptions::new()
981 .write(true)
982 .open(&path)
983 .and_then(|table| table.set_len(length))
984 };
985 repaired.map_err(|source| OutputError::Write { path, source })?;
986 }
987 Ok(())
988 }
989}
990
991#[cfg(not(target_arch = "wasm32"))]
1008pub(crate) fn run_search_into_directory(
1009 inputs: &SweepInputs<'_>,
1010 plan: Option<Arc<SearchPlan>>,
1011 output_dir: &Path,
1012 progress: &mut dyn Progress,
1013) -> Result<SweepRecord, ExploreError> {
1014 let preparation = SearchPreparation::new(inputs, plan, None)?;
1015 preparation.announce(inputs, progress);
1016 let record = preparation.write_directory(inputs, output_dir, progress)?;
1017 progress.report(&ProgressEvent::Ended(&record.report));
1018 Ok(record)
1019}
1020
1021#[cfg(not(target_arch = "wasm32"))]
1030pub(crate) fn run_search_in_memory(
1031 inputs: &SweepInputs<'_>,
1032 plan: Option<Arc<SearchPlan>>,
1033 progress: &mut dyn Progress,
1034) -> Result<SweepRecord, ExploreError> {
1035 let preparation = SearchPreparation::new(inputs, plan, None)?;
1036 preparation.announce(inputs, progress);
1037 let (mut output, manifest) = preparation.memory_output(inputs)?;
1038 let executor = preparation.executor(inputs)?;
1039 let end = preparation.run_all(&executor, &mut output, progress)?;
1040 let record = preparation.finish_in_memory(output, manifest, end)?;
1041 progress.report(&ProgressEvent::Ended(&record.report));
1042 Ok(record)
1043}
1044
1045pub(crate) struct SearchPreparation {
1047 plan: Arc<SearchPlan>,
1048 probe: ProbeReport,
1049 measure: Arc<MeasurePlan>,
1050 watched_reducers: Vec<usize>,
1052 resumed: Option<RecordedSearch>,
1054 #[cfg_attr(
1057 target_arch = "wasm32",
1058 expect(dead_code, reason = "a search in a browser writes no directory")
1059 )]
1060 locked: Option<OutputDir>,
1061 outline: SweepOutline,
1062 started: Instant,
1064 started_unix_ms: u64,
1065}
1066
1067pub(crate) struct SearchWriters<W: Write> {
1069 pub(crate) session: SearchSession,
1070 pub(crate) writer: OutputWriter<W>,
1071 pub(crate) tables: SearchTablesWriter<W>,
1072 pub(crate) columns: ConfigColumns,
1074}
1075
1076impl SearchPreparation {
1077 pub(crate) fn new(
1081 inputs: &SweepInputs<'_>,
1082 plan: Option<Arc<SearchPlan>>,
1083 probe: Option<TimedProbe>,
1084 ) -> Result<Self, ExploreError> {
1085 let (entry, gpu, options) = (inputs.entry, inputs.gpu, inputs.options);
1086 if options.shard != Shard::WHOLE {
1087 return Err(SearchPlanError::Sharded.into());
1088 }
1089 if options.retry_failed {
1090 return Err(SearchPlanError::RetryFailed.into());
1091 }
1092 let (started, started_unix_ms, probe) = match probe {
1093 Some(timed) => (timed.started, timed.started_unix_ms, Some(timed.report)),
1094 None => (Instant::now(), now_unix_ms(), None),
1095 };
1096 let plan = match plan {
1097 Some(plan) => plan,
1098 None => Arc::new(SearchPlan::new(inputs.spec, &entry.schema())?),
1099 };
1100 let resume_dir = inputs
1101 .folder
1102 .filter(|output_dir| options.resume && OutputDir::holds_results(output_dir));
1103 let locked = match (resume_dir, inputs.folder) {
1104 (None, Some(output_dir)) => {
1105 OutputDir::check_free(output_dir)?;
1106 None
1107 }
1108 (Some(output_dir), _) if !inputs.dry_run => Some(OutputDir::open(output_dir)?),
1111 _ => None,
1112 };
1113 let base = plan.base();
1114 if let Some(ctx) = gpu {
1115 check_capacity(entry, base, &ctx.device.limits())?;
1116 }
1117 let probe = match probe {
1118 Some(probe) => probe,
1119 None => ProbeReport::for_plan(entry, gpu, base)?,
1120 };
1121 let measure = MeasurePlan::new(base.run_settings(), base.measure_settings(), probe.columns.clone())?;
1122 let watched_reducers = plan.watched_reducers(&measure)?;
1123 let resumed = resume_dir
1124 .map(|output_dir| {
1125 let params = entry.param_descriptors();
1126 let columns = column_names(params, base.actions(), measure.reducers().names());
1127 RecordedSearch::read(
1128 output_dir,
1129 &plan,
1130 &runs_csv::header_line(&columns),
1131 &series_csv::header_line(measure.columns()),
1132 &watched_reducers,
1133 params.len() + base.actions().len(),
1134 )
1135 })
1136 .transpose()?;
1137 if let Some(recorded) = resumed.as_ref().filter(|recorded| !recorded.runs.is_empty()) {
1138 SearchSession::new(Arc::clone(&plan), watched_reducers.clone(), Arc::clone(&recorded.runs))?
1141 .replay_recorded_runs()?;
1142 }
1143 let skipped = resumed.as_ref().map_or(0, |recorded| recorded.runs.len() as u64);
1144 let pending = plan.run_count().saturating_sub(skipped);
1145 let batch_runs = (plan.search.batch_size as u64).saturating_mul(plan.replicates());
1146 let (layout, projected_bytes) = sized_layout(entry, gpu, options, &probe, pending.min(batch_runs))?;
1147 let outline = SweepOutline {
1148 model: entry.id().to_owned(),
1149 backend: entry.metadata().backend,
1150 configs: None,
1151 replicates: plan.replicates(),
1152 runs: plan.run_count(),
1153 blocks: Vec::new(),
1154 shard: Shard::WHOLE,
1155 skipped,
1156 pending,
1157 layout,
1158 projected_bytes,
1159 series_rows: pending.saturating_mul(measure.series_row_count()),
1160 stat_columns: (0..measure.columns().len())
1161 .map(|column| measure.columns().name(column).to_owned())
1162 .collect(),
1163 reducer_columns: measure.reducers().names().to_vec(),
1164 dry_run: inputs.dry_run,
1165 search: Some(plan.outline()),
1166 };
1167 Ok(Self {
1168 plan,
1169 probe,
1170 measure: Arc::new(measure),
1171 watched_reducers,
1172 resumed,
1173 locked,
1174 outline,
1175 started,
1176 started_unix_ms,
1177 })
1178 }
1179
1180 pub(crate) fn announce(&self, inputs: &SweepInputs<'_>, progress: &mut dyn Progress) {
1183 progress.report(&ProgressEvent::Planned(&self.outline));
1184 let mut warnings: Vec<SweepWarning> = self
1185 .plan
1186 .base
1187 .warnings()
1188 .iter()
1189 .cloned()
1190 .map(SweepWarning::Plan)
1191 .collect();
1192 if let Some(resumed) = &self.resumed {
1193 warnings.extend(build_warnings(&resumed.recorded, inputs.provenance, inputs.entry));
1194 }
1195 for warning in &warnings {
1196 progress.report(&ProgressEvent::Warned(warning));
1197 }
1198 }
1199
1200 #[cfg(any(target_arch = "wasm32", test))]
1201 pub(crate) fn measure(&self) -> &Arc<MeasurePlan> {
1202 &self.measure
1203 }
1204
1205 #[cfg(any(target_arch = "wasm32", test))]
1206 pub(crate) fn plan(&self) -> &Arc<SearchPlan> {
1207 &self.plan
1208 }
1209
1210 fn manifest(&self, inputs: &SweepInputs<'_>) -> Result<Manifest, OutputError> {
1212 running_manifest(
1213 inputs,
1214 &ManifestParts {
1215 mode: ManifestMode::Search,
1216 plan: &self.plan.base,
1217 probe: &self.probe,
1218 outline: &self.outline,
1219 manifest_plan: ManifestPlan {
1220 plan_hash: hex(self.plan.base.plan_hash()),
1221 results_fingerprint: hex(self.plan.base.results_fingerprint()),
1222 configs: self.outline.configs,
1223 replicates: self.outline.replicates,
1224 runs: self.outline.runs,
1225 blocks: Vec::new(),
1226 },
1227 started_unix_ms: self.started_unix_ms,
1228 recorded: self.resumed.as_ref().map(|resumed| &resumed.recorded),
1229 search: Some(manifest_search(&self.plan)),
1230 },
1231 )
1232 }
1233
1234 fn session(&self) -> Result<SearchSession, ExploreError> {
1236 let recorded = self
1237 .resumed
1238 .as_ref()
1239 .map(|resumed| Arc::clone(&resumed.runs))
1240 .unwrap_or_default();
1241 Ok(SearchSession::new(
1242 Arc::clone(&self.plan),
1243 self.watched_reducers.clone(),
1244 recorded,
1245 )?)
1246 }
1247
1248 pub(crate) fn memory_output(
1250 &self,
1251 inputs: &SweepInputs<'_>,
1252 ) -> Result<(SearchWriters<Vec<u8>>, Manifest), ExploreError> {
1253 let params = inputs.entry.param_descriptors();
1254 let writer = memory_writer(&self.plan.base, params, &self.measure)?;
1255 let generations = writes_generations(&self.plan).then(Vec::new);
1256 let tables = SearchTablesWriter::new(
1257 Vec::new(),
1258 Vec::new(),
1259 generations,
1260 params,
1261 self.plan.base.actions(),
1262 self.is_pattern_search(),
1263 )
1264 .map_err(|source| table_error(Path::new(EVALUATIONS_FILE), source))?;
1265 let output = SearchWriters {
1266 session: self.session()?,
1267 writer,
1268 tables,
1269 columns: ConfigColumns::new(params, self.plan.base.actions()),
1270 };
1271 Ok((output, self.manifest(inputs)?))
1272 }
1273
1274 fn is_pattern_search(&self) -> bool {
1275 matches!(self.plan.search.algorithm, SearchAlgorithm::PatternSpaceExploration(_))
1276 }
1277
1278 #[cfg(not(target_arch = "wasm32"))]
1281 fn write_directory(
1282 &self,
1283 inputs: &SweepInputs<'_>,
1284 output_dir: &Path,
1285 progress: &mut dyn Progress,
1286 ) -> Result<SweepRecord, ExploreError> {
1287 let created;
1288 let dir = if let Some(dir) = &self.locked {
1289 dir
1290 } else {
1291 created = OutputDir::create(output_dir)?;
1292 &created
1293 };
1294 let mut manifest = self.manifest(inputs)?;
1295 dir.write_manifest(&manifest)?;
1296 let mut standing = None;
1297 let (end, counts, session, search_report) = match self.run_into(dir, inputs, progress, &mut standing) {
1298 Ok(finished) => finished,
1299 Err(error) => {
1300 manifest.fail(now_unix_ms());
1301 if let Some(standing) = standing {
1302 manifest.search = Some(standing);
1303 }
1304 drop(dir.write_manifest(&manifest));
1306 return Err(error);
1307 }
1308 };
1309 let end = finish_manifest(&mut manifest, end, counts);
1310 manifest.search = Some(session.manifest_search());
1311 dir.write_manifest(&manifest)?;
1312 Ok(SweepRecord {
1313 report: self.report(end, counts, Some(output_dir.to_owned())),
1314 manifest,
1315 files: None,
1316 search: Some(search_report),
1317 })
1318 }
1319
1320 #[cfg(not(target_arch = "wasm32"))]
1325 fn run_into(
1326 &self,
1327 dir: &OutputDir,
1328 inputs: &SweepInputs<'_>,
1329 progress: &mut dyn Progress,
1330 standing: &mut Option<ManifestSearch>,
1331 ) -> Result<(BatchEnd, ResultCounts, SearchSession, SearchReport), ExploreError> {
1332 let params = inputs.entry.param_descriptors();
1333 let executor = self.executor(inputs)?;
1334 let fresh_tables = match &self.resumed {
1335 Some(resumed) => {
1336 resumed.repair(dir)?;
1337 resumed.fresh_tables
1338 }
1339 None => true,
1340 };
1341 let writer = if fresh_tables {
1342 dir.open_writer(&self.plan.base, params, &self.measure)?
1343 } else {
1344 dir.append_writer(&self.plan.base, params)?
1345 };
1346 let generations = if writes_generations(&self.plan) {
1347 Some(dir.create_table(GENERATIONS_FILE)?)
1348 } else {
1349 None
1350 };
1351 let tables = SearchTablesWriter::new(
1352 dir.create_table(EVALUATIONS_FILE)?,
1353 dir.create_table(BATCHES_FILE)?,
1354 generations,
1355 params,
1356 self.plan.base.actions(),
1357 self.is_pattern_search(),
1358 )
1359 .map_err(|source| table_error(&dir.path().join(EVALUATIONS_FILE), source))?;
1360 let mut output = SearchWriters {
1361 session: self.session()?,
1362 writer,
1363 tables,
1364 columns: ConfigColumns::new(params, self.plan.base.actions()),
1365 };
1366 let end = self.run_all(&executor, &mut output, progress);
1367 *standing = Some(output.session.manifest_search());
1368 let end = end?;
1369 let SearchWriters {
1370 session,
1371 writer,
1372 tables,
1373 columns,
1374 } = output;
1375 let written = writer.counts();
1376 writer.finish().map_err(|source| OutputError::Write {
1377 path: dir.path().to_owned(),
1378 source,
1379 })?;
1380 tables
1381 .into_inner()
1382 .map_err(|source| table_error(&dir.path().join(EVALUATIONS_FILE), source))?;
1383 let closing_file = closing_table_file(&self.plan);
1384 let search_report = session.report();
1385 session
1386 .write_closing_table(&search_report, dir.create_table(closing_file)?, &columns)
1387 .map_err(|source| table_error(&dir.path().join(closing_file), source))?;
1388 dir.write_summary()?;
1389 let kept = self.resumed.as_ref().map(|resumed| resumed.counts).unwrap_or_default();
1390 Ok((end, kept + written, session, search_report))
1391 }
1392
1393 #[cfg(not(target_arch = "wasm32"))]
1399 fn executor<'a>(&self, inputs: &SweepInputs<'a>) -> Result<Executor<'a>, ExploreError> {
1400 let options = inputs.options;
1401 let executor = Executor::new(
1402 inputs.entry,
1403 inputs.gpu,
1404 Arc::clone(&self.measure),
1405 self.outline.layout,
1406 options.control.clone(),
1407 )?;
1408 Ok(executor
1409 .with_timeout(self.plan.base.run_settings().timeout)
1410 .with_active_runs(options.active_runs.clone())
1411 .with_gpu_memory_budget(options.gpu_memory_budget))
1412 }
1413
1414 #[cfg(not(target_arch = "wasm32"))]
1419 fn run_all<W: Write>(
1420 &self,
1421 executor: &Executor<'_>,
1422 output: &mut SearchWriters<W>,
1423 progress: &mut dyn Progress,
1424 ) -> Result<BatchEnd, ExploreError> {
1425 let mut meter = ProgressMeter::new(self.outline.pending);
1426 loop {
1427 if executor.control().is_aborted() {
1428 return Ok(BatchEnd::Aborted);
1429 }
1430 let Some(batch) = output.session.ask()? else {
1431 return Ok(BatchEnd::Complete);
1432 };
1433 let requests = batch.requests();
1434 let mut sink = SearchSink {
1435 batch: &batch,
1436 writer: &mut output.writer,
1437 watched_reducers: output.session.watched_reducers(),
1438 values: Vec::with_capacity(requests.len()),
1439 progress: &mut *progress,
1440 meter: &mut meter,
1441 };
1442 let end = if requests.is_empty() {
1443 BatchEnd::Complete
1444 } else {
1445 executor.run_batch(&requests, &mut sink)?
1446 };
1447 let values = sink.values;
1448 drop(requests);
1449 if end != BatchEnd::Complete {
1450 return Ok(end);
1451 }
1452 let update = output.session.tell(batch, values);
1453 output
1454 .tables
1455 .write_batch(&update)
1456 .map_err(|source| table_error(Path::new(EVALUATIONS_FILE), source))?;
1457 progress.report(&ProgressEvent::SearchBatchTold(&update));
1458 }
1459 }
1460
1461 pub(crate) fn finish_in_memory(
1464 &self,
1465 output: SearchWriters<Vec<u8>>,
1466 mut manifest: Manifest,
1467 end: BatchEnd,
1468 ) -> Result<SweepRecord, ExploreError> {
1469 let SearchWriters {
1470 session,
1471 writer,
1472 tables,
1473 columns,
1474 } = output;
1475 let counts = writer.counts();
1476 let end = finish_manifest(&mut manifest, end, counts);
1477 manifest.search = Some(session.manifest_search());
1478 let mut files = SweepFiles::assemble(writer, &manifest)?;
1479 let (evaluations, batches, generations) = tables
1480 .into_inner()
1481 .map_err(|source| table_error(Path::new(EVALUATIONS_FILE), source))?;
1482 let closing_file = closing_table_file(&self.plan);
1483 let search_report = session.report();
1484 let closing = session
1485 .write_closing_table(&search_report, Vec::new(), &columns)
1486 .map_err(|source| table_error(Path::new(closing_file), source))?;
1487 files.search_tables.push((EVALUATIONS_FILE, evaluations));
1488 files.search_tables.push((BATCHES_FILE, batches));
1489 if let Some(generations) = generations {
1490 files.search_tables.push((GENERATIONS_FILE, generations));
1491 }
1492 files.search_tables.push((closing_file, closing));
1493 Ok(SweepRecord {
1494 report: self.report(end, counts, None),
1495 manifest,
1496 files: Some(files),
1497 search: Some(search_report),
1498 })
1499 }
1500
1501 pub(crate) fn report(&self, end: SweepEnd, counts: ResultCounts, output_dir: Option<PathBuf>) -> SweepReport {
1502 SweepReport {
1503 outline: self.outline.clone(),
1504 end,
1505 counts,
1506 elapsed: self.started.elapsed(),
1507 output_dir,
1508 }
1509 }
1510}
1511
1512fn table_error(path: &Path, source: io::Error) -> ExploreError {
1514 ExploreError::Output(OutputError::Write {
1515 path: path.to_owned(),
1516 source,
1517 })
1518}
1519
1520#[cfg(not(target_arch = "wasm32"))]
1522struct SearchSink<'s, W: Write> {
1523 batch: &'s AskedBatch,
1524 writer: &'s mut OutputWriter<W>,
1525 watched_reducers: &'s [usize],
1527 values: Vec<Vec<Option<f64>>>,
1529 progress: &'s mut dyn Progress,
1530 meter: &'s mut ProgressMeter,
1531}
1532
1533#[cfg(not(target_arch = "wasm32"))]
1534impl<W: Write> RunSink for SearchSink<'_, W> {
1535 fn commit(&mut self, outcome: RunOutcome) -> io::Result<()> {
1536 self.writer
1537 .write_config_run(&outcome, self.batch.config_of(&outcome.run))?;
1538 self.values.push(watched_values(&outcome, self.watched_reducers));
1539 self.progress.report(&ProgressEvent::RunCommitted(&outcome));
1540 Ok(())
1541 }
1542
1543 fn finished(&mut self, outcome: &RunOutcome) {
1544 if let Some(update) = self.meter.record_finished_run(outcome.status) {
1545 self.progress.report(&ProgressEvent::Progressed(update));
1546 }
1547 }
1548}