Comienza a usar el SDK de ML Diagnostics

El SDK de ML Diagnostics para Python se puede integrar con cargas de trabajo de AA para supervisar cargas de trabajo, recopilar y administrar métricas, parámetros de configuración y perfiles de modelos en Google Cloud. Usa el SDK para la creación de perfiles programática y para ver las métricas del modelo. Las métricas del modelo no están disponibles si usas la creación de perfiles bajo demanda sin el SDK.

En esta guía, se muestra cómo crear ejecuciones de aprendizaje automático, supervisar cargas de trabajo, recopilar y administrar métricas y parámetros de configuración de cargas de trabajo, implementar recursos de XProf administrados y habilitar la captura de perfiles programática y bajo demanda.

Para tener supervisión de cargas de trabajo, recopilación de métricas y parámetros de configuración, y habilitar la creación de perfiles, usa una versión del SDK que sea, al menos, tan reciente como la 1.0.6.

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

Instala el SDK de ML Diagnostics

Instala la google-cloud-mldiagnostics biblioteca:

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 métricas y la información de configuración. 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 tus propios registros de aplicaciones en la Google Cloud consola de.

Instala la biblioteca 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 de 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', ...}

Crea 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 seguimiento 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 para 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 barrido de tamaño de 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. Este parámetro se puede establecer con un argumento para cada comando o con el comando: gcloud config set compute/region.
configs	Opcional	Son pares clave-valor que contienen parámetros de configuración para la ejecución. Si no se definen los parámetros de configuración, aparecerán los parámetros de configuración predeterminados del software y del sistema, pero no los de la carga de trabajo de AA.
gcs_path	Condicionalmente obligatorio	Es la ubicación de Google Cloud Storage donde 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 bajo demanda. Puedes habilitar la creación de perfiles bajo demanda y la creación de perfiles programática en el mismo código, siempre que no ocurran al mismo tiempo.

El SDK recopila automáticamente los siguientes parámetros de configuración y no es necesario especificarlos en machinelearning_run:

• Parámetros de configuración del software: Framework, versión del framework, marcas de XLA
• Parámetros de configuración del sistema: Tipo de dispositivo, cantidad de porciones, tamaño de porción, cantidad de hosts.

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

Escribe parámetros de configuración

Muchas cargas de trabajo contienen demasiados parámetros de configuración para definirlos directamente en la definición machinelearning_run. En estos casos, puedes escribir parámetros de configuración en 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

Puedes recopilar métricas del modelo, métricas de rendimiento del modelo y métricas del sistema con el SDK. 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 libTPU, psutil y JAX:

• Uso de TensorCore 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 time como eje x predeterminado.

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

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

Las métricas predefinidas, así como otras métricas definidas por el usuario, se pueden registrar con el eje x como time o como 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 de aprendizaje automático específica:

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)

Métricas del sistema

Puedes recopilar automáticamente métricas del sistema en el intervalo predeterminado de 1 minuto o en un intervalo más detallado.

Intervalo predeterminado

Para recopilar métricas del sistema automáticamente con un intervalo de 1 minuto desde Cloud Monitoring, establece log_system_metrics=false. Consulta Supervisa cargas de trabajo para obtener detalles sobre las métricas del sistema que se recopilan automáticamente con este valor predeterminado.

A continuación, se muestra un ejemplo de cómo recibir métricas del sistema predeterminadas de Cloud Monitoring:

machinelearning_run(
      name="my-run-123",
      run_group="my-group",
      project="gcp-project",
      region="us-central1",
      gcs_path="gs://my_bucket_path",
      on_demand_xprof=True,
      log_system_metrics=False
  )

Intervalo detallado

Para recopilar métricas del sistema automáticamente con un intervalo detallado (10 segundos), establece log_system_metrics=true. Con esta opción, el SDK recopila automáticamente las siguientes métricas del sistema de las bibliotecas libTPU, psutil y JAX:

• Uso de TensorCore de TPU
• Ciclo de trabajo de TPU
• Uso de HBM
• Uso de CPU del host
• Uso de memoria del host

A continuación, se muestra un ejemplo de cómo recopilar métricas del sistema detalladas:

machinelearning_run(
      name="my-run-123",
      run_group="my-group",
      project="gcp-project",
      region="us-central1",
      gcs_path="gs://my_bucket_path",
      on_demand_xprof=True,
      log_system_metrics=True
  )

Estas métricas del sistema tienen time como eje x predeterminado.

Captura 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 declarar de forma explícita cuándo iniciar y detener el registro de datos. La captura bajo demanda se produce en tiempo real, en la que activas el analizador 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 creación de perfiles a nivel del framework se integran automáticamente en los comandos de creación de perfiles de ML Diagnostics. Esto significa que tu código de creación de perfiles no depende del framework que uses.

Captura de perfiles 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 programática con el SDK de ML Diagnostics de las siguientes maneras:

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

Puedes usar el mismo código de captura de perfiles en todos los frameworks. Todas las sesiones de perfiles 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 creación de perfiles programática, el SDK comienza a crear 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 seguimiento independientes, una para cada host. Para agrupar seguimientos 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 creació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 bajo demanda

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

Para habilitar la captura de perfiles bajo demanda, configura la ejecución con compatibilidad bajo demanda:

# 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 perfiles en todos los frameworks. Todas las sesiones de perfiles se capturan en el bucket de Cloud Storage definido en la ejecución de aprendizaje automático.

Para la creación de perfiles bajo demanda 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 bajo demanda 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 con tu carga de trabajo, empaqueta la carga de trabajo en una imagen y crea tu archivo YAML con la imagen especificada. Para la creación de perfiles programática, etiqueta la carga de trabajo en el archivo YAML con managed-mldiagnostics-gke=true. La creación de perfiles bajo demanda no requiere esta etiqueta.

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 de la ejecución y el vínculo en tus registros de kubectl pasando el nombre del trabajo y el espacio de nombres. Debes especificar el contenedor de la 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