Serving Models with Axum

use axum::{extract::State, routing::post, Json, Router};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::task;
 
#[derive(Clone)]
struct AppState {
    // session: Arc<ort::session::Session>,
}
 
#[derive(Deserialize)]
struct PredictRequest {
    features: Vec<f32>,
}
 
#[derive(Serialize)]
struct PredictResponse {
    score: f32,
}
 
async fn predict(
    State(_state): State<AppState>,
    Json(req): Json<PredictRequest>,
) -> Json<PredictResponse> {
    let features = req.features;
    let score = task::spawn_blocking(move || run_model(&features))
        .await
        .unwrap();
    Json(PredictResponse { score })
}
 
fn run_model(features: &[f32]) -> f32 {
    features.iter().sum::<f32>() / features.len() as f32
}

use axum::{extract::State, http::StatusCode, routing::post, Json, Router};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Instant;
use tokio::task;
use tracing::info;
 
#[derive(Clone)]
struct AppState {
    model_version: Arc<String>,
}
 
#[derive(Deserialize)]
struct PredictRequest {
    features: Vec<f32>,
}
 
#[derive(Serialize)]
struct PredictResponse {
    score: f32,
    model_version: String,
}
 
async fn predict(
    State(state): State<AppState>,
    Json(req): Json<PredictRequest>,
) -> Result<Json<PredictResponse>, StatusCode> {
    if req.features.is_empty() || req.features.len() > 10_000 {
        return Err(StatusCode::BAD_REQUEST);
    }
    let version = state.model_version.clone();
    let start = Instant::now();
    let features = req.features;
    let score = task::spawn_blocking(move || infer(&features))
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    info!(elapsed_ms = start.elapsed().as_millis(), "inference");
    Ok(Json(PredictResponse {
        score,
        model_version: (*version).clone(),
    }))
}
 
fn infer(features: &[f32]) -> f32 {
    features.iter().copied().fold(0.0f32, f32::max)
}
 
#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();
    let state = AppState {
        model_version: Arc::new("v1.0.0".into()),
    };
    let app = Router::new()
        .route("/predict", post(predict))
        .with_state(state);
    let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}