GKE에서 Ray를 사용하여 멀티 호스트 TPU vLLM 추론 배포

1. 소개

이 Codelab에서는 Google Cloud TPU를 사용하여 Google Kubernetes Engine (GKE)에 고성능 멀티 호스트 vLLM (Virtual Large Language Model) 추론 서비스를 배포하는 방법을 알아봅니다. Ray를 사용하여 분산 추론을 구성하고 LeaderWorkerSets를 사용하여 GKE에서 워크로드를 기본적으로 관리합니다.

이 둘러보기에서는 Qwen 30B와 같은 대규모 모델을 제공하기 위한 프로덕션 설정을 시뮬레이션합니다.

실습할 내용

  • 가속기 트래픽을 위한 커스텀 VPC 네트워크를 만듭니다.
  • Ray Operator 및 GCS Fuse CSI 드라이버로 GKE 클러스터를 프로비저닝합니다.
  • 가속화된 모델 로드를 위해 GCS Rapid Cache를 초기화합니다.
  • 예약된 용량으로 멀티 호스트 TPU v6e 노드 풀을 프로비저닝합니다.
  • 모델 가중치에 안전하게 액세스할 수 있도록 워크로드 아이덴티티를 구성합니다.
  • 30B 매개변수 모델을 제공하는 vLLM 엔진을 배포하고 테스트합니다.

필요한 항목

  • 결제가 사용 설정된 Google Cloud 프로젝트.
  • TPU v6e 리소스 (칩 32개, ct6e-standard-4t)를 위한 Google Cloud 예약.
  • 소스 버킷에서 모델 가중치를 복사할 수 있는 액세스 권한.
  • gcloud, kubectl, helm이 설치된 Cloud Shell 또는 로컬 터미널.
  • 예상 소요 시간: 60분
  • 예상 비용: 60달러 미만 (즉시 삭제하는 경우)

2. 시작하기 전에

Google Cloud 프로젝트 만들기 또는 선택

  1. Google Cloud 콘솔에서 Google Cloud 프로젝트를 생성하거나 선택합니다.
  2. Cloud 프로젝트에 결제가 사용 설정되어 있는지 확인합니다.

Cloud Shell 시작

  1. Google Cloud 콘솔 상단에서 Cloud Shell 활성화 를 클릭합니다.
  2. 인증 확인
gcloud auth list
  1. 프로젝트 확인하기
gcloud config get project
  1. 필요한 경우 설정합니다.
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

환경 변수 설정

명령어 실행을 더 쉽게 하려면 셸에서 다음 변수를 정의하세요. <YOUR_ZONE>을 할당된 TPU 영역으로 바꾸고 <YOUR_RESERVATION_NAME>을 예약 ID로 바꿉니다. 게이트 모델 가중치를 다운로드하려면 Hugging Face 사용자 액세스 토큰을 만들어야 합니다. 만들었으면 <YOUR_HUGGING_FACE_TOKEN>을 새로 만든 토큰으로 바꿉니다.

