使用 TF Serving 在 Kubernetes 上部署 🤗 ViT

在上一篇文章中，我们展示了如何使用 TensorFlow Serving 在本地部署来自 🤗 Transformers 的Vision Transformer (ViT) 模型。我们涵盖了将预处理和后处理操作嵌入 Vision Transformer 模型、处理 gRPC 请求等主题！

虽然本地部署是构建有用内容的绝佳开端，但在实际项目中，您需要执行可以服务大量用户的部署。在这篇文章中，您将学习如何使用 Docker 和 Kubernetes 扩展上一篇文章中的本地部署。因此，我们假设您对 Docker 和 Kubernetes 有一定的熟悉度。

这篇文章是在上一篇文章的基础上编写的，所以我们强烈建议您先阅读该文章。您可以在这个仓库中找到本文中讨论的所有代码。

为什么选择 Docker 和 Kubernetes？

扩展我们部署的基本工作流程包括以下步骤：

• 容器化应用逻辑：应用逻辑涉及一个可以处理请求并返回预测的服务模型。对于容器化，Docker 是行业标准的首选。

• 部署 Docker 容器：您有多种选择。最广泛使用的选项是将 Docker 容器部署到 Kubernetes 集群上。Kubernetes 提供了许多部署友好的功能（例如自动扩缩和安全性）。您可以使用像 Minikube 这样的解决方案在本地管理 Kubernetes 集群，或者像 Elastic Kubernetes Service (EKS) 这样的无服务器解决方案。

您可能想知道，在Sagemaker、Vertex AI 等直接提供 ML 部署特定功能的时代，为什么还要使用这种显式设置呢？这是很合理的想法。

上述工作流程在业界被广泛采用，许多组织从中受益。它已经经过多年实战检验。它还允许您对部署进行更精细的控制，同时抽象掉不重要的部分。

本文使用 Google Kubernetes Engine (GKE) 来配置和管理 Kubernetes 集群。如果您使用 GKE，我们假设您已经拥有一个已启用计费的 GCP 项目。此外，请注意，您需要配置 gcloud 工具才能在 GKE 上执行部署。但本文讨论的概念同样适用于您决定使用 Minikube 的情况。

注意：本文中显示的代码片段可以在 Unix 终端上执行，只要您配置了 `gcloud` 工具以及 Docker 和 `kubectl`。更多说明请参阅随附的存储库。

使用 Docker 进行容器化

该服务模型可以处理原始图像输入（以字节为单位），并能够进行预处理和后处理。

在本节中，您将看到如何使用 基础 TensorFlow Serving 镜像 来容器化该模型。TensorFlow Serving 以 SavedModel 格式使用模型。回想一下您在上一篇文章中是如何获取这种 SavedModel 的。我们假设您的 SavedModel 已压缩为 tar.gz 格式。如果需要，您可以从此处获取。然后，SavedModel 应放置在特殊的目录结构 // 中。这是 TensorFlow Serving 同时管理不同版本模型的多个部署的方式。

准备 Docker 镜像

下面的 shell 脚本将 `SavedModel` 放置在 `models` 父目录下的 `hf-vit/1` 中。在准备 Docker 镜像时，您将复制其中的所有内容。这个例子中只有一个模型，但这是一种更通用的方法。

$ MODEL_TAR=model.tar.gz
$ MODEL_NAME=hf-vit
$ MODEL_VERSION=1
$ MODEL_PATH=models/$MODEL_NAME/$MODEL_VERSION

$ mkdir -p $MODEL_PATH
$ tar -xvf $MODEL_TAR --directory $MODEL_PATH

下面，我们展示了在我们的案例中 `models` 目录的结构

$ find /models
/models
/models/hf-vit
/models/hf-vit/1
/models/hf-vit/1/keras_metadata.pb
/models/hf-vit/1/variables
/models/hf-vit/1/variables/variables.index
/models/hf-vit/1/variables/variables.data-00000-of-00001
/models/hf-vit/1/assets
/models/hf-vit/1/saved_model.pb

自定义的 TensorFlow Serving 镜像应基于官方镜像构建。有多种方法可以实现这一点，但您将通过运行 Docker 容器来完成，如官方文档所示。我们首先在后台模式下运行 `tensorflow/serving` 镜像，然后将整个 `models` 目录复制到正在运行的容器中，如下所示。

$ docker run -d --name serving_base tensorflow/serving
$ docker cp models/ serving_base:/models/

我们使用了 TensorFlow Serving 的官方 Docker 镜像作为基础，但您也可以使用您从源代码构建的镜像。

