מילוי בקשות של מודלים פתוחים של Gemma באמצעות TPU מרובה מארחים ב-GKE עם Ray

במדריך הזה נסביר איך לפרוס שירות הסקת מסקנות TPU מרובה-מארחים באמצעות Ray Serve LLM. בעזרת התמיכה המובנית של Ray ב-TPU, אפשר לתזמן באופן אטומי עובדים של מנוע מבוזר על פני טופולוגיות מורכבות של מאיצים, וכך לפרוס מודלים גדולים על פלח TPU מרובה מארחים לצורך הסקה.

המדריך הזה מיועד למהנדסי למידת מכונה (ML), לאדמינים ולמפעילים של פלטפורמות ולמומחים בתחום הנתונים וה-AI שמעוניינים להשתמש ביכולות של Kubernetes לארגון קונטיינרים כדי להפעיל עומסי עבודה של AI/ML בפרוסות TPU מבוזרות עם כמה מארחים. מידע נוסף על תפקידים נפוצים ומשימות לדוגמה שמוזכרים ב Google Cloud תוכן זמין במאמר תפקידים נפוצים של משתמשים ומשימות ב-GKE.

לפני שקוראים את הדף הזה, חשוב לוודא שמכירים את הנושאים הבאים:

רקע

בקטע הזה מתוארות הטכנולוגיות העיקריות שמופיעות במדריך הזה.

TPUs

יחידות לעיבוד טנסורים (TPU) מאפשרות להאיץ עומסי עבודה ספציפיים שפועלים בצמתים, כמו למידת מכונה ועיבוד נתונים. היתרון העיקרי של יחידות TPU הוא ביצועים בקנה מידה נרחב. במדריך הזה משתמשים ב-TPU Trillium, הדור השישי של Cloud TPU. פרוסות TPU עם כמה מארחים מורכבות מכמה צמתים פיזיים שמתקשרים באמצעות חיבור מהיר בין שבבים (ICI), שמתאים מאוד להצגת נתונים עם תפוקה גבוהה וזמן אחזור נמוך.

‫vLLM on Ray

‫vLLM הוא מנוע להפעלת מודלים מסוג LLM עם תפוקה גבוהה ושימוש יעיל בזיכרון. השילוב עם Ray Serve מאפשר ל-vLLM להתרחב למספר מארחים ולגשת לטופולוגיות של חומרה פיזית באופן טבעי. במדריך הזה נדגים איך להשתמש בפריסות LLMConfig ו-LLMServer של Ray Serve כדי לתזמן הסקה של vLLM בפרוסות של כמה מארחים, וכך לאפשר למסגרת לטפל באופן אוטומטי בהפצה של טופולוגיה ובפיזור של קבוצות מיקום.

מטרות

המדריך הזה מספק בסיס להבנה ולבדיקה של פריסת LLM מעשית להסקת מסקנות בסביבת Kubernetes מנוהלת שמשתמשת ב-TPU מרובי-מארחים.

  1. מכינים את הסביבה עם אשכול GKE במצב Autopilot או Standard.
  2. יוצרים קובץ אימג' של קונטיינר בהתאמה אישית עם תלות מוטמעת.
  3. פריסת סקריפט Python של Ray LLM באשכול כדי לתזמן היקש של vLLM על פלח TPU.
  4. אפשר להשתמש ב-Ray LLM כדי להפעיל את מודל Gemma 4 דרך curl וממשק צ'אט אינטרנטי אופציונלי.

לפני שמתחילים

