Premiers pas avec le SDK ML Diagnostics

Le SDK Python ML Diagnostics peut être intégré aux charges de travail de ML pour collecter et gérer les métriques, les configurations et les profils de charge de travail sur Google Cloud. Ce guide vous explique comment créer des exécutions de machine learning, collecter et gérer les métriques et les configurations de charge de travail, déployer des ressources XProf gérées et activer la capture de profil programmatique et à la demande.

Pour en savoir plus sur l'utilisation du SDK ML Diagnostics, consultez le dépôt google-cloud-mldiagnostics.

Installer le SDK ML Diagnostics

Installez la bibliothèque google-cloud-mldiagnostics :

pip install google-cloud-mldiagnostics

Importez les packages suivants dans le code de votre charge de travail de ML :

from google_cloud_mldiagnostics import machinelearning_run
from google_cloud_mldiagnostics import metrics
from google_cloud_mldiagnostics import xprof

Activer Cloud Logging

Le SDK utilise le module Python standard logging pour générer des métriques et des informations de configuration. Pour acheminer ces journaux vers Cloud Logging, installez et configurez la bibliothèque google-cloud-logging. Cela vous permet d'afficher les journaux du SDK, les métriques enregistrées et les journaux de votre propre application dans la console Google Cloud .

Installez la bibliothèque google-cloud-logging :

pip install google-cloud-logging

Configurez la journalisation dans votre script en associant le gestionnaire Cloud Logging à l'enregistreur racine Python. Ajoutez les lignes suivantes au début de votre script Python :

  import logging
  import google.cloud.logging

  # Instantiate a Cloud Logging client
  logging_client = google.cloud.logging.Client()

  # Attach the Cloud Logging handler to the Python root logger
  logging_client.setup_logging()

  # Standard logging calls will go to Cloud Logging
  logging.info("SDK logs and application logs will appear in Cloud Logging.")

Activer la journalisation détaillée

Par défaut, le niveau de journalisation est défini sur INFO. Pour recevoir des journaux plus détaillés du SDK, tels que les détails d'exécution du machine learning, définissez le niveau de journalisation sur DEBUG après avoir appelé setup_logging() :

import logging
import google.cloud.logging

logging_client = google.cloud.logging.Client()
logging_client.setup_logging()
logging.getLogger().setLevel(logging.DEBUG) # Enable DEBUG level logs

logging.debug("This is a debug message.")
logging.info("This is an info message.")

Lorsque DEBUG est activé, vous recevez des diagnostics SDK supplémentaires dans Cloud Logging. Exemple :

DEBUG:google_cloud_mldiagnostics.core.global_manager:current run details:
{'name': 'projects/my-gcp-project/locations/us-central1/mlRuns/my-run-12345',
'gcs_path': 'gs://my-bucket/profiles', ...}

Créer une exécution de machine learning

Pour utiliser la plate-forme ML Diagnostics, vous devez d'abord créer une exécution de machine learning. Cela implique d'instrumenter votre charge de travail de ML avec le SDK pour effectuer la journalisation, collecter des métriques et activer le traçage de profil.

Voici un exemple de base qui initialise Cloud Logging, crée une exécution de machine learning (MLRun), enregistre des métriques et capture un profil :

import logging
import os
import google.cloud.logging
from google_cloud_mldiagnostics import machinelearning_run, metrics, xprof, metric_types

# 1. Set up Cloud Logging
# Make sure to pip install google-cloud-logging
logging_client = google.cloud.logging.Client()
logging_client.setup_logging()
# Optional: Set logging level to DEBUG for more detailed SDK logs
logging.getLogger().setLevel(logging.DEBUG)

# 2. Define and start machinelearning run
try:
    run = machinelearning_run(
          name="<run_name>",
          run_group="<run_group>",
          configs={ "epochs": 100, "batch_size": 32 },
          project="<some_project>",
          region="<some_zone>",
          gcs_path="gs://<some_bucket>",
          on_demand_xprof=True,
        )
    logging.info(f"MLRun created: {run.name}")

    # 3. Collect metrics during your run
    metrics.record(metric_types.MetricType.LOSS, 0.123, step=1)
    logging.info("Loss metric recorded.")

    # 4. Capture profiles programmatically
    with xprof():
        # ... your code to profile here ...
        pass
    logging.info("Profile captured.")

except Exception as e:
    logging.error(f"Error during MLRun: {e}", exc_info=True)

L'exemple de code utilise les variables suivantes :

