הגדרת 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 .

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

יצירת אשכול 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 מופעל. תוסף Ray TPU webhook מותקן אוטומטית במישור הבקרה של האשכול.

כדי לתקשר עם האשכול, צריך להגדיר את 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 וביומנים הרלוונטיים.

‫Ray מגדיר משאב TPU-{accelerator}-Head כדי לזהות את צומת העובד של Ray שמתאים לערך 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