Tweaks

EricLBuehler · Oct 12, 2024 · a8acde3 · a8acde3
1 parent 818df9d
commit a8acde3
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 17 deletions.
diff --git a/examples/server/stresstest.py b/examples/server/stresstest.py
@@ -2,6 +2,7 @@
 import httpx
 import textwrap
 import json
+import time
 
 
 def log_response(response: httpx.Response):
@@ -40,6 +41,7 @@ def log_response(response: httpx.Response):
 for i in range(1000):
     messages.append({"role": "user", "content": "Hello! How are you? Please write generic binary search function in Rust."})
     print("Sending", i)
+    start = time.time()
     completion = client.chat.completions.create(
         model="mistral",
         messages=messages,
@@ -49,4 +51,4 @@ def log_response(response: httpx.Response):
         temperature=0,
     )
     resp = completion.choices[0].message.content
-    print("Done")
+    print("Done", time.time()-start)
diff --git a/mistralrs-core/src/engine/mod.rs b/mistralrs-core/src/engine/mod.rs
@@ -140,7 +140,6 @@ impl Engine {
                         let current_completion_ids: Vec<usize> =
                             scheduled.completion.iter().map(|seq| *seq.id()).collect();
 
-                        println!("Running...");
                         let res = {
                             let mut pipeline = get_mut_arcmutex!(self.pipeline);
                             let pre_op = if !self.no_kv_cache
@@ -180,7 +179,6 @@ impl Engine {
                                 )
                                 .await
                         };
-                        println!("Ran...");
 
                         handle_pipeline_forward_error!(
                             "completion step",

diff --git a/mistralrs-core/src/pipeline/mod.rs b/mistralrs-core/src/pipeline/mod.rs
@@ -368,7 +368,6 @@ pub trait Pipeline:
 
                 match &logits[0] {
                     ForwardInputsResult::CausalGeneration { .. } => {
-                        println!("Sampling...");
                         self.sample_causal_gen(
                             input_seqs,
                             logits
@@ -389,7 +388,6 @@ pub trait Pipeline:
                             rng,
                         )
                         .await?;
-                        println!("Sampled!");
                     }
                     ForwardInputsResult::Image { .. } => {
                         send_responses(

diff --git a/mistralrs-core/src/sequence.rs b/mistralrs-core/src/sequence.rs
@@ -677,9 +677,7 @@ impl Sequence {
     }
 
     pub fn add_choice_to_group(&self, choice: Choice) {
-        println!("Adding choice");
         get_mut_group!(self).choices.push(choice);
-        println!("Added choice");
         self.update_time_info();
     }
 
@@ -819,9 +817,7 @@ impl SequenceGroup {
         sender: Sender<Response>,
     ) -> Result<(), SendError<Response>> {
         if self.choices.len() == self.n_choices {
-            println!("Sending response");
             sender.send(Response::Done(response)).await?;
-            println!("Sent response");
         }
 
         Ok(())

diff --git a/mistralrs/examples/stresstest/main.rs b/mistralrs/examples/stresstest/main.rs
@@ -1,5 +1,5 @@
 use anyhow::Result;
-use mistralrs::{IsqType, MemoryUsage, TextMessageRole, TextMessages, TextModelBuilder};
+use mistralrs::{IsqType, MemoryUsage, RequestBuilder, TextMessageRole, TextModelBuilder};
 
 const N_ITERS: u64 = 1000;
 const BYTES_TO_MB: usize = 1024 * 1024;
@@ -8,24 +8,28 @@ const BYTES_TO_MB: usize = 1024 * 1024;
 async fn main() -> Result<()> {
     let model = TextModelBuilder::new("microsoft/Phi-3.5-mini-instruct")
         .with_isq(IsqType::Q4K)
+        .with_prefix_cache_n(None)
         .with_logging()
         // .with_paged_attn(|| mistralrs::PagedAttentionMetaBuilder::default().build())?
         .build()
         .await?;
 
     for i in 0..N_ITERS {
-        let messages = TextMessages::new().add_message(
-            TextMessageRole::User,
-            "Hello! How are you? Please write generic binary search function in Rust.",
-        );
+        let messages = RequestBuilder::new()
+            .add_message(
+                TextMessageRole::User,
+                "Hello! How are you? Please write generic binary search function in Rust.",
+            )
+            .set_deterministic_sampler();
 
-        println!("Sending request {}...", i+1);
-        let _response = model.send_chat_request(messages).await?;
+        println!("Sending request {}...", i + 1);
+        let response = model.send_chat_request(messages).await?;
 
         let amount = MemoryUsage.get_memory_available(&model.config().device)? / BYTES_TO_MB;
 
         println!("{amount}");
-
+        println!("{}", response.usage.total_time_sec);
+        println!("{:?}", response.choices[0].message.content);
     }
 
     Ok(())