Cloud 上の Pathways のトラブルシューティング

始める前に

このドキュメントでは、Pathways ワークロードのトラブルシューティング方法について説明します。

ログの表示

Cloud Logging のログ エクスプローラで、プロジェクト、リージョン、クラスタ、ワークロードに合わせて調整した次のクエリを使用します。

resource.type="k8s_container"
resource.labels.project_id="PROJECT"
resource.labels.location="LOCATION"
resource.labels.cluster_name="CLUSTER"
resource.labels.namespace_name="default"
resource.labels.pod_name:"WORKLOAD_NAME"

次のように置き換えます。

  • PROJECT : 実際の Google Cloud プロジェクト ID
  • LOCATION: GKE クラスタを作成したリージョンまたはゾーン
  • CLUSTER : GKE クラスタの名前
  • WORKLOAD_NAME : XPK を使用する場合はワークロードの名前、kubectl を使用する場合は JobSet の名前

このクエリは、pathways-rmpathways-proxypathways-worker などの名前を持つ複数の Pathways Kubernetes コンテナと一致します。resource.labels.container_name:"<container_name>" などのコンテナ名にフィルタを追加すると、問題の原因となっているコンテナを絞り込むことができます。

モニタリング

稼働状況のモニタリング

コンテナログのエントリを確認することで、さまざまな Pathways コンポーネントのヘルス ステータスをモニタリングできます。次に例を示します。

次のログが書き込まれると、pathways-proxy は新しい接続リクエストを処理する準備が整います。

kubectl logs ${HEAD_POD_NAME} --container pathways-proxy
...
I1101 04:51:41.967764       1 proxy_server.cc:125] IFRT proxy server started with status OK

次のログが書き込まれると、pathways-rm は新しい接続リクエストを処理する準備が整います。

kubectl logs $HEAD_POD_NAME --container pathways-rm
...
I1101 04:50:41.967764       1 server_lib.cc:1473] Pathways Server serving on [::]:29001

Pathways Resource Manager に登録されているすべての TPU の準備が完了したことを確認するには、pathways-rm コンテナログで ***<num slices>/<num slices> Pathways Slices Now Ready *** を探します。

kubectl logs $HEAD_POD_NAME --container pathways-rm
...
I1101 04:52:41.967764       1 multi_job_allocator.cc:1063] *** 2/2 Pathways Slices Now Ready ***

すべてのスライスが準備できていなくても、プロキシ サーバーが準備できていれば、Pathways クライアントは IFRT プロキシ サーバーに接続できます。ジョブは、スライスが準備完了になるまで仮想スライスで動作します。仮想スライスを使用すると、TPU が使用できない場合でもコードを実行できます。

次のログが書き込まれると、pathways-worker は新しい接続リクエストを処理する準備が整います。

kubectl logs $WORKER_POD_NAME --container pathways-worker
...
I1101 04:50:41.967764       1 server_lib.cc:1473] Pathways Server serving on [::]:29001

指標の収集

Pathways は、デバッグ用に低レベルのシステム指標を Cloud Monitoring に書き込むことができます。指標には次のものがあります。

  • DCN 転送レイテンシ
  • グループ レイテンシ
  • ホストからデバイスへの転送レイテンシ
  • デバイスからホストへの転送レイテンシ

指標は、GKE クラスタが実行されている {gcp_name} プロジェクトの Cloud Monitoring ダッシュボードで確認できます。

Metrics Explorer でこれらの指標をモニタリングするには:

  1. Metrics Explorer に移動します。
  2. [指標を選択] フィールドを使用して指標名でフィルタする
  3. [フィルタを追加] を選択してワークロード名と期間でフィルタし、pod_name でフィルタします。
  4. モニタリングする指標とワークロードに基づいて、適切な集計タイプを選択します。

DCN 転送レイテンシ

これは、マルチスライス トラフィックのネットワーク転送レイテンシの累積分布を測定する Megascale XLA(MXLA)指標です。レイテンシの測定は、DCN 経由で転送されるデータのリクエストが発行されたときに開始され、データの転送が完了したことを示す確認応答が受信されたときに終了します。この指標をモニタリングするには、dcn_transfer_latencies 指標名でフィルタします。

グループ レイテンシ

これは、マルチスライス トラフィックのエンドツーエンドの全体的なレイテンシの累積分布を測定する MXLA 指標です。レイテンシの測定は、グループのリクエストが発行されると開始し、データ転送が完了したことを示す確認応答を受信すると終了します。この指標をモニタリングするには、collective_e2e_latency 指標名でフィルタします。

ホストからデバイスへの転送レイテンシ

これは、マルチスライス トラフィックのホストからデバイスへの転送レイテンシの累積分布を測定する MXLA 指標です。レイテンシの測定は、DCN 経由で転送されるデータのリクエストが発行されたときに開始され、データ転送が完了したことを示す確認応答が受信されたときに終了します。この指標をモニタリングするには、host_to_device_transfer_latencies 指標名でフィルタします。

