use std::fmt;
use std::sync::mpsc::{self, Receiver, Sender};
use std::sync::{Arc, Mutex, PoisonError};
use std::time::Duration;
use web_time::Instant;
use henad_compute::cpu::sim_thread::WakeFn;
use henad_compute::entry::ModelEntry;
use henad_compute::gpu::GpuContext;
use henad_core::explore::measure::SeriesBuffer;
use henad_core::explore::outcome::RunOutcome;
use henad_core::explore::plan::{Plan, PlanError};
use henad_core::explore::spec::SweepSpec;
use henad_core::metadata::Backend;
use crate::exec::{ActiveRun, ActiveRuns, Concurrency, SweepControl};
use crate::output::OutputError;
use crate::output::manifest::ManifestError;
use crate::progress::{Progress, ProgressEvent};
#[cfg(target_arch = "wasm32")]
use crate::pumped::{PumpedSweep, SweepCommand};
use crate::search_run::{SearchPlan, SearchPlanError, SearchUpdate};
use crate::spec_file::SpecFileError;
use crate::sweep::{Provenance, SpecSource, SweepEnd, SweepOutline, SweepRecord, SweepWarning};
#[cfg(not(target_arch = "wasm32"))]
pub const DEFAULT_SERIES_BUDGET: usize = 256 << 20;
#[cfg(target_arch = "wasm32")]
pub const DEFAULT_SERIES_BUDGET: usize = 64 << 20;
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum SweepOutput {
Memory,
#[cfg(not(target_arch = "wasm32"))]
Directory(std::path::PathBuf),
}
#[derive(Clone)]
#[non_exhaustive]
pub struct SweepRunOptions {
pub concurrency: Concurrency,
pub memory_budget: Option<u64>,
pub gpu_memory_budget: Option<u64>,
pub retry_failed: bool,
pub series_budget: usize,
pub spec_source: SpecSource,
pub provenance: Provenance,
pub wake: Option<WakeFn>,
}
impl std::fmt::Debug for SweepRunOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SweepRunOptions")
.field("concurrency", &self.concurrency)
.field("memory_budget", &self.memory_budget)
.field("gpu_memory_budget", &self.gpu_memory_budget)
.field("retry_failed", &self.retry_failed)
.field("series_budget", &self.series_budget)
.field("spec_source", &self.spec_source)
.field("provenance", &self.provenance)
.field("wake", &self.wake.is_some())
.finish_non_exhaustive()
}
}
impl SweepRunOptions {
pub fn new(provenance: Provenance) -> Self {
Self {
concurrency: Concurrency::Auto,
memory_budget: None,
gpu_memory_budget: None,
retry_failed: false,
series_budget: DEFAULT_SERIES_BUDGET,
spec_source: SpecSource::default(),
provenance,
wake: None,
}
}
}
#[derive(Debug, Clone)]
pub enum SweepEvent {
Planned(Box<SweepOutline>),
Warned(SweepWarning),
RunFinished {
outcome: Box<RunOutcome>,
series_dropped: bool,
},
SearchBatchTold(Arc<SearchUpdate>),
Finished(Box<SweepRecord>),
Failed(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SweepPhase {
Planning,
Running,
Paused,
Ended(SweepEnd),
Failed,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SweepProgress {
pub phase: SweepPhase,
pub runs_total: u64,
pub runs_skipped: u64,
pub runs_done: u64,
pub runs_waiting: u64,
pub runs_failed: u64,
pub elapsed: Duration,
pub remaining: Option<Duration>,
pub active_runs: Vec<ActiveRun>,
}
#[derive(Debug)]
pub enum SweepStartError {
ModelMismatch {
spec_model: String,
entry_model: String,
},
Plan(PlanError),
Search(SearchPlanError),
GpuNeedsNative,
Output(OutputError),
Manifest(ManifestError),
Spec(SpecFileError),
Shard,
Spawn(std::io::Error),
}
impl fmt::Display for SweepStartError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::ModelMismatch {
spec_model,
entry_model,
} => write!(f, "spec is for model '{spec_model}', expected '{entry_model}'"),
Self::Plan(_) => f.write_str("cannot plan the sweep"),
Self::Search(_) => f.write_str("cannot plan the search"),
Self::GpuNeedsNative => f.write_str("a GPU sweep needs a native build"),
Self::Output(_) => f.write_str("cannot write the results"),
Self::Manifest(_) => f.write_str("cannot read the manifest of the sweep to resume"),
Self::Spec(_) => f.write_str("cannot read the spec the manifest records"),
Self::Shard => f.write_str("the manifest records an invalid shard"),
Self::Spawn(_) => f.write_str("cannot start the sweep's thread"),
}
}
}
impl std::error::Error for SweepStartError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Plan(error) => Some(error),
Self::Search(error) => Some(error),
Self::Output(error) => Some(error),
Self::Manifest(error) => Some(error),
Self::Spec(error) => Some(error),
Self::Spawn(error) => Some(error),
Self::ModelMismatch { .. } | Self::GpuNeedsNative | Self::Shard => None,
}
}
}
#[derive(Debug)]
pub struct SweepRun {
plan: Arc<Plan>,
search_plan: Option<Arc<SearchPlan>>,
control: SweepControl,
active_runs: ActiveRuns,
events: Receiver<SweepEvent>,
shared_progress: SharedProgress,
clock: PauseClock,
ended: bool,
#[cfg(not(target_arch = "wasm32"))]
thread: Option<std::thread::JoinHandle<()>>,
#[cfg(target_arch = "wasm32")]
driver: henad_compute::runner::Driver<PumpedSweep>,
}
impl SweepRun {
pub fn start(
entry: ModelEntry,
gpu: Option<GpuContext>,
spec: SweepSpec,
output: SweepOutput,
options: SweepRunOptions,
) -> Result<Self, SweepStartError> {
let (plan, search_plan) = plan_for_start(&entry, &spec)?;
match &output {
SweepOutput::Memory => {}
#[cfg(not(target_arch = "wasm32"))]
SweepOutput::Directory(dir) => {
crate::output::OutputDir::check_free(dir).map_err(SweepStartError::Output)?;
}
}
#[cfg(target_arch = "wasm32")]
drop((gpu, output));
Self::launch(SweepLaunch {
entry,
spec,
plan,
search_plan,
options,
#[cfg(not(target_arch = "wasm32"))]
gpu,
#[cfg(not(target_arch = "wasm32"))]
output,
#[cfg(not(target_arch = "wasm32"))]
shard: henad_core::explore::plan::Shard::WHOLE,
#[cfg(not(target_arch = "wasm32"))]
resume: false,
})
}
#[cfg(not(target_arch = "wasm32"))]
pub fn resume_directory(
entry: ModelEntry,
gpu: Option<GpuContext>,
dir: &std::path::Path,
options: SweepRunOptions,
) -> Result<Self, SweepStartError> {
use crate::output::manifest::Manifest;
use crate::output::{MANIFEST_FILE, OutputDir};
use crate::spec_file::SpecFile;
OutputDir::check_unlocked(dir).map_err(SweepStartError::Output)?;
let manifest = Manifest::read(&dir.join(MANIFEST_FILE)).map_err(SweepStartError::Manifest)?;
let spec = SpecFile::from_json(&manifest.spec)
.and_then(SpecFile::into_spec)
.map_err(SweepStartError::Spec)?;
let (plan, search_plan) = plan_for_start(&entry, &spec)?;
let shard = manifest.shard.to_shard().ok_or(SweepStartError::Shard)?;
let options = SweepRunOptions {
spec_source: SpecSource::from(&manifest.spec_source),
..options
};
Self::launch(SweepLaunch {
entry,
spec,
plan,
search_plan,
options,
gpu,
output: SweepOutput::Directory(dir.to_owned()),
shard,
resume: true,
})
}
pub fn plan(&self) -> &Arc<Plan> {
&self.plan
}
pub fn search_plan(&self) -> Option<&Arc<SearchPlan>> {
self.search_plan.as_ref()
}
pub fn pause(&mut self) {
if self.ended || self.control.is_paused() || self.control.is_aborted() {
return;
}
self.clock.pause();
self.control.pause();
#[cfg(target_arch = "wasm32")]
self.driver.send(SweepCommand::Pause);
}
pub fn resume(&mut self) {
if !self.control.is_paused() {
return;
}
self.clock.resume();
self.control.resume();
#[cfg(target_arch = "wasm32")]
self.driver.send(SweepCommand::Resume);
}
pub fn abort(&mut self) {
self.clock.resume();
self.control.abort();
#[cfg(target_arch = "wasm32")]
self.driver.send(SweepCommand::Abort);
}
pub fn is_paused(&self) -> bool {
self.control.is_paused()
}
pub fn is_ended(&self) -> bool {
self.ended
}
#[cfg(target_arch = "wasm32")]
pub fn update(&mut self, dt: f64) {
self.driver.update(dt);
}
#[cfg(not(target_arch = "wasm32"))]
pub fn update(&mut self, _dt: f64) {}
pub fn try_recv(&mut self) -> Option<SweepEvent> {
let event = self.events.try_recv().ok()?;
if matches!(event, SweepEvent::Finished(_) | SweepEvent::Failed(_)) {
self.ended = true;
}
Some(event)
}
pub fn progress(&self) -> SweepProgress {
let progress_state = self.shared_progress.lock();
let phase = match progress_state.final_phase {
Some(phase) => phase,
None if self.control.is_paused() => SweepPhase::Paused,
None if progress_state.runs_total.is_some() => SweepPhase::Running,
None => SweepPhase::Planning,
};
let ended = progress_state.final_phase.is_some();
let runs_total = progress_state.runs_total.unwrap_or_else(|| {
self.search_plan
.as_ref()
.map_or_else(|| self.plan.run_count(), |search_plan| search_plan.run_count())
});
let runs_done = progress_state.runs_done;
let elapsed = self.clock.elapsed(progress_state.ended_at.unwrap_or_else(Instant::now));
let runs_waiting = if ended {
0
} else {
self.active_runs.waiting_count() as u64
};
let active_runs = self.active_runs.list();
let remaining = if ended || runs_done == 0 {
None
} else {
remaining_time(elapsed, runs_total, runs_done + runs_waiting, &active_runs)
};
SweepProgress {
phase,
runs_total,
runs_skipped: progress_state.runs_skipped,
runs_done,
runs_waiting,
runs_failed: progress_state.runs_failed,
elapsed,
remaining,
active_runs,
}
}
#[cfg(not(target_arch = "wasm32"))]
fn launch(launch: SweepLaunch) -> Result<Self, SweepStartError> {
use crate::output::manifest::ManifestRuntime;
use crate::search_run::{run_search_in_memory, run_search_into_directory};
use crate::sweep::{SweepInputs, SweepOptions, run_in_memory, run_into_directory, sweep_device};
use henad_compute::fault::catching;
let (mut channel, events, shared_progress) = SweepChannel::open(&launch.options);
let control = SweepControl::new();
let active_runs = ActiveRuns::new();
let plan = Arc::clone(&launch.plan);
let search_plan = launch.search_plan.clone();
let mut sweep_options = SweepOptions::new(launch.options.provenance.clone());
sweep_options.concurrency = launch.options.concurrency;
sweep_options.memory_budget = launch.options.memory_budget;
sweep_options.gpu_memory_budget = launch.options.gpu_memory_budget;
sweep_options.control = control.clone();
sweep_options.shard = launch.shard;
sweep_options.resume = launch.resume;
sweep_options.retry_failed = launch.options.retry_failed;
sweep_options.active_runs = Some(active_runs.clone());
sweep_options.spec_source = launch.options.spec_source.clone();
let thread = std::thread::Builder::new()
.name("henad-sweep".to_owned())
.spawn(move || {
let SweepLaunch {
entry,
gpu,
spec,
plan,
search_plan,
output,
..
} = launch;
let device = match catching(ACQUIRING_DEVICE, || sweep_device(&entry, gpu.as_ref())) {
Ok(Ok(device)) => device,
Ok(Err(error)) => return channel.fail(&error),
Err(fault) => return channel.fail(&fault),
};
let gpu = device.as_deref();
let runtime = ManifestRuntime::new(gpu.and_then(GpuContext::runtime_info));
let folder = match &output {
SweepOutput::Memory => None,
SweepOutput::Directory(dir) => Some(dir.as_path()),
};
let inputs = SweepInputs {
entry: &entry,
gpu,
runtime: &runtime,
spec: &spec,
source: &sweep_options.spec_source,
provenance: &sweep_options.provenance,
options: &sweep_options,
folder,
dry_run: false,
};
let ran = catching(RUNNING_SWEEP, || match (folder, search_plan) {
(None, None) => run_in_memory(&inputs, Some(plan), &mut channel),
(Some(dir), None) => run_into_directory(&inputs, Some(plan), dir, &mut channel),
(None, Some(search_plan)) => run_search_in_memory(&inputs, Some(search_plan), &mut channel),
(Some(dir), Some(search_plan)) => {
run_search_into_directory(&inputs, Some(search_plan), dir, &mut channel)
}
});
match ran {
Ok(Ok(record)) => channel.finish(record),
Ok(Err(error)) => channel.fail(&error),
Err(fault) => channel.fail(&fault),
}
})
.map_err(SweepStartError::Spawn)?;
Ok(Self {
plan,
search_plan,
control,
active_runs,
events,
shared_progress,
clock: PauseClock::new(),
ended: false,
thread: Some(thread),
})
}
#[cfg(target_arch = "wasm32")]
fn launch(launch: SweepLaunch) -> Result<Self, SweepStartError> {
let (channel, events, shared_progress) = SweepChannel::open(&launch.options);
let plan = Arc::clone(&launch.plan);
let search_plan = launch.search_plan.clone();
let sweep = PumpedSweep::new(
launch.entry,
launch.spec,
launch.plan,
launch.search_plan,
channel,
launch.options,
)?;
let (control, active_runs) = (sweep.control().clone(), sweep.active_runs().clone());
Ok(Self {
plan,
search_plan,
control,
active_runs,
events,
shared_progress,
clock: PauseClock::new(),
ended: false,
driver: henad_compute::runner::Driver::spawn(sweep, |_| {}),
})
}
}
impl Drop for SweepRun {
fn drop(&mut self) {
self.control.abort();
#[cfg(not(target_arch = "wasm32"))]
if let Some(thread) = self.thread.take() {
drop(thread.join());
}
}
}
#[cfg(not(target_arch = "wasm32"))]
const RUNNING_SWEEP: &str = "running the sweep";
#[cfg(not(target_arch = "wasm32"))]
const ACQUIRING_DEVICE: &str = "acquiring the sweep's GPU device";
struct SweepLaunch {
entry: ModelEntry,
spec: SweepSpec,
plan: Arc<Plan>,
search_plan: Option<Arc<SearchPlan>>,
options: SweepRunOptions,
#[cfg(not(target_arch = "wasm32"))]
gpu: Option<GpuContext>,
#[cfg(not(target_arch = "wasm32"))]
output: SweepOutput,
#[cfg(not(target_arch = "wasm32"))]
shard: henad_core::explore::plan::Shard,
#[cfg(not(target_arch = "wasm32"))]
resume: bool,
}
fn plan_for_start(
entry: &ModelEntry,
spec: &SweepSpec,
) -> Result<(Arc<Plan>, Option<Arc<SearchPlan>>), SweepStartError> {
if spec.model != entry.id() {
return Err(SweepStartError::ModelMismatch {
spec_model: spec.model.clone(),
entry_model: entry.id().to_owned(),
});
}
if cfg!(target_arch = "wasm32") && entry.metadata().backend == Backend::Gpu {
return Err(SweepStartError::GpuNeedsNative);
}
if spec.search.is_some() {
let search_plan = SearchPlan::new(spec, &entry.schema()).map_err(SweepStartError::Search)?;
return Ok((Arc::clone(search_plan.base()), Some(Arc::new(search_plan))));
}
let plan = spec.plan(&entry.schema()).map_err(SweepStartError::Plan)?;
Ok((Arc::new(plan), None))
}
#[derive(Debug, Clone, Copy)]
struct PauseClock {
start: Instant,
paused_since: Option<Instant>,
pause_total: Duration,
}
impl PauseClock {
fn new() -> Self {
Self {
start: Instant::now(),
paused_since: None,
pause_total: Duration::ZERO,
}
}
fn pause(&mut self) {
self.paused_since.get_or_insert_with(Instant::now);
}
fn resume(&mut self) {
if let Some(since) = self.paused_since.take() {
self.pause_total += since.elapsed();
}
}
fn elapsed(&self, now: Instant) -> Duration {
let current_pause = self
.paused_since
.map_or(Duration::ZERO, |since| now.saturating_duration_since(since));
now.saturating_duration_since(self.start)
.saturating_sub(self.pause_total + current_pause)
}
}
fn remaining_time(
elapsed: Duration,
runs_total: u64,
runs_finished: u64,
active_runs: &[ActiveRun],
) -> Option<Duration> {
let stepped: f64 = active_runs
.iter()
.filter(|active| active.end_tick > 0)
.map(|active| (active.tick as f64 / active.end_tick as f64).min(1.0))
.sum();
let done = runs_finished as f64 + stepped;
if done <= 0.0 {
return None;
}
let left = (runs_total as f64 - done).max(0.0);
Duration::try_from_secs_f64(elapsed.as_secs_f64() * left / done).ok()
}
#[derive(Debug, Default)]
struct ProgressState {
runs_total: Option<u64>,
runs_skipped: u64,
runs_done: u64,
runs_failed: u64,
final_phase: Option<SweepPhase>,
ended_at: Option<Instant>,
}
#[derive(Debug, Clone, Default)]
pub(crate) struct SharedProgress(Arc<Mutex<ProgressState>>);
impl SharedProgress {
fn lock(&self) -> std::sync::MutexGuard<'_, ProgressState> {
self.0.lock().unwrap_or_else(PoisonError::into_inner)
}
}
pub(crate) struct SweepChannel {
sender: Sender<SweepEvent>,
shared_progress: SharedProgress,
wake: Option<WakeFn>,
series_budget: usize,
series_bytes: usize,
series_budget_spent: bool,
}
impl SweepChannel {
pub(crate) fn open(options: &SweepRunOptions) -> (Self, Receiver<SweepEvent>, SharedProgress) {
let (sender, events) = mpsc::channel();
let shared_progress = SharedProgress::default();
let channel = Self {
sender,
shared_progress: shared_progress.clone(),
wake: options.wake.clone(),
series_budget: options.series_budget,
series_bytes: 0,
series_budget_spent: false,
};
(channel, events, shared_progress)
}
pub(crate) fn finish(&self, record: SweepRecord) {
{
let mut progress_state = self.shared_progress.lock();
progress_state.final_phase = Some(SweepPhase::Ended(record.report.end));
progress_state.ended_at = Some(Instant::now());
}
self.send(SweepEvent::Finished(Box::new(record)));
}
pub(crate) fn fail(&self, error: &dyn std::error::Error) {
{
let mut progress_state = self.shared_progress.lock();
progress_state.final_phase = Some(SweepPhase::Failed);
progress_state.ended_at = Some(Instant::now());
}
self.send(SweepEvent::Failed(describe(error)));
}
fn apply_series_budget(&mut self, outcome: &RunOutcome) -> (RunOutcome, bool) {
let series = &outcome.series;
let bytes = series.len() * (series.width() + 1) * size_of::<f64>();
self.series_budget_spent |= self.series_bytes + bytes > self.series_budget;
if !self.series_budget_spent {
self.series_bytes += bytes;
return (outcome.clone(), false);
}
let without_series = RunOutcome {
run: outcome.run,
run_key: outcome.run_key,
status: outcome.status,
stop_reason: outcome.stop_reason,
ticks: outcome.ticks,
population: outcome.population,
build_ms: outcome.build_ms,
wall_ms: outcome.wall_ms,
reducers: outcome.reducers.clone(),
series: SeriesBuffer::new(series.width()),
note: outcome.note.clone(),
};
(without_series, !series.is_empty())
}
fn send(&self, event: SweepEvent) {
drop(self.sender.send(event));
if let Some(wake) = &self.wake {
wake();
}
}
}
impl Progress for SweepChannel {
fn report(&mut self, event: &ProgressEvent<'_>) {
match event {
ProgressEvent::Planned(outline) => {
{
let mut progress_state = self.shared_progress.lock();
progress_state.runs_total = Some(outline.pending);
progress_state.runs_skipped = outline.skipped;
}
self.send(SweepEvent::Planned(Box::new((*outline).clone())));
}
ProgressEvent::Warned(warning) => self.send(SweepEvent::Warned((*warning).clone())),
ProgressEvent::RunCommitted(outcome) => {
{
let mut progress_state = self.shared_progress.lock();
progress_state.runs_done += 1;
if outcome.status.is_failure() {
progress_state.runs_failed += 1;
}
}
let (outcome, series_dropped) = self.apply_series_budget(outcome);
self.send(SweepEvent::RunFinished {
outcome: Box::new(outcome),
series_dropped,
});
}
ProgressEvent::SearchBatchTold(update) => {
self.send(SweepEvent::SearchBatchTold(Arc::new((*update).clone())));
}
ProgressEvent::Progressed(_) | ProgressEvent::Ended(_) => {}
}
}
}
fn describe(error: &dyn std::error::Error) -> String {
let mut text = error.to_string();
let mut source = error.source();
while let Some(cause) = source {
text.push_str(": ");
text.push_str(&cause.to_string());
source = cause.source();
}
text
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use henad_core::explore::outcome::PlannedRun;
use super::remaining_time;
use crate::exec::ActiveRun;
fn active(run_id: u64, tick: u64, end_tick: u64) -> ActiveRun {
ActiveRun {
run: PlannedRun {
run_id,
config_id: run_id,
rep: 0,
seed: 1,
},
tick,
end_tick,
}
}
#[test]
fn the_remaining_time_counts_the_runs_in_progress() {
let seconds = |remaining: Option<Duration>| remaining.expect("an estimate").as_secs_f64();
assert!((seconds(remaining_time(Duration::from_secs(8), 10, 4, &[])) - 12.0).abs() < 1e-9);
let nearly_done = [active(1, 90, 100), active(2, 90, 100), active(3, 90, 100)];
let estimate = seconds(remaining_time(Duration::from_secs(37), 8, 1, &nearly_done));
assert!(
(estimate - 43.0).abs() < 1e-6,
"1 run and 2.7 of the next three took 37 s, and 4.3 runs are left, not {estimate}"
);
assert_eq!(
remaining_time(Duration::from_secs(3), 10, 0, &[active(0, 0, 100)]),
None
);
assert_eq!(
remaining_time(Duration::from_secs(3), 10, 10, &[]),
Some(Duration::ZERO)
);
}
#[test]
fn a_remaining_time_past_the_largest_duration_reads_as_none() {
assert_eq!(remaining_time(Duration::from_secs(2), u64::MAX, 1, &[]), None);
}
}