fix block size in health check (#742)

predibase · Jan 18, 2025 · 223c554 · 223c554
1 parent 30c97cb
commit 223c554
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 12 deletions.
diff --git a/router/src/health.rs b/router/src/health.rs
@@ -1,6 +1,6 @@
 use lorax_client::{
-    Batch, NextTokenChooserParameters, Request, ShardInfo, ShardedClient,
-    StoppingCriteriaParameters,
+    input_chunk, Batch, InputChunk, NextTokenChooserParameters, Request, ShardInfo, ShardedClient,
+    StoppingCriteriaParameters, TokenizedInputs,
 };
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
@@ -40,7 +40,12 @@ impl Health {
         let generation_liveness_request = Request {
             id: LIVENESS_ID,
             inputs: "liveness".to_string(),
-            tokenized_inputs: None,
+            tokenized_inputs: Some(TokenizedInputs {
+                ids: vec![75],
+                input_chunks: vec![InputChunk {
+                    chunk: Some(input_chunk::Chunk::Text("liveness".to_string())),
+                }],
+            }),
             truncate: 10,
             prefill_logprobs: false,
             parameters: Some(NextTokenChooserParameters {
@@ -66,7 +71,7 @@ impl Health {
             adapter_index: 0,
             // Block 0 is reserved for health checks
             blocks: vec![0],
-            slots: (0..16).collect(),
+            slots: (0..self.shard_info.block_size).collect(),
             cache_len: 0,
             chunk_len: None,
         };
@@ -84,15 +89,20 @@ impl Health {
     pub(crate) async fn check_classification(&mut self) -> bool {
         let classify_request = Request {
             id: LIVENESS_ID,
-            inputs: "San Francisco".to_string(),
-            tokenized_inputs: None,
+            inputs: "liveness".to_string(),
+            tokenized_inputs: Some(TokenizedInputs {
+                ids: vec![75],
+                input_chunks: vec![InputChunk {
+                    chunk: Some(input_chunk::Chunk::Text("liveness".to_string())),
+                }],
+            }),
             truncate: 10,
             prefill_logprobs: false,
             parameters: None,
             stopping_parameters: None,
             adapter_index: 0,
             blocks: vec![0],
-            slots: (0..16).collect(),
+            slots: (0..self.shard_info.block_size).collect(),
             cache_len: 0,
             chunk_len: None,
         };
@@ -109,15 +119,20 @@ impl Health {
     pub(crate) async fn check_embeddings(&mut self) -> bool {
         let embed_request = Request {
             id: LIVENESS_ID,
-            inputs: "San Francisco".to_string(),
-            tokenized_inputs: None,
+            inputs: "liveness".to_string(),
+            tokenized_inputs: Some(TokenizedInputs {
+                ids: vec![75],
+                input_chunks: vec![InputChunk {
+                    chunk: Some(input_chunk::Chunk::Text("liveness".to_string())),
+                }],
+            }),
             truncate: 10,
             prefill_logprobs: false,
             parameters: None,
             stopping_parameters: None,
             adapter_index: 0,
             blocks: vec![0],
-            slots: (0..16).collect(),
+            slots: (0..self.shard_info.block_size).collect(),
             cache_len: 0,
             chunk_len: None,
         };

diff --git a/rust-toolchain.toml b/rust-toolchain.toml
@@ -1,3 +1,3 @@
 [toolchain]
-channel = "1.79.0"
-components = ["rustfmt", "clippy"]
+channel = "1.83.0"
+components = ["rustfmt", "clippy"]