在 GKE 上執行批次推論作業的最佳做法

本文提供在 Google Kubernetes Engine (GKE) 執行批次推論工作負載的最佳做法。批次推論是指使用機器學習模型，對大型資料集生成預測結果，優先考量高處理量和成本效益，而非即時、低延遲的回應。

本指南會區分批次推論和要求批次處理 (或動態批次處理)。後者是 vLLM 或 SGLang 等引擎的伺服器端技術，可將並行的即時要求分組，盡量提升加速器效率。您可以將要求批次處理套用至批次推論工作負載。

本指南的最佳做法涵蓋兩種常見的批次推論模式：

  • 非同步推論：在資料產生後不久，以區塊形式處理資料。這種做法的延遲時間通常為幾秒到幾分鐘，可兼顧最新資料的需求，以及同時處理多個項目的效率。非同步推論有時也稱為近乎即時推論。
  • 批次推論：按排定的時間間隔 (例如每晚或每週) 處理大量累積資料。延遲時間通常介於數小時到數天之間，因為這類工作通常會安排在離峰時段，盡可能提高資源可用性。

這些建議是根據在 GKE 上進行推論的最佳做法簡介所述的基礎知識，建構而成的專業最佳化層。在針對批次工作負載進行最佳化之前，請務必遵循模型選取、量化和加速器選擇的核心最佳做法。

選擇批次推論處理的架構模式

選擇正確的架構模式，是部署批次推論工作負載時最重要的決策，因為這會影響延遲時間、處理量和成本之間的取捨。為維持效率，請確保離峰時段的推論輸送量超過傳入查詢的速率，以免佇列無限增長。

針對突發工作使用非同步推論

非同步推論非常適合需要頻繁增量更新的用途，例如：

  • 根據最近的互動，每隔幾分鐘更新使用者推薦設定檔。
  • 每隔一分鐘處理社群媒體提及內容，進行即時監控。
  • 從高頻率的金融資料串流中偵測市場變動信號。
  • 對收到的顧客意見回饋或新聞動態消息執行情緒分析。

如果工作負載可容許幾秒到幾分鐘的延遲，請選擇這個模式。

導入非同步推論時，請考量下列特性：

  • 延遲：您可能會需要等待幾十秒到幾分鐘，才能收到第一個權杖。
  • 資料來源：您通常會處理從 MB 到 GB 不等的資料集，例如 Pub/Sub 訊息或 Cloud Storage 檔案，這些資料會在短時間內累積。
  • 運算模式：基礎架構應支援可處理頻繁工作負載的持續性服務。
  • 成本最佳化：這種模式可在低延遲的即時推論和高處理量的批次處理之間取得平衡。

使用批次推論處理大型資料集

批次推論非常適合大規模的週期性工作，可容許數小時或數天的延遲，例如：

  • 根據前一天的金融交易，每晚生成風險評估報表。
  • 為整個目錄建立產品嵌入，以支援下游搜尋和推薦系統。
  • 為大量圖片資料集加上標籤，用於模型訓練或封存分類。

如果您要處理大量資料，且可容忍數小時到數天的延遲時間，請選擇這個模式。

實作批次推論時，請考量下列特性：

  • 延遲：工作負載啟動延遲通常介於幾分鐘到幾天，因為系統通常會在離峰時段排定工作。
  • 資料來源：處理從 GB 到 PB 的大型資料集，通常儲存在 Cloud Storage 或 BigQuery 資料表中。
  • 運算模式：您使用偶發、突發的工作，這些工作會初始化、處理資料，然後終止。
  • 成本最佳化：這個模式非常適合採用按用量計費模式。由於批次工作完成時間彈性，建議使用 Spot VM 來降低成本。

盡可能提高處理量和成本效益

批次推論工作負載非常適合用於可節省成本的基礎架構，即使中斷也無妨。

使用 Spot VM 降低運算費用

