開始使用 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

在 ML 工作負載程式碼中匯入下列套件：

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	必填	特定執行的 ID。SDK 會自動建立 machine-learning-run-id，確保執行名稱不重複。
run_group	選用	這個 ID 可協助您將屬於同一項實驗的多個執行作業分組。舉例來說，與 TPU 節點大小掃描相關聯的所有執行作業，都可能屬於同一個群組。
project	選用	如未指定，系統會從 Google Cloud CLI 擷取專案。
region	必填	系統支援所有 Cluster Director 位置，但 us-east5 除外。這個旗標可透過每個指令的引數設定，或使用 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。

預先定義的指標和其他使用者定義的指標，可以記錄為以 time 為 x 軸，或同時以 time 和 step 為 x 軸。您可以記錄工作負載中的任何自訂指標。

以下範例會擷取工作負載的單一指標，您可以在特定機器學習執行的「模型指標」分頁中查看：

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 剖析指令。也就是說，剖析程式碼不會受到所用架構影響。

以程式輔助方式擷取設定檔

如要以程式輔助方式擷取設定檔，您必須註解模型程式碼，並指定要擷取設定檔的位置。通常，您會擷取幾個訓練步驟的設定檔，或剖析模型中的特定程式碼區塊。

您可以使用 ML Diagnostics SDK，透過下列方式以程式輔助方式擷取設定檔：

• 以 API 為基礎的收集作業：使用 start() 和 stop() 方法控管剖析作業。
• 以裝飾器為基礎的收集作業：使用 @xprof(run) 為函式加上註解，即可自動進行剖析。
• 內容管理工具：搭配 xprof() 使用，進行以範圍為準的剖析，自動處理 start() 和 stop() 作業。

您可以在所有架構中使用相同的設定檔擷取程式碼。所有設定檔工作階段都會擷取至機器學習執行作業中定義的 Cloud Storage bucket。

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

根據預設，在多個主機上呼叫 prof.start() 方法時，如果沒有 session_id 引數，系統會為每個主機分別建立追蹤工作階段。如要將不同主機的追蹤記錄分組到 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 bucket。

如要在 GKE 中進行隨選剖析，請將 GKE connection-operator 和 injection-webhook 部署至 GKE 叢集。這可確保機器學習執行作業能找到執行的 GKE 節點，且「隨選擷取」下拉式選單可自動填入這些節點。詳情請參閱「設定 GKE 叢集」。

封裝 GKE 工作負載

您可以使用 Dockerfile 封裝使用 ML Diagnostics 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 檔案。在 YAML 檔案中，使用 managed-mldiagnostics-gke=true 為工作負載加上標籤。

GKE：

kubectl apply -f YAML_FILE_NAME

如果是 Compute Engine，請使用 SSH 連線至 VM，然後執行工作負載的 Python 程式碼：

source venv/bin/activate
python3.11 WORKLOAD_FILE_NAME

部署工作負載後，請搜尋工作負載命名空間，找出工作名稱：

kubectl get job -n YOUR_NAMESPACE

傳遞工作名稱和命名空間後，您可以在 kubectl 記錄檔中找到執行名稱和連結。您必須指定工作負載容器 (例如：-c workload)，因為 ML Diagnostics Sidecar 會處理自己的記錄。

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