export PROJECT_ID=$(gcloud config get-value project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export ZONE="<YOUR_ZONE>" # e.g., us-east5-a
export REGION=${ZONE%-*}
export CLUSTER_NAME="qwen-serving-cluster"
export GVNIC_NETWORK_PREFIX="qwen-serving"
export BUCKET_NAME="inf-demo-model-storage-${PROJECT_NUMBER}"
export RESERVATION_NAME="<YOUR_RESERVATION_NAME>"
export NODE_POOL_NAME="tpu-v6e-32-resvd-pool"
export MULTIHOST_COLLECTION_NAME="tpu-6-collection"
export HF_TOKEN="<YOUR_HUGGING_FACE_TOKEN>" # Token with access to Qwen model if restricted

API 사용 설정

필요한 Google Cloud 서비스를 사용 설정합니다.

gcloud services enable \
    container.googleapis.com \
    compute.googleapis.com \
    iam.googleapis.com \
    cloudresourcemanager.googleapis.com

3. 커스텀 네트워킹 만들기

멀티 호스트 TPU 워크로드에는 효율적인 가속기 통신을 위한 더 높은 MTU 크기를 비롯한 특정 네트워크 구성이 필요합니다. 클러스터의 커스텀 VPC 네트워크를 만듭니다.

  1. 큰 MTU (8896)로 VPC 네트워크를 만듭니다.
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. 클러스터의 서브넷을 만듭니다.
    gcloud compute --project=${PROJECT_ID} \
    networks subnets create ${GVNIC_NETWORK_PREFIX}-tpu \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --region=${REGION} \
    --range=192.168.100.0/24
  3. 작업자가 통신할 수 있도록 내부 트래픽을 허용하는 방화벽 규칙을 만듭니다.
    gcloud compute --project=${PROJECT_ID} firewall-rules create ${GVNIC_NETWORK_PREFIX}-allow-internal \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --allow=all \
    --source-ranges=172.16.0.0/12,192.168.0.0/16,10.0.0.0/8 \
    --description="Allow all internal traffic within the network."

4. GKE 클러스터 프로비저닝

GCS Fuse 마운트 및 Ray Operator 워크로드를 지원하도록 구성된 표준 GKE 클러스터 설정을 만듭니다.

  1. 클러스터를 만듭니다:
    gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${REGION} \
    --release-channel=rapid \
    --machine-type=e2-standard-4 \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --subnetwork=${GVNIC_NETWORK_PREFIX}-tpu \
    --num-nodes=1 \
    --gateway-api=standard \
    --enable-managed-prometheus \
    --enable-dataplane-v2 \
    --enable-dataplane-v2-metrics \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --addons=GcsFuseCsiDriver,RayOperator \
    --enable-ip-alias
  2. 클러스터 사용자 인증 정보를 검색합니다:
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Hugging Face 보안 비밀을 만듭니다.
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. Helm을 통해 LeaderWorkerSet (LWS)를 설치합니다. LWS는 함께 예약해야 하는 포드 그룹을 관리합니다.
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. GCS Rapid Cache 사용 설정

서빙 중에 Cloud Storage에서 수십 GB의 가중치를 읽는 속도를 높이려면 GCS 버킷을 만들고 영역에서 GCS Rapid Cache를 사용 설정합니다.

  1. 버킷을 만듭니다:
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. TPU 영역에서 Rapid Cache를 초기화합니다.
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. 워크로드 아이덴티티 및 스토리지 권한 설정

장기 실행 키를 삽입하지 않고 가중치 버킷을 GKE 포드에 안전하게 마운트하도록 아이덴티티 링크를 구성합니다.

  1. 전용 IAM 서비스 계정을 만듭니다:
    gcloud iam service-accounts create tpu-reader-sa
  2. 버킷 읽기 권한을 부여합니다:
    gcloud storage buckets add-iam-policy-binding gs://${BUCKET_NAME} \
    --member="serviceAccount:tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role="roles/storage.objectAdmin"
  3. default 네임스페이스 Kubernetes 서비스 계정에 워크로드 아이덴티티 결합을 만듭니다.
    gcloud iam service-accounts add-iam-policy-binding tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com \
    --role="roles/iam.workloadIdentityUser" \
    --member="serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"
  4. Kubernetes SA에 주석을 추가합니다:
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. 모델 가중치 설정

30B 매개변수 모델을 제공하려면 Hugging Face에서 GCS 버킷으로 가중치를 다운로드해야 합니다. Cloud Shell 디스크 할당량 한도 (5GB)를 무시하려면 표준 Kubernetes 작업 을 사용하여 클러스터 내부에서 직접 다운로드하고 마운트된 GCS Fuse 볼륨에 안전하게 씁니다.

  1. 모델 다운로더 작업을 배포합니다.
    cat <<EOF | kubectl apply -f -
apiVersion: batch/v1
kind: Job
metadata:
  name: model-downloader
spec:
  ttlSecondsAfterFinished: 60
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
        gke-gcsfuse/memory-limit: "0"
    spec:
      serviceAccountName: default
      restartPolicy: OnFailure
      containers:
      - name: downloader
        image: python:3.10-slim
        command: ["/bin/sh", "-c"]
        args:
        - |
          pip install -U "huggingface_hub[hf_transfer]" filelock
          export HF_HUB_ENABLE_HF_TRANSFER=1

          python -c '
          import filelock

          class DummyLock:
              def __init__(self, *args, **kwargs): pass
              def __enter__(self): return self
              def __exit__(self, *args): pass
              def acquire(self, *args, **kwargs): pass
              def release(self, *args, **kwargs): pass

          filelock.FileLock = DummyLock

          from huggingface_hub import snapshot_download
          snapshot_download(
              repo_id="Qwen/Qwen3-30B-A3B", 
              local_dir="/models/qwen3-30b-weights",
              local_dir_use_symlinks=False
          )
          '
        env:
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        volumeMounts:
        - name: model-weights
          mountPath: /models
      volumes:
      - name: model-weights
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: ${BUCKET_NAME}
            mountOptions: "implicit-dirs"
EOF
  2. 다운로드 모니터링: 다운로더 포드의 로그를 확인하여 진행 상황을 확인합니다.
    kubectl logs -f job/model-downloader
    작업이 성공 상태로 완료될 때까지 기다립니다.

8. 예약된 TPU 노드 풀 만들기

기존 용량 예약을 사용하여 실제 멀티 호스트 TPU 슬라이스를 프로비저닝합니다.

  1. 생성 명령어를 실행합니다:
    gcloud beta container node-pools create ${NODE_POOL_NAME} \
    --project=${PROJECT_ID} \
    --cluster=${CLUSTER_NAME} \
    --region=${REGION} \
    --node-locations=${ZONE} \
    --machine-type=ct6e-standard-4t \
    --tpu-topology=4x8 \
    --num-nodes=8 \
    --scopes=https://www.googleapis.com/auth/cloud-platform \
    --reservation-affinity=specific \
    --reservation=${RESERVATION_NAME} \
    --accelerator-network-profile=auto \
    --node-labels=cloud.google.com/gke-nodepool-group-name=${MULTIHOST_COLLECTION_NAME} \
    --node-labels=cloud.google.com/gke-workload-type=HIGH_AVAILABILITY \
    --node-labels=cloud.google.com/gke-networking-dra-driver=true
  2. 노드가 조인될 때까지 기다립니다. ct6e가 포함된 노드 8개가 kubectl get nodes에 조인될 때까지 기다립니다.

9. vLLM 서비스 배포

  1. 네트워크 클레임 만들기: 네트워크 환경을 요청해야 합니다.
    cat <<EOF | kubectl apply -f -
apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: all-netdev
spec:
  spec:
    devices:
      requests:
      - name: req-netdev
        exactly:
          deviceClassName: netdev.google.com
          allocationMode: All
EOF
  2. 부하 분산기 API 엔드포인트 배포:
    cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: Service
metadata:
  name: vllm-tpu-service
spec:
  type: LoadBalancer
  selector:
    leaderworkerset.sigs.k8s.io/name: vllm-tpu-qwen
    leaderworkerset.sigs.k8s.io/worker-index: "0"
  ports:
  - protocol: TCP
    port: 8000
    targetPort: 8000
EOF
  3. LeaderWorkerSet 워크로드 배포: 이 매니페스트는 8개의 슬라이스 호스트에서 Ray 헤드/작업자 집계를 동적으로 시작합니다.
    cat <<EOF | kubectl apply -f -
apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: vllm-tpu-qwen
spec:
  replicas: 1
  leaderWorkerTemplate:
    size: 8
    restartPolicy: RecreateGroupOnPodRestart
    workerTemplate:
      metadata:
        annotations:
          gke-gcsfuse/volumes: "true"
          gke-gcsfuse/memory-limit: "0"
        labels:
          leaderworkerset.sigs.k8s.io/name: vllm-tpu-qwen
          gke-gcsfuse/volumes: "true"
      spec:
        hostname: vllm-tpu-qwen
        serviceAccountName: default
        containers:
        - name: vllm-tpu
          image: vllm/vllm-tpu:nightly
          command: ["sh", "-c"]
          args:
          - |
            MY_TPU_IP=\$(hostname -I | awk '{print \$1}')
            echo "My TPU Network IP is: \$MY_TPU_IP"

            LEADER_DNS="vllm-tpu-qwen-0.vllm-tpu-qwen"
            until getent hosts \$LEADER_DNS; do
              echo "DNS not ready. Sleeping 5s..."
              sleep 5
              done
            LEADER_IP=\$(getent hosts \$LEADER_DNS | awk '{print \$1}')

            export JAX_PLATFORMS=''
            export SCAN_TPU_CHIPS=True
            export TPU_MULTIHOST_BACKEND=ray
            export JAX_DISTRIBUTED_INITIALIZATION_TIMEOUT=300
            export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/lib
            export VLLM_HOST_IP=\$MY_TPU_IP

            if [ "\$LWS_WORKER_INDEX" = "0" ]; then
              echo "Starting Ray Head..."
              ray start --head --port=6379 --node-ip-address=\$MY_TPU_IP --resources='{"TPU": 4}' --block &
              sleep 20
              until ray status; do sleep 5; done

              echo "Starting vLLM API Server..."
              python3 -m vllm.entrypoints.openai.api_server \
                --model=/models/qwen3-30b-weights \
                --tensor-parallel-size=32 \
                --pipeline-parallel-size=1 \
                --distributed-executor-backend=ray \
                --host=0.0.0.0 --port=8000 \
                --enforce-eager \
                --gpu-memory-utilization=0.90
            else
              ray start --address=\${LEADER_IP}:6379 --node-ip-address=\$MY_TPU_IP --resources='{"TPU": 4}' --block
            fi
          ports:
          - containerPort: 8000
          - containerPort: 6379
          volumeMounts:
          - name: model-weights
            mountPath: /models
            readOnly: true
          - name: dshm
            mountPath: /dev/shm
          resources:
            claims:
            - name: net-resources
            limits:
              google.com/tpu: 4
              memory: "100Gi"
            requests:
              google.com/tpu: 4
              memory: "100Gi"
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x8
          gke.networks.io/accelerator-network-profile: auto
        resourceClaims:
        - name: net-resources
          resourceClaimTemplateName: all-netdev
        volumes:
        - name: model-weights
          csi:
            driver: gcsfuse.csi.storage.gke.io
            volumeAttributes:
              bucketName: ${BUCKET_NAME}
              mountOptions: "implicit-dirs"
        - name: dshm
          emptyDir:
            medium: Memory
EOF

10. 테스트 배포 응답

LeaderWorkerSet의 모든 포드가 컨테이너 이미지를 가져오고 Ray를 초기화하고 완전히 Ready 상태가 되는 데 5~10분이 걸릴 수 있습니다. 포드 초기화를 확인하여 상태를 추적할 수 있습니다.

kubectl get pods -l leaderworkerset.sigs.k8s.io/name=vllm-tpu-qwen -w

모든 8개의 vllm-tpu-qwen- 포드가 STATUS를 Running으로, READY를 2/2로 표시하고 부하 분산기가 외부 IP를 수신했는지 확인한 후 계속합니다. 7~10분이 걸릴 수 있습니다.

  1. 외부 IP 검색: 
    export EXTERNAL_IP=$(kubectl get svc vllm-tpu-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
echo $EXTERNAL_IP

주의: 프로덕션 서비스에서는 이 엔드포인트를 Identity Aware Proxy (IAP)와 같은 것으로 보호해야 합니다.

  1. curl을 사용하여 추론 요청을 제출합니다.
        curl -N -s http://$EXTERNAL_IP:8000/v1/chat/completions \
        -H "Content-Type: application/json" \
        -d '{
            "model": "/models/qwen3-30b-weights",
            "messages": [{"role": "user", "content": "Write a haiku about high-performance computing on TPUs."}],
            "temperature": 0.7,
            "max_tokens": 100,
            "stream": true
        }' | sed 's/^data: //' | grep -v '\[DONE\]' | grep -v '^$' | jq -rj '.choices[0].delta.content // empty' ; echo ""
    생성된 추론을 나타내는 텍스트가 포함된 JSON 응답과 유사한 출력이 표시됩니다.

11. 삭제

Google Cloud 계정에 지속적으로 비용이 청구되지 않도록 하려면 이 Codelab 중에 만든 리소스를 삭제하세요.

  1. 노드 풀 삭제:
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. 클러스터 삭제:
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. 네트워크 및 방화벽 설정 삭제:
    gcloud compute firewall-rules delete \
    "${GVNIC_NETWORK_PREFIX}-allow-internal" \
    --project="${PROJECT_ID}" --quiet

gcloud compute networks subnets delete "${GVNIC_NETWORK_PREFIX}-tpu" \
    --region="${REGION}" --quiet

gcloud compute networks delete "${GVNIC_NETWORK_PREFIX}-main" --quiet
  4. 서비스 계정 결합 해제 및 삭제:
        # 1. Create the cleanup script
    cat << 'EOF' > clean_up_sa.sh
    #!/bin/bash

    # Validate that PROJECT_ID is available
    if [ -z "$PROJECT_ID" ]; then
      echo "Error: PROJECT_ID environment variable is not set."
      exit 1
    fi

    SA_EMAIL="tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com"
    SA_MEMBER="serviceAccount:${SA_EMAIL}"

    echo "Gathering IAM policy for ${SA_EMAIL}..."

    # Fetch roles assigned to this specific SA
    ROLES=$(gcloud projects get-iam-policy ${PROJECT_ID} \
        --flatten="bindings[].members" \
        --filter="bindings.members:${SA_MEMBER}" \
        --format="value(bindings.role)")

    if [ -z "$ROLES" ]; then
        echo "No IAM bindings found for this service account."
    else
        for ROLE in $ROLES; do
            echo "Removing binding for: ${ROLE}..."
            gcloud projects remove-iam-policy-binding ${PROJECT_ID} \
                --member="${SA_MEMBER}" \
                --role="${ROLE}" --quiet > /dev/null
        done
        echo "Successfully unbound all roles."
    fi

    # 2. Delete the service account itself
    echo "Deleting service account..."
    gcloud iam service-accounts delete ${SA_EMAIL} --project=${PROJECT_ID} --quiet

    echo "Cleanup complete."
    EOF

    # 2. Make the script executable and run it
    chmod +x clean_up_sa.sh
    ./clean_up_sa.sh
  5. GCS 버킷 삭제 Cloud 콘솔로 이동하여 Cloud Storage -> 버킷을 선택하고 inf-demo-model-storage를 선택한 후 '삭제'를 선택합니다.

12. 축하합니다

축하합니다. Google Kubernetes Engine을 통해 Ray를 기본적으로 활용하는 멀티 호스트 TPU 고추론 속도 vLLM 스택을 배포했습니다.

학습한 내용

  • 고속 TPU 트래픽에 맞게 맞춤설정된 맞춤 경로 프로비저닝
  • GCS Fuse 및 리전별 Rapid Cache를 활용한 가중치 마운트
  • LeaderWorkerSets를 통해 기본적으로 동기화된 멀티 호스트 워크로드 슬라이스 조정
  • 자세한 내용은 vLLM 사용자 가이드 및 llm-d 배포 가이드를 참고하세요.