注意：TensorFlow Serving 受益于利用 AVX512 等指令集的硬件优化。这些指令集可以加速深度学习模型推理。因此，如果您知道模型将部署在什么硬件上，通常最好获取 TensorFlow Serving 镜像的优化版本并在整个过程中使用它。

现在，运行中的容器已经拥有了所有必需的文件并以适当的目录结构排列，我们需要创建一个包含这些更改的新 Docker 镜像。这可以通过下面的 `docker commit` 命令完成，您将获得一个名为 `$NEW_IMAGE` 的新 Docker 镜像。需要注意的一点是，您需要将 `MODEL_NAME` 环境变量设置为模型名称，在本例中为 `hf-vit`。这会告诉 TensorFlow Serving 要部署哪个模型。

$ NEW_IMAGE=tfserving:$MODEL_NAME

$ docker commit \ 
    --change "ENV MODEL_NAME $MODEL_NAME" \ 
    serving_base $NEW_IMAGE

在本地运行 Docker 镜像

最后，您可以在本地运行新构建的 Docker 镜像，看看它是否正常工作。下面您可以看到 `docker run` 命令的输出。由于输出非常详细，我们将其精简，只关注重要部分。此外，值得注意的是，它分别打开了 `8500` 和 `8501` 端口，用于 gRPC 和 HTTP/REST 端点。

$ docker run -p 8500:8500 -p 8501:8501 -t $NEW_IMAGE &


---------OUTPUT---------
(Re-)adding model: hf-vit
Successfully reserved resources to load servable {name: hf-vit version: 1}
Approving load for servable version {name: hf-vit version: 1}
Loading servable version {name: hf-vit version: 1}
Reading SavedModel from: /models/hf-vit/1
Reading SavedModel debug info (if present) from: /models/hf-vit/1
Successfully loaded servable version {name: hf-vit version: 1}
Running gRPC ModelServer at 0.0.0.0:8500 ...
Exporting HTTP/REST API at:localhost:8501 ...

推送 Docker 镜像

这里的最后一步是将 Docker 镜像推送到镜像仓库。您将使用 Google Container Registry (GCR) 来实现此目的。以下代码行可以为您完成此操作

$ GCP_PROJECT_ID=<GCP_PROJECT_ID>
$ GCP_IMAGE=gcr.io/$GCP_PROJECT_ID/$NEW_IMAGE

$ gcloud auth configure-docker
$ docker tag $NEW_IMAGE $GCP_IMAGE
$ docker push $GCP_IMAGE

由于我们使用 GCR，您需要在 Docker 镜像标签（请注意其他格式）前加上 `gcr.io/`。准备好 Docker 镜像并推送到 GCR 后，您现在可以继续将其部署到 Kubernetes 集群上。

在 Kubernetes 集群上部署

在 Kubernetes 集群上部署需要以下内容：

• 配置 Kubernetes 集群，本文使用 Google Kubernetes Engine (GKE) 完成。但是，欢迎您使用其他平台和工具，如 EKS 或 Minikube。

• 连接到 Kubernetes 集群以执行部署。

• 编写 YAML 清单。

• 使用实用工具 `kubectl` 使用清单执行部署。

让我们逐一回顾这些步骤。

在 GKE 上配置 Kubernetes 集群

您可以使用类似这样的 shell 脚本来实现此目的（可在此处获取）

$ GKE_CLUSTER_NAME=tfs-cluster
$ GKE_CLUSTER_ZONE=us-central1-a
$ NUM_NODES=2
$ MACHINE_TYPE=n1-standard-8

$ gcloud container clusters create $GKE_CLUSTER_NAME \
    --zone=$GKE_CLUSTER_ZONE \
    --machine-type=$MACHINE_TYPE \
    --num-nodes=$NUM_NODES

GCP 提供了多种机器类型，可让您按需配置部署。我们鼓励您参考文档以了解更多信息。

集群配置完成后，您需要连接到它以执行部署。由于此处使用了 GKE，您还需要进行身份验证。您可以使用如下 shell 脚本来完成这两个操作：

$ GCP_PROJECT_ID=<GCP_PROJECT_ID>

$ export USE_GKE_GCLOUD_AUTH_PLUGIN=True

$ gcloud container clusters get-credentials $GKE_CLUSTER_NAME \
    --zone $GKE_CLUSTER_ZONE \
    --project $GCP_PROJECT_ID

`gcloud container clusters get-credentials` 命令负责连接到集群和身份验证。完成此操作后，您就可以编写清单了。

编写 Kubernetes 清单

Kubernetes 清单使用 YAML 文件编写。虽然可以使用单个清单文件来执行部署，但创建单独的清单文件通常有利于委派关注点分离。通常使用三个清单文件来实现这一点：

