אופטימיזציה של אימון AI במעבדי TPU באמצעות DWS ו-Kueue

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

ההגדרה שמוגדרת במסמך הזה יוצרת מאגר ראשי של צמתים ומאגר גיבוי של צמתים:

  • על פי דרישה (Plan A): מאגר הצמתים הזה הוא הבחירה הראשונה שלכם לצמתים. המערכת מנסה קודם לתזמן משימות במכונות האלה לפי דרישה. המונח על פי דרישה מציין שברגע שהמכונות האלה מתחילות לפעול, יש לכם גישה אמינה ורציפה אליהן.
  • ‫DWS flex-start (תוכנית ב'): זהו מאגר הצמתים לגיבוי. אם מכונות של תוכנית א' לא זמינות, תוכנית התזמון Kueue מקצה אוטומטית את העבודה למאגר הזה של תוכנית ב'. אחר כך, DWS יחפש את החומרה של תוכנית ב', אבל לא מובטחת גישה מיידית כי יכול להיות שהחומרה הזו לא תהיה זמינה. אבל DWS לא מוותר: הוא מחזיק את הבקשה שלכם בתור למשך עד 7 ימים ומספק את המכונות באופן אוטומטי ברגע שהן זמינות.

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

סקירה כללית של שלבי ההגדרה

כדי להגדיר את מערכת הגיבוי האוטומטית, צריך לבצע כמה שלבי הגדרה. כדאי לחלק את ההגדרה הזו לשתי קטגוריות:

  • משימות של אדמין אשכול: הגדרה חד-פעמית של התשתית, כמו יצירת אשכול GKE, הקצאת מאגרי צמתים והתקנת בקר התזמון של Kueue.
  • משימות של מפתחי AI: תהליכי עבודה יומיים שחוזרים על עצמם, כמו הגדרת הדרישות של משימת האימון ושליחת עומס העבודה.

גם אם אתם מבצעים את כל השלבים האלה בעצמכם, חשוב לזכור את ההבחנה הזו כדי להבין את התהליך הכולל.

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

נושא משימה
הגדרת התשתית (אדמין של אשכול) 1. יוצרים אשכול GKE
2. יוצרים את מאגרי הצמתים
3. אימות הסטטוס של flex-start במאגר הצמתים
התקנה והגדרה של Kueue (אדמין של אשכול) 1. מתקינים את Kueue
2. הגדרת כללי ההגדרה
הפעלת משימת הדרכה (מפתח AI) 1. יוצרים את ConfigMap
2. מגדירים את מניפסט RayJob
3. שליחת עומס העבודה
4. מתחברים אל RayJob
5. בדיקת היומנים

מושגים מרכזיים

  • מאגר צמתים על פי דרישה (תוכנית א'): מאגר הצמתים הראשי בעדיפות גבוהה. המשימה תמיד תנסה להשתמש במאגר הזה קודם.
  • מאגר צמתים של DWS flex-start (Plan B): מאגר הצמתים של הגיבוי. אם המכונות במאגר הראשי לא זמינות, המערכת משתמשת במאגר הזה באופן אוטומטי כדי לחפש חומרה זמינה.
  • ‫Kueue: תוכנת תזמון שמנהלת את תור העבודות. הוא מיירט את בקשת המשימה ומחליט באיזה מאגר צמתים להשתמש (תוכנית א' או תוכנית ב').
  • משימה: עומס העבודה של אימון ה-AI שרוצים להריץ. במסמך הזה, תגדירו אותו באמצעות מניפסט של RayJob.

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

  1. בדף לבחירת הפרויקט במסוף Google Cloud , בוחרים פרויקט ב- Google Cloud או יוצרים אותו.

    תפקידים שנדרשים כדי לבחור או ליצור פרויקט

    • Select a project: כדי לבחור פרויקט לא צריך תפקיד IAM ספציפי – אפשר לבחור כל פרויקט שקיבלתם בו תפקיד.
    • יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים?

    כניסה לדף לבחירת הפרויקט

  2. מוודאים שהחיוב מופעל בפרויקט Google Cloud .

  3. מפעילים את Google Kubernetes Engine ואת Cloud TPU API, אם הם עדיין לא מופעלים.

    התפקידים שנדרשים כדי להפעיל ממשקי API

    כדי להפעיל ממשקי API, צריך את ההרשאה serviceusage.services.enable. אם יצרתם את הפרויקט, סביר להניח שכבר יש לכם את ההרשאה הזו דרך התפקיד 'בעלים' (roles/owner). אחרת, תוכלו לקבל את ההרשאה הזו דרך התפקיד 'אדמין בממשק 'שימוש בשירות'' (roles/serviceusage.serviceUsageAdmin). איך מקצים תפקידים

    הפעלת ממשקי ה-API

  4. במסוף Google Cloud , מפעילים את Cloud Shell.

    הפעלת Cloud Shell

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

הגדרת משתני סביבה

כדי לפשט את הפקודות שאתם מריצים במסמך הזה, אתם יכולים להגדיר משתני סביבה ב-Cloud Shell. המשתנים האלה מכילים ערכים כמו מזהה הפרויקט שלכם Google Cloud , שמות מאגרי הצמתים והמיקום של אשכול GKE.

אחרי שמגדירים את המשתנים האלה, אפשר להשתמש בהם שוב בכמה פקודות על ידי הפניה לשם המשתנה (לדוגמה, $CLUSTER_NAME) במקום להקליד מחדש או להחליף ערכים בכל פעם. הגישה הזו מקלה על המעקב אחרי התהליך ומפחיתה את הסיכון לשגיאות.

כדי להגדיר את משתני הסביבה השימושיים הבאים ב-Cloud Shell, מריצים את הפקודות הבאות:

export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export ZONE="us-east5-b"
export REGION="us-east5"
export CLUSTER_NAME="tpu-cluster"
export GKE_VERSION="1.34"
export ONDEMAND_NODEPOOL="on-demand-pool"
export DWS_NODEPOOL="dws-pool"

הנה הסבר על משתני הסביבה האלה:

  • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
  • ‫PROJECT_NUMBER: מספר המזהה הייחודי של הפרויקט (לדוגמה, 123456789012).
  • ‫ZONE: אזור המחשוב של האשכול (לדוגמה, us-east5-b). צריך לבחור אזור שבו סוג המאיץ שבחרתם זמין. למידע על זמינות, אפשר לעיין במאמרים בנושא מכסות של Cloud TPU או מכסות של GPU.
  • ‫REGION: האזור שבו יוצרים את משאבי האשכול (לדוגמה, us-east5).
  • ‫CLUSTER_NAME: השם שתבחרו לאשכול GKE.
  • ‫GKE_VERSION: גרסת ה-GKE של האשכול. צריך להשתמש בגרסה 1.34 ואילך.
  • ‫ONDEMAND_NODEPOOL: השם של מאגר הצמתים הרגיל על פי דרישה. (זהו מאגר הצמתים של תוכנית א').
  • ‫DWS_NODEPOOL: השם של מאגר הצמתים של DWS flex-start. (זהו מאגר הצמתים של תוכנית ב').

הגדרת התשתית (אדמין של האשכול)

כמנהלי האשכול, אתם מגדירים את אשכול GKE ואת מאגרי הצמתים כך שיתמכו במנגנון החלופי.

יצירת אשכול GKE

קודם יוצרים את אשכול GKE. האשכול הזה הוא הסביבה שבה מתקינים את בקר Kueue, מגדירים את מאגרי הצמתים ומריצים את משימות האימון של ה-AI. כדי ליצור אשכול ולהתחבר אליו, מבצעים את השלבים הבאים:

  1. יוצרים את האשכול:

    gcloud container clusters create ${CLUSTER_NAME} \
  --cluster-version=${GKE_VERSION} \
  --machine-type=n2-standard-16 \
  --location=${ZONE} \
  --enable-image-streaming \
  --addons=RayOperator \
  --project=${PROJECT_ID}

    בפקודה הזו נעשה שימוש בדגלים הבאים:

    • ‫--addons=RayOperator: התקנת Ray Operator באשכול. אתם צריכים את האופרטור הזה כדי לנהל את עומס העבודה של RayJob שתשלחו בהמשך המסמך.
    • ‫--enable-image-streaming: מאפשר לאשכול לשלוף קובצי אימג' של קונטיינרים מהר יותר. התכונה הזו מקצרת משמעותית את הזמן שנדרש כדי להפעיל תמונות גדולות של קונטיינרים של AI.

  2. מאחזרים את פרטי הכניסה של האשכול כדי שממשק שורת הפקודה של kubectl יוכל להתחבר אליו. הפקודה הזו מעדכנת את קובץ התצורה של Kubernetes, שמאוחסן כברירת מחדל בספרייה ~/.kube/config:

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --location=${ZONE} \
  --project=${PROJECT_ID}

יצירת מאגרי הצמתים

יוצרים את מאגרי הצמתים הראשיים והגיבוי לסביבה: מאגר הצמתים על פי דרישה (תוכנית א') ומאגר הצמתים של DWS flex-start (תוכנית ב'):

  1. יוצרים את מאגר הצמתים על פי דרישה: המאגר הזה משמש כמשאב העיקרי למשימות אימון:

    gcloud container node-pools create ${ONDEMAND_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    בדוגמה הזו, המכונות בעדיפות ראשונה הן מאיצי TPU v6e. מציינים את החומרה הזו באמצעות הדגל --machine-type=ct6e-standard-4t. אתם יכולים לשנות את סוג המכונה הזה כך שיתאים לחומרה, כמו GPUs או TPUs שונים, שאתם רוצים למודל ה-AI שלכם.

  2. יצירת מאגר צמתים עם הפעלה גמישה של DWS: בדוגמה הזו, בוחרים את אותו סוג מכונה (--machine-type=ct6e-standard-4t) שבחרתם למאגר הראשי לפי דרישה. לא חייבים להשתמש בסוג מכונה שונה במאגר הצמתים של תוכנית ב'. בגלל שאתם ממש רוצים את החומרה הספציפית הזו, אם היא לא זמינה באופן מיידי, אתם יכולים להגדיר אותה כתוכנית חלופית ולעבור לשיטה אחרת כדי להשיג אותה. בשיטה החלופית הזו נעשה שימוש ב-DWS, שמחפש באופן רציף את החומרה הזמינה למשך עד 7 ימים:

    gcloud container node-pools create ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --enable-autoscaling \
  --enable-queued-provisioning \
  --flex-start \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

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

    • ‫--num-nodes=0, --min-nodes=0, --max-nodes=4 ו---enable-autoscaling: השילוב הזה מאפשר להגדיל את מספר הצמתים במאגרי הצמתים מאפס צמתים כשנדרש צומת למשימה, ולהקטין את מספר הצמתים כשאין צורך בצומת, וכך לחסוך בעלויות.
    • ‫--tpu-topology: מגדיר את הסידור הפיזי של שבבי ה-TPU. אתם מציינים את הפריסה הזו כי הסידור הפיזי של הצ'יפים משפיע על מהירות ההרצה של משימת האימון המבוזר.
    • ‫--reservation-affinity=none: (תכנון רק של מאגר) הדגל הזה מונע ממאגר הצמתים על פי דרישה לצרוך את הקיבולת השמורה שלכם. מאגר הצמתים של DWS flex-start לא צורך הזמנות.
    • ‫--enable-queued-provisioning ו---flex-start: (רק למאגר של תוכנית ב') הדגלים האלה מאפשרים ל-DWS להקצות צמתים למאגר של תוכנית ב' מקיבולת גמישה כשהיא זמינה.

אימות הסטטוס של התחלה גמישה במאגר הצמתים

בודקים את מאגר הצמתים של DWS flex-start ומוודאים ש-flex-start מופעל:

gcloud container node-pools describe ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --format="get(config.flexStart)"

אם האפשרות flex-start מופעלת, הפלט הוא True.

התקנה והגדרה של Kueue (אדמין של אשכול)

בקטע הזה מתקינים את בקר Kueue באשכול. תזכורת: Kueue היא תוכנת תזמון שמנהלת את תור העבודות. הוא מיירט את בקשת העבודה, מחליט באיזה מאגר צמתים להשתמש (על פי דרישה או DWS flex-start), ואז מקצה את העבודה.

התקנת Kueue

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

helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
  --namespace kueue-system \
  --create-namespace \
  --set "controllerManager.featureGates[0].name=ElasticJobsViaWorkloadSlices" \
  --set "controllerManager.featureGates[0].enabled=true"

הגדרת כללי ההגדרה

יוצרים קובץ מניפסט בפורמט YAML שמגדיר את כללי העדיפות. הכללים האלה אומרים ל-Kueue להשתמש קודם במאגר לפי דרישה ואחר כך במאגר DWS flex-start:

  1. יוצרים קובץ בשם dws-tpu-queue.yaml עם התוכן הבא. הקובץ הזה מגדיר שני סוגים של משאבים (על פי דרישה ו-DWS flex-start) ותור אשכול שמתעדף אותם. קובץ ההגדרות הזה מגדיר את הלוגיקה שבה Kueue משתמש כדי לטפל בעבודות:

    • ‫ResourceFlavor: בתחילת המסמך הזה יצרתם שני מאגרי צמתים והקציתם להם שמות באמצעות משתני הסביבה ${ONDEMAND_NODEPOOL} ו-${DWS_NODEPOOL}. כש-GKE יצר את מאגרי הצמתים האלה, הוא תייג אוטומטית כל צומת במאגרים בשם שבחרתם למשתני הסביבה האלה. הקטע ResourceFlavor אומר ל-Kueue לחפש צמתים עם התוויות האלה.
    • ‫ClusterQueue: בקטע הזה של המניפסט מוגדר כלל העדיפות. הוא מפרט קודם את הטעם לפי דרישה, כך ש-Kueue מנסה להקצות קודם מכונות לפי דרישה. אם Kueue לא מצליח להשיג את המכונות האלה, הוא מנסה להקצות במקומן מכונות DWS flex-start.
    • ‫Quotas: הקובץ מגדיר מכסה, שהיא מגבלה על סך המשאבים (כמו מעבד, זיכרון ושבבי TPU) שהמשימות יכולות להשתמש בהם בכל זמן נתון במאגר הצמתים על פי דרישה. כשהמשימות מגיעות למגבלה הזו, Kueue מנסה באופן אוטומטי להקצות מכונות DWS flex-start (המכונות של תוכנית ב'), שהגדרתם ב-dws-tpu-queue.yaml עם מכסת שימוש גבוהה בהרבה. 
      apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "default-cpu"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: default-pool
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "on-demand"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${ONDEMAND_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "dws"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${DWS_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  namespaceSelector: {}
  resourceGroups:
    - coveredResources: ["cpu", "memory", "google.com/tpu"]
      flavors:
        - name: "default-cpu" # Used for Ray Head Pod.
          resources:
            - name: "cpu"
              nominalQuota: 10
            - name: "memory"
              nominalQuota: 20Gi
            - name: "google.com/tpu"
              nominalQuota: 0
        - name: "on-demand" # First choice: on-demand node-pool.
          resources:
            - name: "cpu"
              nominalQuota: 40
            - name: "memory"
              nominalQuota: 75Gi
            - name: "google.com/tpu"
              nominalQuota: 16
        - name: "dws" # If on-demand is unavailable, fallback to DWS.
          resources:
            - name: "cpu"
              nominalQuota: 1000000000
            - name: "memory"
              nominalQuota: 1000000000Gi
            - name: "google.com/tpu"
              nominalQuota: 1000000000 # "Infinite" quota
  admissionChecksStrategy:
    admissionChecks:
      - name: "dws-prov"
        onFlavors: [dws]
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: "default"
  name: "user-queue"
spec:
  clusterQueue: "cluster-queue"
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: AdmissionCheck
metadata:
  name: dws-prov
spec:
  controllerName: kueue.x-k8s.io/provisioning-request
  parameters:
    apiGroup: kueue.x-k8s.io
    kind: ProvisioningRequestConfig
    name: dws-config
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ProvisioningRequestConfig
metadata:
  name: dws-config
spec:
  provisioningClassName: queued-provisioning.gke.io
  managedResources:
    - google.com/tpu

  2. מחילים את ההגדרה על האשכול. הפקודה הבאה משתמשת בכלי שורת פקודה בשם envsubst כדי להחליף משתני placeholder שמופיעים בקובץ dws-tpu-queue.yaml. ‫envsubst מחליף את הפלייס הולדרים בערכים של משתני הסביבה שהגדרתם קודם:

    envsubst < dws-tpu-queue.yaml | kubectl apply -f -

הרצת משימת אימון (מפתח AI)

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

בקטע הזה תבצעו את השלבים הבאים:

  • יוצרים סקריפט אימון ב-Python.
  • מאחסנים את הסקריפט ב-Kubernetes ConfigMap.
  • פורסים RayJob שטוען את ConfigMap בתור נפח, כדי שאפשר יהיה להריץ את סקריפט האימון בצמתים.

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

סקריפט האימון

מעתיקים את סקריפט Python הבא ומדביקים אותו בקובץ בשם train.py:

import time
import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.v2.jax import JaxTrainer
import jax
import jax.numpy as jnp

def train_func():
    # JaxTrainer handles JAX distributed setup.
    print(f"Local Devices: {jax.local_devices()}")

    # Simple Linear Regression Training Loop.
    key = jax.random.PRNGKey(0)
    x = jax.random.normal(key, (1000, 10))
    w_true = jax.random.normal(key, (10, 1))
    y = jnp.dot(x, w_true)

    # Initialize weights
    w = jnp.zeros((10, 1))
    learning_rate = 0.1

    @jax.jit
    def update(w, x, y):
        y_pred = jnp.dot(x, w)
        loss = jnp.mean((y_pred - y) ** 2)
        grad = jax.grad(lambda w: jnp.mean((jnp.dot(x, w) - y) ** 2))(w)
        return w - learning_rate * grad, loss

    # Training loop
    print("Starting training...")
    for epoch in range(50):
        w, loss = update(w, x, y)
        if epoch % 10 == 0:
            train.report({"loss": loss.item(), "epoch": epoch})
            print(f"Epoch {epoch}: Loss {loss:.4f}")

    print("Training Complete!")
    # Allow metrics to sync before closing
    time.sleep(3)

def main():
    scaling_config = ScalingConfig(
        num_workers=4,
        resources_per_worker={"TPU": 4},
        use_tpu=True,
        topology="4x4",
        accelerator_type="TPU-V6E"
    )

    trainer = JaxTrainer(
        train_loop_per_worker=train_func,
        scaling_config=scaling_config
    )

    result = trainer.fit()
    print(f"Run Result: {result.metrics}")

if __name__ == "__main__":
    main()

סקריפט האימון משתמש ב-JAX, ספריית Python לחישובים מספריים עם ביצועים גבוהים, כדי לאמן מודל רגרסיה ליניארית. הסקריפט הזה הוא דוגמה פשוטה שממחישה איך להשתמש ב-DWS וב-Kueue כדי להגדיר מעבר אוטומטי לגיבוי, והוא לא מבצע מקביליות נתונים או מקביליות מודלים.

שימו לב שבקטע ScalingConfig של סקריפט ההדרכה מוגדרות דרישות החומרה למשימת ההדרכה. בקטע הזה מוגדרת טופולוגיה של TPU בגודל 4x4, שתואמת לפריסה הפיזית של מאגרי הצמתים שהגדרתם קודם.

יצירת ה-ConfigMap

מעלים את התוכן של סקריפט train.py לאובייקט Kubernetes ConfigMap. כך האשכול יכול לאחסן את הסקריפט ולהפוך אותו לזמין ל-RayJob:

kubectl create configmap jax-train-script --from-file=train.py

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

החלת המניפסט של RayJob

יוצרים קובץ בשם rayjob-tpu-v6e-dws.yaml עם התוכן הבא. קובץ המניפסט הזה מגדיר את משימת האימון ומסביר למערכת איך לנתב אותה:

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: rayjob-tpu-v6e-dws-${JOB_ID}
  labels:
    kueue.x-k8s.io/queue-name: user-queue
  annotations:
    kueue.x-k8s.io/elastic-job: "true"
spec:
  shutdownAfterJobFinishes: true
  entrypoint: python /app/train.py
  runtimeEnvYAML: |
    pip:
      - jax[tpu]==0.8.2
      - pandas==2.3.3
  rayClusterSpec:
    enableInTreeAutoscaling: true
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-nodepool: default-pool
          containers:
            - name: ray-head
              image: rayproject/ray:2.53.0-py311
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "2"
                  memory: "4Gi"
                requests:
                  cpu: "2"
                  memory: "4Gi"
              volumeMounts:
                - mountPath: /app
                  name: train-script-volume
          volumes:
            - name: train-script-volume
              configMap:
                name: jax-train-script
    workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 2
        numOfHosts: 4
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            tolerations:
              - key: "google.com/tpu"
                operator: "Exists"
                effect: "NoSchedule"
              - key: "cloud.google.com/gke-queued"
                operator: "Exists"
                effect: "NoSchedule"
            containers:
              - name: ray-worker
                image: rayproject/ray:2.53.0-py311
                resources:
                  limits:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                  requests:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                volumeMounts:
                  - mountPath: /app
                    name: train-script-volume
            volumes:
              - name: train-script-volume
                configMap:
                  name: jax-train-script
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

המניפסט הזה כולל שלוש הגדרות שמאפשרות למערכת החלופית לפעול:

  • בקשה לחומרה ספציפית: בקטע nodeSelector מצוינת החומרה שהסקריפט דורש (בדוגמה הזו, tpu-v6e-slice עם טופולוגיה של 4x4).
  • בחירת התור: התווית kueue.x-k8s.io/queue-name מעבירה את העבודה ישירות ל-Kueue. הפעולה הזו מפעילה את הלוגיקה האוטומטית של החלופה.
  • סובלני כלפי צמתים עם הפעלה גמישה של DWS: בקטע tolerations אפשר להגדיר שהעבודה תפעל במאגר הצמתים של תוכנית ב'. מכיוון שצמתי Flex-start של DWS מסומנים במיוחד (tainted) על ידי GKE כדי שעומסי עבודה רגילים לא יפעלו עליהם בטעות, המשימה צריכה לסבול באופן מפורש את ה-taint של cloud.google.com/gke-queued.

שליחת עומס העבודה

כדי להוכיח שמערכת הגיבוי פועלת, צריך לשלוח שתי משימות. המשימה הראשונה צורכת את הקיבולת לפי דרישה של תוכנית א', מה שמאלץ את המשימה השנייה לחזור לקיבולת של תוכנית ב' עם התחלה גמישה של DWS.

מריצים את הפקודה הבאה כדי לשלוח את שתי המשימות. הפקודה משתמשת בלולאת for וב-envsubst כדי להוסיף מזהה משימה ייחודי למניפסט בכל הרצה:

for i in 1 2; do
  export JOB_ID=$i
  envsubst < rayjob-tpu-v6e-dws.yaml | kubectl apply -f -
  echo "Submitted Job $i"
  sleep 2
done

אחרי ששולחים את העבודות, המערכת מטפלת בעומס העבודה באופן הבא:

  1. יירוט: Kueue מזהה את המשימות באמצעות תווית התור ומשהה אותן באופן זמני.
  2. החלטה: Kueue בודק את זמינות המשאבים בהתאם לכללים של האדמין. המערכת בודקת קודם את המאגר של תוכנית א'.
  3. במטלה:
    • מכיוון שהמשאבים של תוכנית א' זמינים לעבודה הראשונה, Kueue מקצה את עבודה 1 לתוכנית הזו.
    • מכיוון שמשימה 1 צורכת את המשאבים של תוכנית א', Kueue מקצה אוטומטית את משימה 2 למאגר של תוכנית ב' (DWS flex-start).
  4. הפעלה: Kueue מבטל את ההשעיה של המשימות. הפעולה הזו מפעילה את המידרוג האוטומטי של אשכול GKE כדי להקצות את הצמתים ולהפעיל את הסקריפטים לאימון.

התחברות ל-RayJob

כשלב אימות סופי, אפשר להשתמש בפקודה kubectl port-forward כדי להתחבר ל-Ray Dashboard ולצפות בהרצת המשימות.

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

kubectl port-forward service/rayjob-tpu-v6e-dws-1-head-svc 8265:8265 &

אחרי שמריצים את הפקודה, פותחים דפדפן אינטרנט ונכנסים לכתובת http://localhost:8265. במרכז הבקרה של Ray, אפשר לראות את סטטוס העבודה ואת המדדים שדווחו כדי לוודא ששתי העבודות הושלמו בהצלחה במאגרי הצמתים המתאימים.

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

kubectl logs job/rayjob-tpu-v6e-dws-1

הפלט הקטום של סקריפט ההדרכה אמור להיראות כך. ההודעות Training Complete! ו-Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded אמורות להופיע לקראת סוף הפלט:

(pid=, ip=10.68.3.4) 5] XLA::TPU program HBM usage: 52.5K / 31.25G
(pid=, ip=10.68.9.4) :2152] XLA::TPU program VMEM usage: 141.0K / 128.00M [repeated 5x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722540     855 deepsea_compiler_backend.cc:2163] Total hbm usage >= 260.14M: [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167]     reserved           204B [repeated 19x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722542     855 deepsea_compiler_backend.cc:2163]     program           70.0K  [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     arguments            0B  [repeated 12x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777627     888 deepsea_compiler_backend.cc:2163] Output size 0B; shares 0B with arguments. [repeated 14x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777625     888 deepsea_compiler_backend.cc:2163] Total host usage >= 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     program         unknown size  [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167] Program sflag requirement 224B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167]     scoped              40B [repeated 21x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167] Program vmem requirement 141.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program smem requirement 40B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program host requirement 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program hbm requirement 70.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2167]     overlays          70.0K [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2175] XLA::TPU program SMEM usage: 1.9K / 1.00M (3 parameters) [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167]     HLO temp          76.0K (0.0% utilization: Unpadded (0B) Padded (0B), 100.0% fragmentation (76.0K)) [repeated 14x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Training Complete! [repeated 3x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Epoch 40: Loss 0.0000 [repeated 3x across cluster]
2026-03-20 03:59:51,008 SUCC cli.py:65 -- ------------------------------------------
2026-03-20 03:59:51,008 SUCC cli.py:66 -- Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded
2026-03-20 03:59:51,008 SUCC cli.py:67 -- ------------------------------------------

הסרת המשאבים

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

מחיקת הפרויקט

  • במסוף Google Cloud , נכנסים לדף Manage resources.

    כניסה לדף Manage resources

  • ברשימת הפרויקטים, בוחרים את הפרויקט שרוצים למחוק ולוחצים על Delete.
  • כדי למחוק את הפרויקט, כותבים את מזהה הפרויקט בתיבת הדו-שיח ולוחצים על Shut down.

    • מחיקת המשאבים הבודדים

    אם רוצים לשמור את פרויקט GGoogle Cloud שבו השתמשתם במדריך הזה, מריצים את הפקודה הבאה כדי למחוק את האשכול:

    gcloud container clusters delete ${CLUSTER_NAME} \
    --location=${ZONE} \
    --project=${PROJECT_ID} \
    --quiet

    סיכום

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

    כדי שהפעולה תתבצע, צריך לבצע את השלבים הבאים:

    1. יצירת אשכול GKE: הגדרת הסביבה לאירוח מאגרי הצמתים וכלי התזמון.
    2. הגדרתם את מאגרי הצמתים: יצרתם מאגר צמתים לפי דרישה (תוכנית א') ומאגר צמתים של DWS (תוכנית ב').
    3. התקנה והגדרה של Kueue: פריסת בקר Kueue והחלת כללי עדיפות שמנחים את המערכת לנסות קודם את תוכנית א' ואז לחזור לתוכנית ב'.
    4. יצירת ConfigMap: פריסת סקריפט אימון פשוט של JAX באשכול, שישמש כעומס העבודה לבדיקה.
    5. הגדרתם מניפסט של RayJob: הגדרתם את הג'וב כך שיבקש חומרה ספציפית, ינותב לבקר Kueue ויאפשר שימוש בצמתי DWS.
    6. הגשת עומס העבודה: הגשת שתי משימות כדי לאלץ את Kueue לנתב אוטומטית את המשימה השנייה לתוכנית ב', כשהמשאבים של תוכנית א' מוצו.
    7. אימות התוצאות: נעשה שימוש בהפניית יציאות כדי להתחבר ל-Ray Dashboard ולוודא ששתי המשימות פעלו בהצלחה.

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