use std::collections::{HashMap, VecDeque};
use std::fmt::Debug;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::{Receiver, Sender, channel};
use std::thread;
use std::time::{Duration, Instant};
use metrics::{Counter, Gauge, Histogram, Label, Unit};
use crate::internal::*;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LaneTable {
taken: Vec<bool>,
}
impl LaneTable {
pub fn new(max_lanes: usize) -> TractResult<LaneTable> {
ensure!(max_lanes > 0, "A laned state needs at least one lane");
Ok(LaneTable { taken: vec![false; max_lanes] })
}
pub fn max_lanes(&self) -> usize {
self.taken.len()
}
pub fn taken(&self) -> usize {
self.taken.iter().filter(|t| **t).count()
}
pub fn take(&mut self) -> Option<LaneId> {
let lane = self.taken.iter().position(|t| !t)?;
self.taken[lane] = true;
Some(LaneId(lane))
}
pub fn give_back(&mut self, lane: LaneId) -> TractResult<()> {
ensure!(self.is_taken(lane), "Lane {} is not taken, so it can not be given back", lane.0);
self.taken[lane.0] = false;
Ok(())
}
pub fn is_taken(&self, lane: LaneId) -> bool {
self.taken.get(lane.0).copied().unwrap_or(false)
}
pub fn seat(&self, lanes: impl IntoIterator<Item = LaneId>) -> TractResult<Seating> {
let lanes: Vec<LaneId> = lanes.into_iter().collect();
for lane in &lanes {
ensure!(self.is_taken(*lane), "Seating lane {}, which no stream took", lane.0);
}
Seating::new(self.max_lanes(), lanes)
}
}
crate::declare_knob!(
scopable TRACT_MAX_SEATS,
usize,
256,
"Most streams a laned runtime serves in one turn, clamped to the state's lanes."
);
crate::declare_knob!(
scopable TRACT_AUTOBATCH_LANES,
Option<usize>,
None,
"Sessions a deployment config's subnet is autobatched for; unset, it is not autobatched."
);
crate::declare_knob!(
scopable TRACT_TURN_LINGER_US,
usize,
0,
"How long a laned runtime waits for more streams once one is ready to run."
);
#[derive(Clone)]
pub struct LanedRunnable {
shared: Arc<Shared>,
}
struct Shared {
requests: Mutex<Option<Sender<Request>>>,
worker: Mutex<Option<thread::JoinHandle<()>>>,
inner: Arc<dyn Runnable>,
model: Option<Arc<TypedModel>>,
plan: Option<Arc<TypedSimplePlan>>,
batch: Symbol,
max_lanes: usize,
counts: Arc<Counts>,
}
impl Drop for Shared {
fn drop(&mut self) {
if let Ok(mut requests) = self.requests.lock() {
requests.take();
}
let worker = self.worker.lock().ok().and_then(|mut worker| worker.take());
if let Some(worker) = worker {
let _ = worker.join();
}
}
}
#[derive(Debug, Default)]
struct Counts {
turns: AtomicU64,
seats: AtomicU64,
}
struct Meters {
turns: Counter,
seats: Counter,
seats_per_turn: Histogram,
lanes_taken: Gauge,
queue_wait: Histogram,
turn: Histogram,
}
impl Meters {
fn new(name: Option<&str>) -> Meters {
metrics::describe_counter!(
"tract_autobatch_turns_total",
"Turns an autobatched runnable ran"
);
metrics::describe_counter!(
"tract_autobatch_seats_total",
"Seats an autobatched runnable filled over its turns"
);
metrics::describe_histogram!(
"tract_autobatch_seats_per_turn",
"Seats one turn of an autobatched runnable filled"
);
metrics::describe_gauge!(
"tract_autobatch_lanes_taken",
"Lanes held by the streams of an autobatched runnable"
);
metrics::describe_histogram!(
"tract_autobatch_queue_wait_seconds",
Unit::Seconds,
"From a seat queued to the start of the turn seating it, linger included"
);
metrics::describe_histogram!(
"tract_autobatch_turn_seconds",
Unit::Seconds,
"One turn of an autobatched runnable, stacking and slicing included"
);
let labels: Vec<Label> =
name.map(|name| Label::new("runnable", name.to_string())).into_iter().collect();
Meters {
turns: metrics::counter!("tract_autobatch_turns_total", labels.clone()),
seats: metrics::counter!("tract_autobatch_seats_total", labels.clone()),
seats_per_turn: metrics::histogram!("tract_autobatch_seats_per_turn", labels.clone()),
lanes_taken: metrics::gauge!("tract_autobatch_lanes_taken", labels.clone()),
queue_wait: metrics::histogram!("tract_autobatch_queue_wait_seconds", labels.clone()),
turn: metrics::histogram!("tract_autobatch_turn_seconds", labels),
}
}
}
impl LanedRunnable {
pub fn wrap(inner: Arc<dyn Runnable>, max_lanes: usize) -> TractResult<LanedRunnable> {
let model = inner.typed_model().cloned();
let plan = inner.typed_plan().cloned();
let mut symbols: Vec<Symbol> = vec![];
let mut facts: Vec<(String, TypedFact)> = vec![];
let mut batch_in: Vec<bool> = vec![];
for ix in 0..inner.input_count() {
let fact = inner.input_fact(ix)?;
let symbol = batch_symbol(fact);
batch_in.push(symbol.is_some());
symbols.extend(symbol);
facts.push((format!("input {ix}"), fact.clone()));
}
let mut batch_out: Vec<bool> = vec![];
for ix in 0..inner.output_count() {
let fact = inner.output_fact(ix)?;
let symbol = batch_symbol(fact);
batch_out.push(symbol.is_some());
symbols.extend(symbol);
facts.push((format!("output {ix}"), fact.clone()));
}
symbols.sort();
symbols.dedup();
if symbols.len() != 1 {
bail!(off_axis_zero(&facts, &symbols));
}
ensure!(batch_out.iter().any(|b| *b), "A laned model must batch one output at least");
let batch = symbols.remove(0);
let counts = Arc::new(Counts::default());
let meters = Meters::new(model.as_ref().and_then(|model| model.name()));
let max_seats = TRACT_MAX_SEATS.get().min(max_lanes);
let linger = Duration::from_micros(TRACT_TURN_LINGER_US.get() as u64);
let lanes = LaneTable::new(max_lanes)?;
let (requests, queue) = channel::<Request>();
let (spawned, ready) = channel::<TractResult<()>>();
let worker_counts = counts.clone();
let worker_inner = inner.clone();
let worker_thread = thread::Builder::new().name("tract-lanes".into()).spawn(move || {
let state = match worker_inner.spawn().and_then(|mut state| {
let every: Vec<LaneId> = (0..max_lanes).map(LaneId).collect();
state.reset_lanes(&every).context("Preparing a laned model")?;
Ok(state)
}) {
Ok(state) => {
let _ = spawned.send(Ok(()));
state
}
Err(e) => {
let _ = spawned.send(Err(e));
return;
}
};
Worker {
state,
lanes,
queued: Queue::default(),
batch_in,
batch_out,
max_seats,
linger,
counts: worker_counts,
meters,
}
.work(queue);
})?;
ready.recv().map_err(|_| format_err!("The laned worker died spawning the state"))??;
Ok(LanedRunnable {
shared: Arc::new(Shared {
requests: Mutex::new(Some(requests)),
worker: Mutex::new(Some(worker_thread)),
inner,
model,
plan,
batch,
max_lanes,
counts,
}),
})
}
pub fn max_lanes(&self) -> usize {
self.shared.max_lanes
}
pub fn inner(&self) -> &Arc<dyn Runnable> {
&self.shared.inner
}
pub fn batch_symbol(&self) -> &Symbol {
&self.shared.batch
}
pub fn turns_and_seats(&self) -> (u64, u64) {
(
self.shared.counts.turns.load(Ordering::Relaxed),
self.shared.counts.seats.load(Ordering::Relaxed),
)
}
fn request(&self) -> TractResult<Sender<Request>> {
self.shared
.requests
.lock()
.map_err(|_| format_err!("Poisoned laned sender"))?
.clone()
.context("The laned runnable is gone")
}
}
fn batch_symbol(fact: &TypedFact) -> Option<Symbol> {
match fact.shape.dims().first() {
Some(TDim::Sym(sym)) => Some(sym.clone()),
_ => None,
}
}
fn off_axis_zero(facts: &[(String, TypedFact)], symbols: &[Symbol]) -> TractError {
let elsewhere: Vec<String> = facts
.iter()
.filter_map(|(what, fact)| {
let dims = fact.shape.dims();
let axis = dims.iter().skip(1).position(|dim| matches!(dim, TDim::Sym(_)))? + 1;
Some(format!("{what} carries {} on axis {axis}", dims[axis]))
})
.collect();
let elsewhere = elsewhere.join(", ");
if symbols.is_empty() {
let found = if elsewhere.is_empty() {
"none of them carries a symbol at all".to_string()
} else {
elsewhere
};
format_err!(
"A laned model carries its batch symbol on axis 0 of the inputs and outputs it \
batches, and {found}. Axis 0 is where a seat's values are a contiguous run, so move \
the batch there as a graph edit -- Batchify, or an AddAxis/MoveAxis the optimiser \
can absorb -- rather than have every turn transpose"
)
} else {
let symbols: Vec<String> = symbols.iter().map(|s| s.to_string()).collect();
format_err!(
"A laned model carries one batch symbol on axis 0, this one carries {}. One symbol \
has to stand for the whole turn's occupancy, so share it across the batched inputs \
and outputs in the export",
symbols.join(" and ")
)
}
}
impl Debug for LanedRunnable {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "LanedRunnable({} lanes)", self.shared.max_lanes)
}
}
impl Runnable for LanedRunnable {
fn spawn(&self) -> TractResult<Box<dyn State>> {
let requests = self.request()?;
let (taken, lane) = channel();
requests
.send(Request::Spawn(taken))
.map_err(|_| format_err!("The laned worker is gone"))?;
let lane = lane.recv().map_err(|_| format_err!("The laned worker dropped a lane"))??;
Ok(Box::new(LanedStateHandle {
lease: Arc::new(Lease { lane, requests }),
runnable: self.clone(),
}))
}
fn typed_plan(&self) -> Option<&Arc<TypedSimplePlan>> {
self.shared.plan.as_ref()
}
fn typed_model(&self) -> Option<&Arc<TypedModel>> {
self.shared.model.as_ref()
}
}
#[derive(Clone, Debug)]
pub struct LanedStateHandle {
lease: Arc<Lease>,
runnable: LanedRunnable,
}
#[derive(Debug)]
struct Lease {
lane: LaneId,
requests: Sender<Request>,
}
impl Drop for Lease {
fn drop(&mut self) {
let _ = self.requests.send(Request::Drop(self.lane));
}
}
impl State for LanedStateHandle {
fn run(&mut self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
let (done, outputs) = channel();
self.lease
.requests
.send(Request::Call(Call { leased: self.lease.lane, inputs, done }))
.map_err(|_| format_err!("The laned worker is gone"))?;
outputs.recv().map_err(|_| format_err!("The laned worker dropped a turn"))?
}
fn runnable(&self) -> &dyn Runnable {
&self.runnable
}
}
enum Request {
Spawn(Sender<TractResult<LaneId>>),
Call(Call),
Drop(LaneId),
}
struct Call {
leased: LaneId,
inputs: TVec<TValue>,
done: Sender<TractResult<TVec<TValue>>>,
}
struct Seat {
call: u64,
lane: LaneId,
ix: usize,
inputs: TVec<TValue>,
queued: Instant,
}
struct Completer {
served: Vec<Option<TVec<TValue>>>,
done: Sender<TractResult<TVec<TValue>>>,
}
impl Call {
fn explode(&self, batch_in: &[bool]) -> TractResult<Vec<TVec<TValue>>> {
ensure!(
self.inputs.len() == batch_in.len(),
"A call feeds {} inputs, the model takes {}",
self.inputs.len(),
batch_in.len()
);
let mut seats: Option<usize> = None;
for (ix, input) in self.inputs.iter().enumerate().filter(|(ix, _)| batch_in[*ix]) {
ensure!(
input.rank() > 0,
"Input {ix} carries the batch axis, so it can not be a scalar"
);
let asked = input.shape()[0];
ensure!(
seats.is_none_or(|seats| seats == asked),
"A call asks for as many seats as its batched inputs carry, and input {ix} carries \
{asked} against {} before it",
seats.unwrap_or(0)
);
seats = Some(asked);
}
let seats = seats.unwrap_or(1);
ensure!(seats > 0, "A call asks for one seat at least");
if seats == 1 {
return Ok(vec![self.inputs.clone()]);
}
(0..seats)
.map(|seat| {
self.inputs
.iter()
.zip(batch_in)
.map(|(input, is_batched)| {
if *is_batched {
Ok(input.slice(0, seat, seat + 1)?.into_tvalue())
} else {
Ok(input.clone())
}
})
.collect()
})
.collect()
}
}
impl Completer {
fn is_full(&self) -> bool {
self.served.iter().all(Option::is_some)
}
fn answer(self, batch_out: &[bool]) {
let served: Vec<TVec<TValue>> =
self.served.into_iter().map(|outputs| outputs.unwrap()).collect();
let answer = if served.len() == 1 {
Ok(served.into_iter().next().unwrap())
} else {
let first = &served[0];
(0..first.len())
.map(|ix| {
if batch_out.get(ix).copied().unwrap_or(false) {
let seats: TVec<&Tensor> =
served.iter().map(|outputs| &*outputs[ix]).collect();
Ok(Tensor::stack_tensors(0, &seats)?.into_tvalue())
} else {
Ok(first[ix].clone())
}
})
.collect()
};
let _ = self.done.send(answer);
}
}
#[derive(Default)]
struct Queue {
seats: VecDeque<Seat>,
completers: HashMap<u64, Completer>,
calls: u64,
}
impl Queue {
fn push(&mut self, call: Call, batch_in: &[bool]) {
let id = self.calls;
self.calls += 1;
let queued = Instant::now();
match call.explode(batch_in) {
Ok(seats) => {
self.completers
.insert(id, Completer { served: vec![None; seats.len()], done: call.done });
for (ix, inputs) in seats.into_iter().enumerate() {
self.seats.push_back(Seat { call: id, lane: call.leased, ix, inputs, queued });
}
}
Err(e) => {
let _ = call.done.send(Err(e));
}
}
}
fn serve(&mut self, seat: Seat, outputs: TractResult<TVec<TValue>>, batch_out: &[bool]) {
if !self.completers.contains_key(&seat.call) {
return;
}
let outputs = match outputs {
Ok(outputs) => outputs,
Err(e) => {
let completer = self.completers.remove(&seat.call).unwrap();
self.seats.retain(|queued| queued.call != seat.call);
let _ = completer.done.send(Err(e));
return;
}
};
let completer = self.completers.get_mut(&seat.call).unwrap();
completer.served[seat.ix] = Some(outputs);
if completer.is_full() {
self.completers.remove(&seat.call).unwrap().answer(batch_out);
}
}
}
struct Worker {
state: Box<dyn State>,
lanes: LaneTable,
queued: Queue,
batch_in: Vec<bool>,
batch_out: Vec<bool>,
max_seats: usize,
linger: Duration,
counts: Arc<Counts>,
meters: Meters,
}
impl Worker {
fn work(mut self, queue: Receiver<Request>) {
loop {
while self.queued.seats.is_empty() {
match queue.recv() {
Ok(request) => self.serve(request),
Err(_) => return,
}
if !self.queued.seats.is_empty() && !self.linger.is_zero() {
thread::sleep(self.linger);
}
}
while let Ok(request) = queue.try_recv() {
self.serve(request);
}
let (seated, borrowed) = self.fill();
if seated.is_empty() {
continue;
}
self.counts.turns.fetch_add(1, Ordering::Relaxed);
self.counts.seats.fetch_add(seated.len() as u64, Ordering::Relaxed);
let started = Instant::now();
self.meters.turns.increment(1);
self.meters.seats.increment(seated.len() as u64);
self.meters.seats_per_turn.record(seated.len() as f64);
for seat in &seated {
self.meters.queue_wait.record(started.duration_since(seat.queued));
}
let served = self.run_turn(&seated);
self.meters.turn.record(started.elapsed());
for lane in borrowed {
let _ = self.lanes.give_back(lane);
}
match served {
Ok(per_seat) => {
for (seat, outputs) in seated.into_iter().zip(per_seat) {
self.queued.serve(seat, Ok(outputs), &self.batch_out);
}
}
Err(e) => {
let e = format!("{e:#}");
for seat in seated {
self.queued.serve(
seat,
Err(format_err!("Laned turn failed: {e}")),
&self.batch_out,
);
}
}
}
}
}
fn serve(&mut self, request: Request) {
match request {
Request::Spawn(taken) => {
let lane = match self.lanes.take() {
None => Err(format_err!(
"Every one of the {} lanes is taken",
self.lanes.max_lanes()
)),
Some(lane) => match self.state.reset_lanes(&[lane]) {
Ok(()) => Ok(lane),
Err(e) => {
let _ = self.lanes.give_back(lane);
Err(e)
}
},
};
self.meters.lanes_taken.set(self.lanes.taken() as f64);
let _ = taken.send(lane);
}
Request::Call(call) => self.queued.push(call, &self.batch_in),
Request::Drop(lane) => {
let _ = self.lanes.give_back(lane);
self.meters.lanes_taken.set(self.lanes.taken() as f64);
}
}
}
fn fill(&mut self) -> (Vec<Seat>, Vec<LaneId>) {
let mut seated: Vec<Seat> = vec![];
let mut taken: Vec<LaneId> = vec![];
let mut borrowed: Vec<LaneId> = vec![];
let mut waiting: VecDeque<Seat> = VecDeque::new();
while let Some(mut seat) = self.queued.seats.pop_front() {
if seated.len() >= self.max_seats {
waiting.push_back(seat);
continue;
}
let lane = if taken.contains(&seat.lane) {
match self.lanes.take() {
Some(lane) => match self.state.reset_lanes(&[lane]) {
Ok(()) => {
borrowed.push(lane);
Some(lane)
}
Err(_) => {
let _ = self.lanes.give_back(lane);
None
}
},
None => None,
}
} else {
Some(seat.lane)
};
match lane {
Some(lane) => {
seat.lane = lane;
taken.push(lane);
seated.push(seat);
}
None => waiting.push_back(seat),
}
}
self.queued.seats = waiting;
(seated, borrowed)
}
fn run_turn(&mut self, seated: &[Seat]) -> TractResult<Vec<TVec<TValue>>> {
let seating = self.lanes.seat(seated.iter().map(|seat| seat.lane))?;
let mut batched: TVec<TValue> = tvec!();
for (ix, is_batched) in self.batch_in.iter().enumerate() {
if *is_batched {
let seats: TVec<&Tensor> = seated.iter().map(|seat| &*seat.inputs[ix]).collect();
for (seat, input) in seats.iter().enumerate() {
ensure!(
input.rank() > 0 && input.shape()[0] == 1,
"Seat {seat} feeds {:?} of input {ix}, which a turn seats one at a time",
input.shape()
);
}
batched.push(Tensor::stack_tensors(0, &seats)?.into_tvalue());
} else {
let shared = &seated[0].inputs[ix];
for (seat, turn) in seated.iter().enumerate().skip(1) {
ensure!(
turn.inputs[ix] == *shared,
"Input {ix} carries no batch axis, so one value of it serves the whole \
turn, and seats 0 and {seat} feed it different ones"
);
}
batched.push(shared.clone());
}
}
self.state.seat(seating)?;
let outputs = self.state.run(batched)?;
let mut per_seat: Vec<TVec<TValue>> = seated.iter().map(|_| tvec!()).collect();
for (ix, output) in outputs.into_iter().enumerate() {
if self.batch_out.get(ix).copied().unwrap_or(false) {
ensure!(
output.shape()[0] == seated.len(),
"The turn fills {} seats, output {ix} carries {:?}",
seated.len(),
output.shape()
);
for (seat, outputs) in per_seat.iter_mut().enumerate() {
outputs.push(output.slice(0, seat, seat + 1)?.into_tvalue());
}
} else {
for outputs in per_seat.iter_mut() {
outputs.push(output.clone());
}
}
}
Ok(per_seat)
}
}
#[cfg(all(test, not(target_family = "wasm")))]
mod laned_test {
use super::*;
use crate::ops::math::{add, mul};
fn doubler(max_lanes: usize) -> TractResult<LanedRunnable> {
let mut model = TypedModel::default();
model.set_name("doubler");
let batch = model.symbols.sym("B");
let input = model.add_source("input", f32::fact(dims!(batch, 3)))?;
let two = model.add_const("two", tensor2(&[[2f32]]))?;
let doubled = model.wire_node("doubled", mul(), &[input, two])?;
model.select_output_outlets(&doubled)?;
let inner = DefaultRuntime.prepare(model)?;
LanedRunnable::wrap(inner.into(), max_lanes)
}
fn turn(handle: &mut Box<dyn State>, stream: usize, turn: usize) -> TractResult<()> {
let input = tensor2(&[[stream as f32, turn as f32, 1.]]);
let output = handle.run(tvec!(input.into_tvalue()))?;
assert_eq!(&*output[0], &tensor2(&[[2. * stream as f32, 2. * turn as f32, 2.]]));
Ok(())
}
static LINGER: Mutex<()> = Mutex::new(());
fn spawn_once_free(runnable: &LanedRunnable) -> TractResult<Box<dyn State>> {
for _ in 0..100 {
if let Ok(handle) = runnable.spawn() {
return Ok(handle);
}
std::thread::sleep(Duration::from_millis(10));
}
runnable.spawn()
}
#[test]
fn one_stream_at_a_time() -> TractResult<()> {
let runnable = doubler(2)?;
let mut handle = runnable.spawn()?;
for t in 0..4 {
turn(&mut handle, 0, t)?;
}
Ok(())
}
fn reported(
recorder: &metrics_util::debugging::DebuggingRecorder,
) -> HashMap<String, metrics_util::debugging::DebugValue> {
recorder
.snapshotter()
.snapshot()
.into_vec()
.into_iter()
.filter(|(key, ..)| {
key.key().labels().any(|l| l.key() == "runnable" && l.value() == "doubler")
})
.map(|(key, _, _, value)| (key.key().name().to_string(), value))
.collect()
}
#[test]
fn the_worker_reports_what_it_serves() -> TractResult<()> {
use metrics_util::debugging::{DebugValue, DebuggingRecorder};
let recorder = DebuggingRecorder::new();
let runnable = metrics::with_local_recorder(&recorder, || doubler(4))?;
let mut handles: Vec<Box<dyn State>> =
(0..2).map(|_| runnable.spawn()).collect::<TractResult<_>>()?;
for t in 0..3 {
for (stream, handle) in handles.iter_mut().enumerate() {
turn(handle, stream, t)?;
}
}
let (turns, seats) = runnable.turns_and_seats();
let reported = reported(&recorder);
let samples = |name: &str| match &reported[name] {
DebugValue::Histogram(samples) => samples.iter().map(|s| s.0).collect::<Vec<f64>>(),
other => panic!("{name} is {other:?}"),
};
assert_eq!(reported["tract_autobatch_turns_total"], DebugValue::Counter(turns));
assert_eq!(reported["tract_autobatch_seats_total"], DebugValue::Counter(seats));
assert_eq!(reported["tract_autobatch_lanes_taken"], DebugValue::Gauge(2f64.into()));
let per_turn = samples("tract_autobatch_seats_per_turn");
assert_eq!(per_turn.len() as u64, turns);
assert_eq!(per_turn.iter().sum::<f64>() as u64, seats);
assert_eq!(samples("tract_autobatch_queue_wait_seconds").len() as u64, seats);
let runs = samples("tract_autobatch_turn_seconds");
assert_eq!(runs.len() as u64, turns);
assert!(runs.iter().all(|run| *run > 0.));
Ok(())
}
#[test]
fn every_stream_gets_its_own_seat() -> TractResult<()> {
let runnable = doubler(8)?;
let streams: Vec<_> = (0..8)
.map(|stream| {
let runnable = runnable.clone();
std::thread::spawn(move || -> TractResult<()> {
let mut handle = runnable.spawn()?;
for t in 0..32 {
turn(&mut handle, stream, t)?;
}
Ok(())
})
})
.collect();
for stream in streams {
stream.join().unwrap()?;
}
Ok(())
}
#[test]
fn a_turn_seats_the_streams_that_are_ready() -> TractResult<()> {
let _linger = LINGER.lock().unwrap_or_else(|e| e.into_inner());
TRACT_TURN_LINGER_US.set(20_000);
let runnable = doubler(8);
TRACT_TURN_LINGER_US.clear();
let runnable = runnable?;
let streams: Vec<_> = (0..8)
.map(|stream| {
let runnable = runnable.clone();
std::thread::spawn(move || -> TractResult<()> {
let mut handle = runnable.spawn()?;
for t in 0..4 {
turn(&mut handle, stream, t)?;
}
Ok(())
})
})
.collect();
for stream in streams {
stream.join().unwrap()?;
}
let (turns, seats) = runnable.turns_and_seats();
assert!(seats > turns, "{seats} seats over {turns} turns, none of them shared");
Ok(())
}
fn biased(max_lanes: usize) -> TractResult<LanedRunnable> {
let mut model = TypedModel::default();
let batch = model.symbols.sym("B");
let input = model.add_source("input", f32::fact(dims!(batch, 3)))?;
let bias = model.add_source("bias", f32::fact(dims!(1, 1)))?;
let two = model.add_const("two", tensor2(&[[2f32]]))?;
let doubled = model.wire_node("doubled", mul(), &[input, two])?;
let biased = model.wire_node("biased", add(), &[doubled[0], bias])?;
model.select_output_outlets(&biased)?;
let inner = DefaultRuntime.prepare(model)?;
LanedRunnable::wrap(inner.into(), max_lanes)
}
fn biased_turns(runnable: &LanedRunnable, biases: &[f32]) -> TractResult<Vec<TractResult<()>>> {
let handles: Vec<Box<dyn State>> =
biases.iter().map(|_| runnable.spawn()).collect::<TractResult<_>>()?;
let streams: Vec<_> = handles
.into_iter()
.zip(biases.iter().copied())
.map(|(mut handle, bias)| {
std::thread::spawn(move || -> TractResult<()> {
handle.run(tvec!(
tensor2(&[[1f32, 2., 3.]]).into_tvalue(),
tensor2(&[[bias]]).into_tvalue()
))?;
Ok(())
})
})
.collect();
Ok(streams.into_iter().map(|stream| stream.join().unwrap()).collect())
}
#[test]
fn seats_agreeing_on_a_shared_input_share_a_turn() -> TractResult<()> {
let _linger = LINGER.lock().unwrap_or_else(|e| e.into_inner());
TRACT_TURN_LINGER_US.set(100_000);
let runnable = biased(2);
TRACT_TURN_LINGER_US.clear();
let runnable = runnable?;
let served = biased_turns(&runnable, &[7., 7.])?;
assert!(served.iter().all(|s| s.is_ok()), "{served:?}");
assert_eq!(runnable.turns_and_seats(), (1, 2));
Ok(())
}
#[test]
fn seats_disagreeing_on_a_shared_input_fail_the_turn() -> TractResult<()> {
let _linger = LINGER.lock().unwrap_or_else(|e| e.into_inner());
TRACT_TURN_LINGER_US.set(100_000);
let runnable = biased(2);
TRACT_TURN_LINGER_US.clear();
let runnable = runnable?;
let served = biased_turns(&runnable, &[7., 8.])?;
assert_eq!(runnable.turns_and_seats(), (1, 2));
for stream in &served {
let error = format!("{:#}", stream.as_ref().unwrap_err());
assert!(error.contains("seats 0 and 1 feed it different ones"), "{error}");
}
Ok(())
}
fn wide_call(handle: &mut Box<dyn State>, seats: usize) -> TractResult<()> {
let fed: Vec<f32> =
(0..seats).flat_map(|s| [s as f32, s as f32 + 1., s as f32 + 2.]).collect();
let doubled: Vec<f32> = fed.iter().map(|x| 2. * x).collect();
let output = handle.run(tvec!(Tensor::from_shape(&[seats, 3], &fed)?.into_tvalue()))?;
assert_eq!(&*output[0], &Tensor::from_shape(&[seats, 3], &doubled)?);
Ok(())
}
#[test]
fn a_call_asks_for_as_many_seats_as_it_feeds() -> TractResult<()> {
let runnable = doubler(4)?;
let mut handle = runnable.spawn()?;
wide_call(&mut handle, 3)?;
assert_eq!(runnable.turns_and_seats(), (1, 3));
Ok(())
}
#[test]
fn a_call_wider_than_the_lanes_is_split_across_turns() -> TractResult<()> {
let runnable = doubler(2)?;
let mut handle = runnable.spawn()?;
wide_call(&mut handle, 5)?;
let (turns, seats) = runnable.turns_and_seats();
assert_eq!(seats, 5);
assert_eq!(turns, 3, "5 seats over 2 lanes are 3 turns, got {turns}");
Ok(())
}
#[test]
fn a_seat_borrows_a_free_lane_only() -> TractResult<()> {
let runnable = doubler(2)?;
let mut handle = runnable.spawn()?;
let mut other = runnable.spawn()?;
turn(&mut other, 1, 0)?;
wide_call(&mut handle, 4)?;
let (turns, seats) = runnable.turns_and_seats();
assert_eq!(seats, 5);
assert_eq!(turns, 5, "another stream holds the second lane, so each seat is a turn");
Ok(())
}
fn adder(max_lanes: usize) -> TractResult<LanedRunnable> {
let mut model = TypedModel::default();
let batch = model.symbols.sym("B");
let left = model.add_source("left", f32::fact(dims!(batch, 3)))?;
let right = model.add_source("right", f32::fact(dims!(batch, 3)))?;
let sum = model.wire_node("sum", add(), &[left, right])?;
model.select_output_outlets(&sum)?;
let inner = DefaultRuntime.prepare(model)?;
LanedRunnable::wrap(inner.into(), max_lanes)
}
#[test]
fn a_call_whose_batched_inputs_disagree_on_seats_fails() -> TractResult<()> {
let runnable = adder(4)?;
let mut handle = runnable.spawn()?;
let error = handle
.run(tvec!(
tensor2(&[[1f32, 2., 3.], [4., 5., 6.]]).into_tvalue(),
tensor2(&[[7f32, 8., 9.]]).into_tvalue()
))
.unwrap_err();
let error = format!("{error:#}");
assert!(error.contains("input 1 carries 1 against 2"), "{error}");
Ok(())
}
#[test]
fn a_batch_off_axis_zero_says_which_axis_it_sits_on() -> TractResult<()> {
let mut model = TypedModel::default();
let batch = model.symbols.sym("B");
let input = model.add_source("input", f32::fact(dims!(3, batch)))?;
let two = model.add_const("two", tensor2(&[[2f32]]))?;
let doubled = model.wire_node("doubled", mul(), &[input, two])?;
model.select_output_outlets(&doubled)?;
let inner = DefaultRuntime.prepare(model)?;
let error = format!("{:#}", LanedRunnable::wrap(inner.into(), 2).unwrap_err());
assert!(error.contains("input 0 carries B on axis 1"), "{error}");
assert!(error.contains("Batchify"), "{error}");
Ok(())
}
#[test]
fn two_batch_symbols_are_named() -> TractResult<()> {
let mut model = TypedModel::default();
let left_batch = model.symbols.sym("L");
let right_batch = model.symbols.sym("R");
let left = model.add_source("left", f32::fact(dims!(left_batch, 3)))?;
let right = model.add_source("right", f32::fact(dims!(right_batch, 3)))?;
let sum = model.wire_node("sum", add(), &[left, right])?;
model.select_output_outlets(&sum)?;
let inner = DefaultRuntime.prepare(model)?;
let error = format!("{:#}", LanedRunnable::wrap(inner.into(), 2).unwrap_err());
assert!(error.contains("L and R"), "{error}");
Ok(())
}
#[test]
fn a_dropped_stream_gives_its_lane_back() -> TractResult<()> {
let runnable = doubler(1)?;
let mut handle = runnable.spawn()?;
turn(&mut handle, 0, 0)?;
assert!(runnable.spawn().is_err());
let clone = dyn_clone::clone_box(&*handle);
drop(handle);
assert!(runnable.spawn().is_err());
drop(clone);
let mut handle = spawn_once_free(&runnable)?;
turn(&mut handle, 1, 0)?;
Ok(())
}
}
#[cfg(test)]
mod lane_table_test {
use super::*;
#[test]
fn takes_the_lowest_free_lane() -> TractResult<()> {
let mut table = LaneTable::new(3)?;
assert_eq!(table.take(), Some(LaneId(0)));
assert_eq!(table.take(), Some(LaneId(1)));
table.give_back(LaneId(0))?;
assert_eq!(table.take(), Some(LaneId(0)));
assert_eq!(table.taken(), 2);
Ok(())
}
#[test]
fn runs_out_of_lanes() -> TractResult<()> {
let mut table = LaneTable::new(1)?;
assert_eq!(table.take(), Some(LaneId(0)));
assert_eq!(table.take(), None);
Ok(())
}
#[test]
fn gives_back_a_taken_lane_only() -> TractResult<()> {
let mut table = LaneTable::new(2)?;
assert!(table.give_back(LaneId(0)).is_err());
table.take();
table.give_back(LaneId(0))?;
assert!(table.give_back(LaneId(0)).is_err());
assert!(table.give_back(LaneId(7)).is_err());
Ok(())
}
#[test]
fn seats_taken_lanes_in_order() -> TractResult<()> {
let mut table = LaneTable::new(4)?;
table.take();
table.take();
table.take();
table.give_back(LaneId(1))?;
let seating = table.seat([LaneId(2), LaneId(0)])?;
assert_eq!(seating.max_lanes(), 4);
assert_eq!(seating.occupancy(), 2);
assert_eq!(seating.address(0), (Some(0), Some(2)));
assert_eq!(seating.address(1), (Some(1), Some(0)));
assert!(table.seat([LaneId(0), LaneId(1)]).is_err());
assert!(table.seat([LaneId(0), LaneId(0)]).is_err());
Ok(())
}
}