Image for post MLflow desde Cero: Rastrea Experimentos y Modelos para Proyectos de ML Reproducibles en Python

MLflow desde Cero: Rastrea Experimentos y Modelos para Proyectos de ML Reproducibles en Python

Escrito por Sergio Márquez

Contexto del Problema

Como desarrolladores que inician o se consolidan en el mundo del Machine Learning (ML), rápidamente nos enfrentamos a un desafío común: la gestión del caos. Imagina que estás entrenando un modelo. Pruebas diferentes algoritmos, ajustas hiperparámetros, experimentas con distintas técnicas de preprocesamiento de datos. Cada intento genera un conjunto de resultados: métricas de rendimiento, el modelo entrenado, los parámetros utilizados, e incluso los datos de entrada. Sin una estrategia clara, este proceso se convierte en un laberinto de archivos con nombres confusos (modelo_final_v2_mejorado_este_si.pkl), hojas de cálculo manuales y la constante pregunta: "¿Qué versión de este modelo fue la que dio el mejor resultado? ¿Y con qué parámetros la entrené?".

Esta falta de trazabilidad y reproducibilidad no solo ralentiza el desarrollo, sino que también dificulta la colaboración en equipo, la depuración de errores y, en última instancia, la puesta en producción de modelos fiables. Necesitamos una forma sistemática de registrar cada "experimento" que realizamos, capturando todos los detalles relevantes de manera automática y organizada. Aquí es donde herramientas como MLflow se vuelven indispensables.

Conceptos Clave

MLflow es una plataforma de código abierto para gestionar el ciclo de vida completo del Machine Learning, incluyendo experimentación, reproducibilidad y despliegue. Aunque es una herramienta robusta con varios componentes, en este artículo nos centraremos en su módulo de Tracking, que es el corazón de la gestión de experimentos.

¿Qué es MLflow Tracking?

MLflow Tracking es una API y una interfaz de usuario (UI) para registrar y consultar información sobre tus experimentos de ML. Piensa en ello como un "cuaderno de laboratorio" digital y automatizado para tus modelos.

Componentes Principales de MLflow Tracking:

  • Experimentos (Experiments): Un experimento es una colección de "runs" (ejecuciones). Puedes agrupar runs relacionados bajo un mismo experimento. Por ejemplo, todos los intentos de entrenar un modelo de clasificación para un problema específico podrían pertenecer al mismo experimento.
  • Ejecuciones (Runs): Una ejecución corresponde a una única ejecución de tu código de ML. Cada run registra los parámetros de entrada, las métricas de salida, el código fuente y cualquier artefacto generado (como el modelo entrenado o gráficos).
  • Parámetros (Parameters): Son los valores de entrada clave para tu código de ML. Esto incluye hiperparámetros del modelo (ej. tasa de aprendizaje, número de épocas), rutas a los datos, etc. Se registran como pares clave-valor.
  • Métricas (Metrics): Son los valores numéricos que quieres evaluar para cada run. Esto puede ser la precisión, el F1-score, el error cuadrático medio (MSE), la pérdida (loss), etc. MLflow permite registrar métricas a lo largo del tiempo (útil para ver la evolución de la pérdida durante el entrenamiento).
  • Artefactos (Artifacts): Son los archivos de salida de tu run. Esto incluye el modelo entrenado serializado, gráficos, imágenes, archivos de texto, o cualquier otro archivo que sea relevante para el experimento. MLflow los almacena y los asocia con el run específico.

El objetivo principal es que, al finalizar un run, tengas toda la información necesaria para entender qué se hizo, cómo se hizo y cuáles fueron los resultados, permitiendo la reproducibilidad y la comparación sencilla entre diferentes intentos.

Implementación Paso a Paso

Vamos a empezar con un ejemplo práctico. Entrenaremos un modelo de Regresión Logística con el famoso dataset Iris y usaremos MLflow para registrar todo el proceso.

Paso 1: Instalación de MLflow

Primero, asegúrate de tener MLflow y las librerías necesarias instaladas. Si no las tienes, puedes instalarlas con pip:


pip install mlflow scikit-learn pandas matplotlib

Paso 2: Configuración Básica y Primer Run

Para este ejemplo, MLflow guardará los datos de tracking localmente en un directorio llamado mlruns/. No necesitamos una configuración compleja para empezar.

Crea un archivo llamado train_iris.py y añade el siguiente código:


import mlflow
import mlflow.sklearn
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, f1_score
from sklearn.datasets import load_iris
import warnings

