Deploying 🤗 ViT on Kubernetes with TF Serving
chansung, sayakpaul
In the [<u>previous post</u>](https://huggingface.co/blog/tf-serving-vision), we showed how
to deploy a [<u>Vision Transformer (ViT)</u>](https://huggingface.co/docs/transformers/main/en/model_doc/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 [<u>previous post</u>](https://huggingface.co/blog/tf-serving-vision), so, we highly
recommend reading it first. You can find all the code
discussed throughout this post in [<u>this repository</u>](https://github.com/sayakpaul/deploy-hf-tf-vision-models/tree/main/hf_vision_model_onnx_gke).

## 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
  [<u>Minikube</u>](https://minikube.sigs.k8s.io/docs/start/) to
  manage Kubernetes clusters locally or a serverless solution like
  [<u>Elastic Kubernetes Service (EKS)</u>](https://docs.aws.amazon.com/eks/latest/userguide/what-is-eks.html).

You might be wondering why use an explicit setup like this in the age
of [<u>Sagemaker,</u>](https://aws.amazon.com/sagemaker/) [<u>Vertex AI</u>](https://cloud.google.com/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 [<u>Google Kubernetes Engine (GKE)</u>](https://cloud.google.com/kubernetes-engine)
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`](https://cloud.google.com/sdk/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](https://github.com/sayakpaul/deploy-hf-tf-vision-models/tree/main/hf_vision_model_onnx_gke). 

## 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
[<u>base TensorFlow Serving Image</u>](http://hub.docker.com/r/tensorflow/serving/tags/). TensorFlow Serving consumes models
in the [`SavedModel`](https://www.tensorflow.org/guide/saved_model) format. Recall how you
obtained such a `SavedModel` in the [<u>previous post</u>](https://huggingface.co/blog/tf-serving-vision). We assume that
you have the `SavedModel` compressed in `tar.gz` format. You can fetch
it from [<u>here</u>](https://huggingface.co/deploy-hf-tf-vit/vit-base16-extended/resolve/main/saved_model.tar.gz)
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.

```bash
$ 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:

```bash
$ 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](http://hub.docker.com/r/tensorflow/serving/tags/).
There are various approaches for this, but you’ll do this by running a Docker container as illustrated in the
[<u>official document</u>](https://www.tensorflow.org/tfx/serving/serving_kubernetes#commit_image_for_deployment). We start by running `tensorflow/serving` image in background mode, then the entire `models` directory is copied to the running container
as below.

```bash
$ 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 [<u>built from source</u>](https://github.com/tensorflow/serving/blob/master/tensorflow_serving/g3doc/setup.md#building-from-source)
as well.

**Note**: TensorFlow Serving benefits from hardware optimizations that leverage instruction sets such as
[<u>AVX512</u>](https://en.wikipedia.org/wiki/AVX-512). These
instruction sets can [<u>speed up deep learning model inference</u>](https://huggingface.co/blog/bert-cpu-scaling-part-1). 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`](https://docs.docker.com/engine/reference/commandline/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.

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

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