אופטימיזציה של תעדוף עומסי עבודה של AI/ML ב-GKE

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

במסמך הזה מתוארים כלים ושיטות מומלצות למיקסום השימוש במשאבים ולצמצום זמן ההשבתה של עומסי עבודה הטרוגניים של AI/ML ב-Google Kubernetes Engine ‏ (GKE), במיוחד כשאין קיבולת בהזמנות או דרך משאבים על פי דרישה. עומסי עבודה הטרוגניים הם סוגים שונים של עומסי עבודה של AI/ML שפועלים בו-זמנית באותו אשכול GKE. לדוגמה, יכול להיות שתפעילו שירות הסקה אונליין שרגיש לזמן האחזור, לצד סדרה של משימות אימון באצווה שניתן להפריע להן.

המדריך הזה כולל המלצות לאדמינים ולמפעילים של הפלטפורמה, ולמומחים בתחום הנתונים וה-AI.

היתרונות של תעדוף עומסי עבודה של AI/ML

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

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

רקע

‫GKE תומך בכלים הבאים בקוד פתוח לאופטימיזציה של ניצול המשאבים.

‫Kueue: מערכת לתור עומסי עבודה (workload) ב-Kubernetes, שנועדה לעומסי עבודה באצווה, ל-AI ולמחשוב עתיר ביצועים. אפשר להרחיב את Kueue כדי לנהל סוגים אחרים של עומסי עבודה, כמו אלה שמוגדרים על ידי הגדרות של משאבים בהתאמה אישית כמו leaderworkerset. ‫Kueue מנהל את המכסות ואת האופן שבו עומסי העבודה צורכים אותן באשכול Kubernetes. ‫Kueue מקבל החלטות לגבי מתי עומס עבודה ממתין, מתי עומס עבודה מתחיל (למשל, על ידי יצירת ה-Pod) ומתי מתבצעת קדימות ל-Pod ששייך לעומס עבודה.

מידע נוסף על Kueue זמין במאמר מושגים ב-Kueue.
החלפה בזמן פעולה: טכניקה שמקצרת את הזמן הממוצע עד להחזרה לפעולה (MTTR). החלפה חמה מאפשרת קדימות על סמך העדיפות של עומס העבודה, כשמשאבי האשכול נמצאים בשימוש מלא ואין קיבולת נוספת זמינה, ממופעים לפי דרישה או מהזמנות קיימות.
- כשצומת שמארח עומס עבודה הופך ללא תקין, עומס העבודה מתוזמן מחדש בצמתים חלופיים מתאימים. אם אין צמתים חלופיים זמינים, אפשר להשתמש ב-Hotswap כדי להפסיק לפני הזמן עומס עבודה בעדיפות נמוכה יותר, כדי לפנות מקום לעומס העבודה שמשוחזר.
- אם מגדירים את ה-Pods באמצעות PriorityClass, עומס העבודה שהוגדר עם עדיפות גבוהה יותר מפנה עומס עבודה בעדיפות נמוכה שפועל כדי לקבל את המשאבים שלו. תהליך הפינוי הזה נקרא הפסקה זמנית.

תרחישים לדוגמה

בטבלה הבאה מפורטות השיטות המומלצות לכל תרחיש שימוש:

תרחיש לדוגמה שיטה מומלצת תיאור

עומסי עבודה מרובים עם עדיפויות שונות

אתם יכולים להשתמש ב-Kueue כדי להגדיר תורים ולהקצות עדיפויות לעומסי עבודה על סמך החשיבות שלהם. מערכת Kueue יכולה לנהל את המכסות כדי שלצוותים או לפרויקטים מסוימים תהיה גישה לכמות מסוימת של משאבים.

