‫Google משתמשת בטכנולוגיית AI כדי לתרגם תוכן לשפה המועדפת עליך. בתרגומים כאלו עשויות להיות שגיאות.

הגדרת KubeRay עם TPU Trillium

טייס אוטומטי רגילה

במדריך הזה נסביר איך להגדיר את KubeRay עם TPU Trillium ב-Google Kubernetes Engine ‏ (GKE). במאמר הזה מוסבר איך להגדיר תצורות של TPU עם מארח יחיד ועם כמה מארחים, כולל משתני הסביבה הדרושים ומפרטי ה-Pod ל-TPU Trillium.

המדריך הזה מיועד לאדמינים ולמפעילים של פלטפורמות, ולמומחים בתחום הנתונים וה-AI שרוצים ללמוד איך להגדיר את ההפעלה של TPU Trillium באמצעות KubeRay עבור מאגרי צמתים עם מארח יחיד וכמה מארחים. במדריך הזה מוצג איך להריץ סקריפט עם Jax שמאמת את ההפעלה של TPU. במדריך הזה לא נפרס מודל.

לפני שמגדירים את KubeRay ב-GKE, חשוב להכיר את ההגדרות והטרמינולוגיה של Ray ב-GKE.

סקירה כללית

במדריך הזה מוסבר איך להריץ סקריפט Python עם Jax כדי לוודא שההפעלה של TPU Trillium עם KubeRay בוצעה בהצלחה. ‫Jax היא ספרייה של חישובים מספריים עם ביצועים גבוהים שתומכת בעומסי עבודה של למידת מכונה. ‫KubeRay הוא אופרטור של Kubernetes שמספק דרך מאוחדת לפריסה, לניהול ולניטור של אפליקציות Ray ב-Kubernetes.

ב-TPU מדור Trillium‏ (v6e) נדרשים משתני סביבה ומפרטי Pod ספציפיים, ששונים מאלה של דורות קודמים של TPU. במדריך הזה מוסבר איך לבצע את ההגדרות הנדרשות כדי לפרוס בהצלחה עומס עבודה באמצעות KubeRay ב-TPU של Trillium.

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

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

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

הפעלת Google Kubernetes Engine API

כדי להשתמש ב-CLI של Google Cloud למשימה הזו, צריך להתקין ואז לאתחל את ה-CLI של gcloud. אם התקנתם בעבר את ה-CLI של gcloud, מריצים את הפקודה gcloud components update כדי לקבל את הגרסה העדכנית. יכול להיות שגרסאות קודמות של ה-CLI של gcloud לא יתמכו בהרצת הפקודות שמופיעות במסמך הזה.
הערה: אם כבר התקנתם את ה-CLI של gcloud, הקפידו להגדיר את compute/region המאפיין. אם אתם משתמשים בעיקר באשכולות אזוריים, צריך להגדיר את compute/zone במקום זאת. אם מגדירים מיקום ברירת מחדל, אפשר להימנע משגיאות ב-CLI של gcloud כמו השגיאה הבאה: One of [--zone, --region] must be supplied: Please specify location. יכול להיות שתצטרכו לציין את המיקום בפקודות מסוימות אם המיקום של האשכול שונה מברירת המחדל שהגדרתם.

מוודאים שמותקן Ray CLI (גרסה 2.37.0).

הפעלת Cloud Shell

Cloud Shell מגיע עם כלי שורת הפקודה gcloud, helm ו-kubectl שמשמשים במדריך הזה.

עוברים אל Google Cloud המסוף.
בחלק העליון של חלון המסוף Google Cloud , לוחצים על הלחצן Activate Cloud Shell .

במסוף Google Cloud ייפתח סשן של Cloud Shell בתוך מסגרת חדשה, ותופיע הודעה של שורת הפקודה.

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

אפשר להגדיר את KubeRay ב-TPU באשכול GKE Autopilot או באשכול רגיל. מומלץ להשתמש באשכול Autopilot כדי ליהנות מחוויית Kubernetes מנוהלת באופן מלא. כדי לבחור את מצב הפעולה של GKE שהכי מתאים לעומסי העבודה שלכם, אפשר לעיין במאמר מידע על מצבי הפעולה של GKE.

טייס אוטומטי

