Skip to content

English · Español

Lab 04 — Cablear mlflow en la disciplina de manifest existente

Objetivo: hacer que los runs de la Fase 18 sean navegables en mlflow ui sin reemplazar nunca manifest.json como fuente de verdad.

Tiempo estimado: 60–90 minutos.

Requisito previo: labs 01-03 hechos.

Lo que produces

  • src/minitrain/mlflow_wrap.py — un wrapper fino gestionado por context manager sobre mlflow.start_run().
  • Modificaciones a scripts/train_mini.py — llamadas al wrapper.
  • experiments/18-train-mini/ re-ejecutado con logging de mlflow.
  • mlruns.db (SQLite, gitignored).
  • experiments/18-train-mini/mlflow_screenshot.png — una captura del run en la UI de mlflow mostrando curvas + params + artifacts.
  • experiments/18-train-mini/README.md actualizado con el URI del run de mlflow.

Antecedentes que debes haber leído

  • theory/04-checkpoints-and-mlflow.md §"mlflow — qué te da y qué no".

TODOs

Bloque A — setup de mlflow

  • Instala mlflow según el opt group experiments de pyproject.toml: 
    uv pip install --group experiments
  • Fija el tracking URI: 
    export MLFLOW_TRACKING_URI=sqlite:///mlruns.db
  • Añade mlruns.db a .gitignore.
  • Arranca la UI en otra terminal: mlflow ui --backend-store-uri sqlite:///mlruns.db --port 5000. Abre http://localhost:5000 para confirmar que carga.

Bloque B — escribe el wrapper

# src/minitrain/mlflow_wrap.py
import mlflow
from contextlib import contextmanager
from pathlib import Path

@contextmanager
def tracking_run(experiment_name: str, run_name: str, manifest_path: Path, config: dict):
    """Open an mlflow run, log params from config, log manifest.json as artifact,
    yield a logger that the training loop can call to log metrics."""
    mlflow.set_experiment(experiment_name)
    with mlflow.start_run(run_name=run_name) as run:
        # Always log the manifest first — it's the source of truth.
        mlflow.log_artifact(str(manifest_path))
        mlflow.log_params(_flatten_config(config))

        def log_metrics(metrics: dict[str, float], step: int):
            for k, v in metrics.items():
                mlflow.log_metric(k, v, step=step)

        def log_artifact(path: Path):
            mlflow.log_artifact(str(path))

        # Update the manifest with the mlflow URI so the manifest can find this run.
        _patch_manifest(manifest_path, mlflow_run_id=run.info.run_id)

        yield SimpleNamespace(log_metrics=log_metrics, log_artifact=log_artifact, run=run)
  • Implementa _flatten_config — aplana la config anidada a claves dotted.keys para los params planos de mlflow.
  • Implementa _patch_manifest — reescribe manifest.json con el campo adicional mlflow_run_uri. Escritura atómica.
  • Testea que el wrapper sale limpio en KeyboardInterrupt — el bloque with cierra el run, el manifest queda consistente.

Bloque C — integra con train_mini.py

Modifica el script de entrenamiento para envolver el bucle:

with tracking_run(
    experiment_name="phase18_minigpt",
    run_name=f"phase18_{git_sha[:7]}",
    manifest_path=out_dir / "manifest.json",
    config=config,
) as logger:
    for step in range(total_steps):
        ...
        if step % log_every == 0:
            logger.log_metrics(
                {"train_loss": loss, "lr": lr, "g_norm": g_norm}, step=step
            )
        if step % val_every == 0:
            logger.log_metrics({"val_loss": val_loss, "val_ppl": val_ppl}, step=step)

    logger.log_artifact(checkpoint_path)
    logger.log_artifact(out_dir / "loss_curve.png")
    logger.log_artifact(out_dir / "results.json")
  • Ejecuta el entrenamiento de nuevo (puedes acortar a 500 steps si re-ejecutar el run completo es caro — anótalo en el README).
  • Tras el run, verifica en la UI:
  • El run aparece en el experimento phase18_minigpt.
  • Los parámetros muestran la config aplanada.
  • Las métricas tienen curvas indexadas por step (train_loss, lr, g_norm, val_loss, val_ppl).
  • Los artifacts incluyen manifest.json, loss_curve.png, results.json, weights.safetensors.

Bloque D — round-trip manifest ↔ mlflow

El manifest es la fuente de verdad. Comprueba que:

def test_manifest_can_locate_mlflow_run():
    manifest = json.loads(Path("experiments/18-train-mini/manifest.json").read_text())
    run_id = manifest["mlflow_run_uri"].split("/")[-1]
    run = mlflow.get_run(run_id)
    assert run.data.params["optimizer.lr_max"] == "0.0003"
    assert run.data.metrics["val_ppl"] < manifest["metrics"]["ngram_baseline_val_ppl"]
  • Este test confirma el ciclo: manifest → URI de mlflow → API de mlflow → métricas que coinciden con lo que el manifest afirma.

Bloque E — haz la captura

experiments/18-train-mini/mlflow_screenshot.png:

  • Abre el run en mlflow ui.
  • Captura la vista del run mostrando el panel de curvas + parámetros + lista de artifacts.
  • Commitea el PNG (un archivo de ~200 KB).

Bloque F — actualización del README

En experiments/18-train-mini/README.md, añade:

## mlflow

- Tracking URI: `sqlite:///mlruns.db` (local, gitignored).
- Experiment: `phase18_minigpt`.
- Run ID: `<run_id>` (from manifest.json).
- Screenshot: `mlflow_screenshot.png`.
- To re-open the UI: `mlflow ui --backend-store-uri sqlite:///mlruns.db --port 5000`.

Restricciones

  • Sin autologging. No actives mlflow.autolog(). Solo logging explícito. El wrapper son ~30 líneas; no importes la cocina entera.
  • El manifest es la verdad. Si la vista de mlflow no concuerda con el manifest, el bug está en el cableado de mlflow, no en el manifest. Testea esto.
  • Un experimento por fase. La Fase 18 usa phase18_minigpt. La Fase 19 creará phase19_dynamics. No los mezcles.

Condiciones de parada

Hecho cuando:

  1. mlflow ui muestra el run con las tres curvas de métricas, todos los parámetros de config y ≥ 4 artifacts.
  2. experiments/18-train-mini/mlflow_screenshot.png está commiteado.
  3. El test de round-trip manifest-a-mlflow pasa.
  4. Puedes re-encontrar un run vía manifest.json o mlflow ui — ambos llevan a los mismos artifacts.

Escollos

  • mlflow.start_run() sin context manager. Deja runs abiertos colgados en Ctrl+C. Siempre with.
  • Hacer log de métricas sin step=. Mlflow entonces auto-incrementa, y una métrica logueada desde varios sitios se reordena. Siempre pasa step=.
  • Hacer log del archivo safetensors como un parámetro en vez de como artifact. Los parámetros son strings clave-valor, máximo 500 caracteres. Usa log_artifact.
  • Fijar MLFLOW_TRACKING_URI a una ruta relativa sin el prefijo sqlite:///. Mlflow entonces escribe un file store, no SQLite, y te comes el dolor del file store lento.

Cuándo consultar solutions/

Después de que la captura esté commiteada. Solución en solutions/04-mlflow-wiring-ref.md (escrita al abrir la fase).

El trabajo de lab de la Fase 18 está completo. Continúa con /quiz 18, después PHASE_18_REPORT.md, después learners/borja/phase-18/reflections.md, después proceed.