Capítulo 1: Introducción a MLflow

¿Qué es MLflow?

MLflow es una plataforma de código abierto para gestionar el ciclo de vida completo del aprendizaje automático. Proporciona cuatro componentes principales:

Configuración de MLflow

Instalación

# Instalación básica
pip install mlflow

# Con extras
pip install mlflow[extras]

# Iniciar servidor local de tracking
mlflow server --host 0.0.0.0 --port 5000

Configuración para producción

# docker-compose.yml
version: '3.8'
services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.10.0
    ports:
      - "5000:5000"
    environment:
      - MLFLOW_BACKEND_STORE_URI=postgresql://user:pass@postgres:5432/mlflow
      - MLFLOW_DEFAULT_ARTIFACT_ROOT=s3://mlflow-artifacts/
      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
    command: >
      mlflow server
      --backend-store-uri postgresql://user:pass@postgres:5432/mlflow
      --default-artifact-root s3://mlflow-artifacts/
      --host 0.0.0.0
      --port 5000
    depends_on:
      - postgres

  postgres:
    image: postgres:15
    environment:
      - POSTGRES_USER=user
      - POSTGRES_PASSWORD=pass
      - POSTGRES_DB=mlflow
    volumes:
      - postgres_data:/var/lib/postgresql/data

volumes:
  postgres_data:

Seguimiento con MLflow

Seguimiento básico de experimentos

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score

# Establecer URI de tracking
mlflow.set_tracking_uri("http://localhost:5000")

# Establecer experimento
mlflow.set_experiment("customer-churn-prediction")

# Iniciar una ejecución
with mlflow.start_run(run_name="random-forest-baseline"):
    # Registrar parámetros
    n_estimators = 100
    max_depth = 10
    mlflow.log_param("n_estimators", n_estimators)
    mlflow.log_param("max_depth", max_depth)
    mlflow.log_param("model_type", "RandomForest")
    
    # Entrenar modelo
    model = RandomForestClassifier(
        n_estimators=n_estimators,
        max_depth=max_depth,
        random_state=42
    )
    model.fit(X_train, y_train)
    
    # Realizar predicciones
    predictions = model.predict(X_test)
    probabilities = model.predict_proba(X_test)[:, 1]
    
    # Registrar métricas
    mlflow.log_metric("accuracy", accuracy_score(y_test, predictions))
    mlflow.log_metric("f1_score", f1_score(y_test, predictions))
    mlflow.log_metric("auc_roc", roc_auc_score(y_test, probabilities))
    
    # Registrar modelo
    mlflow.sklearn.log_model(model, "model")
    
    # Registrar artefactos (archivos)
    mlflow.log_artifact("feature_importance.csv")

Autologging

# Habilitar autolog para sklearn
mlflow.sklearn.autolog()

# Ahora el entrenamiento registra todo automáticamente
model = RandomForestClassifier(n_estimators=100)
model.fit(X_train, y_train)
# ¡Parámetros, métricas y modelo se registran automáticamente!

# Autolog para otros frameworks
mlflow.pytorch.autolog()
mlflow.tensorflow.autolog()
mlflow.xgboost.autolog()
mlflow.lightgbm.autolog()

Ejecuciones anidadas para ajuste de hiperparámetros

from sklearn.model_selection import ParameterGrid

param_grid = {
    'n_estimators': [100, 200, 300],
    'max_depth': [5, 10, 15, None],
    'min_samples_split': [2, 5, 10]
}

with mlflow.start_run(run_name="hyperparameter-search"):
    best_auc = 0
    best_params = None
    
    for params in ParameterGrid(param_grid):
        with mlflow.start_run(nested=True):
            # Registrar parámetros
            mlflow.log_params(params)
            
            # Entrenar modelo
            model = RandomForestClassifier(**params, random_state=42)
            model.fit(X_train, y_train)
            
            # Evaluar
            auc = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
            mlflow.log_metric("auc_roc", auc)
            
            # Rastrear el mejor
            if auc > best_auc:
                best_auc = auc
                best_params = params
    
    # Registrar mejores resultados en la ejecución principal
    mlflow.log_params({f"best_{k}": v for k, v in best_params.items()})
    mlflow.log_metric("best_auc_roc", best_auc)

MLflow Projects

Estructura del proyecto

ml-project/
├── MLproject          # Definición del proyecto
├── conda.yaml         # Especificación del entorno
├── train.py           # Script de entrenamiento
└── requirements.txt   # Alternativa a conda

Archivo MLproject

# MLproject
name: customer-churn-model

conda_env: conda.yaml

entry_points:
  main:
    parameters:
      n_estimators: {type: int, default: 100}
      max_depth: {type: int, default: 10}
      data_path: {type: string, default: "data/processed/"}
    command: "python train.py --n-estimators {n_estimators} --max-depth {max_depth} --data-path {data_path}"
  
  evaluate:
    parameters:
      model_path: string
      test_data: string
    command: "python evaluate.py --model-path {model_path} --test-data {test_data}"

Ejecución de proyectos

# Ejecutar localmente
mlflow run . -P n_estimators=200 -P max_depth=15

# Ejecutar desde Git
mlflow run https://github.com/company/ml-project -P n_estimators=200

# Ejecutar con entorno específico
mlflow run . --env-manager conda

Modelos MLflow

Firma del modelo y ejemplo de entrada

from mlflow.models.signature import infer_signature
import pandas as pd

# Inferir firma a partir de los datos
signature = infer_signature(X_train, model.predict(X_train))

# Crear ejemplo de entrada
input_example = X_train.head(5)

# Registrar modelo con firma
mlflow.sklearn.log_model(
    model,
    "model",
    signature=signature,
    input_example=input_example,
    registered_model_name="customer-churn-model"
)

Sabores de modelo

# Diferentes formas de registrar modelos

# Sklearn
mlflow.sklearn.log_model(sklearn_model, "model")

# PyTorch
mlflow.pytorch.log_model(pytorch_model, "model")

# TensorFlow/Keras
mlflow.tensorflow.log_model(tf_model, "model")

# Modelos personalizados con pyfunc
class CustomModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # Cargar artefactos
        self.model = load_model(context.artifacts["model_path"])
    
    def predict(self, context, model_input):
        return self.model.predict(model_input)

mlflow.pyfunc.log_model(
    "model",
    python_model=CustomModel(),
    artifacts={"model_path": "path/to/model"}
)

Puntos clave

Summary

1. MLflow Tracking registra parámetros, métricas y artefactos
2. Autologging simplifica el seguimiento para frameworks compatibles
3. MLflow Projects asegura código de ML reproducible
4. MLflow Models proporciona empaquetado estándar para despliegue
5. Firmas de modelo definen entradas y salidas esperadas

---

Siguiente: Capítulo 2: Weights and Biases