Something went wrong. Try again.
Routing embedding stack for Qwen3-Embedding-0.6B on Apple Silicon: orchestrator + MLX + llama.cpp backends
Something went wrong. Try again.
7.0 kB · 208 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209mod dto;mod prompt;mod server;mod worker;
use std::collections::HashMap;use std::path::PathBuf;use std::sync::atomic::Ordering;use std::sync::Arc;
use anyhow::{bail, Context, Result};use tracing::{error, info};
use crate::dto::ModelSize;use crate::server::{AppState, ModelWorker};use crate::worker::{WorkerConfig, WorkRequest};
// ---------------------------------------------------------------------------// Env-var helpers// ---------------------------------------------------------------------------
fn env_or<T: std::str::FromStr>(key: &str, default: T) -> T { std::env::var(key) .ok() .and_then(|v| v.parse().ok()) .unwrap_or(default)}
/// Parse `EMBED_MODELS` env var (comma-separated, e.g. "0.6B,2B,4B")./// Defaults to just 0.6B.fn parse_model_sizes() -> Vec<ModelSize> { let raw = std::env::var("EMBED_MODELS").unwrap_or_else(|_| "0.6B".to_string()); raw.split(',') .filter_map(|s| { let s = s.trim(); match s { "0.6B" => Some(ModelSize::S06), "2B" => Some(ModelSize::S2), "4B" => Some(ModelSize::S4), other => { error!("Unknown model size '{other}', skipping. Valid: 0.6B, 2B, 4B"); None } } }) .collect()}
// ---------------------------------------------------------------------------// Model path resolution// ---------------------------------------------------------------------------
/// Resolve the GGUF file path for a given model size.////// Priority:/// 1. Size-specific env var: `MODEL_PATH_0_6B`, `MODEL_PATH_2B`, `MODEL_PATH_4B`/// 2. Generic `MODEL_PATH` (only practical when loading a single model)/// 3. Auto-download from HuggingFace via `hf-hub`fn resolve_model_path(size: ModelSize) -> Result<PathBuf> { // 1. Size-specific env var let size_key = match size { ModelSize::S06 => "MODEL_PATH_0_6B", ModelSize::S2 => "MODEL_PATH_2B", ModelSize::S4 => "MODEL_PATH_4B", }; if let Ok(path) = std::env::var(size_key) { return Ok(PathBuf::from(path)); }
// 2. Generic MODEL_PATH fallback if let Ok(path) = std::env::var("MODEL_PATH") { return Ok(PathBuf::from(path)); }
// 3. Auto-download from HuggingFace info!( repo = size.repo_id(), file = size.gguf_filename(), "Downloading model from HuggingFace…" );
let api = hf_hub::api::sync::ApiBuilder::new() .with_progress(true) .build() .context("unable to create HuggingFace API client")?;
let path = api .model(size.repo_id().to_string()) .get(size.gguf_filename()) .with_context(|| format!("unable to download {}", size.gguf_filename()))?;
Ok(path)}
// ---------------------------------------------------------------------------// Main// ---------------------------------------------------------------------------
#[tokio::main]async fn main() -> Result<()> { // Init tracing tracing_subscriber::fmt() .with_env_filter( tracing_subscriber::EnvFilter::try_from_default_env() .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), ) .init();
// Parse config let port: u16 = env_or("EMBED_PORT", 8080); let max_batch: usize = env_or("EMBED_MAX_BATCH", 64); let max_len: usize = env_or("EMBED_MAX_LEN", 512); let wait_ms: u64 = env_or("EMBED_WAIT_MS", 10);
let sizes = parse_model_sizes(); if sizes.is_empty() { bail!("No valid model sizes configured. Set EMBED_MODELS=0.6B,2B,4B"); }
let default_size = sizes[0]; info!( ?sizes, %default_size, port, max_batch, max_len, wait_ms, "Starting llama-embed-server" );
// ── Init llama backend ONCE ──────────────────────────────────────── // LlamaBackend::init() is a process-global singleton. We leak it to // obtain a &'static reference that all worker threads can share. let backend: &'static llama_cpp_2::llama_backend::LlamaBackend = { let b = llama_cpp_2::llama_backend::LlamaBackend::init() .context("LlamaBackend::init failed")?; Box::leak(Box::new(b)) };
// ── Resolve model paths and spawn workers ────────────────────────── let mut senders: HashMap<ModelSize, std::sync::mpsc::SyncSender<WorkRequest>> = HashMap::new(); let mut ready_receivers = Vec::new();
for &size in &sizes { let model_path = resolve_model_path(size)?; info!(size = %size, path = %model_path.display(), "Resolved model path");
// Bounded channel (cap = 1024) let (tx, rx) = std::sync::mpsc::sync_channel::<WorkRequest>(1024);
// Readiness oneshot let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
// Worker config let config = WorkerConfig { model_path, max_batch, max_len, wait_ms, };
// Spawn GPU worker on a dedicated OS thread let label = size.label().to_string(); std::thread::Builder::new() .name(format!("gpu-worker-{label}")) .spawn(move || { worker::run_worker(backend, config, rx, ready_tx); }) .with_context(|| format!("failed to spawn worker for {label}"))?;
senders.insert(size, tx); ready_receivers.push((size, ready_rx)); }
// ── Await all workers to be ready ────────────────────────────────── let mut model_workers = HashMap::new(); for (size, ready_rx) in ready_receivers { let n_embd = ready_rx .await .with_context(|| format!("{size} worker died during init"))? .with_context(|| format!("{size} worker init error"))?;
info!(size = %size, n_embd, "Worker ready");
let tx = senders.remove(&size).unwrap(); model_workers.insert(size, ModelWorker { tx, n_embd }); }
// ── Build shared state and serve ─────────────────────────────────── let state = Arc::new(AppState { workers: model_workers, ready: std::sync::atomic::AtomicBool::new(true), default_size, });
state.ready.store(true, Ordering::Release);
let app = server::router(state); let addr = format!("0.0.0.0:{port}"); let listener = tokio::net::TcpListener::bind(&addr) .await .with_context(|| format!("failed to bind {addr}"))?;
info!(%addr, "Listening"); axum::serve(listener, app).await.context("axum::serve")?;
Ok(())}