MLflow en Python: experimentos y versionado de modelos

Compara tres configuraciones de Random Forest con Iris, guarda sus parámetros y métricas de validación, y recupera el modelo elegido desde MLflow. El ejemplo registra un test final separado y crea una versión con el alias champion. Sus resultados son archivos locales y registros del catálogo; servir predicciones a usuarios requiere otra aplicación.

Configuración y almacenamiento local

El programa se comprobó con Python 3.13.5 en Windows y las siguientes dependencias directas. Crea una carpeta de proyecto e instálalas en un entorno virtual:

python -m venv .venv
.\.venv\Scripts\python.exe -m pip install `
    mlflow==3.14.0 scikit-learn==1.9.0 pandas==2.3.3 `
    numpy==2.4.6 skops==0.14.0

Los comandos usan el ejecutable de Python del entorno de Windows. En Linux o macOS, usa .venv/bin/python en los mismos comandos. La validación empleó un entorno existente con estas versiones, sin probar una instalación nueva ni un lock completo de dependencias transitivas.

El script elige explícitamente una base SQLite local para los metadatos del tracking y del registry. En esta configuración, mlflow.db conserva los registros y mlruns/ guarda los modelos y archivos CSV. Conserva ambos para volver a abrir el experimento. MLflow actual usa por defecto un backend de base de datos; el almacenamiento de metadatos sólo en archivos es una opción heredada. Documentación del backend.

Copia las cuatro secciones Python siguientes en train.py, en orden. Ejecuta el archivo desde la carpeta del proyecto:

.\.venv\Scripts\python.exe train.py

Registrar tres ejecuciones comparables

La partición estratificada reserva 30 filas para el test final y divide las restantes en 90 de entrenamiento y 30 de validación. Cada bosque aprende sólo con las filas de entrenamiento. Las métricas de validación permiten elegir un candidato; el test final queda fuera de esa decisión.

batch_id identifica las tres ejecuciones creadas en este intento. Así, la consulta posterior evita elegir una ejecución antigua de otro lote de experimentos. Cada ejecución guarda un CSV de importancia de variables y un modelo con ejemplo de entrada y firma inferida. Los requisitos explícitos del modelo registran las dependencias directas probadas y evitan recorrer el entorno para inferirlas.

from pathlib import Path
from uuid import uuid4

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

REQUIREMENTS = [
    "mlflow==3.14.0",
    "scikit-learn==1.9.0",
    "pandas==2.3.3",
    "numpy==2.4.6",
    "skops==0.14.0",
]
tracking_uri = "sqlite:///" + Path("mlflow.db").resolve().as_posix()
mlflow.set_tracking_uri(tracking_uri)
mlflow.set_registry_uri(tracking_uri)
mlflow.set_experiment("iris-random-forest")
batch_id = uuid4().hex

iris = load_iris()
X = pd.DataFrame(iris.data, columns=iris.feature_names)
y = iris.target
X_dev, X_test, y_dev, y_test = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=42
)
X_train, X_validation, y_train, y_validation = train_test_split(
    X_dev, y_dev, test_size=0.25, stratify=y_dev, random_state=42
)

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},
]
model_uris = {}
for config in configs:
    with mlflow.start_run(tags={"batch_id": batch_id}) as run:
        mlflow.log_params({
            **config,
            "random_state": 42,
            "training_rows": len(X_train),
            "validation_rows": len(X_validation),
            "test_rows": len(X_test),
        })
        model = RandomForestClassifier(**config, random_state=42, n_jobs=1)
        model.fit(X_train, y_train)
        prediction = model.predict(X_validation)
        mlflow.log_metrics({
            "validation_accuracy": accuracy_score(y_validation, prediction),
            "validation_f1": f1_score(y_validation, prediction, average="weighted"),
            "validation_precision": precision_score(
                y_validation, prediction, average="weighted"
            ),
            "validation_recall": recall_score(y_validation, prediction, average="weighted"),
        })
        importance = pd.DataFrame({
            "feature": iris.feature_names,
            "importance": model.feature_importances_,
        }).sort_values("importance", ascending=False)
        importance.to_csv("feature_importance.csv", index=False)
        mlflow.log_artifact("feature_importance.csv")
        info = mlflow.sklearn.log_model(
            model,
            name="random_forest_model",
            input_example=X_train.iloc[:3],
            pip_requirements=REQUIREMENTS,
            serialization_format="skops",
        )
        model_uris[run.info.run_id] = info.model_uri

