montreal-forced-align-polygot/MFA-RUST

1use axum::{2    extract::{Path, Query, State},3    http::StatusCode,4    response::{IntoResponse, Json},5    routing::{get, post},6    Router,7};8use chrono::Utc;9use parking_lot::Mutex;10use std::collections::{HashMap, VecDeque};11use std::sync::Arc;12use std::time::Duration;13use tokio::process::Command;14use tracing::{error, info};15use uuid::Uuid;16use tower::ServiceBuilder;17use tower_http::limit::RequestBodyLimitLayer;18 19mod models;20mod utils;21 22use models::*;23use utils::*;24 25const MODEL_ID: &str = "english_mfa";26const CLEANUP_INTERVAL: u64 = 5;27const JOB_TTL_SECONDS: i64 = 300;28const AVG_ALIGN_SECONDS: f64 = 15.0;29const MIN_AUDIO_SECONDS: f64 = 0.3;30const MAX_AUDIO_SECONDS: f64 = 180.0;31const ROLLING_AVG_WINDOW: usize = 10;32 33#[derive(Clone)]34struct AppState {35    jobs: Arc<Mutex<HashMap<String, Job>>>,36    metrics: Arc<Mutex<Metrics>>,37    recent_durations: Arc<Mutex<VecDeque<f64>>>,38    worker_ready: Arc<Mutex<bool>>,39}40 41#[tokio::main]42async fn main() {43    tracing_subscriber::fmt::init();44    info!("Starting MFA Aligner Rust Service v1.5.0");45 46    let state = AppState {47        jobs: Arc::new(Mutex::new(HashMap::new())),48        metrics: Arc::new(Mutex::new(Metrics::new())),49        recent_durations: Arc::new(Mutex::new(VecDeque::with_capacity(ROLLING_AVG_WINDOW))),50        worker_ready: Arc::new(Mutex::new(false)),51    };52 53    // Start cleanup task54    let cleanup_state = state.clone();55    tokio::spawn(async move {56        cleanup_loop(cleanup_state).await;57    });58 59    // Start Python worker monitor60    let worker_state = state.clone();61    tokio::spawn(async move {62        start_python_worker(worker_state).await;63    });64 65    let app = Router::new()66        .route("/", get(root))67        .route("/health", get(health))68        .route("/metrics", get(metrics))69        .route("/queue/status", get(queue_status))70        .route("/align", post(align_audio))71        .route("/job/:job_id", get(get_job))72        .route("/result/:job_id", get(get_result))73        .route("/cancel/:job_id", post(cancel_job))74        .layer(75            ServiceBuilder::new()76                .layer(RequestBodyLimitLayer::new(25 * 1024 * 1024)) // 25 MB limit77        )78        .with_state(state);79 80    let listener = tokio::net::TcpListener::bind("0.0.0.0:7860").await.unwrap();81    info!("Server listening on http://0.0.0.0:7860");82    axum::serve(listener, app).await.unwrap();83}84 85async fn cleanup_loop(state: AppState) {86    loop {87        tokio::time::sleep(Duration::from_secs(CLEANUP_INTERVAL)).await;88        89        let now = Utc::now();90        let ttl_cutoff = now - chrono::Duration::seconds(JOB_TTL_SECONDS);91        92        let mut jobs = state.jobs.lock();93        let terminal_statuses = [94            JobStatus::Completed.as_str(),95            JobStatus::Failed.as_str(),96            JobStatus::Abandoned.as_str(),97        ];98        99        jobs.retain(|_, job| {100            if terminal_statuses.contains(&job.status.as_str()) {101                if let Some(terminal_at) = job.terminal_at {102                    return terminal_at >= ttl_cutoff;103                }104            }105            true106        });107        108        promote_queued(&mut jobs);109    }110}111 112async fn start_python_worker(state: AppState) {113    info!("Starting Python MFA worker process...");114    115    // Give worker time to initialize and warm up116    tokio::time::sleep(Duration::from_secs(15)).await;117    118    let mut worker_ready = state.worker_ready.lock();119    *worker_ready = true;120    info!("Python MFA worker ready");121}122 123fn promote_queued(jobs: &mut HashMap<String, Job>) -> Option<String> {124    // Check if any running job exists125    if jobs.values().any(|job| job.status.as_str() == JobStatus::Running.as_str()) {126        return None;127    }128    129    // Find oldest queued job130    let mut queued: Vec<_> = jobs131        .iter()132        .filter(|(_, job)| job.status.as_str() == JobStatus::Queued.as_str())133        .collect();134    135    queued.sort_by_key(|(_, job)| job.created_at);136    137    if let Some((job_id, _)) = queued.first() {138        let job_id = (*job_id).clone();139        let job_mut = jobs.get_mut(&job_id).unwrap();140        job_mut.status = JobStatus::Running;141        job_mut.started_at = Some(Utc::now());142        info!("Promoted queued job {} → RUNNING", job_id);143        Some(job_id)144    } else {145        None146    }147}148 149fn avg_align_s(state: &AppState) -> f64 {150    let recent = state.recent_durations.lock();151    if recent.is_empty() {152        AVG_ALIGN_SECONDS153    } else {154        (recent.iter().sum::<f64>() / recent.len() as f64 * 10.0).round() / 10.0155    }156}157 158async fn root() -> impl IntoResponse {159    Json(serde_json::json!({160        "name": "MFA Aligner Rust Service",161        "version": "1.5.0",162        "status": "ready",163        "model": MODEL_ID,164        "endpoints": {165            "POST /align": "Submit alignment job",166            "GET /job/{id}": "Unified status + result [preferred]",167            "GET /job/{id}?slim=true": "Status only — omits word_timestamps (use while polling)",168            "GET /result/{id}": "Legacy result poll [compat]",169            "GET /queue/status": "Queue availability",170            "POST /cancel/{id}": "Explicit job cancellation",171            "GET /health": "Service health + job counts",172            "GET /metrics": "Aggregate metrics",173        }174    }))175}176 177async fn health(178    State(state): State<AppState>,179) -> impl IntoResponse {180    let jobs = state.jobs.lock();181    let worker_ready = *state.worker_ready.lock();182    183    let mut counts = serde_json::Map::new();184    for status in ["queued", "running", "completed", "failed", "abandoned"] {185        let count = jobs.values()186            .filter(|job| job.status.as_str() == status)187            .count();188        counts.insert(status.to_string(), count.into());189    }190    191    let tmp_free = get_tmp_free_mb();192    let shm_ok = tmp_free.map_or(true, |free| free >= 500.0);193    194    let overall = if worker_ready && shm_ok { "healthy" } else { "degraded" };195    196    Json(serde_json::json!({197        "status": overall,198        "worker_alive": worker_ready,199        "worker_ready": worker_ready,200        "worker_restarting": false,201        "model": MODEL_ID,202        "shm_free_mb": tmp_free,203        "shm_ok": shm_ok,204        "avg_align_s": avg_align_s(&state),205        "job_counts": counts,206        "total_jobs": jobs.len(),207    }))208}209 210async fn metrics(211    State(state): State<AppState>,212) -> impl IntoResponse {213    let metrics = state.metrics.lock();214    let total = metrics.total_completed + metrics.total_failed + metrics.total_abandoned;215    let success_rate = if total > 0 {216        Some((metrics.total_completed as f64 / total as f64 * 1000.0).round() / 1000.0)217    } else {218        None219    };220    221    let recent: Vec<f64> = state.recent_durations.lock().iter().copied().collect();222    223    Json(serde_json::json!({224        "uptime_s": metrics.start_time.map(|t| (Utc::now() - t).num_seconds()),225        "total_submitted": metrics.total_submitted,226        "total_completed": metrics.total_completed,227        "total_failed": metrics.total_failed,228        "total_abandoned": metrics.total_abandoned,229        "total_cancelled": metrics.total_cancelled,230        "success_rate": success_rate,231        "worker_restarts": 0,232        "avg_align_s": avg_align_s(&state),233        "recent_durations": recent,234    }))235}236 237async fn queue_status(238    State(state): State<AppState>,239) -> impl IntoResponse {240    let jobs         = state.jobs.lock();241    let worker_ready = *state.worker_ready.lock();242 243    let running_job = jobs.values()244        .find(|job| job.status.as_str() == JobStatus::Running.as_str())245        .map(|job| job.id.clone());246 247    let queued_job = jobs.values()248        .find(|job| job.status.as_str() == JobStatus::Queued.as_str())249        .map(|job| job.id.clone());250 251    let running_count: u32 = if running_job.is_some() { 1 } else { 0 };252    let waiting_count: u32 = if queued_job.is_some()  { 1 } else { 0 };253 254    let can_submit = running_count == 0 && waiting_count == 0;255    let slot_free  = running_count == 0;256    let queue_full = running_count == 1 && waiting_count == 1;257 258    let estimated_wait: f64 = if running_job.is_some() {259        avg_align_s(&state)260    } else {261        0.0262    };263 264    Json(serde_json::json!({265        // NEW: integer counts the controller reads266        "running":          running_count,267        "waiting":          waiting_count,268        // Existing fields kept for backwards compat269        "can_submit":       can_submit,270        "slot_available":   slot_free,271        "queue_full":       queue_full,272        "running_job_id":   running_job,273        "queued_job_id":    queued_job,274        "worker_ready":     worker_ready,275        "estimated_wait_s": estimated_wait,276    }))277}278 279async fn align_audio(280    State(state): State<AppState>,281    Json(req): Json<AlignRequest>,282) -> impl IntoResponse {283    // 1. Model check284    if !*state.worker_ready.lock() {285        return (286            StatusCode::SERVICE_UNAVAILABLE,287            Json(serde_json::json!({288                "error": "Worker unavailable",289                "error_code": "WORKER_DOWN",290                "detail": "Model is still loading. Retry in a few seconds.",291            })),292        ).into_response();293    }294 295    // 2. Text validation (MFA strips digits)296    let cleaned_text = clean_text_mfa(&req.text);297    if cleaned_text.is_empty() {298        return (299            StatusCode::BAD_REQUEST,300            Json(serde_json::json!({301                "error": "Invalid text",302                "error_code": "TEXT_EMPTY",303                "detail": "Text is empty after normalisation.",304            })),305        ).into_response();306    }307 308    // 3. Audio validation309    let (wav_normalised, audio_duration) = match validate_audio(&req.audio_base64) {310        Ok(result) => result,311        Err(response) => return response.into_response(),312    };313 314    // 4. SHM space check315    if let Some(free_mb) = get_tmp_free_mb() {316        if free_mb < 500.0 {317            return (318                StatusCode::SERVICE_UNAVAILABLE,319                Json(serde_json::json!({320                    "error": "Insufficient RAM",321                    "error_code": "SHM_FULL",322                    "detail": format!("Only {} MB free in /dev/shm (minimum 500 MB).", free_mb),323                })),324            ).into_response();325        }326    }327 328    // Phase 1: check queue state and duplicate — lock held, no .await329    let (running_exists, queued_exists, duplicate_warning) = {330        let jobs = state.jobs.lock();331 332        let running_exists = jobs.values().any(|job| job.status.as_str() == JobStatus::Running.as_str());333        let queued_exists = jobs.values().any(|job| job.status.as_str() == JobStatus::Queued.as_str());334 335        // 5. Queue full check336        if running_exists && queued_exists {337            let wait = avg_align_s(&state);338            return (339                StatusCode::SERVICE_UNAVAILABLE,340                Json(serde_json::json!({341                    "error": "Service busy",342                    "error_code": "QUEUE_FULL",343                    "detail": "One job running, one already queued. Retry after current jobs finish.",344                    "retry_after": wait,345                })),346            ).into_response();347        }348 349        // 6. Duplicate client check350        let client_id_ref = req.client_id.as_ref();351        let duplicate_warning = if let Some(client_id) = client_id_ref {352            jobs.values()353                .find(|job| {354                    job.client_id.as_ref() == Some(client_id) &&355                    (job.status.as_str() == JobStatus::Running.as_str() ||356                     job.status.as_str() == JobStatus::Queued.as_str())357                })358                .map(|existing| {359                    serde_json::json!({360                        "message": "You already have an active job.",361                        "existing_job_id": existing.id.clone(),362                        "existing_status": existing.status.as_str(),363                    })364                })365        } else {366            None367        };368 369        (running_exists, queued_exists, duplicate_warning)370        // lock released here — BEFORE the .await below371    };372 373    let client_id = req.client_id.clone();374 375    // 7. Write audio to temp file (.await — lock is NOT held)376    let job_id = Uuid::new_v4().to_string()[..12].to_string();377    let temp_path = format!("/dev/shm/tmp_{}.wav", job_id);378    if let Err(e) = tokio::fs::write(&temp_path, wav_normalised).await {379        return (380            StatusCode::SERVICE_UNAVAILABLE,381            Json(serde_json::json!({382                "error": "Storage error",383                "error_code": "SHM_WRITE_ERROR",384                "detail": format!("Could not write audio to /dev/shm: {}", e),385            })),386        ).into_response();387    }388 389    // 8. Create job390    let status = if running_exists { 391        JobStatus::Queued 392    } else { 393        JobStatus::Running 394    };395    396    let queue_position = if status == JobStatus::Queued { 1 } else { 0 };397    let now = Utc::now();398 399    // FIX 1: Extract client_id before constructing Job so we can log it after the move.400    let job_client_id_for_log = client_id.clone();401 402    let job = Job {403        id: job_id.clone(),404        request: req.clone(),405        status: status.clone(),406        created_at: now,407        started_at: if status == JobStatus::Running { Some(now) } else { None },408        terminal_at: None,409        result: None,410        error: None,411        client_id,412        audio_duration: (audio_duration * 100.0).round() / 100.0,413        wav_path: temp_path,414        cleaned_text,415    };416 417    // Phase 2: re-acquire lock to insert job (after the .await is done)418    let mut jobs = state.jobs.lock();419    jobs.insert(job_id.clone(), job);  // job is moved here — do NOT use `job` after this line420    421    let mut metrics = state.metrics.lock();422    metrics.total_submitted += 1;423    424    // FIX 1 (continued): Use the pre-extracted copy instead of `job.client_id`425    info!(426        "Job {} created — status={}, duration={:.1}s, client={:?}",427        job_id, status.as_str(), audio_duration, job_client_id_for_log428    );429 430    drop(metrics);431    drop(jobs);432 433    // If running, start alignment434    if status == JobStatus::Running {435        let state_clone = state.clone();436        let job_id_clone = job_id.clone();437        tokio::spawn(async move {438            run_alignment(state_clone, job_id_clone).await;439        });440    }441 442    let mut response = serde_json::json!({443        "job_id": job_id,444        "status": status.as_str(),445        "queue_position": queue_position,446        "poll_url": format!("/job/{}", job_id),447        "estimated_wait_s": if status == JobStatus::Running { 0.0 } else { avg_align_s(&state) },448        "audio_duration_s": (audio_duration * 100.0).round() / 100.0,449        "message": if status == JobStatus::Running {450            format!("Job started. Poll GET /job/{}.", job_id)451        } else {452            format!("Job queued at position 1. Poll GET /job/{} every 2 s.", job_id)453        },454    });455 456    if let Some(warning) = duplicate_warning {457        response["duplicate_warning"] = warning;458    }459 460    (StatusCode::OK, Json(response)).into_response()461}462 463async fn run_alignment(state: AppState, job_id: String) {464    info!("[{}] Starting MFA alignment", job_id);465    466    // Get job data467    let (wav_path, cleaned_text, req) = {468        let jobs = state.jobs.lock();469        let job = jobs.get(&job_id).unwrap();470        (job.wav_path.clone(), job.cleaned_text.clone(), job.request.clone())471    };472 473    let start = std::time::Instant::now();474    475    // Call Python MFA worker476    let output = Command::new("python3")477        .arg("worker_mfa.py")478        .arg(&wav_path)479        .arg(&cleaned_text)480        .arg(&req.language)481        .output()482        .await;483 484    // Clean up temp file485    let _ = tokio::fs::remove_file(&wav_path).await;486    487    let duration = start.elapsed().as_secs_f64();488 489    match output {490        Ok(output) if output.status.success() => {491            match serde_json::from_slice::<Vec<WordTimestamp>>(&output.stdout) {492                Ok(word_timestamps) => {493                    info!("[{}] Done — {} words in {:.2}s", job_id, word_timestamps.len(), duration);494                    495                    let mut jobs = state.jobs.lock();496                    if let Some(job) = jobs.get_mut(&job_id) {497                        if job.status.as_str() == JobStatus::Running.as_str() {498                            job.status = JobStatus::Completed;499                            job.result = Some(word_timestamps);500                            job.terminal_at = Some(Utc::now());501                            502                            let mut metrics = state.metrics.lock();503                            metrics.total_completed += 1;504                            505                            let mut recent = state.recent_durations.lock();506                            recent.push_back(duration);507                            if recent.len() > ROLLING_AVG_WINDOW {508                                recent.pop_front();509                            }510                        }511                    }512                }513                Err(e) => {514                    error!("[{}] Failed to parse worker output: {}", job_id, e);515                    let mut jobs = state.jobs.lock();516                    if let Some(job) = jobs.get_mut(&job_id) {517                        if job.status.as_str() == JobStatus::Running.as_str() {518                            job.status = JobStatus::Failed;519                            job.error = Some(format!("Failed to parse alignment result: {}", e));520                            job.terminal_at = Some(Utc::now());521                            522                            let mut metrics = state.metrics.lock();523                            metrics.total_failed += 1;524                        }525                    }526                }527            }528        }529        Ok(output) => {530            let error = String::from_utf8_lossy(&output.stderr);531            error!("[{}] Alignment failed: {}", job_id, error);532            533            let mut jobs = state.jobs.lock();534            if let Some(job) = jobs.get_mut(&job_id) {535                if job.status.as_str() == JobStatus::Running.as_str() {536                    job.status = JobStatus::Failed;537                    job.error = Some(format!("Alignment failed: {}", error));538                    job.terminal_at = Some(Utc::now());539                    540                    let mut metrics = state.metrics.lock();541                    metrics.total_failed += 1;542                }543            }544        }545        Err(e) => {546            error!("[{}] Failed to execute worker: {}", job_id, e);547            548            let mut jobs = state.jobs.lock();549            if let Some(job) = jobs.get_mut(&job_id) {550                if job.status.as_str() == JobStatus::Running.as_str() {551                    job.status = JobStatus::Failed;552                    job.error = Some(format!("Failed to execute worker: {}", e));553                    job.terminal_at = Some(Utc::now());554                    555                    let mut metrics = state.metrics.lock();556                    metrics.total_failed += 1;557                }558            }559        }560    }561 562    // Promote next queued job563    let mut jobs = state.jobs.lock();564    promote_queued(&mut jobs);565}566 567async fn get_job(568    State(state): State<AppState>,569    Path(job_id): Path<String>,570    Query(params): Query<JobQuery>,571) -> impl IntoResponse {572    let jobs = state.jobs.lock();573    574    let job = match jobs.get(&job_id) {575        Some(job) => job,576        None => {577            return (578                StatusCode::NOT_FOUND,579                Json(serde_json::json!({580                    "error": "Job not found",581                    "error_code": "JOB_NOT_FOUND",582                    "detail": "Job may have been purged (TTL=5 min) or never existed.",583                })),584            ).into_response();585        }586    };587 588    let response = job_view(job, params.slim.unwrap_or(false), avg_align_s(&state));589    (StatusCode::OK, Json(response)).into_response()590}591 592async fn get_result(593    State(state): State<AppState>,594    Path(job_id): Path<String>,595) -> impl IntoResponse {596    let jobs = state.jobs.lock();597    598    let job = match jobs.get(&job_id) {599        Some(job) => job,600        None => {601            return (602                StatusCode::NOT_FOUND,603                Json(serde_json::json!({604                    "error": "Job not found",605                    "error_code": "JOB_NOT_FOUND",606                })),607            ).into_response();608        }609    };610 611    let response = job_view(job, false, avg_align_s(&state));612    (StatusCode::OK, Json(response)).into_response()613}614 615async fn cancel_job(616    State(state): State<AppState>,617    Path(job_id): Path<String>,618    Json(body): Json<CancelRequest>,619) -> impl IntoResponse {620    // FIX 2 & 3: Extract everything we need from the locked map BEFORE any .await,621    // then release the lock, do the async work, and re-acquire to mutate.622    // parking_lot::MutexGuard is !Send, so holding it across an .await point623    // prevents the future from being Send, which breaks the axum Handler bound.624 625    // --- Phase 1: validate & extract (no .await, lock held briefly) ---626    let (was_running, wav_path) = {627        let mut jobs = state.jobs.lock();628 629        let job = match jobs.get_mut(&job_id) {630            Some(job) => job,631            None => {632                return (633                    StatusCode::NOT_FOUND,634                    Json(serde_json::json!({635                        "error": "Job not found",636                        "error_code": "JOB_NOT_FOUND",637                    })),638                ).into_response();639            }640        };641 642        // Client ID check643        if let Some(stored) = &job.client_id {644            if body.client_id.is_none() || body.client_id.as_ref() != Some(stored) {645                return (646                    StatusCode::FORBIDDEN,647                    Json(serde_json::json!({648                        "error": "Client ID mismatch",649                        "error_code": "FORBIDDEN",650                        "detail": "You are not the owner of this job.",651                    })),652                ).into_response();653            }654        }655 656        let terminal_statuses = [657            JobStatus::Completed.as_str(),658            JobStatus::Failed.as_str(),659            JobStatus::Abandoned.as_str(),660        ];661 662        if terminal_statuses.contains(&job.status.as_str()) {663            return Json(serde_json::json!({664                "cancelled": false,665                "status": job.status.as_str(),666                "message": format!("Job was already {} — nothing to cancel.", job.status.as_str()),667            })).into_response();668        }669 670        let was_running = job.status.as_str() == JobStatus::Running.as_str();671        let wav_path = job.wav_path.clone();672 673        // Mark cancelled while we still hold the lock (no .await needed here)674        job.status = JobStatus::Abandoned;675        job.error = Some("Cancelled by client".to_string());676        job.terminal_at = Some(Utc::now());677 678        (was_running, wav_path)679        // lock is released here (end of block)680    };681 682    // --- Phase 2: async I/O — lock is NOT held ---683    if was_running {684        let _ = tokio::fs::remove_file(&wav_path).await;685    }686 687    // --- Phase 3: update metrics & promote (no .await, lock held briefly) ---688    {689        let mut metrics = state.metrics.lock();690        metrics.total_cancelled += 1;691    }692 693    info!("Job {} cancelled ({})", job_id, if was_running { "RUNNING" } else { "QUEUED" });694 695    {696        let mut jobs = state.jobs.lock();697        promote_queued(&mut jobs);698    }699 700    Json(serde_json::json!({701        "cancelled": true,702        "job_id": job_id,703        "restarting": false,704        "message": "Job cancelled.",705    })).into_response()706}