לפני שמתחילים, חשוב לוודא שביצעתם את הפעולות הבאות:

  • מפעילים את ממשק Google Kubernetes Engine API.
    • הפעלת Google Kubernetes Engine API
  • כדי להשתמש ב-CLI של Google Cloud למשימה הזו, צריך להתקין ואז להפעיל את gcloud CLI. אם התקנתם בעבר את ה-CLI של gcloud, מריצים את הפקודה gcloud components update כדי לקבל את הגרסה העדכנית. יכול להיות שגרסאות קודמות של ה-CLI של gcloud לא יתמכו בהרצת הפקודות שמופיעות במסמך הזה.
  • מוודאים שלפרויקט יש מכסת קיבולת מספקת ל-TPU Trillium ‏ (v6e) באזור שנבחר. מידע נוסף זמין במאמר בנושא מכסות של Cloud TPU.
  • מוודאים שאשכול GKE משתמש ב-GKE Dataplane V2 ועומד בדרישות הגרסה של DRANET: 1.35.2-gke.1842000 ואילך גם ב-Standard וגם ב-Autopilot.
  • צריך לוודא שיש לכם את התפקידים הבאים ב-IAM:
    • roles/container.admin
    • roles/iam.serviceAccountAdmin

הכנת הסביבה

במדריך הזה תשתמשו ב-Cloud Shell כדי לנהל משאבים שמתארחים ב- Google Cloud. ב-Cloud Shell מותקן מראש התוכנה שדרושה למדריך הזה, כולל kubectl ו-gcloud CLI.

כדי להגדיר את הסביבה באמצעות Cloud Shell:

  1. במסוף Google Cloud , מפעילים סשן של Cloud Shell על ידי לחיצה על Activate Cloud Shell (הפעלת Cloud Shell) כפתור הפעלת Shell. סשן יופעל בחלונית התחתונה של מסוף Google Cloud .

  2. יוצרים ומפעילים סביבה וירטואלית של Python:

    python3 -m venv ray-env
source ray-env/bin/activate

  3. מתקינים את Ray CLI:

    pip install "ray"

  4. מגדירים את משתני הסביבה שמוגדרים כברירת מחדל:

    export PROJECT_ID=$(gcloud config get project)
export CLUSTER_NAME=ray-llm-cluster
export REGION=REGION
export ZONE=ZONE
export NAMESPACE=default
export KSA_NAME=ray-ksa
export GSA_NAME=tpu-reader-sa
export NETWORK_NAME=${CLUSTER_NAME}-net
export GS_BUCKET=BUCKET_NAME
export REPO_NAME=ray-repo
export CUSTOM_IMAGE_URI=REGION-docker.pkg.dev/PROJECT_ID/REPOSITORY/vllm-tpu-ray:vllm-tpu

    מחליפים את מה שכתוב בשדות הבאים:

    • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
    • ‫CLUSTER_NAME: השם של האשכול.
    • ‫REGION: האזור שבו קיבולת TPU Trillium זמינה.
    • ‫ZONE: האזור שבו קיבולת ה-TPU Trillium שלכם זמינה. מידע נוסף זמין במאמר זמינות של TPU ב-GKE.
    • ‫REPOSITORY: השם של מאגר Artifact Registry.
    • ‫BUCKET_NAME: השם של קטגוריית האחסון.

יצירה והגדרה של Google Cloud משאבים

כדי ליצור את המשאבים הנדרשים, פועלים לפי ההוראות הבאות.

יצירת אשכול GKE ומאגר צמתים

אפשר להפעיל את Gemma ב-TPU באשכול GKE במצב Autopilot או במצב Standard. ‫DRANET מנוהל על ידי GKE ומבקש ומנהל באופן דינמי משאבי רשת בעלי ביצועים גבוהים עבור ה-Pods המבוזרים שלכם. כך, GKE יכולה להקצות אוטומטית רשתות משניות במהירות גבוהה לתקשורת בין מאיצים, בלי שתצטרכו להגדיר VPC באופן ידני.

טייס אוטומטי

  1. ב-Cloud Shell, יוצרים את אשכול Autopilot:

    gcloud container clusters create-auto ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --enable-ray-operator \
    --location=${REGION}

  2. מגדירים את kubectl לתקשורת עם האשכול:

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
    --location=${REGION}

  3. כדי להשתמש ב-GKE managed DRANET במצב Autopilot, צריך לפרוס את משאב ComputeClass המותאם אישית שמופיע במאגר כדי להצטרף לשימוש ברשת דינמית:

    apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: dranet-compute-class
