Visualiser les jobs avec Vertex AI TensorBoard

Si vous êtes intéressé par les clusters d'entraînement gérés sur Gemini Enterprise Agent Platform, contactez votre conseiller commercial pour y accéder.

Avec les clusters d'entraînement gérés sur Gemini Enterprise Agent Platform, 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 ensuite automatiquement diffusés vers l'interface TensorBoard pour analyse.

Prérequis

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

  • Un cluster de clusters d'entraînement gérés 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 Créer un bucket Cloud Storage.
  • Une instance Vertex AI TensorBoard. Pour obtenir des instructions de création, consultez 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 de TensorBoard pour votre job, transmettez les arguments suivants à l'aide de l'option --extra dans l'envoi de votre job Slurm :

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

  • tensorboard_url : spécifie l'URL de l'instance, du test ou de l'exécution Vertex AI TensorBoard. Si seule une instance est fournie, un test et une exécution sont créés. Si elle est omise, 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 job 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 où 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 le job et le script d'entraînement Python qui lit la variable d'environnement pour écrire ses journaux.

Script de job 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'un job en cours d'exécution, vous devez régulièrement fermer et recréer le rédacteur de résumé dans votre code d'entraînement. Cette étape 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" garantit que les résultats intermédiaires sont visibles dans la console TensorBoard avant la fin du job.

Afficher Vertex AI TensorBoard

Une fois votre job envoyé, vous pouvez suivre sa progression en accédant à la page Tests sur Gemini Enterprise Agent Platform dans la console Google Cloud .

Étapes suivantes

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.