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

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

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

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

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

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

רקע

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

  • ‫Kueue: מערכת תורים לעומסי עבודה (workload) ב-Kubernetes, שנועדה לעומסי עבודה באצווה, ל-AI ולמחשוב עתיר ביצועים. אפשר להרחיב את 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 מתזמן מחדש את עומס העבודה בעדיפות נמוכה במאגר הצמתים שהשתקם.

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

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