El modelo se guarda en formato skops, compatible con la versión fijada de MLflow. El ejemplo carga artefactos que él mismo creó. Revisa la procedencia y los tipos que debas autorizar antes de cargar un modelo externo, y conserva dependencias compatibles entre entrenamiento y carga. La firma registra la estructura esperada de los datos; por sí sola no determina que una entrada futura tenga sentido o sea segura para un servicio productivo. Formato de modelos scikit-learn.

Elegir mediante métricas de validación

Consulta el lote actual y ordena por accuracy de validación. Los empates se resuelven mediante la hora de inicio más reciente, sin establecer que uno de esos modelos generalice mejor. La URI elegida procede del ModelInfo que devolvió el registro del modelo; no se reconstruye una ruta de artefactos.

from mlflow.tracking import MlflowClient

client = MlflowClient()
experiment = client.get_experiment_by_name("iris-random-forest")
if experiment is None:
    raise RuntimeError("Run the training section first.")
runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    filter_string=f"tags.batch_id = '{batch_id}'",
    order_by=["metrics.validation_accuracy DESC", "attributes.start_time DESC"],
    max_results=1,
)
if not runs:
    raise RuntimeError("No runs found for this training batch.")
best_run = runs[0]
model_uri = model_uris[best_run.info.run_id]
print("Selected parameters:", best_run.data.params)
print("Validation accuracy:", best_run.data.metrics["validation_accuracy"])

Los scores de validación forman parte de la selección. Presentar el mayor como estimación final insesgada reutilizaría las mismas observaciones para elegir y evaluar.

Cargar y evaluar una vez en el test final

Recupera el artefacto elegido, calcula su accuracy de test y añade esa métrica sólo a la ejecución seleccionada. Las dos flores del ejemplo muestran cómo enviar columnas con nombre al modelo recuperado; el programa imprime sus predicciones reales de especie.

loaded_model = mlflow.sklearn.load_model(model_uri)
test_prediction = loaded_model.predict(X_test)
test_accuracy = accuracy_score(y_test, test_prediction)
with mlflow.start_run(run_id=best_run.info.run_id):
    mlflow.log_metric("test_accuracy", test_accuracy)
print("Final test accuracy:", test_accuracy)

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)
print("Predicted species:", iris.target_names[predictions].tolist())

Este dataset pequeño y una partición ilustran el flujo. No establecen calidad productiva ni una clasificación fiable de rendimiento. Cambiar los modelos después de mirar el test convertiría esas observaciones en nuevos datos de selección. Conserva una evaluación externa o usa una evaluación exterior adecuada cuando iteres.

Registrar una versión y asignar un alias

El registro crea una versión en IrisClassifier; champion es una referencia modificable a esa versión. Reasignarlo cambia el modelo que recupera una consulta posterior del alias. Los stages están deprecados desde MLflow 2.9.0; aliases y tags son el flujo recomendado, sin atribuir la sustitución inicial de stages a la versión 3. Flujo del registry y migración.

registered = mlflow.register_model(model_uri=model_uri, name="IrisClassifier")
client.set_registered_model_alias(
    name="IrisClassifier", alias="champion", version=registered.version
)
champion_model = mlflow.sklearn.load_model("models:/IrisClassifier@champion")
print("Registered version:", registered.version)

El registry conserva la versión y su relación con la ejecución de origen. Una aplicación aún debe cargarla y exponer inferencia, con controles de acceso y verificaciones operativas propias. Asignar champion no inicia esa aplicación.

Inspeccionar el experimento local

Desde la misma carpeta del proyecto, inicia la interfaz con la base usada por el script:

.\.venv\Scripts\python.exe -m mlflow server `
    --backend-store-uri sqlite:///mlflow.db `
    --host 127.0.0.1 --port 5000

Abre http://127.0.0.1:5000, busca iris-random-forest y revisa sus tres ejecuciones: parámetros, métricas validation_*, CSV y modelos registrados. Sólo la elegida tiene test_accuracy. El catálogo de modelos muestra la versión y el alias. El servidor escucha en tu propia máquina; es una configuración de inspección local.

El ejemplo registra los datos descritos, sin constituir por sí solo un sistema completo de reproducibilidad o gobernanza. Un equipo también necesita procedencia de datos, versiones de código y entorno, políticas de evaluación y decisiones de despliegue revisadas. Puedes continuar con pipelines programados mediante Prefect e inferencia real con FastAPI y Docker.