Support send + sync in loader

mistralrs-core/src/pipeline/loaders/diffusion_loaders.rs

@@ -45,7 +45,7 @@ pub trait DiffusionModel {
     fn max_seq_len(&self) -> usize;
 }
 
-pub trait DiffusionModelLoader {
+pub trait DiffusionModelLoader: Send + Sync {
     /// If the model is being loaded with `load_model_from_hf` (so manual paths not provided), this will be called.
     fn get_model_paths(&self, api: &ApiRepo, model_id: &Path) -> Result<Vec<PathBuf>>;
     /// If the model is being loaded with `load_model_from_hf` (so manual paths not provided), this will be called.

mistralrs-core/src/pipeline/loaders/mod.rs

@@ -377,7 +377,7 @@ impl ModelKind {
 ///     None,
 /// ).unwrap();
 /// ```
-pub trait Loader {
+pub trait Loader: Send + Sync {
     /// If `revision` is None, then it defaults to `main`.
     /// If `dtype` is None, then it defaults to the model default (usually BF16).
     /// If model is not found on HF, will attempt to resolve locally.

mistralrs-core/src/pipeline/loaders/normal_loaders.rs

@@ -87,7 +87,7 @@ pub struct NormalLoadingMetadata {
     pub real_device: Device,
 }
 
-pub trait NormalModelLoader: IsqModelLoader {
+pub trait NormalModelLoader: IsqModelLoader + Send + Sync {
     fn load(
         &self,
         config: &str,

mistralrs-core/src/pipeline/loaders/vision_loaders.rs

@@ -58,7 +58,7 @@ pub trait VisionModel: IsqModel + AnyMoeBaseModelMixin {
     fn default_model_specific_args(&self, input_ids: &Tensor) -> Box<dyn Any>;
 }
 
-pub trait VisionModelLoader: IsqModelLoader {
+pub trait VisionModelLoader: IsqModelLoader + Send + Sync {
     fn load(
         &self,
         config: &str,

mistralrs/Cargo.toml

@@ -150,3 +150,7 @@ required-features = []
 [[example]]
 name = "imatrix"
 required-features = []
+
+[[example]]
+name = "async"
+required-features = []

mistralrs/examples/async/main.rs

@@ -0,0 +1,59 @@
+use anyhow::Result;
+use mistralrs::{
+    IsqType, PagedAttentionMetaBuilder, RequestBuilder, TextMessageRole, TextMessages,
+    TextModelBuilder,
+};
+
+async fn run() -> Result<()> {
+    let model = TextModelBuilder::new("microsoft/Phi-3.5-mini-instruct")
+        .with_isq(IsqType::Q8_0)
+        .with_logging()
+        .with_paged_attn(|| PagedAttentionMetaBuilder::default().build())?
+        .build()
+        .await?;
+
+    let messages = TextMessages::new()
+        .add_message(
+            TextMessageRole::System,
+            "You are an AI agent with a specialty in programming.",
+        )
+        .add_message(
+            TextMessageRole::User,
+            "Hello! How are you? Please write generic binary search function in Rust.",
+        );
+
+    let response = model.send_chat_request(messages).await?;
+
+    println!("{}", response.choices[0].message.content.as_ref().unwrap());
+    dbg!(
+        response.usage.avg_prompt_tok_per_sec,
+        response.usage.avg_compl_tok_per_sec
+    );
+
+    // Next example: Return some logprobs with the `RequestBuilder`, which enables higher configurability.
+    let request = RequestBuilder::new().return_logprobs(true).add_message(
+        TextMessageRole::User,
+        "Please write a mathematical equation where a few numbers are added.",
+    );
+
+    let response = model.send_chat_request(request).await?;
+
+    println!(
+        "Logprobs: {:?}",
+        &response.choices[0]
+            .logprobs
+            .as_ref()
+            .unwrap()
+            .content
+            .as_ref()
+            .unwrap()[0..3]
+    );
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let task = tokio::task::spawn(async move { run().await.unwrap() });
+    task.await?;
+    Ok(())
+}