removed version constraints for candle deps

Cargo.toml

@@ -16,9 +16,9 @@ tracing = "0.1.37"
 tracing-subscriber = "0.3.18"
 uuid = { version = "1.6.1", features = ["v4"] }
 serde_json = "1.0.111"
-candle-core = { git = "https://github.com/huggingface/candle.git", version = "0.4.0" }
-candle-nn = { git = "https://github.com/huggingface/candle.git", version = "0.4.0" }
-candle-transformers = { git = "https://github.com/huggingface/candle.git", version = "0.4.0" }
+candle-core = { git = "https://github.com/huggingface/candle.git" }
+candle-nn = { git = "https://github.com/huggingface/candle.git" }
+candle-transformers = { git = "https://github.com/huggingface/candle.git" }
 tokenizers = "0.15.0"
 hf-hub = { version = "0.3.2", features = ["tokio"] }
 anyhow = "1.0.79"
@@ -27,7 +27,7 @@ tracing-chrome = "0.7.1"
 tower-http = { version = "0.5.1", features = ["trace", "timeout"] }
 once_cell = "1.19.0"
 clap = { version = "4.4.18", features = ["derive"] }
-pin-project = "1.1.3"
+workerpool = "1.2.1"
 
 [features]
 default = []

src/bin/server.rs

@@ -5,10 +5,10 @@ use clap::Parser;
 use tokio::net::TcpListener;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
-use glowrs::utils::{device::print_device_info, port_in_range};
+use glowrs::model::device::print_device_info;
 use glowrs::server::utils;
 use glowrs::server::{init_router, RouterArgs};
-
+use glowrs::server::utils::port_in_range;
 
 
 #[derive(Debug, Parser)]

src/infer/handler.rs

@@ -70,7 +70,7 @@ mod test {
     use super::*;
     use crate::infer::client::Client;
     use crate::infer::Queue;
-    use crate::utils::device::DEVICE;
+    use crate::model::device::DEVICE;
 
     fn append_str(s_in: String) -> String {
         format!("{}-processed", s_in)

src/infer/mod.rs

@@ -3,6 +3,7 @@ mod client;
 pub mod embed;
 mod handler;
 pub mod batch;
+pub(crate) mod pool;
 
 use uuid::Uuid;
 pub use queue::Queue;

src/infer/pool.rs

@@ -0,0 +1 @@
+

src/lib.rs

@@ -1,4 +1,3 @@
 pub mod model;
 pub mod infer;
-pub mod utils;
 pub mod server;

src/model/device.rs

@@ -0,0 +1,28 @@
+use once_cell::sync::Lazy;
+use candle_core::Device;
+
+#[cfg(all(feature = "metal", feature = "cuda"))]
+compile_error!("feature \"metal\" and feature \"cuda\" cannot be enabled at the same time");
+
+#[cfg(feature = "metal")]
+pub static DEVICE: Lazy<Device> =
+	Lazy::new(|| Device::new_metal(0).expect("No Metal device found."));
+
+#[cfg(feature = "cuda")]
+pub static DEVICE: Lazy<Device> =
+	Lazy::new(|| Device::new_cuda(0).expect("No CUDA device found."));
+
+#[cfg(not(any(feature = "metal", feature = "cuda")))]
+pub static DEVICE: Lazy<Device> = Lazy::new(|| Device::Cpu);
+
+pub fn print_device_info() {
+
+	#[cfg(not(any(feature = "metal", feature = "cuda")))]
+		tracing::info!("Using CPU");
+
+	#[cfg(feature = "cuda")]
+		tracing::info!("Using CUDA");
+
+	#[cfg(feature = "metal")]
+	tracing::info!("Using Metal");
+}

src/model/embedder.rs

@@ -10,9 +10,10 @@ use tokenizers::Tokenizer;
 
 // Re-exports
 pub use candle_transformers::models::{bert::BertModel, jina_bert::BertModel as JinaBertModel};
+use crate::model::utils::normalize_l2;
 
 use crate::server::data_models::{Sentences, Usage};
-use crate::utils::{normalize_l2, device::DEVICE};
+use crate::model::device::DEVICE;
 
 
 pub trait LoadableModel: Sized {

src/model/mod.rs

@@ -1,2 +1,4 @@
 pub mod embedder;
 pub mod sentence_transformer;
+mod utils;
+pub mod device;

src/model/utils.rs

@@ -0,0 +1,5 @@
+use candle_core::Tensor;
+
+pub fn normalize_l2(v: &Tensor) -> candle_core::Result<Tensor> {
+	v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)
+}

src/server/utils.rs

@@ -1,5 +1,7 @@
 use anyhow::Result;
 use tokio::signal;
+use std::result;
+use std::ops::RangeInclusive;
 
 type Nullary = fn() -> Result<()>;
 
@@ -32,3 +34,20 @@ pub async fn shutdown_signal(shutdown_fns_opt: Option<&[Nullary]>) {
         _ = terminate => {},
     }
 }
+
+const PORT_RANGE: RangeInclusive<u16> = 1..=65535;
+
+pub fn port_in_range(s: &str) -> result::Result<u16, String> {
+    let port: u16 = s
+        .parse()
+        .map_err(|_| format!("`{s}` isn't a port number"))?;
+    if PORT_RANGE.contains(&port) {
+        Ok(port)
+    } else {
+        Err(format!(
+            "port not in range {}-{}",
+            PORT_RANGE.start(),
+            PORT_RANGE.end()
+        ))
+    }
+}

tests/test_async_openai_client.py

@@ -21,7 +21,7 @@ async def create_embeddings() -> CreateEmbeddingResponse:
 
 
 async def call_health() -> None:
-	# Call the /health endpoint 10k times
+	# Call the /health endpoint 100 times
 	async with AsyncClient() as client:
 		for _ in range(100):
 			response = await client.get("http://127.0.0.1:3000/health")