Skip to main content

henad_explore/
handle.rs

1//! Handles on sweeps and searches that run beside a host's frame loop, with the same API on every target.
2//!
3//! On native a sweep runs on its own thread, and a pause holds its runs between two slices of steps. In a
4//! browser a pumped sweep steps one CPU run at a time from the host's frames. Either way the host reads the sweep's
5//! events from a channel that loses no events, and a progress record where the latest write wins.
6
7use std::fmt;
8use std::sync::mpsc::{self, Receiver, Sender};
9use std::sync::{Arc, Mutex, PoisonError};
10use std::time::Duration;
11
12use web_time::Instant;
13
14use henad_compute::cpu::sim_thread::WakeFn;
15use henad_compute::entry::ModelEntry;
16use henad_compute::gpu::GpuContext;
17use henad_core::explore::measure::SeriesBuffer;
18use henad_core::explore::outcome::RunOutcome;
19use henad_core::explore::plan::{Plan, PlanError};
20use henad_core::explore::spec::SweepSpec;
21use henad_core::metadata::Backend;
22
23use crate::exec::{ActiveRun, ActiveRuns, Concurrency, SweepControl};
24use crate::output::OutputError;
25use crate::output::manifest::ManifestError;
26use crate::progress::{Progress, ProgressEvent};
27#[cfg(target_arch = "wasm32")]
28use crate::pumped::{PumpedSweep, SweepCommand};
29use crate::search_run::{SearchPlan, SearchPlanError, SearchUpdate};
30use crate::spec_file::SpecFileError;
31use crate::sweep::{Provenance, SpecSource, SweepEnd, SweepOutline, SweepRecord, SweepWarning};
32
33/// Maximum total size in bytes of the series that the events of a sweep carry, unless the host requests a different
34/// budget.
35#[cfg(not(target_arch = "wasm32"))]
36pub const DEFAULT_SERIES_BUDGET: usize = 256 << 20;
37
38/// Maximum total size in bytes of the series that the events of a sweep carry, unless the host requests a different
39/// budget.
40#[cfg(target_arch = "wasm32")]
41pub const DEFAULT_SERIES_BUDGET: usize = 64 << 20;
42
43/// Place a sweep writes its four files.
44///
45/// The variants differ between targets, and a match outside this crate ends in a wildcard arm.
46#[derive(Debug, Clone, PartialEq, Eq)]
47#[non_exhaustive]
48pub enum SweepOutput {
49    /// Memory, returned in the [`SweepRecord`] once the sweep ends.
50    Memory,
51    /// A directory that holds no results, created when missing.
52    #[cfg(not(target_arch = "wasm32"))]
53    Directory(std::path::PathBuf),
54}
55
56/// Settings of a sweep that a [`SweepRun`] runs. None of them change its results.
57///
58/// [`Self::new`] returns the defaults, and a caller sets the other fields by assignment. A spec file's `[execution]`
59/// settings go in before the caller's settings, in the order
60/// [`SweepOptions::apply_execution`](crate::sweep::SweepOptions::apply_execution) sets them.
61#[derive(Clone)]
62#[non_exhaustive]
63pub struct SweepRunOptions {
64    /// Number of runs stepped at once on native. A browser steps one run at a time regardless of this setting.
65    pub concurrency: Concurrency,
66    /// Host memory budget in bytes for all live runs together, `None` for no limit.
67    pub memory_budget: Option<u64>,
68    /// GPU memory budget in bytes for all live runs of a GPU model together, `None` for the device's largest buffer.
69    pub gpu_memory_budget: Option<u64>,
70    /// Whether a resume reruns the runs that ended on a fault.
71    pub retry_failed: bool,
72    /// Maximum total size in bytes of the series that the [`SweepEvent::RunFinished`] events carry.
73    ///
74    /// Once a run's series would pass it, that run and every later one arrive without their series. The files keep
75    /// every series.
76    pub series_budget: usize,
77    /// Spec file the sweep was read from, for the manifest.
78    pub spec_source: SpecSource,
79    /// Build of the host, for the manifest.
80    pub provenance: Provenance,
81    /// Called after each event, so that an idle host wakes up to collect it. It must not block.
82    pub wake: Option<WakeFn>,
83}
84
85/// Prints whether a wake callback is set, instead of the callback.
86impl std::fmt::Debug for SweepRunOptions {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        f.debug_struct("SweepRunOptions")
89            .field("concurrency", &self.concurrency)
90            .field("memory_budget", &self.memory_budget)
91            .field("gpu_memory_budget", &self.gpu_memory_budget)
92            .field("retry_failed", &self.retry_failed)
93            .field("series_budget", &self.series_budget)
94            .field("spec_source", &self.spec_source)
95            .field("provenance", &self.provenance)
96            .field("wake", &self.wake.is_some())
97            .finish_non_exhaustive()
98    }
99}
100
101impl SweepRunOptions {
102    /// Returns the options at their defaults, recording `provenance` in every manifest the sweep writes.
103    pub fn new(provenance: Provenance) -> Self {
104        Self {
105            concurrency: Concurrency::Auto,
106            memory_budget: None,
107            gpu_memory_budget: None,
108            retry_failed: false,
109            series_budget: DEFAULT_SERIES_BUDGET,
110            spec_source: SpecSource::default(),
111            provenance,
112            wake: None,
113        }
114    }
115}
116
117/// Event of a sweep, in the order it happened.
118#[derive(Debug, Clone)]
119pub enum SweepEvent {
120    /// The sweep is planned and probed, and its runs start.
121    Planned(Box<SweepOutline>),
122    /// A warning. The sweep still runs, but likely not as intended.
123    Warned(SweepWarning),
124    /// A run was written. Runs arrive in plan order, and a search's runs in the order that it requests them.
125    RunFinished {
126        /// Outcome of the run as written, with an empty series once the series budget is spent.
127        outcome: Box<RunOutcome>,
128        /// Whether the run's series passed the budget of [`SweepRunOptions::series_budget`] and was left out of
129        /// `outcome`.
130        series_dropped: bool,
131    },
132    /// A search was told the evaluations of one batch, after the runs of the batch.
133    SearchBatchTold(Arc<SearchUpdate>),
134    /// The sweep ran to its end or was aborted, as its record says. Nothing follows.
135    Finished(Box<SweepRecord>),
136    /// The sweep ended on its own error, outside any run, with the error and its causes. Nothing follows.
137    Failed(String),
138}
139
140/// Stage a sweep has reached.
141#[derive(Debug, Clone, Copy, PartialEq, Eq)]
142pub enum SweepPhase {
143    /// Planning and probing, before the first run starts.
144    Planning,
145    /// Stepping its runs.
146    Running,
147    /// Held between two slices of steps. Runs in progress keep their state.
148    Paused,
149    /// Ran to its end or was aborted.
150    Ended(SweepEnd),
151    /// Ended on its own error, outside any run.
152    Failed,
153}
154
155/// Progress of a sweep at the moment it was read.
156#[derive(Debug, Clone, PartialEq, Eq)]
157pub struct SweepProgress {
158    /// Stage the sweep has reached.
159    pub phase: SweepPhase,
160    /// Number of runs this sweep executes. Until the sweep is planned, the number of runs in the whole plan, or in a
161    /// search's whole budget.
162    pub runs_total: u64,
163    /// Number of runs a resumed directory held already. The sweep keeps them.
164    pub runs_skipped: u64,
165    /// Number of runs written so far.
166    pub runs_done: u64,
167    /// Number of runs that finished and wait for an earlier run to be written. An abort drops them.
168    pub runs_waiting: u64,
169    /// Number of written runs that ended on a fault or a timeout.
170    pub runs_failed: u64,
171    /// Time since the sweep started, pauses left out.
172    pub elapsed: Duration,
173    /// Time left at the pace so far, `None` before the first run is written and once the sweep ends.
174    ///
175    /// The pace counts the runs that finished and the share of its ticks each run in progress has stepped. A time too
176    /// long for a [`Duration`], such as the projection for a search with a budget of `u64::MAX` evaluations, is
177    /// `None`.
178    pub remaining: Option<Duration>,
179    /// Runs in progress, in order of their ids.
180    pub active_runs: Vec<ActiveRun>,
181}
182
183/// A sweep that cannot start.
184#[derive(Debug)]
185pub enum SweepStartError {
186    /// A spec written for the model `spec_model`, given the entry for `entry_model`.
187    ModelMismatch {
188        /// Model id in the spec.
189        spec_model: String,
190        /// Id of the model entry passed to the sweep.
191        entry_model: String,
192    },
193    /// A spec that the model rejects.
194    Plan(PlanError),
195    /// A search spec that the model rejects.
196    Search(SearchPlanError),
197    /// A GPU model on a target that cannot step a GPU model in a sweep.
198    GpuNeedsNative,
199    /// An output directory that holds the results of a sweep, or that another sweep, search or merge is writing to.
200    Output(OutputError),
201    /// A manifest that cannot be read.
202    Manifest(ManifestError),
203    /// A manifest whose spec cannot be read back.
204    Spec(SpecFileError),
205    /// A directory whose manifest specifies a shard outside its plan.
206    Shard,
207    /// Starting the sweep's thread failed.
208    Spawn(std::io::Error),
209}
210
211impl fmt::Display for SweepStartError {
212    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
213        match self {
214            Self::ModelMismatch {
215                spec_model,
216                entry_model,
217            } => write!(f, "spec is for model '{spec_model}', expected '{entry_model}'"),
218            Self::Plan(_) => f.write_str("cannot plan the sweep"),
219            Self::Search(_) => f.write_str("cannot plan the search"),
220            Self::GpuNeedsNative => f.write_str("a GPU sweep needs a native build"),
221            Self::Output(_) => f.write_str("cannot write the results"),
222            Self::Manifest(_) => f.write_str("cannot read the manifest of the sweep to resume"),
223            Self::Spec(_) => f.write_str("cannot read the spec the manifest records"),
224            Self::Shard => f.write_str("the manifest records an invalid shard"),
225            Self::Spawn(_) => f.write_str("cannot start the sweep's thread"),
226        }
227    }
228}
229
230impl std::error::Error for SweepStartError {
231    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
232        match self {
233            Self::Plan(error) => Some(error),
234            Self::Search(error) => Some(error),
235            Self::Output(error) => Some(error),
236            Self::Manifest(error) => Some(error),
237            Self::Spec(error) => Some(error),
238            Self::Spawn(error) => Some(error),
239            Self::ModelMismatch { .. } | Self::GpuNeedsNative | Self::Shard => None,
240        }
241    }
242}
243
244/// Handle on a running sweep or search.
245///
246/// Dropping the handle aborts the sweep. On native it then waits for the sweep's thread to write its files.
247#[derive(Debug)]
248pub struct SweepRun {
249    plan: Arc<Plan>,
250    /// Plan of a search, `None` for a sweep.
251    search_plan: Option<Arc<SearchPlan>>,
252    control: SweepControl,
253    active_runs: ActiveRuns,
254    events: Receiver<SweepEvent>,
255    shared_progress: SharedProgress,
256    clock: PauseClock,
257    /// Whether the last event, [`SweepEvent::Finished`] or [`SweepEvent::Failed`], has been received.
258    ended: bool,
259    #[cfg(not(target_arch = "wasm32"))]
260    thread: Option<std::thread::JoinHandle<()>>,
261    #[cfg(target_arch = "wasm32")]
262    driver: henad_compute::runner::Driver<PumpedSweep>,
263}
264
265impl SweepRun {
266    /// Plans `spec` against `entry` and starts the sweep, or the search when `spec` has a `[search]` table, writing
267    /// its files to `output`.
268    ///
269    /// `gpu` is a device the host shares with the sweep, its [`FaultSink`](henad_compute::fault::FaultSink) included.
270    /// Note that the sink holds one fault, and whichever side reads it first takes it. A fault that the sweep takes
271    /// first ends every live run, whichever side raised it, and a fault that the host takes first leaves the runs
272    /// going. When `gpu` is `None`, the sweep acquires a device on its own thread for a GPU model, sized to the entry's
273    /// [`gpu_needs`](ModelEntry::gpu_needs), and builds `entry` on it. The manifest records the adapter of a shared
274    /// device only when its context carries [`RuntimeInfo`](henad_compute::runtime_info::RuntimeInfo), as
275    /// [`GpuContext::with_runtime_info`] attaches it.
276    ///
277    /// Planning happens before this returns. The probe build and the runs happen after it, on native on the
278    /// sweep's own thread and in a browser in [`Self::update`].
279    ///
280    /// # Errors
281    ///
282    /// Returns [`SweepStartError`] when `spec` specifies a different model, cannot be planned, or targets a GPU
283    /// model in a browser, when `output` is a directory that holds results, or when the sweep's thread cannot start.
284    /// A device the sweep cannot acquire fails the sweep with a [`SweepEvent::Failed`].
285    pub fn start(
286        entry: ModelEntry,
287        gpu: Option<GpuContext>,
288        spec: SweepSpec,
289        output: SweepOutput,
290        options: SweepRunOptions,
291    ) -> Result<Self, SweepStartError> {
292        let (plan, search_plan) = plan_for_start(&entry, &spec)?;
293        match &output {
294            SweepOutput::Memory => {}
295            #[cfg(not(target_arch = "wasm32"))]
296            SweepOutput::Directory(dir) => {
297                crate::output::OutputDir::check_free(dir).map_err(SweepStartError::Output)?;
298            }
299        }
300        // A browser steps CPU models only.
301        #[cfg(target_arch = "wasm32")]
302        drop((gpu, output));
303        Self::launch(SweepLaunch {
304            entry,
305            spec,
306            plan,
307            search_plan,
308            options,
309            #[cfg(not(target_arch = "wasm32"))]
310            gpu,
311            #[cfg(not(target_arch = "wasm32"))]
312            output,
313            #[cfg(not(target_arch = "wasm32"))]
314            shard: henad_core::explore::plan::Shard::WHOLE,
315            #[cfg(not(target_arch = "wasm32"))]
316            resume: false,
317        })
318    }
319
320    /// Resumes the sweep whose results the directory `dir` holds, running only the runs it lacks.
321    ///
322    /// The spec, source and shard come from the directory's manifest, and the sweep runs as
323    /// [`crate::sweep::run_spec`] resumes a sweep. Note that the execution settings come from `options` alone. The
324    /// `[execution]` table the manifest records is not applied. `gpu` is the device a GPU model runs on, as in
325    /// [`Self::start`].
326    ///
327    /// # Errors
328    ///
329    /// Returns [`SweepStartError`] when another sweep, search or merge is writing to `dir`, the manifest or its spec
330    /// cannot be read, the spec specifies a different model or cannot be planned, or the sweep's thread cannot start. A
331    /// directory whose runs do not fit the plan, or a device the sweep cannot acquire, fails the sweep with a
332    /// [`SweepEvent::Failed`].
333    #[cfg(not(target_arch = "wasm32"))]
334    pub fn resume_directory(
335        entry: ModelEntry,
336        gpu: Option<GpuContext>,
337        dir: &std::path::Path,
338        options: SweepRunOptions,
339    ) -> Result<Self, SweepStartError> {
340        use crate::output::manifest::Manifest;
341        use crate::output::{MANIFEST_FILE, OutputDir};
342        use crate::spec_file::SpecFile;
343
344        OutputDir::check_unlocked(dir).map_err(SweepStartError::Output)?;
345        let manifest = Manifest::read(&dir.join(MANIFEST_FILE)).map_err(SweepStartError::Manifest)?;
346        let spec = SpecFile::from_json(&manifest.spec)
347            .and_then(SpecFile::into_spec)
348            .map_err(SweepStartError::Spec)?;
349        let (plan, search_plan) = plan_for_start(&entry, &spec)?;
350        let shard = manifest.shard.to_shard().ok_or(SweepStartError::Shard)?;
351        let options = SweepRunOptions {
352            spec_source: SpecSource::from(&manifest.spec_source),
353            ..options
354        };
355        Self::launch(SweepLaunch {
356            entry,
357            spec,
358            plan,
359            search_plan,
360            options,
361            gpu,
362            output: SweepOutput::Directory(dir.to_owned()),
363            shard,
364            resume: true,
365        })
366    }
367
368    /// Plan of the sweep's configs and runs, or of a search's fixed values and actions alone.
369    pub fn plan(&self) -> &Arc<Plan> {
370        &self.plan
371    }
372
373    /// Plan of a search, `None` for a sweep.
374    ///
375    /// A search's configs are chosen as it runs. Each config arrives in a [`SweepEvent::SearchBatchTold`], and
376    /// [`SearchPlan::replay`] rebuilds any of its runs.
377    pub fn search_plan(&self) -> Option<&Arc<SearchPlan>> {
378        self.search_plan.as_ref()
379    }
380
381    /// Holds every run at its next slice of steps.
382    pub fn pause(&mut self) {
383        if self.ended || self.control.is_paused() || self.control.is_aborted() {
384            return;
385        }
386        self.clock.pause();
387        self.control.pause();
388        #[cfg(target_arch = "wasm32")]
389        self.driver.send(SweepCommand::Pause);
390    }
391
392    /// Lets a paused sweep continue.
393    pub fn resume(&mut self) {
394        if !self.control.is_paused() {
395            return;
396        }
397        self.clock.resume();
398        self.control.resume();
399        #[cfg(target_arch = "wasm32")]
400        self.driver.send(SweepCommand::Resume);
401    }
402
403    /// Ends the sweep at its next slice of steps, paused runs included. The runs written so far stay in its files.
404    pub fn abort(&mut self) {
405        self.clock.resume();
406        self.control.abort();
407        #[cfg(target_arch = "wasm32")]
408        self.driver.send(SweepCommand::Abort);
409    }
410
411    /// Returns whether the sweep is paused.
412    pub fn is_paused(&self) -> bool {
413        self.control.is_paused()
414    }
415
416    /// Returns whether the last event, [`SweepEvent::Finished`] or [`SweepEvent::Failed`], has been received.
417    pub fn is_ended(&self) -> bool {
418        self.ended
419    }
420
421    /// Steps the sweep within a frame's budget in a browser. On native the sweep runs on its own thread.
422    #[cfg(target_arch = "wasm32")]
423    pub fn update(&mut self, dt: f64) {
424        self.driver.update(dt);
425    }
426
427    /// Steps the sweep within a frame's budget in a browser. On native the sweep runs on its own thread.
428    #[cfg(not(target_arch = "wasm32"))]
429    pub fn update(&mut self, _dt: f64) {}
430
431    /// Returns the next event, or `None` when no event is waiting.
432    pub fn try_recv(&mut self) -> Option<SweepEvent> {
433        let event = self.events.try_recv().ok()?;
434        if matches!(event, SweepEvent::Finished(_) | SweepEvent::Failed(_)) {
435            self.ended = true;
436        }
437        Some(event)
438    }
439
440    /// Returns the progress of the sweep as it stands.
441    pub fn progress(&self) -> SweepProgress {
442        let progress_state = self.shared_progress.lock();
443        let phase = match progress_state.final_phase {
444            Some(phase) => phase,
445            None if self.control.is_paused() => SweepPhase::Paused,
446            None if progress_state.runs_total.is_some() => SweepPhase::Running,
447            None => SweepPhase::Planning,
448        };
449        let ended = progress_state.final_phase.is_some();
450        let runs_total = progress_state.runs_total.unwrap_or_else(|| {
451            self.search_plan
452                .as_ref()
453                .map_or_else(|| self.plan.run_count(), |search_plan| search_plan.run_count())
454        });
455        let runs_done = progress_state.runs_done;
456        let elapsed = self.clock.elapsed(progress_state.ended_at.unwrap_or_else(Instant::now));
457        let runs_waiting = if ended {
458            0
459        } else {
460            self.active_runs.waiting_count() as u64
461        };
462        let active_runs = self.active_runs.list();
463        let remaining = if ended || runs_done == 0 {
464            None
465        } else {
466            remaining_time(elapsed, runs_total, runs_done + runs_waiting, &active_runs)
467        };
468        SweepProgress {
469            phase,
470            runs_total,
471            runs_skipped: progress_state.runs_skipped,
472            runs_done,
473            runs_waiting,
474            runs_failed: progress_state.runs_failed,
475            elapsed,
476            remaining,
477            active_runs,
478        }
479    }
480
481    #[cfg(not(target_arch = "wasm32"))]
482    fn launch(launch: SweepLaunch) -> Result<Self, SweepStartError> {
483        use crate::output::manifest::ManifestRuntime;
484        use crate::search_run::{run_search_in_memory, run_search_into_directory};
485        use crate::sweep::{SweepInputs, SweepOptions, run_in_memory, run_into_directory, sweep_device};
486        use henad_compute::fault::catching;
487
488        let (mut channel, events, shared_progress) = SweepChannel::open(&launch.options);
489        let control = SweepControl::new();
490        let active_runs = ActiveRuns::new();
491        let plan = Arc::clone(&launch.plan);
492        let search_plan = launch.search_plan.clone();
493        let mut sweep_options = SweepOptions::new(launch.options.provenance.clone());
494        sweep_options.concurrency = launch.options.concurrency;
495        sweep_options.memory_budget = launch.options.memory_budget;
496        sweep_options.gpu_memory_budget = launch.options.gpu_memory_budget;
497        sweep_options.control = control.clone();
498        sweep_options.shard = launch.shard;
499        sweep_options.resume = launch.resume;
500        sweep_options.retry_failed = launch.options.retry_failed;
501        sweep_options.active_runs = Some(active_runs.clone());
502        sweep_options.spec_source = launch.options.spec_source.clone();
503        let thread = std::thread::Builder::new()
504            .name("henad-sweep".to_owned())
505            .spawn(move || {
506                let SweepLaunch {
507                    entry,
508                    gpu,
509                    spec,
510                    plan,
511                    search_plan,
512                    output,
513                    ..
514                } = launch;
515                let device = match catching(ACQUIRING_DEVICE, || sweep_device(&entry, gpu.as_ref())) {
516                    Ok(Ok(device)) => device,
517                    Ok(Err(error)) => return channel.fail(&error),
518                    Err(fault) => return channel.fail(&fault),
519                };
520                let gpu = device.as_deref();
521                let runtime = ManifestRuntime::new(gpu.and_then(GpuContext::runtime_info));
522                let folder = match &output {
523                    SweepOutput::Memory => None,
524                    SweepOutput::Directory(dir) => Some(dir.as_path()),
525                };
526                let inputs = SweepInputs {
527                    entry: &entry,
528                    gpu,
529                    runtime: &runtime,
530                    spec: &spec,
531                    source: &sweep_options.spec_source,
532                    provenance: &sweep_options.provenance,
533                    options: &sweep_options,
534                    folder,
535                    dry_run: false,
536                };
537                let ran = catching(RUNNING_SWEEP, || match (folder, search_plan) {
538                    (None, None) => run_in_memory(&inputs, Some(plan), &mut channel),
539                    (Some(dir), None) => run_into_directory(&inputs, Some(plan), dir, &mut channel),
540                    (None, Some(search_plan)) => run_search_in_memory(&inputs, Some(search_plan), &mut channel),
541                    (Some(dir), Some(search_plan)) => {
542                        run_search_into_directory(&inputs, Some(search_plan), dir, &mut channel)
543                    }
544                });
545                match ran {
546                    Ok(Ok(record)) => channel.finish(record),
547                    Ok(Err(error)) => channel.fail(&error),
548                    Err(fault) => channel.fail(&fault),
549                }
550            })
551            .map_err(SweepStartError::Spawn)?;
552        Ok(Self {
553            plan,
554            search_plan,
555            control,
556            active_runs,
557            events,
558            shared_progress,
559            clock: PauseClock::new(),
560            ended: false,
561            thread: Some(thread),
562        })
563    }
564
565    #[cfg(target_arch = "wasm32")]
566    fn launch(launch: SweepLaunch) -> Result<Self, SweepStartError> {
567        let (channel, events, shared_progress) = SweepChannel::open(&launch.options);
568        let plan = Arc::clone(&launch.plan);
569        let search_plan = launch.search_plan.clone();
570        let sweep = PumpedSweep::new(
571            launch.entry,
572            launch.spec,
573            launch.plan,
574            launch.search_plan,
575            channel,
576            launch.options,
577        )?;
578        let (control, active_runs) = (sweep.control().clone(), sweep.active_runs().clone());
579        Ok(Self {
580            plan,
581            search_plan,
582            control,
583            active_runs,
584            events,
585            shared_progress,
586            clock: PauseClock::new(),
587            ended: false,
588            driver: henad_compute::runner::Driver::spawn(sweep, |_| {}),
589        })
590    }
591}
592
593impl Drop for SweepRun {
594    fn drop(&mut self) {
595        self.control.abort();
596        #[cfg(not(target_arch = "wasm32"))]
597        if let Some(thread) = self.thread.take() {
598            drop(thread.join());
599        }
600    }
601}
602
603/// Task that a panic outside any run reports as its `during`.
604#[cfg(not(target_arch = "wasm32"))]
605const RUNNING_SWEEP: &str = "running the sweep";
606
607/// Task that a panic reports as its `during` while the sweep acquires its device.
608#[cfg(not(target_arch = "wasm32"))]
609const ACQUIRING_DEVICE: &str = "acquiring the sweep's GPU device";
610
611/// Sweep to launch, planned.
612///
613/// A browser runs the whole plan of a CPU model into memory, so the fields that choose a device, an output directory,
614/// a shard or a resume are native only.
615struct SweepLaunch {
616    entry: ModelEntry,
617    spec: SweepSpec,
618    plan: Arc<Plan>,
619    search_plan: Option<Arc<SearchPlan>>,
620    options: SweepRunOptions,
621    #[cfg(not(target_arch = "wasm32"))]
622    gpu: Option<GpuContext>,
623    #[cfg(not(target_arch = "wasm32"))]
624    output: SweepOutput,
625    #[cfg(not(target_arch = "wasm32"))]
626    shard: henad_core::explore::plan::Shard,
627    #[cfg(not(target_arch = "wasm32"))]
628    resume: bool,
629}
630
631/// Returns the plan of `spec` over `entry` that a handle starts, and for a search its search plan. For a search, the
632/// first plan is the search plan's base.
633///
634/// # Errors
635///
636/// Returns [`SweepStartError`] when the spec specifies a different model, targets a GPU model in a browser,
637/// or cannot be planned.
638fn plan_for_start(
639    entry: &ModelEntry,
640    spec: &SweepSpec,
641) -> Result<(Arc<Plan>, Option<Arc<SearchPlan>>), SweepStartError> {
642    if spec.model != entry.id() {
643        return Err(SweepStartError::ModelMismatch {
644            spec_model: spec.model.clone(),
645            entry_model: entry.id().to_owned(),
646        });
647    }
648    if cfg!(target_arch = "wasm32") && entry.metadata().backend == Backend::Gpu {
649        return Err(SweepStartError::GpuNeedsNative);
650    }
651    if spec.search.is_some() {
652        let search_plan = SearchPlan::new(spec, &entry.schema()).map_err(SweepStartError::Search)?;
653        return Ok((Arc::clone(search_plan.base()), Some(Arc::new(search_plan))));
654    }
655    let plan = spec.plan(&entry.schema()).map_err(SweepStartError::Plan)?;
656    Ok((Arc::new(plan), None))
657}
658
659/// Time a sweep has run, pauses left out.
660#[derive(Debug, Clone, Copy)]
661struct PauseClock {
662    start: Instant,
663    /// Start of the pause in progress.
664    paused_since: Option<Instant>,
665    /// Length of the pauses that have ended.
666    pause_total: Duration,
667}
668
669impl PauseClock {
670    fn new() -> Self {
671        Self {
672            start: Instant::now(),
673            paused_since: None,
674            pause_total: Duration::ZERO,
675        }
676    }
677
678    fn pause(&mut self) {
679        self.paused_since.get_or_insert_with(Instant::now);
680    }
681
682    fn resume(&mut self) {
683        if let Some(since) = self.paused_since.take() {
684            self.pause_total += since.elapsed();
685        }
686    }
687
688    /// Returns the time run up to `now`.
689    fn elapsed(&self, now: Instant) -> Duration {
690        let current_pause = self
691            .paused_since
692            .map_or(Duration::ZERO, |since| now.saturating_duration_since(since));
693        now.saturating_duration_since(self.start)
694            .saturating_sub(self.pause_total + current_pause)
695    }
696}
697
698/// Returns the time left for `runs_total` runs after `elapsed`, at the pace of the `runs_finished` runs that finished
699/// and the share of its ticks each run of `active_runs` has stepped.
700///
701/// Returns `None` before any run has made progress, and for a time too long for a [`Duration`].
702fn remaining_time(
703    elapsed: Duration,
704    runs_total: u64,
705    runs_finished: u64,
706    active_runs: &[ActiveRun],
707) -> Option<Duration> {
708    let stepped: f64 = active_runs
709        .iter()
710        .filter(|active| active.end_tick > 0)
711        .map(|active| (active.tick as f64 / active.end_tick as f64).min(1.0))
712        .sum();
713    let done = runs_finished as f64 + stepped;
714    if done <= 0.0 {
715        return None;
716    }
717    let left = (runs_total as f64 - done).max(0.0);
718    Duration::try_from_secs_f64(elapsed.as_secs_f64() * left / done).ok()
719}
720
721/// Counts and stage of a sweep, written by its channel and read by its handle.
722#[derive(Debug, Default)]
723struct ProgressState {
724    /// Number of runs this sweep executes, set once the sweep is planned.
725    runs_total: Option<u64>,
726    /// Number of runs a resumed directory held already, set once the sweep is planned.
727    runs_skipped: u64,
728    runs_done: u64,
729    runs_failed: u64,
730    /// Stage the sweep ended in, `None` while it runs.
731    final_phase: Option<SweepPhase>,
732    ended_at: Option<Instant>,
733}
734
735/// Progress state shared by a sweep's channel and its handle.
736#[derive(Debug, Clone, Default)]
737pub(crate) struct SharedProgress(Arc<Mutex<ProgressState>>);
738
739impl SharedProgress {
740    fn lock(&self) -> std::sync::MutexGuard<'_, ProgressState> {
741        self.0.lock().unwrap_or_else(PoisonError::into_inner)
742    }
743}
744
745/// Sender of a sweep's events to its handle, and writer of the progress the handle reads.
746pub(crate) struct SweepChannel {
747    sender: Sender<SweepEvent>,
748    shared_progress: SharedProgress,
749    wake: Option<WakeFn>,
750    series_budget: usize,
751    /// Size in bytes of the series the events have carried.
752    series_bytes: usize,
753    /// Whether a series has passed the budget. Every later run arrives without its series.
754    series_budget_spent: bool,
755}
756
757impl SweepChannel {
758    /// Returns a channel with the wake and series budget of `options`, the end its events arrive at, and the progress
759    /// state it writes.
760    pub(crate) fn open(options: &SweepRunOptions) -> (Self, Receiver<SweepEvent>, SharedProgress) {
761        let (sender, events) = mpsc::channel();
762        let shared_progress = SharedProgress::default();
763        let channel = Self {
764            sender,
765            shared_progress: shared_progress.clone(),
766            wake: options.wake.clone(),
767            series_budget: options.series_budget,
768            series_bytes: 0,
769            series_budget_spent: false,
770        };
771        (channel, events, shared_progress)
772    }
773
774    /// Sends the record of a sweep that ran to its end or was aborted.
775    pub(crate) fn finish(&self, record: SweepRecord) {
776        {
777            let mut progress_state = self.shared_progress.lock();
778            progress_state.final_phase = Some(SweepPhase::Ended(record.report.end));
779            progress_state.ended_at = Some(Instant::now());
780        }
781        self.send(SweepEvent::Finished(Box::new(record)));
782    }
783
784    /// Sends `error`, with its causes, as the end of a sweep that failed.
785    pub(crate) fn fail(&self, error: &dyn std::error::Error) {
786        {
787            let mut progress_state = self.shared_progress.lock();
788            progress_state.final_phase = Some(SweepPhase::Failed);
789            progress_state.ended_at = Some(Instant::now());
790        }
791        self.send(SweepEvent::Failed(describe(error)));
792    }
793
794    /// Returns `outcome` with its series left out once the series have passed the budget, and whether the series was
795    /// left out.
796    fn apply_series_budget(&mut self, outcome: &RunOutcome) -> (RunOutcome, bool) {
797        let series = &outcome.series;
798        let bytes = series.len() * (series.width() + 1) * size_of::<f64>();
799        self.series_budget_spent |= self.series_bytes + bytes > self.series_budget;
800        if !self.series_budget_spent {
801            self.series_bytes += bytes;
802            return (outcome.clone(), false);
803        }
804        let without_series = RunOutcome {
805            run: outcome.run,
806            run_key: outcome.run_key,
807            status: outcome.status,
808            stop_reason: outcome.stop_reason,
809            ticks: outcome.ticks,
810            population: outcome.population,
811            build_ms: outcome.build_ms,
812            wall_ms: outcome.wall_ms,
813            reducers: outcome.reducers.clone(),
814            series: SeriesBuffer::new(series.width()),
815            note: outcome.note.clone(),
816        };
817        (without_series, !series.is_empty())
818    }
819
820    fn send(&self, event: SweepEvent) {
821        // A dropped handle receives no more events, and the sweep runs on to write its files.
822        drop(self.sender.send(event));
823        if let Some(wake) = &self.wake {
824            wake();
825        }
826    }
827}
828
829impl Progress for SweepChannel {
830    fn report(&mut self, event: &ProgressEvent<'_>) {
831        match event {
832            ProgressEvent::Planned(outline) => {
833                {
834                    let mut progress_state = self.shared_progress.lock();
835                    progress_state.runs_total = Some(outline.pending);
836                    progress_state.runs_skipped = outline.skipped;
837                }
838                self.send(SweepEvent::Planned(Box::new((*outline).clone())));
839            }
840            ProgressEvent::Warned(warning) => self.send(SweepEvent::Warned((*warning).clone())),
841            ProgressEvent::RunCommitted(outcome) => {
842                {
843                    let mut progress_state = self.shared_progress.lock();
844                    progress_state.runs_done += 1;
845                    if outcome.status.is_failure() {
846                        progress_state.runs_failed += 1;
847                    }
848                }
849                let (outcome, series_dropped) = self.apply_series_budget(outcome);
850                self.send(SweepEvent::RunFinished {
851                    outcome: Box::new(outcome),
852                    series_dropped,
853                });
854            }
855            ProgressEvent::SearchBatchTold(update) => {
856                self.send(SweepEvent::SearchBatchTold(Arc::new((*update).clone())));
857            }
858            ProgressEvent::Progressed(_) | ProgressEvent::Ended(_) => {}
859        }
860    }
861}
862
863/// Returns `error` followed by each of its causes, joined by colons.
864fn describe(error: &dyn std::error::Error) -> String {
865    let mut text = error.to_string();
866    let mut source = error.source();
867    while let Some(cause) = source {
868        text.push_str(": ");
869        text.push_str(&cause.to_string());
870        source = cause.source();
871    }
872    text
873}
874
875#[cfg(test)]
876mod tests {
877    use std::time::Duration;
878
879    use henad_core::explore::outcome::PlannedRun;
880
881    use super::remaining_time;
882    use crate::exec::ActiveRun;
883
884    /// Returns run `run_id` in progress at `tick` of `end_tick`.
885    fn active(run_id: u64, tick: u64, end_tick: u64) -> ActiveRun {
886        ActiveRun {
887            run: PlannedRun {
888                run_id,
889                config_id: run_id,
890                rep: 0,
891                seed: 1,
892            },
893            tick,
894            end_tick,
895        }
896    }
897
898    #[test]
899    fn the_remaining_time_counts_the_runs_in_progress() {
900        let seconds = |remaining: Option<Duration>| remaining.expect("an estimate").as_secs_f64();
901        assert!((seconds(remaining_time(Duration::from_secs(8), 10, 4, &[])) - 12.0).abs() < 1e-9);
902        // Four runs started together, and one finished as the other three neared their end.
903        let nearly_done = [active(1, 90, 100), active(2, 90, 100), active(3, 90, 100)];
904        let estimate = seconds(remaining_time(Duration::from_secs(37), 8, 1, &nearly_done));
905        assert!(
906            (estimate - 43.0).abs() < 1e-6,
907            "1 run and 2.7 of the next three took 37 s, and 4.3 runs are left, not {estimate}"
908        );
909        assert_eq!(
910            remaining_time(Duration::from_secs(3), 10, 0, &[active(0, 0, 100)]),
911            None
912        );
913        assert_eq!(
914            remaining_time(Duration::from_secs(3), 10, 10, &[]),
915            Some(Duration::ZERO)
916        );
917    }
918
919    #[test]
920    fn a_remaining_time_past_the_largest_duration_reads_as_none() {
921        assert_eq!(remaining_time(Duration::from_secs(2), u64::MAX, 1, &[]), None);
922    }
923}