Create a GKE cluster

You can serve LLMs on TPUs in a GKE Autopilot or Standard cluster. We recommend that you use a Autopilot cluster for a fully managed Kubernetes experience. To choose the GKE mode of operation that's the best fit for your workloads, seeChoose a GKE mode of operation.

Configure kubectl to communicate with your cluster

To configure kubectl to communicate with your cluster, run the followingcommand:

gcloudcontainerclustersget-credentials${CLUSTER_NAME}--location=${CONTROL_PLANE_LOCATION}

Create a Kubernetes Secret for Hugging Face credentials

1. Create a namespace. You can skip this step if you are using thedefault namespace:

   kubectlcreatenamespace${NAMESPACE}

2. Create a Kubernetes Secret that contains the Hugging Face token, run thefollowing command:

   kubectlcreatesecretgenerichf-secret\--from-literal=hf_api_token=${HF_TOKEN}\--namespace${NAMESPACE}

Create a Cloud Storage bucket

In Cloud Shell, run the following command:

gcloudstoragebucketscreategs://${GSBUCKET}\--uniform-bucket-level-access

This creates a Cloud Storage bucket to store the model files youdownload from Hugging Face.

Set up a Kubernetes ServiceAccount to access the bucket

1. Create the Kubernetes ServiceAccount:

   kubectlcreateserviceaccount${KSA_NAME}--namespace${NAMESPACE}

2. Grant read-write access to the Kubernetes ServiceAccount in order to access the Cloud Storage bucket:

   gcloudstoragebucketsadd-iam-policy-bindinggs://${GSBUCKET}\--member"principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}"\--role"roles/storage.objectUser"

3. Alternatively, you can grant read-write access to all Cloud Storagebuckets in the project:

   gcloudprojectsadd-iam-policy-binding${PROJECT_ID}\--member"principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}"\--role"roles/storage.objectUser"

   GKE creates the following resources for the LLM:

   1. A Cloud Storage bucket to store the downloaded model and thecompilation cache. ACloud Storage FUSE CSI driverreads the content of the bucket.
   2. Volumes with file caching enabled and theparallel download feature of Cloud Storage FUSE.
   Best practice:

   Use a file cache backed bytmpfsorHyperdisk / Persistent Diskdepending on the expected size of the model contents, for example, weight files.In this tutorial, you use Cloud Storage FUSE file cache backed by RAM.

Deploy vLLM model server

To deploy the vLLM model server, this tutorial uses a Kubernetes Deployment. ADeployment is a Kubernetes API object that lets you run multiple replicas of Pods that are distributed among the nodes in a cluster..

1. Inspect the following Deployment manifest saved asvllm-llama3-70b.yaml, that uses a single replica:

   apiVersion:apps/v1kind:Deploymentmetadata:name:vllm-tpuspec:replicas:1selector:matchLabels:app:vllm-tputemplate:metadata:labels:app:vllm-tpuannotations:gke-gcsfuse/volumes:"true"gke-gcsfuse/cpu-limit:"0"gke-gcsfuse/memory-limit:"0"gke-gcsfuse/ephemeral-storage-limit:"0"spec:serviceAccountName:KSA_NAMEnodeSelector:cloud.google.com/gke-tpu-topology:2x4cloud.google.com/gke-tpu-accelerator:tpu-v6e-slicecontainers:-name:vllm-tpuimage:vllm/vllm-tpu:latestcommand:["python3","-m","vllm.entrypoints.openai.api_server"]args:---host=0.0.0.0---port=8000---tensor-parallel-size=8---max-model-len=4096---model=meta-llama/Llama-3.1-70B---download-dir=/data---max-num-batched-tokens=512---max-num-seqs=128env:-name:HUGGING_FACE_HUB_TOKENvalueFrom:secretKeyRef:name:hf-secretkey:hf_api_token-name:VLLM_XLA_CACHE_PATHvalue:"/data"-name:VLLM_USE_V1value:"1"ports:-containerPort:8000resources:limits:google.com/tpu:8readinessProbe:tcpSocket:port:8000initialDelaySeconds:15periodSeconds:10volumeMounts:-name:gcs-fuse-csi-ephemeralmountPath:/data-name:dshmmountPath:/dev/shmvolumes:-name:gke-gcsfuse-cacheemptyDir:medium:Memory-name:dshmemptyDir:medium:Memory-name:gcs-fuse-csi-ephemeralcsi:driver:gcsfuse.csi.storage.gke.iovolumeAttributes:bucketName:GSBUCKETmountOptions:"implicit-dirs,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:100,file-cache:max-parallel-downloads:-1,file-cache:download-chunk-size-mb:10,file-cache:max-size-mb:-1"---apiVersion:v1kind:Servicemetadata:name:vllm-servicespec:selector:app:vllm-tputype:LoadBalancerports:-name:httpprotocol:TCPport:8000targetPort:8000

   If you scale up the Deployment to multiple replicas, the concurrent writes to theVLLM_XLA_CACHE_PATH will cause the error:RuntimeError: filesystem error: cannot create directories. To prevent this error, you have two options:

   1. Remove the XLA cache location by removing the following block from the Deployment YAML. This means all replicas will recompile the cache.

      - name: VLLM_XLA_CACHE_PATH  value: "/data"

   2. Scale the Deployment to1, and wait for the first replica to become ready and write to the XLA cache. Then scale to additional replicas. This allows the remainder of the replicas to read the cache, without attempting to write it.

