Deploying 🤗 ViT on Kubernetes with TF Serving

In the previous post, we showed how to deploy a Vision Transformer (ViT) model from 🤗 Transformers locally with TensorFlow Serving. We covered topics like embedding preprocessing and postprocessing operations within the Vision Transformer model, handling gRPC requests, and more!

While local deployments are an excellent head start to building something useful, you’d need to perform deployments that can serve many users in real-life projects. In this post, you’ll learn how to scale the local deployment from the previous post with Docker and Kubernetes. Therefore, we assume some familiarity with Docker and Kubernetes.

This post builds on top of the previous post, so, we highly recommend reading it first. You can find all the code discussed throughout this post in this repository.

Why go with Docker and Kubernetes?

The basic workflow of scaling up a deployment like ours includes the following steps:

• Containerizing the application logic: The application logic involves a served model that can handle requests and return predictions. For containerization, Docker is the industry-standard go-to.

• Deploying the Docker container: You have various options here. The most widely used option is deploying the Docker container on a Kubernetes cluster. Kubernetes provides numerous deployment-friendly features (e.g. autoscaling and security). You can use a solution like Minikube to manage Kubernetes clusters locally or a serverless solution like Elastic Kubernetes Service (EKS).

You might be wondering why use an explicit setup like this in the age of Sagemaker, Vertex AI that provides ML deployment-specific features right off the bat. It is fair to think about it.

The above workflow is widely adopted in the industry, and many organizations benefit from it. It has already been battle-tested for many years. It also lets you have more granular control of your deployments while abstracting away the non-trivial bits.

This post uses Google Kubernetes Engine (GKE) to provision and manage a Kubernetes cluster. We assume you already have a billing-enabled GCP project if you’re using GKE. Also, note that you’d need to configure the gcloud utility for performing the deployment on GKE. But the concepts discussed in this post equally apply should you decide to use Minikube.

Note: The code snippets shown in this post can be executed on a Unix terminal as long as you have configured the gcloud utility along with Docker and kubectl. More instructions are available in the accompanying repository.

Containerization with Docker

The serving model can handle raw image inputs as bytes and is capable of preprocessing and postprocessing.

In this section, you’ll see how to containerize that model using the base TensorFlow Serving Image. TensorFlow Serving consumes models in the SavedModel format. Recall how you obtained such a SavedModel in the previous post. We assume that you have the SavedModel compressed in tar.gz format. You can fetch it from here just in case. Then SavedModel should be placed in the special directory structure of <MODEL_NAME>/<VERSION>/<SavedModel>. This is how TensorFlow Serving simultaneously manages multiple deployments of different versioned models.

Preparing the Docker image

The shell script below places the SavedModel in hf-vit/1 under the parent directory models. You'll copy everything inside it when preparing the Docker image. There is only one model in this example, but this is a more generalizable approach.

$ 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

Below, we show how the models directory is structured in our case:

$ 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

The custom TensorFlow Serving image should be built on top of the base one. There are various approaches for this, but you’ll do this by running a Docker container as illustrated in the official document. We start by running tensorflow/serving image in background mode, then the entire models directory is copied to the running container as below.

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

We used the official Docker image of TensorFlow Serving as the base, but you can use ones that you have built from source as well.

Note: TensorFlow Serving benefits from hardware optimizations that leverage instruction sets such as AVX512. These instruction sets can speed up deep learning model inference. So, if you know the hardware on which the model will be deployed, it’s often beneficial to obtain an optimized build of the TensorFlow Serving image and use it throughout.

Now that the running container has all the required files in the appropriate directory structure, we need to create a new Docker image that includes these changes. This can be done with the docker commit command below, and you'll have a new Docker image named $NEW_IMAGE. One important thing to note is that you need to set the MODEL_NAME environment variable to the model name, which is hf-vit in this case. This tells TensorFlow Serving what model to deploy.

