Começar a usar o SDK ML Diagnostics

O SDK do Python do ML Diagnostics pode ser integrado a cargas de trabalho de ML para coletar e gerenciar métricas, configurações e perfis de carga de trabalho no Google Cloud. Este guia mostra como criar execuções de aprendizado de máquina, coletar e gerenciar métricas e configurações de carga de trabalho, implantar recursos gerenciados do XProf e ativar a captura de perfil programática e sob demanda.

Para mais informações sobre o uso do SDK ML Diagnostics, consulte o repositório google-cloud-mldiagnostics.

Instalar o SDK do ML Diagnostics

Instale a biblioteca google-cloud-mldiagnostics:

pip install google-cloud-mldiagnostics

Importe os seguintes pacotes no código da sua carga de trabalho de ML:

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

Ativar Cloud Logging

O SDK usa o módulo padrão do Python logging para gerar métricas e informações de configuração. Para rotear esses registros para o Cloud Logging, instale e configure a biblioteca google-cloud-logging. Assim, é possível conferir registros do SDK, métricas registradas e registros do seu próprio aplicativo no console do Google Cloud .

Instale a biblioteca google-cloud-logging:

pip install google-cloud-logging

Configure a geração de registros no script anexando o gerenciador do Cloud Logging ao logger raiz do Python. Adicione as seguintes linhas ao início do 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.")

Ativar o registro detalhado

Por padrão, o nível de geração de registros é definido como INFO. Para receber registros mais detalhados do SDK, como detalhes da execução de aprendizado de máquina, defina o nível de registro como DEBUG depois de chamar 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.")

Com o DEBUG ativado, você recebe outros diagnósticos do SDK no Cloud Logging. Exemplo:

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

Criar uma execução de machine learning

Para usar a plataforma ML Diagnostics, primeiro crie uma execução de machine learning. Isso envolve instrumentar sua carga de trabalho de ML com o SDK para realizar geração de registros, coletar métricas e ativar o rastreamento de perfis.

Confira um exemplo básico que inicializa o Cloud Logging, cria uma execução de aprendizado de máquina (MLRun), registra métricas e captura um 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)

O exemplo de código usa as seguintes variáveis:

Variável	Requisito	Descrição
name	Obrigatório	Um identificador da execução específica. O SDK cria automaticamente um machine-learning-run-id para garantir que os nomes das execuções sejam exclusivos.
run_group	Opcional	Um identificador que pode ajudar a agrupar várias execuções pertencentes ao mesmo experimento. Por exemplo, todas as execuções associadas a uma varredura de tamanho de fração de TPU podem pertencer ao mesmo grupo.
project	Opcional	Se não for especificado, o projeto será extraído da Google Cloud CLI.
region	Obrigatório	Todos os locais do Cluster Director são compatíveis, exceto us-east5. Essa flag pode ser definida por um argumento para cada comando ou com o comando: gcloud config set compute/region.
configs	Opcional	Pares de chave-valor que contêm parâmetros de configuração para a execução. Se as configurações não forem definidas, as configurações padrão de software e sistema vão aparecer, mas não as de carga de trabalho de ML.
gcs_path	Obrigatório sob certas condições	O local de armazenamento Google Cloud em que todos os perfis são salvos. Por exemplo, gs://my-bucket ou gs://my-bucket/folder1. Necessário apenas se o SDK for usado para captura de perfil.
on-demand-xprof	Opcional	Inicia xprofz daemon na porta 9999 para ativar a criação de perfil sob demanda. É possível ativar a criação de perfis on demand e programática no mesmo código, desde que não ocorram ao mesmo tempo.

As configurações a seguir são coletadas automaticamente pelo SDK e não precisam ser especificadas no machinelearning_run:

• Configurações de software: framework, versão do framework, flags XLA.
• Configurações do sistema: tipo de dispositivo, número de segmentos, tamanho do segmento, número de hosts.

As informações do projeto e da região são armazenadas como metadados de execução de machine learning. A região usada para a execução de machine learning não precisa corresponder à região usada para a execução da carga de trabalho.

Gravar configurações

Muitas cargas de trabalho têm configurações demais para serem definidas diretamente na definição de machinelearning_run. Nesses casos, é possível gravar configurações na sua execução usando 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",
)

Coletar métricas

É possível coletar métricas de modelo, de performance do modelo e do sistema com o SDK. É possível criar visualizações dessas métricas como valores médios e com gráficos de série temporal.

O SDK oferece duas funções para registrar métricas: metrics.record() para capturar pontos de dados individuais e metrics.record_metrics() para registrar várias métricas em um único lote. As duas funções gravam métricas no Cloud Logging, permitindo visualização e análise.

Para gravar uma única métrica:

# 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 várias métricas:

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)

As seguintes métricas do sistema são coletadas automaticamente pelo SDK das bibliotecas libTPU, psutil e JAX:

• Utilização do TensorCore da TPU
• Ciclo de trabalho da TPU
• Utilização de HBM
• Utilização da CPU do host
• Utilização de memória do host

Não é necessário especificar essas métricas manualmente. Essas métricas do sistema têm tempo como o eixo x padrão.

