ML Diagnostics SDK 使用入门

ML Diagnostics Python SDK 可与机器学习工作负载集成，以在 Google Cloud上收集和管理工作负载指标、配置和配置文件。本指南将向您展示如何创建机器学习运行、收集和管理工作负载指标和配置、部署受管理的 XProf 资源，以及启用程序化和按需配置文件捕获。

如需详细了解如何使用 ML Diagnostics SDK，请参阅 google-cloud-mldiagnostics 代码库。

安装 ML Diagnostics SDK

安装 google-cloud-mldiagnostics 库：

pip install google-cloud-mldiagnostics

在机器学习工作负载代码中导入以下软件包：

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

启用 Cloud Logging

该 SDK 使用标准 Python logging 模块来输出指标和配置信息。如需将这些日志路由到 Cloud Logging，请安装并配置 google-cloud-logging 库。这样一来，您就可以在 Google Cloud 控制台中查看 SDK 日志、记录的指标以及您自己的应用日志。

安装google-cloud-logging库：

pip install google-cloud-logging

通过将 Cloud Logging 处理程序附加到 Python 根日志记录器，在脚本中配置日志记录。将以下行添加到 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.")

启用详细日志记录

默认情况下，日志记录级别设置为 INFO。如需从 SDK 接收更详细的日志（例如机器学习运行详情），请在调用 setup_logging() 后将日志记录级别设置为 DEBUG：

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.")

启用 DEBUG 后，您会在 Cloud Logging 中收到额外的 SDK 诊断信息。例如：

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

创建机器学习运行作业

如需使用 ML 诊断平台，您需要先创建机器学习运行。这需要使用 SDK 对机器学习工作负载进行插桩，以执行日志记录、收集指标和启用配置文件跟踪。

以下是一个基本示例，用于初始化 Cloud Logging、创建机器学习运行 (MLRun)、记录指标和捕获配置文件：

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)

此代码示例使用以下变量：

变量	要求	说明
name	必需	特定运行的标识符。SDK 会自动创建 machine-learning-run-id，以确保运行名称是唯一的。
run_group	可选	一种标识符，可帮助对属于同一实验的多次运行进行分组。例如，与 TPU 切片大小扫描相关联的所有运行都可以属于同一组。
project	可选	如果未指定，则从 Google Cloud CLI 中提取项目。
region	必需	除 us-east5 之外，所有 Cluster Director 位置均受支持。可以通过每个命令的实参或使用命令 gcloud config set compute/region 来设置此标志。
configs	可选	包含运行配置参数的键值对。如果未定义配置，则会显示默认的软件和系统配置，但不会显示 ML 工作负载配置。
gcs_path	必需（有条件）	保存所有配置文件的 Google Cloud 存储位置。 例如：gs://my-bucket 或 gs://my-bucket/folder1。 仅当 SDK 用于捕获个人资料时才需要。
on-demand-xprof	可选	在端口 9999 上启动 xprofz daemon 以启用按需分析。您可以在同一代码中同时启用按需分析和程序化分析，前提是它们不会同时发生。

以下配置由 SDK 自动收集，无需在 machinelearning_run 中指定：

• 软件配置：框架、框架版本、XLA 标志。
• 系统配置：设备类型、切片数量、切片大小、主机数量。

项目和区域信息存储为机器学习运行元数据。用于机器学习运行的区域不必与用于工作负载运行的区域一致。

写入配置

许多工作负载包含的配置过多，无法直接在 machinelearning_run 定义中定义。在这些情况下，您可以使用 JSON 或 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",
)

收集指标

您可以使用 SDK 收集模型指标、模型性能指标和系统指标。您可以创建这些指标的可视化图表，以平均值和时序图表的形式呈现。

SDK 提供了两个用于记录指标的函数：metrics.record() 用于捕获单个数据点，metrics.record_metrics() 用于在单个批次中记录多个指标。这两个函数都会将指标写入 Cloud Logging，以便进行可视化和分析。

记录单个指标：

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

如需记录多个指标，请执行以下操作：

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)

SDK 会自动从 libTPU、psutil 和 JAX 库收集以下系统指标：

• TPU TensorCore 利用率
• TPU 占空比
• HBM 利用率
• 主机 CPU 利用率
• 主机内存利用率

您无需手动指定这些指标。这些系统指标以时间作为默认的 x 轴。

如果分配了以下预定义指标键，它们将自动显示在Google Cloud 控制台中。这些指标不会自动计算，而是您可以为其分配值的预定义键。

