main.rs (2183B)
1 use anyhow::Result; 2 use axum::{body::Bytes, extract::State, http::StatusCode, routing::post, serve, Json, Router}; 3 use clap::Parser; 4 use parse_duration::parse; 5 use safetensors::SafeTensors; 6 use std::sync::Arc; 7 use tensor::{fit, normalize, to_tensor}; 8 use tokio::{net::TcpListener, spawn, task::spawn_blocking}; 9 10 mod tensor; 11 12 mod model; 13 use model::{Label, Model}; 14 15 mod queue; 16 use queue::Queue; 17 18 #[derive(Debug, Parser)] 19 #[command(about, long_about = None, version)] 20 struct Args { 21 #[arg(short, long, default_value = "0.0.0.0:8000")] 22 address: String, 23 24 #[arg(short, long, default_value = "./ast/model.pt")] 25 model_path: String, 26 27 #[arg(short, long, default_value_t = 1)] 28 batch_size: usize, 29 30 #[arg(short, long, default_value = "100ms")] 31 timeout: String, 32 } 33 34 async fn handler( 35 State(queue): State<Arc<Queue>>, 36 body: Bytes, 37 ) -> Result<Json<Label>, (StatusCode, String)> { 38 let tensor = spawn_blocking(move || { 39 let tensors = SafeTensors::deserialize(&body[..]) 40 .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; 41 42 let fbank = tensors 43 .tensor("fbank") 44 .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; 45 46 let mut tensor = to_tensor(fbank)?; 47 fit(&mut tensor)?; 48 normalize(&mut tensor)?; 49 50 Ok(tensor) 51 }) 52 .await 53 .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))??; 54 55 let label = queue 56 .add(tensor) 57 .await 58 .map_err(|e| (StatusCode::SERVICE_UNAVAILABLE, e.to_string()))? 59 .await 60 .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; 61 62 Ok(Json(label)) 63 } 64 65 #[tokio::main] 66 async fn main() -> Result<()> { 67 let args = Args::parse(); 68 69 let timeout = parse(&args.timeout)?; 70 let (queue, receiver) = Queue::new(args.batch_size); 71 let queue = Arc::new(queue); 72 73 let model = Model::new(args.model_path)?; 74 spawn(async move { 75 queue::run(receiver, model, args.batch_size, timeout).await; 76 }); 77 78 let router = Router::new().route("/", post(handler)).with_state(queue); 79 80 let listener = TcpListener::bind(args.address).await?; 81 Ok(serve(listener, router).await?) 82 }