janitor

janitor
git clone https://git.prime8.dev/janitor.git
Log | Files | Refs | README | LICENSE

queue.rs (2649B)

      1 use spinoff::{spinners, Spinner};
      2 use std::time::Duration;
      3 use tch::Tensor;
      4 use tokio::{
      5     sync::{
      6         mpsc::{channel, Receiver, Sender},
      7         oneshot::{
      8             channel as oneshot_channel, Receiver as OneshotReceiver, Sender as OneshotSender,
      9         },
     10     },
     11     time::{timeout, Instant},
     12 };
     13 
     14 use crate::model::{Label, Model};
     15 
     16 type Job = (Tensor, OneshotSender<Label>);
     17 
     18 pub struct Queue {
     19     sender: Sender<Job>,
     20 }
     21 
     22 impl Queue {
     23     pub fn new(batch_size: usize) -> (Self, Receiver<Job>) {
     24         let (sender, receiver) = channel(batch_size * 2);
     25         (Self { sender }, receiver)
     26     }
     27 
     28     pub async fn add(&self, tensor: Tensor) -> anyhow::Result<OneshotReceiver<Label>> {
     29         let (result_tx, result_rx) = oneshot_channel();
     30         self.sender
     31             .send((tensor, result_tx))
     32             .await
     33             .map_err(|_| anyhow::anyhow!("Queue is closed, cannot accept new jobs"))?;
     34         Ok(result_rx)
     35     }
     36 }
     37 
     38 pub async fn run(
     39     mut receiver: Receiver<Job>,
     40     model: Model,
     41     batch_size: usize,
     42     batch_timeout: Duration,
     43 ) {
     44     let mut spinner = Spinner::new(spinners::Line, "Waiting for jobs", None);
     45 
     46     loop {
     47         let (first_tensor, first_tx) = match receiver.recv().await {
     48             Some(job) => job,
     49             None => break,
     50         };
     51 
     52         let mut tensors = vec![first_tensor];
     53         let mut transmitters = vec![first_tx];
     54 
     55         let start = Instant::now();
     56         while tensors.len() < batch_size {
     57             let elapsed = start.elapsed();
     58             if elapsed >= batch_timeout {
     59                 break;
     60             }
     61 
     62             match timeout(batch_timeout - elapsed, receiver.recv()).await {
     63                 Ok(Some((tensor, tx))) => {
     64                     tensors.push(tensor);
     65                     transmitters.push(tx);
     66                 }
     67                 _ => break,
     68             }
     69         }
     70 
     71         spinner.update_text(format!(
     72             "Executing {} job(s) with {} still queued",
     73             tensors.len(),
     74             receiver.len()
     75         ));
     76 
     77         if let Err(e) = process_batch(&model, tensors, transmitters).await {
     78             eprintln!("Batch error: {}", e);
     79         }
     80 
     81         spinner.update_text("Waiting for jobs");
     82     }
     83 }
     84 
     85 async fn process_batch(
     86     model: &Model,
     87     tensors: Vec<Tensor>,
     88     transmitters: Vec<OneshotSender<Label>>,
     89 ) -> anyhow::Result<()> {
     90     let stacked = Tensor::f_stack(&tensors, 0)?;
     91 
     92     let labels = model.label(&stacked)?;
     93 
     94     for (label, tx) in labels.iter().zip(transmitters) {
     95         if tx.send(*label).is_err() {
     96             eprintln!("Receiver dropped for a job, result lost");
     97         }
     98     }
     99 
    100     Ok(())
    101 }