Skip to main content

tract_core/
lanes.rs

1//! Serving many streams through one state.
2//!
3//! A prepared model serves one stream at a time: `spawn()` hands out a state,
4//! and each `run()` on it is one turn of that stream. [`LanedRunnable::wrap`]
5//! makes one state serve many streams at once, by batching the turns that
6//! happen to be ready into a single `run()` on a single state. Callers see
7//! nothing of it: they still spawn a state each and run it per turn.
8//!
9//! This is the one corner of `core` where threads, queues and a promise per
10//! caller appear, so its plumbing carries more comment than the rest of the
11//! crate: the types are small, and what they mean is where the bugs are.
12//!
13//! # Two batch axes, and the words for them
14//!
15//! `doc/lexicon.md` is the reference; the four words this module lives on:
16//!
17//! - **lane** -- where one stream's state sits inside the shared state: axis 0
18//!   of a laned op state's buffers, of extent `max_lanes`, addressed through
19//!   the turn's [`Seating`].
20//! - **seat** -- a position in one turn's batch: axis 0 of the turn's input and
21//!   output tensors, of extent the turn's occupancy.
22//! - **turn** -- one `run()` on the state. It has one seating, and its seats
23//!   are the streams it serves.
24//! - **call** -- one `run()` on a [`LanedStateHandle`]. Usually one seat of one
25//!   turn, but a caller feeding several at once asks for several seats, and is
26//!   answered over as many turns as they took to seat.
27//!
28//! One lane takes at most one seat per turn: a stream's state is sequential.
29//!
30//! # The actors
31//!
32//! ```text
33//!     caller thread                     caller thread
34//!      (stream A)                        (stream B)
35//!          |                                 |
36//!   LanedStateHandle                  LanedStateHandle
37//!    lane 0, cloned                    lane 1, cloned
38//!          |                                 |
39//!          |   Request::{Spawn, Call, Drop}  |
40//!          +---------------+-----------------+
41//!                          |   one mpsc queue, cloned per handle
42//!                          v
43//!       +--------------------------------------------+
44//!       |        worker thread "tract-lanes"         |
45//!       |                                            |
46//!       |  Worker  state   the one state, laned      |
47//!       |          lanes   which of them are taken   |
48//!       |          queued  seats waiting, and the    |
49//!       |                  calls they answer         |
50//!       +--------------------------------------------+
51//!                          |
52//!                          |   one answer channel per call
53//!                          v
54//!                    back to the callers
55//! ```
56//!
57//! The worker owns the state and the [`LaneTable`] both, and is the only thread
58//! which touches either. That is not tidiness: taking a lane **resets** it,
59//! which writes the state -- device memory for a state on a GPU -- so it has to
60//! happen where the state lives. Hence a handle asks for a lane rather than
61//! taking one, and `Lease` only knows how to send it back.
62//!
63//! # The life of a handle, in requests
64//!
65//! - `Request::Spawn` -- `LanedRunnable::spawn` asks for a lane and blocks
66//!   until the worker resets one and answers, or until it answers that every
67//!   lane is taken.
68//! - `Request::Call` -- `LanedStateHandle::run` sends its inputs and a
69//!   one-shot channel, then blocks on that channel.
70//! - `Request::Drop` -- the last clone of a handle dropped sends its lane
71//!   back, and the next stream can have it.
72//!
73//! # The life of a turn
74//!
75//! 1. The worker blocks until a request arrives. If it put a seat in the queue,
76//!    the turn **lingers** for [`TRACT_TURN_LINGER_US`], so that streams whose
77//!    pulses land within a hair of each other share a turn instead of taking
78//!    one each. Zero by default.
79//! 2. Everything else pending is drained, so the turn sees every seat ready.
80//! 3. `fill` seats the head of the queue: at most [`TRACT_MAX_SEATS`] seats, at
81//!    most one per lane.
82//! 4. `run_turn` stacks the batched inputs along axis 0 in seat order, checks
83//!    that the shared ones agree across seats, publishes the seating, and runs
84//!    the state **once**.
85//! 5. Its outputs are sliced back per seat, borrowed lanes go back to the
86//!    table, and each seat is handed to its call's `Completer`, which answers
87//!    the caller once its last seat has landed.
88//!
89//! A turn that fails fails every seat of it, and a seat that fails fails its
90//! whole call and drops that call's seats still queued.
91//!
92//! # A call asking for several seats
93//!
94//! A beam decoder hands a stateless model its k hypotheses in one `run()`. Such
95//! a call **explodes** into one `Seat` per slice of its batched inputs; the
96//! queue is therefore a queue of seats, not of calls, and a call wider than the
97//! free lanes is split at the turn boundary rather than held until it fits
98//! whole -- so a wide call cannot starve the one-seat calls behind it. Its first
99//! seat sits in its caller's own lane and the rest **borrow** free lanes, which
100//! resets them, so no seat ever reads what another caller left.
101//!
102//! With `max_lanes` 4, A holding lane 0 and B lane 1, A calling for 5 seats and
103//! B for one:
104//!
105//! ```text
106//!   A: run([5, ..])                      B: run([1, ..])
107//!         |                                    |
108//!         | explode                            | explode
109//!         v                                    v
110//!      A0 A1 A2 A3 A4                          B0
111//!         |                                    |
112//!         +----------------+-------------------+
113//!                          v
114//!            queue: A0 A1 A2 A3 A4 B0
115//!
116//!   turn 1                        lane 0  lane 1  lane 2  lane 3
117//!     A0 -> its own lane          [ A0 ]  [ B0 ]  [ A1 ]  [ A2 ]
118//!     A1 -> borrows lane 2          |       |       |       |
119//!     A2 -> borrows lane 3          +-------+---+---+-------+
120//!     A3 -> nothing free, waits                 |  one run(), occupancy 4
121//!     A4 -> waits                               v
122//!     B0 -> its own lane              B answered; A has 3 of 5 seats
123//!
124//!   turn 2                        lane 0  lane 1  lane 2  lane 3
125//!     A3 -> its own lane          [ A3 ]    --    [ A4 ]    --
126//!     A4 -> borrows lane 2          |               |
127//!                                   +-------+-------+
128//!                                           |  one run(), occupancy 2
129//!                                           v
130//!                                 A's 5 seats stack back into one [5, ..]
131//! ```
132//!
133//! A one-seat call pays nothing for any of this: `explode` does not slice it and
134//! `assemble` hands its outputs straight back.
135//!
136//! # What a laned model must satisfy
137//!
138//! - **The batch axis is axis 0**, on at least one input and one output, and it
139//!   is one symbol for all of them. It is the only position where a seat's
140//!   values are a contiguous run, which is what makes stacking a memcpy and
141//!   slicing a view, and the only one whose seats are independent under a
142//!   liquid schedule. A model wanting it elsewhere gets it moved by a graph
143//!   edit, never by a per-turn transpose here.
144//! - **An input or output whose axis 0 is not a symbol is shared**: one value of
145//!   it serves the whole turn, so seats disagreeing about it fail the turn, and
146//!   a shared output is handed back to every seat.
147//! - **Axis 0 of every stateful node** is that batch axis. `wrap` walks the I/O
148//!   facts and resets every lane once, which fails a state that cannot serve
149//!   several streams; the interior is otherwise on trust.
150//!
151//! # Traps
152//!
153//! - `Seat::lane` is the lane the seat **sits in**, not the lane its caller
154//!   holds: `fill` overwrites it when it borrows, or a turn ends up seating one
155//!   lane twice.
156//! - The knobs are read at `wrap` and are process-wide, so tests which set
157//!   [`TRACT_TURN_LINGER_US`] have to serialize.
158//! - Dropping the runnable closes the queue and **joins** the worker. The worker
159//!   holds per-thread device state whose destructors must not run against
160//!   libraries already tearing themselves down.
161//!
162//! # Not here
163//!
164//! Admission is not a policy yet: `spawn` fails when every lane is taken rather
165//! than waiting for one. A caller holding a lane for the life of a handle is
166//! what makes that felt, and what a server wants of it is still open.
167
168use 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/// The lanes of one laned state: which are taken, and which of them a turn
181/// seats.
182///
183/// Plain data. Taking a lane does not touch the state's buffers, and clearing
184/// what a stream left in a lane it gave up is the table's caller's, since it
185/// writes the state -- device memory for a state on a GPU -- and must run where
186/// the state lives. So a lane handed to a new stream carries the previous one's
187/// history until that caller resets it.
188#[derive(Debug, Clone, PartialEq, Eq)]
189pub struct LaneTable {
190    /// One flag per lane, true while a stream holds it. A lane's index is its
191    /// [`LaneId`], and the length is the state's fixed lane count.
192    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    /// The extent of the lane axis of the state's per-lane buffers, fixed for
202    /// the life of the state.
203    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    /// The lowest free lane, `None` when every lane is taken -- whether that
212    /// blocks the new stream or fails it is the caller's policy. Lowest first,
213    /// so that a turn seating every lane seats a run of consecutive lanes.
214    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    /// Hand `lane` back, for [`LaneTable::take`] to give to another stream.
221    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    /// Seat `lanes`, in that order: seat `ix` of the coming turn carries the
232    /// `ix`th of them. Every one must be taken, so that a stream which ended
233    /// can not be seated by a stale handle of it.
234    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/// A model prepared to serve many streams at once: one state, one lane per
265/// stream, and turns seating whoever is ready.
266///
267/// `spawn` hands out a [`LanedStateHandle`] per stream, each holding a lane, and
268/// every `run` on a handle is a request to the worker thread which owns the
269/// state and the [`LaneTable`] both. The worker takes the turns queued at that
270/// moment, at most one per lane and at most [`TRACT_MAX_SEATS`] of them,
271/// concatenates their inputs along axis 0, publishes the seating and runs the
272/// state once, then hands each stream back its own seat.
273///
274/// A stream feeds one seat per turn: axis 0 carries streams, not data. Inputs and
275/// outputs whose axis 0 is a symbol are the batched ones; the rest are shared,
276/// so one value of such an input serves the whole turn and every seat must feed
277/// the same one, and such an output is handed back to every stream.
278#[derive(Clone)]
279pub struct LanedRunnable {
280    shared: Arc<Shared>,
281}
282
283struct Shared {
284    /// [`std::sync::mpsc::Sender`] is not `Sync`, and a `Runnable` is: handles
285    /// take their own clone of it, under the lock, once. `None` once the
286    /// runnable is being dropped, which is what closes the queue.
287    requests: Mutex<Option<Sender<Request>>>,
288    /// Joined when the runnable is dropped, so the worker is gone before
289    /// whatever the caller does next.
290    worker: Mutex<Option<thread::JoinHandle<()>>>,
291    /// The one-stream model the worker spawned its state from.
292    inner: Arc<dyn Runnable>,
293    /// `inner`'s model and plan, carried so that a laned runnable answers
294    /// [`Runnable::typed_model`] and [`Runnable::typed_plan`] like any other.
295    model: Option<Arc<TypedModel>>,
296    plan: Option<Arc<TypedSimplePlan>>,
297    /// The symbol axis 0 of the batched tensors carries: the turn's occupancy,
298    /// never a stream's own shapes.
299    batch: Symbol,
300    /// Lanes the state was reset for at `wrap`, hence the most streams that can
301    /// hold a handle at once and the widest a turn can be.
302    max_lanes: usize,
303    /// Turns and seats served, shared with the worker which is what counts them.
304    counts: Arc<Counts>,
305}
306
307/// The worker owns per-thread device state -- a CUDA stream, its cuBLAS and
308/// cuDNN handles -- whose destructors run as the thread exits. Nothing joined
309/// it before, so a process that returned from `main` while the worker was
310/// still winding down ran those destructors against libraries already tearing
311/// themselves down in their own `atexit` handlers, and segfaulted in
312/// `cudnnDestroy` about one run in fifteen. Closing the queue is what stops
313/// the worker, so the sender goes first and the join waits for at most the
314/// turn in flight.
315impl 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/// What the worker has served, for whoever tunes the turn policy: mean
328/// occupancy is `seats / turns`.
329#[derive(Debug, Default)]
330struct Counts {
331    /// Turns run since the model was prepared.
332    turns: AtomicU64,
333    /// Seats filled over those turns, so `seats / turns` is mean occupancy.
334    seats: AtomicU64,
335}
336
337/// What the worker reports through the `metrics` facade, labelled `runnable`
338/// with the model's `tract.name` when it has one. tract installs no recorder:
339/// the host does, before the runnable is wrapped, which is when the handles
340/// bind to it.
341struct 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    /// Serve `max_lanes` streams through `inner`, which must be prepared from a
393    /// model carrying a batch axis: at least one input and one output with a
394    /// symbol on axis 0, and one symbol for all of them.
395    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    /// The model as it was prepared, serving one stream at a time: what a turn
480    /// of one seat has to agree with.
481    pub fn inner(&self) -> &Arc<dyn Runnable> {
482        &self.shared.inner
483    }
484
485    /// The symbol axis 0 of the batched tensors carries. A stream feeds one seat
486    /// per turn, so it stands for the turn's occupancy, never for a stream's
487    /// own shapes.
488    pub fn batch_symbol(&self) -> &Symbol {
489        &self.shared.batch
490    }
491
492    /// Turns run and seats filled since the model was prepared: how wide the
493    /// turns the queue actually offers are.
494    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
511/// The symbol axis 0 of `fact` carries, or `None` for a tensor every seat
512/// shares. A stored fact can claim a symbol on an axis of extent one, so this
513/// says how the caller talks, not what the graph does with it.
514fn 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
521/// Why `facts` do not carry one batch symbol on axis 0: where the symbols they
522/// do carry sit, and that putting one on axis 0 is a graph edit.
523fn 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/// One stream's view of a [`LanedRunnable`]: the lane it holds, and the queue to
586/// the worker. Cloning it shares the lane -- clones are the same stream, and the
587/// lane goes back to the table once the last of them is dropped.
588#[derive(Clone, Debug)]
589pub struct LanedStateHandle {
590    /// The lane this stream holds, shared by the handle's clones: the last of
591    /// them dropped is what gives the lane back.
592    lease: Arc<Lease>,
593    /// The runnable the handle came from, for [`State::runnable`].
594    runnable: LanedRunnable,
595}
596
597#[derive(Debug)]
598struct Lease {
599    /// The lane the worker handed this stream at spawn.
600    lane: LaneId,
601    /// Where to send the lane back, which is all `Drop` needs.
602    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
626/// What a handle asks the worker for, one variant per event of the handle's
627/// life: a lane when it is spawned, a turn per call, and its lane back when it
628/// is dropped.
629enum Request {
630    /// A new stream wants a lane; the worker answers with one, reset, or with
631    /// the error that every lane is taken.
632    Spawn(Sender<TractResult<LaneId>>),
633    /// A stream wants a turn, or several at once.
634    Call(Call),
635    /// A stream is over and its lane goes back to the table.
636    Drop(LaneId),
637}
638
639/// One `run()` on a handle. Its inputs carry one seat or several, and it is
640/// answered once every one of them has been served -- over as many turns as
641/// the free lanes took to seat them all.
642struct Call {
643    /// The lane the caller holds for the life of its handle, which this call's
644    /// first seat sits in. It is a reservation, not a home: a call asking for
645    /// several seats borrows the rest.
646    leased: LaneId,
647    /// What the caller fed, batched inputs still stacked as they came.
648    inputs: TVec<TValue>,
649    /// Where the assembled answer goes; the caller blocks on the other end.
650    done: Sender<TractResult<TVec<TValue>>>,
651}
652
653/// One seat of one call, waiting for a lane: the lane it sits in -- its
654/// caller's own until a turn borrows another for it -- its place in the call,
655/// and the slice of the call's inputs it feeds.
656struct Seat {
657    /// The call this seat answers, keying its [`Completer`].
658    call: u64,
659    /// The lane the seat sits in: the call's leased one, overwritten by `fill`
660    /// with a borrowed one when that lane already took a seat this turn.
661    lane: LaneId,
662    /// Its place in the call, so the answer stacks back in the order asked.
663    ix: usize,
664    /// One seat's slice of the call's inputs, shared inputs whole.
665    inputs: TVec<TValue>,
666    /// When its call reached the worker.
667    queued: Instant,
668}
669
670/// The half of a call's answer the worker holds: its seats' outputs as they
671/// land, in the order it asked for them, and where to send them once the last
672/// of them has. The caller holds the other half and blocks on it.
673struct Completer {
674    /// One slot per seat the call asked for, filled as its seats land -- over
675    /// several turns when the call was wider than the free lanes.
676    served: Vec<Option<TVec<TValue>>>,
677    /// The call's own answer channel, moved here from the [`Call`].
678    done: Sender<TractResult<TVec<TValue>>>,
679}
680
681impl Call {
682    /// The inputs of each seat of this call: the batched ones sliced along
683    /// axis 0, the shared ones as they came. Every batched input says how many
684    /// seats the call asks for, so they must agree, and a call with none asks
685    /// for one.
686    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    /// Whether every seat the call asked for has landed.
733    fn is_full(&self) -> bool {
734        self.served.iter().all(Option::is_some)
735    }
736
737    /// Answer the call with its seats' outputs: the batched ones stacked back
738    /// along axis 0 in the order it asked for its seats, the shared ones as the
739    /// first seat got them. Every seat must have landed.
740    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/// The seats the worker has to seat, and the calls they answer. A call explodes
764/// into seats as it arrives, so what a turn picks from is a queue of seats: it
765/// fills to the free lanes from the head, and a call wider than they are is
766/// split at the boundary rather than held until it fits whole.
767#[derive(Default)]
768struct Queue {
769    /// Seats waiting for a lane, oldest first: what a turn fills from.
770    seats: VecDeque<Seat>,
771    /// The calls with seats still to land, by call id.
772    completers: HashMap<u64, Completer>,
773    /// The id the next call gets, so seats of different calls never collide.
774    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    /// Hand `seat` its turn's outputs, and answer the call when that was the
797    /// last seat it was waiting for. A failed seat fails the whole call at
798    /// once, and drops the seats of it still queued.
799    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
820/// The worker thread's world: the one state it runs, the lanes it hands to
821/// streams, the seats it has still to run, and what it was set up with.
822///
823/// It is built on the calling thread and moved to the worker, which is the only
824/// thread to touch it afterwards -- `state` and `lanes` go together because
825/// taking a lane resets it, and a reset writes the state.
826struct Worker {
827    /// The one state every turn runs, spawned from the model as prepared.
828    state: Box<dyn State>,
829    /// Which lanes streams hold, and what a turn's seating is drawn from.
830    lanes: LaneTable,
831    /// Seats waiting for a lane, and the calls they answer.
832    queued: Queue,
833    /// One flag per input, true where axis 0 carries the batch symbol: those
834    /// are sliced per seat, the rest serve the whole turn.
835    batch_in: Vec<bool>,
836    /// The same per output: batched ones are sliced back per seat, shared ones
837    /// handed to every seat of the turn.
838    batch_out: Vec<bool>,
839    /// The widest turn to run, [`TRACT_MAX_SEATS`] clamped to the state's lanes.
840    max_seats: usize,
841    /// How long a turn waits for latecomers once its first seat is queued
842    /// ([`TRACT_TURN_LINGER_US`]), zero to run as soon as one is ready.
843    linger: Duration,
844    /// Turns and seats served, shared with the runnable the caller reads them
845    /// from.
846    counts: Arc<Counts>,
847    meters: Meters,
848}
849
850impl Worker {
851    /// Serve `queue` until it closes, which is what dropping the runnable does.
852    fn work(mut self, queue: Receiver<Request>) {
853        loop {
854            // The linger belongs to the turn a request opens, not to the
855            // request: taking or giving back a lane must not delay the turns
856            // behind it, and turns left waiting by a full one have lingered
857            // already.
858            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    /// Hand out a lane or take it back there and then; queue the seats of a
909    /// call for the coming turns. Taking a lane resets it, which is why it
910    /// happens here rather than in the handle: it writes the state.
911    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    /// Seat the head of the queue, and the lanes borrowed to do it. A seat sits
939    /// in the lane its call leased, and a call's second and later seats of one
940    /// turn in lanes borrowed from whatever is free -- which resets them, so a
941    /// seat is always served by a lane holding nothing of another caller's.
942    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    /// Run `seated` as one turn: their batched inputs stacked along axis 0 in
983    /// seat order, their shared ones checked to agree, and the outputs sliced
984    /// back per seat.
985    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// The suite is disabled on Wasm because a laned runnable spawns a thread.
1036#[cfg(all(test, not(target_family = "wasm")))]
1037mod laned_test {
1038    use super::*;
1039    use crate::ops::math::{add, mul};
1040
1041    /// `[BATCH, 3] * 2`, prepared on the cpu runtime: stateless, so its lanes
1042    /// address nothing and only the seating of the batch axis is exercised.
1043    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    /// `TRACT_TURN_LINGER_US` is process-wide, so the tests which widen the
1063    /// turns hold this while they build their runnable and run their streams.
1064    static LINGER: Mutex<()> = Mutex::new(());
1065
1066    /// A dropped handle hands its lane back through the queue, so the lane is
1067    /// free at some point after the drop rather than at it.
1068    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    /// What the worker reported to `recorder`, by metric name, for the
1089    /// runnable named `doubler`.
1090    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    /// `[B, 3] * 2 + bias`, `bias` carrying no batch axis: the shape of a
1185    /// shared input, which one value of serves the whole turn.
1186    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    /// One turn per stream, all of them at once, the `stream`th feeding
1200    /// `biases[stream]`. Every lane is taken before any turn is queued, so the
1201    /// linger has the turns to seat together rather than a `spawn` to serve.
1202    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    /// A call feeding `seats` seats of the doubler at once: seat `s` feeds
1251    /// `[s, s + 1, s + 2]`, and the call is answered in that order.
1252    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    /// `[B, 3] + [B, 3]`: two batched inputs, which a call has to feed the same
1295    /// number of seats of.
1296    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}