spec:
  nodePoolAutoCreation:
    enabled: true
  nodePoolConfig:
    dra:
      networking:
        enabled: true
  priorities:
  - machineType: ct6e-standard-4t
    acceleratorNetworkProfile: auto

  4. מחילים את המניפסט על האשכול:

    kubectl apply -f ai-ml/gke-ray/rayserve/llm/tpu/networking/dranet-compute-class.yaml

רגילה

  1. ב-Cloud Shell, יוצרים אשכול Standard שמאפשר את Ray operator ומשתמש ב-GKE Dataplane V2:

    gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --addons=RayOperator,GcsFuseCsiDriver \
    --machine-type=n2-standard-8 \
    --enable-dataplane-v2 \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --location=${ZONE}

  2. יוצרים מאגר של צומתי TPU עם כמה מארחים, כשהדרייבר DRANET מופעל:

    gcloud container node-pools create v6e-16 \
    --location=${ZONE} \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct6e-standard-4t \
    --tpu-topology=4x4 \
    --num-nodes=4 \
    --enable-gvnic \
    --scopes=https://www.googleapis.com/auth/cloud-platform \
    --accelerator-network-profile=auto \
    --node-labels=cloud.google.com/gke-networking-dra-driver=true

הגדרת אחסון ואימות

יוצרים קטגוריה ב-Cloud Storage ומפעילים את Rapid Cache כדי להאיץ את טעינת המודל, ואז מגדירים אימות ל-Hugging Face:

  1. באזור ה-TPU, יוצרים קטגוריית אחסון ומפעילים את מכונת Rapid Cache:

    gcloud storage buckets create gs://${GS_BUCKET} --project=${PROJECT_ID} --default-storage-class=STANDARD --location=${REGION}

gcloud storage buckets anywhere-caches create gs://${GS_BUCKET} ${ZONE} \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

  2. כדי להטמיע בצורה מאובטחת את מאגר המשקלים ב-Pods של GKE, צריך להגדיר קישורי זהויות. קודם כול, יוצרים חשבון שירות ייעודי ב-IAM ומעניקים לו הרשאות קריאה לקטגוריה:

    gcloud iam service-accounts create ${GSA_NAME}

