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

הפעלת עומס עבודה אינטראקטיבי באמצעות Pathways

עומסי עבודה אינטראקטיביים של Pathways הם עומסי עבודה מרוחקים של JAX שפועלים מחוץ לרכיבי הבקר של Pathways באשכול, למשל במכונה וירטואלית. בניגוד לעומסי עבודה באצווה, כשעומס עבודה אינטראקטיבי מסתיים, הוא לא משבית את רכיבי אשכול Pathways, שנשארים זמינים לחיבורים מלקוחות JAX אחרים. במאמר הזה אנחנו משתמשים ב-notebook של Jupyter כדוגמה לעומסי עבודה אינטראקטיביים. אתם יכולים להריץ עומסי עבודה אינטראקטיביים עם Pathways באשכול של Google Kubernetes Engine ‏(GKE) או באשכול של Cluster Director עם Slurm.

בדף הזה יש הוראות לשתי הפלטפורמות. בכל קטע, בוחרים בכרטיסייה GKE או Cluster Director.

באמצעות ממשק IFRT, משתמשי JAX שולחים פקודות לאשכול Pathways. קוד JAX, בין אם הוא מופעל ממסוף, מ-notebook או מכל סביבה שתואמת ל-Python, יכול לקיים אינטראקציה עם משאבי Pathways.

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

צריך לבחור את הכרטיסייה הרלוונטית לאופן שבו תכננתם להשתמש בדוגמאות בדף הזה:
gcloud

התקינו את ה-CLI של Google Cloud ואז היכנסו ל-CLI של gcloud באמצעות הזהות המאוחדת שלכם. אחרי שנכנסתם לחשבון, אתחלו את ה-CLI של Google Cloud באמצעות הפקודה הבאה:
```
gcloud init
```
REST

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

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

GKE

מבצעים את הפעולות הבאות:

יצירת אשכול GKE עם Pathways.
מתקינים את Cluster Toolkit.
התקנת כלי Kubernetes.
מתקינים את Google Cloud CLI.
מפעילים את Kubernetes Engine API.
מוודאים ש-Pathways מופעל בפרויקט Google Cloud .

Cluster Director

מבצעים את הפעולות הבאות:

יוצרים אשכול ב-Cluster Director ומגדירים את סביבת Pathways.
מתקינים את ה-CLI של gcloud.
מוודאים שיש לכם גישת SSH לצומת ההתחברות של האשכול.
מתקינים את Python 3 עם המודול venv במארח שמריץ את לקוח JAX. ב-Debian וב-Ubuntu, מתקינים את החבילה python3-venv.

הפעלת Pathways במצב אינטראקטיבי

GKE

מגדירים את משתני הסביבה הבאים:
```
export WORKLOAD=WORKLOAD
export WORKLOAD_NODEPOOL_COUNT=WORKLOAD_NODEPOOL_COUNT
export COMPUTE_TYPE="COMPUTE_TYPE"
export TOPOLOGY="TOPOLOGY"
export PROJECT_ID=PROJECT_ID
export REGION=REGION
export ZONE=ZONE
export CLUSTER=CLUSTER
export BUCKET_NAME=BUCKET_NAME
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫WORKLOAD: שם ייחודי לזיהוי עומס העבודה (עד 22 תווים, בגלל מגבלות התוויות של DNS).
- ‫WORKLOAD_NODEPOOL_COUNT: מספר מאגרי הצמתים שמשמשים לעומס עבודה ב-Pathways.
- ‫COMPUTE_TYPE: סוג המכונה של ה-TPU, לדוגמה ct6e-standard-4t. למידע נוסף על סוגי ה-TPU שנתמכים בכל גרסת TPU, קראו את המאמר גרסאות TPU.
- ‫TOPOLOGY: טופולוגיית המיקום של ה-TPU, לדוגמה 2x4.
- ‫PROJECT_ID: מזהה הפרויקט ב-Google Cloud .
- ‫REGION: האזור שבו נמצאים הקטגוריה והאשכול.
- ‫ZONE: האזור שבו מתכננים להריץ את עומס העבודה.
- ‫CLUSTER: השם של אשכול GKE.
- ‫BUCKET_NAME: הקטגוריה של Cloud Storage לנכסי Pathways.
אם עדיין לא יצרתם קטגוריה של Cloud Storage בשביל פריטי המידע שנוצרו בתהליך של Pathways, אתם צריכים ליצור אחת:
```
gcloud storage buckets create gs://${BUCKET_NAME} \
    --project=${PROJECT_ID} \
    --location=${REGION}
