From b4975819ebeaaf7cfc1cf7cd854bb17d1e93579e Mon Sep 17 00:00:00 2001 From: Asep Haryana Saputra <90584806+MythEclipse@users.noreply.github.com> Date: Sat, 23 May 2026 09:48:54 +0000 Subject: [PATCH] feat: serve ML inference endpoints with Axum --- apps/ml-service/src/main.rs | 43 +++++++++++- apps/ml-service/src/routes.rs | 124 ++++++++++++++++++++++++++++++++++ 2 files changed, 165 insertions(+), 2 deletions(-) diff --git a/apps/ml-service/src/main.rs b/apps/ml-service/src/main.rs index c45b34a..7621cfa 100644 --- a/apps/ml-service/src/main.rs +++ b/apps/ml-service/src/main.rs @@ -4,6 +4,45 @@ mod image; mod model; mod routes; -fn main() { - println!("zeavis-ml-service"); +use anyhow::Result; +use config::Config; +use model::ModelService; +use routes::{router, AppState}; +use std::sync::Arc; +use tokio::net::TcpListener; +use tracing_subscriber::EnvFilter; + +#[tokio::main] +async fn main() -> Result<()> { + // Initialize tracing subscriber with EnvFilter + tracing_subscriber::fmt() + .with_env_filter(EnvFilter::from_default_env()) + .init(); + + // Load configuration from environment + let config = Config::from_env()?; + + // Create ModelService and wrap in Arc + let model = Arc::new(ModelService::new(&config.model_path, config.input_size)); + + // Log model status + tracing::info!( + model_loaded = model.is_loaded(), + model_path = ?config.model_path, + "Model service initialized" + ); + + // Create AppState + let state = AppState { model }; + + // Bind TcpListener to host:port + let addr = format!("{}:{}", config.host, config.port); + let listener = TcpListener::bind(&addr).await?; + + tracing::info!(address = %addr, "Server listening"); + + // Serve with Axum + axum::serve(listener, router(state)).await?; + + Ok(()) } diff --git a/apps/ml-service/src/routes.rs b/apps/ml-service/src/routes.rs index 892745f..1fb1d06 100644 --- a/apps/ml-service/src/routes.rs +++ b/apps/ml-service/src/routes.rs @@ -1,5 +1,15 @@ use serde::{Deserialize, Serialize}; use crate::config::{LABELS, SERVICE_NAME, SERVICE_VERSION}; +use crate::model::{ModelService, Prediction}; +use crate::error::ServiceError; +use crate::image::preprocess_image; +use axum::{ + extract::{State, Multipart}, + http::StatusCode, + routing::{get, post}, + Json, Router, +}; +use std::sync::Arc; #[derive(Debug, Serialize, Deserialize)] pub struct HealthResponse { @@ -17,6 +27,18 @@ pub struct MetadataResponse { pub labels: Vec, } +#[derive(Debug, Serialize, Deserialize)] +pub struct PredictionResponse { + pub label: String, + pub confidence: f32, + pub probabilities: std::collections::BTreeMap, +} + +#[derive(Clone)] +pub struct AppState { + pub model: Arc, +} + pub fn health_response(model_loaded: bool) -> HealthResponse { HealthResponse { status: "ok".to_string(), @@ -35,6 +57,82 @@ pub fn metadata_response(model_path: String, model_loaded: bool, input_size: u32 } } +pub fn prediction_response(prediction: Prediction) -> PredictionResponse { + PredictionResponse { + label: prediction.label, + confidence: prediction.confidence, + probabilities: prediction.probabilities, + } +} + +pub async fn health(State(state): State) -> Json { + Json(health_response(state.model.is_loaded())) +} + +pub async fn metadata(State(state): State) -> Json { + Json(metadata_response( + state.model.model_path().to_string_lossy().to_string(), + state.model.is_loaded(), + state.model.input_size(), + )) +} + +pub async fn predict( + State(state): State, + mut multipart: Multipart, +) -> Result, ServiceError> { + // Extract the file field from multipart + let mut file_data = None; + while let Ok(Some(field)) = multipart.next_field().await { + if field.name() == Some("file") { + if let Some(content_type) = field.content_type() { + if !content_type.starts_with("image/") { + return Err(ServiceError::BadRequest( + "Uploaded file must be an image".to_string(), + )); + } + } else { + return Err(ServiceError::BadRequest( + "Uploaded file must be an image".to_string(), + )); + } + file_data = Some(field.bytes().await); + break; + } + } + + // Handle missing or multipart read errors + let bytes = match file_data { + Some(Ok(b)) => b, + Some(Err(_)) => { + return Err(ServiceError::BadRequest( + "Uploaded file must be an image".to_string(), + )) + } + None => { + return Err(ServiceError::BadRequest( + "Uploaded file must be an image".to_string(), + )) + } + }; + + // Preprocess the image + let input = preprocess_image(&bytes, state.model.input_size())?; + + // Run prediction + let prediction = state.model.predict(input)?; + + Ok(Json(prediction_response(prediction))) +} + +pub fn router(state: AppState) -> Router { + Router::new() + .route("/health", get(health)) + .route("/metadata", get(metadata)) + .route("/predict", post(predict)) + .with_state(state) +} + #[cfg(test)] mod tests { use super::*; @@ -74,4 +172,30 @@ mod tests { vec!["Bercak Daun", "Daun Sehat", "Karat Daun", "Hawar Daun"] ); } + + #[test] + fn prediction_response_matches_prediction_contract() { + let prediction = Prediction { + label: "Karat Daun".to_string(), + confidence: 0.6, + probabilities: { + let mut map = std::collections::BTreeMap::new(); + map.insert("Bercak Daun".to_string(), 0.1); + map.insert("Daun Sehat".to_string(), 0.2); + map.insert("Karat Daun".to_string(), 0.6); + map.insert("Hawar Daun".to_string(), 0.1); + map + }, + }; + + let response = prediction_response(prediction); + + assert_eq!(response.label, "Karat Daun"); + assert_eq!(response.confidence, 0.6); + assert_eq!(response.probabilities.len(), 4); + assert_eq!(response.probabilities.get("Bercak Daun"), Some(&0.1)); + assert_eq!(response.probabilities.get("Daun Sehat"), Some(&0.2)); + assert_eq!(response.probabilities.get("Karat Daun"), Some(&0.6)); + assert_eq!(response.probabilities.get("Hawar Daun"), Some(&0.1)); + } }