feat: add MLflow inference logging to all Ray Serve apps
All checks were successful
Build and Publish ray-serve-apps / build-and-publish (push) Successful in 16s

- Add mlflow_logger.py: lightweight REST-based MLflow logger (no mlflow dep)
- Instrument serve_llm.py with latency, token counts, tokens/sec metrics
- Instrument serve_embeddings.py with latency, batch_size, total_tokens
- Instrument serve_whisper.py with latency, audio_duration, realtime_factor
- Instrument serve_tts.py with latency, audio_duration, text_chars
- Instrument serve_reranker.py with latency, num_pairs, top_k
This commit is contained in:
2026-02-12 06:14:30 -05:00
parent 2edafc33c0
commit 7ec2107e0c
6 changed files with 346 additions and 4 deletions

View File

@@ -6,10 +6,13 @@ Runs on: elminster (RTX 2070 8GB, CUDA)
import base64
import io
import os
import time
from typing import Any
from ray import serve
from ray_serve.mlflow_logger import InferenceLogger
@serve.deployment(name="WhisperDeployment", num_replicas=1)
class WhisperDeployment:
@@ -38,6 +41,17 @@ class WhisperDeployment:
print("Whisper model loaded successfully")
# MLflow metrics
self._mlflow = InferenceLogger(
experiment_name="ray-serve-whisper",
run_name=f"whisper-{self.model_size}",
tags={"model.name": f"whisper-{self.model_size}", "model.framework": "faster-whisper", "device": self.device},
flush_every=5,
)
self._mlflow.initialize(
params={"model_size": self.model_size, "device": self.device, "compute_type": self.compute_type}
)
async def __call__(self, request: dict[str, Any]) -> dict[str, Any]:
"""
Handle transcription requests.
@@ -59,6 +73,7 @@ class WhisperDeployment:
}
"""
_start = time.time()
language = request.get("language")
task = request.get("task", "transcribe") # transcribe or translate
response_format = request.get("response_format", "json")
@@ -130,6 +145,14 @@ class WhisperDeployment:
"segments": segment_list,
}
# Log to MLflow
self._mlflow.log_request(
latency_s=time.time() - _start,
audio_duration_s=info.duration,
segments=len(segment_list),
realtime_factor=(time.time() - _start) / info.duration if info.duration > 0 else 0,
)
# Default JSON format (OpenAI-compatible)
return {
"text": full_text.strip(),