GKE でのバッチ推論のベスト プラクティス

このドキュメントでは、Google Kubernetes Engine（GKE）でバッチ推論ワークロードを実行するためのベスト プラクティスについて説明します。バッチ推論は、機械学習モデルを使用して大規模なデータセットに対する予測を生成するプロセスです。このプロセスでは、即時性や低レイテンシのレスポンスよりも、高スループットと費用対効果が優先されます。

このガイドでは、バッチ推論とリクエスト バッチ処理（または動的バッチ処理）を区別します。リクエスト バッチ処理は、vLLM や SGLang などのエンジンで、同時実行のリアルタイム リクエストをグループ化してアクセラレータの効率を最適化するサーバーサイドの手法です。リクエスト バッチ処理は、バッチ推論ワークロードに適用できます。

このガイドのベスト プラクティスは、次の 2 つの一般的なバッチ推論パターンを対象としています。

  • 準リアルタイム バッチ推論: データが生成されてからすぐにチャンク単位で処理します。通常、数秒から数分のレイテンシで、このアプローチは最新のデータの必要性と複数のアイテムを同時に処理する効率性のバランスを取ります。
  • オフライン バッチ推論: 大量の蓄積データをスケジュールされた間隔（毎晩や毎週など）で処理します。これらのジョブは、リソースの可用性を最大限に高めるために、ピーク時以外の時間にスケジュールされることが多いため、通常、レイテンシは数時間から数日になります。

これらの推奨事項は、GKE での推論に関するベスト プラクティスの概要で説明されている基盤の上に構築された最適化の特殊なレイヤです。バッチ ワークロード用に最適化する前に、モデルの選択、量子化、アクセラレータの選択に関するコアのベスト プラクティスに従っていることを確認してください。

バッチ推論処理のアーキテクチャ パターンを選択する

バッチ推論ワークロードをデプロイするうえで最も重要な決定は、適切なアーキテクチャ パターンを選択することです。これは、レイテンシ、スループット、費用のトレードオフに影響するためです。効率を維持するには、オフピーク時に推論スループットが受信クエリのレートを超えていることを確認して、キューが無限に増大しないようにします。

バースト的なワークロードに準リアルタイムのバッチ推論を使用する

ニア リアルタイム バッチ処理は、次のような頻繁な増分更新が必要なユースケースに適しています。

  • 最近のインタラクションに基づいて、ユーザーのレコメンデーション プロファイルを数分ごとに更新します。
  • リアルタイム モニタリングのために、ソーシャル メディアでの言及を 1 分間隔で処理します。
  • 高頻度の金融データ ストリームから市場を動かすシグナルを検出する。
  • 受信した顧客からのフィードバックやニュース フィードの感情分析を行う。

ワークロードで数秒から数分のレイテンシを許容できる場合は、このパターンを選択します。

ほぼリアルタイムのバッチ推論を実装する場合は、次の特性を考慮してください。

  • レイテンシ: 最初のトークンまでの時間は、数十秒から数分程度です。
  • データソース: 通常、Pub/Sub からのメッセージや、短期間に Cloud Storage から収集されたファイルなど、数メガバイトから数ギガバイトのデータセットを処理します。
  • コンピューティング パターン: インフラストラクチャは、頻繁に発生するワークロードのバーストを処理する継続的なサービスをサポートする必要があります。
  • 費用の最適化: このパターンは、低レイテンシのリアルタイム推論と高スループットのオフライン処理のバランスを取ります。

大規模なデータセットにオフライン バッチ推論を使用する

オフライン バッチ処理は、次のような数時間または数日間の遅延を許容できる大規模なエピソード ジョブに最適です。

  • 前日の金融取引に基づいて、夜間のリスク評価レポートを生成します。
  • カタログ全体の商品エンベディングを作成して、下流の検索システムとレコメンデーション システムを強化します。
  • モデルのトレーニングやアーカイブ分類のために、大規模な画像データセットにラベルを付けます。

大量のデータを処理し、数時間から数日程度のレイテンシを許容できる場合は、このパターンを選択します。

オフライン バッチ推論を実装する場合は、次の特性を考慮してください。

  • レイテンシ: ジョブは通常、オフピーク時にスケジュールされるため、ワークロードの開始レイテンシは通常、数分から数日です。
  • データソース: 通常、Cloud Storage または BigQuery テーブルに保存されているギガバイトからペタバイトまでの大規模なデータセットを処理します。
  • コンピューティング パターン: 初期化、データの処理、終了を行うエピソード型のバースト ジョブを使用します。
  • 費用の最適化: このパターンは、従量課金モデルで最適化できます。オフライン ジョブには柔軟な完了ウィンドウがあるため、Spot VM を使用してコストを削減することをおすすめします。

