ML Diagnostics SDK を使ってみる

ML Diagnostics Python SDK は、ML ワークロードと統合して、 Google Cloudでワークロード指標、構成、プロファイルを収集して管理できます。このガイドでは、ML 実行の作成、ワークロード指標と構成の収集と管理、マネージド 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 からより詳細なログ（ML 実行の詳細など）を受け取るには、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 の実行を作成

ML 診断プラットフォームを使用するには、まず ML 実行を作成する必要があります。これには、SDK を使用して ML ワークロードを計測し、ロギングの実行、指標の収集、プロファイル トレースの有効化を行います。

次の基本的な例では、Cloud Logging を初期化し、ML 実行（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 を除くすべてのクラスタ ディレクターのロケーションがサポートされています。このフラグは、各コマンドの引数で設定するか、コマンド gcloud config set compute/region で設定できます。
configs	省略可	実行の構成パラメータを含む Key-Value ペア。構成が定義されていない場合、デフォルトのソフトウェア構成とシステム構成は表示されますが、ML ワークロード構成は表示されません。
gcs_path	条件付きで必須	すべてのプロファイルが保存される Google Cloud ストレージの場所。たとえば、gs://my-bucket や gs://my-bucket/folder1 です。SDK がプロファイル キャプチャに使用される場合にのみ必要です。
on-demand-xprof	省略可	オンデマンド プロファイリングを有効にするため、ポート 9999 で xprofz daemon を開始します。オンデマンド プロファイリングとプログラマティック プロファイリングが同時に発生しない限り、同じコードで両方を有効にできます。

次の構成は SDK によって自動的に収集されるため、machinelearning_run 内で指定する必要はありません。

• ソフトウェア構成: フレームワーク、フレームワーク バージョン、XLA フラグ。
• システム構成: デバイスタイプ、スライスの数、スライスサイズ、ホストの数。

プロジェクトとリージョン情報は、ML 実行メタデータとして保存されます。ML 実行に使用されるリージョンは、ワークロード実行に使用されるリージョンと一致している必要はありません。

構成を書き込む

多くのワークロードには、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 には、指標を記録するための 2 つの関数（個々のデータポイントをキャプチャする metrics.record() と、複数の指標を 1 つのバッチで記録する 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 の両方として記録できます。ワークロードのカスタム指標を記録できます。

次の例では、ワークロードの単一の指標を取得します。この指標は、特定の ML 実行の [モデル指標] タブで確認できます。

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

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)

プロファイルのキャプチャ

ML ワークロードの XProf プロファイルは、プログラムによるキャプチャまたはオンデマンド キャプチャ（手動キャプチャ）でキャプチャできます。プログラムによるキャプチャでは、プロファイリング コマンドを ML コードに直接埋め込み、データの記録を開始および停止するタイミングを明示的に指定します。オンデマンド キャプチャはリアルタイムで行われます。ワークロードがすでにアクティブに実行されているときに、プロファイラをトリガーします。

すべてのフレームワーク レベルのプロファイリング コマンドは ML Diagnostics プロファイリング コマンドに自動的に統合されるため、プロファイルをキャプチャする 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 は ML ワークロード コードが実行されている各ホスト（プロセス）でプロファイリングを開始します。ノードのリストが指定されていない場合、すべてのホストが対象となります。

# 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 Diagnostics SDK を使用するアプリケーションをパッケージ化できます。Cloud Logging 統合用の google-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

ワークロードをデプロイしたら、ワークロードの Namespace を検索してジョブ名を見つけます。

kubectl get job -n YOUR_NAMESPACE

ジョブ名と Namespace を渡すことで、kubectl ログで実行名とリンクを確認できます。ML 診断サイドカーは独自のロギングを処理するため、ワークロード コンテナ（-c workload など）を指定する必要があります。

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