zeavis-edu/apps/ml-service/src/main.rs

mod config;
mod error;
mod image;
mod model;
mod routes;
mod telemetry;

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 with calibration and wrap in Arc
    let model = Arc::new(ModelService::with_calibration(
        &config.model_path,
        config.input_size,
        config.temperature,
        config.conf_threshold_high,
        config.conf_threshold_low,
    ));
    tracing::info!(
        model_loaded = model.is_loaded(),
        model_path = ?config.model_path,
        temperature = config.temperature,
        conf_high = config.conf_threshold_high,
        conf_low = config.conf_threshold_low,
        "Model service initialized"
    );

    // Set model load status metric
    telemetry::model_load_status().set(if model.is_loaded() { 1.0 } else { 0.0 });

    // 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(())
}