warnings.filterwarnings("ignore") # Para ignorar warnings de convergencia de sklearn

# 1. Cargar el dataset Iris
iris = load_iris()
X = pd.DataFrame(iris.data, columns=iris.feature_names)
y = iris.target

# 2. Dividir los datos en conjuntos de entrenamiento y prueba
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 3. Definir parámetros del modelo
# Usaremos una variable de entorno para simular un parámetro configurable
# En un entorno real, esto podría venir de un archivo de configuración o CLI
alpha = float(os.environ.get("ALPHA", 0.1)) # Ejemplo de parámetro configurable
l1_ratio = float(os.environ.get("L1_RATIO", 0.5)) # Ejemplo de parámetro configurable

# 4. Iniciar un run de MLflow
# El "with" statement asegura que el run se cierre correctamente
with mlflow.start_run():
    # Loguear parámetros
    mlflow.log_param("solver", "saga")
    mlflow.log_param("max_iter", 1000)
    mlflow.log_param("alpha", alpha)
    mlflow.log_param("l1_ratio", l1_ratio)

    # 5. Entrenar el modelo
    model = LogisticRegression(
        solver="saga",
        max_iter=1000,
        penalty="elasticnet",
        l1_ratio=l1_ratio,
        C=alpha, # C es el inverso de la fuerza de regularización
        random_state=42
    )
    model.fit(X_train, y_train)

    # 6. Realizar predicciones y calcular métricas
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred, average='weighted')

    # 7. Loguear métricas
    mlflow.log_metric("accuracy", accuracy)
    mlflow.log_metric("f1_score", f1)

    # 8. Loguear el modelo
    # Esto guarda el modelo y sus dependencias para poder cargarlo después
    mlflow.sklearn.log_model(model, "iris_logistic_regression_model")

    print(f"Run MLflow completado. ID del Run: {mlflow.active_run().info.run_id}")
    print(f"Accuracy: {accuracy:.4f}")
    print(f"F1-Score: {f1:.4f}")

print("Script finalizado.")

```

Paso 3: Ejecutar el Script y Ver la UI de MLflow

Para ejecutar el script, simplemente abre tu terminal en el directorio donde guardaste train_iris.py y ejecuta:


python train_iris.py

Verás una salida similar a:


Run MLflow completado. ID del Run: [algún_id_alfanumérico]
Accuracy: 1.0000
F1-Score: 1.0000
Script finalizado.

Ahora, para ver los resultados en la interfaz de usuario de MLflow, ejecuta en la misma terminal:


mlflow ui

Abre tu navegador y ve a http://localhost:5000. Deberías ver la interfaz de MLflow con tu primer experimento y run registrado. Podrás ver los parámetros, métricas y el modelo como un artefacto descargable.

Mini Proyecto / Aplicación Sencilla: Comparando Hiperparámetros

Para demostrar el verdadero poder de MLflow Tracking, vamos a modificar nuestro script para ejecutar múltiples runs, probando diferentes valores para el parámetro de regularización C (que hemos mapeado a alpha en nuestro código) y el l1_ratio de la Regresión Logística. Esto simulará una búsqueda de hiperparámetros.

Crea un nuevo archivo llamado hyperparameter_tuning.py:


import mlflow
import mlflow.sklearn
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, f1_score
from sklearn.datasets import load_iris
import warnings
import os
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore")

# Cargar el dataset Iris
iris = load_iris()
X = pd.DataFrame(iris.data, columns=iris.feature_names)
y = iris.target

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Definir el nombre del experimento
mlflow.set_experiment("Iris_Logistic_Regression_Tuning")

# Valores a probar para los hiperparámetros
alphas = [0.01, 0.1, 1.0, 10.0]
l1_ratios = [0.0, 0.25, 0.5, 0.75, 1.0] # 0.0 es L2, 1.0 es L1

results = []

for alpha in alphas:
    for l1_ratio in l1_ratios:
        with mlflow.start_run():
            # Loguear parámetros
            mlflow.log_param("solver", "saga")
            mlflow.log_param("max_iter", 1000)
            mlflow.log_param("C_alpha", alpha)
            mlflow.log_param("l1_ratio", l1_ratio)

            # Entrenar el modelo
            # Cuidado: l1_ratio solo es aplicable con penalty='elasticnet'
            # Para l1_ratio=0.0 (L2) o l1_ratio=1.0 (L1), penalty debe ser 'l2' o 'l1' respectivamente
            # Aquí simplificamos usando elasticnet y ajustando l1_ratio
            if l1_ratio == 0.0: # L2 regularization
                penalty_type = "l2"
                current_l1_ratio = None # l1_ratio no se usa con penalty='l2'
            elif l1_ratio == 1.0: # L1 regularization
                penalty_type = "l1"
                current_l1_ratio = None # l1_ratio no se usa con penalty='l1'
            else: # Elastic Net regularization
                penalty_type = "elasticnet"
                current_l1_ratio = l1_ratio

            model = LogisticRegression(
                solver="saga",
                max_iter=1000,
                penalty=penalty_type,
                l1_ratio=current_l1_ratio,
                C=alpha,
                random_state=42
            )
            model.fit(X_train, y_train)

            # Realizar predicciones y calcular métricas
            y_pred = model.predict(X_test)
            accuracy = accuracy_score(y_test, y_pred)
            f1 = f1_score(y_test, y_pred, average='weighted')

            # Loguear métricas
            mlflow.log_metric("accuracy", accuracy)
            mlflow.log_metric("f1_score", f1)

            # Loguear el modelo
            mlflow.sklearn.log_model(model, "iris_logistic_regression_model")

            # Guardar un gráfico como artefacto
            fig, ax = plt.subplots()
            ax.bar(iris.feature_names, model.coef_)
            ax.set_title(f"Coeficientes del modelo (C={alpha}, L1_ratio={l1_ratio})")
            plt.tight_layout()
            plt.savefig("feature_coefficients.png")
            mlflow.log_artifact("feature_coefficients.png")
            plt.close(fig) # Cerrar la figura para liberar memoria
            os.remove("feature_coefficients.png") # Limpiar el archivo local

            print(f"Run completado: C={alpha}, L1_ratio={l1_ratio}, Accuracy={accuracy:.4f}, F1={f1:.4f}")
            results.append({"C": alpha, "l1_ratio": l1_ratio, "accuracy": accuracy, "f1_score": f1, "run_id": mlflow.active_run().info.run_id})

print("Todos los runs de ajuste de hiperparámetros han finalizado.")

# Opcional: Imprimir los mejores resultados
best_run = max(results, key=lambda x: x['f1_score'])
print(f"\nMejor F1-Score: {best_run['f1_score']:.4f} con C={best_run['C']} y L1_ratio={best_run['l1_ratio']} (Run ID: {best_run['run_id']})")

```

