Inizia a utilizzare l'SDK ML Diagnostics

L'SDK Python ML Diagnostics può essere integrato con i workload ML per raccogliere e gestire metriche, configurazioni e profili dei modelli su Google Cloud. Utilizza l'SDK per la profilazione programmatica e per visualizzare le metriche dei modelli. Le metriche dei modelli non sono disponibili se utilizzi la profilazione on demand senza l'SDK.

Questa guida mostra come creare esecuzioni di machine learning, raccogliere e gestire metriche e configurazioni dei workload, eseguire il deployment delle risorse XProf gestite e abilitare l'acquisizione di profili programmatici e on demand.

Per ulteriori informazioni sull'utilizzo dell'SDK ML Diagnostics, consulta il repository google-cloud-mldiagnostics.

Installa l'SDK ML Diagnostics

Installa la google-cloud-mldiagnostics libreria:

pip install google-cloud-mldiagnostics

Importa i seguenti pacchetti nel codice del workload ML:

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

Abilita Cloud Logging

L'SDK utilizza il modulo logging Python standard per generare metriche e informazioni di configurazione. Per indirizzare questi log a Cloud Logging, installa e configura la libreria google-cloud-logging. In questo modo, puoi visualizzare i log dell'SDK, le metriche registrate e i log delle tue applicazioni all'interno della Google Cloud console.

Installa la libreria google-cloud-logging:

pip install google-cloud-logging

Configura il logging nello script collegando il gestore di Cloud Logging al logger root di Python. Aggiungi le seguenti righe all'inizio dello 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.")

Abilita il logging dettagliato

Per impostazione predefinita, il livello di logging è impostato su INFO. Per ricevere log più dettagliati dall'SDK, ad esempio i dettagli dell'esecuzione di machine learning, imposta il livello di logging su DEBUG dopo aver chiamato 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.")

Con DEBUG abilitato, ricevi ulteriori diagnostiche dell'SDK in Cloud Logging. Ad esempio:

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', ...}

Crea un'esecuzione di machine learning

Per utilizzare la piattaforma ML Diagnostics, devi prima creare un'esecuzione di machine learning. Ciò comporta l'instrumentazione del workload ML con l'SDK per eseguire il logging, raccogliere metriche e abilitare la tracciabilità dei profili.

Di seguito è riportato un esempio di base che inizializza Cloud Logging, crea un'esecuzione di machine learning (MLRun), registra le metriche e acquisisce un profilo:

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'esempio di codice utilizza le seguenti variabili:

Variabile Requisito Descrizione
name Obbligatorio Un identificatore per l'esecuzione specifica. L'SDK crea automaticamente un machine-learning-run-id per garantire che i nomi delle esecuzioni siano univoci.
run_group Facoltativo Un identificatore che può aiutarti a raggruppare più esecuzioni appartenenti allo stesso esperimento. Ad esempio, tutte le esecuzioni associate a una scansione delle dimensioni della slice TPU sweep potrebbero appartenere allo stesso gruppo.
project Facoltativo Se non specificato, il progetto viene estratto da Google Cloud CLI.
region Obbligatorio Sono supportate tutte le località di Cluster Director , ad eccezione di us-east5. Questo flag può essere impostato da un argomento per ogni comando o con il comando: gcloud config set compute/region.
configs Facoltativo Coppie chiave-valore contenenti i parametri di configurazione per l'esecuzione. Se le configurazioni non sono definite, vengono visualizzate le configurazioni software e di sistema predefinite ma non quelle del workload ML.
gcs_path Obbligatorio condizionalmente La Google Cloud località di archiviazione in cui vengono salvati tutti i profili. Ad esempio: gs://my-bucket o gs://my-bucket/folder1. Obbligatorio solo se l'SDK viene utilizzato per l'acquisizione dei profili.
on-demand-xprof Facoltativo Avvia xprofz daemon sulla porta 9999 per abilitare la profilazione on demand. Puoi abilitare sia la profilazione on demand sia quella programmatica nello stesso codice, a condizione che non si verifichino contemporaneamente.

