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

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

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

  • ‫On-demand (תוכנית 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 (תוכנית ב'): מאגר הצמתים של הגיבוי. אם מכונות במאגר הראשי לא זמינות, המערכת משתמשת במאגר הזה באופן אוטומטי כדי לחפש חומרה זמינה.
  • ‫Kueue: תוכנת תזמון שמנהלת את תור העבודות. הוא מיירט את בקשת העבודה ומחליט באיזה מאגר צמתים להשתמש (תוכנית א' או תוכנית ב').
  • משימה: עומס העבודה של אימון ה-AI שרוצים להריץ. במסמך הזה, תגדירו אותו באמצעות מניפסט RayJob.

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

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

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

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

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

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

  3. מפעילים את ממשקי ה-API של Google Kubernetes Engine ו-Cloud TPU.

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

    כדי להפעיל ממשקי API, צריך את תפקיד ה-IAM 'אדמין של Service Usage' (roles/serviceusage.serviceUsageAdmin), שכולל את ההרשאה serviceusage.services.enable. איך מקצים תפקידים

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

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

    הפעלת Cloud Shell

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

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

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

אחרי שמגדירים את המשתנים האלה, אפשר להשתמש בהם מחדש בכמה פקודות על ידי הפניה לשם המשתנה (לדוגמה, $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: השם של מאגר הצמתים הסטנדרטי על פי דרישה. (זוהי מאגר הצמתים של תוכנית A).
  • ‫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 CLI יוכל להתחבר אליו. הפקודה הזו מעדכנת את קובץ התצורה של 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 \
  --reservation-affinity=none \
  --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: עוזר לוודא שמאגר הצמתים לא צורך את החומרה שהוזמנה מראש. Google Cloud מאפשר לכם להזמין מכונות ספציפיות כדי לוודא שהן זמינות. הגדרת הדגל הזה לערך none אומרת למערכת לעקוף את ההזמנות האלה ולבקש באופן דינמי מכונות לא שמורות במקום זאת.
    • ‫--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)"

אם האפשרות 'התחלה גמישה' מופעלת, הפלט הוא 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: בקובץ מוגדרת מכסה, שהיא מגבלה על סך המשאבים (כמו CPU, זיכרון ושבבי 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.
  • מאחסנים את הסקריפט הזה ב-ConfigMap של Kubernetes.
  • פורסים 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 אפשר להגדיר שהעבודה תפעל במאגר הצמתים של תוכנית ב'. צמתים של DWS עם התחלה גמישה מסומנים במיוחד (מוכתמים) על ידי GKE, כדי שעומסי עבודה רגילים לא יפעלו עליהם בטעות. לכן, העבודה שלכם צריכה להיות סובלנית במפורש לכתם 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 לתוכנית הזו.
    • מכיוון שהמשימה Job 1 צורכת את המשאבים של Plan A, ‏ Kueue מקצה אוטומטית את המשימה Job 2 למאגר Plan B (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 ולוודא ששתי המשימות פעלו בהצלחה.

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