Variable	Exigence	Description
name	Obligatoire	Identifiant de l'exécution spécifique. Le SDK crée automatiquement un machine-learning-run-id pour s'assurer que les noms d'exécution sont uniques.
run_group	Facultatif	Identifiant permettant de regrouper plusieurs exécutions appartenant au même test. Par exemple, toutes les exécutions associées à un balayage de la taille de tranche TPU peuvent appartenir au même groupe.
project	Facultatif	Si aucun projet n'est spécifié, il est extrait de Google Cloud CLI.
region	Obligatoire	Toutes les zones Cluster Director sont acceptées, à l'exception de us-east5. Ce flag peut être défini par un argument pour chaque commande ou avec la commande : gcloud config set compute/region.
configs	Facultatif	Paires clé/valeur contenant les paramètres de configuration de l'exécution. Si les configurations ne sont pas définies, les configurations logicielles et système par défaut s'affichent, mais pas celles de la charge de travail de ML.
gcs_path	Obligatoire sous certaines conditions	Emplacement de stockage Google Cloud où tous les profils sont enregistrés. Par exemple, gs://my-bucket ou gs://my-bucket/folder1. Obligatoire uniquement si le SDK est utilisé pour la capture de profil.
on-demand-xprof	Facultatif	Démarre xprofz daemon sur le port 9999 pour activer le profilage à la demande. Vous pouvez activer le profilage à la demande et le profilage programmatique dans le même code, à condition qu'ils ne se produisent pas en même temps.

Les configurations suivantes sont collectées automatiquement par le SDK et n'ont pas besoin d'être spécifiées dans machinelearning_run :

• Configurations logicielles : framework, version du framework, flags XLA.
• Configurations système : type d'appareil, nombre de tranches, taille des tranches, nombre d'hôtes.

Les informations sur le projet et la région sont stockées en tant que métadonnées d'exécution du machine learning. La région utilisée pour l'exécution du machine learning ne doit pas nécessairement correspondre à celle utilisée pour l'exécution de la charge de travail.

Écrire des configurations

De nombreuses charges de travail contiennent trop de configurations pour être définies directement dans la définition machinelearning_run. Dans ce cas, vous pouvez écrire des configurations pour votre exécution à l'aide de JSON ou YAML.

import yaml
import json

# Read the YAML file
with open('config.yaml', 'r') as yaml_file:
  # Parse YAML into a Python dictionary
  yaml_data = yaml.safe_load(yaml_file)

# Define machinelearning run
machinelearning_run(
  name="RUN_NAME",
  run_group="GROUP_NAME",
  configs=yaml_data,
  project="PROJECT_NAME",
  region="ZONE",
  gcs_path="gs://BUCKET_NAME",
)

Collecter des métriques

Vous pouvez collecter des métriques de modèle, des métriques de performances de modèle et des métriques système avec le SDK. Vous pouvez créer des visualisations de ces métriques sous forme de valeurs moyennes et de graphiques de séries temporelles.

Le SDK fournit deux fonctions pour enregistrer les métriques : metrics.record() pour capturer des points de données individuels et metrics.record_metrics() pour enregistrer plusieurs métriques dans un même lot. Les deux fonctions écrivent des métriques dans Cloud Logging, ce qui permet de les visualiser et de les analyser.

Pour enregistrer une seule métrique :

# Record a metric only with time as the x-axis
metrics.record(metric_types.MetricType.LOSS, 0.123)

# Record a metric with time and step as the x-axis
metrics.record(metric_types.MetricType.LOSS, 0.123, step=1)

Pour enregistrer plusieurs métriques :

from google_cloud_mldiagnostics import metric_types
# User codes
# machinelearning_run should be called
# ......

for step in range(num_steps):
  if (step + 1) % 10 == 0:
    metrics.record_metrics([
        # Model quality metrics
        {"metric_name": metric_types.MetricType.LEARNING_RATE, "value": step_size},
        {"metric_name": metric_types.MetricType.LOSS, "value": loss},
        {"metric_name": metric_types.MetricType.GRADIENT_NORM, "value": gradient},
        {"metric_name": metric_types.MetricType.TOTAL_WEIGHTS, "value": total_weights},
        # Model performance metrics
        {"metric_name": metric_types.MetricType.STEP_TIME, "value": step_time},
        {"metric_name": metric_types.MetricType.THROUGHPUT, "value": throughput},
        {"metric_name": metric_types.MetricType.LATENCY, "value": latency},
        {"metric_name": metric_types.MetricType.TFLOPS, "value": tflops},
        {"metric_name": metric_types.MetricType.MFU, "value": mfu},
    ], step=step+1)

