Acompanhe as execuções com o MLflow e a página de execução de Jobs

Important

A CLI do ambiente de execução de IA está em fase Beta.

Cada carga de trabalho que submetes com air run é simultaneamente uma execução de uma tarefa do Databricks e uma execução do MLflow:

  • A execução da tarefa (visível na página Jobs & Pipelines do espaço de trabalho) monitoriza a execução: estado, recursos de computação, novas tentativas e saída do controlador.
  • A execução do MLflow acompanha o experimento: parâmetros, métricas, métricas do sistema e artefactos.

Uma submissão cria uma execução de tarefa e uma execução do MLflow. Uma nova tentativa cria uma nova execução do MLflow.

Experiências e execuções

Dois campos YAML de carga de trabalho controlam como a execução aparece no MLflow:

experiment_name: my-training # Creates or appends to this MLflow experiment
mlflow_run_name: baseline-lr3e5 # Names the MLflow run for this submission
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
command: torchrun --nproc_per_node=8 train.py
max_retries: 2
  • experiment_name (Obrigatório): Cria uma experiência MLflow com este nome se não existir, ou adiciona uma nova execução à experiência existente. Um experimento tem muitas corridas.
  • mlflow_run_name (Opcional): Define o nome da corrida. Se for omitido, o nome da execução assume por predefinição o nome da experiência (experiment_name).
  • max_retries (Opcional): Cada nova tentativa é uma nova execução no MLflow no mesmo experimento, para poder comparar as tentativas. A submissão original e as suas repetições partilham a mesma execução da tarefa.

Página de execução MLflow que mostra métricas

Podes chegar a uma corrida a partir de três locais:

  • A página de execuções de Jobs lista as suas execuções e cada execução remete para a respetiva execução e experiência do MLflow.
  • MLflow: A página de Experiências lista as tuas experiências MLflow.
  • Execuções anteriores: air get run <job-run-id> apresenta ligações clicáveis para a tarefa, a experiência e a execução do MLflow desta execução. air list runs Lista as tuas corridas anteriores e permite-te filtrar para encontrar uma corrida específica.
air get run <job-run-id> # Links to the job, experiment, and MLflow run
air list runs # List previous runs; filter to find a specific run

Métricas do sistema

As métricas do sistema de GPU, CPU e memória são capturadas automaticamente em cada execução. Nenhuma configuração é necessária. Consulte-as no separador Métricas do sistema da execução do MLflow.

Aba de métricas do sistema de uma execução do MLflow (GPU/CPU/memória)

Registar métricas personalizadas

A plataforma cria a execução do MLflow e expõe o seu ID ao seu processo de treino através da variável de ambiente MLFLOW_RUN_ID. Usa a API de rastreamento MLflow para registar os teus próprios parâmetros, métricas e artefactos nessa execução.

Em cargas de trabalho distribuídas (multinó), todos os nós partilham a mesma execução do MLflow. Registar apenas a partir do processo rank-0, para que cada métrica seja registada uma vez:

import os

import mlflow

# Log from rank 0 only; all nodes share the same MLFLOW_RUN_ID.
if os.environ.get("RANK", "0") == "0":
    with mlflow.start_run(run_id=os.environ["MLFLOW_RUN_ID"]):
        mlflow.log_param("learning_rate", 3e-4)
        for step, loss in enumerate(training_losses):
            mlflow.log_metric("train_loss", loss, step=step)

Registos e artefactos

Transmita ou descarregue os registos de uma corrida com air logs:

air logs <job-run-id> # Stream logs from node 0
air logs <job-run-id> --node 2 # Logs from a specific node
air logs <job-run-id> --download-to ./logs/ # Download instead of streaming

Os registos estão também disponíveis como artefactos na execução do MLflow. Para guardar os checkpoints do modelo, grave-os num volume do Unity Catalog. Para padrões de pontos de verificação e gestão de volumes, consulte Rastreio de experiências e observabilidade.

Recursos adicionais