2. Apply the manifest by running the following command:

   kubectlapply-fvllm-llama3-70b.yaml-n${NAMESPACE}

3. View the logs from the running model server:

   kubectllogs-f-lapp=vllm-tpu-n${NAMESPACE}

   The output should look similar to the following:

   INFO:     Started server process [1]INFO:     Waiting for application startup.INFO:     Application startup complete.INFO:     Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)

Serve the model

1. To get the external IP address of the VLLM service, run the following command:

   exportvllm_service=$(kubectlgetservicevllm-service-ojsonpath='{.status.loadBalancer.ingress[0].ip}'-n${NAMESPACE})

2. Interact with the model usingcurl:

   curlhttp://$vllm_service:8000/v1/completions\-H"Content-Type: application/json"\-d'{    "model": "meta-llama/Llama-3.1-70B",    "prompt": "San Francisco is a",    "max_tokens": 7,    "temperature": 0}'

   The output should be similar to the following:

   {"id":"cmpl-6b4bb29482494ab88408d537da1e608f","object":"text_completion","created":1727822657,"model":"meta-llama/Llama-3-8B","choices":[{"index":0,"text":" top holiday destination featuring scenic beauty and","logprobs":null,"finish_reason":"length","stop_reason":null,"prompt_logprobs":null}],"usage":{"prompt_tokens":5,"total_tokens":12,"completion_tokens":7}}

Set up the custom autoscaler

In this section, you set up horizontal Pod autoscaling using customPrometheus metrics. You use the Google Cloud Managed Service for Prometheus metrics from the vLLM server.

To learn more, seeGoogle Cloud Managed Service for Prometheus.This should be enabled by default on the GKE cluster.

1. Set up the Custom Metrics Stackdriver Adapter onyour cluster:

   kubectlapply-fhttps://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml

2. Add the Monitoring Viewer role to the service account that the Custom Metrics Stackdriver Adapter uses:

   gcloudprojectsadd-iam-policy-bindingprojects/${PROJECT_ID}\--roleroles/monitoring.viewer\--member=principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/custom-metrics/sa/custom-metrics-stackdriver-adapter

   Note: Ensure that the service account that your GKE clusteruses has theMonitoring Metric Writer role. This tutorialuses the default Compute Engine service account.

3. Save the following manifest asvllm_pod_monitor.yaml:

   apiVersion:monitoring.googleapis.com/v1kind:PodMonitoringmetadata:name:vllm-pod-monitoringspec:selector:matchLabels:app:vllm-tpuendpoints:-path:/metricsport:8000interval:15s

4. Apply it to the cluster:

   kubectlapply-fvllm_pod_monitor.yaml-n${NAMESPACE}