使用 Spot VM 的折扣處理批次工作。由於批次推論工作負載通常可容許延遲和中斷，因此很適合使用 Spot 容量的折扣價格。

請確保批次推論程式碼會實作檢查點，以處理可能發生的先占事件。如果 Spot VM 遭到先占，您可以建立新節點，並從上次處理的批次繼續執行工作負載，不必從頭開始。

調整工作負載批次大小和要求批次大小

為避免資源爭用和工作逾時，請確保傳送至引擎的項目數量 (工作負載批次) 至少與伺服器可處理的並行要求數量 (要求批次) 相同，以免加速器使用率不足。

調整工作負載批次大小

工作負載批次大小是指在單一工作單元中，傳送至推論引擎的項目總數。您可以在用戶端提交邏輯或 Kubernetes 工作設定中，透過資料分割或將多個項目分組為單一要求，來設定這項功能。

如要判斷最佳工作負載批次大小，請使用下列界線：

  • 計算最小批次大小：確保工作負載批次大小至少與要求批次大小相同。舉例來說，如果伺服器可同時處理 256 個項目，但您只傳送一個項目，就會造成資源嚴重閒置。如要找出最小大小，請檢查推論伺服器設定，例如 vLLM 中的 max_num_seqs 引數。您可以設定用戶端邏輯，將多個項目分組到單一要求中，也可以將資料分片，讓每項工作收到的資料量達到或超過要求批次大小。
  • 計算最大批次大小：確保工作負載批次大小能讓 Pod 在達到 Kubernetes 工作中定義的 activeDeadlineSeconds 逾時時間之前完成。預估處理一個要求批次所需的時間，並設定工作負載大小，確保 Pod 能在期限內完成作業。舉例來說，如果 activeDeadlineSeconds 為 3,600 秒，啟動負擔為 600 秒，請確保執行時間上限允許 Pod 在 3,000 秒內完成。

如果工作負載批次大小太小，工作會浪費時間在 Pod 啟動的額外負荷上 (下載權重、佈建、初始化加速器)；如果太大，工作可能會因 activeDeadlineSeconds 超時而遭 GKE 終止，導致工作失敗並失去進度。

調整要求批次大小

要求批次大小是指推論伺服器在加速器上同時處理的並行要求數量。如要最佳化這項參數，請調整推論伺服器設定中的伺服器專屬旗標 (例如 vLLM 的 --max-num-seqs 旗標)。

您的目標是盡量提高 GPU 使用率，但不要觸發記憶體不足 (OOM) 錯誤。如果要求批次大小未經校正，系統會無法充分利用加速器，或導致模型伺服器當機。如果是 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) 調整推論引擎資源。
  • dead-letter 主題：在設定的指數輪詢重試次數後，擷取處理失敗的訊息。

詳情請參閱 GitHub 上的參考實作。

緩衝及匯總要求

如要管理要求流程，請按照下列步驟操作：

  • 使用 Pub/Sub 做為持久緩衝區：實作 Pub/Sub，以持久儲存推論要求。這項設定會做為 FIFO 緩衝區，保留要求，直到消費者有能力處理要求為止，避免流量暴增期間伺服器超載。
  • 使用提取訂閱項目搭配用戶端流量控制：設定提取訂閱模式。訂閱端應用程式可藉此明確要求訊息，但前提是應用程式有能力處理這些訊息，讓您完全掌控取用率。
  • 匯總訊息，填滿伺服器批次大小：避免將一則 Pub/Sub 訊息做為一個推論要求傳送。訂閱者應將多則訊息組合成單一批次要求，並符合推論伺服器的最佳批次大小 (例如與 vLLM 中的 max_num_seqs 設定相符)。這種做法有助於確保加速器充分飽和，並盡可能提高處理量。具體來說，請將訂閱者的 max_messages 提取設定設為 max_num_seqs 的倍數，確保每個模型轉送傳遞都完全飽和。

自動調度訂閱者和伺服器