Les métriques système suivantes sont automatiquement collectées par le SDK à partir des bibliothèques libTPU, psutil et JAX :

• Utilisation des TPU Tensor Core
• Cycle d'utilisation du TPU
• Utilisation de la mémoire HBM
• Utilisation du processeur de l'hôte
• Utilisation de la mémoire de l'hôte

Vous n'avez pas besoin de spécifier manuellement ces métriques. Ces métriques système ont time (temps) comme axe X par défaut.

Les clés de métriques prédéfinies suivantes s'affichent automatiquement dans la consoleGoogle Cloud si elles sont attribuées. Ces métriques ne sont pas calculées automatiquement. Il s'agit de clés prédéfinies auxquelles vous pouvez attribuer des valeurs.

• Clés des métriques de qualité du modèle : LEARNING_RATE, LOSS, GRADIENT_NORM, TOTAL_WEIGHTS.
• Clés des métriques de performances du modèle : STEP_TIME, THROUGHPUT, LATENCY, MFU, TFLOPS.

Les métriques prédéfinies, ainsi que d'autres métriques définies par l'utilisateur, peuvent être enregistrées avec l'axe X comme time, ou à la fois time et step. Vous pouvez enregistrer n'importe quelle métrique personnalisée dans la charge de travail.

L'exemple suivant capture une seule métrique pour la charge de travail, que vous pouvez afficher dans l'onglet Métriques du modèle pour l'exécution spécifique du machine learning :

metrics.record("custom_metrics_1", step_size, step=step + 1)

Pour enregistrer plusieurs métriques en un seul appel, utilisez la méthode record_metrics. Exemple :

metrics.record_metrics([
        # Model quality metrics
        {"metric_name": metric_types.MetricType.LEARNING_RATE, "value": step_size},
        {"metric_name": metric_types.MetricType.LOSS, "value": loss},
        {"metric_name": metric_types.MetricType.GRADIENT_NORM, "value": gradient},
        {"metric_name": metric_types.MetricType.TOTAL_WEIGHTS, "value": total_weights},
        # Model performance metrics
        {"metric_name": metric_types.MetricType.STEP_TIME, "value": step_time},
        {"metric_name": metric_types.MetricType.THROUGHPUT, "value": throughput},
        {"metric_name": metric_types.MetricType.LATENCY, "value": latency},
        {"metric_name": metric_types.MetricType.TFLOPS, "value": tflops},
        {"metric_name": metric_types.MetricType.MFU, "value": mfu},
     # Custom metrics
     {"custom_metrics_1", "value":<value>},
     {"custom_metrics_2", "value":<value>},
     {"avg_mtp_acceptance_rate_percent", "value":<value>},
     {"dpo_reward_accuracy", "value":<value>},
    ], step=step+1)

Capturer des profils

Vous pouvez capturer des profils XProf de votre charge de travail de ML avec la capture programmatique ou la capture à la demande (capture manuelle). La capture programmatique consiste à intégrer des commandes de profilage directement dans votre code de machine learning et à indiquer explicitement quand commencer et arrêter l'enregistrement des données. La capture à la demande se produit en temps réel. Vous déclenchez le profileur alors que la charge de travail est déjà en cours d'exécution.

Les commandes du SDK permettant de capturer des profils sont indépendantes du framework, car toutes les commandes de profilage au niveau du framework sont automatiquement intégrées aux commandes de profilage ML Diagnostics. Cela signifie que votre code de profilage ne dépend pas du framework que vous utilisez.

Capture programmatique de profils

La capture programmatique vous oblige à annoter le code de votre modèle et à spécifier les endroits où capturer les profils. En règle générale, il convient de capturer un profil pour quelques étapes d'entraînement ou de profiler un bloc de code spécifique dans le modèle.

Vous pouvez effectuer la capture programmatique de profils avec le SDK ML Diagnostics de différentes manières :

• Collecte basée sur l'API : contrôlez le profilage avec les méthodes start() et stop().
• Collecte basée sur les décorateurs : annotez les fonctions avec @xprof(run) pour le profilage automatique.
• Gestionnaire de contexte : à utiliser avec xprof() pour le profilage basé sur la portée qui gère automatiquement les opérations start() et stop().

Vous pouvez utiliser le même code de capture de profil dans tous les frameworks. Toutes les sessions de profil sont capturées dans le bucket Cloud Storage défini dans l'exécution du machine learning.

# Support collection via APIs
prof = xprof()  # Updates metadata and starts xprofz collector
prof.start()  # Collects traces to bucket
# ..... Your code execution here
# ....
prof.stop()