Le seguenti configurazioni vengono raccolte automaticamente dall'SDK e non devono essere specificate in machinelearning_run:

  • Configurazioni software: framework, versione del framework, flag XLA.
  • Configurazioni di sistema: tipo di dispositivo, numero di slice, dimensioni delle slice, numero di host.

Le informazioni su progetto e regione vengono archiviate come metadati dell'esecuzione di machine learning. La regione utilizzata per l'esecuzione di machine learning non deve necessariamente corrispondere alla regione utilizzata per l'esecuzione del workload.

Scrivi le configurazioni

Molti workload contengono troppe configurazioni per essere definite direttamente nella definizione machinelearning_run. In questi casi, puoi scrivere le configurazioni per l'esecuzione utilizzando JSON o 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",
)

Raccogli le metriche

Puoi raccogliere le metriche dei modelli, le metriche sul rendimento dei modelli e le metriche di sistema con l'SDK. Puoi creare visualizzazioni di queste metriche come valori medi e con grafici delle serie temporali.

L'SDK fornisce due funzioni per la registrazione delle metriche: metrics.record() per l'acquisizione di singoli punti dati e metrics.record_metrics() per la registrazione di più metriche in un singolo batch. Entrambe le funzioni scrivono le metriche in Cloud Logging, consentendo la visualizzazione e l'analisi.

Per registrare una singola metrica:

# 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)

Per registrare più metriche:

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)

Le seguenti metriche di sistema vengono raccolte automaticamente dall'SDK dalle librerie libTPU, psutil e JAX:

  • Utilizzo di Tensor Core delle TPU
  • Ciclo di servizio TPU
  • Utilizzo HBM
  • Utilizzo CPU host
  • Utilizzo memoria host

Non è necessario specificare manualmente queste metriche. Queste metriche di sistema hanno time come asse x predefinito.

Le seguenti chiavi metriche predefinite verranno visualizzate automaticamente nella Google Cloud console se assegnate. Queste metriche non vengono calcolate automaticamente; sono chiavi predefinite a cui puoi assegnare valori.

  • Chiavi metriche di qualità del modello: LEARNING_RATE, LOSS, GRADIENT_NORM, TOTAL_WEIGHTS.
  • Chiavi metriche sul rendimento del modello: STEP_TIME, THROUGHPUT, LATENCY, MFU, TFLOPS.

Le metriche predefinite, nonché altre metriche definite dall'utente, possono essere registrate con l'asse x come time o sia time che step. Puoi registrare qualsiasi metrica personalizzata nel workload.

L'esempio seguente acquisisce una singola metrica per il workload, che puoi visualizzare nella scheda Metriche del modello per l'esecuzione di machine learning specifica:

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

Per registrare più metriche in una singola chiamata, utilizza il metodo record_metrics. Ad esempio:

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)

Acquisisci i profili

Puoi acquisire i profili XProf del tuo workload ML con l'acquisizione programmatica o l'acquisizione on demand (acquisizione manuale). L'acquisizione programmatica comporta l'incorporamento dei comandi di profilazione direttamente nel codice di machine learning e la dichiarazione esplicita di quando avviare e interrompere la registrazione dei dati. L'acquisizione on demand avviene in tempo reale, in cui attivi il profiler mentre il workload è già in esecuzione attiva.

I comandi dell'SDK per l'acquisizione dei profili sono indipendenti dal framework, poiché tutti i comandi di profilazione a livello di framework sono integrati automaticamente nei comandi di profilazione di ML Diagnostics. Ciò significa che il codice di profilazione non dipende dal framework che utilizzi.

Acquisizione programmatica dei profili

L'acquisizione programmatica richiede di annotare il codice del modello e specificare dove vuoi acquisire i profili. In genere, acquisisci un profilo per alcuni passaggi di addestramento o profili un blocco di codice specifico all'interno del modello.

