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

Importante

A CLI do AI Runtime está em Beta.

Cada carga de trabalho enviada com air run corresponde tanto a uma execução de job no Databricks quanto a uma execução no MLflow:

  • A execução do trabalho (visível na página Trabalhos e Pipelines do workspace) acompanha a execução: status, computação, novas tentativas e saída do driver.
  • Uma execução no MLflow registra o experimento: parâmetros, métricas, métricas do sistema e artefatos.

Um envio cria uma execução de trabalho e uma execução do MLflow. Uma repetição cria uma nova execução do MLflow.

Experimentos e execuções

Dois campos YAML de carga de trabalho controlam como a execução é exibida 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 um experimento do MLflow com esse nome se um não existir ou acrescenta uma nova execução ao experimento existente. Um experimento contém muitas execuções.
  • mlflow_run_name (Opcional): define o nome da execução. Se omitido, o nome da execução será definido como o nome do experimento (experiment_name).
  • max_retries (Opcional): cada tentativa de repetição é uma nova execução de MLflow no mesmo experimento, para que você possa comparar tentativas. A submissão original e suas novas tentativas compartilham uma execução da tarefa.

Página de execução do MLflow mostrando métricas

Você pode acessar uma execução a partir de três lugares:

  • Jobs: A página de execuções dos Jobs lista suas execuções, e cada execução contém links para sua execução e seu experimento no MLflow.
  • MLflow: a página Experimentos lista seus experimentos do MLflow.
  • Execuções anteriores: air get run <job-run-id> exibe links clicáveis para o job, o experimento e a execução no MLflow da execução. air list runs lista as execuções anteriores e permite filtrar para localizar uma execução 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 para cada execução. Nenhuma configuração é necessária. Veja-os na guia Métricas do sistema da execução do MLflow.

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

Registrar métricas personalizadas

A plataforma cria a execução do MLflow e expõe sua ID ao processo de treinamento por meio da variável de MLFLOW_RUN_ID ambiente. Use a API de acompanhamento do MLflow para registrar seus próprios parâmetros, métricas e artefatos nessa execução.

Em cargas de trabalho distribuídas (vários nós), cada nó compartilha a mesma execução do MLflow. Faça logon somente do processo de classificação 0, para que cada métrica seja registrada 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)

Logs e artefatos

Transmitir ou baixar os logs de uma execução 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 logs também estão disponíveis como artefatos na execução do MLflow. Para persistir pontos de verificação de modelo, escreva-os em um volume do Catálogo do Unity. Para padrões de ponto de verificação e gerenciamento de volumes, consulte Acompanhamento e observabilidade de experimentos.

Recursos adicionais