Como usar o MLflow para rastreamento de experimentos e versionamento de modelos

1. Introdução ao MLflow e sua importância no ciclo de vida de ML

No desenvolvimento de modelos de machine learning, um dos maiores desafios é o rastreamento manual de experimentos. Cientistas de dados frequentemente se veem perdendo anotações sobre quais hiperparâmetros geraram determinados resultados, ou qual versão de dataset produziu um modelo específico. Essa falta de reprodutibilidade compromete a confiabilidade do trabalho e dificulta a colaboração entre equipes.

O MLflow surge como uma solução open-source que organiza todo o ciclo de vida de ML. Seus quatro componentes principais são:

  • MLflow Tracking: Registra parâmetros, métricas, artefatos e código de cada execução
  • MLflow Projects: Empacota código de forma reprodutível
  • MLflow Models: Padroniza o formato de modelos para deploy
  • MLflow Registry: Gerencia versões e estágios dos modelos

Os cenários de uso vão desde prototipação em notebooks até produção em larga escala com milhares de experimentos.

2. Configuração do ambiente e primeiros passos com MLflow Tracking

A instalação do MLflow é simples:

pip install mlflow

Para iniciar o servidor de tracking local:

mlflow ui

Isso abre a interface em http://localhost:5000. Para configuração remota, defina a variável de ambiente:

export MLFLOW_TRACKING_URI=http://meu-servidor:5000

Exemplo prático: rastreando um experimento de regressão linear:

import mlflow
import mlflow.sklearn
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import numpy as np

# Gerar dados sintéticos
np.random.seed(42)
X = np.random.rand(100, 1)
y = 2 * X + 1 + 0.1 * np.random.randn(100, 1)

with mlflow.start_run():
    # Registrar parâmetros
    mlflow.log_param("alpha", 0.01)
    mlflow.log_param("fit_intercept", True)

    # Treinar modelo
    model = LinearRegression()
    model.fit(X, y)
    predictions = model.predict(X)

    # Registrar métricas
    mse = mean_squared_error(y, predictions)
    mlflow.log_metric("mse", mse)

    # Salvar modelo
    mlflow.sklearn.log_model(model, "modelo_regressao")

    print(f"Run ID: {mlflow.active_run().info.run_id}")

Cada execução gera um run_id único que permite recuperar todos os dados posteriormente.

3. Registro e organização de experimentos com tags e metadados

Tags são essenciais para categorizar experimentos. Exemplo:

with mlflow.start_run():
    mlflow.set_tag("projeto", "precificacao_imoveis")
    mlflow.set_tag("equipe", "data_science_a")
    mlflow.set_tag("objetivo", "testar_regularizacao")
    mlflow.set_tag("dataset_version", "v3.2")

Para criar uma hierarquia, use nomes de experimentos:

mlflow.set_experiment("precificacao_imoveis/experimentos_iniciais")

Boas práticas de nomenclatura:
- Use experimento/tipo_algoritmo/data para organização
- Documente hiperparâmetros como tags, não apenas como parâmetros
- Inclua hash do commit do código como tag para rastreabilidade

4. Versionamento de modelos com MLflow Models

O MLflow Model possui três componentes: sabor (flavor), assinatura (signature) e artefatos. Exemplo com XGBoost:

import mlflow.xgboost
import xgboost as xgb
from mlflow.models.signature import infer_signature

X_train = [[1, 2], [3, 4], [5, 6]]
y_train = [0, 1, 0]

model = xgb.XGBClassifier(n_estimators=100)
model.fit(X_train, y_train)

# Inferir assinatura (schema de entrada/saída)
signature = infer_signature(X_train, model.predict(X_train))

with mlflow.start_run():
    mlflow.xgboost.log_model(
        model, 
        "xgboost_model",
        signature=signature,
        registered_model_name="classificador_xgb"
    )

A assinatura permite validar dados de entrada durante o deploy, prevenindo erros silenciosos.

5. Gerenciamento do ciclo de vida do modelo com MLflow Registry

O Registry permite transições entre estágios:

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Transicionar modelo para Staging
client.transition_model_version_stage(
    name="classificador_xgb",
    version=1,
    stage="Staging"
)

# Para Production
client.transition_model_version_stage(
    name="classificador_xgb",
    version=1,
    stage="Production"
)

# Para Archived (rollback)
client.transition_model_version_stage(
    name="classificador_xgb",
    version=1,
    stage="Archived"
)

O versionamento semântico é automático: cada novo registro incrementa a versão. Para rollback seguro, mantenha a versão anterior em Production enquanto testa a nova em Staging.

6. Automação de pipelines com MLflow Projects

Crie um arquivo MLproject para definir o pipeline:

name: pipeline_treinamento

conda_env: conda.yaml

entry_points:
  main:
    parameters:
      learning_rate: {type: float, default: 0.01}
      n_estimators: {type: int, default: 100}
    command: "python treinamento.py --learning_rate {learning_rate} --n_estimators {n_estimators}"

Arquivo conda.yaml:

name: mlflow-env
channels:
  - conda-forge
dependencies:
  - python=3.9
  - pip
  - pip:
    - mlflow
    - scikit-learn
    - xgboost

Execute o pipeline:

mlflow run . -P learning_rate=0.05 -P n_estimators=200

7. Monitoramento e comparação avançada de experimentos

Para comparar múltiplas execuções via API:

runs = mlflow.search_runs(
    experiment_ids=["1"],
    filter_string="metrics.mse < 0.1",
    order_by=["metrics.mse ASC"]
)
print(runs[["run_id", "metrics.mse", "params.alpha"]])

Para visualizar gráficos comparativos, use a interface web do MLflow (http://localhost:5000). Selecione múltiplos runs e compare métricas como acurácia, MSE ou tempo de treinamento ao longo do tempo.

Para alertas de degradação:

def verificar_degradacao(experiment_id, metrica_limite=0.1):
    runs = mlflow.search_runs(experiment_ids=[experiment_id])
    ultimo_mse = runs.iloc[-1]["metrics.mse"]
    if ultimo_mse > metrica_limite:
        enviar_alerta(f"MSE alto detectado: {ultimo_mse}")

8. Integração do MLflow com ferramentas de produção e deploy

Exemplo de serviço FastAPI que carrega modelo do Registry:

from fastapi import FastAPI
import mlflow.pyfunc

app = FastAPI()

# Carregar modelo da Production
model = mlflow.pyfunc.load_model(
    "models:/classificador_xgb/Production"
)

@app.post("/predict")
async def predict(data: dict):
    prediction = model.predict([list(data.values())])
    return {"prediction": prediction.tolist()}

Para CI/CD com GitHub Actions:

name: Deploy Model
on:
  release:
    types: [published]

jobs:
  deploy:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v2
      - name: Train and register model
        run: |
          pip install mlflow
          mlflow run . -P learning_rate=0.01
          mlflow models transition-model-version \
            --name classificador_xgb \
            --version ${{ github.event.release.tag_name }} \
            --stage Production

Estratégias de A/B testing podem usar diferentes versões registradas:

# Versão A (atual)
model_a = mlflow.pyfunc.load_model("models:/classificador_xgb/Production")

# Versão B (candidata)
model_b = mlflow.pyfunc.load_model("models:/classificador_xgb/Staging")

# Roteamento: 90% tráfego para A, 10% para B
if random.random() < 0.1:
    resultado = model_b.predict(dados)
else:
    resultado = model_a.predict(dados)

O MLflow oferece uma plataforma completa para rastrear experimentos, versionar modelos e gerenciar o ciclo de vida completo de machine learning, desde a prototipação até o deploy em produção com governança e reprodutibilidade.

Referências