תרחיש לדוגמה	שיטה מומלצת	תיאור
עומסי עבודה מרובים עם עדיפויות שונות	אתם יכולים להשתמש ב-Kueue כדי להגדיר תורים ולהקצות עדיפויות לעומסי עבודה על סמך החשיבות שלהם. מערכת Kueue יכולה לנהל את המכסות כדי שלצוותים או לפרויקטים מסוימים תהיה גישה לכמות מסוימת של משאבים.	באמצעות Kueue אפשר להחיל את ההגדרות הבאות: כדי לתעדף משימות בעדיפות גבוהה, מקצים להן ערך גבוה יותר של `WorkloadPriority` ב-Kueue. מפעילים את התור של Kueue עם חלוקה הוגנת של משאבים, כדי שכל עומסי העבודה יקבלו בסופו של דבר משאבים, גם אלה עם עדיפות נמוכה. כדי לבדוק את ההגדרה של השיטה המומלצת, אפשר לעיין בדוגמה ל-Kueue במסמך הזה.
צריך להקטין את ה-MTTR הנוכחי.	אפשר להשתמש ב-Hotswap כדי לתזמן מחדש עומסי עבודה במשאבים תקינים כשמתרחשת הפרעה, ולהפסיק לפני הזמן (preempt) עומסי עבודה בעדיפות נמוכה לטובת עומסי עבודה בעדיפות גבוהה.	החלפה חמה מאפשרת להחיל את ההגדרות הבאות: מגדירים PriorityClasses כדי להגדיר רמות עדיפות לעומסי העבודה. הקצאת `PriorityClasses` גבוה יותר לעומסי עבודה קריטיים. תזמון מחדש אוטומטי של עומסי עבודה בצמתים תקינים כשמתרחשות הפרעות. כדי לבדוק את ההגדרה של השיטה המומלצת, אפשר לעיין בדוגמה להחלפה חמה במסמך הזה.
כמה עומסי עבודה של AI מתחרים על משאבים מוגבלים	משלבים את Kueue עם Hotswap. השילוב הזה מספק מערכת חזקה שמתעדפת עומסי עבודה קריטיים גם במהלך התזמון הראשוני וגם במהלך זמן הריצה.	בעזרת Kueue ו-Hotswap אפשר להחיל את ההגדרות הבאות: משתמשים ב-Kueue כדי לנהל את התזמון הראשוני ואת ההרשאה של עומסי עבודה על סמך עדיפות. משתמשים ב-Hotswap כדי לטפל בהפרעות בעומס העבודה ולאפשר שחזור מהיר. החלפה חמה עוזרת לצמצם את זמן ההתאוששות של עומס עבודה בעדיפות גבוהה כשמתרחשת הפרעה. כדי לבדוק את ההגדרה של השיטה המומלצת, אפשר לעיין בדוגמה של Kueue ו-Hotswap במסמך הזה.

באמצעות Kueue אפשר להחיל את ההגדרות הבאות:

כדי לתעדף משימות בעדיפות גבוהה, מקצים להן ערך גבוה יותר של WorkloadPriority ב-Kueue.
מפעילים את התור של Kueue עם חלוקה הוגנת של משאבים, כדי שכל עומסי העבודה יקבלו בסופו של דבר משאבים, גם אלה עם עדיפות נמוכה.

כדי לבדוק את ההגדרה של השיטה המומלצת, אפשר לעיין בדוגמה ל-Kueue במסמך הזה.

צריך להקטין את ה-MTTR הנוכחי.

אפשר להשתמש ב-Hotswap כדי לתזמן מחדש עומסי עבודה במשאבים תקינים כשמתרחשת הפרעה, ולהפסיק לפני הזמן (preempt) עומסי עבודה בעדיפות נמוכה לטובת עומסי עבודה בעדיפות גבוהה.

החלפה חמה מאפשרת להחיל את ההגדרות הבאות:

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

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

כמה עומסי עבודה של AI מתחרים על משאבים מוגבלים

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

בעזרת 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 מתזמן מחדש את עומס העבודה בעדיפות נמוכה במאגר הצמתים ששוקם.

‫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-trillium JobSets לתור של האשכול על סמך העדיפויות המוגדרות והמשאבים הזמינים.
אם high-jax-trillium JobSet נקטע בגלל כשל בתשתית, Hotswap מבצע preempt ל-low-jax-trillium JobSet כדי לתזמן מחדש את high-jax-trillium JobSet בעדיפות גבוהה.
החלפה מהירה מבטיחה שה-JobSet בעדיפות גבוהה יופעל מחדש במהירות, וכך יצומצם זמן ההמתנה שלו.
כשהתשתית משתקמת, Hotswap מתזמן מחדש את JobSet בעדיפות נמוכה במאגר הצמתים המשוקם.

אופטימיזציה של תעדוף עומסי עבודה של AI/ML ב-GKE קל לארגן דפים בעזרת אוספים אפשר לשמור ולסווג תוכן על סמך ההעדפות שלך.