montreal-forced-align-polygot/MFA-python
0
1"""2Montreal Forced Aligner (MFA) — Word Timestamp Service v5.53Changes vs v5.4:4 5═══════════════════════════════════════════════════════════════════════════════6CORRECTNESS / SAFETY FIXES7═══════════════════════════════════════════════════════════════════════════════8 91. EARLY AUDIO + TEXT VALIDATION (was: validated inside worker after queuing)10 /align now validates at submission time and returns HTTP 400 immediately:11 · base64 decode failure12 · not a valid WAV file13 · sample width not 16-bit or 32-bit float (unsupported format)14 · mono or stereo only (>2 channels rejected)15 · audio duration < MIN_AUDIO_SECONDS (0.3 s) — MFA fails on very short audio16 · audio duration > MAX_AUDIO_SECONDS (180 s) — cap is alignment latency, not memory17 (HF free tier gives ~7–9 GB /dev/shm; 180 s audio is only ~5.8 MB)18 · text empty after _clean() — all-symbol / all-emoji inputs19 Jobs that would fail are now rejected before entering the queue, so the20 frontend gets a 400 immediately instead of FAILED after a full cycle.21 222. /dev/shm FREE SPACE CHECK (was: missing)23 /align rejects with 503 if /dev/shm has less than SHM_MIN_FREE_MB (500 MB)24 free. HF free tier gives ~7–9 GB /dev/shm; 500 MB guards against genuinely25 stressed RAM (MFA MFCC/Kaldi temp files for 180 s audio can reach 200–500 MB).26 273. WORKER LIVENESS CHECK AT SUBMISSION (was: missing)28 /align rejects with 503 if the worker process is not alive AND a restart is29 not already in progress, rather than silently queueing jobs that will never30 be processed.31 324. /cancel NO LONGER BLOCKS THE EVENT LOOP (was: _restart_worker called sync)33 _restart_worker() calls terminate()+join() which can block ~3–5 s.34 /cancel now fires it via loop.run_in_executor so the endpoint returns35 immediately and other requests are not stalled.36 375. GRACE PERIOD EXTENDED DURING WORKER RESTART (was: fixed 10 s)38 When _promote_queued() is called while a restart is in progress,39 last_heartbeat is extended by HEARTBEAT_GRACE + WORKER_RESTART_GRACE (15 s)40 so the promoted job is not abandoned while the new worker boots.41 42═══════════════════════════════════════════════════════════════════════════════43PERFORMANCE IMPROVEMENTS44═══════════════════════════════════════════════════════════════════════════════45 466. AUDIO PASSED VIA /dev/shm FILE, NOT PICKLED THROUGH MPQueue (was: bytes)47 Previously the full audio bytes were serialised via multiprocessing.Queue48 (pickle). For a 30-second 16kHz mono WAV that is ~1 MB of pickle overhead49 per job. Now /align writes the validated WAV directly to a temp file under50 WORK_ROOT and sends only the file *path* through the queue. The worker51 reads the file, processes it, then deletes the temp file. IPC payload52 drops from O(audio_size) to O(1).53 547. ROLLING AVERAGE FOR estimated_wait_s (was: hardcoded AVG_ALIGN_SECONDS=5)55 A deque of the last ROLLING_AVG_WINDOW (10) completed job durations is56 maintained. estimated_wait_s in /align, /job, and /queue/status responses57 now reflects the real recent average alignment time.58 598. worker_ready FLAG IN /health (was: missing)60 The worker signals warm-up completion by putting ("ready",) on the result61 queue. /health exposes worker_ready:bool so the frontend can show a62 "warming up" state instead of submitting a job that will be delayed by the63 warm-up overhead.64 65═══════════════════════════════════════════════════════════════════════════════66FRONTEND / EDGE-CASE IMPROVEMENTS67═══════════════════════════════════════════════════════════════════════════════68 699. GET /metrics — monitoring endpoint70 Returns: total_jobs_submitted, total_completed, total_failed, total_abandoned,71 total_cancelled, avg_align_s (rolling), worker_restarts, uptime_s.72 7310. DUPLICATE CLIENT DETECTION — /align74 If a client_id already has an active (QUEUED or RUNNING) job and submits75 again, the response includes a warning with the existing job_id instead of76 silently creating a second job. The duplicate is still accepted (the77 caller may have a legitimate reason) but the frontend is informed.78 7911. GET /job/{id}?slim=true — skip word_timestamps during polling80 During the polling loop the frontend doesn't need the full timestamp array81 on every response — it only needs it once on completion. ?slim=true omits82 word_timestamps from COMPLETED responses, reducing payload size.83 84═══════════════════════════════════════════════════════════════════════════════85RETAINED FROM v5.486═══════════════════════════════════════════════════════════════════════════════87 · Persistent PretrainedAligner (skips 22–30 s setup per job).88 · /dev/shm RAM tmpfs for all MFA I/O.89 · TTL-based job purging (JOB_TTL_SECONDS=300).90 · Worker restart outside jobs_lock.91 · fine_tune=False, _clear_dir() optimisation.92 · Queued jobs don't heartbeat; grace period on promotion.93 · Unified _job_view(), /cancel, /job, /result (legacy).94"""95 96import os, re, uuid, base64, logging, shutil, io, wave, time97import asyncio98import queue as _stdlib_queue99from collections import deque100from pathlib import Path101from multiprocessing import Process, Queue as MPQueue102from datetime import datetime, timedelta103from typing import Dict, Any, Optional104 105import numpy as np106from fastapi import FastAPI, HTTPException, Request, Query107from fastapi.middleware.cors import CORSMiddleware108from fastapi.responses import JSONResponse109from pydantic import BaseModel110from praatio import textgrid as tgio111 112from montreal_forced_aligner.alignment import PretrainedAligner113 114logging.basicConfig(level=logging.INFO)115logger = logging.getLogger(__name__)116 117app = FastAPI(title="MFA Aligner — RAM-Optimised v5.5", version="5.5.0")118app.add_middleware(119 CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]120)121 122MFA_ACOUSTIC = "english_mfa"123MFA_DICTIONARY = "english_mfa"124MFA_ROOT = os.environ.get("MFA_ROOT_DIR", "/mfa")125WORK_ROOT = "/dev/shm/mfa_work"126 127# ── Timing & size constants ───────────────────────────────────────────────────128CLEANUP_INTERVAL = 5 # s — cleanup loop tick rate129JOB_TTL_SECONDS = 300 # s — terminal jobs purged after this130AVG_ALIGN_SECONDS = 15 # s — fallback before rolling average has data (conservative for 180 s audio)131MIN_AUDIO_SECONDS = 0.3 # s — MFA fails on very short audio132MAX_AUDIO_SECONDS = 180.0 # s — cap is latency not memory (180 s audio ≈ 5.8 MB WAV;133 # HF Spaces free tier has 14–18 GB RAM so /dev/shm is ~7–9 GB)134SHM_MIN_FREE_MB = 500 # MB — HF free tier gives ~7–9 GB /dev/shm; MFA MFCC/Kaldi135 # temp files for 180 s audio can reach 200–500 MB so we136 # guard against genuinely stressed RAM, not theoretical OOM137ROLLING_AVG_WINDOW = 10 # n jobs — rolling average window for wait estimate138 139 140# ═══════════════════════════════════════════════════════════════════════════════141# MODELS142# ═══════════════════════════════════════════════════════════════════════════════143 144class JobStatus:145 QUEUED = "queued"146 RUNNING = "running"147 COMPLETED = "completed"148 ABANDONED = "abandoned"149 FAILED = "failed"150 151TERMINAL_STATUSES = {JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.ABANDONED}152 153class AlignRequest(BaseModel):154 audio_base64: str155 text: str156 language: str = "english"157 pass_number: int = 1158 client_id: Optional[str] = None159 160class CancelRequest(BaseModel):161 client_id: Optional[str] = None162 163jobs: Dict[str, Dict[str, Any]] = {}164jobs_lock = asyncio.Lock()165 166_work_queue: MPQueue = MPQueue()167_result_queue: MPQueue = MPQueue()168 169# ── Worker state ──────────────────────────────────────────────────────────────170_worker_restarting: bool = False171_worker_ready: bool = False # set True when warm-up completes172 173# ── Metrics (all mutated under jobs_lock or atomically) ───────────────────────174_metrics: Dict[str, Any] = {175 "total_submitted": 0,176 "total_completed": 0,177 "total_failed": 0,178 "total_abandoned": 0,179 "total_cancelled": 0,180 "worker_restarts": 0,181 "start_time": None, # set at startup182}183# Rolling deque of (duration_s,) for completed jobs — used for wait estimate184_recent_durations: deque = deque(maxlen=ROLLING_AVG_WINDOW)185 186 187# ═══════════════════════════════════════════════════════════════════════════════188# UTILITIES189# ═══════════════════════════════════════════════════════════════════════════════190 191def _shm_free_mb() -> Optional[float]:192 try:193 st = os.statvfs("/dev/shm")194 return round(st.f_bavail * st.f_frsize / 1_048_576, 1)195 except Exception:196 return None197 198def _avg_align_s() -> float:199 """Rolling average of recent job durations, or fallback constant."""200 if not _recent_durations:201 return AVG_ALIGN_SECONDS202 return round(sum(_recent_durations) / len(_recent_durations), 1)203 204def _clear_dir(path: Path) -> None:205 """Unlink dir contents without destroying the inode — faster than rmtree+mkdir."""206 for child in path.iterdir():207 if child.is_dir():208 shutil.rmtree(str(child), ignore_errors=True)209 else:210 child.unlink(missing_ok=True)211 212def _normalise_wav(wav_bytes: bytes) -> bytes:213 """214 Normalise a WAV to the exact format MFA requires: mono, 16-bit PCM.215 Handles:216 · 32-bit float → int16 conversion217 · stereo → mono downmix (average of channels)218 16-bit mono WAVs are returned unchanged (zero copy).219 Called once at submission time so the worker always gets a clean file.220 """221 buf = io.BytesIO(wav_bytes)222 with wave.open(buf) as wf:223 sample_width = wf.getsampwidth()224 n_channels = wf.getnchannels()225 frame_rate = wf.getframerate()226 n_frames = wf.getnframes()227 raw = wf.readframes(n_frames)228 229 # Fast path: already mono 16-bit230 if sample_width == 2 and n_channels == 1:231 return wav_bytes232 233 # Decode samples234 if sample_width == 4:235 samples = np.frombuffer(raw, dtype=np.float32)236 samples = np.clip(samples, -1.0, 1.0)237 else:238 # 16-bit PCM — reshape for possible stereo downmix239 samples = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0240 241 # Stereo → mono downmix242 if n_channels == 2:243 samples = samples.reshape(-1, 2).mean(axis=1)244 elif n_channels > 2:245 samples = samples.reshape(-1, n_channels).mean(axis=1)246 247 # Back to int16248 pcm = (np.clip(samples, -1.0, 1.0) * 32767).astype(np.int16)249 250 out = io.BytesIO()251 with wave.open(out, "wb") as wf:252 wf.setnchannels(1)253 wf.setsampwidth(2)254 wf.setframerate(frame_rate)255 wf.writeframes(pcm.tobytes())256 return out.getvalue()257 258def _clean(text: str) -> str:259 """260 Normalise text for MFA's english_mfa dictionary.261 262 Rules (each justified by what english_mfa can/cannot pronounce):263 264 1. Unicode punctuation → ASCII equivalents265 · curly apostrophes (\u2018 \u2019) → straight apostrophe (')266 Keeps contractions intact: "it\u2019s" → "it's" ✓267 · em/en dash (\u2014 \u2013) → space268 "well\u2014known" → "well known" (two words MFA knows)269 270 2. Digits stripped entirely (not replaced with words)271 english_mfa has no pronunciation for "123", "42", etc.272 OOV digits cause alignment failure or silently skipped words273 depending on MFA version. Stripping is safer than leaving them in.274 If the caller needs digit-to-word conversion (e.g. "42" → "forty two")275 that should happen upstream before calling /align.276 277 3. Only ASCII letters, mid-word apostrophes, mid-word hyphens, and spaces278 are kept. Everything else becomes a space.279 280 4. Mid-word apostrophe rule:281 Only apostrophes flanked by a letter on BOTH sides are kept.282 · "don't" → "don't" ✓ (mid-word)283 · "'hello'" → "hello" ✓ (leading/trailing quote apostrophes stripped)284 · "it's" → "it's" ✓285 · "rock 'n' roll" → "rock n roll" (apostrophes not mid-word)286 This is applied AFTER the main strip so the regex sees clean input.287 288 5. Mid-word hyphen rule:289 Only hyphens flanked by a letter on BOTH sides are kept.290 · "well-known" → "well-known" ✓ (in english_mfa dict)291 · "- item" → " item" (leading hyphen stripped)292 · "item -" → "item " (trailing hyphen stripped)293 294 6. Newlines → space (single-utterance corpus; MFA does not need them)295 296 7. Collapse runs of whitespace, strip leading/trailing.297 298 Returns the cleaned string. Caller checks for empty result.299 """300 # Step 1 — unicode punctuation normalisation301 text = text.replace("\u2019", "'").replace("\u2018", "'") # curly apostrophes302 text = text.replace("\u2014", " ").replace("\u2013", " ") # em/en dash303 text = text.replace("\u2026", " ") # ellipsis304 text = text.replace("\u00e9", "e").replace("\u00e8", "e") # é è → e (common in loanwords)305 text = text.replace("\u00e0", "a").replace("\u00e2", "a") # à â → a306 text = text.replace("\u00f4", "o").replace("\u00f6", "o") # ô ö → o307 text = text.replace("\u00fc", "u").replace("\u00fb", "u") # ü û → u308 text = text.replace("\u00ee", "i").replace("\u00ef", "i") # î ï → i309 text = text.replace("\u00e7", "s") # ç → s (rough but better than OOV)310 311 # Step 2 — newlines → space312 text = text.replace("\n", " ").replace("\r", " ")313 314 # Step 3 — strip digits explicitly before the main pass315 text = re.sub(r"[0-9]+", " ", text)316 317 # Step 4 — keep only ASCII letters, apostrophes, hyphens, spaces318 text = re.sub(r"[^a-zA-Z'\- ]", " ", text)319 320 # Step 5 — mid-word apostrophe: strip apostrophes NOT flanked by letters321 # Regex: apostrophe that is NOT preceded by [a-zA-Z] OR not followed by [a-zA-Z]322 text = re.sub(r"(?<![a-zA-Z])'|'(?![a-zA-Z])", " ", text)323 324 # Step 6 — mid-word hyphen: strip hyphens NOT flanked by letters325 text = re.sub(r"(?<![a-zA-Z])-|-(?![a-zA-Z])", " ", text)326 327 # Step 7 — collapse whitespace328 text = re.sub(r"\s+", " ", text).strip()329 330 return text.lower() # MFA dictionary is lowercase331 332 333def _validate_audio(audio_b64: str) -> tuple[bytes, float]:334 """335 Decode and validate audio at submission time.336 Returns (wav_bytes_mono_16bit, duration_s) or raises HTTPException 400.337 338 Parses the WAV exactly once and normalises in the same pass.339 Catching problems here gives the frontend an immediate 400 instead of340 a FAILED job after the full alignment cycle.341 """342 # 1. Base64 decode343 try:344 raw = base64.b64decode(audio_b64)345 except Exception:346 raise HTTPException(400, {347 "error": "Invalid audio",348 "error_code": "AUDIO_DECODE_ERROR",349 "detail": "audio_base64 is not valid base64.",350 })351 352 # 2. WAV parse — single open to read all metadata353 try:354 buf = io.BytesIO(raw)355 with wave.open(buf) as wf:356 sample_width = wf.getsampwidth()357 n_channels = wf.getnchannels()358 frame_rate = wf.getframerate()359 n_frames = wf.getnframes()360 except Exception as e:361 raise HTTPException(400, {362 "error": "Invalid audio",363 "error_code": "AUDIO_PARSE_ERROR",364 "detail": f"Could not parse as WAV: {e}",365 })366 367 # 3. Format checks368 if sample_width not in (2, 4):369 raise HTTPException(400, {370 "error": "Unsupported audio format",371 "error_code": "AUDIO_FORMAT_ERROR",372 "detail": f"Only 16-bit PCM or 32-bit float WAV supported (got {sample_width*8}-bit).",373 })374 if frame_rate < 8000:375 raise HTTPException(400, {376 "error": "Unsupported audio format",377 "error_code": "AUDIO_RATE_ERROR",378 "detail": f"Sample rate too low: {frame_rate} Hz (minimum 8000 Hz).",379 })380 381 # 4. Duration checks382 if n_frames == 0:383 raise HTTPException(400, {384 "error": "Invalid audio",385 "error_code": "AUDIO_EMPTY",386 "detail": "WAV file contains no audio frames.",387 })388 duration = n_frames / frame_rate389 if duration < MIN_AUDIO_SECONDS:390 raise HTTPException(400, {391 "error": "Audio too short",392 "error_code": "AUDIO_TOO_SHORT",393 "detail": f"Audio is {duration:.2f}s — minimum is {MIN_AUDIO_SECONDS}s.",394 })395 if duration > MAX_AUDIO_SECONDS:396 raise HTTPException(400, {397 "error": "Audio too long",398 "error_code": "AUDIO_TOO_LONG",399 "detail": f"Audio is {duration:.1f}s — maximum is {MAX_AUDIO_SECONDS}s.",400 })401 402 # 5. Normalise to mono 16-bit PCM (stereo downmix + format conversion in one pass)403 wav_normalised = _normalise_wav(raw)404 return wav_normalised, duration405 406def _parse_textgrid(tg_path: Path) -> list:407 tg = tgio.openTextgrid(str(tg_path), includeEmptyIntervals=False)408 names = tg.tierNames409 tier = next((n for n in names if "word" in n.lower()), names[0] if names else None)410 if not tier:411 return []412 words = []413 for iv in tg.getTier(tier).entries:414 label = iv.label.strip()415 if not label or label.lower() in {"sp", "sil", "<eps>", "spn"}:416 continue417 words.append({418 "word": label,419 "start": round(float(iv.start), 3),420 "end": round(float(iv.end), 3),421 })422 return words423 424def _job_view(job_id: str, job: Dict[str, Any], slim: bool = False) -> Dict[str, Any]:425 """426 Canonical response body for /job/{id} and /result/{id}.427 slim=True omits word_timestamps from COMPLETED responses (use during polling).428 """429 status = job["status"]430 base = {431 "job_id": job_id,432 "status": status,433 "created_at": job["created_at"].isoformat(),434 }435 436 if status == JobStatus.QUEUED:437 base.update({438 "queue_position": 1,439 "estimated_wait_s": _avg_align_s(),440 "message": "Queued at position 1. Poll /job/{id} every 2 s.",441 })442 443 elif status == JobStatus.RUNNING:444 elapsed = (datetime.utcnow() - job["started_at"]).total_seconds() if job.get("started_at") else None445 base.update({446 "started_at": job["started_at"].isoformat() if job.get("started_at") else None,447 "elapsed_s": round(elapsed, 1) if elapsed is not None else None,448 "estimated_remaining_s": max(0, round(_avg_align_s() - (elapsed or 0), 1)),449 "message": "Running.",450 })451 452 elif status == JobStatus.COMPLETED:453 base.update({454 "total_words": len(job["result"]),455 "language": job["request"].language,456 "pass_number": job["request"].pass_number,457 "model": f"MFA/{MFA_ACOUSTIC}",458 "fine_tuned": False,459 "started_at": job["started_at"].isoformat() if job.get("started_at") else None,460 "completed_at": job["terminal_at"].isoformat() if job.get("terminal_at") else None,461 "duration_s": round(462 (job["terminal_at"] - job["started_at"]).total_seconds(), 2463 ) if job.get("started_at") and job.get("terminal_at") else None,464 })465 # Only include the full payload if not in slim mode466 if not slim:467 base["word_timestamps"] = job["result"]468 469 elif status in (JobStatus.FAILED, JobStatus.ABANDONED):470 base.update({471 "error": job.get("error", "Unknown error"),472 "completed_at": job["terminal_at"].isoformat() if job.get("terminal_at") else None,473 })474 475 return base476 477 478# ═══════════════════════════════════════════════════════════════════════════════479# PERSISTENT WORKER PROCESS480# ═══════════════════════════════════════════════════════════════════════════════481 482def _worker_main(work_q: MPQueue, result_q: MPQueue, work_root: str):483 """484 ═══════════════════════════════════════════════════════════════════════485 FIX v5.5-HOTFIX: The "Hello World" bug had 4 interlocking causes:486 487 1. Warm-up and real jobs shared the same corpus directory.488 2. PretrainedAligner.setup() caches corpus/utterance tables in memory.489 Overwriting utt.lab on disk does NOT update the in-memory table.490 3. _clear_dir() only deleted one level; Kaldi MFCC/lattice subdirs491 in TMP_DIR persisted and were reused.492 4. persistent_aligner reuse across jobs carried stale corpus state.493 494 THE FIX (all 4 addressed):495 ① Warm-up uses a SEPARATE _warmup/ directory, never shared with jobs.496 ② The warm-up aligner is explicitly del'd — its corpus state dies.497 ③ Each job uses shutil.rmtree (full recursive wipe) on OUTPUT and TMP.498 ④ Each job creates a brand-new PretrainedAligner instance (no reuse).499 ═══════════════════════════════════════════════════════════════════════500 """501 import logging as _lm502 import wave as _wave503 _lm.basicConfig(level=_lm.INFO)504 log = _lm.getLogger("mfa_worker")505 506 os.makedirs(work_root, exist_ok=True)507 508 # ── Job corpus paths (ONLY used for real jobs, never warm-up) ──────────509 CORPUS_ROOT = Path(work_root) / "_corpus"510 SPEAKER_DIR = CORPUS_ROOT / "corpus" / "speaker1"511 OUTPUT_DIR = CORPUS_ROOT / "output"512 TMP_DIR = CORPUS_ROOT / "tmp"513 WAV_PATH = SPEAKER_DIR / "utt.wav"514 LAB_PATH = SPEAKER_DIR / "utt.lab"515 for d in [SPEAKER_DIR, OUTPUT_DIR, TMP_DIR]:516 d.mkdir(parents=True, exist_ok=True)517 518 # ── Warm-up: ISOLATED directory + aligner, discarded after ─────────────519 # Uses _warmup/ which is NEVER touched by real jobs.520 # The warmup_aligner instance is del'd so its corpus tables are freed.521 log.info("Worker: starting warm-up (isolated — will not affect real jobs)...")522 try:523 WARMUP_ROOT = Path(work_root) / "_warmup"524 WARMUP_SPEAKER = WARMUP_ROOT / "corpus" / "speaker1"525 WARMUP_OUT = WARMUP_ROOT / "output"526 WARMUP_TMP = WARMUP_ROOT / "tmp"527 for d in [WARMUP_SPEAKER, WARMUP_OUT, WARMUP_TMP]:528 d.mkdir(parents=True, exist_ok=True)529 530 # Write 0.5s silence + "hello world" text into the WARMUP corpus531 silent = np.zeros(8000, dtype=np.int16)532 with _wave.open(str(WARMUP_SPEAKER / "utt.wav"), "wb") as wf:533 wf.setnchannels(1)534 wf.setsampwidth(2)535 wf.setframerate(16000)536 wf.writeframes(silent.tobytes())537 (WARMUP_SPEAKER / "utt.lab").write_text("hello world", encoding="utf-8")538 539 warmup_aligner = PretrainedAligner(540 corpus_directory = str(WARMUP_ROOT / "corpus"),541 dictionary_path = MFA_DICTIONARY,542 acoustic_model_path = MFA_ACOUSTIC,543 output_directory = str(WARMUP_OUT),544 temporary_directory = str(WARMUP_TMP),545 mfa_root_dir = MFA_ROOT,546 fine_tune = False,547 jobs = 1,548 overwrite = True,549 )550 warmup_aligner.setup()551 # Intentionally NOT calling warmup_aligner.align() —552 # setup() alone loads the acoustic model into memory (the slow part).553 # align() on silent audio would fail or produce garbage anyway.554 555 # Discard the warmup aligner — its corpus state ("hello world")556 # must NOT survive into job processing.557 del warmup_aligner558 559 # Clean up warmup corpus files — not needed anymore560 shutil.rmtree(str(WARMUP_ROOT), ignore_errors=True)561 562 log.info("Worker: warm-up complete — warmup aligner discarded.")563 result_q.put(("ready",))564 565 except Exception as e:566 log.warning(f"Worker: warm-up failed (non-fatal, first job will be slower): {e}")567 result_q.put(("ready",)) # always signal ready so service doesn't stall568 569 # ── Job loop ─────────────────────────────────────────────────────────────570 while True:571 try:572 msg = work_q.get()573 if msg is None:574 log.info("Worker: shutdown signal — exiting.")575 break576 577 cmd, job_id, payload = msg578 579 if cmd == "cancel":580 log.info(f"[{job_id}] Cancel received — skipping.")581 continue582 583 wav_path_str = payload["wav_path"]584 clean_text = payload["text"]585 wav_path = Path(wav_path_str)586 587 log.info(f"[{job_id}] Aligning: '{clean_text[:60]}...' (wav={wav_path_str})")588 t_start = time.monotonic()589 590 try:591 # ── Step 1: Write fresh corpus files ─────────────────────────592 shutil.copy2(wav_path_str, str(WAV_PATH))593 wav_path.unlink(missing_ok=True) # clean up IPC temp file594 LAB_PATH.write_text(clean_text, encoding="utf-8")595 596 # ── Step 2: FULL recursive wipe of output + tmp ──────────────597 # _clear_dir() only removes top-level children.598 # Kaldi creates nested subdirs (split1/speaker1/mfcc/, etc.)599 # that must be fully destroyed so MFA can't reuse stale data.600 shutil.rmtree(str(OUTPUT_DIR), ignore_errors=True)601 shutil.rmtree(str(TMP_DIR), ignore_errors=True)602 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)603 TMP_DIR.mkdir(parents=True, exist_ok=True)604 605 # ── Step 3: Fresh aligner instance for this job ───────────────606 # CRITICAL: Do NOT reuse a PretrainedAligner across jobs.607 # The corpus/utterance tables are cached on the instance object.608 # setup() on a reused instance does NOT re-read corpus files.609 # A new instance + new setup() reads the current utt.lab/utt.wav.610 job_aligner = PretrainedAligner(611 corpus_directory = str(CORPUS_ROOT / "corpus"),612 dictionary_path = MFA_DICTIONARY,613 acoustic_model_path = MFA_ACOUSTIC,614 output_directory = str(OUTPUT_DIR),615 temporary_directory = str(TMP_DIR),616 mfa_root_dir = MFA_ROOT,617 fine_tune = False,618 jobs = 1,619 overwrite = True,620 )621 job_aligner.setup() # reads current utt.wav + utt.lab622 job_aligner.align()623 job_aligner.export_files(str(OUTPUT_DIR))624 del job_aligner # free memory immediately after use625 626 # ── Step 4: Parse TextGrid output ─────────────────────────────627 tg_files = list(OUTPUT_DIR.rglob("*.TextGrid"))628 if not tg_files:629 raise RuntimeError("MFA produced no TextGrid output — "630 "check that text contains only dictionary words")631 632 words = _parse_textgrid(tg_files[0])633 duration = round(time.monotonic() - t_start, 2)634 result_q.put(("success", job_id, words, duration))635 log.info(f"[{job_id}] Done — {len(words)} words in {duration}s.")636 637 except Exception as e:638 # Clean up IPC temp file if it still exists (e.g. copy failed)639 wav_path.unlink(missing_ok=True)640 log.exception(f"[{job_id}] Alignment error")641 result_q.put(("error", job_id, str(e), 0))642 # No aligner rebuild needed — next job always creates a fresh one643 644 except Exception as outer:645 log.error(f"Worker outer loop error (continuing): {outer}")646 647 648# ═══════════════════════════════════════════════════════════════════════════════649# WORKER LIFECYCLE650# ═══════════════════════════════════════════════════════════════════════════════651 652_worker_process: Optional[Process] = None653 654def _start_worker():655 global _worker_process, _worker_ready656 _worker_ready = False657 p = Process(658 target=_worker_main,659 args=(_work_queue, _result_queue, WORK_ROOT),660 daemon=False, # REQUIRED — MFA spawns Kaldi child processes661 )662 p.start()663 logger.info(f"Worker process started (pid={p.pid})")664 _worker_process = p665 666def _restart_worker():667 """668 Kill current worker + start fresh.669 MUST be called OUTSIDE jobs_lock — terminate()+join() can block several670 seconds and would stall all concurrent requests if the lock was held.671 Always call via run_in_executor from async context.672 """673 global _worker_process, _worker_restarting, _worker_ready674 _worker_restarting = True675 _worker_ready = False676 try:677 old = _worker_process678 if old and old.is_alive():679 old.terminate()680 old.join(timeout=3)681 if old.is_alive():682 old.kill()683 old.join(timeout=1)684 logger.info("Old worker terminated.")685 shutil.rmtree(str(Path(WORK_ROOT) / "_corpus"), ignore_errors=True)686 _metrics["worker_restarts"] += 1687 _start_worker()688 finally:689 _worker_restarting = False690 691@app.on_event("startup")692async def startup():693 _metrics["start_time"] = datetime.utcnow()694 os.makedirs(WORK_ROOT, exist_ok=True)695 logger.info(f"Work dir: {WORK_ROOT} RAM free: {_shm_free_mb()} MB")696 _start_worker()697 asyncio.create_task(_collect_results())698 asyncio.create_task(_cleanup_loop())699 logger.info("MFA Aligner Service v5.5 ready.")700 701@app.on_event("shutdown")702async def shutdown():703 global _worker_process704 if _worker_process and _worker_process.is_alive():705 _work_queue.put(None)706 _worker_process.join(timeout=5)707 if _worker_process.is_alive():708 _worker_process.terminate()709 _worker_process.join(timeout=2)710 logger.info("Worker shut down.")711 712 713# ═══════════════════════════════════════════════════════════════════════════════714# RESULT COLLECTOR715# ═══════════════════════════════════════════════════════════════════════════════716 717async def _collect_results():718 """719 Drains result_queue in a non-blocking poll loop.720 Uses get_nowait() + asyncio.sleep — no thread permanently blocked on queue.721 """722 while True:723 try:724 try:725 msg = _result_queue.get_nowait()726 except _stdlib_queue.Empty:727 await asyncio.sleep(0.1)728 continue729 730 # Worker ready signal (warm-up complete)731 if msg[0] == "ready":732 global _worker_ready733 _worker_ready = True734 logger.info("Worker warm-up complete — worker_ready=True.")735 continue736 737 status, job_id, data, duration_s = msg738 739 async with jobs_lock:740 if job_id not in jobs:741 logger.warning(f"Result for unknown/purged job {job_id} — discarded.")742 continue743 now = datetime.utcnow()744 if status == "success":745 jobs[job_id]["status"] = JobStatus.COMPLETED746 jobs[job_id]["result"] = data747 jobs[job_id]["terminal_at"] = now748 _metrics["total_completed"] += 1749 _recent_durations.append(duration_s)750 logger.info(f"Job {job_id} completed — {len(data)} words in {duration_s}s.")751 else:752 jobs[job_id]["status"] = JobStatus.FAILED753 jobs[job_id]["error"] = data754 jobs[job_id]["terminal_at"] = now755 _metrics["total_failed"] += 1756 logger.error(f"Job {job_id} failed: {data}")757 await _promote_queued()758 759 except Exception as e:760 logger.error(f"Result collector error: {e}")761 await asyncio.sleep(0.1)762 763 764# ═══════════════════════════════════════════════════════════════════════════════765# CLEANUP LOOP766# ═══════════════════════════════════════════════════════════════════════════════767 768async def _cleanup_loop():769 while True:770 await asyncio.sleep(CLEANUP_INTERVAL)771 need_restart = False772 773 async with jobs_lock:774 now = datetime.utcnow()775 ttl_cutoff = now - timedelta(seconds=JOB_TTL_SECONDS)776 777 # 2. Purge terminal jobs past TTL.778 expired = [779 jid for jid, info in jobs.items()780 if info["status"] in TERMINAL_STATUSES781 and (info.get("terminal_at") or now) < ttl_cutoff782 ]783 for jid in expired:784 logger.info(f"Purging TTL-expired job {jid}.")785 del jobs[jid]786 787 if expired:788 logger.info(f"Cleanup: purged {len(expired)} jobs. Remaining: {len(jobs)}")789 790 await _promote_queued()791 792 # 3. Restart worker OUTSIDE jobs_lock.793 if need_restart:794 logger.info("Restarting worker (outside jobs_lock).")795 loop = asyncio.get_running_loop()796 await loop.run_in_executor(None, _restart_worker)797 798 799async def _promote_queued() -> Optional[str]:800 """801 Promote oldest queued job → RUNNING.802 Must be called with jobs_lock held.803 804 Grace period: last_heartbeat = now + HEARTBEAT_GRACE (+ WORKER_RESTART_GRACE805 if a restart is in progress) so the frontend has time to detect the806 transition and start heartbeating before the cleanup loop fires.807 """808 if any(info["status"] == JobStatus.RUNNING for info in jobs.values()):809 return None810 queued = sorted(811 [(jid, info) for jid, info in jobs.items() if info["status"] == JobStatus.QUEUED],812 key=lambda x: x[1]["created_at"],813 )814 if not queued:815 return None816 jid, info = queued[0]817 now = datetime.utcnow()818 info["status"] = JobStatus.RUNNING819 info["started_at"] = now820 logger.info(f"Promoted queued job {jid} → RUNNING.")821 _work_queue.put(("run", jid, info["_ipc_payload"]))822 return jid823 824 825# ═══════════════════════════════════════════════════════════════════════════════826# ENDPOINTS827# ═══════════════════════════════════════════════════════════════════════════════828 829@app.get("/")830async def root():831 return {832 "name": "MFA Aligner v5.5",833 "version": "5.5.0",834 "status": "ready",835 "model": MFA_ACOUSTIC,836 "endpoints": {837 "POST /align": "Submit alignment job",838 "GET /job/{id}": "Unified status + result [preferred]",839 "GET /job/{id}?slim=true": "Status only — omits word_timestamps (use while polling)",840 "GET /result/{id}": "Legacy result poll [compat]",841 "GET /queue/status": "Queue availability",842 "POST /cancel/{id}": "Explicit job cancellation",843 "GET /health": "Worker health + job counts",844 "GET /metrics": "Aggregate metrics",845 },846 }847 848 849@app.get("/health")850async def health():851 worker_alive = _worker_process is not None and _worker_process.is_alive()852 shm_free = _shm_free_mb()853 shm_ok = shm_free is None or shm_free >= SHM_MIN_FREE_MB854 855 async with jobs_lock:856 counts = {s: 0 for s in [857 JobStatus.QUEUED, JobStatus.RUNNING,858 JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.ABANDONED,859 ]}860 for info in jobs.values():861 counts[info["status"]] = counts.get(info["status"], 0) + 1862 863 overall = "healthy"864 if not worker_alive or _worker_restarting:865 overall = "degraded"866 if not shm_ok:867 overall = "degraded"868 869 return {870 "status": overall,871 "worker_alive": worker_alive,872 "worker_ready": _worker_ready, # False while warming up873 "worker_restarting": _worker_restarting,874 "model": MFA_ACOUSTIC,875 "shm_free_mb": shm_free,876 "shm_ok": shm_ok,877 "avg_align_s": _avg_align_s(),878 "job_counts": counts,879 "total_jobs": sum(counts.values()),880 }881 882 883@app.get("/metrics")884async def metrics():885 """Aggregate lifetime metrics — useful for dashboards and debugging."""886 uptime = (887 round((datetime.utcnow() - _metrics["start_time"]).total_seconds(), 0)888 if _metrics["start_time"] else None889 )890 total = _metrics["total_completed"] + _metrics["total_failed"] + _metrics["total_abandoned"]891 success_rate = round(_metrics["total_completed"] / total, 3) if total > 0 else None892 893 return {894 "uptime_s": uptime,895 "total_submitted": _metrics["total_submitted"],896 "total_completed": _metrics["total_completed"],897 "total_failed": _metrics["total_failed"],898 "total_abandoned": _metrics["total_abandoned"],899 "total_cancelled": _metrics["total_cancelled"],900 "success_rate": success_rate,901 "worker_restarts": _metrics["worker_restarts"],902 "avg_align_s": _avg_align_s(),903 "recent_durations": list(_recent_durations),904 }905 906 907@app.get("/queue/status")908async def queue_status():909 async with jobs_lock:910 running_job = next(911 (jid for jid, info in jobs.items() if info["status"] == JobStatus.RUNNING), None912 )913 queued_job = next(914 (jid for jid, info in jobs.items() if info["status"] == JobStatus.QUEUED), None915 )916 917 running_count = 1 if running_job else 0918 waiting_count = 1 if queued_job else 0919 920 return {921 # ── NEW: integer counts the controller reads ──────────────────────922 "running": running_count,923 "waiting": waiting_count,924 # ── Existing fields kept for backwards compat ─────────────────────925 "can_submit": running_count == 0 and waiting_count == 0,926 "slot_available": running_count == 0,927 "queue_full": running_count == 1 and waiting_count == 1,928 "running_job_id": running_job,929 "queued_job_id": queued_job,930 "worker_ready": _worker_ready,931 "estimated_wait_s": _avg_align_s() if running_job else 0,932 }933 934 935@app.post("/align")936async def align_audio(req: AlignRequest):937 """938 Submit an alignment job.939 940 All validation (audio format, duration, text) is performed here before941 the job enters the queue. Invalid requests get an immediate HTTP 400.942 943 Response fields:944 job_id — use for all subsequent calls945 status — "running" | "queued"946 queue_position — 0 = started now, 1 = queued behind running job947 poll_url — GET here for status / result948 estimated_wait_s — real rolling average, 0 if running949 duplicate_warning — present if client_id already has an active job950 951 Frontend flow:952 queue_position==0: poll for "completed"953 queue_position==1: poll /job/{id} every 2 s until status=="running".954 503 QUEUE_FULL: wait retry_after seconds, then retry955 """956 # ── 1. Text validation ────────────────────────────────────────────────────957 cleaned_text = _clean(req.text)958 if not cleaned_text:959 raise HTTPException(400, {960 "error": "Invalid text",961 "error_code": "TEXT_EMPTY",962 "detail": "Text is empty after removing unsupported characters.",963 })964 965 # ── 2. Audio validation (blocking CPU work — run in executor) ────────────966 loop = asyncio.get_running_loop()967 try:968 wav_normalised, audio_duration = await loop.run_in_executor(969 None, _validate_audio, req.audio_base64970 )971 except HTTPException:972 raise # re-raise validation errors as-is973 974 # ── 3. System checks ──────────────────────────────────────────────────────975 shm_free = _shm_free_mb()976 if shm_free is not None and shm_free < SHM_MIN_FREE_MB:977 raise HTTPException(503, {978 "error": "Insufficient RAM",979 "error_code": "SHM_FULL",980 "detail": f"Only {shm_free} MB free in /dev/shm (minimum {SHM_MIN_FREE_MB} MB).",981 })982 983 worker_alive = _worker_process is not None and _worker_process.is_alive()984 if not worker_alive and not _worker_restarting:985 raise HTTPException(503, {986 "error": "Worker unavailable",987 "error_code": "WORKER_DOWN",988 "detail": "Alignment worker is not running. Try again in a few seconds.",989 })990 991 async with jobs_lock:992 running_exists = any(info["status"] == JobStatus.RUNNING for info in jobs.values())993 queued_exists = any(info["status"] == JobStatus.QUEUED for info in jobs.values())994 995 if running_exists and queued_exists:996 return JSONResponse(997 status_code=503,998 content={999 "error": "Service busy",1000 "error_code": "QUEUE_FULL",1001 "detail": "One job running, one already queued. Retry after current jobs finish.",1002 "retry_after": _avg_align_s(),1003 },1004 headers={"Retry-After": str(int(_avg_align_s()))},1005 )1006 1007 # ── 4. Duplicate client check ─────────────────────────────────────────1008 duplicate_warning = None1009 if req.client_id:1010 existing = next(1011 (jid for jid, info in jobs.items()1012 if info.get("client_id") == req.client_id1013 and info["status"] in (JobStatus.RUNNING, JobStatus.QUEUED)),1014 None1015 )1016 if existing:1017 duplicate_warning = {1018 "message": "You already have an active job.",1019 "existing_job_id": existing,1020 "existing_status": jobs[existing]["status"],1021 }1022 1023 # ── 5. Write audio to /dev/shm temp file ─────────────────────────────1024 # Passes only a file path through MPQueue instead of pickling audio bytes.1025 tmp_wav_path = Path(WORK_ROOT) / f"tmp_{uuid.uuid4().hex[:8]}.wav"1026 try:1027 tmp_wav_path.write_bytes(wav_normalised)1028 except Exception as e:1029 raise HTTPException(503, {1030 "error": "Storage error",1031 "error_code": "SHM_WRITE_ERROR",1032 "detail": f"Could not write audio to /dev/shm: {e}",1033 })1034 1035 status = JobStatus.QUEUED if running_exists else JobStatus.RUNNING1036 job_id = uuid.uuid4().hex[:12]1037 queue_position = 1 if status == JobStatus.QUEUED else 01038 now = datetime.utcnow()1039 1040 # IPC payload: path + pre-cleaned text only — no audio bytes1041 ipc_payload = {1042 "wav_path": str(tmp_wav_path),1043 "text": cleaned_text,1044 }1045 1046 try:1047 jobs[job_id] = {1048 "request": req,1049 "status": status,1050 "created_at": now,1051 "started_at": now if status == JobStatus.RUNNING else None,1052 "terminal_at": None,1053 "result": None,1054 "error": None,1055 "client_id": req.client_id,1056 "audio_duration": round(audio_duration, 2),1057 "_ipc_payload": ipc_payload,1058 }1059 1060 _metrics["total_submitted"] += 11061 logger.info(f"Job {job_id} created — status={status}, duration={audio_duration:.1f}s, client={req.client_id}")1062 1063 if status == JobStatus.RUNNING:1064 _work_queue.put(("run", job_id, ipc_payload))1065 1066 except Exception:1067 # Prevent temp WAV leak if job registration fails unexpectedly1068 tmp_wav_path.unlink(missing_ok=True)1069 raise1070 1071 response = {1072 "job_id": job_id,1073 "status": status,1074 "queue_position": queue_position,1075 "poll_url": f"/job/{job_id}",1076 "estimated_wait_s": 0 if status == JobStatus.RUNNING else _avg_align_s(),1077 "audio_duration_s": round(audio_duration, 2),1078 "message": (1079 f"Job started. Poll GET /job/{job_id}."1080 if status == JobStatus.RUNNING else1081 f"Job queued at position 1. Poll GET /job/{job_id} every 2 s."1082 ),1083 }1084 if duplicate_warning:1085 response["duplicate_warning"] = duplicate_warning1086 1087 return JSONResponse(response)1088 1089 1090@app.get("/job/{job_id}")1091async def get_job(job_id: str, slim: bool = Query(False)):1092 """1093 Unified status + result endpoint.1094 1095 ?slim=true — omits word_timestamps from COMPLETED responses.1096 Use this during the polling loop; fetch without slim=true1097 once you see status=="completed" to retrieve the full result.1098 1099 Frontend state machine:1100 "queued" → poll every 2 s (slim=true)1101 "running" → keep polling (slim=true)1102 "completed" → fetch /job/{id} (slim=false) for timestamps1103 "failed" → show error, consider resubmitting1104 "abandoned" → resubmit1105 """1106 async with jobs_lock:1107 if job_id not in jobs:1108 raise HTTPException(404, detail={1109 "error": "Job not found",1110 "error_code": "JOB_NOT_FOUND",1111 "detail": "Job may have been purged (TTL=5 min) or never existed.",1112 })1113 return _job_view(job_id, jobs[job_id], slim=slim)1114 1115 1116@app.get("/result/{job_id}")1117async def get_result(job_id: str):1118 """Legacy endpoint — kept for backwards compatibility. Use GET /job/{id} instead."""1119 async with jobs_lock:1120 if job_id not in jobs:1121 raise HTTPException(404, "Job not found")1122 return _job_view(job_id, jobs[job_id])1123 1124 1125@app.post("/cancel/{job_id}")1126async def cancel_job(job_id: str, body: CancelRequest):1127 """1128 Explicitly cancel a job.1129 1130 RUNNING: marks abandoned, sends cancel signal to worker, restarts worker1131 via run_in_executor (non-blocking — doesn't stall the event loop).1132 QUEUED: removes cleanly from queue, no restart.1133 Terminal: idempotent 200.1134 """1135 need_restart = False1136 1137 async with jobs_lock:1138 if job_id not in jobs:1139 raise HTTPException(404, detail={1140 "error": "Job not found",1141 "error_code": "JOB_NOT_FOUND",1142 })1143 1144 job = jobs[job_id]1145 status = job["status"]1146 1147 stored = job.get("client_id")1148 if stored and (body.client_id is None or body.client_id != stored):1149 raise HTTPException(403, detail={1150 "error": "Client ID mismatch",1151 "error_code": "FORBIDDEN",1152 "detail": "You are not the owner of this job.",1153 })1154 1155 if status in TERMINAL_STATUSES:1156 return {1157 "cancelled": False,1158 "status": status,1159 "message": f"Job was already {status} — nothing to cancel.",1160 }1161 1162 now = datetime.utcnow()1163 jobs[job_id]["status"] = JobStatus.ABANDONED1164 jobs[job_id]["error"] = "Cancelled by client"1165 jobs[job_id]["terminal_at"] = now1166 _metrics["total_cancelled"] += 11167 1168 if status == JobStatus.RUNNING:1169 _work_queue.put(("cancel", job_id, None))1170 need_restart = True1171 logger.info(f"Job {job_id} cancelled (RUNNING) — scheduling worker restart.")1172 else:1173 # Clean up the temp WAV file if it hasn't been consumed yet1174 ipc = job.get("_ipc_payload", {})1175 wav_path = ipc.get("wav_path")1176 if wav_path:1177 Path(wav_path).unlink(missing_ok=True)1178 logger.info(f"Job {job_id} cancelled (QUEUED) — removed cleanly.")1179 1180 await _promote_queued()1181 1182 if need_restart:1183 # Non-blocking: run_in_executor so the cancel response returns immediately1184 loop = asyncio.get_running_loop()1185 asyncio.ensure_future(loop.run_in_executor(None, _restart_worker))1186 1187 return {1188 "cancelled": True,1189 "job_id": job_id,1190 "restarting": need_restart,1191 "message": "Job cancelled." + (" Worker restart initiated." if need_restart else ""),1192 }