• deployment.yaml 通过提供 Docker 镜像的名称、运行 Docker 镜像时的额外参数、用于外部访问的端口以及资源限制来定义 Deployment 的期望状态。

• service.yaml 定义外部客户端与 Kubernetes 集群内部 Pod 之间的连接。

• hpa.yaml 定义了根据 CPU 利用率等规则来扩缩 Deployment 中 Pod 数量的规则。

您可以在这里找到本文的相关清单。下面，我们提供了一个图片概览，说明这些清单是如何使用的。

接下来，我们介绍每个清单的重要部分。

deployment.yaml:

apiVersion: apps/v1
kind: Deployment
metadata:
  labels:
    app: tfs-server
  name: tfs-server
...
spec:
  containers:
  - image: gcr.io/$GCP_PROJECT_ID/tfserving-hf-vit:latest
    name: tfs-k8s
    imagePullPolicy: Always
    args: ["--tensorflow_inter_op_parallelism=2", 
           "--tensorflow_intra_op_parallelism=8"] 
    ports:
    - containerPort: 8500
      name: grpc
    - containerPort: 8501
      name: restapi
    resources:
      limits:
        cpu: 800m
      requests:
        cpu: 800m
...

您可以随意配置 `tfs-server`、`tfs-k8s` 等名称。在 `containers` 下，您指定部署将使用的 Docker 镜像 URI。通过设置容器 `resources` 的允许范围来监控当前资源利用率。这可以使 Horizontal Pod Autoscaler（稍后讨论）决定扩缩容器数量。`requests.cpu` 是操作员设置的确保容器正常工作的最小 CPU 资源量。这里 800m 表示整个 CPU 资源的 80%。因此，HPA 监控所有 Pod 的 `requests.cpu` 总和的平均 CPU 利用率，以做出扩缩决策。

除了 Kubernetes 特定的配置外，您还可以在 `args` 中指定 TensorFlow Serving 特定的选项。在这种情况下，您有两个：

• `tensorflow_inter_op_parallelism`，它设置并行运行以执行独立操作的线程数。建议值为 2。

• `tensorflow_intra_op_parallelism`，它设置并行运行以执行单个操作的线程数。建议值为部署 CPU 的物理核心数。

您可以从此处和此处了解有关这些选项（以及其他选项）的更多信息以及调整它们以进行部署的技巧。

service.yaml:

apiVersion: v1
kind: Service
metadata:
  labels:
    app: tfs-server
  name: tfs-server
spec:
  ports:
  - port: 8500
    protocol: TCP
    targetPort: 8500
    name: tf-serving-grpc
  - port: 8501
    protocol: TCP
    targetPort: 8501
    name: tf-serving-restapi
  selector:
    app: tfs-server
  type: LoadBalancer

我们将服务类型设置为“LoadBalancer”，以便端点可以外部暴露给 Kubernetes 集群。它选择“tfs-server”部署，通过指定的端口与外部客户端建立连接。我们为 gRPC 和 HTTP/REST 连接分别开放了“8500”和“8501”两个端口。

hpa.yaml:

apiVersion: autoscaling/v1
kind: HorizontalPodAutoscaler
metadata:
 name: tfs-server

spec:
 scaleTargetRef:
   apiVersion: apps/v1
   kind: Deployment
   name: tfs-server
 minReplicas: 1
 maxReplicas: 3
 targetCPUUtilizationPercentage: 80

HPA 代表横向 Pod 自动扩缩。它设置了决定何时扩缩目标部署中 Pod 数量的标准。您可以在此处了解更多关于 Kubernetes 内部使用的自动扩缩算法的信息。

在这里，您指定 Kubernetes 应如何处理自动扩缩。特别是，您定义了它应该执行自动扩缩的副本范围——`minReplicas` 和 `maxReplicas`，以及目标 CPU 利用率。`targetCPUUtilizationPercentage` 是一个重要的自动扩缩指标。以下帖子恰当地总结了它的含义（摘自此处）

CPU 利用率是部署中所有 Pod 在过去一分钟内的平均 CPU 使用率除以该部署的请求 CPU。如果 Pod 的平均 CPU 利用率高于您定义的目标，则您的副本将被调整。

回想一下在部署清单中指定 `resources`。通过指定 `resources`，Kubernetes 控制平面开始监控指标，因此 `targetCPUUtilization` 才能正常工作。否则，HPA 不知道 Deployment 的当前状态。

您可以根据自己的需求进行实验并将其设置为所需的数字。但是请注意，自动扩缩将取决于您在 GCP 上可用的配额，因为 GKE 内部使用 Google Compute Engine 来管理这些资源。

执行部署