デバイスからホストへの転送レイテンシ

これは、マルチスライス トラフィックのデバイスからホストへの転送レイテンシの累積分布を測定する MXLA 指標です。レイテンシの測定は、DCN 経由で転送されるデータのリクエストが発行されたときに開始され、データの転送が完了したことを示す確認応答が受信されたときに終了します。この指標をモニタリングするには、device_to_host_transfer_latencies 指標名でフィルタします。

一般的なエラーのデバッグ

アクセラレータ構成のハッシュ化に関する警告を表示できない

次のメッセージは警告であり、JAX コードのパフォーマンスには影響しません。

INFO:jax._src.cache_key:get (_hash_accelerator_config): unable to hash accelerator config, falling back to hashing devices + platform: UNIMPLEMENTED: GetTopologyForDevices is not supported for the IFRT proxy client. (type <class 'jaxlib.xla_extension.XlaRuntimeError'>)

リソースで権限 logging.logEntries.create が拒否されました

次のエラーが表示された場合は、Vertex AI Workbench で使用される Compute Engine サービス アカウントに、 Google Cloud ロギング システムにログエントリを書き込む権限があることを確認してください。

INFO:absl:Created 'ArrayHandler' with primary_host=0, replica-id=0
WARNING:absl:pathwaysutils: Detected Pathways-on-Cloud backend. Applying changes.
Failed to submit 1 logs.
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/google/api_core/grpc_helpers.py", line 65, in error_remapped_callable
    return callable_(**args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/grpc/_channel.py", line 1181, in __call__
    return _end_unary_response_blocking(state, call, False, None)
  File "/opt/conda/lib/python3.10/site-packages/google/api_core/grpc_helpers.py", line 1006, in _end_unary_response_blocking
    raise _InactiveRpcError(state) # pytype: disable-not-instantiable
grpc._channel._InactiveRpcError: <_InactiveRpcError of RPC that terminated with:
  status = StatusCode.PERMISSION_DENIED
  details = "Permission 'logging.logEntries.create' denied on resource (or it may not exist)."
  debug_error_string = "UNKNOWN:Error received from peer ipv4:216.239.34.174:443 {created_time:"2024-10-03T20:30:44.820425276+00.00", grpc_status:7, grpc_message:"Permission \'logging.logEntries.create\' denied on resource (or it may not exist)."}"
>

この問題を解決するには、Compute Engine サービス アカウントに ログ書き込みロールを追加します。

IFRT プロキシ サーバーに接続できない

このエラーが表示された場合は、IFRT プロキシ クライアントが IFRT プロキシ サーバーに接続できません。

  • VPC ネットワークが正しく構成されていることを確認する
  • 接続を許可するようにファイアウォールが構成されていることを確認する
  • Pathways クラスタがプロビジョニングされていることを確認する
  • 起動時のエラー メッセージを確認する

IFRT コマンドを Pathways クラスタに送信する最初の JAX コマンドが実行されようとすると、約 1 分間応答が停止し、次のような RuntimeError が表示されます。

RuntimeError: Unable to initialize backend 'proxy': UNAVAILABLE: Unable to establish connection to ifrt_proxy server, please check provided address example-workload-proxy-0-0.example-workload.default.svc.example-cluster-domain.:38676'; detailed error: DNS resolution failed (set JAX_PLATFORMS='' to automatically choose an available backend)

Pathways クラスタへの既存の接続がある

Pathways クラスタは、一度に 1 つのクライアントとのセッションのみを維持できます。2 つの別々のノートブックが同じ Pathways クラスタに接続しようとすると、1 つは接続できますが、もう 1 つは次のエラーが表示されます。

INFO:absl:Created ArrayHandler with primary_host=0, replica_id=0
WARNING:absl:pathwaysutils: Detected Pathways-on-Cloud backend. Applying changes.
E0927 21:19:52.919607   37624 rpc_helper.cc:56] Connection to IFRT proxy server was terminated: CANCELLED: Cancelled
E0927 21:20:03.467547   37719 rpc_helper.cc:56] Connection to IFRT proxy server was terminated: CANCELLED: Cancelled
E0927 21:20:14.011645   37807 rpc_helper.cc:56] Connection to IFRT proxy server was terminated: CANCELLED: Cancelled
E0927 21:20:24.557955   37924 rpc_helper.cc:56] Connection to IFRT proxy server was terminated: CANCELLED: Cancelled
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:887, in backends()
    885   continue
--> 887 backend = _init_backend(platform)
    888 _backends[platform] = backend

File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:973, in _init_backend(platform)
    972 logger.debug("Initializing backend '%s'", platform)
--> 973 backend = registration.factory()
    974 # TODO: consider raising more descriptive errors directly from backend
    975 # factories instead of returning None.