Puoi eseguire l'acquisizione programmatica dei profili con l'SDK ML Diagnostics nei seguenti modi:

  • Raccolta basata su API: controlla la profilazione con start() e stop() metodi.
  • Raccolta basata su decoratore: annota le funzioni con @xprof(run) per la profilazione automatica.
  • Gestore di contesto: utilizza xprof() per la profilazione basata sull'ambito che gestisce automaticamente le operazioni start() e stop().

Puoi utilizzare lo stesso codice di acquisizione dei profili in tutti i framework. Tutte le sessioni di profilazione vengono acquisite nel bucket Cloud Storage definito nell'esecuzione di 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()

Profilazione multi-host (processo)

Durante la profilazione programmatica, l'SDK avvia la profilazione su ogni host (processo) in cui viene eseguito il codice del workload ML. Se l'elenco dei nodi non viene fornito, vengono inclusi tutti gli host.

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

Per impostazione predefinita, la chiamata del metodo prof.start() senza l'argomento session_id su più host genera sessioni di traccia separate, una per ogni host. Per raggruppare le tracce di host diversi in una singola sessione multi-host unificata in XProf, assicurati che il metodo prof.start() venga chiamato con lo stesso argomento session_id su tutti gli host partecipanti. Ad esempio:

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

Per abilitare la profilazione per host specifici:

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

Acquisizione di profili on demand

Utilizza l'acquisizione di profili on demand quando vuoi acquisire i profili in modo ad hoc o quando l'acquisizione programmatica dei profili non è già abilitata. L'acquisizione on demand è utile quando si verificano problemi con il rendimento del modello durante l'esecuzione e vuoi acquisire i profili in quei momenti per diagnosticare i problemi.

Le metriche dei modelli (generate dall'SDK) non sono disponibili quando si utilizza la profilazione on demand. Le metriche di sistema TPU sono comunque disponibili nelle dashboard TPU di GKE.

Per abilitare l'acquisizione di profili on demand, configura l'esecuzione con il supporto on demand:

# 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
)

Puoi utilizzare lo stesso codice di acquisizione dei profili in tutti i framework. Tutte le sessioni di profilazione vengono acquisite nel bucket Cloud Storage definito nell'esecuzione di machine learning.

Per la profilazione on demand su GKE, esegui il deployment di connection-operator e injection-webhook di GKE nel cluster GKE. In questo modo, l'esecuzione di machine learning può individuare i nodi GKE su cui è in esecuzione e il menu a discesa di acquisizione on demand può popolare automaticamente questi nodi. Per ulteriori informazioni, consulta Configura il cluster GKE.

Pacchetto del workload per GKE

Puoi utilizzare un Dockerfile per creare il pacchetto di un'applicazione che utilizza l'SDK ML Diagnostics. Installa il pacchetto google-cloud-logging per l'integrazione di Cloud Logging. Ad esempio:

# 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"]

Esegui il deployment del workload

Dopo aver integrato l'SDK con il workload, crea il pacchetto del workload in un'immagine e crea il file YAML con l'immagine specificata. Per la profilazione programmatica, etichetta il workload nel file YAML con managed-mldiagnostics-gke=true. La profilazione on demand non richiede questa etichetta.

Per GKE:

kubectl apply -f YAML_FILE_NAME

Per Compute Engine, connettiti alla VM tramite SSH ed esegui il codice Python per il workload:

source venv/bin/activate
python3.11 WORKLOAD_FILE_NAME

Dopo aver eseguito il deployment del workload, trova il nome del job cercando lo spazio dei nomi del workload:

kubectl get job -n YOUR_NAMESPACE

Puoi trovare il nome e il link dell'esecuzione nei log kubectl passando il nome e lo spazio dei nomi del job. Devi specificare il container del workload (ad esempio: -c workload) perché il sidecar ML Diagnostics gestisce il proprio logging.

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