Merge pull request #24 from MapleEve/fix/model-load-timing

chore: log model load timings

Author: MapleEve

app/application/transcription_jobs.py

@@ -1,6 +1,7 @@
 """Application-level transcription job orchestration."""
 
 import logging
+import time
 from pathlib import Path
 
 from config import (
@@ -39,6 +40,7 @@ def _record_status(status: str) -> None:
         extra_filename = audio_path.name if status == "converting" else None
         _write_status(job_id, status, filename=extra_filename)
 
+    job_started = time.perf_counter()
     try:
 
         def _process_pipeline():
@@ -69,16 +71,20 @@ def _process_pipeline():
         jobs[job_id]["result"] = tr
         _write_status(job_id, "completed")
         logger.info(
-            "Job %s completed: %d segments, %d speakers",
-            job_id,
+            "transcription_job_timing status=completed elapsed_s=%.3f segment_count=%d speaker_count=%d",
+            time.perf_counter() - job_started,
             len(tr.get("segments", [])),
             len(tr.get("speaker_map", {})),
         )
         if file_hash:
             unregister_in_flight(file_hash, job_id)
 
     except Exception as e:
-        logger.exception("Job %s failed", job_id)
+        logger.exception(
+            "transcription_job_timing status=failed elapsed_s=%.3f error_type=%s",
+            time.perf_counter() - job_started,
+            e.__class__.__name__,
+        )
         jobs[job_id]["status"] = "failed"
         jobs[job_id]["error"] = str(e)
         _write_status(job_id, "failed", error=str(e))

app/pipeline/orchestrator.py

@@ -14,6 +14,7 @@
 import json
 import re
 import tempfile
+import time
 from contextlib import nullcontext
 from pathlib import Path
 from typing import Any
@@ -336,11 +337,23 @@ def whisper(self):
                 whisper_device,
                 compute_type,
             )
+            load_started = time.perf_counter()
             self._whisper = WhisperModel(
                 model_ref,
                 **_faster_whisper_device_kwargs(whisper_device),
                 compute_type=compute_type,
             )
+            logger.info(
+                "Loaded faster-whisper model in %.2fs (cold_load=True, device=%s, compute_type=%s)",
+                time.perf_counter() - load_started,
+                whisper_device,
+                compute_type,
+            )
+        else:
+            logger.info(
+                "Reusing faster-whisper model (hot reuse, device=%s)",
+                getattr(self, "_whisper_device", None) or getattr(self, "device", ""),
+            )
         return self._whisper
 
     @property
@@ -363,6 +376,7 @@ def diarization(self):
                     token=self.hf_token,
                 )
             logger.info("Loading pyannote diarization model")
+            load_started = time.perf_counter()
             self._diarization = _load_trusted_pyannote_model(
                 PyannotePipeline.from_pretrained,
                 model_ref,
@@ -371,6 +385,11 @@ def diarization(self):
             _dev = diarization_device if ":" in diarization_device else "cuda:0"
             if diarization_device.startswith("cuda"):
                 self._diarization.to(torch.device(_dev))
+            logger.info(
+                "Loaded pyannote diarization model in %.2fs (cold_load=True, device=%s)",
+                time.perf_counter() - load_started,
+                diarization_device,
+            )
             # Suppress over-segmentation of short backchannel turns
             try:
                 if hasattr(self._diarization, "_binarize") and hasattr(
@@ -385,6 +404,12 @@ def diarization(self):
                     )
             except Exception as exc:
                 logger.warning("Could not set min_duration_off: %s", exc)
+        else:
+            logger.info(
+                "Reusing pyannote diarization model (hot reuse, device=%s)",
+                getattr(self, "_diarization_device", None)
+                or getattr(self, "device", ""),
+            )
         return self._diarization
 
     @property
@@ -400,6 +425,7 @@ def embedding_model(self):
             )
             model_ref = _resolve_local_pyannote_file(model_ref, "pytorch_model.bin")
             logger.info("Loading WeSpeaker speaker encoder")