• 模型质量指标键：LEARNING_RATE、LOSS、GRADIENT_NORM、TOTAL_WEIGHTS。
• 模型效果指标键：STEP_TIME、THROUGHPUT、LATENCY、MFU、TFLOPS。

预定义指标以及其他用户定义的指标可以记录为 x 轴为 time 或同时为 time 和 step。您可以记录工作负载中的任何自定义指标。

以下示例捕获了工作负载的单个指标，您可以在特定机器学习运行的模型指标标签页中查看该指标：

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

如需在一次调用中记录多个指标，请使用 record_metrics 方法。例如：

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)

捕获配置文件

您可以通过程序化捕获或按需捕获（手动捕获）来捕获机器学习工作负载的 XProf 性能分析文件。程序化捕获是指将分析命令直接嵌入到机器学习代码中，并明确说明何时开始和停止记录数据。按需捕获是实时进行的，即在工作负载已在积极运行时触发性能分析器。

用于捕获配置信息的 SDK 命令与框架无关，因为所有框架级配置分析命令都会自动集成到 ML Diagnostics 配置分析命令中。这意味着您的分析代码不依赖于您使用的框架。

程序化配置文件捕获

程序化捕获要求您对模型代码进行注释，并指定要在哪个位置捕获性能分析文件。通常，您会捕获几个训练步的性能分析文件，或者对模型中的特定代码块进行性能分析。

您可以通过以下方式使用机器学习诊断 SDK 以程序化方式捕获配置文件：

• 基于 API 的收集：使用 start() 和 stop() 方法控制分析。
• 基于装饰器的收集：使用 @xprof(run) 为函数添加注解，以进行自动分析。
• 上下文管理器：与 xprof() 搭配使用，可实现基于范围的分析，自动处理 start() 和 stop() 操作。

您可以在所有框架中使用相同的配置文件捕获代码。所有分析会话都捕获在机器学习运行中定义的 Cloud Storage 存储桶中。

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

多主机（进程）分析

在程序化分析期间，SDK 会在执行机器学习工作负载代码的每个主机（进程）上开始分析。如果未提供节点列表，则包含所有主机。

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

默认情况下，在多个主机上调用不带 session_id 参数的 prof.start() 方法会导致单独的跟踪记录会话，每个主机对应一个跟踪记录会话。如需在 XProf 中将来自不同主机的轨迹分组到单个统一的多主机会话中，请确保在所有参与的主机上使用prof.start()方法时，session_id实参都相同。例如：

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

如需为特定主机启用分析，请执行以下操作：

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

按需捕获性能剖析文件

如果您想临时捕获性能分析文件，或者尚未启用程序化性能分析文件捕获，可以使用按需性能分析文件捕获。如果在运行期间模型指标出现问题，并且您希望在这些时刻捕获性能分析文件以诊断问题，那么按需捕获功能会很有用。

如需启用按需配置文件捕获，请配置支持按需运行：

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

您可以在所有框架中使用相同的配置文件捕获代码。所有分析会话都捕获在机器学习运行中定义的 Cloud Storage 存储桶中。

如需在 GKE 上进行按需分析，请将 GKE connection-operator 和 injection-webhook 部署到 GKE 集群中。这样可确保机器学习运行能够找到其运行所在的 GKE 节点，并且“按需捕获”下拉菜单可以自动填充这些节点。如需了解详情，请参阅配置 GKE 集群。

为 GKE 打包工作负载

您可以使用 Dockerfile 来打包使用 ML 诊断 SDK 的应用。安装 google-cloud-logging 软件包以实现 Cloud Logging 集成。例如：

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

部署工作负载

将 SDK 与工作负载集成后，将工作负载打包到映像中，并使用指定的映像创建 YAML 文件。使用 managed-mldiagnostics-gke=true 在 YAML 文件中为工作负载添加标签。

对于 GKE：

kubectl apply -f YAML_FILE_NAME

对于 Compute Engine，请使用 SSH 连接到虚拟机，然后运行工作负载的 Python 代码：

source venv/bin/activate
python3.11 WORKLOAD_FILE_NAME

部署工作负载后，通过搜索工作负载命名空间来查找作业名称：

kubectl get job -n YOUR_NAMESPACE

您可以通过传递作业名称和命名空间，在 kubectl 日志中找到运行名称和链接。您必须指定工作负载容器（例如：-c workload），因为 ML 诊断边车会处理自己的日志记录。

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