Serve scalable LLMs on GKE using TorchServe

Prepare the environment

Clone the example repository and open the tutorial directory:

git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
cd kubernetes-engine-samples/ai-ml/t5-model-serving

Create the cluster

Run the following command:

gcloud container clusters create-auto ml-cluster \
    --release-channel=RELEASE_CHANNEL \
    --cluster-version=CLUSTER_VERSION \
    --location=us-central1

Replace the following:

• RELEASE_CHANNEL: the release channel for your cluster. Must be one of rapid, regular, or stable. Choose a channel that has GKE version 1.28.3-gke.1203000 or later to use L4 GPUs. To see the versions available in a specific channel, see View the default and available versions for release channels.
• CLUSTER_VERSION: the GKE version to use. Must be 1.28.3-gke.1203000 or later.

This operation takes several minutes to complete.

Create an Artifact Registry repository

1. Create a new Artifact Registry standard repository with the Docker format in the same region as your cluster:

   gcloud artifacts repositories create models \
       --repository-format=docker \
       --location=us-central1 \
       --description="Repo for T5 serving image"
   

2. Verify the repository name:

   gcloud artifacts repositories describe models \
       --location=us-central1
   

   The output is similar to the following:

   Encryption: Google-managed key
   Repository Size: 0.000MB
   createTime: '2023-06-14T15:48:35.267196Z'
   description: Repo for T5 serving image
   format: DOCKER
   mode: STANDARD_REPOSITORY
   name: projects/PROJECT_ID/locations/us-central1/repositories/models
   updateTime: '2023-06-14T15:48:35.267196Z'
   

Package the model

In this section, you package the model and the serving framework in a single container image using Cloud Build and push the resulting image to the Artifact Registry repository.

1. Review the Dockerfile for the container image:

   # Copyright 2023 Google LLC
   #
   # Licensed under the Apache License, Version 2.0 (the "License");
   # you may not use this file except in compliance with the License.
   # You may obtain a copy of the License at
   #
   #     https://www.apache.org/licenses/LICENSE-2.0
   #
   # Unless required by applicable law or agreed to in writing, software
   # distributed under the License is distributed on an "AS IS" BASIS,
   # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   # See the License for the specific language governing permissions and
   # limitations under the License.
   
   ARG BASE_IMAGE=pytorch/torchserve:0.12.0-cpu
   
   FROM alpine/git
   
   ARG MODEL_NAME=t5-small
   ARG MODEL_REPO=https://huggingface.co/${MODEL_NAME}
   ENV MODEL_NAME=${MODEL_NAME}
   ENV MODEL_VERSION=${MODEL_VERSION}
   
   RUN git clone "${MODEL_REPO}" /model
   
   FROM ${BASE_IMAGE}
   
   ARG MODEL_NAME=t5-small
   ARG MODEL_VERSION=1.0
   ENV MODEL_NAME=${MODEL_NAME}
   ENV MODEL_VERSION=${MODEL_VERSION}
   
   COPY --from=0 /model/. /home/model-server/
   COPY handler.py \
        model.py \
        requirements.txt \
        setup_config.json /home/model-server/
   
   RUN  torch-model-archiver \
        --model-name="${MODEL_NAME}" \
        --version="${MODEL_VERSION}" \
        --model-file="model.py" \
        --serialized-file="pytorch_model.bin" \
        --handler="handler.py" \
        --extra-files="config.json,spiece.model,tokenizer.json,setup_config.json" \
        --runtime="python" \
        --export-path="model-store" \
        --requirements-file="requirements.txt"
   
   FROM ${BASE_IMAGE}
   
   ENV PATH /home/model-server/.local/bin:$PATH
   ENV TS_CONFIG_FILE /home/model-server/config.properties
   # CPU inference will throw a warning cuda warning (not error)
   # Could not load dynamic library 'libnvinfer_plugin.so.7'
   # This is expected behaviour. see: https://stackoverflow.com/a/61137388
   ENV TF_CPP_MIN_LOG_LEVEL 2
   
   COPY --from=1 /home/model-server/model-store/ /home/model-server/model-store
   COPY config.properties /home/model-server/
   

   This Dockerfile defines the following multiple stage build process:

   1. Download the model artifacts from the Hugging Face repository.
   2. Package the model using the PyTorch Serving Archive tool. This creates a model archive (.mar) file that the inference server uses to load the model.
   3. Build the final image with PyTorch Serve.

