1use std::collections::{HashMap, VecDeque};
169use std::fmt::Debug;
170use std::sync::Mutex;
171use std::sync::atomic::{AtomicU64, Ordering};
172use std::sync::mpsc::{Receiver, Sender, channel};
173use std::thread;
174use std::time::{Duration, Instant};
175
176use metrics::{Counter, Gauge, Histogram, Label, Unit};
177
178use crate::internal::*;
179
180#[derive(Debug, Clone, PartialEq, Eq)]
189pub struct LaneTable {
190 taken: Vec<bool>,
193}
194
195impl LaneTable {
196 pub fn new(max_lanes: usize) -> TractResult<LaneTable> {
197 ensure!(max_lanes > 0, "A laned state needs at least one lane");
198 Ok(LaneTable { taken: vec![false; max_lanes] })
199 }
200
201 pub fn max_lanes(&self) -> usize {
204 self.taken.len()
205 }
206
207 pub fn taken(&self) -> usize {
208 self.taken.iter().filter(|t| **t).count()
209 }
210
211 pub fn take(&mut self) -> Option<LaneId> {
215 let lane = self.taken.iter().position(|t| !t)?;
216 self.taken[lane] = true;
217 Some(LaneId(lane))
218 }
219
220 pub fn give_back(&mut self, lane: LaneId) -> TractResult<()> {
222 ensure!(self.is_taken(lane), "Lane {} is not taken, so it can not be given back", lane.0);
223 self.taken[lane.0] = false;
224 Ok(())
225 }
226
227 pub fn is_taken(&self, lane: LaneId) -> bool {
228 self.taken.get(lane.0).copied().unwrap_or(false)
229 }
230
231 pub fn seat(&self, lanes: impl IntoIterator<Item = LaneId>) -> TractResult<Seating> {
235 let lanes: Vec<LaneId> = lanes.into_iter().collect();
236 for lane in &lanes {
237 ensure!(self.is_taken(*lane), "Seating lane {}, which no stream took", lane.0);
238 }
239 Seating::new(self.max_lanes(), lanes)
240 }
241}
242
243crate::declare_knob!(
244 scopable TRACT_MAX_SEATS,
245 usize,
246 256,
247 "Most streams a laned runtime serves in one turn, clamped to the state's lanes."
248);
249
250crate::declare_knob!(
251 scopable TRACT_AUTOBATCH_LANES,
252 Option<usize>,
253 None,
254 "Sessions a deployment config's subnet is autobatched for; unset, it is not autobatched."
255);
256
257crate::declare_knob!(
258 scopable TRACT_TURN_LINGER_US,
259 usize,
260 0,
261 "How long a laned runtime waits for more streams once one is ready to run."
262);
263
264#[derive(Clone)]
279pub struct LanedRunnable {
280 shared: Arc<Shared>,
281}
282
283struct Shared {
284 requests: Mutex<Option<Sender<Request>>>,
288 worker: Mutex<Option<thread::JoinHandle<()>>>,
291 inner: Arc<dyn Runnable>,
293 model: Option<Arc<TypedModel>>,
296 plan: Option<Arc<TypedSimplePlan>>,
297 batch: Symbol,
300 max_lanes: usize,
303 counts: Arc<Counts>,
305}
306
307impl Drop for Shared {
316 fn drop(&mut self) {
317 if let Ok(mut requests) = self.requests.lock() {
318 requests.take();
319 }
320 let worker = self.worker.lock().ok().and_then(|mut worker| worker.take());
321 if let Some(worker) = worker {
322 let _ = worker.join();
323 }
324 }
325}
326
327#[derive(Debug, Default)]
330struct Counts {
331 turns: AtomicU64,
333 seats: AtomicU64,
335}
336
337struct Meters {
342 turns: Counter,
343 seats: Counter,
344 seats_per_turn: Histogram,
345 lanes_taken: Gauge,
346 queue_wait: Histogram,
347 turn: Histogram,
348}
349
350impl Meters {
351 fn new(name: Option<&str>) -> Meters {
352 metrics::describe_counter!(
353 "tract_autobatch_turns_total",
354 "Turns an autobatched runnable ran"
355 );
356 metrics::describe_counter!(
357 "tract_autobatch_seats_total",
358 "Seats an autobatched runnable filled over its turns"
359 );
360 metrics::describe_histogram!(
361 "tract_autobatch_seats_per_turn",
362 "Seats one turn of an autobatched runnable filled"
363 );
364 metrics::describe_gauge!(
365 "tract_autobatch_lanes_taken",
366 "Lanes held by the streams of an autobatched runnable"
367 );
368 metrics::describe_histogram!(
369 "tract_autobatch_queue_wait_seconds",
370 Unit::Seconds,
371 "From a seat queued to the start of the turn seating it, linger included"
372 );
373 metrics::describe_histogram!(
374 "tract_autobatch_turn_seconds",
375 Unit::Seconds,
376 "One turn of an autobatched runnable, stacking and slicing included"
377 );
378 let labels: Vec<Label> =
379 name.map(|name| Label::new("runnable", name.to_string())).into_iter().collect();
380 Meters {
381 turns: metrics::counter!("tract_autobatch_turns_total", labels.clone()),
382 seats: metrics::counter!("tract_autobatch_seats_total", labels.clone()),
383 seats_per_turn: metrics::histogram!("tract_autobatch_seats_per_turn", labels.clone()),
384 lanes_taken: metrics::gauge!("tract_autobatch_lanes_taken", labels.clone()),
385 queue_wait: metrics::histogram!("tract_autobatch_queue_wait_seconds", labels.clone()),
386 turn: metrics::histogram!("tract_autobatch_turn_seconds", labels),
387 }
388 }
389}
390
391impl LanedRunnable {
392 pub fn wrap(inner: Arc<dyn Runnable>, max_lanes: usize) -> TractResult<LanedRunnable> {
396 let model = inner.typed_model().cloned();
397 let plan = inner.typed_plan().cloned();
398 let mut symbols: Vec<Symbol> = vec![];
399 let mut facts: Vec<(String, TypedFact)> = vec![];
400 let mut batch_in: Vec<bool> = vec![];
401 for ix in 0..inner.input_count() {
402 let fact = inner.input_fact(ix)?;
403 let symbol = batch_symbol(fact);
404 batch_in.push(symbol.is_some());
405 symbols.extend(symbol);
406 facts.push((format!("input {ix}"), fact.clone()));
407 }
408 let mut batch_out: Vec<bool> = vec![];
409 for ix in 0..inner.output_count() {
410 let fact = inner.output_fact(ix)?;
411 let symbol = batch_symbol(fact);
412 batch_out.push(symbol.is_some());
413 symbols.extend(symbol);
414 facts.push((format!("output {ix}"), fact.clone()));
415 }
416 symbols.sort();
417 symbols.dedup();
418 if symbols.len() != 1 {
419 bail!(off_axis_zero(&facts, &symbols));
420 }
421 ensure!(batch_out.iter().any(|b| *b), "A laned model must batch one output at least");
422 let batch = symbols.remove(0);
423 let counts = Arc::new(Counts::default());
424 let meters = Meters::new(model.as_ref().and_then(|model| model.name()));
425 let max_seats = TRACT_MAX_SEATS.get().min(max_lanes);
426 let linger = Duration::from_micros(TRACT_TURN_LINGER_US.get() as u64);
427 let lanes = LaneTable::new(max_lanes)?;
428 let (requests, queue) = channel::<Request>();
429 let (spawned, ready) = channel::<TractResult<()>>();
430 let worker_counts = counts.clone();
431 let worker_inner = inner.clone();
432 let worker_thread = thread::Builder::new().name("tract-lanes".into()).spawn(move || {
433 let state = match worker_inner.spawn().and_then(|mut state| {
434 let every: Vec<LaneId> = (0..max_lanes).map(LaneId).collect();
435 state.reset_lanes(&every).context("Preparing a laned model")?;
436 Ok(state)
437 }) {
438 Ok(state) => {
439 let _ = spawned.send(Ok(()));
440 state
441 }
442 Err(e) => {
443 let _ = spawned.send(Err(e));
444 return;
445 }
446 };
447 Worker {
448 state,
449 lanes,
450 queued: Queue::default(),
451 batch_in,
452 batch_out,
453 max_seats,
454 linger,
455 counts: worker_counts,
456 meters,
457 }
458 .work(queue);
459 })?;
460 ready.recv().map_err(|_| format_err!("The laned worker died spawning the state"))??;
461 Ok(LanedRunnable {
462 shared: Arc::new(Shared {
463 requests: Mutex::new(Some(requests)),
464 worker: Mutex::new(Some(worker_thread)),
465 inner,
466 model,
467 plan,
468 batch,
469 max_lanes,
470 counts,
471 }),
472 })
473 }
474
475 pub fn max_lanes(&self) -> usize {
476 self.shared.max_lanes
477 }
478
479 pub fn inner(&self) -> &Arc<dyn Runnable> {
482 &self.shared.inner
483 }
484
485 pub fn batch_symbol(&self) -> &Symbol {
489 &self.shared.batch
490 }
491
492 pub fn turns_and_seats(&self) -> (u64, u64) {
495 (
496 self.shared.counts.turns.load(Ordering::Relaxed),
497 self.shared.counts.seats.load(Ordering::Relaxed),
498 )
499 }
500
501 fn request(&self) -> TractResult<Sender<Request>> {
502 self.shared
503 .requests
504 .lock()
505 .map_err(|_| format_err!("Poisoned laned sender"))?
506 .clone()
507 .context("The laned runnable is gone")
508 }
509}
510
511fn batch_symbol(fact: &TypedFact) -> Option<Symbol> {
515 match fact.shape.dims().first() {
516 Some(TDim::Sym(sym)) => Some(sym.clone()),
517 _ => None,
518 }
519}
520
521fn off_axis_zero(facts: &[(String, TypedFact)], symbols: &[Symbol]) -> TractError {
524 let elsewhere: Vec<String> = facts
525 .iter()
526 .filter_map(|(what, fact)| {
527 let dims = fact.shape.dims();
528 let axis = dims.iter().skip(1).position(|dim| matches!(dim, TDim::Sym(_)))? + 1;
529 Some(format!("{what} carries {} on axis {axis}", dims[axis]))
530 })
531 .collect();
532 let elsewhere = elsewhere.join(", ");
533 if symbols.is_empty() {
534 let found = if elsewhere.is_empty() {
535 "none of them carries a symbol at all".to_string()
536 } else {
537 elsewhere
538 };
539 format_err!(
540 "A laned model carries its batch symbol on axis 0 of the inputs and outputs it \
541 batches, and {found}. Axis 0 is where a seat's values are a contiguous run, so move \
542 the batch there as a graph edit -- Batchify, or an AddAxis/MoveAxis the optimiser \
543 can absorb -- rather than have every turn transpose"
544 )
545 } else {
546 let symbols: Vec<String> = symbols.iter().map(|s| s.to_string()).collect();
547 format_err!(
548 "A laned model carries one batch symbol on axis 0, this one carries {}. One symbol \
549 has to stand for the whole turn's occupancy, so share it across the batched inputs \
550 and outputs in the export",
551 symbols.join(" and ")
552 )
553 }
554}
555
556impl Debug for LanedRunnable {
557 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
558 write!(f, "LanedRunnable({} lanes)", self.shared.max_lanes)
559 }
560}
561
562impl Runnable for LanedRunnable {
563 fn spawn(&self) -> TractResult<Box<dyn State>> {
564 let requests = self.request()?;
565 let (taken, lane) = channel();
566 requests
567 .send(Request::Spawn(taken))
568 .map_err(|_| format_err!("The laned worker is gone"))?;
569 let lane = lane.recv().map_err(|_| format_err!("The laned worker dropped a lane"))??;
570 Ok(Box::new(LanedStateHandle {
571 lease: Arc::new(Lease { lane, requests }),
572 runnable: self.clone(),
573 }))
574 }
575
576 fn typed_plan(&self) -> Option<&Arc<TypedSimplePlan>> {
577 self.shared.plan.as_ref()
578 }
579
580 fn typed_model(&self) -> Option<&Arc<TypedModel>> {
581 self.shared.model.as_ref()
582 }
583}
584
585#[derive(Clone, Debug)]
589pub struct LanedStateHandle {
590 lease: Arc<Lease>,
593 runnable: LanedRunnable,
595}
596
597#[derive(Debug)]
598struct Lease {
599 lane: LaneId,
601 requests: Sender<Request>,
603}
604
605impl Drop for Lease {
606 fn drop(&mut self) {
607 let _ = self.requests.send(Request::Drop(self.lane));
608 }
609}
610
611impl State for LanedStateHandle {
612 fn run(&mut self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
613 let (done, outputs) = channel();
614 self.lease
615 .requests
616 .send(Request::Call(Call { leased: self.lease.lane, inputs, done }))
617 .map_err(|_| format_err!("The laned worker is gone"))?;
618 outputs.recv().map_err(|_| format_err!("The laned worker dropped a turn"))?
619 }
620
621 fn runnable(&self) -> &dyn Runnable {
622 &self.runnable
623 }
624}
625
626enum Request {
630 Spawn(Sender<TractResult<LaneId>>),
633 Call(Call),
635 Drop(LaneId),
637}
638
639struct Call {
643 leased: LaneId,
647 inputs: TVec<TValue>,
649 done: Sender<TractResult<TVec<TValue>>>,
651}
652
653struct Seat {
657 call: u64,
659 lane: LaneId,
662 ix: usize,
664 inputs: TVec<TValue>,
666 queued: Instant,
668}
669
670struct Completer {
674 served: Vec<Option<TVec<TValue>>>,
677 done: Sender<TractResult<TVec<TValue>>>,
679}
680
681impl Call {
682 fn explode(&self, batch_in: &[bool]) -> TractResult<Vec<TVec<TValue>>> {
687 ensure!(
688 self.inputs.len() == batch_in.len(),
689 "A call feeds {} inputs, the model takes {}",
690 self.inputs.len(),
691 batch_in.len()
692 );
693 let mut seats: Option<usize> = None;
694 for (ix, input) in self.inputs.iter().enumerate().filter(|(ix, _)| batch_in[*ix]) {
695 ensure!(
696 input.rank() > 0,
697 "Input {ix} carries the batch axis, so it can not be a scalar"
698 );
699 let asked = input.shape()[0];
700 ensure!(
701 seats.is_none_or(|seats| seats == asked),
702 "A call asks for as many seats as its batched inputs carry, and input {ix} carries \
703 {asked} against {} before it",
704 seats.unwrap_or(0)
705 );
706 seats = Some(asked);
707 }
708 let seats = seats.unwrap_or(1);
709 ensure!(seats > 0, "A call asks for one seat at least");
710 if seats == 1 {
711 return Ok(vec![self.inputs.clone()]);
712 }
713 (0..seats)
714 .map(|seat| {
715 self.inputs
716 .iter()
717 .zip(batch_in)
718 .map(|(input, is_batched)| {
719 if *is_batched {
720 Ok(input.slice(0, seat, seat + 1)?.into_tvalue())
721 } else {
722 Ok(input.clone())
723 }
724 })
725 .collect()
726 })
727 .collect()
728 }
729}
730
731impl Completer {
732 fn is_full(&self) -> bool {
734 self.served.iter().all(Option::is_some)
735 }
736
737 fn answer(self, batch_out: &[bool]) {
741 let served: Vec<TVec<TValue>> =
742 self.served.into_iter().map(|outputs| outputs.unwrap()).collect();
743 let answer = if served.len() == 1 {
744 Ok(served.into_iter().next().unwrap())
745 } else {
746 let first = &served[0];
747 (0..first.len())
748 .map(|ix| {
749 if batch_out.get(ix).copied().unwrap_or(false) {
750 let seats: TVec<&Tensor> =
751 served.iter().map(|outputs| &*outputs[ix]).collect();
752 Ok(Tensor::stack_tensors(0, &seats)?.into_tvalue())
753 } else {
754 Ok(first[ix].clone())
755 }
756 })
757 .collect()
758 };
759 let _ = self.done.send(answer);
760 }
761}
762
763#[derive(Default)]
768struct Queue {
769 seats: VecDeque<Seat>,
771 completers: HashMap<u64, Completer>,
773 calls: u64,
775}
776
777impl Queue {
778 fn push(&mut self, call: Call, batch_in: &[bool]) {
779 let id = self.calls;
780 self.calls += 1;
781 let queued = Instant::now();
782 match call.explode(batch_in) {
783 Ok(seats) => {
784 self.completers
785 .insert(id, Completer { served: vec![None; seats.len()], done: call.done });
786 for (ix, inputs) in seats.into_iter().enumerate() {
787 self.seats.push_back(Seat { call: id, lane: call.leased, ix, inputs, queued });
788 }
789 }
790 Err(e) => {
791 let _ = call.done.send(Err(e));
792 }
793 }
794 }
795
796 fn serve(&mut self, seat: Seat, outputs: TractResult<TVec<TValue>>, batch_out: &[bool]) {
800 if !self.completers.contains_key(&seat.call) {
801 return;
802 }
803 let outputs = match outputs {
804 Ok(outputs) => outputs,
805 Err(e) => {
806 let completer = self.completers.remove(&seat.call).unwrap();
807 self.seats.retain(|queued| queued.call != seat.call);
808 let _ = completer.done.send(Err(e));
809 return;
810 }
811 };
812 let completer = self.completers.get_mut(&seat.call).unwrap();
813 completer.served[seat.ix] = Some(outputs);
814 if completer.is_full() {
815 self.completers.remove(&seat.call).unwrap().answer(batch_out);
816 }
817 }
818}
819
820struct Worker {
827 state: Box<dyn State>,
829 lanes: LaneTable,
831 queued: Queue,
833 batch_in: Vec<bool>,
836 batch_out: Vec<bool>,
839 max_seats: usize,
841 linger: Duration,
844 counts: Arc<Counts>,
847 meters: Meters,
848}
849
850impl Worker {
851 fn work(mut self, queue: Receiver<Request>) {
853 loop {
854 while self.queued.seats.is_empty() {
859 match queue.recv() {
860 Ok(request) => self.serve(request),
861 Err(_) => return,
862 }
863 if !self.queued.seats.is_empty() && !self.linger.is_zero() {
864 thread::sleep(self.linger);
865 }
866 }
867 while let Ok(request) = queue.try_recv() {
868 self.serve(request);
869 }
870 let (seated, borrowed) = self.fill();
871 if seated.is_empty() {
872 continue;
873 }
874 self.counts.turns.fetch_add(1, Ordering::Relaxed);
875 self.counts.seats.fetch_add(seated.len() as u64, Ordering::Relaxed);
876 let started = Instant::now();
877 self.meters.turns.increment(1);
878 self.meters.seats.increment(seated.len() as u64);
879 self.meters.seats_per_turn.record(seated.len() as f64);
880 for seat in &seated {
881 self.meters.queue_wait.record(started.duration_since(seat.queued));
882 }
883 let served = self.run_turn(&seated);
884 self.meters.turn.record(started.elapsed());
885 for lane in borrowed {
886 let _ = self.lanes.give_back(lane);
887 }
888 match served {
889 Ok(per_seat) => {
890 for (seat, outputs) in seated.into_iter().zip(per_seat) {
891 self.queued.serve(seat, Ok(outputs), &self.batch_out);
892 }
893 }
894 Err(e) => {
895 let e = format!("{e:#}");
896 for seat in seated {
897 self.queued.serve(
898 seat,
899 Err(format_err!("Laned turn failed: {e}")),
900 &self.batch_out,
901 );
902 }
903 }
904 }
905 }
906 }
907
908 fn serve(&mut self, request: Request) {
912 match request {
913 Request::Spawn(taken) => {
914 let lane = match self.lanes.take() {
915 None => Err(format_err!(
916 "Every one of the {} lanes is taken",
917 self.lanes.max_lanes()
918 )),
919 Some(lane) => match self.state.reset_lanes(&[lane]) {
920 Ok(()) => Ok(lane),
921 Err(e) => {
922 let _ = self.lanes.give_back(lane);
923 Err(e)
924 }
925 },
926 };
927 self.meters.lanes_taken.set(self.lanes.taken() as f64);
928 let _ = taken.send(lane);
929 }
930 Request::Call(call) => self.queued.push(call, &self.batch_in),
931 Request::Drop(lane) => {
932 let _ = self.lanes.give_back(lane);
933 self.meters.lanes_taken.set(self.lanes.taken() as f64);
934 }
935 }
936 }
937
938 fn fill(&mut self) -> (Vec<Seat>, Vec<LaneId>) {
943 let mut seated: Vec<Seat> = vec![];
944 let mut taken: Vec<LaneId> = vec![];
945 let mut borrowed: Vec<LaneId> = vec![];
946 let mut waiting: VecDeque<Seat> = VecDeque::new();
947 while let Some(mut seat) = self.queued.seats.pop_front() {
948 if seated.len() >= self.max_seats {
949 waiting.push_back(seat);
950 continue;
951 }
952 let lane = if taken.contains(&seat.lane) {
953 match self.lanes.take() {
954 Some(lane) => match self.state.reset_lanes(&[lane]) {
955 Ok(()) => {
956 borrowed.push(lane);
957 Some(lane)
958 }
959 Err(_) => {
960 let _ = self.lanes.give_back(lane);
961 None
962 }
963 },
964 None => None,
965 }
966 } else {
967 Some(seat.lane)
968 };
969 match lane {
970 Some(lane) => {
971 seat.lane = lane;
972 taken.push(lane);
973 seated.push(seat);
974 }
975 None => waiting.push_back(seat),
976 }
977 }
978 self.queued.seats = waiting;
979 (seated, borrowed)
980 }
981
982 fn run_turn(&mut self, seated: &[Seat]) -> TractResult<Vec<TVec<TValue>>> {
986 let seating = self.lanes.seat(seated.iter().map(|seat| seat.lane))?;
987 let mut batched: TVec<TValue> = tvec!();
988 for (ix, is_batched) in self.batch_in.iter().enumerate() {
989 if *is_batched {
990 let seats: TVec<&Tensor> = seated.iter().map(|seat| &*seat.inputs[ix]).collect();
991 for (seat, input) in seats.iter().enumerate() {
992 ensure!(
993 input.rank() > 0 && input.shape()[0] == 1,
994 "Seat {seat} feeds {:?} of input {ix}, which a turn seats one at a time",
995 input.shape()
996 );
997 }
998 batched.push(Tensor::stack_tensors(0, &seats)?.into_tvalue());
999 } else {
1000 let shared = &seated[0].inputs[ix];
1001 for (seat, turn) in seated.iter().enumerate().skip(1) {
1002 ensure!(
1003 turn.inputs[ix] == *shared,
1004 "Input {ix} carries no batch axis, so one value of it serves the whole \
1005 turn, and seats 0 and {seat} feed it different ones"
1006 );
1007 }
1008 batched.push(shared.clone());
1009 }
1010 }
1011 self.state.seat(seating)?;
1012 let outputs = self.state.run(batched)?;
1013 let mut per_seat: Vec<TVec<TValue>> = seated.iter().map(|_| tvec!()).collect();
1014 for (ix, output) in outputs.into_iter().enumerate() {
1015 if self.batch_out.get(ix).copied().unwrap_or(false) {
1016 ensure!(
1017 output.shape()[0] == seated.len(),
1018 "The turn fills {} seats, output {ix} carries {:?}",
1019 seated.len(),
1020 output.shape()
1021 );
1022 for (seat, outputs) in per_seat.iter_mut().enumerate() {
1023 outputs.push(output.slice(0, seat, seat + 1)?.into_tvalue());
1024 }
1025 } else {
1026 for outputs in per_seat.iter_mut() {
1027 outputs.push(output.clone());
1028 }
1029 }
1030 }
1031 Ok(per_seat)
1032 }
1033}
1034
1035#[cfg(all(test, not(target_family = "wasm")))]
1037mod laned_test {
1038 use super::*;
1039 use crate::ops::math::{add, mul};
1040
1041 fn doubler(max_lanes: usize) -> TractResult<LanedRunnable> {
1044 let mut model = TypedModel::default();
1045 model.set_name("doubler");
1046 let batch = model.symbols.sym("B");
1047 let input = model.add_source("input", f32::fact(dims!(batch, 3)))?;
1048 let two = model.add_const("two", tensor2(&[[2f32]]))?;
1049 let doubled = model.wire_node("doubled", mul(), &[input, two])?;
1050 model.select_output_outlets(&doubled)?;
1051 let inner = DefaultRuntime.prepare(model)?;
1052 LanedRunnable::wrap(inner.into(), max_lanes)
1053 }
1054
1055 fn turn(handle: &mut Box<dyn State>, stream: usize, turn: usize) -> TractResult<()> {
1056 let input = tensor2(&[[stream as f32, turn as f32, 1.]]);
1057 let output = handle.run(tvec!(input.into_tvalue()))?;
1058 assert_eq!(&*output[0], &tensor2(&[[2. * stream as f32, 2. * turn as f32, 2.]]));
1059 Ok(())
1060 }
1061
1062 static LINGER: Mutex<()> = Mutex::new(());
1065
1066 fn spawn_once_free(runnable: &LanedRunnable) -> TractResult<Box<dyn State>> {
1069 for _ in 0..100 {
1070 if let Ok(handle) = runnable.spawn() {
1071 return Ok(handle);
1072 }
1073 std::thread::sleep(Duration::from_millis(10));
1074 }
1075 runnable.spawn()
1076 }
1077
1078 #[test]
1079 fn one_stream_at_a_time() -> TractResult<()> {
1080 let runnable = doubler(2)?;
1081 let mut handle = runnable.spawn()?;
1082 for t in 0..4 {
1083 turn(&mut handle, 0, t)?;
1084 }
1085 Ok(())
1086 }
1087
1088 fn reported(
1091 recorder: &metrics_util::debugging::DebuggingRecorder,
1092 ) -> HashMap<String, metrics_util::debugging::DebugValue> {
1093 recorder
1094 .snapshotter()
1095 .snapshot()
1096 .into_vec()
1097 .into_iter()
1098 .filter(|(key, ..)| {
1099 key.key().labels().any(|l| l.key() == "runnable" && l.value() == "doubler")
1100 })
1101 .map(|(key, _, _, value)| (key.key().name().to_string(), value))
1102 .collect()
1103 }
1104
1105 #[test]
1106 fn the_worker_reports_what_it_serves() -> TractResult<()> {
1107 use metrics_util::debugging::{DebugValue, DebuggingRecorder};
1108 let recorder = DebuggingRecorder::new();
1109 let runnable = metrics::with_local_recorder(&recorder, || doubler(4))?;
1110 let mut handles: Vec<Box<dyn State>> =
1111 (0..2).map(|_| runnable.spawn()).collect::<TractResult<_>>()?;
1112 for t in 0..3 {
1113 for (stream, handle) in handles.iter_mut().enumerate() {
1114 turn(handle, stream, t)?;
1115 }
1116 }
1117 let (turns, seats) = runnable.turns_and_seats();
1118 let reported = reported(&recorder);
1119 let samples = |name: &str| match &reported[name] {
1120 DebugValue::Histogram(samples) => samples.iter().map(|s| s.0).collect::<Vec<f64>>(),
1121 other => panic!("{name} is {other:?}"),
1122 };
1123 assert_eq!(reported["tract_autobatch_turns_total"], DebugValue::Counter(turns));
1124 assert_eq!(reported["tract_autobatch_seats_total"], DebugValue::Counter(seats));
1125 assert_eq!(reported["tract_autobatch_lanes_taken"], DebugValue::Gauge(2f64.into()));
1126 let per_turn = samples("tract_autobatch_seats_per_turn");
1127 assert_eq!(per_turn.len() as u64, turns);
1128 assert_eq!(per_turn.iter().sum::<f64>() as u64, seats);
1129 assert_eq!(samples("tract_autobatch_queue_wait_seconds").len() as u64, seats);
1130 let runs = samples("tract_autobatch_turn_seconds");
1131 assert_eq!(runs.len() as u64, turns);
1132 assert!(runs.iter().all(|run| *run > 0.));
1133 Ok(())
1134 }
1135
1136 #[test]
1137 fn every_stream_gets_its_own_seat() -> TractResult<()> {
1138 let runnable = doubler(8)?;
1139 let streams: Vec<_> = (0..8)
1140 .map(|stream| {
1141 let runnable = runnable.clone();
1142 std::thread::spawn(move || -> TractResult<()> {
1143 let mut handle = runnable.spawn()?;
1144 for t in 0..32 {
1145 turn(&mut handle, stream, t)?;
1146 }
1147 Ok(())
1148 })
1149 })
1150 .collect();
1151 for stream in streams {
1152 stream.join().unwrap()?;
1153 }
1154 Ok(())
1155 }
1156
1157 #[test]
1158 fn a_turn_seats_the_streams_that_are_ready() -> TractResult<()> {
1159 let _linger = LINGER.lock().unwrap_or_else(|e| e.into_inner());
1160 TRACT_TURN_LINGER_US.set(20_000);
1161 let runnable = doubler(8);
1162 TRACT_TURN_LINGER_US.clear();
1163 let runnable = runnable?;
1164 let streams: Vec<_> = (0..8)
1165 .map(|stream| {
1166 let runnable = runnable.clone();
1167 std::thread::spawn(move || -> TractResult<()> {
1168 let mut handle = runnable.spawn()?;
1169 for t in 0..4 {
1170 turn(&mut handle, stream, t)?;
1171 }
1172 Ok(())
1173 })
1174 })
1175 .collect();
1176 for stream in streams {
1177 stream.join().unwrap()?;
1178 }
1179 let (turns, seats) = runnable.turns_and_seats();
1180 assert!(seats > turns, "{seats} seats over {turns} turns, none of them shared");
1181 Ok(())
1182 }
1183
1184 fn biased(max_lanes: usize) -> TractResult<LanedRunnable> {
1187 let mut model = TypedModel::default();
1188 let batch = model.symbols.sym("B");
1189 let input = model.add_source("input", f32::fact(dims!(batch, 3)))?;
1190 let bias = model.add_source("bias", f32::fact(dims!(1, 1)))?;
1191 let two = model.add_const("two", tensor2(&[[2f32]]))?;
1192 let doubled = model.wire_node("doubled", mul(), &[input, two])?;
1193 let biased = model.wire_node("biased", add(), &[doubled[0], bias])?;
1194 model.select_output_outlets(&biased)?;
1195 let inner = DefaultRuntime.prepare(model)?;
1196 LanedRunnable::wrap(inner.into(), max_lanes)
1197 }
1198
1199 fn biased_turns(runnable: &LanedRunnable, biases: &[f32]) -> TractResult<Vec<TractResult<()>>> {
1203 let handles: Vec<Box<dyn State>> =
1204 biases.iter().map(|_| runnable.spawn()).collect::<TractResult<_>>()?;
1205 let streams: Vec<_> = handles
1206 .into_iter()
1207 .zip(biases.iter().copied())
1208 .map(|(mut handle, bias)| {
1209 std::thread::spawn(move || -> TractResult<()> {
1210 handle.run(tvec!(
1211 tensor2(&[[1f32, 2., 3.]]).into_tvalue(),
1212 tensor2(&[[bias]]).into_tvalue()
1213 ))?;
1214 Ok(())
1215 })
1216 })
1217 .collect();
1218 Ok(streams.into_iter().map(|stream| stream.join().unwrap()).collect())
1219 }
1220
1221 #[test]
1222 fn seats_agreeing_on_a_shared_input_share_a_turn() -> TractResult<()> {
1223 let _linger = LINGER.lock().unwrap_or_else(|e| e.into_inner());
1224 TRACT_TURN_LINGER_US.set(100_000);
1225 let runnable = biased(2);
1226 TRACT_TURN_LINGER_US.clear();
1227 let runnable = runnable?;
1228 let served = biased_turns(&runnable, &[7., 7.])?;
1229 assert!(served.iter().all(|s| s.is_ok()), "{served:?}");
1230 assert_eq!(runnable.turns_and_seats(), (1, 2));
1231 Ok(())
1232 }
1233
1234 #[test]
1235 fn seats_disagreeing_on_a_shared_input_fail_the_turn() -> TractResult<()> {
1236 let _linger = LINGER.lock().unwrap_or_else(|e| e.into_inner());
1237 TRACT_TURN_LINGER_US.set(100_000);
1238 let runnable = biased(2);
1239 TRACT_TURN_LINGER_US.clear();
1240 let runnable = runnable?;
1241 let served = biased_turns(&runnable, &[7., 8.])?;
1242 assert_eq!(runnable.turns_and_seats(), (1, 2));
1243 for stream in &served {
1244 let error = format!("{:#}", stream.as_ref().unwrap_err());
1245 assert!(error.contains("seats 0 and 1 feed it different ones"), "{error}");
1246 }
1247 Ok(())
1248 }
1249
1250 fn wide_call(handle: &mut Box<dyn State>, seats: usize) -> TractResult<()> {
1253 let fed: Vec<f32> =
1254 (0..seats).flat_map(|s| [s as f32, s as f32 + 1., s as f32 + 2.]).collect();
1255 let doubled: Vec<f32> = fed.iter().map(|x| 2. * x).collect();
1256 let output = handle.run(tvec!(Tensor::from_shape(&[seats, 3], &fed)?.into_tvalue()))?;
1257 assert_eq!(&*output[0], &Tensor::from_shape(&[seats, 3], &doubled)?);
1258 Ok(())
1259 }
1260
1261 #[test]
1262 fn a_call_asks_for_as_many_seats_as_it_feeds() -> TractResult<()> {
1263 let runnable = doubler(4)?;
1264 let mut handle = runnable.spawn()?;
1265 wide_call(&mut handle, 3)?;
1266 assert_eq!(runnable.turns_and_seats(), (1, 3));
1267 Ok(())
1268 }
1269
1270 #[test]
1271 fn a_call_wider_than_the_lanes_is_split_across_turns() -> TractResult<()> {
1272 let runnable = doubler(2)?;
1273 let mut handle = runnable.spawn()?;
1274 wide_call(&mut handle, 5)?;
1275 let (turns, seats) = runnable.turns_and_seats();
1276 assert_eq!(seats, 5);
1277 assert_eq!(turns, 3, "5 seats over 2 lanes are 3 turns, got {turns}");
1278 Ok(())
1279 }
1280
1281 #[test]
1282 fn a_seat_borrows_a_free_lane_only() -> TractResult<()> {
1283 let runnable = doubler(2)?;
1284 let mut handle = runnable.spawn()?;
1285 let mut other = runnable.spawn()?;
1286 turn(&mut other, 1, 0)?;
1287 wide_call(&mut handle, 4)?;
1288 let (turns, seats) = runnable.turns_and_seats();
1289 assert_eq!(seats, 5);
1290 assert_eq!(turns, 5, "another stream holds the second lane, so each seat is a turn");
1291 Ok(())
1292 }
1293
1294 fn adder(max_lanes: usize) -> TractResult<LanedRunnable> {
1297 let mut model = TypedModel::default();
1298 let batch = model.symbols.sym("B");
1299 let left = model.add_source("left", f32::fact(dims!(batch, 3)))?;
1300 let right = model.add_source("right", f32::fact(dims!(batch, 3)))?;
1301 let sum = model.wire_node("sum", add(), &[left, right])?;
1302 model.select_output_outlets(&sum)?;
1303 let inner = DefaultRuntime.prepare(model)?;
1304 LanedRunnable::wrap(inner.into(), max_lanes)
1305 }
1306
1307 #[test]
1308 fn a_call_whose_batched_inputs_disagree_on_seats_fails() -> TractResult<()> {
1309 let runnable = adder(4)?;
1310 let mut handle = runnable.spawn()?;
1311 let error = handle
1312 .run(tvec!(
1313 tensor2(&[[1f32, 2., 3.], [4., 5., 6.]]).into_tvalue(),
1314 tensor2(&[[7f32, 8., 9.]]).into_tvalue()
1315 ))
1316 .unwrap_err();
1317 let error = format!("{error:#}");
1318 assert!(error.contains("input 1 carries 1 against 2"), "{error}");
1319 Ok(())
1320 }
1321
1322 #[test]
1323 fn a_batch_off_axis_zero_says_which_axis_it_sits_on() -> TractResult<()> {
1324 let mut model = TypedModel::default();
1325 let batch = model.symbols.sym("B");
1326 let input = model.add_source("input", f32::fact(dims!(3, batch)))?;
1327 let two = model.add_const("two", tensor2(&[[2f32]]))?;
1328 let doubled = model.wire_node("doubled", mul(), &[input, two])?;
1329 model.select_output_outlets(&doubled)?;
1330 let inner = DefaultRuntime.prepare(model)?;
1331 let error = format!("{:#}", LanedRunnable::wrap(inner.into(), 2).unwrap_err());
1332 assert!(error.contains("input 0 carries B on axis 1"), "{error}");
1333 assert!(error.contains("Batchify"), "{error}");
1334 Ok(())
1335 }
1336
1337 #[test]
1338 fn two_batch_symbols_are_named() -> TractResult<()> {
1339 let mut model = TypedModel::default();
1340 let left_batch = model.symbols.sym("L");
1341 let right_batch = model.symbols.sym("R");
1342 let left = model.add_source("left", f32::fact(dims!(left_batch, 3)))?;
1343 let right = model.add_source("right", f32::fact(dims!(right_batch, 3)))?;
1344 let sum = model.wire_node("sum", add(), &[left, right])?;
1345 model.select_output_outlets(&sum)?;
1346 let inner = DefaultRuntime.prepare(model)?;
1347 let error = format!("{:#}", LanedRunnable::wrap(inner.into(), 2).unwrap_err());
1348 assert!(error.contains("L and R"), "{error}");
1349 Ok(())
1350 }
1351
1352 #[test]
1353 fn a_dropped_stream_gives_its_lane_back() -> TractResult<()> {
1354 let runnable = doubler(1)?;
1355 let mut handle = runnable.spawn()?;
1356 turn(&mut handle, 0, 0)?;
1357 assert!(runnable.spawn().is_err());
1358 let clone = dyn_clone::clone_box(&*handle);
1359 drop(handle);
1360 assert!(runnable.spawn().is_err());
1361 drop(clone);
1362 let mut handle = spawn_once_free(&runnable)?;
1363 turn(&mut handle, 1, 0)?;
1364 Ok(())
1365 }
1366}
1367
1368#[cfg(test)]
1369mod lane_table_test {
1370 use super::*;
1371
1372 #[test]
1373 fn takes_the_lowest_free_lane() -> TractResult<()> {
1374 let mut table = LaneTable::new(3)?;
1375 assert_eq!(table.take(), Some(LaneId(0)));
1376 assert_eq!(table.take(), Some(LaneId(1)));
1377 table.give_back(LaneId(0))?;
1378 assert_eq!(table.take(), Some(LaneId(0)));
1379 assert_eq!(table.taken(), 2);
1380 Ok(())
1381 }
1382
1383 #[test]
1384 fn runs_out_of_lanes() -> TractResult<()> {
1385 let mut table = LaneTable::new(1)?;
1386 assert_eq!(table.take(), Some(LaneId(0)));
1387 assert_eq!(table.take(), None);
1388 Ok(())
1389 }
1390
1391 #[test]
1392 fn gives_back_a_taken_lane_only() -> TractResult<()> {
1393 let mut table = LaneTable::new(2)?;
1394 assert!(table.give_back(LaneId(0)).is_err());
1395 table.take();
1396 table.give_back(LaneId(0))?;
1397 assert!(table.give_back(LaneId(0)).is_err());
1398 assert!(table.give_back(LaneId(7)).is_err());
1399 Ok(())
1400 }
1401
1402 #[test]
1403 fn seats_taken_lanes_in_order() -> TractResult<()> {
1404 let mut table = LaneTable::new(4)?;
1405 table.take();
1406 table.take();
1407 table.take();
1408 table.give_back(LaneId(1))?;
1409 let seating = table.seat([LaneId(2), LaneId(0)])?;
1410 assert_eq!(seating.max_lanes(), 4);
1411 assert_eq!(seating.occupancy(), 2);
1412 assert_eq!(seating.address(0), (Some(0), Some(2)));
1413 assert_eq!(seating.address(1), (Some(1), Some(0)));
1414 assert!(table.seat([LaneId(0), LaneId(1)]).is_err());
1415 assert!(table.seat([LaneId(0), LaneId(0)]).is_err());
1416 Ok(())
1417 }
1418}