Deploying 🤗 ViT on Kubernetes with TF Serving
chansung park, Sayak Paul
In the previous post, we showed howto deploy a Vision Transformer (ViT)model from 🤗 Transformers locally with TensorFlow Serving. We coveredtopics like embedding preprocessing and postprocessing operations withinthe Vision Transformer model, handling gRPC requests, and more!While local deployments are an excellent head start to buildingsomething useful, you’d need to perform deployments that can serve manyusers in real-life projects. In this post, you’ll learn how to scale thelocal 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 highlyrecommend reading it first. You can find all the codediscussed throughout this post in this repository.Why go with Docker and Kubernetes?The basic workflow of scaling up a deployment like ours includes thefollowing steps:Containerizing the application logic: The application logicinvolves a served model that can handle requests and returnpredictions. For containerization, Docker is the industry-standardgo-to.Deploying the Docker container: You have various options here. The mostwidely used option is deploying the Docker container on a Kubernetescluster. Kubernetes provides numerous deployment-friendly features(e.g. autoscaling and security). You can use a solution likeMinikube tomanage Kubernetes clusters locally or a serverless solution likeElastic Kubernetes Service (EKS).You might be wondering why use an explicit setup like this in the ageof Sagemaker, Vertex AIthat provides ML deployment-specific features right off the bat. It is fair to thinkabout it.The above workflow is widely adopted in the industry, and manyorganizations benefit from it. It has already been battle-tested formany years. It also lets you have more granular control of yourdeployments 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 abilling-enabled GCP project if you’re using GKE. Also, note that you’d need toconfigure the gcloud utility forperforming the deployment on GKE. But the concepts discussed in this postequally apply should you decide to use Minikube.Note: The code snippets shown in this post can be executed on a Unix terminalas long as you have configured the gcloud utility along with Docker and kubectl.More instructions are available in the accompanying repository. Containerization with DockerThe serving model can handle raw image inputs as bytes and is capable of preprocessing andpostprocessing.In this section, you’ll see how to containerize that model using thebase TensorFlow Serving Image. TensorFlow Serving consumes modelsin the SavedModel format. Recall how youobtained such a SavedModel in the previous post. We assume thatyou have the SavedModel compressed in tar.gz format. You can fetchit from herejust in case. Then SavedModel should be placed in the special directorystructure of <MODEL_NAME>/<VERSION>/<SavedModel>. This is how TensorFlow Serving simultaneously manages multiple deployments of different versioned models.Preparing the Docker imageThe shell script below places the SavedModel in hf-vit/1 under theparent directory models. You'll copy everything inside it when preparingthe Docker image. There is only one model in this example, but thisis 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_PATHBelow, 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.pbThe 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 theofficial document. We start by running tensorflow/serving image in background mode, then the entire models directory is copied to the running containeras 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, butyou can use ones that you have built from sourceas well.Note: TensorFlow Serving benefits from hardware optimizations that leverage instruction sets such asAVX512. Theseinstruction sets can speed up deep learning model inference. So,if you know the hardware on which the model will be deployed, it’s oftenbeneficial to obtain an optimized build of the TensorFlow Serving imageand use it throughout.Now that the running container has all the required files in theappropriate directory structure, we need to create a new Docker imagethat 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_NAMEenvironment variable to the model name, which is hf-vit in thiscase. This tells TensorFlow Serving what model to deploy.$ NEW_IMAGE=tfserving:$MODEL_NAME$ docker commit \ --change "ENV MODEL_NAME $MODEL_NAME" \ serving_base $NEW_IMAGERunning the Docker image locallyLastly, you can run the newly built Docker image locally to see if itworks fine. Below you see the output of the docker run command. Sincethe output is verbose, we trimmed it down to focus on the importantbits. Also, it is worth noting that it opens up 8500 and 8501ports for gRPC and HTTP/REST endpoints, respectively.$ docker run -p 8500:8500 -p 8501:8501 -t $NEW_IMAGE &---------OUTPUT---------(Re-)adding model: hf-vitSuccessfully 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/1Reading SavedModel debug info (if present) from: /models/hf-vit/1Successfully 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 imageThe final step here is to push the Docker image to an image repository.You'll use Google Container Registry (GCR) for thispurpose. 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_IMAGESince we’re using GCR, you need to prefix theDocker 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 aKubernetes cluster.Deploying on a Kubernetes clusterDeployment on a Kubernetes cluster requires the following:Provisioning a Kubernetes cluster, done with Google Kubernetes Engine (GKE) inthis post. However, you’re welcome to use other platforms and toolslike 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 GKEYou can use a shell script like so for this (availablehere):$ 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_NODESGCP offers a variety of machine types to configure the deployment in away you want. We encourage you to refer to thedocumentationto learn more about it.Once the cluster is provisioned, you need to connect to it to performthe deployment. Since GKE is used here, you also need to authenticateyourself. 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_IDThe gcloud container clusters get-credentials command takes care ofboth connecting to the cluster and authentication. Once this is done,you’re ready to write the manifests.Writing Kubernetes manifestsKubernetes manifests are written in YAMLfiles. While it’s possible to use a single manifest file to perform thedeployment, creating separate manifest files is often beneficial fordelegating the separation of concerns. It’s common to use three manifestfiles for achieving this:deployment.yaml defines the desired state of the Deployment byproviding the name of the Docker image, additional arguments whenrunning the Docker image, the ports to open for external accesses,and the limits of resources.service.yaml defines connections between external clients andinside Pods in the Kubernetes cluster.hpa.yaml defines rules to scale up and down the number of Podsconsisting of the Deployment, such as the percentage of CPUutilization.You can find the relevant manifests for this posthere.Below, we present a pictorial overview of how these manifests areconsumed.Next, we go through the important parts of each of these manifests.deployment.yaml:apiVersion: apps/v1kind: Deploymentmetadata:labels:app: tfs-servername: tfs-server...spec:containers:- image: gcr.io/$GCP_PROJECT_ID/tfserving-hf-vit:latestname: tfs-k8simagePullPolicy: Alwaysargs: ["--tensorflow_inter_op_parallelism=2", "--tensorflow_intra_op_parallelism=8"] ports:- containerPort: 8500name: grpc- containerPort: 8501name: restapiresources:limits:cpu: 800mrequests:cpu: 800m...You can configure the names like tfs-server, tfs-k8s any way youwant. Under containers, you specify the Docker image URI thedeployment will use. The current resource utilization gets monitored bysetting the allowed bounds of the resources for the container. Itcan let Horizontal Pod Autoscaler (discussed later) decide to scale up or down the number ofcontainers. requests.cpu is the minimal amount of CPU resources tomake the container work correctly set by operators. Here 800m means 80%of the whole CPU resource. So, HPA monitors the average CPU utilizationout of the sum of requests.cpu across all Pods to make scalingdecisions.Besides Kubernetes specific configuration, you can specify TensorFlowServing specific options in args.In this case, you have two:tensorflow_inter_op_parallelism, which sets the number of threadsto run in parallel to execute independent operations. Therecommended value for this is 2.tensorflow_intra_op_parallelism, which sets the number of threadsto run in parallel to execute individual operations. The recommendedvalue is the number of physical cores the deployment CPU has.You can learn more about these options (and others) and tips on tuningthem for deployment fromhere andhere.service.yaml:apiVersion: v1kind: Servicemetadata:labels:app: tfs-servername: tfs-serverspec:ports:- port: 8500protocol: TCPtargetPort: 8500name: tf-serving-grpc- port: 8501protocol: TCPtargetPort: 8501name: tf-serving-restapiselector:app: tfs-servertype: LoadBalancerWe made the service type ‘LoadBalancer’ so the endpoints areexposed externally to the Kubernetes cluster. It selects the‘tfs-server’ Deployment to make connections with external clients viathe specified ports. We open two ports of ‘8500’ and ‘8501’ for gRPC andHTTP/REST connections respectively.hpa.yaml:apiVersion: autoscaling/v1kind: HorizontalPodAutoscalermetadata:name: tfs-serverspec:scaleTargetRef:apiVersion: apps/v1kind: Deploymentname: tfs-serverminReplicas: 1maxReplicas: 3targetCPUUtilizationPercentage: 80HPA stands for Horizontal Pod Autoscaler. It sets criteriato decide when to scale the number of Pods in the target Deployment. Youcan learn more about the autoscaling algorithm internally used byKubernetes here.Here you specify how Kubernetes should handle autoscaling. Inparticular, you define the replica bound within which it should performautoscaling – minReplicas\ and maxReplicas and the target CPUutilization. targetCPUUtilizationPercentage is an important metricfor autoscaling. The following thread aptly summarizes what it means(taken from here):The CPU utilization is the average CPU usage of all Pods in adeployment across the last minute divided by the requested CPU of thisdeployment. If the mean of the Pods' CPU utilization is higher than thetarget you defined, your replicas will be adjusted.Recall specifying resources in the deployment manifest. Byspecifying the resources, the Kubernetes control plane startsmonitoring 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 yourrequirements. Note, however, that autoscaling will be contingent on thequota you have available on GCP since GKE internally uses Google Compute Engineto manage these resources.Performing the deploymentOnce the manifests are ready, you can apply them to the currentlyconnected Kubernetes cluster with the kubectl applycommand.$ kubectl apply -f deployment.yaml$ kubectl apply -f service.yaml$ kubectl apply -f hpa.yamlWhile using kubectl is fine for applying each of the manifests toperform the deployment, it can quickly become harder if you have manydifferent manifests. This is where a utility likeKustomize can be helpful. You simplydefine another specification named kustomization.yaml like so:commonLabels:app: tfs-serverresources:- deployment.yaml- hpa.yaml- service.yamlapiVersion: kustomize.config.k8s.io/v1beta1kind: KustomizationThen it’s just a one-liner to perform the actual deployment:$ kustomize build . | kubectl apply -f -Complete instructions are availablehere.Once the deployment has been performed, we can retrieve the endpoint IPlike so:$ kubectl rollout status deployment/tfs-server$ kubectl get svc tfs-server --watch---------OUTPUT---------NAME TYPE CLUSTER-IP EXTERNAL-IP PORT(S) AGEtfs-server LoadBalancer xxxxxxxxxx xxxxxxxxxx 8500:30869/TCP,8501:31469/TCP xxxNote down the external IP when it becomes available.And that sums up all the steps you need to deploy your model onKubernetes! Kubernetes elegantly provides abstractions for complex bitslike autoscaling and cluster management while letting you focus onthe crucial aspects you should care about while deploying a model. Theseinclude resource utilization, security (we didn’t cover that here),performance north stars like latency, etc.Testing the endpointGiven that you got an external IP for the endpoint, you can use thefollowing listing to test it:import tensorflow as tf import jsonimport base64image_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 itmeets more traffic then we recommend you to check this article.Refer to the corresponding repositoryto know more about running load tests with Locust and visualize the results.Notes on different TF Serving configurationsTensorFlow Servingprovidesvarious options to tailor the deployment based on your application usecase. Below, we briefly discuss some of them.enable_batching enables the batch inference capability thatcollects incoming requests with a certain amount of timing window,collates them as a batch, performs a batch inference, and returns theresults of each request to the appropriate clients. TensorFlow Servingprovides a rich set of configurable options (such as max_batch_size,num_batch_threads) to tailor your deployment needs. You can learnmore about themhere. Batching isparticularly beneficial for applications where you don't need predictions from a modelinstantly. In those cases, you'd typically gather together multiple samples for prediction in batches andthen send those batches for prediction. Lucky for us, TensorFlow Serving can configure all of theseautomatically when we enable its batching capabilities.enable_model_warmup warms up some of the TensorFlow componentsthat are lazily instantiated with dummy input data. This way, you canensure everything is appropriately loaded up and that there will be nolags during the actual service time.ConclusionIn this post and the associated repository, you learned about deploying the Vision Transformer modelfrom 🤗 Transformers on a Kubernetes cluster. If you’re doing this forthe first time, the steps may appear to be a little daunting, but onceyou get the grasp, they’ll soon become an essential component of yourtoolbox. If you were already familiar with this workflow, we hope this post was still beneficialfor you. We applied the same deployment workflow for an ONNX-optimized version of the sameVision 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 withsignificantly less code with Vertex AI – more likemodel.deploy(autoscaling_config=...) and boom! We hope you’re just asexcited as we are.AcknowledgementThanks to the ML Developer Relations Program team at Google, whichprovided us with GCP credits for conducting the experiments.