Suivre les exécutions avec MLflow et la page d’exécution de Jobs

Important

Cette fonctionnalité est disponible en préversion publique.

Chaque charge de travail que vous soumettez avec air run est à la fois une exécution de tâche Databricks et une exécution MLflow :

  • L’exécution de la tâche (visible sur la page Jobs & Pipelines de l’espace de travail) permet de suivre l’exécution : état, ressources de calcul, nouvelles tentatives et sortie du pilote.
  • Le run MLflow assure le suivi de l’expérience : paramètres, métriques, métriques du système et artefacts.

Une soumission crée une exécution de travail et une exécution MLflow. Une nouvelle tentative crée une exécution MLflow.

Expériences et exécutions

Deux champs YAML de charge de travail contrôlent l’affichage de l’exécution dans 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 (Obligatoire) : crée une expérience MLflow avec ce nom s’il n’existe pas ou ajoute une nouvelle exécution à l’expérience existante. Une expérience comprend de nombreuses exécutions.
  • mlflow_run_name (Facultatif) : définit le nom d’exécution. S’il est omis, le nom d’exécution est défini par défaut sur le nom de l’expérience (experiment_name).
  • max_retries (Facultatif) : chaque nouvelle tentative est une nouvelle exécution MLflow dans la même expérience. Vous pouvez donc comparer les tentatives. La soumission d’origine et ses nouvelles tentatives partagent la même exécution de travail.

Page d’exécution MLflow montrant les métriques

Vous pouvez accéder à une exécution à partir de trois emplacements :

  • Travaux : la page d’exécution des travaux répertorie vos exécutions, et chaque exécution est liée à son exécution et à son expérience MLflow.
  • MLflow : la page Expériences répertorie vos expériences MLflow.
  • Exécutions précédentes : air get run <job-run-id> affiche des liens cliquables vers la tâche, l’expérience et l’exécution MLflow associés à cette exécution. air list runs répertorie vos exécutions précédentes et vous permet de filtrer pour rechercher une exécution spécifique.
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étriques du système

Les métriques système gpu, processeur et mémoire sont capturées automatiquement pour chaque exécution. Aucune configuration n’est requise. Affichez-les sous l’onglet Métriques système de l’exécution MLflow.

Onglet Métriques système d’une exécution MLflow (GPU/PROCESSEUR/mémoire)

Journaliser les métriques personnalisées

La plateforme crée l’exécution de MLflow et expose son ID à votre processus d’apprentissage via la variable d’environnement MLFLOW_RUN_ID . Utilisez l’API de suivi MLflow pour consigner vos propres paramètres, métriques et artefacts à cette exécution.

Sur les charges de travail distribuées (à plusieurs nœuds), chaque nœud partage la même exécution MLflow. Journal à partir du processus de rang 0 uniquement. Chaque métrique n'est donc enregistrée qu'une seule fois :

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)

Journaux et artefacts

Lisez en streaming ou téléchargez les logs d’une exécution avec 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

Les journaux d’activité sont également disponibles en tant qu’artefacts sur l’exécution de MLflow. Pour conserver les points de contrôle de modèle, écrivez-les dans un volume de catalogue Unity. Pour connaître les modèles de point de contrôle et la gestion des volumes, consultez Suivi des expériences et observabilité.

Ressources supplémentaires