MLflow en Databricks: trazabilidad de experimentos de Machine Learning

MLflow es la plataforma open-source de MLOps integrada en Databricks que resuelve uno de los mayores problemas en Machine Learning: rastrear lo que se probó, comparar resultados y reproducir el mejor modelo en producción.

Estructura de MLflow

• Experiment — agrupa runs relacionados (p. ej. "Modelo_Churn_v2").
• Run — una ejecución de entrenamiento con sus parámetros y métricas.
• Artifact — archivos generados (modelo serializado, gráficos, datasets).
• Model Registry — versionado y promoción de modelos (Staging → Production).

Registrar un experimento

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier

mlflow.set_experiment("/Experiments/Churn_Prediction")

with mlflow.start_run(run_name="RF_100trees"):
    params = {"n_estimators": 100, "max_depth": 8}
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)

    accuracy = model.score(X_test, y_test)

    mlflow.log_params(params)
    mlflow.log_metric("accuracy", accuracy)
    mlflow.sklearn.log_model(model, "random_forest_model")

Autologging

# MLflow registra automáticamente los parámetros y métricas de sklearn
mlflow.sklearn.autolog()
model.fit(X_train, y_train)  # todo registrado automáticamente

Model Registry y despliegue

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

# Promover a Production vía UI o API
client = mlflow.tracking.MlflowClient()
client.transition_model_version_stage("ChurnPrediction", version=3, stage="Production")

# Cargar el modelo de producción en cualquier notebook
model = mlflow.sklearn.load_model("models:/ChurnPrediction/Production")

Conclusión

MLflow es indispensable para equipos de ciencia de datos que quieren reproducibilidad y gobernanza. En Databricks, está integrado por defecto — cada notebook tiene un experiment asociado y los mejores modelos pueden promoverse a producción con un clic.