Gestisci un LLM utilizzando TPU Trillium (v6e) su GKE con vLLM

Creare un cluster GKE

Puoi pubblicare LLM su TPU in un cluster GKE Autopilot o Standard. Ti consigliamo di utilizzare un cluster Autopilot per un'esperienza Kubernetes completamente gestita. Per scegliere la modalità operativa GKE più adatta ai tuoi carichi di lavoro, consulta Scegliere una modalità operativa GKE.

Configura kubectl per comunicare con il cluster

Per configurare kubectl in modo che comunichi con il cluster, esegui il seguente comando:

  gcloud container clusters get-credentials CLUSTER_NAME --location=ZONE

Crea un secret di Kubernetes per le credenziali di Hugging Face

1. Crea uno spazio dei nomi. Puoi saltare questo passaggio se utilizzi lo spazio dei nomi default:

   kubectl create namespace NAMESPACE
   

2. Crea un secret di Kubernetes contenente il token di Hugging Face, esegui il seguente comando:

   kubectl create secret generic hf-secret \
       --from-literal=hf_api_token=HUGGING_FACE_TOKEN \
       --namespace NAMESPACE
   

Crea un bucket Cloud Storage

In Cloud Shell, esegui questo comando:

gcloud storage buckets create gs://GSBUCKET \
    --uniform-bucket-level-access

Viene creato un bucket Cloud Storage per archiviare i file del modello che hai scaricato da Hugging Face.

Configura un account di servizio Kubernetes per accedere al bucket

1. Crea l'account di servizio Kubernetes:

   kubectl create serviceaccount KSA_NAME --namespace NAMESPACE
   

2. Concedi l'accesso in lettura/scrittura all'account di servizio Kubernetes per accedere al bucket Cloud Storage:

   gcloud storage buckets add-iam-policy-binding gs://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. In alternativa, puoi concedere l'accesso in lettura e scrittura a tutti i bucket Cloud Storage del progetto:

   gcloud projects add-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 crea le seguenti risorse per l'LLM:

   1. Un bucket Cloud Storage per archiviare il modello scaricato e la cache di compilazione. Un driver CSI di Cloud Storage FUSE legge i contenuti del bucket.
   2. Volumi con la memorizzazione nella cache dei file abilitata e la funzionalità di download parallelo di Cloud Storage FUSE.
   Best practice:

   Utilizza una cache di file basata su tmpfs o Hyperdisk / Persistent Disk a seconda delle dimensioni previste dei contenuti del modello, ad esempio i file di peso. In questo tutorial utilizzi la cache dei file Cloud Storage FUSE basata su RAM.

Crea ed esegui il deployment dell'immagine TPU

Contieni il server vLLM:

1. Clona il repository vLLM e crea l'immagine:

   git clone https://github.com/vllm-project/vllm && cd vllm && git reset --hard cd34029e91ad2d38a58d190331a65f9096c0b157 && docker build -f Dockerfile.tpu . -t vllm-tpu
   

2. Esegui il push dell'immagine ad Artifact Registry:

   gcloud artifacts repositories create vllm-tpu --repository-format=docker --location=REGION_NAME && \
   gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
   docker image tag vllm-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-tpu/vllm-tpu:latest && \
   docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-tpu/vllm-tpu:latest
   

Esegui il deployment del server di modelli vLLM

Per eseguire il deployment del server di modelli vLLM:

1. Controlla il file manifest del deployment salvato come vllm-llama3-70b.yaml. Un deployment è un oggetto dell'API Kubernetes che ti consente di eseguire più repliche di pod distribuite tra i nodi di un cluster:

   
   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: vllm-tpu
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: vllm-tpu
     template:
       metadata:
         labels:
           app: vllm-tpu
         annotations:
           gke-gcsfuse/volumes: "true"
           gke-gcsfuse/cpu-limit: "0"
           gke-gcsfuse/memory-limit: "0"
           gke-gcsfuse/ephemeral-storage-limit: "0"
       spec:
         serviceAccountName: KSA_NAME
         nodeSelector:
           cloud.google.com/gke-tpu-topology: 2x4
           cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
         containers:
         - name: vllm-tpu
           image: IMAGE_NAME
           command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
           args:
           - --host=0.0.0.0
           - --port=8000
           - --tensor-parallel-size=8
           - --max-model-len=8192
           - --model=meta-llama/Meta-Llama-3.1-70B
           - --download-dir=/data
           env: 
           - name: HUGGING_FACE_HUB_TOKEN
             valueFrom:
               secretKeyRef:
                 name: hf-secret
                 key: hf_api_token
           - name: VLLM_XLA_CACHE_PATH
             value: "/data"
           ports:
           - containerPort: 8000
           resources:
             limits:
               google.com/tpu: 8
           readinessProbe:
             tcpSocket:
               port: 8000
             initialDelaySeconds: 15
             periodSeconds: 10
           volumeMounts:
           - name: gcs-fuse-csi-ephemeral
             mountPath: /data
           - name: dshm
             mountPath: /dev/shm
         volumes:
         - name: gke-gcsfuse-cache
           emptyDir:
             medium: Memory
         - name: dshm
           emptyDir:
             medium: Memory
         - name: gcs-fuse-csi-ephemeral
           csi:
             driver: gcsfuse.csi.storage.gke.io
             volumeAttributes:
               bucketName: GSBUCKET
               mountOptions: "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: v1
   kind: Service
   metadata:
     name: vllm-service
   spec:
     selector:
       app: vllm-tpu
     type: LoadBalancer	
     ports:
       - name: http
         protocol: TCP
         port: 8000  
         targetPort: 8000
   

2. Applica il manifest eseguendo il seguente comando:

   kubectl apply -f vllm-llama3-70b.yaml -n NAMESPACE
   

3. Visualizza i log del server del modello in esecuzione:

   kubectl logs -f -l app=vllm-tpu -n NAMESPACE
   

   L'output dovrebbe essere simile al seguente:

   INFO 09-20 19:03:48 launcher.py:19] Available routes are:
   INFO 09-20 19:03:48 launcher.py:27] Route: /openapi.json, Methods: GET, HEAD
   INFO 09-20 19:03:48 launcher.py:27] Route: /docs, Methods: GET, HEAD
   INFO 09-20 19:03:48 launcher.py:27] Route: /docs/oauth2-redirect, Methods: GET, HEAD
   INFO 09-20 19:03:48 launcher.py:27] Route: /redoc, Methods: GET, HEAD
   INFO 09-20 19:03:48 launcher.py:27] Route: /health, Methods: GET
   INFO 09-20 19:03:48 launcher.py:27] Route: /tokenize, Methods: POST
   INFO 09-20 19:03:48 launcher.py:27] Route: /detokenize, Methods: POST
   INFO 09-20 19:03:48 launcher.py:27] Route: /v1/models, Methods: GET
   INFO 09-20 19:03:48 launcher.py:27] Route: /version, Methods: GET
   INFO 09-20 19:03:48 launcher.py:27] Route: /v1/chat/completions, Methods: POST
   INFO 09-20 19:03:48 launcher.py:27] Route: /v1/completions, Methods: POST
   INFO 09-20 19:03:48 launcher.py:27] Route: /v1/embeddings, Methods: POST
   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)
   (RayWorkerWrapper pid=25987) INFO 09-20 19:03:46 tpu_model_runner.py:290] Compilation for decode done in 202.93 s.
   (RayWorkerWrapper pid=25987) INFO 09-20 19:03:46 tpu_model_runner.py:283] batch_size: 256, seq_len: 1 [repeated 7x across cluster]
   INFO 09-20 19:03:53 metrics.py:351] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 0.0 tokens/s, Running: 0 reqs, Swapped: 0 reqs, Pending: 0 reqs, GPU KV cache usage: 0.0%, CPU KV cache usage: 0.0%.
   

Pubblica il modello

1. Esegui questo comando per recuperare l'indirizzo IP esterno del servizio:

   export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
   