Ejecuta este script:


python hyperparameter_tuning.py

Mientras se ejecuta, puedes mantener mlflow ui abierto en tu navegador. Verás cómo se van añadiendo nuevos runs al experimento "Iris_Logistic_Regression_Tuning". Una vez que todos los runs hayan terminado, podrás:

  • Seleccionar múltiples runs y compararlos lado a lado.
  • Ordenar los runs por métricas (ej. F1-Score) para encontrar el mejor modelo.
  • Descargar los modelos entrenados o los gráficos de coeficientes de cada run.

Esto te permite visualizar rápidamente qué combinación de hiperparámetros produjo los mejores resultados, sin tener que rastrear manualmente cada intento.

Cargar un Modelo Logueado para Inferencia

Una de las grandes ventajas es que puedes cargar un modelo directamente desde MLflow para usarlo en inferencia, sin preocuparte por las dependencias o la serialización.

Crea un archivo predict_model.py:


import mlflow
import pandas as pd
from sklearn.datasets import load_iris

# ID del run del modelo que quieres cargar (reemplaza con uno de tus runs)
# Puedes obtener este ID de la UI de MLflow o de la salida de tu script
RUN_ID = "[REEMPLAZA_CON_EL_ID_DEL_MEJOR_RUN]"

# Ruta al artefacto del modelo dentro del run
# Por defecto, mlflow.sklearn.log_model guarda el modelo en "model"
MODEL_PATH = f"runs:/{RUN_ID}/iris_logistic_regression_model"

print(f"Cargando modelo desde: {MODEL_PATH}")

try:
    # Cargar el modelo usando la sintaxis de URI de MLflow
    loaded_model = mlflow.sklearn.load_model(MODEL_PATH)
    print("Modelo cargado exitosamente.")

    # Preparar nuevos datos para inferencia (usaremos una muestra del dataset Iris)
    iris = load_iris()
    new_data = pd.DataFrame(iris.data[:5], columns=iris.feature_names)

    print("\nNuevos datos para predicción:")
    print(new_data)

    # Realizar predicciones
    predictions = loaded_model.predict(new_data)
    probabilities = loaded_model.predict_proba(new_data)

    print("\nPredicciones:")
    print(predictions)
    print("\nProbabilidades:")
    print(probabilities)

