在 Kubernetes 上使用 TF Serving 部署 Hugging Face ViT

Hugging Face 詳細說明了一套工作流程，用於透過 Docker 與 Kubernetes 搭配 TensorFlow Serving，擴展來自 🤗 Transformers 函式庫的 Vision Transformer（ViT）模型的部署。此方法讓開發者能夠從本機原型轉移至能夠透過容器化與自動化編排處理大量使用者流量的生產環境。

使用 Docker 與 Kubernetes 擴展模型部署

為了在實務專案中擴展模型部署，Hugging Face 建議採用兩步驟的工作流程：將應用程式邏輯容器化，並將這些容器部署至 Kubernetes 叢集。雖然有 SageMaker 或 Vertex AI 等受管服務可用，使用 Docker 與 Kubernetes 能在抽象化複雜基礎設施管理的同時，提供更細緻的部署控制。

在此實作中，我們使用 Google Kubernetes Engine（GKE）來佈建與管理叢集，儘管這些概念同樣適用於其他平台，如 Amazon EKS 或本機工具 Minikube。

使用 Docker 的容器化流程

TensorFlow Serving 要求模型必須採用 SavedModel 格式，並以特定目錄結構組織：<MODEL_NAME>/<VERSION>/<SavedModel>。此結構讓 TensorFlow Serving 能同時管理模型的多個版本。

建立自訂映像檔

建立可投入生產的映像檔的流程包括：

  1. 目錄設定：將 SavedModel 放置於結構化路徑（例如 models/hf-vit/1）中。
  2. 基礎映像整合：執行官方的 tensorflow/serving 映像，並將 models 目錄複製到執行中的容器內。
  3. 映像提交：使用 docker commit 建立新映像，並設定 MODEL_NAME 環境變數（例如 hf-vit），以告訴 TensorFlow Serving 要部署哪個模型。

本機測試與發佈

在叢集部署之前，會在本機測試映像，將埠 8500（gRPC）與 8501（HTTP/REST）映射。驗證無誤後，使用 gcloud auth configure-docker 與 docker push 將映像推送至註冊表，例如 Google Container Registry（GCR）。

Kubernetes 部署與編排

部署至 Kubernetes 叢集需要佈建基礎設施，並透過 YAML manifest 定義應用程式的期望狀態。

基礎設施佈建

使用 GKE 時，會以特定機器類型（例如 n1-standard-8）與設定好的節點數量佈建叢集。認證與連線透過 gcloud container clusters get-credentials 指令處理。

Kubernetes Manifest

使用三個主要的 manifest 檔案來分離關注點：

  • deployment.yaml：定義 Docker 映像 URI、資源限制（例如 cpu: 800m）以及 TensorFlow Serving 的特定參數。關鍵調校參數包括 tensorflow_inter_op_parallelism（建議值：2）與 tensorflow_intra_op_parallelism（建議值：實體 CPU 核心數）。
  • service.yaml：設定 LoadBalancer，將 gRPC（埠 8500）與 REST（埠 8501）端點向外部暴露。
  • hpa.yaml：實作水平 Pod 自動擴縮器（HPA），根據 targetCPUUtilizationPercentage（例如 80%）自動調整副本數量（例如 1 到 3 之間）。

為了管理多個 manifest，Hugging Face 建議使用 Kustomize，它允許透過單一指令 kustomize build . | kubectl apply -f - 進行部署。

測試與最佳化

端點驗證

透過 kubectl get svc 取得外部 IP 後，可將 base64 編碼的影像位元組傳送至 REST API 以測試端點。成功的請求會回傳預測的標籤與信心分數。

TensorFlow Serving 設定

為了進一步優化效能，TensorFlow Serving 提供了多項設定：

  • enable_batching：在時間窗口內收集進入的請求，以執行批次推論，對於不需要即時單筆預測的應用程式而言，效率極高。
  • enable_model_warmup：使用虛擬輸入資料延遲載入 TensorFlow 元件，消除實際服務時的初始延遲。

硬體最佳化

TensorFlow Serving 可透過硬體特定的指令集（如 AVX512）加速。使用針對部署硬體優化的 TensorFlow Serving 映像版本，可顯著提升深度學習模型的推論速度。

Sources