File /opt/conda/lib/python3.10/site-packages/pathwaysutils/proxy_backend.py:24, in register_backend_factory.<locals>.<lambda>()
     21 def register_backend_factory():
     22   xla_bridge.register_backend_factory(
     23       "proxy",
---> 24       lambda: ifrt_proxy.get_client(
     25           jax.config.read("jax_backend_target"),
     26           ifrt_proxy.ClientConnectionOptions(),
     27       ),
     28       priority=-1,
     29   )

XlaRuntimeError: UNAVAILABLE: Unable to establish connection to ifrt_proxy server, please check provided address example-workload-proxy-0-0.example-workload.default.svc.example-cluster-domain.:38676'; detailed error: Connection to IFRT proxy server was terminated: CANCELLED: Cancelled

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
Cell In[2], line 4
      1 import pathwaysutils
      3 import jax
----> 4 print(jax.devices())

File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:1085, in devices(backend)
   1060 def devices(
   1061     backend: str | xla_client.Client | None = None
   1062 ) -> list[xla_client.Device]:
   1063   """Returns a list of all devices for a given backend.
   1064
   1065   .. currentmodule:: jaxlib.xla_extension
   (...)
   1083     List of Device subclasses.
   1084   """
-> 1085   return get_backend(backend).devices()

File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:1019, in get_backend(platform)
   1015 @lru_cache(maxsize=None)  # don't use util.memoize because there is no X64 dependence.
   1016 def get_backend(
   1017     platform: None | str | xla_client.Client = None
   1018 ) -> xla_client.Client:
-> 1019   return _get_backend_uncached(platform)

File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:998, in _get_backend_uncached(platform)
    994   return platform
    996 platform = (platform or _XLA_BACKEND.value or _PLATFORM_NAME.value or None)
--> 998 bs = backends()
    999 if platform is not None:
   1000   platform = canonicalize_platform(platform)

File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:903, in backends()
    901       else:
    902         err_msg += " (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)"
--> 903       raise RuntimeError(err_msg)
    905 assert _default_backend is not None
    906 if not config.jax_platforms.value:

RuntimeError: Unable to initialize backend 'proxy': UNAVAILABLE: Unable to establish connection to ifrt_proxy server, please check provided address 'example-workload-proxy-0-0.example-workload.default.svc.example-cluster-domain.:38676'; detailed error: Connection to IFRT proxy server was terminated: CANCELLED: Cancelled (set JAX_PLATFORMS='' to automatically choose an available backend)

元のクライアントが切断されると、2 番目のクライアントが接続できるようになります。予期しない切断が発生した場合は、他のクライアントが接続できるように Pathways クラスタを再起動する必要があります。

LocalProxy.init() で予期しないキーワード引数 'unbound_message' が渡されました

pathwaysutils のインポート後にこのエラーが表示された場合は、環境に古いバージョンの Flask または Werkzeug がインストールされているかどうかを確認します。

 pip3 list --outdated (replacing pip with pip3 as needed)

Flask または Werkzeug がリストされている場合は、アップグレードを検討してください。ただし、アップグレードするとプロジェクト内の他のパッケージや依存関係が破損する可能性があることに注意してください。

 pip install flask Werkzeug --upgrade ()

プロキシ サーバーからの内部エラー

"Internal error from proxy server during Array::IsDeleted(): UNAVAILABLE:Connection to IFRT proxy server was terminated: FAILED_PRECONDITION:
GrpcClientSession: writes no longer allowed."

このエラーは、IFRT プロキシ サーバーがクライアントから切断されたことを示します。この問題は、クライアントを再起動することで解決できます。ノートブックの場合は、ノートブック カーネルを再起動してノートブックを再実行できます。

SIGTERMS と HBM OOM

RESOURCE_EXHAUSTED エラーに関連付けられたログで SIGTERM が見つかった場合は、HBM OOM を示している可能性があります。この場合は、JAX コードで使用される HBM メモリの量を減らすことができます。

INVALID_ARGUMENT

"INVALID_ARGUMENT : Permanent error, with a last message of Lifecycle
matches_prefix cannot specify more than 50 prefixes per config.; Error while
initializing persistent cache storage Cloud Storage"

このエラーは、--pathways_gcs_location フラグに渡された Cloud Storage バケットがライフサイクル ポリシーの上限に達した場合に発生します。この場合は、使用されなくなった Cloud Storage ライフサイクル ポリシーをクリーンアップします。

永続的なエラー

Permanent error, with a last message of The specified bucket does not exist.; Error while initializing persistent cache storage gcs

このエラーは、無効な Cloud Storage ロケーションを Pathways コンテナに指定すると、Resource Manager または Pathways ワーカーに表示されます。

デーモンからのエラー レスポンス

Error response from daemon: dockerfile parse error line 16: Unknown flag: exclude

これは古いバージョンの Docker が原因です。Docker のバージョンをアップグレードしてください。

IFRT プロキシ クライアントとサーバーが合意に失敗しました

IFRT Proxy client and server failed to agree on the protocol version; supported versions: client = [1, 1], server = [3, 14]

これは、古いバージョンの MaxText が使用されていることを示しています。最新の MaxText イメージを再構築してください。

次のステップ