Comienza a usar el SDK de ML Diagnostics

El SDK de ML Diagnostics para Python se puede integrar en cargas de trabajo de AA para recopilar y administrar métricas, parámetros de configuración y perfiles de cargas de trabajo en Google Cloud. En esta guía, se muestra cómo crear ejecuciones de aprendizaje automático, recopilar y administrar métricas y configuraciones de cargas de trabajo, implementar recursos administrados de XProf y habilitar la captura de perfiles programática y a pedido.

Para obtener más información sobre cómo usar el SDK de ML Diagnostics, consulta el repositorio de google-cloud-mldiagnostics.

Instala el SDK de ML Diagnostics

Instala la biblioteca de google-cloud-mldiagnostics:

pip install google-cloud-mldiagnostics

Importa los siguientes paquetes en el código de tu carga de trabajo de AA:

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

Habilita Cloud Logging

El SDK usa el módulo logging estándar de Python para generar información de configuración y métricas. Para enrutar estos registros a Cloud Logging, instala y configura la biblioteca google-cloud-logging. Esto te permite ver los registros del SDK, las métricas registradas y los registros de tu propia aplicación en la consola de Google Cloud .

Instala la biblioteca de google-cloud-logging:

pip install google-cloud-logging

Configura el registro en tu secuencia de comandos adjuntando el controlador de Cloud Logging al registrador raíz de Python. Agrega las siguientes líneas al comienzo de tu secuencia de comandos de 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.")

Habilita el registro detallado

De forma predeterminada, el nivel de registro se establece en INFO. Para recibir registros más detallados del SDK, como los detalles de la ejecución del aprendizaje automático, establece el nivel de registro en DEBUG después de llamar a 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 habilitado, recibirás diagnósticos adicionales del SDK en Cloud Logging. Por ejemplo:

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

Crear una ejecución de aprendizaje automático

Para usar la plataforma de ML Diagnostics, primero debes crear una ejecución de aprendizaje automático. Esto implica instrumentar tu carga de trabajo de AA con el SDK para realizar el registro, recopilar métricas y habilitar el registro de perfiles.

A continuación, se muestra un ejemplo básico que inicializa Cloud Logging, crea una ejecución de aprendizaje automático (MLRun), registra métricas y captura un perfil:

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)

En el ejemplo de código, se usan las siguientes variables:

Variable	Requisito	Descripción
name	Obligatorio	Es un identificador de la ejecución específica. El SDK crea automáticamente un machine-learning-run-id para garantizar que los nombres de ejecución sean únicos.
run_group	Opcional	Es un identificador que puede ayudar a agrupar varias ejecuciones que pertenecen al mismo experimento. Por ejemplo, todas las ejecuciones asociadas con un análisis de variación del tamaño de la porción de TPU podrían pertenecer al mismo grupo.
project	Opcional	Si no se especifica, el proyecto se extrae de Google Cloud CLI.
region	Obligatorio	Se admiten todas las ubicaciones de Cluster Director, excepto us-east5. Esta marca se puede establecer con un argumento para cada comando o con el comando gcloud config set compute/region.
configs	Opcional	Pares clave-valor que contienen parámetros de configuración para la ejecución. Si no se definen configuraciones, aparecerán las configuraciones predeterminadas del software y del sistema, pero no las de la carga de trabajo de AA.
gcs_path	Condicionalmente obligatorio	Es la ubicación de almacenamiento Google Cloud en la que se guardan todos los perfiles. Por ejemplo: gs://my-bucket o gs://my-bucket/folder1. Solo se requiere si el SDK se usa para la captura de perfiles.
on-demand-xprof	Opcional	Inicia xprofz daemon en el puerto 9999 para habilitar la creación de perfiles a pedido. Puedes habilitar la generación de perfiles bajo demanda y la generación de perfiles programática en el mismo código, siempre y cuando no se produzcan al mismo tiempo.

El SDK recopila automáticamente las siguientes configuraciones y no es necesario especificarlas en machinelearning_run:

• Configuraciones de software: Framework, versión del framework y marcas de XLA
• Configuraciones del sistema: Tipo de dispositivo, cantidad de segmentos, tamaño de segmento y cantidad de hosts.

La información del proyecto y la región se almacena como metadatos de la ejecución del aprendizaje automático. La región que se usa para la ejecución del aprendizaje automático no tiene que coincidir con la región que se usa para la ejecución de la carga de trabajo.

Escribe archivos de configuración

Muchas cargas de trabajo contienen demasiadas configuraciones para definirlas directamente en la definición de machinelearning_run. En estos casos, puedes escribir archivos de configuración para tu ejecución con 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",
)

Recopila métricas

Con el SDK, puedes recopilar métricas del modelo, métricas de rendimiento del modelo y métricas del sistema. Puedes crear visualizaciones de estas métricas como valores promedio y con gráficos de series temporales.

El SDK proporciona dos funciones para registrar métricas: metrics.record() para capturar puntos de datos individuales y metrics.record_metrics() para registrar varias métricas en un solo lote. Ambas funciones escriben métricas en Cloud Logging, lo que permite la visualización y el análisis.

Para registrar una sola métrica, haz lo siguiente:

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

Para registrar varias métricas, haz lo siguiente:

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)

El SDK recopila automáticamente las siguientes métricas del sistema de las bibliotecas de libTPU, psutil y JAX:

• Uso de Tensor Core de TPU
• Ciclo de trabajo de TPU
• Uso de HBM
• Uso de CPU del host
• Uso de memoria del host

No es necesario que especifiques estas métricas de forma manual. Estas métricas del sistema tienen el tiempo como eje X predeterminado.