except Exception as e:
    print(f"Error al cargar o usar el modelo: {e}")
    print("Asegúrate de que el RUN_ID sea correcto y que el modelo exista en esa ruta.")

```

Importante: Reemplaza [REEMPLAZA_CON_EL_ID_DEL_MEJOR_RUN] con el ID de un run real de tu experimento (puedes copiarlo de la UI de MLflow).

Ejecuta el script:


python predict_model.py

Esto demuestra cómo MLflow no solo te ayuda a rastrear, sino también a reutilizar tus modelos de manera sencilla y reproducible.

Errores Comunes y Depuración

  • "No active run" error: Si intentas loguear parámetros o métricas fuera de un bloque with mlflow.start_run(): o sin haber llamado a mlflow.start_run() explícitamente, MLflow no sabrá a qué run asociar la información. Siempre asegúrate de que haya un run activo.

    
        # Incorrecto
        # mlflow.log_param("param", 1)

        # Correcto
        with mlflow.start_run():
            mlflow.log_param("param", 1)

  • Problemas con el Tracking URI: Por defecto, MLflow usa ./mlruns. Si quieres usar una ubicación diferente (por ejemplo, una base de datos remota o un servidor de MLflow), debes configurarlo con mlflow.set_tracking_uri("tu_uri_aqui"). Si no puedes ver tus runs, verifica que el URI sea el correcto y que el servidor de tracking esté activo si es remoto.

    
        # Ejemplo para un servidor remoto
        # os.environ["MLFLOW_TRACKING_URI"] = "http://localhost:5000"
        # mlflow.set_tracking_uri("http://localhost:5000")

  • Dependencias del modelo: Cuando logueas un modelo con mlflow.sklearn.log_model() (o sus equivalentes para otras librerías), MLflow intenta inferir y guardar las dependencias del entorno. Sin embargo, a veces puede haber problemas si el entorno de carga no tiene las mismas versiones de librerías. Asegúrate de que tu entorno de inferencia sea lo más parecido posible al de entrenamiento, o considera usar MLflow Projects para empaquetar tu código y entorno.

  • Archivos de artefactos no encontrados: Si logueas un artefacto con mlflow.log_artifact("mi_archivo.png"), el archivo debe existir en la ruta especificada en el momento de la llamada. Si el archivo se elimina antes de que MLflow lo copie, o si la ruta es incorrecta, el artefacto no se guardará. Recuerda que MLflow copia el archivo, no lo mueve.

  • Confusión entre Experiments y Runs: Recuerda que un experimento es una colección de runs. Puedes establecer el experimento actual con mlflow.set_experiment("NombreDeMiExperimento"). Si no lo haces, MLflow usará un experimento por defecto llamado "Default". Organizar tus runs en experimentos lógicos te ayudará mucho a navegar la UI.

Aprendizaje Futuro / Próximos Pasos

MLflow Tracking es solo la punta del iceberg. Aquí hay algunas áreas para explorar a medida que te sientas más cómodo:

  • MLflow Projects: Aprende a empaquetar tu código de ML en un formato reproducible usando MLflow Projects. Esto te permite especificar dependencias, puntos de entrada y ejecutar tu código en diferentes entornos con un solo comando (mlflow run).

  • MLflow Models: Explora cómo MLflow estandariza el formato de los modelos para que puedan ser desplegados en diversas plataformas (Docker, Azure ML, SageMaker, etc.) con herramientas integradas.

  • MLflow Model Registry: Una característica crucial para la gestión del ciclo de vida del modelo. Te permite gestionar versiones de modelos, transicionar modelos entre etapas (Staging, Production), y anotar modelos con descripciones y tags.

  • Tracking Remoto: Configura MLflow para usar una base de datos (como PostgreSQL) y un almacenamiento de artefactos (como S3 o Azure Blob Storage) para centralizar el tracking de experimentos en un equipo o en la nube. Esto es fundamental para entornos de producción.

  • Integración con otras herramientas: MLflow se integra bien con librerías populares como Keras, PyTorch, XGBoost, y plataformas de orquestación como Apache Airflow o Kubeflow.

Dominar MLflow te proporcionará una base sólida para construir flujos de trabajo de ML más robustos, reproducibles y colaborativos, llevándote un paso más cerca de las buenas prácticas de MLOps.

Volver al blog