```
יוצרים את מאגרי Pathways באשכול. כדי להריץ עומס עבודה ללא ממשק משתמש, משתמשים בפקודה gcluster job submit (למידע נוסף, עיינו במדריך לשליחת משימות ב-Cluster Toolkit):
```
gcluster job submit \
    --pathways \
    --pathways-gcs-location="gs://${BUCKET_NAME}/pathways-artifacts" \
    --name=${WORKLOAD} \
    --cluster=${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
    --compute-type=${COMPUTE_TYPE} \
    --topology=${TOPOLOGY} \
    --pathways-headless
```
הפקודה gcluster job submit מחפשת מאגר צומתי CPU בשם cpu-np או pathways-np כדי להריץ את ה-Pod הראשי של Pathways. אם למאגר הצמתים של המעבד (CPU) יש שם אחר (כמו cpu-pathways-np מהמדריך ליצירת אשכול), הוסיפו את הדגל --pathways-head-np=NODE_POOL_NAME לפקודות gcluster job submit שבדף הזה. אחרת, הפקודה תיכשל ותתקבל השגיאה failed to auto-detect a suitable CPU node pool.

בשלב הזה, עומס העבודה של JAX יכול להתחבר לשרת ה-proxy של IFRT.

Cluster Director

מתחברים לצומת ההתחברות של האשכול:
```
gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫CLUSTER_ID: השם של האשכול ב-Cluster Director.
- ‫PROJECT_ID: מזהה הפרויקט ב-Google Cloud .
- ‫ZONE: האזור שבו נמצא האשכול.
מגדירים את משתני הסביבה הבאים:
```
export SLICE_SHAPE=TOPOLOGY
export MAX_SLICES=NUM_SLICES
export MIN_SLICES=NUM_SLICES
export GCS_BUCKET=gs://BUCKET_NAME
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫TOPOLOGY: הטופולוגיה של TPU בכל פרוסה, לדוגמה 2x4. הוא צריך להתאים לטופולוגיה של קבוצת הצמתים של ה-TPU.
- ‫NUM_SLICES: מספר פרוסות ה-TPU. מגדירים את MIN_SLICES ו-MAX_SLICES לאותו ערך כדי שאשכול Pathways יהיה מוכן רק כשכל הפרוסות פועלות.
- ‫BUCKET_NAME: קטגוריה קיימת של Cloud Storage ל-Pathways artifacts.
אם מחיצת ה-TPU היא לא מחיצת ברירת המחדל של האשכול, צריך גם להגדיר את PARTITION לשם שלה.
מפעילים את אשכול Pathways במצב headless:
```
cd ~/pathways
sbatch start_pathways.sh
```
מחכים עד שיופיע פלט דומה לזה שמופיע בהמשך ביומן השירות, pathways-service-JOB_ID.out (כאשר JOB_ID הוא מזהה המשימה ב-Slurm):
```
Pathways Runtime Service is READY!
 Discovery file    : /home/USER/pathways/cluster.env
 IFRT Proxy Target : grpc://RM_IP:29000
```
‫RM_IP היא כתובת ה-IP הפנימית של שרת ה-proxy. הקובץ ~/pathways/cluster.env מכיל את אותה כתובת.

בשלב הזה, עומס העבודה של JAX יכול להתחבר לשרת ה-proxy של IFRT.

התחברות לאשכול Pathways במצב אינטראקטיבי

אתם יכולים להתחבר לאשכול Pathways עם או בלי העברה ליציאה אחרת. כדי להתחבר לאשכול Pathways, השתמשו באחד מהקטעים הבאים.

ב-Cluster Director, לקוח JAX חייב להשתמש באותה גרסת JAX כמו שרת Pathways, ‏0.11.1.

התחברות באמצעות העברה ליציאה אחרת

GKE

מכל מארח עם גישה למישור הבקרה של האשכול, מעבירים את היציאה המקומית 29000 ל-Pod הראשי של Pathways כדי לגשת לשרת ה-proxy:

PROXY_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
PROXY_PORT=29000
kubectl port-forward ${PROXY_POD} ${PROXY_PORT}:${PROXY_PORT}

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

Forwarding from 127.0.0.1:29000 -> 29000
Forwarding from [::1]:29000 -> 29000

באותו מארח, פותחים חלון טרמינל חדש. באמצעות Python 3.11 ואילך, יוצרים סביבה וירטואלית, מתקינים את החבילות pathwaysutils ו-jax, מגדירים את משתני הסביבה JAX_PLATFORMS ו-JAX_BACKEND_TARGET ומריצים סקריפט Python שמייבא את pathwaysutils ו-jax:

python3 -m venv .venv
source .venv/bin/activate
pip install pathwaysutils jax

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

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

[device(144,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(145,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(146,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(147,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(148,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(149,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(150,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(151,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(162,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(163,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(164,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(165,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(166,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(167,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(168,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(169,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
Waiting up to 5 seconds.
Sent all pending logs.
2024-11-13 21:38:51.267523: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled

Cluster Director

מכל מארח שיכול להתחבר באמצעות SSH לצומת ההתחברות של האשכול, כמו תחנת העבודה או מכונה וירטואלית מחוץ לרשת המשנה של האשכול, מעבירים את יציאת ה-proxy דרך צומת ההתחברות:

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 29000:RM_IP:29000

באותו מארח, פותחים חלון טרמינל חדש. מגדירים את משתני הסביבה JAX_PLATFORMS ו-JAX_BACKEND_TARGET, ומריצים סקריפט Python שמייבא את pathwaysutils ו-jax. למארח אין נתיב אל RM_IP, ולכן הלקוח מתחבר לקצה המקומי של המנהרה, 127.0.0.1:29000:

python3 -m venv .venv
source .venv/bin/activate
pip install pathwaysutils jax==0.11.1

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

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

[device(288,TPU_DEVICE,coords=[0,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 device(289,TPU_DEVICE,coords=[1,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 ...
 device(313,TPU_DEVICE,coords=[1,3,0,0],logical_task=1,slice=1,default_mem=device,mem_spaces=3)]
W1029 22:52:22.897483    2361 grpc_client.cc:91] IFRT proxy server disconnected: CANCELLED: Stream removed (CANCELLED)

האזהרה IFRT proxy server disconnected מופיעה כשהסקריפט יוצא וסוגר את החיבור שלו. זה תקין.

התחברות ממארחים ב-VPC בלי להשתמש בהעברה ליציאה אחרת

GKE

אם אתם לא רוצים להשתמש בהעברה ליציאה אחרת, אתם יכולים להתחבר לאשכול Pathways באמצעות Cloud DNS או מאזן עומסים פנימי.

התחברות באמצעות Cloud DNS

כשמפעילים את Cloud DNS באשכול, ספק ה-DNS משתנה מ-kube-dns ל-Cloud DNS. כשהתכונה מופעלת, GKE יוצר תחום Cloud DNS פרטי ברשת של הענן הווירטואלי הפרטי (VPC) בשביל שמות ה-DNS. מידע נוסף מופיע במאמר שימוש ב-Cloud DNS ל-GKE.

אם מפעילים את Cloud DNS עם היקף VPC או היקף VPC נוסף, אפשר לזהות את שמות ה-DNS של Kubernetes ממכונות וירטואליות שאינן של GKE בתוך רשת ה-VPC. בהיקף של אשכול, רק צמתים באשכול יכולים לפתור את השמות האלה. השמות הם בפורמט SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN. ל-Pod הראשי של Pathways יש שירות בשם WORKLOAD-pathways-head-0-0.WORKLOAD.NAMESPACE.svc.CUSTOM_DNS_DOMAIN.

הפקודות הבאות מראות איך להתחבר לאשכול Pathways באמצעות Cloud DNS:

מוודאים שאפשר לתרגם את רשומת ה-DNS של ה-leader ממארח שאינו GKE:
```
host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫WORKLOAD: השם של עומס העבודה ב-Pathways.
- ‫CUSTOM_DNS_DOMAIN: דומיין ה-DNS בהתאמה אישית שהוגדר לאשכול.
הפלט אמור להיראות כך:
```
WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN has address 10.0.2.75
```

מתחברים לאשכול Pathways באמצעות שם ה-DNS:

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

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

[device(216,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(217,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(218,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(219,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(220,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(221,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(222,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(223,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(234,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(235,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(236,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(237,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(238,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(239,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(240,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(241,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
Waiting up to 5 seconds.
Sent all pending logs.
2024-11-14 00:02:49.882044: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled

חיבור באמצעות מאזן עומסים פנימי

כדי ליצור כתובת IP פרטית ברשת ה-VPC שמפנה לפריסת Pathways, צריך ליצור שירות שמגובה על ידי מאזן עומסים פנימי (ILB). לא צריך להפעיל את Cloud DNS באשכול.

באשכולות עם הרבה מכונות וירטואליות, אם יוצרים מאזני עומסים פנימיים, צריך להפעיל את הגדרת קבוצת המשנה של ILB. מידע נוסף זמין במאמר הפעלת GKE subsetting באשכול קיים. כשלא מפעילים את הגדרת קבוצת המשנה של ILB, כל הצמתים באשכול הם חלק מקבוצת המכונות של השרת העורפי של כל מאזני העומסים הפנימיים, שלא יכולה להתרחב מעבר ל-250 צמתים. כשהגדרת המשנה של ILB מופעלת, GKE יוצר קבוצות של נקודות קצה ברשת במקום קבוצות של מכונות, וכולל רק צמתים שמריצים את אחד מה-Pods שמספקים את השירות. ההפעלה של ILB subsetting כרוכה בהשהיה חד-פעמית של כ-15 דקות בהגדרה. הפקודה הבאה מפעילה את הגדרת קבוצת המשנה של ILB:

gcloud container clusters update ${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=CLUSTER_LOCATION \
    --enable-l4-ilb-subsetting

מחליפים את CLUSTER_LOCATION באזור או בתחום שבהם נמצא האשכול.

אחרי שמפעילים את ILB subsetting, יוצרים שירות Kubernetes מסוג LoadBalancer באמצעות מניפסט ה-YAML הבא. כך GKE יוצר מאזן עומסים פנימי בתוך רשת ה-VPC של האשכול:

apiVersion: v1
kind: Service
metadata:
  name: WORKLOAD-ilb
  annotations:
    networking.gke.io/load-balancer-type: "Internal"
    networking.gke.io/internal-load-balancer-allow-global-access: "true"
spec:
  type: LoadBalancer
  externalTrafficPolicy: Local
  selector:
    jobset.sigs.k8s.io/jobset-name: WORKLOAD
    jobset.sigs.k8s.io/replicatedjob-name: pathways-head
  ports:
  - name: tcp-port
    protocol: TCP
    port: 29000
    targetPort: 29000

מחליפים את WORKLOAD בשם של עומס העבודה ב-Pathways, ושומרים את הקובץ בתור pathways-headless-ilb.yaml.

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

kubectl apply -f pathways-headless-ilb.yaml

אחרי ש-GKE יוצר את מאזן העומסים (בערך דקה לאחר מכן), מוודאים שבעמודה EXTERNAL-IP מופיעה כתובת IP:

kubectl get services

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

NAME             TYPE           CLUSTER-IP      EXTERNAL-IP   PORT(S)           AGE
WORKLOAD       ClusterIP      None            <none>        <none>            30m
WORKLOAD-ilb   LoadBalancer   34.118.232.46   10.0.0.22     29000:31246/TCP   2m41s

אתם יכולים לגשת לפריסה של Pathways בלי העברה ליציאה אחרת ממארחים באותה רשת VPC כמו האשכול:

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://EXTERNAL_IP:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

מחליפים את EXTERNAL_IP בערך של EXTERNAL-IP שמוצג על ידי kubectl get services, לדוגמה 10.0.0.22.

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

[device(288,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(289,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(290,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(291,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(292,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(293,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(294,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(295,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(306,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(307,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(308,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(309,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(310,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(311,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(312,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(313,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
Waiting up to 5 seconds.
Sent all pending logs.
2024-11-14 00:30:07.296917: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled

Cluster Director

מארחים ברשת המשנה של האשכול, כמו מכונה וירטואלית ב-Compute Engine, יכולים להתחבר ישירות לשרת ה-proxy. חומת האש של Cluster Director מאפשרת תנועה מרשת המשנה של האשכול, אבל לא מרשתות משנה אחרות של אותה רשת VPC:

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://RM_IP:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

מחליפים את RM_IP בכתובת ה-proxy של IFRT מיומן השירות.

קובצי Jupyter Notebook

GKE

אתם יכולים ליצור notebook של Jupyter באמצעות Vertex AI, או ליצור notebook של Jupyter באירוח עצמי.

יצירת מכונה של Vertex AI Workbench

אחרי שמגדירים ומאמתים את אשכול Pathways, אפשר לגשת למכונות הווירטואליות של TPU ב-GKE מ-Vertex AI Workbench Jupyter notebook. הוראות ההגדרה הבאות מבוססות על ההנחה שאשכול GKE Pathways נמצא באותה רשת VPC (שהיא רשת ברירת המחדל, אלא אם הגדרתם אחרת).

נכנסים לדף Vertex AI Workbench במסוף Google Cloud .
בכרטיסייה Instances, לוחצים על Create new ומוודאים שהרשת תואמת לרשת של אשכול ה-GKE.

לחלופין, אפשר להשתמש בשורת הפקודה כדי ליצור מכונה חדשה של Vertex AI Workbench:
```
gcloud workbench instances create INSTANCE_NAME \
    --project=PROJECT_ID \
    --machine-type=e2-standard-4 \
    --data-disk-size=100 \
    --location=ZONE \
    --network=NETWORK
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫INSTANCE_NAME: השם של המכונה ב-Vertex AI Workbench.
- ‫PROJECT_ID: מזהה הפרויקט ב-Google Cloud .
- ‫ZONE: האזור שבו רוצים ליצור את המכונה.
- ‫NETWORK: השם של רשת ה-VPC של אשכול GKE. הדגל --network הוא אופציונלי. אם משמיטים אותו, הפקודה משתמשת ברשת default.
אחרי ש-Vertex AI יוצר את המכונה, נכנסים אליה ולוחצים על Open JupyterLab.

יצירת מכונה של Jupyter Notebook באירוח עצמי

הפקודה הבאה מראה איך ליצור מכונה של Jupyter Notebook באירוח עצמי באמצעות הפקודה gcluster job submit:

gcluster job submit \
    --pathways \
    --name=${WORKLOAD} \
    --cluster=${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
    --compute-type=${COMPUTE_TYPE} \
    --topology=${TOPOLOGY} \
    --pathways-gcs-location="gs://${BUCKET_NAME}/pathways-artifacts" \
    --image=jupyter/base-notebook \
    --command="start-notebook.sh"

במניפסט הבא של YAML מוסבר איך ליצור מכונה של מחברת Jupyter באירוח עצמי באמצעות kubectl. אחרי שיוצרים אשכול Pathways ללא GUI, צריך להחיל את מניפסט ה-YAML הבא. מידע נוסף זמין במאמר הפעלת Pathways במצב אינטראקטיבי:

apiVersion: batch/v1
kind: Job
metadata:
  name: WORKLOAD
spec:
  template:
    spec:
      hostNetwork: true
      dnsPolicy: ClusterFirstWithHostNet
      restartPolicy: OnFailure
      containers:
      - name: jupyter-notebook
        image: jupyter/base-notebook  # Use the appropriate Jupyter image
        ports:
        - containerPort: 8888

השדה hostNetwork: true מאפשר ללקוח JAX ב-notebook להשלים את לחיצת היד של אבטחת שכבת התעבורה של אפליקציה (ALTS) עם שרת ה-proxy של Pathways. בלי השדה הזה, החיבור נכשל במאגרי צמתים שמשתמשים באיחוד זהויות של עומסי עבודה ל-GKE. השדה dnsPolicy: ClusterFirstWithHostNet מאפשר ל-notebook לפתור שמות DNS בתוך האשכול, כמו השם של ה-Pod הראשי של Pathways. מידע נוסף זמין במאמר פתרון בעיות ב-Pathways ב-Cloud.

Cluster Director

אתם יכולים להריץ Jupyter Notebook במכונה וירטואלית קיימת של Compute Engine ברשת המשנה של האשכול, או ליצור Jupyter Notebook באירוח עצמי בצומת ההתחברות של האשכול. בשני המקרים, ה-notebook מתחבר לשרת ה-proxy בלי העברה ליציאה אחרת.

שימוש במכונה וירטואלית קיימת ברשת המשנה של האשכול

ב-VM, מתקינים את JupyterLab בסביבה וירטואלית ומפעילים אותו:

python3 -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
jupyter lab --no-browser --port=8888

יצירת מכונה של Jupyter Notebook באירוח עצמי

בצומת ההתחברות של האשכול, מתקינים את JupyterLab ומפעילים אותו. יוצרים את הסביבה הווירטואלית בספריית הבית באמצעות Python של האשכול. גרסת ה-Python של המערכת בצומת ההתחברות ישנה מדי לגרסאות הנוכחיות של JAX, ולכן pip לא יכול למצוא שם את jax==0.11.1:

cd ~
/home/pathways/venv/bin/python -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
source ~/pathways/cluster.env
jupyter lab --no-browser --port=8888

התחברות למכונה של Notebook

GKE

מתחברים ל-notebook מהמכונה המקומית באמצעות העברה ליציאה אחרת:

Cluster Toolkit

MAIN_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
kubectl port-forward pod/${MAIN_POD} 8888:8888

kubectl

MAIN_POD=$(kubectl get pods \
    --selector=batch.kubernetes.io/job-name=WORKLOAD \
    --output=jsonpath='{.items[0].metadata.name}')
kubectl port-forward pod/${MAIN_POD} 8888:8888

בדפדפן המקומי, נכנסים לכתובת http://localhost:8888?token=TOKEN. מחליפים את TOKEN בטוקן מהיומנים של קונטיינר ה-notebook של Jupyter:

kubectl logs ${MAIN_POD}

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

...
Or copy and paste one of these URLs:
  http://NODE_NAME:8888/lab?token=TOKEN
  http://127.0.0.1:8888/lab?token=TOKEN

‫NODE_NAME הוא שם המארח של צומת GKE שמריץ את ה-Pod של ה-notebook (לדוגמה, gke-my-cluster-cpu-np-1a2b3c4d-ef56).

Cluster Director

ב-notebook במכונה וירטואלית או בצומת ההתחברות, מתחברים מהמכונה המקומית באמצעות העברה ליציאה אחרת במארח הזה. לדוגמה, לצומת ההתחברות:

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 8888:localhost:8888

בדפדפן המקומי, נכנסים לכתובת http://localhost:8888?token=TOKEN. מחליפים את TOKEN בטוקן ש-JupyterLab הדפיס כשהוא התחיל.

קישוריות של Notebook לאשכול Pathways

ב-JupyterLab, יוצרים קובץ notebook חדש ב-Python 3.
מתחברים לשרת ה-proxy של Pathways.

ב-notebook, מוסיפים תא כדי להתקין את pathwaysutils, מגדירים את JAX_PLATFORMS ל-proxy ומגדירים את JAX_BACKEND_TARGET ל-grpc://PROXY_ADDRESS:29000:

!pip install pathwaysutils
# Set JAX_PLATFORMS and JAX_BACKEND_TARGET if you're not using `gcluster job submit` to deploy the notebook
%env JAX_PLATFORMS=proxy
  # Replace your proxy address in the following line:
%env JAX_BACKEND_TARGET=grpc://PROXY_ADDRESS:29000

מחליפים את PROXY_ADDRESS במארח או בכתובת ה-IP של שרת ה-proxy של Pathways מהתחברות לאשכול Pathways במצב אינטראקטיבי (לדוגמה, 127.0.0.1, שם המארח של Cloud DNS או כתובת ה-IP של מאזן העומסים הפנימי). ב-Cluster Director, משתמשים ב-127.0.0.1 אם משתמשים בהעברה ליציאה אחרת, או בכתובת ה-proxy של IFRT מיומן השירות. אם פרסתם notebook באירוח עצמי באמצעות kubectl, החליפו את PROXY_ADDRESS בשם ה-DNS בתוך האשכול של ה-Pod הראשי של Pathways, כך שהערך של JAX_BACKEND_TARGET יהיה grpc://WORKLOAD-pathways-head-0-0.WORKLOAD:29000.

מוסיפים תא שני כדי לאמת את הקישוריות ולהדפיס את המכשירים באשכול Pathways:

import pathwaysutils
import jax

pathwaysutils.initialize()
print(jax.devices())

אם החיבור מצליח, הפלט כולל הודעה שמציינת ש-JAX זיהה את הקצה העורפי של Pathways-on-Cloud.

מספר מכשירי JAX שמופיעים ברשימה תואם למספר שבבי ה-TPU ולמספר הפרוסות שציינתם כשיצרתם את אשכול Pathways.

הוספת הקוד ל-Notebook

להוסיף קוד JAX משלכם ולהריץ אותו באופן אינטראקטיבי ביחידות ה-TPU באשכול Pathways. הקוד הבא מראה איך לבצע חישובים בשני פרוסות מ-notebook אחד. כדי להריץ את הדוגמה המלאה, יוצרים את עומס העבודה ב-Pathways עם לפחות שתי פרוסות, על ידי הגדרת WORKLOAD_NODEPOOL_COUNT ל-2 או יותר. ב-Cluster Director, מגדירים את MIN_SLICES ואת MAX_SLICES לערך 2 או יותר. עם פרוסה אחת, slice1_devices ריק, והחלק האחרון של הדוגמה נכשל עם ValueError:

import jax
import jax.numpy as jnp
from jax import lax
import numpy as np

# You can use JAX APIs as usual across any of the devices.
x = jax.device_put(np.pi / 2., jax.devices()[-1])
jax.jit(jnp.sin)(x)

# pmap can run across all devices on all slices
num_tpus = jax.device_count()
f = jax.pmap(lambda x: lax.psum(1, 'i'), 'i')
x = jnp.arange(num_tpus)
y = f(x)
print(y)

# You can also target devices from a specific slice
slice0_devices = [d for d in jax.devices() if d.slice_index == 0]
f = jax.pmap(lambda x: lax.psum(1, 'i'), 'i', devices=slice0_devices)
x = jnp.arange(len(slice0_devices))
y = f(x)
print(y)
print(y.global_shards)

# You can send data produced on one slice to another slice
slice1_devices = [d for d in jax.devices() if d.slice_index == 1]
g = jax.pmap(lambda x: x + lax.axis_index('i'), 'i', devices=slice1_devices)
z = g(y)
print(z)
print(z.global_shards)

מחיקת עומס העבודה האינטראקטיבי של Pathways

GKE

Cluster Toolkit

gcluster job cancel WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=ZONE

kubectl

kubectl delete job WORKLOAD

Cluster Director

בצומת ההתחברות של האשכול, מבטלים את משימת השירות של Pathways:

scancel -u ${USER} -n pathways-service

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

נתקלתם בבעיות? קראו את המאמר פתרון בעיות ב-Pathways ב-Cloud.

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

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

gcloud

REST

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

GKE

Cluster Director

הפעלת Pathways במצב אינטראקטיבי

GKE

Cluster Director

התחברות לאשכול Pathways במצב אינטראקטיבי

התחברות באמצעות העברה ליציאה אחרת

GKE

Cluster Director

התחברות ממארחים ב-VPC בלי להשתמש בהעברה ליציאה אחרת

GKE

התחברות באמצעות Cloud DNS

חיבור באמצעות מאזן עומסים פנימי

Cluster Director

קובצי Jupyter Notebook

GKE

יצירת מכונה של Vertex AI Workbench

יצירת מכונה של Jupyter Notebook באירוח עצמי

Cluster Director

התחברות למכונה של Notebook

GKE

Cluster Toolkit

kubectl

Cluster Director

קישוריות של Notebook לאשכול Pathways

הוספת הקוד ל-Notebook

מחיקת עומס העבודה האינטראקטיבי של Pathways

GKE

Cluster Toolkit

kubectl

Cluster Director

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

הפעלת עומס עבודה אינטראקטיבי באמצעות Pathways