Las siguientes claves de métricas predefinidas aparecerán automáticamente en la consola deGoogle Cloud si se asignan. Estas métricas no se calculan automáticamente, sino que son claves predefinidas a las que puedes asignar valores.

• Claves de métricas de calidad del modelo: LEARNING_RATE, LOSS, GRADIENT_NORM y TOTAL_WEIGHTS.
• Claves de métricas de rendimiento del modelo: STEP_TIME, THROUGHPUT, LATENCY, MFU y TFLOPS.

Las métricas predefinidas, así como otras métricas definidas por el usuario, se pueden registrar con el eje X como time o con time y step. Puedes registrar cualquier métrica personalizada en la carga de trabajo.

En el siguiente ejemplo, se captura una sola métrica para la carga de trabajo, que puedes ver en la pestaña Métricas del modelo de la ejecución específica de aprendizaje automático:

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

Para registrar varias métricas en una sola llamada, usa el método record_metrics. Por ejemplo:

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)

Captura de perfiles

Puedes capturar perfiles de XProf de tu carga de trabajo de AA con la captura programática o la captura bajo demanda (captura manual). La captura programática implica incorporar comandos de creación de perfiles directamente en tu código de aprendizaje automático y establecer explícitamente cuándo comenzar y detener el registro de datos. La captura bajo demanda se produce en tiempo real, ya que activas el generador de perfiles mientras la carga de trabajo ya se está ejecutando de forma activa.

Los comandos del SDK para capturar perfiles son independientes del framework, ya que todos los comandos de generación de perfiles a nivel del framework se integran automáticamente en los comandos de generación de perfiles de ML Diagnostics. Esto significa que tu código de generación de perfiles no depende del framework que uses.

Captura de perfil programática

La captura programática requiere que anotes el código del modelo y especifiques dónde quieres capturar perfiles. En general, debes capturar un perfil para unos pocos pasos de entrenamiento o perfilar un bloque de código específico dentro del modelo.

Puedes realizar la captura de perfiles de forma programática con el SDK de ML Diagnostics de las siguientes maneras:

• Recopilación basada en la API: Controla la generación de perfiles con los métodos start() y stop().
• Recopilación basada en decoradores: Anota funciones con @xprof(run) para la generación automática de perfiles.
• Administrador de contexto: Úsalo con xprof() para la generación de perfiles basada en el alcance que controla automáticamente las operaciones de start() y stop().

Puedes usar el mismo código de captura de perfil en todos los frameworks. Todas las sesiones de perfil se capturan en el bucket de Cloud Storage definido en la ejecución de aprendizaje automático.

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

Creación de perfiles de varios hosts (procesos)

Durante la generación de perfiles programática, el SDK comienza a generar perfiles en cada host (proceso) en el que se ejecuta el código de la carga de trabajo de AA. Si no se proporciona la lista de nodos, se incluyen todos los hosts.

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

De forma predeterminada, llamar al método prof.start() sin el argumento session_id en varios hosts genera sesiones de registro separadas, una para cada host. Para agrupar los registros de diferentes hosts en una sola sesión unificada de varios hosts en XProf, asegúrate de que se llame al método prof.start() con el argumento session_id mismo en todos los hosts participantes. Por ejemplo:

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

Para habilitar la generación de perfiles para hosts específicos, haz lo siguiente:

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

Captura de perfiles a pedido

Usa la captura de perfiles bajo demanda cuando quieras capturar perfiles de forma ad hoc o si la captura de perfiles programática aún no está habilitada. La captura a pedido es útil cuando hay problemas con las métricas del modelo durante la ejecución y quieres capturar perfiles en esos momentos para diagnosticar los problemas.

Para habilitar la captura de perfiles a pedido, configura la ejecución con compatibilidad a pedido:

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

Puedes usar el mismo código de captura de perfil en todos los frameworks. Todas las sesiones de perfil se capturan en el bucket de Cloud Storage definido en la ejecución de aprendizaje automático.

Para la generación de perfiles a pedido en GKE, implementa connection-operator y injection-webhook de GKE en el clúster de GKE. Esto garantiza que tu ejecución de aprendizaje automático pueda ubicar los nodos de GKE en los que se ejecuta y que el menú desplegable de captura a pedido pueda propagar automáticamente esos nodos. Para obtener más información, consulta Configura el clúster de GKE.

Empaqueta la carga de trabajo para GKE

Puedes usar un Dockerfile para empaquetar una aplicación que use el SDK de ML Diagnostics. Instala el paquete google-cloud-logging para la integración de Cloud Logging. Por ejemplo:

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

Implementar carga de trabajo

Después de integrar el SDK en tu carga de trabajo, empaquétala en una imagen y crea tu archivo YAML con la imagen especificada. Etiqueta la carga de trabajo en el archivo YAML con managed-mldiagnostics-gke=true.

Para GKE:

kubectl apply -f YAML_FILE_NAME

En el caso de Compute Engine, conéctate a la VM con SSH y ejecuta el código de Python para tu carga de trabajo:

source venv/bin/activate
python3.11 WORKLOAD_FILE_NAME

Después de implementar la carga de trabajo, busca el nombre del trabajo buscando el espacio de nombres de la carga de trabajo:

kubectl get job -n YOUR_NAMESPACE

Puedes encontrar el nombre y el vínculo de la ejecución en tus registros de kubectl pasando el nombre y el espacio de nombres del trabajo. Debes especificar el contenedor de carga de trabajo (por ejemplo, -c workload) porque el sidecar de ML Diagnostics controla su propio registro.

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