gcloud storage buckets add-iam-policy-binding gs://${GS_BUCKET} \
    --member="serviceAccount:${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role="roles/storage.objectAdmin"

  3. יוצרים את הקישור של איחוד הזהויות של עומסי העבודה ל-GKE ומבצעים הערה לאובייקט Kubernetes ServiceAccount:

    gcloud iam service-accounts add-iam-policy-binding ${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
    --role="roles/iam.workloadIdentityUser" \
    --member="serviceAccount:${PROJECT_ID}.svc.id.goog[${NAMESPACE}/${KSA_NAME}]"

kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}
kubectl annotate serviceaccount ${KSA_NAME} --namespace ${NAMESPACE} iam.gke.io/gcp-service-account=${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com

  4. כדי להוריד את משקלי המודל של Gemma 4, צריך לאשר את הסכם הרישיון של Google ב-Hugging Face. עוברים אל דף המודל Gemma 4 ב-Hugging Face.

  5. נכנסים לחשבון ומאשרים את תנאי הרישיון בלחיצה על הסכמה וגישה למאגר.

  6. עוברים להגדרות החשבון ב-Hugging Face ויוצרים טוקן גישה עם התפקיד Read.

  7. מייצאים את האסימון של Hugging Face ויוצרים סוד של Kubernetes כדי ש-Ray יוכל לשלוף את משקלי המודל:

    export HF_TOKEN=YOUR_HUGGING_FACE_TOKEN

kubectl create secret generic hf-secret \
  --from-literal=hf_api_token=${HF_TOKEN}

יצירת קובץ אימג' מותאם אישית של קונטיינר

כדי לוודא שלסביבה מרובת המארחים יש את כל התלויות הנדרשות, צריך ליצור תמונה בהתאמה אישית על סמך תמונת ה-TPU של vLLM ולהעתיק אליה את סקריפט ההפעלה.

  1. יוצרים מאגר Artifact Registry:

    gcloud artifacts repositories create ${REPO_NAME} \
    --repository-format=docker \
    --location=${REGION}

  2. מאמתים את Docker בפרויקט:

    gcloud auth configure-docker ${REGION}-docker.pkg.dev

  3. בודקים את Dockerfile במאגר לדוגמה:

    FROM vllm/vllm-tpu:v0.21.0

ENV VLLM_TARGET_DEVICE=tpu
ENV VLLM_XLA_CACHE_PATH=/data

USER root

RUN pip install --no-cache-dir -U \
    "https://s3-us-west-2.amazonaws.com/ray-wheels/master/75b85027a859439fae5634e49aa6443f6fbecfeb/ray-3.0.0.dev0-cp312-cp312-manylinux2014_x86_64.whl" && \
    pip install --no-cache-dir --no-deps "ray[llm]"

COPY serve_tpu_multihost.py /home/ray/serve_tpu_multihost.py

  4. יוצרים את האימג' ומעבירים אותו בדחיפה ל-Artifact Registry:

    docker build -t ${CUSTOM_IMAGE_URI} .
docker push ${CUSTOM_IMAGE_URI}

הכנה מראש של משקלי מודלים ב-Cloud Storage

לפני שמבצעים פריסה של RayCluster, כדאי לבצע אופטימיזציה של ביצועי טעינת המודל ולוודא זמינות גבוהה בכל פרוסת ה-TPU המבוזרת. לשם כך, אפשר להכין מראש את משקלי המודל ישירות בקטגוריית Cloud Storage באמצעות Kubernetes Job עצמאי. הגישה המופרדת הזו מאפשרת סטרימינג מקביל מתואם, ומקצרת את זמני ההפעלה של האשכול.

  1. קובץ המניפסט של משימת ההורדה זמין במאגר. בודקים את ההגדרות של קובץ המניפסט:

    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: ${KSA_NAME}
      restartPolicy: OnFailure
      containers:
      - name: downloader
        image: python:3.10-slim
        command: ["/bin/sh", "-c"]
        args:
        - |
          pip install -U huggingface_hub filelock

          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="google/gemma-4-31B-it", 
              local_dir="/data/google/gemma-4-31B-it"
          )
          '
        env:
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        volumeMounts:
        - name: gcs-fuse-csi-ephemeral
          mountPath: /data
      volumes:
      - name: gcs-fuse-csi-ephemeral
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: ${GS_BUCKET}
            mountOptions: "implicit-dirs"

  2. יוצרים את משימת ההורדה על ידי החלת הקובץ במאגר:

    envsubst < ai-ml/gke-ray/rayserve/llm/tpu/components/model-downloader-job.yaml | kubectl apply -f -

  3. עוקבים אחרי העבודה עד שזרם ההורדה מדווח על הצלחה:

    kubectl logs -f job/model-downloader

יצירת סקריפט ההסקה

סקריפט Python הבא מגדיר אפליקציית Ray Serve שמבוססת על עטיפה (wrapper) ברמה גבוהה של LLMConfig Ray Serve.

  1. בודקים את הסקריפט serve_tpu_multihost.py במאגר לדוגמה:

    import os
import ray
from ray import serve
from ray.serve.llm import LLMConfig, ModelLoadingConfig, LLMServingArgs, build_openai_app

# Read configurations from environment variables
MODEL_ID = os.environ.get("MODEL_ID", "google/gemma-4-31B-it")
MODEL_SOURCE = os.environ.get("MODEL_SOURCE", "/data/google/gemma-4-31B-it")

# TPU hardware options (i.e. TPU-V6E, TPU-V7X etc.)
ACCELERATOR_TYPE = os.environ.get("ACCELERATOR_TYPE", "TPU-V6E")
TPU_TOPOLOGY = os.environ.get("TPU_TOPOLOGY", "4x4")

