Visualiser les jobs avec Vertex AI TensorBoard

Si vous êtes intéressé par l'entraînement géré de Gemini Enterprise Agent Platform, contactez votre représentant commercial pour y accéder.

Avec l'entraînement géré, vous pouvez visualiser vos journaux d'entraînement en temps quasi réel à l'aide de Vertex AI TensorBoard. Il vous suffit de configurer votre charge de travail pour enregistrer les journaux dans un bucket Cloud Storage. Ils seront automatiquement diffusés dans l'interface TensorBoard pour analyse.

Prérequis

Avant de commencer, assurez-vous de disposer des éléments suivants :

  • Un cluster d'entraînement géré en cours d'exécution.
  • Un bucket Cloud Storage pour stocker vos journaux TensorBoard. Ce bucket doit se trouver dans la même région que votre instance TensorBoard. Pour obtenir des instructions de configuration, consultez la section Créer un bucket Cloud Storage.
  • Une instance Vertex AI TensorBoard. Pour obtenir des instructions de création, consultez la section Créer une instance Vertex AI TensorBoard.
  • Les autorisations IAM appropriées. Pour autoriser Cloud Storage FUSE à lire et à écrire dans le bucket de stockage, le compte de service utilisé par les VM de votre cluster doit disposer du rôle Utilisateur d'objets Storage (roles/storage.objectUser).

Activer l'importation TensorBoard

Pour configurer l'intégration TensorBoard pour votre tâche, transmettez les arguments suivants à l'aide de l'indicateur --extra dans l'envoi de votre tâche Slurm :

  • tensorboard_base_output_dir: spécifie le chemin d'accès Cloud Storage pour importer les journaux. Exemple : gs://my-bucket/my-logs.

  • tensorboard_url: spécifie l'URL de l'instance, de l'expérience ou de l'exécution Vertex AI TensorBoard. Si seule une instance est fournie, une expérience et une exécution sont créées. S'il est omis, l'instance TensorBoard par défaut du projet est utilisée. Exemple : projects/123/locations/us-central1/tensorboards/456.

Exemple

# Using specific tensorboard instance
sbatch --extra="tensorboard_base_output_dir=<your-cloud-storage-dir>,tensorboard_url=projects/<project-id>/locations/<location>/tensorboards/<tensorboard-instance-id>" your_script.sbatch

Écrire des journaux à partir de votre tâche d'entraînement

Dans votre script d'entraînement, accédez à la variable d'environnement AIP_TENSORBOARD_LOG_DIR. Cette variable fournit le chemin d'accès Cloud Storage unique dans lequel votre script doit écrire ses journaux TensorBoard.

Le chemin d'accès suit cette structure :

gs://<your-cloud-storage-path>/<cluster-id>-<cluster-uuid>/tensorboard/job-<job-id>/

L'exemple suivant montre un workflow complet avec deux composants clés : le script d'envoi Slurm qui configure la tâche et le script d'entraînement Python qui lit la variable d'environnement pour écrire ses journaux.

Script de tâche Slurm (simple_job.sbatch) :

#!/bin/bash
#SBATCH --job-name=tensorboard-simple-test
#SBATCH --output=tensorboard-simple-test-%j.out

# Activate your Python virtual environment if needed
# source /path/to/your/venv/bin/activate
python3 simple_logger.py

Script Python (simple_logger.py) :

import tensorflow as tf
import os

# Get the log directory from the environment variable
log_dir = os.environ.get("AIP_TENSORBOARD_LOG_DIR")

print(f"Writing TensorBoard logs to: {log_dir}")
writer = tf.summary.create_file_writer(log_dir)

with writer.as_default():
    for step in range(10):
        # Simulate some metrics
        loss = 1.0 - (step * 0.1)
        accuracy = 0.6 + (step * 0.04)

        # Log the metrics
        tf.summary.scalar('loss', loss, step=step)
        tf.summary.scalar('accuracy', accuracy, step=step)
        writer.flush()
        print(f"Step {step}: loss={loss:.4f}, accuracy={accuracy:.4f}")

writer.close()
print(f"--- Finished writing metrics to {log_dir} ---")

Synchronisation des journaux en temps réel

Pour visualiser les métriques d'une tâche en cours d'exécution, vous devez fermer et recréer périodiquement l'enregistreur de résumé dans votre code d'entraînement. Cela est nécessaire, car gcsfuse ne synchronise les fichiers journaux avec Cloud Storage qu'une fois qu'ils sont fermés. Cette technique de "vidage" permet de s'assurer que les résultats intermédiaires sont visibles dans la console TensorBoard avant la fin de la tâche.

Afficher Vertex AI TensorBoard

Une fois votre tâche envoyée, vous pouvez suivre sa progression en accédant à la page Vertex AI Experiments dans la Cloud de Confiance console.

Étape suivante

Analysez les métriques d'entraînement, visualisez les graphiques de modèle et profilez les performances dans Vertex AI TensorBoard pour déboguer les problèmes, optimiser votre modèle et le préparer à l'intégration dans un workflow MLOps complet.