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.
3.5 kB · 119 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120use std::collections::HashMap;use std::sync::atomic::{AtomicBool, Ordering};use std::sync::Arc;
use axum::extract::State;use axum::http::StatusCode;use axum::response::IntoResponse;use axum::routing::{get, post};use axum::{Json, Router};
use crate::dto::{EmbedRequest, EmbedResponse, ModelSize};use crate::worker::WorkRequest;
// ---------------------------------------------------------------------------// Shared application state// ---------------------------------------------------------------------------
pub struct ModelWorker { pub tx: std::sync::mpsc::SyncSender<WorkRequest>, pub n_embd: usize,}
pub struct AppState { /// One worker channel per loaded model size. pub workers: HashMap<ModelSize, ModelWorker>, /// Set to true once all workers are ready. pub ready: AtomicBool, /// The default model size when the request omits `size`. pub default_size: ModelSize,}
// ---------------------------------------------------------------------------// Router// ---------------------------------------------------------------------------
pub fn router(state: Arc<AppState>) -> Router { Router::new() .route("/healthz", get(healthz)) .route("/embed", post(embed)) .with_state(state)}
// ---------------------------------------------------------------------------// GET /healthz// ---------------------------------------------------------------------------
async fn healthz(State(state): State<Arc<AppState>>) -> impl IntoResponse { if state.ready.load(Ordering::Relaxed) { (StatusCode::OK, "ok") } else { (StatusCode::SERVICE_UNAVAILABLE, "warming") }}
// ---------------------------------------------------------------------------// POST /embed// ---------------------------------------------------------------------------
async fn embed( State(state): State<Arc<AppState>>, Json(req): Json<EmbedRequest>,) -> Result<Json<EmbedResponse>, (StatusCode, String)> { // Validate texts non-empty if req.texts.is_empty() { return Err(( StatusCode::BAD_REQUEST, "texts must be non-empty".to_string(), )); }
// Resolve which model to use let size = req.size.unwrap_or(state.default_size);
let worker = state.workers.get(&size).ok_or_else(|| { let available: Vec<String> = state.workers.keys().map(|s| s.to_string()).collect(); ( StatusCode::BAD_REQUEST, format!( "model size {size} not loaded; available: [{}]", available.join(", ") ), ) })?;
// Create oneshot for the response let (reply_tx, reply_rx) = tokio::sync::oneshot::channel();
let work = WorkRequest { texts: req.texts, input_kind: req.input_type, reply: reply_tx, };
// Send to worker worker.tx.send(work).map_err(|_| { ( StatusCode::INTERNAL_SERVER_ERROR, "worker channel closed".to_string(), ) })?;
// Await response let vectors = reply_rx .await .map_err(|_| { ( StatusCode::INTERNAL_SERVER_ERROR, "worker dropped reply".to_string(), ) })? .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("{e:#}")))?;
let dim = worker.n_embd; Ok(Json(EmbedResponse { vectors, dim, model: size.label().to_string(), }))}