Monitorare le esecuzioni con MLflow e la pagina di esecuzione dei processi

Importante

Questa funzionalità è in Anteprima Pubblica.

Ogni carico di lavoro inviato con databricks air run è sia un'esecuzione del processo Databricks che un'esecuzione MLflow:

  • L'esecuzione del job (visibile nella pagina Jobs & Pipelines del workspace) monitora l'esecuzione: stato, risorse di calcolo, tentativi e output del driver.
  • L'esecuzione di MLflow tiene traccia dell'esperimento: parametri, metriche, metriche di sistema e artefatti.

Ogni invio crea un'esecuzione del job e un'esecuzione di MLflow. Un nuovo tentativo crea una nuova esecuzione MLflow.

Esperimenti ed esecuzioni

I seguenti campi YAML del carico di lavoro controllano l'esperimento MLflow, l'esecuzione e l'archiviazione degli artefatti:

experiment_name: my-training # Creates or appends to this MLflow experiment
mlflow_run_name: baseline-lr3e5 # Names the MLflow run for this submission
# Stores artifacts in a UC volume
mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
command: torchrun --nproc_per_node=8 train.py
max_retries: 2
  • experiment_name (Obbligatorio): crea un esperimento MLflow con questo nome se non esiste o aggiunge una nuova esecuzione all'esperimento esistente. Un esperimento comprende molte esecuzioni.
  • mlflow_run_name (Facoltativo): imposta il nome dell'esecuzione. Se omesso, il nome dell'esecuzione viene impostato per impostazione predefinita sul nome dell'esperimento (experiment_name).
  • mlflow_artifact_location (Opzionale): Imposta la posizione radice degli artefatti. Se omesso, un nuovo esperimento utilizza la posizione DBFS predefinita, e un esperimento esistente utilizza la posizione legata al experiment_name. Per un volume di Unity Catalog, usa /Volumes/<catalog>/<schema>/<volume>/.... Per un esperimento esistente, la posizione specificata deve corrispondere alla posizione dell'artefatto dell'esperimento o essere omessa.
  • max_retries (Facoltativo): ogni tentativo di ripetizione è una nuova esecuzione di MLflow nello stesso esperimento, in modo da poter confrontare i tentativi. L'invio originale e i relativi tentativi condividono un'unica esecuzione del job.

Pagina di esecuzione di MLflow che mostra le metriche

È possibile passare a un'esecuzione da tre posizioni:

  • Processi: la pagina delle esecuzioni dei processi elenca le tue esecuzioni e ogni esecuzione rimanda alla relativa esecuzione e al relativo esperimento in MLflow.
  • MLflow: la pagina Esperimenti elenca gli esperimenti MLflow.
  • Carichi di lavoro precedenti: databricks air get <job-run-id> stampa collegamenti selezionabili al processo, all'esperimento e all'esecuzione di MLflow dell'esecuzione. databricks air list --all-status elenca le esecuzioni precedenti e consente di filtrare per trovare un'esecuzione specifica.
databricks air get <job-run-id> # Links to the job, experiment, and MLflow run
databricks air list --all-status # List previous runs; filter to find a specific run

Metriche di sistema

Le metriche del sistema GPU, CPU e memoria vengono acquisite automaticamente per ogni esecuzione. Non è necessaria alcuna configurazione. Visualizzarli nella scheda Metriche di sistema dell'esecuzione di MLflow.

Scheda Metriche di sistema di un'esecuzione MLflow (GPU/CPU/memoria)

Registrare le metriche personalizzate

La piattaforma crea l'esecuzione di MLflow ed espone il relativo ID al processo di training tramite la MLFLOW_RUN_ID variabile di ambiente. Utilizza l'API di tracciamento di MLflow per registrare i tuoi parametri, metriche e artefatti in quell'esecuzione.

Nei carichi di lavoro distribuiti (multinodo) ogni nodo condivide la stessa esecuzione di MLflow. Registra i log solo dal processo rank-0, in modo che ogni metrica venga registrata una sola volta:

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)

Log e artefatti

Trasmettere o scaricare i log di un'esecuzione con databricks air logs:

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

I log sono disponibili anche come artefatti nell'esecuzione di MLflow. Per salvare in modo permanente i checkpoint del modello, scrivili in un volume di Unity Catalog. Per i pattern di checkpointing, vedi Migliorare le prestazioni di addestramento e la resilienza su AI Runtime.

Risorse aggiuntive