במסמך הזה נסביר על כלים ושיטות מומלצות למיקסום השימוש במשאבים ולצמצום זמן ההשבתה של עומסי עבודה הטרוגניים של AI/ML ב-Google Kubernetes Engine (GKE), במיוחד כשאין קיבולת בהזמנות או דרך משאבים על פי דרישה. עומסי עבודה הטרוגניים הם סוגים שונים של עומסי עבודה של AI/ML שפועלים בו-זמנית באותו אשכול GKE. לדוגמה, יכול להיות שתפעילו שירות הסקה אונליין שרגיש לזמן אחזור, לצד סדרה של משימות אימון באצווה שאפשר להפריע להן.
המדריך הזה מיועד לאדמינים ולאופרטורים של פלטפורמות, ולמומחים בתחום הנתונים וה-AI. לסקירה מרוכזת של כל השיטות המומלצות ל-GKE, אפשר לעיין במאמר שיטות מומלצות ל-GKE.
היתרונות של תעדוף עומסי עבודה של AI/ML
לעומסי עבודה הטרוגניים יש עדיפויות שונות, והם חולקים קיבולת ומשאבים מוגבלים. השיטות המומלצות בדף הזה מתארות איך להגדיר את GKE ואת הכלים של קוד פתוח כדי ליהנות מהיתרונות הבאים:
- צמצום זמן ההשבתה של עומסי עבודה בעדיפות גבוהה.
- להריץ במהירות עומסי עבודה בעדיפות גבוהה.
- אופטימיזציה של צריכת המשאבים.
רקע
GKE תומך בכלים הבאים של קוד פתוח לאופטימיזציה של השימוש במשאבים.
Kueue: מערכת תורים של עומסי עבודה (workloads) ב-Kubernetes, שמיועדת לעומסי עבודה של אצווה, AI ומחשוב עתיר ביצועים (HPC). אפשר להרחיב את Kueue כדי לנהל סוגים אחרים של עומסי עבודה, כמו אלה שמוגדרים על ידי הגדרות של משאבים בהתאמה אישית כמו
leaderworkerset. Kueue מנהל את המכסות ואת האופן שבו עומסי העבודה צורכים אותן באשכול Kubernetes. Kueue מקבל החלטות לגבי מתי עומס עבודה ממתין, מתי עומס עבודה מתחיל (למשל, על ידי יצירת ה-Pod) ומתי מתבצעת קדימות ל-Pod ששייך לעומס עבודה.מידע נוסף על Kueue זמין במאמר מושגים ב-Kueue.
החלפה חמה: טכניקה שמפחיתה את הזמן הממוצע עד לשחזור (MTTR). החלפה חמה מאפשרת הפקעה על סמך עדיפות של עומס העבודה, כשמשאבי האשכול מנוצלים במלואם ואין קיבולת נוספת זמינה, ממכונות לפי דרישה או מהזמנות קיימות.
- כשצומת שמארח עומס עבודה הופך ללא תקין, עומס העבודה מתוזמן מחדש בצמתים חלופיים מתאימים. אם אין צמתים פנויים, Hotswap יכול להקדים עומס עבודה בעדיפות נמוכה יותר כדי לפנות מקום לעומס העבודה שמשוחזר.
- אם מגדירים את ה-Pods באמצעות
PriorityClass, עומס העבודה שמוגדר עם עדיפות גבוהה יותר מפנה עומס עבודה פעיל עם עדיפות נמוכה יותר כדי לקבל את המשאבים שלו. תהליך ההוצאה הזה נקרא קדימות.
תרחישים לדוגמה
בטבלה הבאה מפורטות השיטות המומלצות לכל תרחיש שימוש:
| תרחיש שימוש | שיטה מומלצת | תיאור |
|---|---|---|
| עומסי עבודה מרובים עם עדיפויות שונות | אתם יכולים להשתמש ב-Kueue כדי להגדיר תורים ולהקצות עדיפויות לעומסי עבודה על סמך החשיבות שלהם. מערכת Kueue יכולה לנהל את המכסות כדי שצוותים או פרויקטים מסוימים יוכלו לגשת לכמות מסוימת של משאבים. |
באמצעות Kueue אפשר להחיל את ההגדרות הבאות:
כדי לבדוק את ההגדרה של השיטה המומלצת, אפשר לעיין בדוגמה של Kueue במסמך הזה. |
| אתם צריכים לצמצם את ה-MTTR הנוכחי. | אפשר להשתמש ב-Hotswap כדי לתזמן מחדש עומסי עבודה במשאבים תקינים כשמתרחשת הפרעה, ולבטל עומסי עבודה בעדיפות נמוכה לטובת עומסי עבודה בעדיפות גבוהה. |
החלפה חמה מאפשרת להחיל את ההגדרות הבאות:
כדי לבדוק את ההגדרה של השיטה המומלצת, אפשר לעיין בדוגמה להחלפה חמה במאמר הזה. |
| עומסי עבודה מרובים של AI מתחרים על משאבים מוגבלים | שילוב של Kueue ו-Hotswap. השילוב הזה מספק מערכת חזקה שנותנת עדיפות לעומסי עבודה קריטיים גם במהלך התזמון הראשוני וגם במהלך זמן הריצה. |
בעזרת Kueue ו-Hotswap אפשר להחיל את ההגדרות הבאות:
כדי לבדוק את ההגדרה של השיטה המומלצת, אפשר לעיין בדוגמה של Kueue ו-Hotswap במסמך הזה. |
דוגמאות להטמעות של שיטות מומלצות
בדוגמאות הבאות אפשר לראות איך מטמיעים את Kueue ואת Hotswap, ואיך משלבים ביניהם כדי ליישם את השיטות המומלצות שמתוארות בקטע הקודם.
Kueue
קובץ המניפסט הבא מציג הגדרת Kueue:
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
name: tpu-v6e-slice
spec:
nodeLabels:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
name: tpu-training-cq
spec:
resourceGroups:
- flavors:
- name: tpu-v6e-slice
resources:
- name: google.com/tpu
nominalQuota: 32
queueingStrategy: BestEffortFIFO
preemption:
reclaimWithinCohort: Never
reclaimOutOfCohort:
enable: true
reclaimMoreThanNominalQuota: false
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
name: default-queue
namespace: default
spec:
clusterQueue: tpu-training-cq
קובץ המניפסט הזה:
- הגדרה של
ResourceFlavorבשםtpu-v6e-sliceשמציינת את תוויות הצומת לפרוסות TPU v6e. - מגדיר
ClusterQueueבשםtpu-training-cqשמנהל את המכסה של משאבי TPU. - הגדרת
LocalQueueבשםdefault-queueשמאפשרת לעומסי עבודה במרחב השמותdefaultלהשתמש בתור של אשכולtpu-training-cq.
החלפה חמה
בדוגמה הבאה מוצגת הגדרת Hotswap שמגדירה שתי מחלקות עדיפות, low-priority-job ו-high-priority-job. הגדרת ה-Hotswap הזו יוצרת עומס עבודה של JobSet בעדיפות גבוהה ומשתמשת ב-MaxText.
apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
name: low-priority-job
value: 1000000
globalDefault: false
description: "This priority class should be used for low priority pods only."
---
apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
name: high-priority-job
value: 2000000
globalDefault: false
description: "This priority class should be used for critical pods only."
---
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: high-jax-trillium
annotations:
alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
failurePolicy:
maxRestarts: 10
restartStrategy: BlockingRecreate
replicatedJobs:
- name: slice
replicas: 2
template:
spec:
backoffLimit: 0
completions: 4
parallelism: 4
template:
spec:
nodeSelector:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
cloud.google.com/gke-tpu-topology: 4x4
hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet
priorityClassName: high-priority-job
containers:
- name: jax-program
image: <IMAGE LOCATION>
command:
- python3
- MaxText/train.py
- MaxText/configs/base.yml
- model_name=llama2-7b
- run_name=<UNIQUE RUN NAME>
- steps=300
- base_output_directory=gs://<OUTPUT BUCKET>
- dataset_path=gs://max-datasets-rogue
- max_target_length=4096
- dataset_type=synthetic
- enable_checkpointing=False
resources:
limits:
google.com/tpu: 4
בהתאם להגדרה הזו, Hotswap מבצע את הפעולות הבאות:
- אם כשל בתשתית משבש את עומס העבודה בעדיפות גבוהה, המערכת מפעילה מחדש את JobSet. החלפה חמה (hotswap) מפקיעה את עומס העבודה בעדיפות נמוכה כדי לתזמן מחדש את עומס העבודה בעדיפות גבוהה לפני שהתשתית מתאוששת. עומס העבודה בעדיפות נמוכה נשאר בסטטוס 'נכשל'. התהליך הזה מקצר משמעותית את זמן חוסר הפעילות של עומס העבודה.
- כשהתשתית מתאוששת, Hotswap מתזמן מחדש את עומס העבודה בעדיפות נמוכה במאגר הצמתים שהתאושש.
Kueue ו-Hotswap
כדאי לשלב בין Kueue ל-Hotswap כשפועלים בסביבה מורכבת עם משאבים מוגבלים. השילוב הזה מספק מערכת חזקה שמתעדפת עומסי עבודה קריטיים במהלך התזמון הראשוני ובזמן הריצה.
בדוגמה הבאה מוצגת הגדרה משולבת של Kueue ו-Hotswap. בדוגמה הזו נעשה שימוש בפונקציה MaxText:
apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
name: low-priority-job
value: 1000000
globalDefault: false
description: "This priority class should be used for low priority pods only."
---
apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
name: high-priority-job
value: 2000000
globalDefault: false
description: "This priority class should be used for critical pods only."
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
name: tpu-v6e-slice
spec:
nodeLabels:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
name: tpu-training-cq
spec:
resourceGroups:
- flavors:
- name: tpu-v6e-slice
resources:
- name: google.com/tpu
nominalQuota: 32
queueingStrategy: BestEffortFIFO
preemption:
reclaimWithinCohort: Never
reclaimOutOfCohort:
enable: true
reclaimMoreThanNominalQuota: false
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
name: default-queue
namespace: default
spec:
clusterQueue: tpu-training-cq
---
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: low-jax-trillium
annotations:
kueue.x-k8s.io/queue-name: default-queue
alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
failurePolicy:
maxRestarts: 10
restartStrategy: BlockingRecreate
replicatedJobs:
- name: slice
replicas: 2
template:
spec:
backoffLimit: 0
completions: 4
parallelism: 4
template:
metadata:
labels:
kueue.x-k8s.io/managed-by: kueue
kueue.x-k8s.io/priority-class: low-priority-job
spec:
nodeSelector:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
cloud.google.com/gke-tpu-topology: 4x4
hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet
priorityClassName: low-priority-job
containers:
- name: jax-program
image: <IMAGE LOCATION>
command:
- python3
- MaxText/train.py
- MaxText/configs/base.yml
- model_name=llama2-7b
- run_name=low-priority-run
- steps=30000
- base_output_directory=gs://<OUTPUT BUCKET>
- dataset_path=gs://max-datasets-rogue
- max_target_length=4096
- dataset_type=synthetic
- enable_checkpointing=False
resources:
limits:
google.com/tpu: 4
---
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: high-jax-trillium
annotations:
kueue.x-k8s.io/queue-name: default-queue
alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
failurePolicy:
maxRestarts: 10
restartStrategy: BlockingRecreate
replicatedJobs:
- name: slice
replicas: 2
template:
spec:
backoffLimit: 0
completions: 4
parallelism: 4
template:
metadata:
labels:
kueue.x-k8s.io/managed-by: kueue
kueue.x-k8s.io/priority-class: high-priority-job
spec:
nodeSelector:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
cloud.google.com/gke-tpu-topology: 4x4
hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet
priorityClassName: high-priority-job
containers:
- name: jax-program
image: <IMAGE LOCATION>
command:
- python3
- MaxText/train.py
- MaxText/configs/base.yml
- model_name=llama2-7b
- run_name=high-priority-run
- steps=300
- base_output_directory=gs://<OUTPUT BUCKET>
- dataset_path=gs://max-datasets-rogue
- max_target_length=4096
- dataset_type=synthetic
- enable_checkpointing=False
resources:
limits:
google.com/tpu: 4
בהתאם להגדרות האלה, Kueue משולב עם Hotswap ומבצע את הפעולות הבאות:
- Kueue מנהל את ההוספה של
low-jax-trilliumושלhigh-jax-trilliumJobSets לתור של האשכול על סמך העדיפויות המוגדרות והמשאבים הזמינים. - אם העבודה של
high-jax-trilliumJobSet מופסקת בגלל כשל בתשתית, Hotswap מבצעת קדימה שלlow-jax-trilliumJobSet כדי לתזמן מחדש את העבודה של JobSet בעדיפות גבוהה. - החלפה חמה מבטיחה שה-JobSet בעדיפות גבוהה יופעל מחדש במהירות, וכך יצומצם זמן ההמתנה שלו.
- כשהתשתית מתאוששת, התכונה "החלפה חמה" מתזמנת מחדש את ה-JobSet בעדיפות נמוכה במאגר הצמתים שהתאושש.