Gestiona el Ciclo de Vida de Machine Learning con MLflow en Python

Si alguna vez has entrenado más de tres modelos de machine learning, ya conoces el caos que se genera: carpetas llamadas modelo_final, modelo_final_v2, modelo_final_ESTE_SI, archivos de parámetros dispersos en comentarios de notebooks, y la angustiante pregunta de qué hiperparámetros usaste exactamente en ese modelo que funcionó tan bien hace dos semanas. MLflow existe precisamente para resolver ese problema. En este artículo rastrearás experimentos con parámetros, métricas y artefactos, consultarás la mejor ejecución programáticamente y promoverás un modelo a producción con el Model Registry — todo en unas pocas líneas de Python.

El caos sin rastreo de experimentos

Cuando un proyecto de ML crece, la falta de un sistema de rastreo produce problemas concretos:

• Reproducibilidad rota: no puedes recrear un resultado porque no recuerdas qué versión de los datos usaste o qué valor tenía max_depth.
• Comparación imposible: comparar 15 ejecuciones manualmente revisando notas dispersas es propenso a errores.
• Colaboración difícil: un compañero no puede retomar tu trabajo si los artefactos y parámetros viven solo en tu máquina.
• Pérdida de conocimiento: el modelo que descartaste por "malo" tal vez era el mejor para un caso de uso diferente, pero ya no tienes cómo recuperarlo.

Un sistema de rastreo no es un lujo; es higiene básica de ingeniería para proyectos de ML serios.

MLflow: una visión general

MLflow es una plataforma open-source creada por Databricks para gestionar el ciclo de vida completo de ML. Se divide en cuatro componentes principales:

• MLflow Tracking: registra y consulta experimentos: parámetros, métricas, artefactos y modelos.
• MLflow Projects: empaqueta código de ML en un formato reproducible que puede ejecutarse en cualquier plataforma.
• MLflow Models: define un formato estándar para empaquetar modelos con soporte para múltiples "flavors" (sklearn, tensorflow, pytorch, etc.).
• MLflow Model Registry: almacén centralizado de modelos con versionado, staging y comentarios.

Para la mayoría de los proyectos individuales o de equipo pequeño, empezarás con Tracking y, cuando tengas modelos en producción, añadirás el Registry.

Instalación y puesta en marcha

La instalación es directa:

pip install mlflow scikit-learn pandas numpy

Para lanzar la interfaz web local:

mlflow ui

Esto arranca un servidor en http://127.0.0.1:5000 donde podrás visualizar todos tus experimentos. Por defecto, MLflow guarda los datos en un directorio mlruns/ en tu directorio de trabajo actual.

Rastreando experimentos con Python

Usaremos el dataset Iris y un Random Forest para ilustrar todas las funciones de tracking. El objetivo es entrenar varias configuraciones y registrar todo automáticamente.

import mlflow
import mlflow.sklearn
import numpy as np
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score

# Cargar datos
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 nombre del experimento
mlflow.set_experiment("iris-random-forest")

# Configuraciones a probar
configs = [
    {"n_estimators": 50,  "max_depth": 3, "min_samples_split": 2},
    {"n_estimators": 100, "max_depth": 5, "min_samples_split": 5},
    {"n_estimators": 200, "max_depth": 10, "min_samples_split": 2},
]

for config in configs:
    with mlflow.start_run():
        # Registrar parámetros
        mlflow.log_param("n_estimators", config["n_estimators"])
        mlflow.log_param("max_depth", config["max_depth"])
        mlflow.log_param("min_samples_split", config["min_samples_split"])
        mlflow.log_param("test_size", 0.2)
        mlflow.log_param("random_state", 42)

        # Entrenar modelo
        model = RandomForestClassifier(
            n_estimators=config["n_estimators"],
            max_depth=config["max_depth"],
            min_samples_split=config["min_samples_split"],
            random_state=42
        )
        model.fit(X_train, y_train)

        # Evaluar
        y_pred = model.predict(X_test)
        acc = accuracy_score(y_test, y_pred)
        f1  = f1_score(y_test, y_pred, average="weighted")
        prec = precision_score(y_test, y_pred, average="weighted")
        rec  = recall_score(y_test, y_pred, average="weighted")

        # Registrar métricas
        mlflow.log_metric("accuracy", acc)
        mlflow.log_metric("f1_score", f1)
        mlflow.log_metric("precision", prec)
        mlflow.log_metric("recall", rec)

        # Registrar importancia de features como artefacto
        feature_importance = pd.DataFrame({
            "feature": iris.feature_names,
            "importance": model.feature_importances_
        }).sort_values("importance", ascending=False)
        feature_importance.to_csv("feature_importance.csv", index=False)
        mlflow.log_artifact("feature_importance.csv")

        # Guardar el modelo con ejemplo de entrada para documentación automática
        mlflow.sklearn.log_model(
            model,
            "random_forest_model",
            input_example=X_train.iloc[:3]
        )

        print(f"Run completado — accuracy: {acc:.4f}, f1: {f1:.4f}")

Veamos qué hace cada llamada clave:

• mlflow.set_experiment("iris-random-forest"): agrupa todas las ejecuciones bajo un mismo experimento. Si no existe, lo crea.
• mlflow.start_run(): abre un contexto de ejecución. Al salir del bloque with, la ejecución se cierra automáticamente.
• mlflow.log_param(key, value): registra un hiperparámetro o configuración (generalmente valores que no cambian durante el entrenamiento).
• mlflow.log_metric(key, value): registra una métrica de evaluación (puede llamarse múltiples veces con step para trazar curvas de entrenamiento).
• mlflow.log_artifact(path): sube cualquier archivo al almacén de artefactos.
• mlflow.sklearn.log_model(model, name): serializa y guarda el modelo con metadatos suficientes para cargarlo con cualquier "flavor" compatible.

La interfaz de MLflow

Después de ejecutar el script, abre http://127.0.0.1:5000 en tu navegador. Verás:

• Lista de experimentos en el panel izquierdo.
• Tabla de ejecuciones con todas las corridas, mostrando parámetros y métricas en columnas configurables.
• Filtros y ordenación para encontrar rápidamente la mejor ejecución por cualquier métrica.
• Vista de ejecución individual con todos los parámetros, métricas, artefactos descargables y la URI del modelo.

La UI también renderiza gráficos de métricas en el tiempo si las registras con step, lo que es muy útil para el entrenamiento iterativo (épocas en redes neuronales, por ejemplo).

Comparando ejecuciones y encontrando el mejor modelo

Puedes comparar ejecuciones desde la UI seleccionando varias y haciendo clic en "Compare". También puedes hacerlo programáticamente:

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Obtener el experimento por nombre
experiment = client.get_experiment_by_name("iris-random-forest")
experiment_id = experiment.experiment_id

# Buscar la mejor ejecución por accuracy
best_run = client.search_runs(
    experiment_ids=[experiment_id],
    order_by=["metrics.accuracy DESC"],
    max_results=1
)[0]

print(f"Mejor run ID: {best_run.info.run_id}")
print(f"Accuracy: {best_run.data.metrics['accuracy']:.4f}")
print(f"Parámetros: {best_run.data.params}")

Cargando un modelo guardado y haciendo predicciones

Una vez identificada la mejor ejecución, cargar el modelo es sencillo:

import mlflow.sklearn
import pandas as pd

# Usar el run_id de la mejor ejecución
run_id = best_run.info.run_id
model_uri = f"runs:/{run_id}/random_forest_model"

# Cargar el modelo
loaded_model = mlflow.sklearn.load_model(model_uri)

# Nuevos datos para predecir
new_data = pd.DataFrame({
    "sepal length (cm)": [5.1, 6.7],
    "sepal width (cm)":  [3.5, 3.0],
    "petal length (cm)": [1.4, 5.2],
    "petal width (cm)":  [0.2, 2.3]
})

predictions = loaded_model.predict(new_data)
species = ["setosa", "versicolor", "virginica"]
print([species[p] for p in predictions])
# ['setosa', 'virginica']

Registrando el modelo en el Model Registry

Para equipos que manejan modelos en producción, el Registry añade una capa de gobernanza. MLflow 3 reemplazó el sistema de stages (Production, Staging, Archived) con aliases de modelo, que son más flexibles: puedes tener múltiples aliases por versión y reasignarlos sin restricciones de transición de estado.

# Registrar el mejor modelo
result = mlflow.register_model(
    model_uri=f"runs:/{run_id}/random_forest_model",
    name="IrisClassifier"
)

# Asignar alias "champion" a la versión en producción
client.set_registered_model_alias(
    name="IrisClassifier",
    alias="champion",
    version=result.version
)

# Cargar siempre la versión con alias champion
prod_model = mlflow.sklearn.load_model("models:/IrisClassifier@champion")

Con esto, tu pipeline de inferencia siempre apunta al alias @champion sin hardcodear run IDs ni números de versión. Cuando reentrenas y tienes un modelo mejor, simplemente mueves el alias a la nueva versión.

Conclusión

MLflow transforma la forma en que trabajas con machine learning: de una colección de scripts y archivos desorganizados a un flujo reproducible y auditable. Con apenas cuatro llamadas (log_param, log_metric, log_artifact, log_model) obtienes rastreo completo de tus experimentos.

El siguiente paso natural es integrar MLflow en un pipeline de CI/CD: ejecutar entrenamiento automático ante nuevos datos, registrar la ejecución, y promover automáticamente el modelo al Registry si supera un umbral de calidad. Herramientas como GitHub Actions o Azure Pipelines se integran bien con la API de MLflow para este propósito.

El código de este artículo está pensado para ejecutarse directamente; solo necesitas instalar las dependencias y tendrás un sistema de experimentación funcional en minutos.

Una vez rastreados los experimentos, el siguiente paso es envolver tu código de entrenamiento en un workflow confiable y programado: Orquesta Pipelines de ML con Prefect. Y cuando tu modelo @champion esté listo para recibir tráfico, revisa Sirve un Modelo de Machine Learning con FastAPI y Docker.