+            load_started = time.perf_counter()
             model = _load_trusted_pyannote_model(
                 Model.from_pretrained,
                 model_ref,
@@ -409,6 +435,16 @@ def embedding_model(self):
             # window="whole" returns one embedding vector per full chunk —
             # exactly what we need for per-turn embeddings.
             self._embedding_model = Inference(model, window="whole")
+            logger.info(
+                "Loaded WeSpeaker speaker encoder in %.2fs (cold_load=True, device=%s)",
+                time.perf_counter() - load_started,
+                embedding_device,
+            )
+        else:
+            logger.info(
+                "Reusing WeSpeaker speaker encoder (hot reuse, device=%s)",
+                getattr(self, "_embedding_device", None) or getattr(self, "device", ""),
+            )
         return self._embedding_model
 
     def transcribe(

app/pipeline/runner.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import time
 from typing import Any
 
 from infra.audio import cleanup_generated_files
@@ -15,6 +16,35 @@
 DEFAULT_STAGE_ORDER = available_stage_slots()
 
 
+def _safe_stage_metrics(context: PipelineContext, stage_name: str) -> dict[str, Any]:
+    metrics: dict[str, Any] = {}
+    stage_metadata = context.metadata.get(stage_name)
+    if isinstance(stage_metadata, dict):
+        for key in (
+            "status",
+            "model",
+            "language",
+            "segment_count",
+            "speaker_count",
+            "turn_count",
+            "applied",
+            "reason",
+            "persisted",
+        ):
+            if key in stage_metadata:
+                metrics[key] = stage_metadata[key]
+    if "segment_count" not in metrics and context.aligned_segments:
+        metrics["segment_count"] = len(context.aligned_segments)
+    if "speaker_count" not in metrics:
+        if context.speaker_embeddings:
+            metrics["speaker_count"] = len(context.speaker_embeddings)
+        elif context.diarization_turns:
+            metrics["speaker_count"] = len(
+                {turn.get("speaker") for turn in context.diarization_turns}
+            )
+    return metrics
+
+
 class PipelineRunner:
     """Execute the stable stage order against the current pipeline implementation."""
 
@@ -42,7 +72,21 @@ def run_context(self, pipeline: Any, request: PipelineRequest) -> PipelineContex
                 context.metadata.setdefault("selected_providers", {})[stage_name] = (
                     request.provider_for(stage_name)
                 )
+                stage_started = time.perf_counter()
                 stage(context)
+                elapsed_s = time.perf_counter() - stage_started
+                metrics = _safe_stage_metrics(context, stage_name)
+                context.metadata.setdefault("stage_timings", {})[stage_name] = round(
+                    elapsed_s,
+                    3,
+                )
+                logger.info(
+                    "pipeline_stage_timing stage=%s elapsed_s=%.3f provider=%s metrics=%s",
+                    stage_name,
+                    elapsed_s,
+                    request.provider_for(stage_name),
+                    metrics,
+                )
             return context
         finally:
             cleanup_generated_files(context.temporary_paths)

app/providers/asr/default.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import time
 from collections import Counter
 from typing import Any
 
@@ -191,7 +192,9 @@ def run_faster_whisper_asr(
     if no_repeat_ngram_size and no_repeat_ngram_size >= 3:
         whisper_kwargs["no_repeat_ngram_size"] = no_repeat_ngram_size
 
-    segments_iter, info = pipeline.whisper.transcribe(audio_path, **whisper_kwargs)
+    whisper_model = pipeline.whisper
+    processing_started = time.perf_counter()
+    segments_iter, info = whisper_model.transcribe(audio_path, **whisper_kwargs)
     raw_segments = [
         {
             "start": round(float(segment.start), 3),
@@ -200,8 +203,18 @@ def run_faster_whisper_asr(
         }
         for segment in segments_iter
     ]
+    processing_elapsed_s = time.perf_counter() - processing_started
     segments, hallucination_guard = suppress_repetition_hallucinations(raw_segments)
     detected_language = info.language
+    audio_duration_s = max((segment["end"] for segment in raw_segments), default=0.0)
+    logger.info(
+        "asr_processing_timing model=faster-whisper elapsed_s=%.3f language=%s segment_count=%d raw_segment_count=%d duration_s=%.3f",
+        processing_elapsed_s,
+        detected_language,
+        len(segments),
+        len(raw_segments),
+        audio_duration_s,
+    )
     logger.info(
         "Transcription done: %d segments, language=%s, repetition_guard=%s",
         len(segments),

app/providers/diarization/default.py

@@ -5,6 +5,7 @@
 import logging
 import os
 import re
+import time
 from contextlib import contextmanager
 from collections.abc import Callable
 from inspect import Parameter, signature
@@ -186,7 +187,9 @@ def run_pyannote_diarization(
     if max_speakers:
         kwargs["max_speakers"] = max_speakers
 
-    result = pipeline.diarization(audio_path, **kwargs)
+    diarization_model = pipeline.diarization
+    processing_started = time.perf_counter()
+    result = diarization_model(audio_path, **kwargs)
     turns: list[dict[str, object]] = []
     for turn, _, speaker in result.itertracks(yield_label=True):
         turns.append(
@@ -196,6 +199,15 @@ def run_pyannote_diarization(
                 "speaker": speaker,
             }
         )
+    elapsed_s = time.perf_counter() - processing_started
+    logger.info(
+        "diarization_processing_timing model=pyannote elapsed_s=%.3f device=%s turn_count=%d speaker_count=%d",
+        elapsed_s,
+        getattr(pipeline, "_diarization_device", None)
+        or getattr(pipeline, "device", ""),
+        len(turns),
+        len({turn["speaker"] for turn in turns}),
+    )
     return turns
 
 
@@ -242,10 +254,20 @@ def align_diarized_segments_with_metadata(
             language,
             pipeline.device,
         )
+        load_started = time.perf_counter()
         with _cache_only_alignment_environment():
             align_model, align_metadata = whisperx.load_align_model(
                 **load_kwargs,
             )
+        logger.info(
+            "Loaded WhisperX alignment model in %.2fs "
+            "(cold_load=True, language=%s, model_source=%s, device=%s)",
+            time.perf_counter() - load_started,
+            language,
+            model_source,
+            pipeline.device,
+        )
+        processing_started = time.perf_counter()
         aligned_result = whisperx.align(
             segments,
             align_model,
@@ -254,7 +276,15 @@ def align_diarized_segments_with_metadata(
             pipeline.device,
             return_char_alignments=False,
         )
+        processing_elapsed_s = time.perf_counter() - processing_started
         segments = aligned_result.get("segments", segments)
+        logger.info(
+            "alignment_processing_timing model=whisperx elapsed_s=%.3f language=%s segment_count=%d device=%s",
+            processing_elapsed_s,
+            language,
+            len(segments),
+            pipeline.device,
+        )
         logger.info("WhisperX forced alignment succeeded for language=%s", language)
         metadata = {
             "status": "succeeded",

app/providers/embedding/default.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import time
 
 import numpy as np
 import torchaudio
@@ -50,8 +51,7 @@ def extract_embeddings_for_turns(
             )
         except Exception as exc:
             logger.warning(
-                "Failed to load segment %s [%d:%d]: %s",
-                speaker,
+                "Failed to load embedding audio segment [%d:%d]: %s",
                 start_sample,
                 end_sample,
                 exc,
@@ -66,18 +66,30 @@ def extract_embeddings_for_turns(
         speaker_segments.setdefault(speaker, []).append(chunk)
 
     embeddings: dict[str, np.ndarray] = {}
+    model_processing_elapsed_s = 0.0
+    processed_chunk_count = 0
     for speaker, chunks in speaker_segments.items():
         emb_list = []
         chunks.sort(key=lambda chunk: chunk.shape[1], reverse=True)
         for chunk in chunks[:10]:
             embedding_model = pipeline.embedding_model
             embedding_device = getattr(pipeline, "embedding_device", pipeline.device)
+            processing_started = time.perf_counter()
             emb = embedding_model(
                 {"waveform": chunk.to(embedding_device), "sample_rate": target_sr}
             )
+            model_processing_elapsed_s += time.perf_counter() - processing_started
+            processed_chunk_count += 1
             emb_list.append(np.asarray(emb))
         if emb_list:
             embeddings[speaker] = np.mean(emb_list, axis=0)
+    logger.info(
+        "embedding_processing_timing model=wespeaker elapsed_s=%.3f device=%s speaker_count=%d chunk_count=%d",
+        model_processing_elapsed_s,
+        getattr(pipeline, "embedding_device", getattr(pipeline, "device", "")),
+        len(embeddings),
+        processed_chunk_count,
+    )
     return embeddings