Something went wrong. Try again.
a tiny pet with an even tinier brain, living in your browser
Something went wrong. Try again.
12 kB · 377 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378//! Bounded, anonymous timing estimates and two-worker queue planning.//! Costs retain only numeric aggregates. The waiting queue holds current record//! IDs in memory; neither it nor timing estimates are saved or logged.use std::collections::HashMap;use std::time::Instant;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]pub enum Kind { Reactive, Idle, Diary,}impl Kind { fn index(self) -> usize { match self { Self::Reactive => 0, Self::Idle => 1, Self::Diary => 2, } }}#[derive(Clone, Copy, Debug, PartialEq, Eq)]pub struct Work { pub kind: Kind, band: usize,}impl Work { pub fn new(kind: Kind, estimated_tokens: usize) -> Self { Self { kind, band: if estimated_tokens < 256 { 0 } else if estimated_tokens < 640 { 1 } else { 2 }, } } fn index(self) -> usize { self.kind.index() * 3 + self.band }}#[derive(Clone, Copy, Default)]struct Sample { count: u64, mean_ms: f64,}#[derive(Clone, Default)]pub struct Costs { samples: [Sample; 9],}impl Costs { pub fn observe(&mut self, work: Work, ms: f64) { if !ms.is_finite() || ms <= 0.0 { return; } let sample = &mut self.samples[work.index()]; // A plain mean for the first five samples, so one unusual first turn // does not dominate; then an exponentially weighted mean (20% new). let weight = (1.0 / (sample.count as f64 + 1.0)).max(0.2); sample.mean_ms += (ms - sample.mean_ms) * weight; sample.count = sample.count.saturating_add(1); } pub fn estimate(&self, work: Work) -> Option<f64> { let exact = self.samples[work.index()]; if exact.count > 0 { return Some(exact.mean_ms); } // A different prompt band is deliberately not borrowed: CPU prefill // scales much more sharply than CUDA. Unseen work retains affinity. None } pub fn sample_count(&self, work: Work) -> u64 { self.samples[work.index()].count }}
/// Per-lane estimates for planning. When only one lane has run this shape,/// the other lane's estimate is that cost times the least favourable speed/// ratio seen on shapes both lanes have run. This lets an unseen GPU shape/// wait for a GPU that is clearly faster instead of always spilling to an/// idle CPU. Without such a shared shape, unseen work keeps its affinity.pub fn estimates(costs: [&Costs; 2], work: Work) -> [Option<f64>; 2] { let ratio = |lane: usize| { costs[lane] .samples .iter() .zip(&costs[1 - lane].samples) .filter(|(a, b)| a.count > 0 && b.count > 0) .map(|(a, b)| a.mean_ms / b.mean_ms) .reduce(f64::max) }; match [costs[0].estimate(work), costs[1].estimate(work)] { [None, Some(gpu)] => [ratio(0).map(|r| gpu * r), Some(gpu)], [Some(cpu), None] => [Some(cpu), ratio(1).map(|r| cpu * r)], exact => exact, }}
pub struct Waiting { pub epoch: u64, pub context: u64, pub work: Work, pub since: Instant,}#[derive(Default)]pub struct Queue { pub waiting: HashMap<String, Waiting>,}impl Queue { pub fn age_ms(&mut self, id: &str, epoch: u64, context: u64, work: Work, now: Instant) -> f64 { let entry = self.waiting.entry(id.into()).or_insert(Waiting { epoch, context, work, since: now, }); if entry.epoch != epoch || entry.work.kind != work.kind { *entry = Waiting { epoch, context, work, since: now, }; } entry.context = context; entry.work = work; now.duration_since(entry.since).as_secs_f64() * 1000.0 }}#[derive(Clone, Copy)]pub struct Request { pub work: Work, pub age_ms: f64,}pub const BACKGROUND_MAX_WAIT_MS: f64 = 30000.0;
pub fn priority(request: Request) -> u8 { if request.work.kind != Kind::Reactive && request.age_ms >= BACKGROUND_MAX_WAIT_MS { 0 } else if request.work.kind == Kind::Reactive { 1 } else { 2 }}
// Returns a lane for every request, in the supplied priority/FIFO order.// A lane can be busy: its planned request waits instead of spilling blindly.pub fn plan( costs: [&Costs; 2], remaining: [f64; 2], available: [bool; 2], free: [bool; 2], requests: &[Request],) -> Vec<usize> { let mut finish = remaining; let mut cold_free = free; requests .iter() .map(|request| { let times = estimates(costs, request.work); let preferred = if request.work.kind == Kind::Reactive { 1 } else { 0 }; let lane = if !available[0] { 1 } else if !available[1] { 0 } else if let [Some(cpu), Some(gpu)] = times { // Prefer a clear improvement before sending work to a slower // lane. The 10% margin limits oscillation from timing noise. let cpu_end = finish[0] + cpu; let gpu_end = finish[1] + gpu; if cpu_end * 1.1 < gpu_end { 0 } else { 1 } } else if cold_free[preferred] { preferred } else if cold_free[1 - preferred] { 1 - preferred } else { preferred }; cold_free[lane] = false; finish[lane] += times[lane].unwrap_or(remaining[lane].max(1.0)); lane }) .collect()}
#[cfg(test)]mod tests { use super::*; fn request() -> Request { Request { work: Work::new(Kind::Reactive, 100), age_ms: 0.0, } } fn close(value: Option<f64>, expected: f64) -> bool { value.is_some_and(|v| (v - expected).abs() < 1e-9) } fn costs() -> [Costs; 2] { let mut c = [Costs::default(), Costs::default()]; c[0].observe(request().work, 140.0); c[1].observe(request().work, 30.0); c } #[test] fn small_queue_waits_for_gpu_and_large_queue_uses_both() { let c = costs(); assert_eq!( plan( [&c[0], &c[1]], [0.0; 2], [true; 2], [true; 2], &vec![request(); 4] ), vec![1; 4] ); let large = plan( [&c[0], &c[1]], [0.0; 2], [true; 2], [true; 2], &vec![request(); 16], ); assert!(large.contains(&0) && large.contains(&1)); assert_eq!( plan( [&c[0], &c[1]], [0.0, 20.0], [true; 2], [true, false], &vec![request(); 2] ), vec![1; 2] ); assert_eq!( plan( [&c[0], &c[1]], [0.0, f64::INFINITY], [true; 2], [true, false], &[request()] ), vec![0], "a busy unseen GPU shape cannot promise an immediate finish" ); } #[test] fn unknown_shapes_keep_affinity_and_unavailable_gpu_uses_cpu() { let c = costs(); let requests = [ request(), Request { work: Work::new(Kind::Diary, 700), age_ms: 0.0, }, ]; assert_eq!( plan([&c[0], &c[1]], [0.0; 2], [true; 2], [true; 2], &requests), vec![1, 0] ); assert_eq!( plan( [&c[0], &c[1]], [0.0; 2], [true, false], [true, false], &requests ), vec![0, 0] ); } #[test] fn estimates_are_finite_bounded_and_separate_by_work_shape() { let mut c = Costs::default(); let work = request().work; c.observe(work, 100.0); c.observe(work, 200.0); c.observe(work, f64::NAN); c.observe(work, -1.0); assert_eq!(c.estimate(work), Some(150.0)); assert_eq!(c.sample_count(work), 2); assert_eq!(c.estimate(Work::new(Kind::Reactive, 900)), None); assert_eq!(c.estimate(Work::new(Kind::Idle, 100)), None); } #[test] fn an_unusual_first_sample_is_averaged_out_quickly() { // Measured on the GPU: a two-call clock turn, then single-call turns. let mut c = Costs::default(); for ms in [64.0, 30.0, 28.0, 28.0, 33.0] { c.observe(request().work, ms); } assert!(close(c.estimate(request().work), 36.6)); c.observe(request().work, 31.0); assert!(close(c.estimate(request().work), 35.48)); } #[test] fn a_shape_seen_on_one_lane_is_estimated_from_shared_shapes() { let long = Work::new(Kind::Reactive, 900); let mut c = costs(); c[0].observe(long, 1070.0); // Measured 4-pet mixed burst: the GPU was busy and had never run the // long prompt, so the idle CPU took it for about one second. let busy_gpu = [0.0, 30.0]; let long_request = Request { work: long, age_ms: 0.0, }; assert!(close( estimates([&c[0], &c[1]], long)[1], 1070.0 * 30.0 / 140.0 )); assert_eq!( plan( [&c[0], &c[1]], busy_gpu, [true; 2], [true, false], &[long_request] ), vec![1] ); // The CPU side uses the least favourable CPU/GPU ratio. let mut d = costs(); d[1].observe(long, 100.0); assert!(close( estimates([&d[0], &d[1]], long)[0], 100.0 * 140.0 / 30.0 )); assert_eq!( plan( [&d[0], &d[1]], [0.0; 2], [true; 2], [true; 2], &[long_request; 4] ), vec![1; 4] ); assert!( plan( [&d[0], &d[1]], [0.0; 2], [true; 2], [true; 2], &[long_request; 16] ) .contains(&0) ); // Without a shape seen on both lanes, unseen work keeps its affinity. let mut e = [Costs::default(), Costs::default()]; e[0].observe(long, 1070.0); assert_eq!(estimates([&e[0], &e[1]], long), [Some(1070.0), None]); } #[test] fn background_ages_ahead_of_new_reactive_work_and_context_edits_keep_its_place() { let work = Work::new(Kind::Diary, 100); let start = Instant::now(); let mut queue = Queue::default(); assert_eq!(queue.age_ms("test-record", 1, 1, work, start), 0.0); let later = start + std::time::Duration::from_secs(31); let age = queue.age_ms("test-record", 1, 2, Work::new(Kind::Diary, 900), later); assert_eq!(age, 31000.0); assert!(priority(Request { work, age_ms: age }) < priority(request())); assert_eq!( queue.age_ms("test-record", 2, 3, work, later), 0.0, "a new executor starts a new queue wait" ); }}