スループットと費用対効果を最適化する

バッチ推論ワークロードは、中断が発生する可能性のある費用対効果の高いインフラストラクチャに最適です。

Spot VM を使用してコンピューティング費用を削減する

バッチジョブには Spot VM の割引を使用します。バッチ推論ワークロードは通常、レイテンシと中断を許容できるため、Spot 容量の割引料金の対象となります。

バッチ推論コードで、潜在的なプリエンプション イベントを処理するためのチェックポイントが実装されていることを確認します。Spot VM がプリエンプトされた場合、新しいノードを作成し、ゼロから再起動するのではなく、最後に処理されたバッチからワークロードを再開できます。

ワークロードのバッチサイズとリクエストのバッチサイズを調整する

リソースの競合とジョブのタイムアウトを回避するには、アクセラレータの利用率が低下しないように、エンジンに送信されるアイテム数（ワークロード バッチ）が、サーバーが処理できる同時リクエスト数（リクエスト バッチ）以上になるようにします。

ワークロードのバッチサイズを調整する

ワークロード バッチサイズは、1 つの作業単位で推論エンジンに送信されるアイテムの合計数です。これは、データをシャーディングするか、複数のアイテムを 1 つのリクエストにグループ化することで、クライアント送信ロジックまたは Kubernetes Job 構成で構成します。

最適なワークロード バッチサイズを決定するには、次の境界を使用します。

  • 最小バッチサイズを計算する: ワークロードのバッチサイズがリクエストのバッチサイズ以上であることを確認します。たとえば、256 個のアイテムを同時に処理できるサーバーに 1 個のアイテムを送信すると、使用率が大幅に低下します。最小サイズを確認するには、vLLM の max_num_seqs 引数など、推論サーバーの構成を確認します。複数のアイテムを 1 つのリクエストにグループ化するようにクライアント ロジックを構成するか、各ジョブがリクエスト バッチサイズを満たすか超える最小量のデータを受け取るようにデータをシャーディングできます。
  • 最大バッチサイズを計算する: ワークロード バッチサイズにより、Kubernetes Job で定義された activeDeadlineSeconds タイムアウトに達する前に Pod が完了することを確認します。1 つのリクエスト バッチの処理に必要な時間を見積もり、Pod が期限内に完了するようにワークロード サイズを設定します。たとえば、activeDeadlineSeconds が 3,600 秒で、起動オーバーヘッドが 600 秒の場合、最大実行時間で Pod が 3,000 秒以内に完了するようにします。

ワークロードのバッチサイズが小さすぎると、ジョブは Pod の起動オーバーヘッド（重みのダウンロード、プロビジョニング、アクセラレータの初期化）に時間を費やします。大きすぎると、activeDeadlineSeconds タイムアウトにより GKE によってジョブが終了し、ジョブが失敗して進行状況が失われる可能性があります。

リクエスト バッチサイズを調整する

リクエスト バッチサイズは、推論サーバーがアクセラレータで同時に処理する同時リクエストの数です。このパラメータを最適化するには、推論サーバーの構成でサーバー固有のフラグ（vLLM の --max-num-seqs フラグなど）を調整します。

目標は、メモリ不足（OOM）エラーをトリガーせずに GPU 使用率を最大化することです。リクエスト バッチサイズが調整されていない場合、システムはアクセラレータを十分に活用できないか、モデルサーバーがクラッシュします。vLLM の場合は、vLLM auto_tune スクリプトなどのツールを使用して、特定のハードウェアの max_num_seqs と max_num_batched_tokens の設定に最適な値を見つけることができます。詳細については、GKE での推論に関するベスト プラクティスの概要ガイドの推論サーバーの構成を最適化するをご覧ください。

ほぼリアルタイムのバッチ処理用に非同期コンポーネントを実装する

準リアルタイム バッチ処理では、メッセージ バッファを使用して取り込みレイヤを推論レイヤから切り離すことをおすすめします。

次のアーキテクチャ図は、準リアルタイム バッチ推論プラットフォームの例を示しています。このアーキテクチャは、推論サーバーをトラフィックの急増から保護し、ワーク バックログを管理し、アクセラレータの使用率を高く維持します。

