在 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 清单定义应用的期望状态。

基础设施供应

使用 GKE 时，会使用特定机器类型（例如 n1-standard-8）和固定数量的节点来创建集群。认证和连接通过 gcloud container clusters get-credentials 命令完成。

Kubernetes 清单

使用三个主要的清单文件来分离关注点：

  • 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 之间）。

为了管理多个清单，Hugging Face 建议使用 Kustomize，它可以通过 kustomize build . | kubectl apply -f - 实现单命令部署。

测试与优化

端点验证

通过 kubectl get svc 获取外部 IP 后，可向 REST API 发送 base64 编码的图像字节以测试端点。成功的请求会返回预测标签和置信度分数。

TensorFlow Serving 配置

为了进一步优化性能，TensorFlow Serving 提供了多种配置：

  • enable_batching：在时间窗口内收集传入请求，以进行批量推理，对于不需要即时单独预测的应用非常高效。
  • enable_model_warmup：使用虚拟输入数据惰性实例化 TensorFlow 组件，消除实际服务时的初始延迟。

硬件优化

TensorFlow Serving 可以利用硬件特定的指令集（如 AVX512）进行加速。使用针对部署硬件优化构建的 TensorFlow Serving 镜像，可显著提升深度学习模型推理速度。

Sources