Merge pull request #236 from lucasavila00/async_chan3

Async channels
EricLBuehler · Apr 30, 2024 · dd8ffc6 · dd8ffc6
2 parents 3b10a9a + 60a4242
commit dd8ffc6
Show file tree

Hide file tree

Showing 19 changed files with 261 additions and 214 deletions.
diff --git a/mistralrs-bench/Cargo.toml b/mistralrs-bench/Cargo.toml
@@ -21,6 +21,7 @@ mistralrs-core = { version = "0.1.1", path = "../mistralrs-core" }
 tracing.workspace = true
 tracing-subscriber.workspace = true
 either.workspace = true
+tokio.workspace = true
 cli-table = "0.4.7"
 
 [features]

diff --git a/mistralrs-bench/src/main.rs b/mistralrs-bench/src/main.rs
@@ -6,8 +6,9 @@ use mistralrs_core::{
  ModelSelected, Request, RequestMessage, Response, SamplingParams, SchedulerMethod, TokenSource,
  Usage,
 };
+use std::fmt::Display;
 use std::sync::Arc;
-use std::{fmt::Display, sync::mpsc::channel};
+use tokio::sync::mpsc::channel;
 use tracing::info;
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::EnvFilter;
@@ -65,7 +66,7 @@ fn run_bench(
  n_choices: 1,
  };
  let sender = mistralrs.get_sender();
- let (tx, rx) = channel();
+ let (tx, mut rx) = channel(10_000);
 
  let req = Request {
  id: mistralrs.next_request_id(),
@@ -82,11 +83,13 @@ fn run_bench(
 
  for _ in 0..repetitions {
  for _ in 0..concurrency {
- sender.send(req.clone()).expect("Expected receiver.");
+ sender
+ .blocking_send(req.clone())
+ .expect("Expected receiver.");
  }
  for _ in 0..concurrency {
- match rx.recv() {
- Ok(r) => match r {
+ match rx.blocking_recv() {
+ Some(r) => match r {
  Response::InternalError(e) => {
  unreachable!("Got an internal error: {e:?}");
  }
@@ -105,7 +108,7 @@ fn run_bench(
  usages.push(res.usage);
  }
  },
- Err(e) => unreachable!("Expected a Done response, got: {:?}", e),
+ None => unreachable!("Expected a Done response, got None",),
  }
  }
  }