如要有效執行批次推論，訂閱者 (受 CPU 限制) 和推論伺服器 (受 GPU 或 TPU 限制) 的擴充方式必須不同。

  • 根據工作積壓量調整訂閱端規模：根據 Pub/Sub 的 num_undelivered_messages 指標，為訂閱端部署作業設定 HorizontalPodAutoscaler (HPA)。詳情請參閱 Pub/Sub HPA 範例。使用下列方程式計算要使用的副本數：

    \[ 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) 期限，避免重新傳送迴圈和重複處理。這是因為推論工作通常會超過預設逾時時間範圍，一般來說，延長期限應長於最差情況下的批次推論時間。
  • 使用 dead-letter 主題隔離失敗：啟用 dead-letter 主題，自動隔離無法重複傳送的格式錯誤訊息。這樣做可避免「毒藥丸」訊息阻塞佇列，導致整個管道停止運作。
  • 實作輪詢策略：如果推論伺服器傳回 429 (要求過多) 或 503 (服務無法使用) 錯誤，訂閱者必須擷取這些錯誤並實作指數輪詢策略，暫時停止從 Pub/Sub 服務取用資料，直到伺服器恢復運作為止。

大規模自動調度管理批次工作

處理大量資料集時，請遵循下列最佳做法，盡可能提高處理量、確保成本效益、實作完整的稽核追蹤功能，以及套用進階配額管理和工作優先順序設定。

使用 JobSet 進行多節點分散式推論

建議您使用 Kubernetes JobSet 資源，協調需要多個節點合作的分散式推論工作負載，例如在 TPU Pod 或多節點 GPU 叢集上執行的大型模型。標準 Kubernetes 工作無法保證所有必要的 Pod 同時啟動，這可能會導致分散式工作負載發生死結。

JobSet 是 Kubernetes 原生 API，可將多個 Job 視為單一單位管理，並為批次推論提供下列優點：

  • 群組排程：確保在啟動工作負載前，所有必要資源 (例如 TPU 節點或 GPU 節點) 皆可使用，避免發生死結。
  • 專屬放置位置：確保單一 JobSet 專屬存取網路拓撲 (例如 TPU 配量)，盡可能提升互連網路效能。
  • 失敗復原：如果 worker 發生故障，您可以根據設定重新啟動特定複製工作或整組工作。

使用已編列索引的工作進行資料分割

使用 JobSet 時，請將 ReplicatedJob 設為使用 completionMode: Indexed 設定。這項設定會自動將 JOB_COMPLETION_INDEX 環境變數插入每個 Pod。推論程式碼可以使用這個索引，確定要選取哪個資料分片進行處理。

舉例來說，如果您有一個包含 100,000 張圖片的 Cloud Storage bucket，並部署平行處理量為 10 的 JobSet，則 10 個 Pod 各自會在啟動時讀取索引 (0-9)。Pod 0 接著會計算出應處理的圖片為 0 到 9,999，而 Pod 1 則處理 10,000 到 19,999。這種方法可減少對獨立工作佇列服務的需求。

使用 Sidecar 模式處理伺服器飽和問題

如要盡量提高加速器使用率，請使用 Sidecar 模式，透過兩個容器設定 JobSet Pod：

  • 推論伺服器：經過最佳化的伺服器 (例如 vLLM)，完全專注於 GPU 或 TPU 運算。
  • 用戶端驅動程式：邏輯容器，會將大量要求非同步傳送至本機主機上的伺服器。

這種解耦方式可確保 GPU 或 TPU 保持忙碌狀態，不會在等待網路 I/O 或資料預先處理時閒置。如果不採用這種方法，模型依序載入資料時，加速器可能會等待 I/O 作業完成，導致使用率偏低。舉例來說，用戶端驅動程式可以預先擷取資料，並持續將非同步要求傳送至推論伺服器，確保加速器的要求佇列保持飽和，不必等待資料處理完畢。

檢查清單摘要

類別 最佳做法
架構模式
費用和輸送量
訊息和擴展
自動調度管理