Cómo versionar modelos ML con MLflow en Databricks: paso a paso
Este tutorial muestra cómo registrar, versionar y cargar un modelo ML usando MLflow en Databricks, facilitando promover modelos del desarrollo a producción. La capacidad de versionar modelos en Databricks con MLflow ayuda a rastrear experimentos, comparar métricas y gestionar despliegues de forma controlada — fundamental cuando existen múltiples iteraciones del mismo modelo (por ejemplo, 5–20 runs por experimento) y cuando se pretende reproducir resultados en producción.
Requisitos previos
- Cuenta Databricks con un Workspace configurado y permisos para crear modelos/experimentos. Idealmente, un usuario con permisos de Developer o Admin.
- Cluster Databricks en ejecución con runtime que incluya MLflow (ej.: Databricks Runtime ML). Un cluster con 4 cores y 16 GB de RAM es suficiente para este ejemplo.
- Conocimientos básicos de Python y scikit-learn o similar. Conocimientos básicos de línea de comandos y de la UI de Databricks ayudan en la navegación.
Paso 1: Crear un experimento MLflow
Un experimento agrupa ejecuciones (runs). Crear un experimento permite organizar y comparar modelos por objetivo (ej.: "classificacao_iris_v1"). En entornos reales, cree un experimento por proyecto o por pipeline (por ejemplo, "fraude_cartoes_production").
import mlflow
mlflow.set_experiment("exemplo_classificacao_iris")
Por qué: con un experimento puede filtrar fácilmente runs por fecha, autor o hiperparámetros. Ej.: si genera 10 runs con n_estimators = [10, 50, 100], puede comparar rápidamente las métricas agregadas (mean/std) entre grupos.
Paso 2: Entrenar y registrar un modelo con una run
Dentro de una run, anotamos métricas, parámetros y generamos el artefacto del modelo. Aquí usamos scikit-learn para un ejemplo sencillo con el dataset Iris — útil para demostrar el flujo completo antes de aplicar en datos reales.
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
import mlflow.sklearn
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
with mlflow.start_run(run_name="rf_iris_v1") as run:
params = {"n_estimators": 50, "random_state": 42}
mlflow.log_params(params)
model = RandomForestClassifier(**params)
model.fit(X_train, y_train)
preds = model.predict(X_test)
acc = accuracy_score(y_test, preds)
mlflow.log_metric("accuracy", acc)
# Registar o modelo no MLflow Model Registry (name define o modelo versionado)
mlflow.sklearn.log_model(model, artifact_path="model", registered_model_name="IrisRF")
print(f"Run ID: {run.info.run_id}, accuracy: {acc}")
Detalles prácticos: el artefacto queda almacenado en el almacenamiento del workspace (DBFS o OneLake según la configuración). El campo registered_model_name crea una entrada única en el Model Registry llamada "IrisRF" y genera automáticamente una nueva versión (ej.: v1, v2) por cada log_model con el mismo nombre.
Paso 3: Ver las versiones en el Model Registry
Después de registrar, el modelo aparece en el Model Registry de Databricks. La UI muestra versiones, stages (Staging, Production, Archived) y metadatos como la fecha de creación, autor y notas. Usar la API permite automatizar comprobaciones y generar informes (por ejemplo, listar todas las versiones creadas en los últimos 30 días).
from mlflow.tracking import MlflowClient
client = MlflowClient()
# Listar versiones del modelo registrado
for v in client.get_latest_versions("IrisRF"):
print(v.version, v.current_stage, v.status)
# Obtener información de una versión específica
ver = client.get_model_version("IrisRF", "1")
print(ver.version, ver.current_stage, ver.source)
Ejemplo de output esperado: "1 Staging READY" o "1 None READY". Si ve estados como PENDING_REGISTRATION o FAILED, revise los logs de la run para errores en la subida del artefacto.
Paso 4: Promover una versión a Staging/Production
Promover versiones permite workflows de validación y despliegue. La práctica común es promover primero a Staging para ejecutar pruebas A/B o validación por lotes, y solo después mover a Production cuando esté validado. Use políticas de aprobación manual o pipelines CI/CD para controlar quién puede promover.
# Promover la versión 1 a Staging
client.transition_model_version_stage(name="IrisRF", version="1", stage="Staging")
# Tras las pruebas, promover a Production
# client.transition_model_version_stage(name="IrisRF", version="1", stage="Production")
Buenas prácticas: incluya una tag o descripción con el commit del código y el hash del conjunto de datos usado (ej.: git_sha: abc123, data_snapshot: 2026-10-01) para trazabilidad total.
Paso 5: Cargar una versión específica para predicción
Para hacer inferencia, cargue la versión deseada del Model Registry. Puede usar el URI 'models:/<name>/<stage>' para apuntar siempre a la versión en Production o Staging, o 'models:/<name>/<version>' para enlazar a una versión inmutable (ej.: v1).
import mlflow
# Cargar la versión en Staging
model = mlflow.pyfunc.load_model("models:/IrisRF/Staging")
# Hacer una predicción de ejemplo
import numpy as np
sample = X_test[0:2]
pred = model.predict(sample)
print(pred)
Tenga en cuenta que cargar por stage es útil en despliegue automático (servicio de inferencia apunta a "models:/IrisRF/Production"). Si necesita reproducibilidad total, cargue por versión.
Verificar el resultado
Confirme en la UI de Databricks: en la sección Models debería aparecer 'IrisRF' con versiones y stages. En la consola del notebook verá el Run ID, métricas (accuracy) y la predicción de ejemplo. Verifique también el tamaño del artefacto (por ejemplo 2–10 MB para modelos simples) y los metadatos. Si la predicción devuelve clases válidas y la versión aparece en el Registry, el proceso funcionó.
Conclusión
Ahora sabe registrar, versionar, promover y cargar modelos con MLflow en Databricks, lo que facilita trazabilidad y despliegue. Próximos pasos prácticos: automatizar evaluaciones mediante CI/CD, añadir tests unitarios de predicción e integrar pruebas A/B en las versiones. Consejo final: mantenga convenciones de nombres claras (ej.: proyecto_modelo_version) y registre siempre parámetros, métricas y hashes de datos para evitar regresiones a lo largo del tiempo.