# vLLM engine parameters
TENSOR_PARALLEL_SIZE = int(os.environ.get("TENSOR_PARALLEL_SIZE", "16"))
MAX_MODEL_LEN = int(os.environ.get("MAX_MODEL_LEN", "8192"))
MAX_NUM_BATCHED_TOKENS = int(os.environ.get("MAX_NUM_BATCHED_TOKENS", "4096"))

# Define the multi-host TPU LLM config
llm_config = LLMConfig(
    model_loading_config=dict(
        model_id=MODEL_ID,
        model_source=MODEL_SOURCE
    ),
    accelerator_type=ACCELERATOR_TYPE,
    accelerator_config={"kind": "tpu", "topology": TPU_TOPOLOGY},
    engine_kwargs={
        "tensor_parallel_size": TENSOR_PARALLEL_SIZE,
        "max_model_len": MAX_MODEL_LEN,
        "max_num_batched_tokens": MAX_NUM_BATCHED_TOKENS,
        "distributed_executor_backend": "ray",
    }
)

deployment = build_openai_app(
    LLMServingArgs(
        llm_configs=[llm_config]
    )
)

הסבר על Ray LLM API

הסקריפט משתמש בספרייה המקורית ray.serve.llm של Ray Serve כדי להפשט את המורכבות של תזמור TPU מרובה מארחים. המסגרת Ray Serve LLM עוטפת את מנוע vLLM ומספקת מסגרת ניתנת להתאמה עם ביצועים גבוהים, שנועדה במיוחד לעומסי עבודה של הסקת מסקנות מבוזרת מאוד ב-Production.

לשימוש ב-Ray LLM API יש כמה יתרונות מרכזיים:

  • פריסות מרובות צמתים: Ray Serve LLM מאפשר למשתמשים להפעיל מודלים גדולים שמשתרעים על כמה מארחים מבוזרים (כמו חלוקה של TPU למספר מארחים) עם מיקום, תיאום והפצה של טופולוגיה באופן אוטומטי.
  • תאימות ל-vLLM: Ray Serve LLM מספק API שתואם ל-OpenAI ומתאים לשרת של vLLM. בנוסף, תוכלו לגשת לסט התכונות המתקדמות של vLLM (כמו פלט מובנה, יכולות מולטי-מודאליות ומודלים של חשיבה רציונלית) תוך כדי הרחבת נפח העבודה באשכול Kubernetes.
  • תכונות שמוכנות לשימוש בסביבת ייצור: Ray Serve LLM כולל יכולות ברמה ארגונית כמו התאמה אוטומטית לעומס (autoscaling), ניתוב בקשות מותאם אישית למיקסום פגיעות במטמון ושילובים מובנים למדדים ולניטור.

בסקריפט ההסקה שסופק, הפריסה מוגדרת על ידי שני רכיבים עיקריים:

  • ‫LLMConfig: האובייקט הזה מגדיר את הגדרות ההצגה. הוא מציין את מקור המודל, את פרמטרי המנוע של vLLM ואת accelerator_config. אם מגדירים את {"kind": "tpu", "topology": "4x4"}, ‏ Ray Serve LLM מקצה באופן אוטומטי קבוצת מיקום מבוזרת שמיפוי שלה תואם בדיוק לפרוסת TPU v6e פיזית עם 16 שבבים.
  • ‫build_openai_app: ה-API הזה עוטף באופן אוטומטי את מנוע vLLM המוגדר בשרת FastAPI שתואם ל-OpenAI, ומספק לכם API ל-REST בתקן התעשייה (כמו /v1/chat/completions) בלי שתצטרכו לכתוב קוד שרת בהתאמה אישית.

פריסת RayService

פורסים את הגדרת הרשת של Dynamic Resource Allocation ‏ (DRA) ואת RayService מניפסט ההצגה:

  1. כדי לבקש את כל הממשקים הזמינים של NetDevice בכל צומת, צריך לפרוס את ResourceClaimTemplate שמופיע במאגר:

    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

  2. מחילים את מניפסט התבנית על האשכול:

    kubectl apply -f ai-ml/gke-ray/rayserve/llm/tpu/networking/all-netdev-template.yaml

  3. מניפסט ההגשה RayService זמין במאגר. בודקים את ההגדרות של קובץ המניפסט:

    apiVersion: ray.io/v1