2. In un nuovo terminale, interagisci con il modello utilizzando curl:

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

   L'output dovrebbe essere simile al seguente:

   {"id":"cmpl-6b4bb29482494ab88408d537da1e608f","object":"text_completion","created":1727822657,"model":"meta-llama/Meta-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}}
   

Configurare l'autoscalatore personalizzato

In questa sezione, configuri scalabilità automatica orizzontale dei pod utilizzando le metriche Prometheus personalizzate. Utilizzi le metriche di Google Cloud Managed Service per Prometheus dal server vLLM.

Per saperne di più, consulta Google Cloud Managed Service per Prometheus. Questa opzione dovrebbe essere attiva per impostazione predefinita nel cluster GKE.

1. Configura l'adattatore Stackdriver delle metriche personalizzate nel tuo cluster:

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

2. Aggiungi il ruolo Visualizzatore di monitoraggio all'account di servizio utilizzato dall'adattatore Stackdriver per le metriche personalizzate:

   gcloud projects add-iam-policy-binding projects/PROJECT_ID \
       --role roles/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
   

3. Salva il seguente manifest come vllm_pod_monitor.yaml:

   
   apiVersion: monitoring.googleapis.com/v1
   kind: PodMonitoring
   metadata:
    name: vllm-pod-monitoring
   spec:
    selector:
      matchLabels:
        app: vllm-tpu
    endpoints:
    - path: /metrics
      port: 8000
      interval: 15s
   

4. Applicalo al cluster:

   kubectl apply -f vllm_pod_monitor.yaml
   

Creare un carico sull'endpoint vLLM

Crea un carico sul server vLLM per testare la scalabilità automatica di GKE con una metrica vLLM personalizzata.

1. Esegui uno script bash (load.sh) per inviare N richieste parallele all'endpoint vLLM:

   #!/bin/bash
   N=PARALLEL_PROCESSES
   export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
   for i in $(seq 1 $N); do
     while true; do
       curl http://$vllm_service:8000/v1/completions -H "Content-Type: application/json" -d '{"model": "meta-llama/Meta-Llama-3.1-70B", "prompt": "Write a story about san francisco", "max_tokens": 100, "temperature": 0}'
     done &  # Run in the background
   done
   wait
   

   Sostituisci PARALLEL_PROCESSES con il numero di procedimenti paralleli che vuoi eseguire.

2. Esegui lo script bash:

   nohup ./load.sh &
   

Verifica che Google Cloud Managed Service per Prometheus importi le metriche

Dopo che Google Cloud Managed Service per Prometheus ha estratto le metriche e hai aggiunto il carico all'endpoint vLLM, puoi visualizzare le metriche in Cloud Monitoring.

1. Nella console Google Cloud, vai alla pagina Esplora metriche.

   Vai a Esplora metriche

2. Fai clic su < > PromQL.

3. Inserisci la seguente query per osservare le metriche sul traffico:

   vllm:avg_generation_throughput_toks_per_s{cluster='CLUSTER_NAME_HERE'}
   

Nel grafico a linee, la metrica vLLM aumenta da 0 (precaricamento) a un valore (postcaricamento). Questo grafico conferma che le metriche vLLM vengono importate in Google Cloud Managed Service per Prometheus.

L'immagine seguente è un esempio di grafico dopo l'esecuzione dello script di caricamento. In questo caso, il server del modello pubblica circa 2000 token di generazione al secondo.

Esegui il deployment della configurazione di Horizontal Pod Autoscaler

In questa sezione esegui il deployment della configurazione di Horizontal Pod Autoscaler.

1. Salva il seguente manifest come vllm-hpa.yaml:

   
   apiVersion: autoscaling/v2
   kind: HorizontalPodAutoscaler
   metadata:
    name: vllm-hpa
   spec:
    scaleTargetRef:
      apiVersion: apps/v1
      kind: Deployment
      name: vllm-tpu
    minReplicas: 1
    maxReplicas: 2
    metrics:
      - type: Pods
        pods:
          metric:
            name: prometheus.googleapis.com|vllm:num_requests_waiting|gauge
          target:
            type: AverageValue
            averageValue: 10
   

   Le metriche vLLM in Google Cloud Managed Service per Prometheus seguono il formato vllm:metric_name.

   Best practice:

   Utilizza num_requests_waiting per scalare la velocità in bit. Utilizza gpu_cache_usage_perc per i casi d'uso TPU sensibili alla latenza.

2. Esegui il deployment della configurazione di Horizontal Pod Autoscaler:

   kubectl apply -f vllm-hpa.yaml
   

   GKE pianifica il deployment di un altro pod, che attiva il gestore della scalabilità automatica pool di nodi per aggiungere un secondo nodo prima di eseguire il deployment della seconda replica vLLM.

3. Monitora l'avanzamento della scalabilità automatica dei pod:

   kubectl get hpa --watch
   

   L'output è simile al seguente:

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

4. Attendi 10 minuti e ripeti i passaggi descritti nella sezione Verificare che Google Cloud Managed Service per Prometheus importi le metriche. Google Cloud Managed Service per Prometheus acquisisce le metriche da entrambi gli endpoint vLLM.