use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use web_time::Instant;
use henad_compute::entry::ModelState;
use henad_compute::fault::{BUILDING, Fault, FaultKind, STEPPING, catching};
use henad_compute::gpu::fault::catching_on;
use henad_compute::gpu::{GpuContext, GpuSimState, StatsPoll, stepping};
use henad_core::action::Schedule;
use henad_core::explore::measure::Sampler;
use henad_core::explore::outcome::{PlannedRun, RunOutcome, RunStatus, StopReason};
use super::{
BatchEnd, ExecutionError, Executor, GpuTrackDepth, ReorderBuffer, RunRequest, RunSink, RunWatch, gpu_memory_budget,
};
use crate::cursor::{OutcomeParts, RunEnd, RunFailure, failed_build_outcome, milliseconds, refusal_note, run_outcome};
pub(super) fn run_on_tracks(
executor: &Executor<'_>,
ctx: &GpuContext,
requests: &[RunRequest<'_>],
sink: &mut dyn RunSink,
) -> Result<BatchEnd, ExecutionError> {
let mut interleaver = Interleaver::new(executor, ctx, requests);
loop {
if !executor.control.proceed() {
return Ok(BatchEnd::Aborted);
}
if let Some(end) = interleaver.round(sink)? {
return Ok(end);
}
}
}
struct Interleaver<'x> {
executor: &'x Executor<'x>,
ctx: &'x GpuContext,
requests: &'x [RunRequest<'x>],
unbuilt_requests: VecDeque<usize>,
tracks: Vec<GpuTrack>,
track_cap: usize,
gpu_memory_budget: u64,
admission_held: bool,
submissions: VecDeque<Submission>,
failed_builds: Vec<EndedRun>,
reorder: ReorderBuffer<RunOutcome>,
}
struct EndedRun {
request_index: usize,
outcome: RunOutcome,
watch: Option<RunWatch>,
}
struct Submission {
index: wgpu::SubmissionIndex,
done: Arc<AtomicBool>,
}
impl<'x> Interleaver<'x> {
fn new(executor: &'x Executor<'x>, ctx: &'x GpuContext, requests: &'x [RunRequest<'x>]) -> Self {
Self {
executor,
ctx,
requests,
unbuilt_requests: (0..requests.len()).collect(),
tracks: Vec::new(),
track_cap: executor.layout.gpu_tracks.max(1),
gpu_memory_budget: gpu_memory_budget(executor.gpu_memory_budget, ctx),
admission_held: false,
submissions: VecDeque::new(),
failed_builds: Vec::new(),
reorder: ReorderBuffer::default(),
}
}
fn round(&mut self, sink: &mut dyn RunSink) -> Result<Option<BatchEnd>, ExecutionError> {
let mut moved = self.admit();
let started = Instant::now();
if self.poll_device().is_err() || self.ctx.is_lost() {
return Ok(Some(BatchEnd::DeviceLost));
}
for position in 0..self.tracks.len() {
moved |= self.visit(position);
}
if let Some(fault) = self.ctx.faults.take() {
self.fail_every_track(&RunFailure::from(fault), None);
}
if !moved && self.wait_for_oldest().is_err() {
return Ok(Some(BatchEnd::DeviceLost));
}
self.charge(started.elapsed());
if !self.hand_out_finished(sink) {
return Ok(Some(BatchEnd::DeviceLost));
}
while let Some(outcome) = self.reorder.pop_ready() {
self.executor.commit(sink, outcome)?;
}
if self.unbuilt_requests.is_empty() && self.tracks.is_empty() {
debug_assert_eq!(self.reorder.handed_out(), self.requests.len(), "every run is committed");
return Ok(Some(BatchEnd::Complete));
}
Ok(None)
}
fn admit(&mut self) -> bool {
let mut tried = false;
while !self.admission_held && self.tracks.len() < self.track_cap {
let Some(&request_index) = self.unbuilt_requests.front() else {
break;
};
let request = &self.requests[request_index];
let demand_bytes = self
.executor
.entry
.demand(request.params, &self.ctx.device.limits())
.map_or(0, |demand| demand.bytes());
let live_bytes: u64 = self.tracks.iter().map(|track| track.demand_bytes).sum();
if !self.tracks.is_empty() && live_bytes.saturating_add(demand_bytes) > self.gpu_memory_budget {
break;
}
self.unbuilt_requests.pop_front();
tried = true;
match GpuTrack::build(self.executor, request_index, request, demand_bytes) {
Ok(track) => self.tracks.push(track),
Err(failed) if failed.fault.is_out_of_memory() && !self.tracks.is_empty() => {
self.unbuilt_requests.push_front(request_index);
self.track_cap = self.track_cap.saturating_sub(1).max(1);
self.admission_held = true;
}
Err(failed) => {
let failure = RunFailure::from(failed.fault);
let measure = &self.executor.measure;
let outcome = failed_build_outcome(request.run, request.run_key, measure, failed.build_ms, failure);
self.failed_builds.push(EndedRun {
request_index,
outcome,
watch: failed.watch,
});
}
}
}
tried
}
fn poll_device(&mut self) -> Result<(), Fault> {
let polled = catching(STEPPING, || self.ctx.device.poll(wgpu::PollType::Poll));
polled.and_then(|poll| poll.map(drop).map_err(poll_fault))?;
while self.submissions.front().is_some_and(Submission::is_done) {
self.submissions.pop_front();
}
Ok(())
}
fn visit(&mut self, position: usize) -> bool {
let track = &mut self.tracks[position];
if track.end.is_some() {
return false;
}
let landed = match track.check_sample(self.ctx) {
SampleCheck::Pending => false,
SampleCheck::Landed => true,
SampleCheck::Ended(end) => {
track.end = Some(end);
return true;
}
};
match track.record(self.ctx, self.executor.track_depth, &mut self.submissions) {
Ok(VisitWork::Nothing) => landed,
Ok(VisitWork::Submitted { sampled }) => {
if let Some(fault) = self.ctx.faults.take() {
self.fail_every_track(&RunFailure::from(fault), sampled.then_some(position));
}
true
}
Err(fault) => {
track.fail(fault.into(), false);
true
}
}
}
fn fail_every_track(&mut self, failure: &RunFailure, sampled_position: Option<usize>) {
for (position, track) in self.tracks.iter_mut().enumerate() {
if track.end.is_none() {
track.untraced_failure.get_or_insert_with(|| failure.clone());
track.fail(failure.clone(), sampled_position == Some(position));
}
}
}
fn wait_for_oldest(&self) -> Result<(), Fault> {
let Some(oldest) = self.submissions.front() else {
std::thread::yield_now();
return Ok(());
};
if oldest.is_done() {
return Ok(());
}
let (ctx, submission_index) = (self.ctx, oldest.index.clone());
catching(STEPPING, || stepping::await_submission(ctx, submission_index))?
}
fn charge(&mut self, elapsed: Duration) {
let timeout = self.executor.timeout;
let track_count = u32::try_from(self.tracks.len()).unwrap_or(u32::MAX);
let share = elapsed.checked_div(track_count).unwrap_or_default();
for track in &mut self.tracks {
track.wall += share;
let recording = track.end.is_none() && track.failure.is_none();
if recording && timeout.is_some_and(|timeout| track.wall >= timeout) {
track.end = Some(RunEnd::TimedOut);
}
}
}
fn hand_out_finished(&mut self, sink: &mut dyn RunSink) -> bool {
let mut ended = std::mem::take(&mut self.failed_builds);
let mut position = 0;
while position < self.tracks.len() {
let Some(end) = self.tracks[position].end.take() else {
position += 1;
continue;
};
let track = self.tracks.remove(position);
ended.push(track.finish(end, self.executor.timeout));
self.admission_held = false;
}
if ended.iter().any(|run| run.outcome.status == RunStatus::GpuError) && self.device_lost() {
return false;
}
for run in ended {
if let Some(watch) = run.watch {
watch.finish();
}
sink.finished(&run.outcome);
self.reorder.insert(run.request_index, run.outcome);
}
true
}
fn device_lost(&self) -> bool {
let drained = catching(STEPPING, || self.ctx.device.poll(wgpu::PollType::wait_indefinitely()));
!matches!(drained, Ok(Ok(_))) || self.ctx.is_lost()
}
}
impl Submission {
fn is_done(&self) -> bool {
self.done.load(Ordering::Acquire)
}
}
struct GpuTrack {
request_index: usize,
run: PlannedRun,
run_key: u64,
state: Box<dyn GpuSimState>,
schedule: Schedule,
demand_bytes: u64,
sampler: Sampler,
tick: u64,
actions_recorded: bool,
next_sample: Option<u64>,
readback_tick: Option<u64>,
landed_tick: u64,
population: u64,
refusals: Vec<(u64, String)>,
done_flags: VecDeque<Arc<AtomicBool>>,
build_ms: f64,
wall: Duration,
failure: Option<RunFailure>,
untraced_failure: Option<RunFailure>,
end: Option<RunEnd>,
watch: Option<RunWatch>,
}
struct BuildFailure {
fault: Fault,
build_ms: f64,
watch: Option<RunWatch>,
}
enum SampleCheck {
Pending,
Landed,
Ended(RunEnd),
}
enum VisitWork {
Nothing,
Submitted { sampled: bool },
}
impl GpuTrack {
fn build(
executor: &Executor<'_>,
request_index: usize,
request: &RunRequest<'_>,
demand_bytes: u64,
) -> Result<Self, BuildFailure> {
let watch = executor
.active_runs
.as_ref()
.map(|active_runs| active_runs.watch(request.run, executor.measure.total()));
let started = Instant::now();
let built = executor
.entry
.build(request.params, Some(request.run.seed), executor.gpu)
.and_then(gpu_state);
let build_ms = milliseconds(started.elapsed());
let state = match built {
Ok(state) => state,
Err(fault) => {
return Err(BuildFailure { fault, build_ms, watch });
}
};
Ok(Self {
request_index,
run: request.run,
run_key: request.run_key,
state,
schedule: request.schedule.clone(),
demand_bytes,
sampler: Sampler::new(Arc::clone(&executor.measure)),
tick: 0,
actions_recorded: false,
next_sample: Some(executor.measure.first_sample()),
readback_tick: None,
landed_tick: 0,
population: 0,
refusals: Vec::new(),
done_flags: VecDeque::new(),
build_ms,
wall: Duration::ZERO,
failure: None,
untraced_failure: None,
end: None,
watch,
})
}
fn check_sample(&mut self, ctx: &GpuContext) -> SampleCheck {
let Some(tick) = self.readback_tick else {
return self.failure.take().map_or(SampleCheck::Pending, |failure| {
SampleCheck::Ended(RunEnd::Failed(failure))
});
};
let state = &mut self.state;
let sample = match catching(STEPPING, || state.poll_stats_readback(&ctx.device, false)) {
Ok(StatsPoll::Pending) => return SampleCheck::Pending,
Ok(StatsPoll::Failed) => {
let failure = self.failure.take().unwrap_or_else(|| RunFailure {
status: RunStatus::GpuError,
note: format!("stats readback failed at tick {tick}"),
});
return SampleCheck::Ended(RunEnd::Failed(failure));
}
Ok(StatsPoll::Landed) => catching(STEPPING, || (state.stats(), state.population())),
Err(fault) => Err(fault),
};
let (stats, population) = match sample {
Ok(sample) => sample,
Err(fault) => {
let failure = self.failure.take().unwrap_or_else(|| fault.into());
return SampleCheck::Ended(RunEnd::Failed(failure));
}
};
self.readback_tick = None;
self.population = population;
self.landed_tick = tick;
match self.sampler.push(tick, &stats) {
Err(error) => SampleCheck::Ended(RunEnd::Failed(error.into())),
Ok(true) => SampleCheck::Ended(RunEnd::Stopped(StopReason::Condition)),
Ok(false) if tick == self.sampler.plan().total() => SampleCheck::Ended(RunEnd::Stopped(StopReason::Steps)),
Ok(false) => match self.failure.take() {
Some(failure) => SampleCheck::Ended(RunEnd::Failed(failure)),
None => SampleCheck::Landed,
},
}
}
fn record(
&mut self,
ctx: &GpuContext,
depth: GpuTrackDepth,
submissions: &mut VecDeque<Submission>,
) -> Result<VisitWork, Fault> {
if self.failure.is_some() {
return Ok(VisitWork::Nothing);
}
while self.done_flags.front().is_some_and(|done| done.load(Ordering::Acquire)) {
self.done_flags.pop_front();
}
if self.done_flags.len() >= depth.submissions_per_track {
return Ok(VisitWork::Nothing);
}
let mut work = VisitWork::Nothing;
if !self.actions_recorded {
self.actions_recorded = true;
if self.schedule.due(self.tick).next().is_some() {
let (state, schedule) = (&mut self.state, &self.schedule);
let refused = catching_on(ctx, STEPPING, || stepping::run_due(&mut **state, ctx, schedule))?;
self.refusals
.extend(refused.into_iter().map(|action| (action.tick, refusal_note(action))));
work = VisitWork::Submitted { sampled: false };
}
}
let Some(sample_tick) = self.next_sample else {
return Ok(work);
};
let (count, sampled) = if sample_tick == self.tick {
if self.readback_tick.is_some() {
return Ok(work);
}
(0, true)
} else {
let mut boundary = sample_tick.min(self.tick + u64::from(depth.steps_per_submission));
if let Some(due) = self.schedule.next_due_after(self.tick) {
boundary = boundary.min(due);
}
let actions_due = self.schedule.due(boundary).next().is_some();
let sampled = boundary == sample_tick && !actions_due && self.readback_tick.is_none();
((boundary - self.tick) as u32, sampled)
};
debug_assert!(
!sampled || !self.state.stats_readback_pending(),
"run {} recorded a sample while an earlier one reads back, and its copy would be skipped",
self.run.run_id
);
let state = &mut self.state;
let submission_index = catching_on(ctx, STEPPING, || {
stepping::submit_slice(&mut **state, ctx, count, sampled)
})?;
let done = Arc::new(AtomicBool::new(false));
let callback_done = Arc::clone(&done);
ctx.queue
.on_submitted_work_done(move || callback_done.store(true, Ordering::Release));
self.done_flags.push_back(Arc::clone(&done));
submissions.push_back(Submission {
index: submission_index,
done,
});
if count > 0 {
self.tick += u64::from(count);
self.actions_recorded = self.schedule.due(self.tick).next().is_none();
}
if sampled {
self.readback_tick = Some(self.tick);
self.next_sample = self.sampler.plan().next_sample(self.tick);
}
if let Some(watch) = &self.watch {
watch.reach(self.tick);
}
Ok(VisitWork::Submitted { sampled })
}
fn fail(&mut self, failure: RunFailure, discard_readback: bool) {
if discard_readback {
self.readback_tick = None;
}
self.failure.get_or_insert(failure);
}
fn finish(self, end: RunEnd, timeout: Option<Duration>) -> EndedRun {
let end = match (end, self.untraced_failure) {
(RunEnd::Failed(failure), _) | (_, Some(failure)) => RunEnd::Failed(failure),
(end, None) => end,
};
let ticks = self.landed_tick;
let refusals = self
.refusals
.into_iter()
.filter(|&(tick, _)| tick <= ticks)
.map(|(_, note)| note)
.collect();
let parts = OutcomeParts {
sampler: self.sampler,
ticks,
population: self.population,
build_ms: self.build_ms,
wall: self.wall,
timeout,
refusals,
};
EndedRun {
request_index: self.request_index,
outcome: run_outcome(self.run, self.run_key, end, parts),
watch: self.watch,
}
}
}
fn gpu_state(state: ModelState) -> Result<Box<dyn GpuSimState>, Fault> {
match state {
ModelState::Gpu(state) => Ok(state),
ModelState::Cpu(_) => Err(Fault::refused(BUILDING, "a CPU model cannot run on a GPU track")),
}
}
fn poll_fault(error: wgpu::PollError) -> Fault {
Fault {
during: STEPPING,
kind: FaultKind::Poll(error),
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use henad_compute::entry::ModelEntry;
use henad_compute::gpu::GpuContext;
use henad_core::explore::measure::MeasurePlan;
use henad_core::explore::outcome::{RunStatus, StopReason};
use henad_core::explore::plan::Plan;
use henad_core::explore::spec::SweepSpec;
use super::Interleaver;
use crate::cursor::{RunEnd, RunFailure};
use crate::exec::{Executor, RunRequest, SweepControl};
use crate::tests::support::{entry, headless_device, planned, tracks};
fn four_runs(ctx: &GpuContext) -> (ModelEntry, Plan, Arc<MeasurePlan>) {
let mut spec = SweepSpec::new("gpu_game_of_life");
spec.fixed = vec![
("grid_width".to_owned(), "16".to_owned()),
("grid_height".to_owned(), "16".to_owned()),
];
spec.run.steps = 10;
spec.run.replicates = 4;
let model = entry("gpu_game_of_life", Some(ctx));
let (plan, measure) = planned(&model, Some(ctx), &spec);
(model, plan, measure)
}
fn executor<'e>(
model: &'e ModelEntry,
ctx: &'e GpuContext,
measure: &Arc<MeasurePlan>,
timeout: Option<Duration>,
) -> Executor<'e> {
Executor::new(model, Some(ctx), Arc::clone(measure), tracks(4), SweepControl::new())
.expect("a device")
.with_timeout(timeout)
}
#[test]
fn a_round_is_split_evenly_between_the_live_tracks() {
let Some(ctx) = headless_device() else {
return;
};
let (model, plan, measure) = four_runs(&ctx);
let executor = executor(&model, &ctx, &measure, Some(Duration::from_millis(30)));
let requests: Vec<RunRequest<'_>> = plan.runs().map(|run| RunRequest::planned(&plan, run)).collect();
let mut interleaver = Interleaver::new(&executor, &ctx, &requests);
interleaver.admit();
assert_eq!(interleaver.tracks.len(), 4);
interleaver.charge(Duration::from_millis(100));
for track in &interleaver.tracks {
assert_eq!(track.wall, Duration::from_millis(25), "run {}", track.run.run_id);
assert!(track.end.is_none(), "a share of 25 ms is inside the timeout of 30 ms");
}
interleaver.charge(Duration::from_millis(40));
for track in &interleaver.tracks {
assert_eq!(track.wall, Duration::from_millis(35), "run {}", track.run.run_id);
assert!(matches!(track.end, Some(RunEnd::TimedOut)), "run {}", track.run.run_id);
}
}
#[test]
fn an_untraced_fault_leaves_a_run_that_ended_alone() {
let Some(ctx) = headless_device() else {
return;
};
let (model, plan, measure) = four_runs(&ctx);
let executor = executor(&model, &ctx, &measure, None);
let requests: Vec<RunRequest<'_>> = plan.runs().map(|run| RunRequest::planned(&plan, run)).collect();
let mut interleaver = Interleaver::new(&executor, &ctx, &requests);
interleaver.admit();
interleaver.tracks[0].end = Some(RunEnd::Stopped(StopReason::Steps));
let failure = RunFailure {
status: RunStatus::GpuError,
note: "a validation error on another track".to_owned(),
};
interleaver.fail_every_track(&failure, None);
let mut tracks = std::mem::take(&mut interleaver.tracks).into_iter();
let mut ended = tracks.next().expect("four tracks are live");
let end = ended.end.take().expect("the run ended before the fault");
let outcome = ended.finish(end, None).outcome;
assert_eq!(
(outcome.status, outcome.stop_reason),
(RunStatus::Ok, StopReason::Steps)
);
for live in tracks {
assert!(live.untraced_failure.is_some(), "run {}", live.run.run_id);
assert!(live.failure.is_some(), "run {}", live.run.run_id);
}
}
}