עומסי עבודה אינטראקטיביים של Pathways הם עומסי עבודה של JAX מרחוק שפועלים במכונה וירטואלית שלא שייכת לאשכול GKE שמארח את אשכול Pathways. בניגוד לעומסי עבודה של אצווה, השלמת הפעולה של עומס עבודה אינטראקטיבי לא גורמת להשבתה של רכיבי אשכול Pathways, שנשארים זמינים לחיבור על ידי לקוחות JAX אחרים. במסמך הזה אנחנו משתמשים ב-Jupyter notebook כדוגמה כדי להדגים עומסי עבודה אינטראקטיביים.
משתמשי JAX שולחים פקודות לאשכול Pathways באמצעות הממשק IFRT. קוד JAX, בין אם הוא מופעל ממסוף, מ-notebook או מכל סביבה שתואמת ל-Python, יכול ליצור אינטראקציה חלקה עם משאבי Pathways.
לפני שמתחילים
חשוב לוודא שיש לכם:
- יצירת אשכול GKE באמצעות XPK
- קובץ XPK שהותקן
- כלים של Kubernetes שהותקנו
- התקנתם את ה-CLI של gcloud
- הפעלתם את Google Kubernetes Engine API
- מוודאים שהאפשרות Pathways (נתיבי למידה) מופעלת בפרויקט Google Cloud
הפעלת מסלולי למידה במצב אינטראקטיבי
מגדירים את משתני הסביבה הבאים:
export WORKLOAD=WORKLOAD export WORKLOAD_NODEPOOL_COUNT=WORKLOAD_NODEPOOL_COUNT export TPU_TYPE=TPU_TYPE export PROJECT_ID=PROJECT export ZONE=ZONE \ export CLUSTER=CLUSTER
מחליפים את מה שכתוב בשדות הבאים:
-
WORKLOAD: מגדירים שם ייחודי כדי לזהות את עומס העבודה -
WORKLOAD_NODEPOOL_COUNT: מספר מאגרי הצמתים שמשמשים את עומס העבודה של Pathways -
TPU_TYPE: סוג ה-TPU מציין את הגרסה והגודל של Cloud TPU שרוצים ליצור. מידע נוסף על סוגי TPU נתמכים לכל גרסת TPU זמין במאמר גרסאות TPU. -
PROJECT: מזהה הפרויקט ב- Google Cloud -
ZONE: האזור שבו אתם מתכננים להריץ את עומס העבודה -
CLUSTER: השם של אשכול GKE
-
יוצרים את מאגרי הנתונים של Pathways באשכול. כדי להריץ עומס עבודה ללא ממשק משתמש, מריצים את הפקודה הבאה:
xpk workload create-pathways \ --headless \ --workload=${WORKLOAD} \ --num-slices=${WORKLOAD_NODEPOOL_COUNT} \ --tpu-type=${TPU_TYPE} \ --project=${PROJECT} \ --zone=${ZONE} \ --cluster=${CLUSTER}
בשלב הזה, עומס העבודה של JAX יכול להתחבר לשרת ה-Proxy של IFRT.
חיבור לאשכול Pathways במצב אינטראקטיבי
אפשר להתחבר לאשכול Pathways עם או בלי העברת פורטים. כדי להתחבר לאשכול Pathways, משתמשים באחד מהקטעים הבאים.
חיבור באמצעות העברת יציאות
בשלב הזה אפשר להשתמש בהעברת פורטים (מכל מארח עם גישה למישור הבקרה של האשכול) כדי לגשת לשרת ה-proxy:
משתמשים בפקודה שמתאימה לעומס העבודה:
XPK
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
kubectl
PROXY_POD=$(kubectl get pods | grep pathways-${USER}-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
באותו מארח, פותחים חלון טרמינל חדש. מגדירים את משתני הסביבה 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
חיבור ממארחים ב-VPC בלי להשתמש בהעברת יציאות
אם אתם לא רוצים להשתמש בהעברת יציאות, אתם יכולים להתחבר לאשכול Pathways באמצעות Cloud DNS או מאזן עומסים פנימי.
התחברות באמצעות Cloud DNS
הפעלת Cloud DNS באשכול מעבירה את ספק ה-DNS של Cloud מ-kube-dns ל-Cloud DNS. כשמפעילים את האפשרות הזו, נוצר תחום DNS פרטי ב-Cloud DNS בענן הווירטואלי הפרטי (VPC) שלכם עבור השמות של Cloud DNS. מידע נוסף מופיע במאמר שימוש ב-Cloud DNS ל-GKE.
אם מפעילים את Cloud DNS עם האשכול, עם VPC מצטבר או עם היקף VPC, אפשר לפתור את שמות ה-Cloud DNS של Kubernetes ממכונות וירטואליות שאינן GKE בתוך הענן הווירטואלי הפרטי. השמות הם בפורמט <service_name>.<namespace>.svc.<custom_dns_domain>. ל-Pod הראשי של Pathways יש שירות בשם <jobset_name>-pathways-head-0-0.<jobset_name>.<namespace>.svc.<custom_dns_domain>.
הפקודות הבאות מראות איך להתחבר לאשכול Pathways באמצעות Cloud DNS:
מוודאים שאפשר ליצור רזולוציה של רשומת Cloud DNS של הלידר ממארח שאינו GKE:
XPK
host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.USERNAME-testהפלט אמור להיראות כך:
<WORKLOAD>-pathways-head-0-0.<WORKLOAD>.default.svc.<user>-test has address 10.0.2.75
kubectl
host pathways-USERNAME-pathways-head-0-0.pathways-USERNAME.default.svc.USERNAME-testהפלט אמור להיראות כך:
pathways-<user>-pathways-head-0-0.pathways-<user>.default.svc.<user>-test has address 10.0.2.75
מתחברים לאשכול Pathways באמצעות שם Cloud DNS:
XPK
JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.USERNAME-test:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'
kubectl
JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://pathways-USERNAME-pathways-head-0-0.pathways-USERNAME.default.svc.USERNAME-test: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, צריך ליצור שירות שמגובה על ידי מאזן עומסים פנימי. לא צריך להפעיל את Cloud DNS באשכול.
אם אתם יוצרים מאזני עומסים פנימיים, מומלץ להפעיל את התכונה 'חלוקת משנה של ILB' באשכולות עם הרבה מכונות וירטואליות. מידע נוסף זמין במאמר הפעלת חלוקת משנה של GKE באשכול קיים. אם לא מפעילים את האפשרות 'חלוקת משנה של איזון עומסים פנימי', כל הצמתים באשכול יהיו חלק מקבוצת מופעים של השרת העורפי לכל מאזני העומסים הפנימיים. הפתרון הזה לא מתאים לשימוש עם יותר מ-250 צמתים. כשהתכונה 'חלוקת משנה של איזון עומסים פנימי' מופעלת, GKE יוצר קבוצות של נקודות קצה ברשת במקום קבוצות של מכונות וירטואליות, וכולל רק צמתים שמריצים אחד מה-Pods של השירות. ההפעלה של חלוקת משנה של ILB כרוכה בעיכוב חד-פעמי בהגדרה (כ-15 דקות). הפקודה הבאה מראה איך להפעיל את חלוקת המשנה של ILB:
gcloud container clusters update ${CLUSTER} \ --project=${PROJECT} \ [--zone=${ZONE} | --region=${REGION}] \ --enable-l4-ilb-subsetting
אחרי שמפעילים את התכונה 'חלוקת משנה של איזון עומסים פנימי', אפשר ליצור שירות Kubernetes מסוג LoadBalancer באמצעות קובץ ה-YAML הבא. הפעולה הזו תגרום ל-GKE ליצור מאזן עומסים פנימי בתוך ה-VPC של האשכול:
apiVersion: v1 kind: Service metadata: name: pathways-USERNAME-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: pathways-USER jobset.sigs.k8s.io/replicatedjob-name: pathways-head ports: - name: tcp-port protocol: TCP port: 29000 targetPort: 29000
מעדכנים את USER עם מזהה המשתמש שלכם Google Cloud ושומרים את הקובץ בשם pathways-headless-ilb.yaml.
החלת המניפסט:
kubectl apply -f pathways-headless-ilb.yaml
אחרי שמאזן העומסים נוצר (כדקה לאחר מכן), בעמודה EXTERNAL-IP יופיע ערך:
kubectl get services
NAME TYPE CLUSTER-IP EXTERNAL-IP PORT(S) AGE
pathways-$USER ClusterIP None <none> <none> 30m
pathways-$USER-ilb LoadBalancer 34.118.232.46 10.0.0.22 80:31246/TCP 2m41s
אפשר לגשת לפריסת נתיבי הלימוד בלי להגדיר העברת נתונים דרך יציאה במארחים באותה רשת וירטואלית פרטית (VPC) כמו האשכול:
JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://10.0.0.22:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'
הפלט אמור להיראות כך:
[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
מחברות Jupyter
אפשר ליצור מחברת Jupyter באמצעות Vertex AI או ליצור מחברת Jupyter באירוח עצמי.
יצירת מכונה של Vertex AI Workbench
אחרי שמגדירים ומאמתים את אשכול Pathways, אפשר לגשת למכונות וירטואליות של GKE TPU ממחברת Vertex AI Jupyter. ההנחיות הבאות להגדרה מניחות שקלאסטר GKE Pathways נמצא באותה רשת ענן וירטואלי פרטי (VPC) (שהיא רשת ברירת המחדל, אלא אם הגדרתם אחרת). עוברים אל מסוף Vertex AI Workbench.
יוצרים מכונת Workbench חדשה (מהכרטיסייה Instances) באמצעות הלחצן Create new. מוודאים שהרשת זהה לרשת של אשכול GKE. אפשר להשתמש בשורת הפקודה כדי ליצור מכונת Workbench חדשה.
gcloud workbench instances create INSTANCE_NAME \ --machine-type=e2-standard-4 \ --data-disk-size=100 \ --location=ZONE \ [--network=NETWORK]
אחרי שהמופע נוצר, עוברים אליו ולוחצים על Open Jupyterlab.
יצירת מכונה של Jupyter Notebook באירוח עצמי
הפקודה הבאה מראה איך ליצור מופע של מחברת Jupyter באירוח עצמי באמצעות XPK:
xpk workload create-pathways \
--workload=${WORKLOAD} \
--num-slices=${WORKLOAD_NODEPOOL_COUNT} \
--tpu-type=${TPU_TYPE} \
--project=${PROJECT} \
--zone=${ZONE} \
--cluster=${CLUSTER} \
--docker-image=jupyter/base-notebook \
--command "start-notebook.sh"
קובץ ה-YAML הבא מראה איך ליצור מופע של מחברת Jupyter באירוח עצמי באמצעות kubectl. אחרי שיוצרים אשכול Pathways ללא GUI, מחילים את קובץ ה-YAML הבא. מידע נוסף זמין במאמר בנושא הפעלת Pathways במצב אינטראקטיבי באמצעות kubectl.
apiVersion: batch/v1 kind: Job metadata: name: jupyter-notebook-USERNAME spec: template: spec: restartPolicy: OnFailure containers: - name: jupyter-notebook image: jupyter/base-notebook # Use the appropriate Jupyter image ports: - containerPort: 8888
ניווט למופע של ה-notebook
מתחברים ל-notebook מהמחשב המקומי באמצעות העברת יציאות:
XPK
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 | grep jupyter-notebook-USERNAME | awk '{print $1}')
kubectl port-forward pod/${MAIN_POD} 8888:8888
בדפדפן המקומי, עוברים אל http://localhost:8888?token=<var>your-token</var>.
מחליפים את <your-token> באסימון מיומני הרישום של קונטיינר ה-notebook של Jupyter.
kubectl logs ${MAIN_POD}
הפלט שיתקבל:
... Or copy and paste one of these URLs: http://jupyter-notebook-<user>-bbbdh:8888/lab?token=<token> http://127.0.0.1:8888/lab?token=<token>
קישוריות של מחברת לאשכול Pathways
- מתוך Jupyterlab, יוצרים נוטבוק חדש של Python 3.
- התחברות לשרת ה-proxy של Pathways
במחברת, מוסיפים תא להתקנת pathwaysutils, מגדירים את JAX_PLATFORMS לערך
proxy ומגדירים את JAX_BACKEND_TARGET לערך PROXY_ADDRESS.
!pip install pathwaysutils %env JAX_PLATFORMS=proxy # Replace your proxy address below: %env JAX_BACKEND_TARGET=PROXY_ADDRESS
מוסיפים תא שני כבדיקה מסוג 'שלום עולם' ומדפיסים את המכשירים באשכול Pathways.
import pathwaysutils
import jax
pathwaysutils.initialize()
print(jax.devices())
אם הכול פועל בצורה תקינה, תוצג הודעה שמציינת שהמערכת האחורית של Pathways-on-Cloud זוהתה.
מספר מכשירי ה-JAX שמופיעים צריך להיות זהה למספר שבבי ה-TPU ולמספר הפרוסות שציינתם כשנוצר אשכול Pathways.
הוספת קוד ל-Notebook
אפשר להוסיף קוד JAX משלכם ולהריץ אותו באופן אינטראקטיבי ב-TPU באשכול Pathways. בדוגמה הבאה אפשר לראות איך לבצע חישובים בשני פלחים מתוך מחברת אחת.
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.
jax.jit(jnp.sin, device=jax.devices()[-1])(np.pi / 2.)
# 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)
מחיקת אשכול אינטראקטיבי של תוכניות לימודים
XPK
xpk workload delete --workload=WORKLOAD --cluster=CLUSTER --project=PROJECT --zone=ZONE
kubectl
kubectl delete -f pathways-headless-workload.yaml
המאמרים הבאים
- יצירת אשכול GKE באמצעות Pathways
- הסקת מסקנות מרובת מארחים באמצעות Pathways
- עומסי עבודה באצווה עם Pathways
- העברת עומסי עבודה של JAX ל-Pathways
- אימון גמיש עם תוכניות לימודים
- פתרון בעיות