1. In the Google Cloud console, go to theMetrics explorer page.

   Go to Metrics explorer

2. Click< > PromQL.

3. Enter the following query to observe traffic metrics:

   vllm:num_requests_waiting{cluster='CLUSTER_NAME'}

A line graph shows your vLLM metric (num_requests_waiting) measured over time. The vLLM metric scales up from 0 (pre-load) to a value(post-load). This graph confirms your vLLM metrics are being ingested intoGoogle Cloud Managed Service for Prometheus. The following example graph shows a startingpre-load value of 0, which reaches a maximum post-load value of close to 400within one minute.

Deploy the Horizontal Pod Autoscaler configuration

When deciding which metric to autoscale on, we recommend the following metrics for vLLM TPU:

• num_requests_waiting: This metric relates to the number of requests waiting in the model server's queue. This number starts to noticeably grow when the kv cache is full.

• gpu_cache_usage_perc: This metric relates to the kv cache utilization, which directly correlates to the number of requests being processed for a given inference cycle on the model server. Note that this metric works the same on GPUs and TPUs, though it is tied to the GPU naming schema.

We recommend that you usenum_requests_waiting when optimizing for throughput and cost, and when your latency targets are achievable with your model server's maximum throughput.

We recommend that you usegpu_cache_usage_perc when you have latency-sensitive workloads where queue-based scaling isn't fast enough to meet your requirements.

For further explanation, check outBest practices for autoscaling large language model (LLM) inference workloads with TPUs.

When selecting anaverageValue target for your HPA config, you will have to determine this experimentally. Check out theSave on GPUs: Smarter autoscaling for your GKE inferencing workloads blog post for additional ideas on how to optimize this part. Theprofile-generator used in this blog post works for vLLM TPU as well.

In the following instructions, you deploy your HPA configuration by using thenum_requests_waiting metric. For demonstration purposes, you set the metric to a low value so that the HPA configuration scales your vLLM replicas to two. To deploy the Horizontal Pod Autoscaler configuration usingnum_requests_waiting, follow these steps:

1. Save the following manifest asvllm-hpa.yaml:

   apiVersion:autoscaling/v2kind:HorizontalPodAutoscalermetadata:name:vllm-hpaspec:scaleTargetRef:apiVersion:apps/v1kind:Deploymentname:vllm-tpuminReplicas:1maxReplicas:2metrics:-type:Podspods:metric:name:prometheus.googleapis.com|vllm:num_requests_waiting|gaugetarget:type:AverageValueaverageValue:10

   The vLLM metrics in Google Cloud Managed Service for Prometheus follow thevllm:metric_name format.

   Best practice:

   Usenum_requests_waiting for scaling throughput. Usegpu_cache_usage_perc for latency-sensitive TPU use cases.

2. Deploy the Horizontal Pod Autoscaler configuration:

   kubectlapply-fvllm-hpa.yaml-n${NAMESPACE}

   GKE schedules another Pod to deploy,which triggers the node pool autoscaler to add a second node before it deploysthe second vLLM replica.

3. Watch the progress of the Pod autoscaling:

   kubectlgethpa--watch-n${NAMESPACE}

   The output is similar to the following:

   NAME       REFERENCE             TARGETS       MINPODS   MAXPODS   REPLICAS   AGEvllm-hpa   Deployment/vllm-tpu   <unknown>/10   1         2         0          6svllm-hpa   Deployment/vllm-tpu   34972m/10      1         2         1          16svllm-hpa   Deployment/vllm-tpu   25112m/10      1         2         2          31svllm-hpa   Deployment/vllm-tpu   35301m/10      1         2         2          46svllm-hpa   Deployment/vllm-tpu   25098m/10      1         2         2          62svllm-hpa   Deployment/vllm-tpu   35348m/10      1         2         2          77s

   Success: GKE scales up the vLLM server using custom Prometheus metrics.

4. Wait for 10 minutes and repeat the steps in theVerify thatGoogle Cloud Managed Service for Prometheus ingests the metrics section.Google Cloud Managed Service for Prometheus ingests the metrics from both vLLM endpoints now.