$ NEW_IMAGE=tfserving:$MODEL_NAME

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

Running the Docker image locally

Lastly, you can run the newly built Docker image locally to see if it works fine. Below you see the output of the docker run command. Since the output is verbose, we trimmed it down to focus on the important bits. Also, it is worth noting that it opens up 8500 and 8501 ports for gRPC and HTTP/REST endpoints, respectively.

$ 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 ...

Pushing the Docker image

The final step here is to push the Docker image to an image repository. You'll use Google Container Registry (GCR) for this purpose. The following lines of code can do this for you:

$ 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

Since we’re using GCR, you need to prefix the Docker image tag (note the other formats too) with gcr.io/<GCP_PROJECT_ID> . With the Docker image prepared and pushed to GCR, you can now proceed to deploy it on a Kubernetes cluster.

Deploying on a Kubernetes cluster

Deployment on a Kubernetes cluster requires the following:

• Provisioning a Kubernetes cluster, done with Google Kubernetes Engine (GKE) in this post. However, you’re welcome to use other platforms and tools like EKS or Minikube.

• Connecting to the Kubernetes cluster to perform a deployment.

• Writing YAML manifests.

• Performing deployment with the manifests with a utility tool, kubectl.

Let’s go over each of these steps.

Provisioning a Kubernetes cluster on GKE

You can use a shell script like so for this (available here):

$ 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 offers a variety of machine types to configure the deployment in a way you want. We encourage you to refer to the documentation to learn more about it.

Once the cluster is provisioned, you need to connect to it to perform the deployment. Since GKE is used here, you also need to authenticate yourself. You can use a shell script like so to do both of these:

$ 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

The gcloud container clusters get-credentials command takes care of both connecting to the cluster and authentication. Once this is done, you’re ready to write the manifests.

Writing Kubernetes manifests

Kubernetes manifests are written in YAML files. While it’s possible to use a single manifest file to perform the deployment, creating separate manifest files is often beneficial for delegating the separation of concerns. It’s common to use three manifest files for achieving this:

• deployment.yaml defines the desired state of the Deployment by providing the name of the Docker image, additional arguments when running the Docker image, the ports to open for external accesses, and the limits of resources.

• service.yaml defines connections between external clients and inside Pods in the Kubernetes cluster.

• hpa.yaml defines rules to scale up and down the number of Pods consisting of the Deployment, such as the percentage of CPU utilization.

You can find the relevant manifests for this post here. Below, we present a pictorial overview of how these manifests are consumed.

Next, we go through the important parts of each of these manifests.

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
...

You can configure the names like tfs-server, tfs-k8s any way you want. Under containers, you specify the Docker image URI the deployment will use. The current resource utilization gets monitored by setting the allowed bounds of the resources for the container. It can let Horizontal Pod Autoscaler (discussed later) decide to scale up or down the number of containers. requests.cpu is the minimal amount of CPU resources to make the container work correctly set by operators. Here 800m means 80% of the whole CPU resource. So, HPA monitors the average CPU utilization out of the sum of requests.cpu across all Pods to make scaling decisions.

Besides Kubernetes specific configuration, you can specify TensorFlow Serving specific options in args.In this case, you have two:

• tensorflow_inter_op_parallelism, which sets the number of threads to run in parallel to execute independent operations. The recommended value for this is 2.

• tensorflow_intra_op_parallelism, which sets the number of threads to run in parallel to execute individual operations. The recommended value is the number of physical cores the deployment CPU has.

You can learn more about these options (and others) and tips on tuning them for deployment from here and here.

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

We made the service type ‘LoadBalancer’ so the endpoints are exposed externally to the Kubernetes cluster. It selects the ‘tfs-server’ Deployment to make connections with external clients via the specified ports. We open two ports of ‘8500’ and ‘8501’ for gRPC and HTTP/REST connections respectively.

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 stands for Horizontal Pod Autoscaler. It sets criteria to decide when to scale the number of Pods in the target Deployment. You can learn more about the autoscaling algorithm internally used by Kubernetes here.

