Spåra körningar med MLflow och sidan Jobbkörning

Important

Den här funktionen finns som allmänt tillgänglig förhandsversion.

Varje arbetsbelastning som du skickar med air run är både en Databricks-jobbkörning och en MLflow-körning:

  • Jobbkörningen (som visas på sidan Jobs & Pipelines i arbetsytan) spårar körningen: status, beräkningsresurser, återförsök och drivrutinsutdata.
  • MLflow-körningen spårar experimentet: parametrar, mått, systemmått och artefakter.

Ett inskick skapar en jobbkörning och en MLflow-körning. Ett nytt försök skapar en ny MLflow-körning.

Experiment och körningar

Två YAML-arbetsbelastningsfält styr hur körningen visas i 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 (Krävs): Skapar ett MLflow-experiment med det här namnet om det inte finns något, eller lägger till en ny körning i det befintliga experimentet. Ett experiment består av många körningar.
  • mlflow_run_name (Valfritt): Anger körningsnamnet. Om det utelämnas är standardvärdet för körningsnamnet experimentnamnet (experiment_name).
  • max_retries (Valfritt): Varje nytt försök är en ny MLflow-körning i samma experiment, så att du kan jämföra försök. Det ursprungliga inskicket och dess återförsök delar samma jobbkörning.

MLflow-körningssidan som visar mätvärden

Du kan komma till en körning från tre platser:

  • Jobb: Sidan Jobbkörning visar en lista över dina körningar och varje körning länkar till MLflow-körningen och experimentet.
  • MLflow: Sidan Experiment visar dina MLflow-experiment.
  • Tidigare arbetsbelastningar: air get run <job-run-id> skriver ut klickbara länkar till körningens jobb, experiment och MLflow-körning. air list runs visar en lista över dina tidigare körningar och låter dig filtrera för att hitta en specifik körning.
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

Systemmått

GPU-, CPU- och minnessystemmått registreras automatiskt för varje körning. Ingen konfiguration krävs. Visa dem på fliken Systemmått för MLflow-körning.

Systemmåttsfliken för en MLflow-körning (GPU/CPU/minne)

Logga anpassade mått

Plattformen skapar MLflow-körningen och exponerar sitt ID för din träningsprocess via MLFLOW_RUN_ID miljövariabeln. Använd MLflow Tracking API för att logga dina egna parametrar, mätvärden och artefakter i den körningen.

På distribuerade arbetsbelastningar (flera noder) delar varje nod samma MLflow-körning. Logga bara från rank-0-processen, så att varje mätvärde registreras en gång:

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)

Loggar och artefakter

Streama eller ladda ned körningsloggar med 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

Loggarna finns också tillgängliga som artefakter i MLflow-körningen. Om du vill spara modellkontrollpunkter skriver du dem till en Unity Catalog-volym. För kontrollpunktsmönster och hantering av volymer, se Experimentspårning och observerbarhet.

Ytterligare resurser

  • Last updated on