2. Build and push the image using Cloud Build:

   gcloud builds submit model/ \
       --region=us-central1 \
       --config=model/cloudbuild.yaml \
       --substitutions=_LOCATION=us-central1,_MACHINE=gpu,_MODEL_NAME=t5-small,_MODEL_VERSION=1.0
   

   The build process takes several minutes to complete. If you use a larger model size than t5-small, the build process might take significantly more time.

3. Check that the image is in the repository:

   gcloud artifacts docker images list us-central1-docker.pkg.dev/PROJECT_ID/models
   

   Replace PROJECT_ID with your Google Cloud project ID.

   The output is similar to the following:

   IMAGE                                                     DIGEST         CREATE_TIME          UPDATE_TIME
   us-central1-docker.pkg.dev/PROJECT_ID/models/t5-small     sha256:0cd...  2023-06-14T12:06:38  2023-06-14T12:06:38
   

Deploy the packaged model to GKE

To deploy the image, modify the Kubernetes manifest in the example repository to match your environment.

1. Review the manifest for the inference workload:

   # Copyright 2023 Google LLC
   #
   # Licensed under the Apache License, Version 2.0 (the "License");
   # you may not use this file except in compliance with the License.
   # You may obtain a copy of the License at
   #
   #     https://www.apache.org/licenses/LICENSE-2.0
   #
   # Unless required by applicable law or agreed to in writing, software
   # distributed under the License is distributed on an "AS IS" BASIS,
   # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   # See the License for the specific language governing permissions and
   # limitations under the License.
   
   ---
   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: t5-inference
     labels:
       model: t5
       version: v1.0
       machine: gpu
   spec:
     replicas: 1
     selector:
       matchLabels:
         model: t5
         version: v1.0
         machine: gpu
     template:
       metadata:
         labels:
           model: t5
           version: v1.0
           machine: gpu
       spec:
         nodeSelector:
           cloud.google.com/gke-accelerator: nvidia-l4
         securityContext:
           fsGroup: 1000
           runAsUser: 1000
           runAsGroup: 1000
         containers:
           - name: inference
             image: us-central1-docker.pkg.dev/PROJECT_ID/models/t5-small:1.0-gpu
             imagePullPolicy: IfNotPresent
             args: ["torchserve", "--start", "--foreground"]
             resources:
               limits:
                 nvidia.com/gpu: "1"
                 cpu: "3000m"
                 memory: 16Gi
                 ephemeral-storage: 10Gi
               requests:
                 nvidia.com/gpu: "1"
                 cpu: "3000m"
                 memory: 16Gi
                 ephemeral-storage: 10Gi
             ports:
               - containerPort: 8080
                 name: http
               - containerPort: 8081
                 name: management
               - containerPort: 8082
                 name: metrics
             readinessProbe:
               httpGet:
                 path: /ping
                 port: http
               initialDelaySeconds: 120
               failureThreshold: 10
             livenessProbe:
               httpGet:
                 path: /models/t5-small
                 port: management
               initialDelaySeconds: 150
               periodSeconds: 5
   ---
   apiVersion: v1
   kind: Service
   metadata:
     name: t5-inference
     labels:
       model: t5
       version: v1.0
       machine: gpu
   spec:
     type: ClusterIP
     selector:
       model: t5
       version: v1.0
       machine: gpu
     ports:
       - port: 8080
         name: http
         targetPort: http
       - port: 8081
         name: management
         targetPort: management
       - port: 8082
         name: metrics
         targetPort: metrics
   

2. Replace PROJECT_ID with your Google Cloud project ID:

   sed -i "s/PROJECT_ID/PROJECT_ID/g" "kubernetes/serving-gpu.yaml"
   

   This ensures that the container image path in the Deployment specification matches the path to your T5 model image in Artifact Registry.

3. Create the Kubernetes resources:

   kubectl create -f kubernetes/serving-gpu.yaml
   

To verify that the model deployed successfully, do the following:

1. Get the status of the Deployment and the Service:

   kubectl get -f kubernetes/serving-gpu.yaml
   

   Wait until the output shows ready Pods, similar to the following. Depending on the size of the image, the first image pull might take several minutes.

   NAME                            READY   UP-TO-DATE    AVAILABLE   AGE
   deployment.apps/t5-inference    1/1     1             0           66s
   
   NAME                    TYPE        CLUSTER-IP        EXTERNAL-IP   PORT(S)                       AGE
   service/t5-inference    ClusterIP   10.48.131.86    <none>        8080/TCP,8081/TCP,8082/TCP    66s
   

2. Open a local port for the t5-inference Service:

   kubectl port-forward svc/t5-inference 8080
   

3. Open a new terminal window and send a test request to the Service:

   curl -v -X POST -H 'Content-Type: application/json' -d '{"text": "this is a test sentence", "from": "en", "to": "fr"}' "http://localhost:8080/predictions/t5-small/1.0"
   

   If the test request fails and the Pod connection closes, check the logs:

   kubectl logs deployments/t5-inference
   

   If the output is similar to the following, TorchServe failed to install some model dependencies:

   org.pytorch.serve.archive.model.ModelException: Custom pip package installation failed for t5-small
   

   To resolve this issue, restart the Deployment:

   kubectl rollout restart deployment t5-inference
   

   The Deployment controller creates a new Pod. Repeat the previous steps to open a port on the new Pod.

Access the deployed model using the web application

1. Build and push the Fast Dash web application as a container image in Artifact Registry:

   gcloud builds submit client-app/ \
       --region=us-central1 \
       --config=client-app/cloudbuild.yaml
   

2. Open kubernetes/application.yaml in a text editor and replace PROJECT_ID in the image: field with your project ID. Alternatively, run the following command:

   sed -i "s/PROJECT_ID/PROJECT_ID/g" "kubernetes/application.yaml"
   

3. Create the Kubernetes resources:

   kubectl create -f kubernetes/application.yaml
   

   The Deployment and Service might take some time to fully provision.

4. To check the status, run the following command:

   kubectl get -f kubernetes/application.yaml
   

   Wait until the output shows ready Pods, similar to the following:

   NAME                       READY   UP-TO-DATE   AVAILABLE   AGE
   deployment.apps/fastdash   1/1     1            0           1m
   
   NAME               TYPE       CLUSTER-IP      EXTERNAL-IP   PORT(S)          AGE
   service/fastdash   NodePort   203.0.113.12    <none>        8050/TCP         1m
   

5. The web application is now running, although it isn't exposed on an external IP address. To access the web application, open a local port:

   kubectl port-forward service/fastdash 8050
   

6. In a browser, open the web interface:

   • If you're using a local shell, open a browser and go to http://127.0.0.1:8050.
   • If you're using Cloud Shell, click Web preview, and then click Change port. Specify port 8050.

7. To send a request to the T5 model, specify values in the TEXT, FROM LANG, and TO LANG fields in the web interface and click Submit. For a list of available languages, see the T5 documentation.

Enable autoscaling for the model

This section shows you how to enable autoscaling for the model based on metrics from Google Cloud Managed Service for Prometheus by doing the following:

1. Install Custom Metrics Stackdriver Adapter
2. Apply PodMonitoring and HorizontalPodAutoscaling configurations

Google Cloud Managed Service for Prometheus is enabled by default in Autopilot clusters running version 1.25 and later.

Install Custom Metrics Stackdriver Adapter

This adapter lets your cluster use metrics from Prometheus to make Kubernetes autoscaling decisions.

1. Deploy the adapter:

   kubectl create -f https://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml
   

2. Create an IAM service account for the adapter to use:

   gcloud iam service-accounts create monitoring-viewer
   

3. Grant the IAM service account the monitoring.viewer role on the project and the iam.workloadIdentityUser role:

   gcloud projects add-iam-policy-binding PROJECT_ID \
       --member "serviceAccount:monitoring-viewer@PROJECT_ID.iam.gserviceaccount.com" \
       --role roles/monitoring.viewer
   gcloud iam service-accounts add-iam-policy-binding monitoring-viewer@PROJECT_ID.iam.gserviceaccount.com \
       --role roles/iam.workloadIdentityUser \
       --member "serviceAccount:PROJECT_ID.svc.id.goog[custom-metrics/custom-metrics-stackdriver-adapter]"
   

   Replace PROJECT_ID with your Google Cloud project ID.