この図は、Pub/Sub からサブスクライバー、推論ゲートウェイ、推論サーバーへのフローを示しています。結果は AlloyDB に保持され、失敗したメッセージはデッドレター トピックに送信されます。

GKE 上の準リアルタイム バッチ推論プラットフォーム。

このアーキテクチャは、次のコンポーネントで構成されています。

  • Pub/Sub トピック: 受信クライアント メッセージの永続バッファとして機能し、保持期間は 7 ～ 31 日です。
  • サブスクライバー: メッセージ バッチを読み取り、推論サーバーにリクエストを送信し、処理を確認するコンポーネント。
  • サブスクライバー HPA: num_undelivered_messages 指標（未確認メッセージの数）に基づいてサブスクライバー デプロイをスケーリングします。
  • ストレージ: データベース（AlloyDB など）またはオブジェクト ストレージ（Cloud Storage など）を使用して、推論結果を永続化します。
  • Inference Gateway: 推論ワークロードをサブスクライバーに公開します。
  • 推論サーバー: バッチ推論リクエストを処理します（vLLM など）。
  • サーバー HPA: vllm:num_requests_waiting などのエンジン固有の指標に基づいて推論エンジンをスケーリングします。
  • デッドレター トピック: 指数バックオフ再試行が設定された回数行われた後、処理に失敗したメッセージをキャプチャします。

詳細については、GitHub のリファレンス実装をご覧ください。

リクエストをバッファリングして集約する

リクエストのフローを管理する手順は次のとおりです。

  • Pub/Sub を永続バッファとして使用する: Pub/Sub を実装して、推論リクエストを永続的に保存します。この設定は、コンシューマーがリクエストを処理できる容量になるまでリクエストを保持する FIFO バッファとして機能し、バースト トラフィック中のサーバーの過負荷を防ぎます。
  • クライアントサイドのフロー制御で pull サブスクリプションを使用する: Pull サブスクリプション モデルを構成します。これにより、サブスクライバー アプリケーションは、メッセージを処理できる場合にのみメッセージを明示的にリクエストできるようになり、消費率を完全に制御できます。
  • メッセージを集約してサーバーのバッチサイズを埋める: 1 つの Pub/Sub メッセージを 1 つの推論リクエストとして送信しないようにします。代わりに、サブスクライバーは、推論サーバーの最適なバッチサイズ（vLLM の max_num_seqs 設定など）に沿って、複数のメッセージを 1 つのバッチ リクエストにバンドルする必要があります。このアプローチにより、アクセラレータが完全に飽和し、スループットが最大化されます。具体的には、すべてのモデルのフォワード パスが完全に飽和するように、サブスクライバーの max_messages プル設定を max_num_seqs の倍数に構成します。

サブスクライバーとサーバーを自動スケーリングする

バッチ推論を効果的に行うには、サブスクライバー（CPU バウンド）を推論サーバー（GPU または TPU バウンド）とは異なる方法でスケーリングする必要があります。

  • 作業バックログに基づいてサブスクライバーをスケーリングする: Pub/Sub の num_undelivered_messages 指標に基づいて、サブスクライバー デプロイの HorizontalPodAutoscaler（HPA）を構成します。詳細については、指標に基づいて Pod の自動スケーリングを最適化するをご覧ください。次の式を使用して、使用するレプリカを計算します。

    \[ desiredReplicas = \frac{num\_undelivered\_messages}{target\_latency\_seconds \times throughput\_per\_replica} \]

  • インフラストラクチャの割り当てを尊重する: HPA で maxReplicas 設定を構成して、サブスクライバーの最大レプリカ数を明示的に制限します。推論サーバーの GPU または TPU の割り当てでサポートできる範囲を超えてサブスクライバーをスケーリングしないでください。サブスクライバーをオーバー プロビジョニングすると、ボトルネックが推論サーバーに移行し、スループットを増やすことなくリソース競合が増加します。

  • エンジン指標に基づいて推論サーバーをスケーリングする: 推論エンジンによって直接エクスポートされた指標（CPU/メモリだけでなく）に基づいて、推論サーバーのデプロイをスケーリングします。たとえば、vLLM に vllm:num_requests_waiting 設定を使用します。これは、モデルサーバー レベルで処理バックログを直接測定します。詳細については、Pod を自動的にスケーリングするをご覧ください。

エラーとタイムアウトを処理する

