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

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

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

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

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

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

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

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