kind: RayService
metadata:
  name: vllm-tpu-multihost
  labels:
    ai.gke.io/model: "gemma-4-31B-it"
    ai.gke.io/inference-server: "vllm"
spec:
  serveConfigV2: |
    http_options:
      host: 0.0.0.0
      port: 8000
    applications:
      - name: llm
        import_path: ai-ml.gke-ray.rayserve.llm.tpu.serve_tpu_multihost:deployment
        runtime_env:
          working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/main.zip"
          env_vars:
            # Use local disk to prevent multi-host GCSFuse race conditions
            VLLM_XLA_CACHE_PATH: "/tmp/vllm_xla_cache"
  rayClusterConfig:
    headGroupSpec:
      rayStartParams: {}
      template:
        metadata:
          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
          containers:
          - name: ray-head
            image: $CUSTOM_IMAGE_URI
            imagePullPolicy: Always
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "2"
                memory: 16Gi
              requests:
                cpu: "2"
                memory: 16Gi
            volumeMounts:
            - name: dshm
              mountPath: /dev/shm
            - name: gcs-fuse-csi-ephemeral
              mountPath: /data
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
          - name: gke-gcsfuse-cache
            emptyDir:
              medium: Memory
          - name: gcs-fuse-csi-ephemeral
            csi:
              driver: gcsfuse.csi.storage.gke.io
              volumeAttributes:
                bucketName: $GS_BUCKET
                mountOptions: "implicit-dirs"
    workerGroupSpecs:
    - groupName: tpu-group
      replicas: 1
      minReplicas: 1
      maxReplicas: 1
      numOfHosts: 4
      rayStartParams: {}
      template:
        metadata:
          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
          containers:
            - name: ray-worker
              image: $CUSTOM_IMAGE_URI
              imagePullPolicy: Always
              resources:
                limits:
                  cpu: "20"
                  google.com/tpu: "4"
                  memory: 200Gi
                requests:
                  cpu: "20"
                  google.com/tpu: "4"
                  memory: 200Gi
                claims:
                - name: netdev
              env:
                - name: HF_HOME
                  value: "/data/huggingface"
                - name: HF_TOKEN
                  valueFrom:
                    secretKeyRef:
                      name: hf-secret
                      key: hf_api_token
                - name: JAX_PLATFORMS
                  value: "tpu,cpu"
                - name: NODE_IP
                  valueFrom:
                    fieldRef:
                      fieldPath: status.hostIP
                - name: VBAR_CONTROL_SERVICE_URL
                  value: $(NODE_IP):8353
                - name: TPU_MULTIHOST_BACKEND
                  value: "ray"
                - name: TPU_BACKEND_TYPE
                  value: "jax"
                - name: ENABLE_PJRT_COMPATIBILITY
                  value: "true"
              volumeMounts:
              - name: dshm
                mountPath: /dev/shm
              - name: gcs-fuse-csi-ephemeral
                mountPath: /data
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
          - name: gke-gcsfuse-cache
            emptyDir:
              medium: Memory
          - name: gcs-fuse-csi-ephemeral
            csi:
              driver: gcsfuse.csi.storage.gke.io
              volumeAttributes:
                bucketName: $GS_BUCKET
                mountOptions: "implicit-dirs"
          resourceClaims:
            - name: netdev
              resourceClaimTemplateName: all-netdev
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
            cloud.google.com/gke-tpu-topology: 4x4

  4. פורסים את השירות באמצעות המניפסט:

    טייס אוטומטי

    1. כדי לפרוס את השירות באשכול Autopilot, קודם צריך להוריד את המניפסט ולערוך אותו באופן מקומי כדי להוסיף את ההסכמה ComputeClass nodeSelector, שנדרשת לרשת DRANET ב-Autopilot:

      curl -O https://raw.githubusercontent.com/GoogleCloudPlatform/kubernetes-engine-samples/main/ai-ml/gke-ray/rayserve/llm/tpu/ray-service.tpu-v6e-multihost.yaml

    2. מוסיפים את התווית מתחת לשדה nodeSelector כך:

      nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
  cloud.google.com/gke-tpu-topology: 4x4
  cloud.google.com/compute-class: dranet-compute-class

    3. לאחר מכן, פורסים את השירות באמצעות מניפסט מקומי שעבר שינוי:

      envsubst < ray-service.tpu-v6e-multihost.yaml | kubectl apply -f -

    רגילה

    כדי לפרוס את השירות באשכול Standard, פורסים את המניפסט ישירות מהמאגר:

    envsubst < ai-ml/gke-ray/rayserve/llm/tpu/ray-service.tpu-v6e-multihost.yaml | kubectl apply -f -

