cervo_core/batcher/
wrapper.rs

1// Author: Tom Solberg <tom.solberg@embark-studios.com>
2// Copyright © 2022, Embark Studios, all rights reserved.
3// Created: 27 July 2022
4
5use super::Batcher;
6use crate::inferer::{Inferer, Response, State};
7use std::collections::HashMap;
8
9/// Wraps an inferer in a batching interface. This'll separate the
10/// data-insertion and execution, which generally improves
11/// performance.
12///
13/// Can be easily constructed using [`InfererExt::into_batched`](crate::prelude::InfererExt::into_batched).
14pub struct Batched<Inf: Inferer> {
15    inner: Inf,
16    batcher: Batcher,
17}
18
19impl<Inf> Batched<Inf>
20where
21    Inf: Inferer,
22{
23    /// Wrap the provided inferer.
24    pub fn wrap(inferer: Inf) -> Self {
25        let batcher = Batcher::new(&inferer);
26        Self {
27            batcher,
28            inner: inferer,
29        }
30    }
31
32    /// Insert a single element into the batch to include in the next execution.
33    pub fn push(&mut self, id: u64, state: State<'_>) -> anyhow::Result<()> {
34        self.batcher.push(id, state)
35    }
36
37    /// Insert a sequence of elements into the batch to include in the next execution.
38    pub fn extend<'a, Iter: IntoIterator<Item = (u64, State<'a>)>>(
39        &mut self,
40        states: Iter,
41    ) -> anyhow::Result<()> {
42        self.batcher.extend(states)
43    }
44
45    /// Execute the model on the data that has been enqueued previously.
46    pub fn execute(&mut self) -> anyhow::Result<HashMap<u64, Response<'_>>> {
47        self.batcher.execute(&self.inner)
48    }
49
50    /// Split the batcher and the inferer.
51    pub fn into_parts(self) -> (Inf, Batcher) {
52        (self.inner, self.batcher)
53    }
54}
cervo_core/batcher/wrapper.rs

cervo_core/batcher/
wrapper.rs