ב-Cloud Shell, מריצים את הפקודה הבאה:
```
gcloud container clusters create-auto CLUSTER_NAME \
    --enable-ray-operator \
    --release-channel=rapid \
    --location=LOCATION
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫CLUSTER_NAME: השם של האשכול החדש.
- ‫LOCATION: האזור שבו קיבולת TPU Trillium זמינה. מידע נוסף מופיע במאמר זמינות של TPU ב-GKE.
‫GKE יוצר אשכול Autopilot עם התוסף Ray operator מופעל. התוסף מתקין באופן אוטומטי את ה-webhook של Ray TPU במישור הבקרה של האשכול.

כדי לתקשר עם האשכול, צריך להגדיר את kubectl :

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

רגילה

ב-Cloud Shell, יוצרים אשכול Standard שמפעיל את תוסף Ray operator על ידי הפעלת הפקודה הבאה :
```
gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --addons=RayOperator \
  --cluster-version=1.33 \
  --machine-type=n1-standard-16
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫CLUSTER_NAME: השם של האשכול החדש.
- ‫LOCATION: האזור שבו קיבולת TPU Trillium זמינה. מידע נוסף זמין במאמר זמינות של TPU ב-GKE.
יצירת האשכול עשויה להימשך כמה דקות.

כדי לתקשר עם האשכול, צריך להגדיר את kubectl :

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

אפשר ליצור מאגר צמתים של פרוסת TPU עם מארח יחיד או עם כמה מארחים:

מארח יחיד

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

gcloud container node-pools create v6e-4 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=1 \
    --threads-per-core=1 \
    --tpu-topology=2x2

כמה מארחים

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

gcloud container node-pools create v6e-16 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=4 \
    --threads-per-core=1 \
    --tpu-topology=4x4

הפעלת משאב מותאם אישית של RayJob

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

יצירת RayCluster: מפרט RayJob כולל rayClusterSpec שמגדיר את ההגדרה של אשכול Ray (קבוצות ראשוניות וקבוצות של עובדים) שרוצים.
הפעלת משימה ספציפית: השדה entrypoint ב-RayJob מציין את הפקודה או הסקריפט להפעלה באשכול Ray שנוצר. במדריך הזה, entrypoint הוא סקריפט Python ‏ (tpu_list_devices.py) שנועד לאמת את האתחול של TPU Trillium.

כדי ליצור משאב מותאם אישית של RayJob:

מארח יחיד

יוצרים את קובץ המניפסט ray-job.tpu-v6e-singlehost.yaml הבא:

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: v6e-4-job
spec:
  entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
  runtimeEnvYAML: |
    working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
    pip:
      - jax[tpu]==0.4.33
      - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
  rayClusterSpec:
    rayVersion: '2.43.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
          -   name: ray-head
              image: rayproject/ray:2.43.0-py310
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "8"
                  memory: 40G
                requests:
                  cpu: "8"
                  memory: 40G
    workerGroupSpecs:
    -   replicas: 1
        minReplicas: 1
        maxReplicas: 1
        numOfHosts: 1
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            containers:
            -   name: ray-worker
                image: rayproject/ray:2.43.0-py310
                resources:
                  limits:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
                  requests:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 2x2

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

kubectl apply -f ray-job.tpu-v6e-singlehost.yaml

מוודאים ש-RayJob נוצר ופועל:

kubectl get rayjobs v6e-4-job

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

NAME      JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME       START TIME  END TIME   AGE
v6e-4-job PENDING      Running             v6e-4-job-raycluster   2024-10-15T23:15:22Z  20s

מדפיסים את הפלט של RayJob.

kubectl logs -l=job-name=v6e-4-job

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

2024-10-15 16:15:40,222 INFO cli.py:300 -- ray job stop v6e-4-job-hzq5q
2024-10-15 16:15:40,246 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
2024-10-15 16:15:40,112 INFO job_manager.py:528 -- Runtime env is setting up.
2024-10-15 16:15:50,181 INFO worker.py:1461 -- Using address 10.84.1.25:6379 set in the environment variable RAY_ADDRESS
2024-10-15 16:15:50,181 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.25:6379...
2024-10-15 16:15:50,186 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.25:8265
['TPU cores:4']
2024-10-15 16:16:12,349 SUCC cli.py:63 -- -------------------------------------
2024-10-15 16:16:12,349 SUCC cli.py:64 -- Job 'v6e-4-job-hzq5q' succeeded
2024-10-15 16:16:12,349 SUCC cli.py:65 -- -------------------------------------

כמה מארחים

יוצרים את קובץ המניפסט ray-job.tpu-v6e-multihost.yaml הבא:

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: v6e-16-job
spec:
  entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
  runtimeEnvYAML: |
    working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
    pip:
      - jax[tpu]==0.4.33
      - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
  rayClusterSpec:
    rayVersion: '2.43.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
          -   name: ray-head
              image: rayproject/ray:2.43.0-py310
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "8"
                  memory: 40G
                requests:
                  cpu: "8"
                  memory: 40G
    workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 1
        numOfHosts: 4
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            containers:
              - name: ray-worker
                image: rayproject/ray:2.43.0-py310
                resources:
                  limits:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
                  requests:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
                env:
                - name: NODE_IP
                  valueFrom:
                    fieldRef:
                      fieldPath: status.hostIP
                - name: VBAR_CONTROL_SERVICE_URL
                  value: $(NODE_IP):8353
                - name: JAX_PLATFORMS
                  value: tpu,cpu
                - name: ENABLE_PJRT_COMPATIBILITY
                  value: "true"
                ports:
                - containerPort: 8081
                  name: mxla
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

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