אימות

  1. מחכים עד ש-RayService יהיה זמין:

    kubectl wait --for=condition=Ready --timeout=1800s rayservice/vllm-tpu-multihost

  2. כדי לוודא שהמודל נטען בהצלחה, מעיינים ביומנים מ-Ray head Pod:

    kubectl logs -f -l ray.io/node-type=head -c ray-head

פרסום המודל

בקטע הזה אתם תנהלו אינטראקציה עם המודל. לפני שממשיכים, צריך לוודא שהמודל הורד במלואו.

הגדרת העברה ליציאה אחרת

מגדירים העברה ליציאה אחרת למודל על ידי הרצת הפקודה הבאה:

kubectl port-forward svc/vllm-tpu-multihost-head-svc 8000:8000 2>&1 >/dev/null &

אינטראקציה עם המודל באמצעות curl

בקטע הזה מוסבר איך לבצע בדיקת עשן בסיסית כדי לאמת את מודל Gemma 4 שפרסתם.

בסשן חדש של מסוף, משתמשים ב-curl כדי לשוחח עם המודל:

curl -X POST http://127.0.0.1:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "google/gemma-4-31B-it",
        "messages": [
            {
              "role": "user",
              "content": "Why is GKE managed DRANET preferred for multi-host TPU networking?"
            }
        ],
        "max_tokens": 256
    }'

הפלט אמור להיראות כך:

‫
{
  "id": "chatcmpl-392692d3-5325-4832-a3a3-0b084c1045b0",
  "object": "chat.completion",
  "created": 1779883255,
  "model": "google/gemma-4-31B-it",
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "content": "To understand why GKE-managed **DRANET** (Distributed RANET) is preferred for multi-host TPU networking, it is first necessary to understand the fundamental challenge of TPU pods: **the need for massive, low-latency, all-to-all communication.**\n\nWhen you scale a model across multiple TPU hosts (multi-host), the hosts must synchronize gradients and weights constantly. Standard TCP/IP networking introduces too much overhead (latency and CPU jitter) for these operations.\n\nHere is the detailed breakdown of why GKE-managed DRANET is the preferred architecture:\n\n### 1. Bypassing the Kernel (Zero-Copy Networking)\nStandard networking requires the operating system kernel to handle packets, moving data from the network card to kernel space and then to user space.\n*   **The DRANET Advantage:** DRANET implements a specialized networking stack that allows for **Kernel Bypass**. It enables the TPU hardware/drivers to write data directly into the memory of the destination host. This reduces latency and eliminates the CPU overhead associated with processing network interrupts.\n\n### 2. High-Bandwidth, Low-Latency Interconnect\nMulti-host TPU training relies on a specialized topology (like a 2D or 3D"
      },
      "finish_reason": "length"
    }
  ]
}

(אופציונלי) אינטראקציה עם המודל באמצעות ממשק צ'אט של Gradio

בקטע הזה נבנה אפליקציית צ'אט לאתר שמאפשרת אינטראקציה עם מודל שעבר כוונון לפי הוראות.

‫Gradio היא ספריית Python עם wrapper של ChatInterface שיוצר ממשקי משתמש לצ'אטבוטים.

פריסת ממשק הצ'אט

קובץ המניפסט של ממשק הצ'אט זמין במאגר. בודקים את ההגדרות של קובץ המניפסט:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: gradio
  labels:
    app: gradio
