Seguimiento de las ejecuciones con MLflow y la página de ejecución de Trabajos

Cada carga de trabajo que se envía con air run es tanto una ejecución de una tarea de Databricks como una ejecución de MLflow:

• La ejecución del trabajo (visible en la página Trabajos y canalizaciones del área de trabajo) realiza un seguimiento de la ejecución: estado, proceso, reintentos y salida del controlador.
• La ejecución de MLflow realiza un seguimiento del experimento: parámetros, métricas, métricas del sistema y artefactos.

Un envío crea una ejecución de trabajo y una ejecución de MLflow. Un reintento crea una nueva ejecución de MLflow.

Experimentos y ejecuciones

Dos campos YAML de carga de trabajo controlan cómo aparece la ejecución en 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 (Obligatorio): crea un experimento de MLflow con este nombre si no existe o anexa una nueva ejecución al experimento existente. Un experimento incluye muchas ejecuciones.
• mlflow_run_name (Opcional): establece el nombre de ejecución. Si se omite, el nombre de ejecución tiene como valor predeterminado el nombre del experimento (experiment_name).
• max_retries (Opcional): cada intento de reintento es una nueva ejecución de MLflow en el mismo experimento, por lo que puede comparar los intentos. La solicitud original y sus reintentos comparten una misma ejecución de tarea.

Navega entre Trabajos, MLflow y cargas de trabajo previas

Puede acceder a una ejecución desde tres sitios:

• Trabajos: en la página Ejecución de trabajos se enumeran las ejecuciones y cada ejecución se vincula a su ejecución y experimento de MLflow.
• MLflow: en la página Experimentos se enumeran los experimentos de MLflow.
• Workloads anteriores: air get run <job-run-id> muestra enlaces en los que se puede hacer clic al trabajo, al experimento y a la ejecución de MLflow de la ejecución. air list runs enumera las ejecuciones anteriores y le permite filtrar para encontrar una ejecución 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 del sistema

Las métricas del sistema de GPU, CPU y memoria se capturan automáticamente para cada ejecución. No se requiere ninguna configuración. Puede verlos en la pestaña Métricas del sistema de la ejecución de MLflow.

Registro de métricas personalizadas

La plataforma crea la ejecución de MLflow y pone su identificador a disposición de tu proceso de entrenamiento a través de la variable de entorno MLFLOW_RUN_ID. Use la API de seguimiento de MLflow para registrar sus propios parámetros, métricas y artefactos en esa ejecución.

En las cargas de trabajo distribuidas (multinodo), cada nodo comparte la misma ejecución de MLflow. Genere registros solo desde el proceso rank-0, de modo que cada métrica se registre una 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)

Registros y artefactos

Transmita o descargue los registros de una ejecución con 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

Los registros también están disponibles como artefactos en la ejecución de MLflow. Para conservar los puntos de control del modelo, escríbalos en un volumen de catálogo de Unity. Para conocer los patrones de punto de control y la administración de volúmenes, consulte Seguimiento de experimentos y observabilidad.

Recursos adicionales

• Inicio rápido de la CLI de AI Runtime
• Referencia de YAML para cargas de trabajo
• Seguimiento de experimentos y observabilidad