kubectl apply -f ray-job.tpu-v6e-multihost.yaml

מוודאים ש-RayJob‏ v6e-16 נוצר ופועל:

kubectl get rayjobs v6e-16-job

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

NAME         JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME              START TIME             END TIME   AGE
v6e-16-job                Running             v6e-16-job-raycluster-qr6vk   2024-10-16T19:28:19Z              66s

מדפיסים את הפלט של v6e-16 RayJob:

kubectl logs -l=job-name=v6e-16-job

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

2024-10-16 12:21:33,986 INFO cli.py:300 -- ray job stop v6e-16-job-z44s7
2024-10-16 12:21:34,011 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
2024-10-16 12:21:33,826 INFO job_manager.py:528 -- Runtime env is setting up.
2024-10-16 12:21:46,327 INFO worker.py:1461 -- Using address 10.84.1.61:6379 set in the environment variable RAY_ADDRESS
2024-10-16 12:21:46,327 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.61:6379...
2024-10-16 12:21:46,333 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.61:8265
['TPU cores:16', 'TPU cores:16', 'TPU cores:16', 'TPU cores:16']
2024-10-16 12:22:12,156 SUCC cli.py:63 -- ---------------------------------
2024-10-16 12:22:12,156 SUCC cli.py:64 -- Job 'v6e-16-job-z44s7' succeeded
2024-10-16 12:22:12,156 SUCC cli.py:65 -- ---------------------------------

הצגת RayJob ב-Ray Dashboard

מוודאים ש-GKE יצר את שירות RayCluster ומתחברים למופע RayCluster.

מארח יחיד

מאחזרים את השם של RayCluster שנוצר עבור RayJob:

export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-4-job -o jsonpath='{.status.rayClusterName}')

מאחזרים את השם של שירות ה-head של RayCluster:

export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')

מתחברים ללוח הבקרה של Ray באמצעות העברת יציאה של שירות הראשי:
```
kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
```
פותחים דפדפן אינטרנט ומזינים את כתובת ה-URL הבאה:
```
http://localhost:8265/#/jobs
```
צפייה בסטטוס של RayJob וביומנים הרלוונטיים.

כמה מארחים

מאחזרים את השם של RayCluster שנוצר עבור RayJob:

export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-16-job -o jsonpath='{.status.rayClusterName}')

מאחזרים את השם של שירות ה-head של RayCluster:

export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')

מתחברים ללוח הבקרה של Ray באמצעות העברת יציאה של שירות הראשי:
```
kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
```
פותחים דפדפן אינטרנט ומזינים את כתובת ה-URL הבאה:
```
http://localhost:8265/#/jobs
```
צפייה בסטטוס של RayJob וביומנים הרלוונטיים.

ריי מגדיר משאב TPU-{accelerator}-Head כדי לזהות את צומת העובד של ריי שמתאים לערך TPU_WORKER_ID=0. בקבוצת TPU עם כמה מארחים, הצומת של Ray עם TPU_WORKER_ID=0 כולל את הערך TPU-v6e-16-head: 1.0 במשאבים שלו. משתנה הסביבה TPU_WORKER_ID הזה מוגדר על ידי webhook לשינוי ב-GKE עבור KubeRay.

הסרת המשאבים

כדי למנוע חיובים לא רצויים בחשבון אחרי שמסיימים את המדריך, צריך למחוק את RayJob:

מארח יחיד

kubectl delete rayjobs v6e-4-job

כמה מארחים

kubectl delete rayjobs v6e-16-job

הגדרת KubeRay עם TPU Trillium קל לארגן דפים בעזרת אוספים אפשר לשמור ולסווג תוכן על סמך ההעדפות שלך.

סקירה כללית

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

הפעלת Cloud Shell

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

טייס אוטומטי

רגילה

מארח יחיד

כמה מארחים

הפעלת משאב מותאם אישית של RayJob

מארח יחיד

כמה מארחים

הצגת RayJob ב-Ray Dashboard

מארח יחיד

כמה מארחים

הסרת המשאבים

מארח יחיד

כמה מארחים

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

הגדרת KubeRay עם TPU Trillium