Experiment Tracking y Model Registry

¿Por Qué Experiment Tracking?

En IA, un cambio en un prompt, un parámetro de chunking o un modelo de embedding puede cambiar drásticamente los resultados. Experiment tracking registra cada variación para reproducir los mejores resultados y entender qué funciona.

Sin tracking:   "Creo que el prompt v3 era mejor... ¿o era v5?"
Con tracking:   Run #47: prompt_v3 + chunk_512 + reranker → faithfulness=0.92 ✓

MLflow: Tracking para LLMOps

Setup e Integración

import mlflow
from datetime import datetime

# Configurar servidor de tracking
mlflow.set_tracking_uri("http://mlflow.internal:5000")
mlflow.set_experiment("rag-pipeline-optimization")


class RAGExperimentTracker:
    """Tracker para experimentos de RAG."""

    def __init__(self, experiment_name: str):
        mlflow.set_experiment(experiment_name)

    def track_rag_experiment(
        self,
        config: dict,
        eval_results: dict,
        sample_queries: list[dict] | None = None,
    ):
        with mlflow.start_run(
            run_name=f"rag-{datetime.now():%Y%m%d-%H%M}"
        ):
            # ── Registrar configuración ──────────────
            mlflow.log_params({
                "embedding_model": config["embedding_model"],
                "chunk_size": config["chunk_size"],
                "chunk_overlap": config["chunk_overlap"],
                "llm_model": config["llm_model"],
                "temperature": config["temperature"],
                "top_k": config["top_k"],
                "reranker": config.get("reranker", "none"),
                "prompt_version": config["prompt_version"],
            })

            # ── Registrar métricas de evaluación ─────
            mlflow.log_metrics({
                "faithfulness": eval_results["faithfulness"],
                "answer_relevancy": eval_results["answer_relevancy"],
                "context_precision": eval_results["context_precision"],
                "context_recall": eval_results["context_recall"],
                "hallucination_rate": eval_results["hallucination_rate"],
                "avg_latency_ms": eval_results["avg_latency_ms"],
                "avg_cost_usd": eval_results["avg_cost_usd"],
                "p95_latency_ms": eval_results["p95_latency_ms"],
            })

            # ── Registrar artefactos ─────────────────
            if sample_queries:
                mlflow.log_dict(
                    {"queries": sample_queries},
                    "sample_queries.json"
                )

            # ── Registrar el prompt como artefacto ───
            mlflow.log_artifact(
                f"prompts/{config['prompt_version']}/system.txt"
            )

            # ── Tags para filtrar ────────────────────
            mlflow.set_tags({
                "team": "ai-engineering",
                "pipeline": "rag",
                "stage": config.get("stage", "development"),
            })


# ── Uso ──────────────────────────────────────────────
tracker = RAGExperimentTracker("rag-optimization-q1")

tracker.track_rag_experiment(
    config={
        "embedding_model": "text-embedding-3-large",
        "chunk_size": 512,
        "chunk_overlap": 50,
        "llm_model": "claude-sonnet-4-20250514",
        "temperature": 0.1,
        "top_k": 5,
        "reranker": "cohere-rerank-v3",
        "prompt_version": "v3",
    },
    eval_results={
        "faithfulness": 0.92,
        "answer_relevancy": 0.88,
        "context_precision": 0.85,
        "context_recall": 0.90,
        "hallucination_rate": 0.03,
        "avg_latency_ms": 1200,
        "avg_cost_usd": 0.008,
        "p95_latency_ms": 2500,
    },
)

Weights & Biases para Evaluación

import wandb

class WandbEvalTracker:
    def __init__(self, project: str):
        self.project = project

    def track_eval_run(self, config: dict, results: list[dict]):
        run = wandb.init(
            project=self.project,
            config=config,
            tags=["evaluation", config.get("pipeline", "rag")],
        )

        # Crear tabla de resultados detallados
        columns = [
            "query", "expected", "actual",
            "faithfulness", "relevancy", "latency_ms"
        ]
        table = wandb.Table(columns=columns)

        for r in results:
            table.add_data(
                r["query"],
                r["expected_answer"],
                r["actual_answer"],
                r["faithfulness"],
                r["relevancy"],
                r["latency_ms"],
            )

        # Log la tabla — permite exploración interactiva en UI
        run.log({"eval_results": table})

        # Métricas agregadas
        avg_faith = sum(r["faithfulness"] for r in results) / len(results)
        avg_rel = sum(r["relevancy"] for r in results) / len(results)
        run.log({
            "avg_faithfulness": avg_faith,
            "avg_relevancy": avg_rel,
            "total_queries": len(results),
        })

        run.finish()

Comparación Sistemática de Experimentos

class ExperimentComparator:
    """Compara múltiples configuraciones de RAG."""

    def __init__(self, tracker: RAGExperimentTracker):
        self.tracker = tracker
        self.configs = []

    def add_variant(self, name: str, config: dict):
        self.configs.append({"name": name, **config})

    def run_comparison(
        self,
        eval_dataset: list[dict],
        rag_pipeline_factory,  # Callable que crea pipeline desde config
    ) -> list[dict]:
        results = []

        for variant in self.configs:
            print(f"\n▶ Running variant: {variant['name']}")

            # Crear pipeline con esta configuración
            pipeline = rag_pipeline_factory(variant)

            # Evaluar
            eval_results = []
            for item in eval_dataset:
                response = pipeline.query(item["question"])
                eval_results.append({
                    "query": item["question"],
                    "expected": item["expected_answer"],
                    "actual": response.answer,
                    "contexts": response.contexts,
                    "latency_ms": response.latency_ms,
                })

            # Calcular métricas
            metrics = self._calculate_metrics(eval_results)

            # Registrar en MLflow
            self.tracker.track_rag_experiment(
                config=variant,
                eval_results=metrics,
                sample_queries=eval_results[:10],
            )

            results.append({"variant": variant["name"], **metrics})

        return self._rank_results(results)

    def _rank_results(self, results: list[dict]) -> list[dict]:
        """Ordenar por score compuesto."""
        for r in results:
            r["composite_score"] = (
                r["faithfulness"] * 0.35 +
                r["answer_relevancy"] * 0.25 +
                r["context_precision"] * 0.20 +
                (1 - r["hallucination_rate"]) * 0.20
            )
        return sorted(results, key=lambda x: x["composite_score"], reverse=True)