As seguintes chaves de métricas predefinidas vão aparecer automaticamente no console doGoogle Cloud se forem atribuídas. Essas métricas não são calculadas automaticamente. Elas são chaves predefinidas a que você pode atribuir valores.

• Chaves de métricas de qualidade do modelo: LEARNING_RATE, LOSS, GRADIENT_NORM, TOTAL_WEIGHTS.
• Chaves de métricas de performance do modelo: STEP_TIME, THROUGHPUT, LATENCY, MFU, TFLOPS.

As métricas predefinidas e outras definidas pelo usuário podem ser registradas com o eixo x como time ou time e step. Você pode gravar qualquer métrica personalizada na carga de trabalho.

O exemplo a seguir captura uma única métrica para a carga de trabalho, que pode ser visualizada na guia Métricas do modelo da execução específica de machine learning:

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

Para gravar várias métricas em uma chamada, use o método record_metrics. Exemplo:

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)

Capturar perfis

É possível capturar perfis XProf da sua carga de trabalho de ML com captura programática ou on demand (manual). A captura programática envolve incorporar comandos de criação de perfil diretamente no código de aprendizado de máquina e declarar explicitamente quando iniciar e interromper a gravação de dados. A captura on demand ocorre em tempo real. Você aciona o criador de perfis enquanto a carga de trabalho já está em execução.

Os comandos do SDK para capturar perfis são independentes de framework, já que todos os comandos de criação de perfil no nível do framework são integrados automaticamente aos comandos de criação de perfil do ML Diagnostics. Isso significa que seu código de criação de perfil não depende do framework usado.

Captura programática de perfil

A captura programática exige que você anote o código do modelo e especifique onde quer capturar perfis. Normalmente, você captura um perfil para algumas etapas de treinamento ou cria um perfil de um bloco específico de código no modelo.

É possível fazer a captura programática de perfis com o SDK ML Diagnostics das seguintes maneiras:

• Coleta baseada em API: controle o perfil com os métodos start() e stop().
• Coleta baseada em decorador: anote funções com @xprof(run) para criação automática de perfis.
• Gerenciador de contexto: use com xprof() para criação de perfil baseada em escopo que processa automaticamente operações start() e stop().

Você pode usar o mesmo código de captura de perfil em todas as estruturas. Todas as sessões de criação de perfil são capturadas no bucket do Cloud Storage definido na execução de 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()

Criação de perfil de vários hosts (processos)

Durante o perfil programático, o SDK inicia a criação de perfis em cada host (processo) em que o código de carga de trabalho de ML está sendo executado. Se a lista de nós não for fornecida, todos os hosts serão incluídos.

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

Por padrão, chamar o método prof.start() sem o argumento session_id em vários hosts resulta em sessões de rastreamento separadas, uma para cada host. Para agrupar rastreamentos de diferentes hosts em uma única sessão unificada de vários hosts no XProf, verifique se o método prof.start() é chamado com o mesmo argumento session_id em todos os hosts participantes. Exemplo:

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

Para ativar a criação de perfis de hosts específicos:

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

Captura de perfil sob demanda

Use a captura de perfil on demand quando quiser capturar perfis on demand ou quando a captura de perfil programática não estiver ativada. A captura sob demanda é útil quando há problemas com as métricas do modelo durante a execução e você quer capturar perfis nesses momentos para diagnosticar os problemas.

Para ativar a captura de perfil sob demanda, configure a execução com suporte sob 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
)

Você pode usar o mesmo código de captura de perfil em todas as estruturas. Todas as sessões de criação de perfil são capturadas no bucket do Cloud Storage definido na execução de machine learning.

Para criação de perfil sob demanda no GKE, implante o GKE connection-operator e o injection-webhook no cluster do GKE. Isso garante que a execução de machine learning possa localizar os nós do GKE em que está sendo executada, e o menu suspenso de captura sob demanda pode preencher automaticamente esses nós. Para mais informações, consulte Configurar cluster do GKE.

Empacotar a carga de trabalho para o GKE

Você pode usar um Dockerfile para empacotar um aplicativo que usa o SDK de diagnóstico de ML. Instale o pacote google-cloud-logging para integração com o Cloud Logging. Exemplo:

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

Implantar carga de trabalho

Depois de integrar o SDK à sua carga de trabalho, empacote-a em uma imagem e crie o arquivo YAML com a imagem especificada. Rotule a carga de trabalho no arquivo YAML com managed-mldiagnostics-gke=true.

Para o GKE:

kubectl apply -f YAML_FILE_NAME

Para o Compute Engine, conecte-se à VM usando SSH e execute o código Python para sua carga de trabalho:

source venv/bin/activate
python3.11 WORKLOAD_FILE_NAME

Depois de implantar a carga de trabalho, encontre o nome do job pesquisando o namespace da carga de trabalho:

kubectl get job -n YOUR_NAMESPACE

Para encontrar o nome e o link da execução nos registros kubectl, transmita o nome e o namespace do job. É necessário especificar o contêiner da carga de trabalho (por exemplo, -c workload) porque o sidecar do ML Diagnostics processa o próprio registro em registros.

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