エラーとタイムアウトを処理するには、次の操作を行います。

  • 確認応答期限を事前に延長する: 処理中のメッセージの Pub/Sub 確認応答（ack）期限を事前に延長するようにサブスクライバーを構成して、再配信ループと重複処理を防ぎます。このアプローチが必要なのは、推論タスクがデフォルトのタイムアウト ウィンドウよりも長くかかることが多いためです。原則として、拡張期間は最悪のケースのバッチ推論時間よりも長く設定します。
  • デッドレター トピックを使用して障害を分離する: デッドレター トピックを有効にして、配信に繰り返し失敗する不正な形式のメッセージを自動的に分離します。このアプローチにより、「毒薬」メッセージがキューをブロックしてパイプライン全体を停止することを防ぐことができます。
  • バックオフ戦略を実装する: 推論サーバーが 429（リクエストが多すぎます）または 503（サービスを利用できません）エラーを返した場合、サブスクライバーはこれらのエラーをキャッチし、指数バックオフ戦略を実装して、サーバーが復旧するまで Pub/Sub からの消費を一時的に一時停止する必要があります。

オフライン バッチジョブを大規模にオーケストレートする

大規模なデータセットを処理する場合は、次のベスト プラクティスに沿って、スループットを最大化し、費用対効果を確保し、監査用の包括的なトレーサビリティを実装し、高度な割り当て管理とジョブの優先順位付けを適用します。

マルチノード分散推論に JobSet を使用する

複数のノードが連携する必要がある分散推論ワークロード（TPU Pod で実行される大規模なモデルやマルチノード GPU クラスタなど）をオーケストレートするには、Kubernetes JobSet リソースを使用することをおすすめします。標準の Kubernetes ジョブでは、必要なすべての Pod が同時に起動することを保証できません。これにより、分散ワークロードでデッドロックが発生する可能性があります。

JobSet は、Job のグループをユニットとして管理する Kubernetes ネイティブ API であり、バッチ推論に次のメリットがあります。

  • ギャング スケジューリング: TPU スライスや GPU ノードなど、必要なリソースがすべて使用可能になってからワークロードを開始し、デッドロックを防ぎます。
  • 排他的な配置: 1 つの JobSet がネットワーク トポロジ（TPU スライスなど）に排他的にアクセスできるようにして、相互接続のパフォーマンスを最大化します。
  • 障害復旧: 構成に応じて、ワーカーが失敗した場合に、特定の複製ジョブまたはジョブのセット全体を再起動できます。

データ シャーディングにインデックス付きジョブを使用する

JobSet を使用する場合は、ReplicatedJob を構成して completionMode: Indexed 設定を使用します。この設定により、JOB_COMPLETION_INDEX 環境変数が各 Pod に自動的に挿入されます。推論コードは、このインデックスを使用して、処理するデータの固有のシャードを決定論的に選択できます。

たとえば、100,000 個の画像を含む Cloud Storage バケットがあり、並列処理が 10 の JobSet をデプロイすると、10 個の Pod のそれぞれが起動時にインデックス（0 ～ 9）を読み取ります。Pod 0 は画像 0 ～ 9,999 を処理し、Pod 1 は 10,000 ～ 19,999 を処理するように計算できます。このアプローチにより、個別のタスクキュー サービスが不要になります。

サーバーの飽和状態にサイドカー パターンを使用する

アクセラレータの使用率を最大化するには、サイドカー パターンを使用して 2 つのコンテナで JobSet Pod を構成します。

  • 推論サーバー: GPU または TPU コンピューティングに完全に焦点を当てた最適化されたサーバー（vLLM など）。
  • クライアント ドライバ: localhost のサーバーに大量のリクエストを非同期で送信するロジック コンテナ。

この分離により、GPU または TPU はネットワーク I/O またはデータ前処理を待機している間も常にビジー状態になり、アイドル状態になることはありません。このアプローチを使用しないと、データを順番に読み込むモデルでアクセラレータが I/O オペレーションの完了を待機し、使用率が低下する可能性があります。たとえば、クライアント ドライバは、データの処理を待つのではなく、データをプリフェッチして、推論サーバーに非同期リクエストを継続的に送信し、アクセラレータのリクエスト キューが飽和状態を維持するようにします。

チェックリストの概要

カテゴリ ベスト プラクティス
アーキテクチャ パターン
費用とスループット
メッセージングとスケーリング
オーケストレーション