# Also supports collection via decorators
@xprof()
def abc(self):
    # does something
    pass

# Use xprof as a context manager to automatically start and stop collection
with xprof() as prof:
    # Your training or execution code here
    train_model()
    evaluate_model()

Profilage multihôte (processus)

Lors du profilage programmatique, le SDK commence le profilage sur chaque hôte (processus) où le code de charge de travail de ML est en cours d'exécution. Si la liste des nœuds n'est pas fournie, tous les hôtes sont inclus.

# starts profiling on all nodes
prof = xprof()
prof.start()
# ...
prof.stop()

Par défaut, l'appel de la méthode prof.start() sans l'argument session_id sur plusieurs hôtes entraîne des sessions de trace distinctes, une pour chaque hôte. Pour regrouper les traces de différents hôtes dans une seule session multihôte unifiée dans XProf, assurez-vous que la méthode prof.start() est appelée avec le même argument session_id sur tous les hôtes participants. Exemple :

# Use the same session_id on all hosts to group traces
prof = xprof()
prof.start(session_id="profiling_session")
# ...
prof.stop()

Pour activer le profilage pour des hôtes spécifiques :

# starts profiling on node with index 0 and 2
prof = xprof(process_index_list=[0,2])
prof.start()
# ...
prof.stop()

Capture de profil à la demande

Utilisez la capture de profil à la demande si vous souhaitez capturer des profils de manière ad hoc ou si la capture de profil programmatique n'est pas déjà activée. La capture à la demande est utile en cas de problèmes avec les métriques du modèle pendant l'exécution, et que vous souhaitez capturer des profils à ces moments-là pour diagnostiquer les problèmes.

Pour activer la capture de profil à la demande, configurez l'exécution avec l'assistance à la demande :

# Define machinelearning run
machinelearning_run(
    name="<run_name>",
    # specify where profiling data is stored
    gcs_path="gs://<bucket>",
    ...
    # enable on demand profiling, starts xprofz daemon on port 9999
    on_demand_xprof=True
)

Vous pouvez utiliser le même code de capture de profil dans tous les frameworks. Toutes les sessions de profil sont capturées dans le bucket Cloud Storage défini dans l'exécution du machine learning.

Pour le profilage à la demande sur GKE, déployez connection-operator et injection-webhook dans le cluster GKE. Cela garantit que votre exécution de machine learning peut localiser les nœuds GKE sur lesquels elle s'exécute, et que le menu déroulant de capture à la demande peut remplir automatiquement ces nœuds. Pour en savoir plus, consultez Configurer un cluster GKE.

Empaqueter une charge de travail pour GKE

Vous pouvez utiliser un fichier Dockerfile pour empaqueter une application qui utilise le SDK ML Diagnostics. Installez le package google-cloud-logging pour l'intégration de Cloud Logging. Exemple :

# Base image (user's choice, e.g., python:3.10-slim, or a base with ML frameworks)
FROM python:3.11-slim

# Install base utilities
RUN pip install --no-cache-dir --upgrade pip

# Install SDK and Logging client
# psutil is installed as a dependency of google-cloud-mldiagnostics
RUN pip install --no-cache-dir \
    google-cloud-mldiagnostics \
    google-cloud-logging

# Optional: For JAX/TPU workloads
# RUN pip install --no-cache-dir "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html &&
#     pip install --no-cache-dir libtpu xprof

# Add your application code
COPY ./app /app
WORKDIR /app

# Run your script
CMD ["python", "your_train_script.py"]

Déployer une charge de travail

Après avoir intégré le SDK à votre charge de travail, packagez la charge de travail dans une image et créez votre fichier YAML avec l'image spécifiée. Ajoutez le libellé managed-mldiagnostics-gke=true à la charge de travail dans le fichier YAML.

Pour GKE :

kubectl apply -f YAML_FILE_NAME

Pour Compute Engine, connectez-vous à la VM à l'aide de SSH et exécutez le code Python pour votre charge de travail :

source venv/bin/activate
python3.11 WORKLOAD_FILE_NAME

Après avoir déployé la charge de travail, recherchez le nom de votre job en recherchant l'espace de noms de votre charge de travail :

kubectl get job -n YOUR_NAMESPACE

Vous pouvez trouver le nom et le lien de l'exécution dans vos journaux kubectl en transmettant le nom et l'espace de noms du job. Vous devez spécifier le conteneur de charge de travail (par exemple, -c workload), car le side-car ML Diagnostics gère sa propre journalisation.

kubectl logs jobs/s5-tpu-slice-0 -n YOUR_NAMESPACE -c workload