4. Annotate the Kubernetes ServiceAccount of the adapter to let it impersonate the IAM service account:

   kubectl annotate serviceaccount custom-metrics-stackdriver-adapter \
       --namespace custom-metrics \
       iam.gke.io/gcp-service-account=monitoring-viewer@PROJECT_ID.iam.gserviceaccount.com
   

5. Restart the adapter to propagate the changes:

   kubectl rollout restart deployment custom-metrics-stackdriver-adapter \
       --namespace=custom-metrics
   

Apply PodMonitoring and HorizontalPodAutoscaling configurations

PodMonitoring is a Google Cloud Managed Service for Prometheus custom resource that enables metrics ingestion and target scraping in a specific namespace.

1. Deploy the PodMonitoring resource in the same namespace as the TorchServe Deployment:

   kubectl apply -f kubernetes/pod-monitoring.yaml
   

2. Review the HorizontalPodAutoscaler manifest:

   # Copyright 2023 Google LLC
   #
   # Licensed under the Apache License, Version 2.0 (the "License");
   # you may not use this file except in compliance with the License.
   # You may obtain a copy of the License at
   #
   #     https://www.apache.org/licenses/LICENSE-2.0
   #
   # Unless required by applicable law or agreed to in writing, software
   # distributed under the License is distributed on an "AS IS" BASIS,
   # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   # See the License for the specific language governing permissions and
   # limitations under the License.
   
   apiVersion: autoscaling/v2
   kind: HorizontalPodAutoscaler
   metadata:
     name: t5-inference
   spec:
     scaleTargetRef:
       apiVersion: apps/v1
       kind: Deployment
       name: t5-inference
     minReplicas: 1
     maxReplicas: 5
     metrics:
     - type: Pods
       pods:
         metric:
           name: prometheus.googleapis.com|ts_queue_latency_microseconds|counter
         target:
           type: AverageValue
           averageValue: "30000"
   

   The HorizontalPodAutoscaler scales the T5 model Pod quantity based on the cumulative duration of the request queue. Autoscaling is based on the ts_queue_latency_microseconds metric, which shows cumulative queue duration in microseconds.

3. Create the HorizontalPodAutoscaler:

   kubectl apply -f kubernetes/hpa.yaml
   

Verify autoscaling using a load generator

To test your autoscaling configuration, generate load for the serving application. This tutorial uses a Locust load generator to send requests to the prediction endpoint for the model.

1. Create the load generator:

   kubectl apply -f kubernetes/loadgenerator.yaml
   

   Wait for the load generator Pods to become ready.

2. Expose the load generator web interface locally:

   kubectl port-forward svc/loadgenerator 8080
   

   If you see an error message, try again when the Pod is running.

3. In a browser, open the load generator web interface:

   • If you're using a local shell, open a browser and go to http://127.0.0.1:8080.
   • If you're using Cloud Shell, click Web preview, and then click Change port. Enter port 8080.

4. Click the Charts tab to observe performance over time.

5. Open a new terminal window and watch the replica count of your horizontal Pod autoscalers:

   kubectl get hpa -w
   

   The number of replicas increases as the load increases. The scaleup might take approximately ten minutes. As new replicas start, the number of successful requests in the Locust chart increases.

   NAME           REFERENCE                 TARGETS           MINPODS   MAXPODS   REPLICAS   AGE
   t5-inference   Deployment/t5-inference   71352001470m/7M   1         5        1           2m11s
   

Recommendations

• Build your model with the same version of the base Docker image that you'll use for serving.
• If your model has special package dependencies, or if the size of your dependencies is large, create a custom version of your base Docker image.
• Watch the tree version of your model dependency packages. Ensure that your package dependencies support each others' versions. For example, Panda version 2.0.3 supports NumPy version 1.20.3 and later.
• Run GPU-intensive models on GPU nodes and CPU-intensive models on CPU. This could improve the stability of model serving and ensures that you're efficiently consuming node resources.