spec:
  replicas: 1
  selector:
    matchLabels:
      app: gradio
  template:
    metadata:
      labels:
        app: gradio
    spec:
      containers:
      - name: gradio
        image: us-docker.pkg.dev/google-samples/containers/gke/gradio-app:v1.0.7
        resources:
          requests:
            cpu: "250m"
            memory: "512Mi"
          limits:
            cpu: "500m"
            memory: "512Mi"
        env:
        - name: CONTEXT_PATH
          value: "/v1/chat/completions"
        - name: HOST
          value: "http://vllm-tpu-multihost-serve-svc:8000"
        - name: LLM_ENGINE
          value: "openai-chat"
        - name: MODEL_ID
          value: "google/gemma-4-31B-it"
        - name: DISABLE_SYSTEM_MESSAGE
          value: "true"
        ports:
        - containerPort: 7860
---
apiVersion: v1
kind: Service
metadata:
  name: gradio
spec:
  selector:
    app: gradio
  ports:
  - protocol: TCP
    port: 8080
    targetPort: 7860
  type: ClusterIP

החלת המניפסט:

kubectl apply -f ai-ml/gke-ray/rayserve/llm/tpu/components/gradio.yaml

מחכים שהפריסה תהיה זמינה:

kubectl wait --for=condition=Available --timeout=900s deployment/gradio

שימוש בממשק הצ'אט

ב-Cloud Shell, מריצים את הפקודה הבאה:

kubectl port-forward service/gradio 8080:8080

הפעולה הזו יוצרת העברת פורטים מ-Cloud Shell לשירות Gradio.

לוחצים על סמל תצוגה מקדימה של האתר כפתור לתצוגה מקדימה של אתר בפינה השמאלית העליונה של סרגל המשימות של Cloud Shell. לוחצים על תצוגה מקדימה ביציאה 8080. תיפתח כרטיסייה חדשה בדפדפן.

מנהלים אינטראקציה עם Gemma באמצעות ממשק הצ'אט של Gradio. מוסיפים הנחיה ולוחצים על שליחה.

מעקב אחר ביצועי המודל

כדי לראות את מרכזי הבקרה של מדדי יכולת הצפייה של מודל שפועל ב-KubeRay, אפשר להשתמש במרכזי הבקרה הייעודיים של Ray ב-GKE.

הוראות מפורטות להגדרת האשכול ולגישה ללוחות הבקרה של יכולת הצפייה מפורטות במאמר איסוף והצגה של יומנים ומדדים עבור RayClusters ב-Google Kubernetes Engine ‏ (GKE).

גישה ללוח הבקרה של Ray

כדי לבדוק את הסטטוס של רכיבי Ray, לצפות ביומני אפליקציות מפורטים ולעקוב אחרי ניצול משאבים ברמת הצומת באופן מקורי ב-Ray, אפשר לגשת ללוח הבקרה של Ray.

  1. מעבירים את שירות צומת הראשי של Ray למכונה המקומית:

    kubectl port-forward svc/vllm-tpu-multihost-head-svc 8265:8265

  2. פותחים את הדפדפן ועוברים אל http://localhost:8265. אם אתם משתמשים ב-Cloud Shell, לוחצים על הלחצן Web Preview (תצוגה מקדימה באינטרנט) ובוחרים באפשרות Preview on port 8265 (תצוגה מקדימה ביציאה 8265).

  3. כדי לראות את הפריסות של vLLM, את תקינות העותקים של המודל ואת זמן האחזור של השאילתות, לוחצים על הכרטיסייה Serve (פרסום).

הסרת המשאבים

כדי להימנע מחיובים בחשבון על המשאבים שבהם השתמשתם במדריך הזה, מוחקים את המשאבים: Google Cloud

  1. מוחקים את RayService:

    kubectl delete rayservice vllm-tpu-multihost

  2. מחיקת אשכול GKE:

    gcloud container clusters delete ${CLUSTER_NAME} --zone=${ZONE}

המאמרים הבאים