Here you specify how Kubernetes should handle autoscaling. In particular, you define the replica bound within which it should perform autoscaling – minReplicas\ and maxReplicas and the target CPU utilization. targetCPUUtilizationPercentage is an important metric for autoscaling. The following thread aptly summarizes what it means (taken from here):

The CPU utilization is the average CPU usage of all Pods in a deployment across the last minute divided by the requested CPU of this deployment. If the mean of the Pods' CPU utilization is higher than the target you defined, your replicas will be adjusted.

Recall specifying resources in the deployment manifest. By specifying the resources, the Kubernetes control plane starts monitoring the metrics, so the targetCPUUtilization works. Otherwise, HPA doesn't know the current status of the Deployment.

You can experiment and set these to the required numbers based on your requirements. Note, however, that autoscaling will be contingent on the quota you have available on GCP since GKE internally uses Google Compute Engine to manage these resources.

Performing the deployment

Once the manifests are ready, you can apply them to the currently connected Kubernetes cluster with the kubectl apply command.

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

While using kubectl is fine for applying each of the manifests to perform the deployment, it can quickly become harder if you have many different manifests. This is where a utility like Kustomize can be helpful. You simply define another specification named kustomization.yaml like so:

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

Then it’s just a one-liner to perform the actual deployment:

$ kustomize build . | kubectl apply -f -

Complete instructions are available here. Once the deployment has been performed, we can retrieve the endpoint IP like so:

$ 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

Note down the external IP when it becomes available.

And that sums up all the steps you need to deploy your model on Kubernetes! Kubernetes elegantly provides abstractions for complex bits like autoscaling and cluster management while letting you focus on the crucial aspects you should care about while deploying a model. These include resource utilization, security (we didn’t cover that here), performance north stars like latency, etc.

Testing the endpoint

Given that you got an external IP for the endpoint, you can use the following listing to test it:

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}]}

If you’re interested to know how this deployment would perform if it meets more traffic then we recommend you to check this article. Refer to the corresponding repository to know more about running load tests with Locust and visualize the results.

Notes on different TF Serving configurations

TensorFlow Serving provides various options to tailor the deployment based on your application use case. Below, we briefly discuss some of them.

enable_batching enables the batch inference capability that collects incoming requests with a certain amount of timing window, collates them as a batch, performs a batch inference, and returns the results of each request to the appropriate clients. TensorFlow Serving provides a rich set of configurable options (such as max_batch_size, num_batch_threads) to tailor your deployment needs. You can learn more about them here. Batching is particularly beneficial for applications where you don't need predictions from a model instantly. In those cases, you'd typically gather together multiple samples for prediction in batches and then send those batches for prediction. Lucky for us, TensorFlow Serving can configure all of these automatically when we enable its batching capabilities.

enable_model_warmup warms up some of the TensorFlow components that are lazily instantiated with dummy input data. This way, you can ensure everything is appropriately loaded up and that there will be no lags during the actual service time.

Conclusion

In this post and the associated repository, you learned about deploying the Vision Transformer model from 🤗 Transformers on a Kubernetes cluster. If you’re doing this for the first time, the steps may appear to be a little daunting, but once you get the grasp, they’ll soon become an essential component of your toolbox. If you were already familiar with this workflow, we hope this post was still beneficial for you.

We applied the same deployment workflow for an ONNX-optimized version of the same Vision Transformer model. For more details, check out this link. ONNX-optimized models are especially beneficial if you're using x86 CPUs for deployment.

In the next post, we’ll show you how to perform these deployments with significantly less code with Vertex AI – more like model.deploy(autoscaling_config=...) and boom! We hope you’re just as excited as we are.

Acknowledgement

Thanks to the ML Developer Relations Program team at Google, which provided us with GCP credits for conducting the experiments.