Model Registry: Promover Configuraciones

# Registrar la mejor configuración como "modelo" en MLflow
import mlflow

def promote_rag_config(run_id: str, model_name: str = "rag-production"):
    """Promover una configuración como production-ready."""

    # Registrar como modelo
    result = mlflow.register_model(
        model_uri=f"runs:/{run_id}/artifacts",
        name=model_name,
    )

    # Transicionar a "Production"
    client = mlflow.tracking.MlflowClient()
    client.transition_model_version_stage(
        name=model_name,
        version=result.version,
        stage="Production",
    )

    print(f"✅ Config promoted: {model_name} v{result.version} → Production")

Prompt Registry

from dataclasses import dataclass

@dataclass
class PromptVersion:
    version: str
    template: str
    metrics: dict
    author: str
    notes: str

class PromptRegistry:
    """Registro versionado de prompts con métricas."""

    def __init__(self):
        self._versions: dict[str, list[PromptVersion]] = {}

    def register(
        self,
        name: str,
        template: str,
        metrics: dict,
        author: str,
        notes: str = "",
    ) -> PromptVersion:
        if name not in self._versions:
            self._versions[name] = []

        version = f"v{len(self._versions[name]) + 1}"
        pv = PromptVersion(
            version=version,
            template=template,
            metrics=metrics,
            author=author,
            notes=notes,
        )
        self._versions[name].append(pv)
        return pv

    def get_best(self, name: str, metric: str = "faithfulness") -> PromptVersion:
        """Obtener la versión con mejor métrica."""
        versions = self._versions.get(name, [])
        return max(versions, key=lambda v: v.metrics.get(metric, 0))

    def get_production(self, name: str) -> PromptVersion:
        """Obtener la versión más reciente."""
        return self._versions[name][-1]

    def compare(self, name: str) -> list[dict]:
        """Tabla comparativa de todas las versiones."""
        return [
            {
                "version": v.version,
                "author": v.author,
                **v.metrics,
                "notes": v.notes,
            }
            for v in self._versions.get(name, [])
        ]

Resumen

Flujo completo:

1. Experimentar → probar variaciones con tracking
2. Comparar → ranking automático de configuraciones
3. Promover → mejor config a Production via Model Registry
4. Monitorear → validar que Production mantiene calidad

🧠 Preguntas de Repaso

1. ¿Qué problema resuelve el experiment tracking en proyectos de IA?

• A) Mejora la velocidad de inferencia de los modelos
• B) Permite comparar sistemáticamente configuraciones (prompt, chunk_size, modelo, reranker) con métricas, evitando el "¿qué versión era mejor?"
• C) Reduce el costo de las APIs de LLM
• D) Automatiza el entrenamiento de modelos

Respuesta: B) — Sin tracking, es imposible saber qué combinación de configuraciones (prompt v3, chunk_size 512, reranker cohere) dio mejor resultado. Con tracking se compara objetivamente: "Run #47: faithfulness=0.92 vs Run #42: faithfulness=0.85".

2. ¿Cómo se calcula el score compuesto (composite_score) para comparar configuraciones de RAG?

• A) Promedio simple de todas las métricas
• B) faithfulness × 0.35 + answer_relevancy × 0.25 + context_precision × 0.20 + (1 − hallucination_rate) × 0.20
• C) Solo faithfulness × 100
• D) (latencia + costo) / calidad

Respuesta: B) — El composite_score pondera: faithfulness (35%) + answer_relevancy (25%) + context_precision (20%) + (1 − hallucination_rate) (20%). Los pesos reflejan prioridades: la fidelidad al contexto es lo más importante, seguida por relevancia.

3. ¿Cuál es la diferencia principal entre MLflow y Weights & Biases (W&B) para experiment tracking?

• A) MLflow es de pago y W&B es gratis
• B) MLflow ofrece tracking + model registry para pipeline RAG completo; W&B destaca en visualización interactiva y tablas detalladas para análisis de evaluaciones
• C) W&B solo funciona con PyTorch, MLflow con TensorFlow
• D) MLflow es cloud-only, W&B es solo local

Respuesta: B) — MLflow proporciona tracking de métricas/parámetros Y model registry (para promover configs a Production). W&B destaca en visualización con tablas interactivas que permiten análisis detallado query-por-query de las evaluaciones.

4. ¿Cuál es el flujo completo de experiment tracking antes de un deployment?

• A) Entrenar → Evaluar → Deployar
• B) Experimentar (probar variaciones con tracking) → Comparar (ranking automático) → Promover (mejor config a Production via Model Registry) → Monitorear
• C) Codificar → Testear → Deployar → Monitorear
• D) Diseñar → Implementar → Validar

Respuesta: B) — El flujo es: (1) Experimentar con variaciones registrando métricas, (2) Comparar con composite score para ranking automático, (3) Promover la mejor configuración a Production via Model Registry, (4) Monitorear que Production mantenga la calidad esperada.