清单准备好后，您可以使用 `kubectl apply` 命令将其应用于当前连接的 Kubernetes 集群。

$ kubectl apply -f deployment.yaml
$ kubectl apply -f service.yaml
$ kubectl apply -f hpa.yaml

虽然使用 `kubectl` 来应用每个清单以执行部署很方便，但如果清单数量很多，它很快就会变得困难。这时，像 Kustomize 这样的实用工具就能派上用场了。您只需定义另一个名为 `kustomization.yaml` 的规范，如下所示：

commonLabels:
  app: tfs-server
resources:
- deployment.yaml
- hpa.yaml
- service.yaml
apiVersion: kustomize.config.k8s.io/v1beta1
kind: Kustomization

然后，只需一行代码即可执行实际部署：

$ kustomize build . | kubectl apply -f -

完整的说明可在此处找到。部署完成后，我们可以像这样检索端点 IP：

$ kubectl rollout status deployment/tfs-server
$ kubectl get svc tfs-server --watch

---------OUTPUT---------
NAME        TYPE          CLUSTER-IP   EXTERNAL-IP  PORT(S)                        AGE
tfs-server  LoadBalancer  xxxxxxxxxx   xxxxxxxxxx   8500:30869/TCP,8501:31469/TCP  xxx

当外部 IP 可用时，请记下它。

这就总结了将模型部署到 Kubernetes 所需的所有步骤！Kubernetes 优雅地提供了复杂部分的抽象，例如自动扩缩和集群管理，同时让您专注于部署模型时应关注的关键方面，包括资源利用率、安全性（我们在此未涉及）、延迟等性能指标。

测试端点

假设您获得了端点的外部 IP，您可以使用以下列表对其进行测试：

import tensorflow as tf 
import json
import base64

image_path = tf.keras.utils.get_file(
    "image.jpg", "http://images.cocodataset.org/val2017/000000039769.jpg"
)
bytes_inputs = tf.io.read_file(image_path)
b64str = base64.urlsafe_b64encode(bytes_inputs.numpy()).decode("utf-8")
data = json.dumps(
    {"signature_name": "serving_default", "instances": [b64str]}
)

json_response = requests.post(
    "http://<ENDPOINT-IP>:8501/v1/models/hf-vit:predict", 
    headers={"content-type": "application/json"}, 
    data=data
)
print(json.loads(json_response.text))

---------OUTPUT---------
{'predictions': [{'label': 'Egyptian cat', 'confidence': 0.896659195}]}

如果您有兴趣了解此部署在面对更多流量时的性能，我们建议您查看这篇文章。请参阅相应的仓库，了解更多关于使用 Locust 运行负载测试和可视化结果的信息。

关于不同 TF Serving 配置的说明

TensorFlow Serving 提供了各种选项，可根据您的应用程序用例定制部署。下面，我们简要讨论其中一些选项。

`enable_batching` 启用批处理推理功能，该功能在一定的时间窗口内收集传入请求，将它们整理为批次，执行批处理推理，并将每个请求的结果返回给相应的客户端。TensorFlow Serving 提供了一组丰富的可配置选项（例如 `max_batch_size`、`num_batch_threads`）以根据您的部署需求进行调整。您可以在此处了解更多信息。批处理对于那些不需要模型即时预测的应用程序特别有利。在这种情况下，您通常会将多个样本批量收集起来进行预测，然后将这些批次发送进行预测。幸运的是，当我们启用 TensorFlow Serving 的批处理功能时，它可以自动配置所有这些功能。

`enable_model_warmup` 使用虚拟输入数据预热一些惰性实例化的 TensorFlow 组件。通过这种方式，您可以确保所有内容都已正确加载，并且在实际服务期间不会出现延迟。

结论

在本文和相关的仓库中，您学习了如何在 Kubernetes 集群上部署 🤗 Transformers 中的 Vision Transformer 模型。如果您是第一次进行此操作，这些步骤可能看起来有点令人望而生畏，但一旦掌握，它们很快就会成为您工具箱中必不可少的一部分。如果您已经熟悉此工作流程，我们希望本文对您仍然有所帮助。

我们对同一个 Vision Transformer 模型的 ONNX 优化版本应用了相同的部署工作流程。欲了解更多详情，请查看此链接。如果您的部署使用 x86 CPU，则 ONNX 优化模型特别有利。

在下一篇文章中，我们将向您展示如何使用 Vertex AI 大幅减少代码量来执行这些部署——更像是 `model.deploy(autoscaling_config=...)`，然后搞定！我们希望您和我们一样兴奋。

致谢

感谢 Google 的 ML 开发者关系项目团队，他们为我们提供了 GCP 积分，用于进行实验。