Prácticas recomendadas para la priorización de cargas de trabajo de IA/AA de GKE

En este documento, se describen las herramientas y las prácticas recomendadas para maximizar el uso de los recursos y minimizar el tiempo de inactividad de las cargas de trabajo de IA/AA heterogéneas en Google Kubernetes Engine (GKE), en especial cuando no hay capacidad en las reservas o a través de recursos a pedido. Las cargas de trabajo heterogéneas hacen referencia a diferentes tipos de cargas de trabajo de IA/AA que se ejecutan de forma simultánea en el mismo clúster de GKE. Por ejemplo, puedes ejecutar un servicio de inferencia en línea sensible a la latencia junto con una serie de trabajos de entrenamiento por lotes interrumpibles.

En esta guía, se proporcionan recomendaciones para administradores y operadores de plataformas, y especialistas en datos y en IA. Para obtener una descripción general consolidada de todas las prácticas recomendadas de GKE, consulta Prácticas recomendadas para GKE.

Beneficios de la priorización de cargas de trabajo de IA/AA

Las cargas de trabajo heterogéneas tienen diferentes prioridades y comparten capacidad y recursos limitados. En las prácticas recomendadas de esta página, se describe cómo configurar GKE y las herramientas de código abierto para ayudarte a obtener los siguientes beneficios:

  • Minimizar el tiempo de inactividad de las cargas de trabajo de prioridad alta
  • Ejecutar rápidamente las cargas de trabajo de prioridad alta
  • Optimizar el consumo de recursos

Fondo

GKE admite las siguientes herramientas de código abierto para optimizar el uso de los recursos.

  • Kueue: Es un sistema de colas de cargas de trabajo nativo de Kubernetes diseñado para cargas de trabajo por lotes, de IA y de computación de alto rendimiento. Kueue se puede extender para administrar otros tipos de cargas de trabajo, como las definidas por las definiciones de recursos personalizados, como leaderworkerset. Kueue administra las cuotas y cómo las cargas de trabajo las consumen en un clúster de Kubernetes. Kueue toma decisiones sobre cuándo espera una carga de trabajo, cuándo comienza una carga de trabajo (por ejemplo, creando el Pod) y cuándo se interrumpe un Pod que pertenece a una carga de trabajo.

    Para obtener más información sobre Kueue, consulta la documentación de conceptos de Kueue.

  • Hotswap: Es una técnica que reduce el tiempo promedio de recuperación (MTTR). Hotswap habilita la interrupción según la prioridad de la carga de trabajo cuando los recursos del clúster están completamente utilizados y no hay capacidad adicional disponible, ya sea de instancias a pedido o de reservas existentes.

    • Cuando un nodo que aloja una carga de trabajo deja de estar en buen estado, la carga de trabajo se reprograma en los nodos de repuesto aptos. Si no hay nodos de repuesto disponibles, Hotswap puede interrumpir una carga de trabajo de menor prioridad para dejar espacio para la carga de trabajo que se recupera.
    • Si configuras tus Pods con PriorityClass, la carga de trabajo configurada con mayor prioridad expulsa una carga de trabajo de prioridad baja en ejecución para adquirir sus recursos. Este proceso de expulsión se conoce como interrupción.

Casos de uso

Usa la siguiente tabla para comprender las prácticas recomendadas para cada caso de uso:

Caso de uso Práctica recomendada Descripción
Varias cargas de trabajo con diferentes prioridades Usa Kueue para definir colas y asignar prioridades a las cargas de trabajo según su importancia. Kueue puede administrar la cuota para que ciertos equipos o proyectos tengan acceso a una cantidad determinada de recursos.

Kueue te permite aplicar los siguientes parámetros de configuración:

  • Prioriza los trabajos de prioridad alta asignándoles un WorkloadPriority de Kueue más alto.
  • Habilita la cola de uso compartido equitativo de Kueue para que todas las cargas de trabajo reciban recursos, incluso las de prioridad baja.

Para probar la configuración de prácticas recomendadas, consulta el ejemplo de Kueue en este documento.

Debes reducir el MTTR actual. Usa Hotswap para reprogramar las cargas de trabajo en recursos en buen estado cuando se produce una interrupción y para interrumpir las cargas de trabajo de prioridad baja en favor de las cargas de trabajo de prioridad alta.

Hotswap te permite aplicar los siguientes parámetros de configuración:

  • Configura PriorityClasses para definir los niveles de prioridad de tus cargas de trabajo.
  • Asigna PriorityClasses más altas a las cargas de trabajo fundamentales.
  • Reprograma automáticamente las cargas de trabajo en nodos en buen estado cuando se producen interrupciones.

Para probar la configuración de prácticas recomendadas, consulta el ejemplo de Hotswap en este documento.

Varias cargas de trabajo de IA que compiten por recursos limitados Combina Kueue y Hotswap. Esta combinación proporciona un sistema sólido que prioriza las cargas de trabajo fundamentales durante la programación inicial y durante el tiempo de ejecución.

Kueue y Hotswap te permiten aplicar los siguientes parámetros de configuración:

  • Usa Kueue para administrar la programación inicial y la admisión de cargas de trabajo según la prioridad.
  • Usa Hotswap para controlar las interrupciones de las cargas de trabajo y habilitar la recuperación rápida. Hotswap ayuda a reducir el tiempo de recuperación de una carga de trabajo de prioridad alta cuando se produce una interrupción.

Para probar la configuración de prácticas recomendadas, consulta el ejemplo de Kueue y Hotswap en este documento.

Ejemplos de implementaciones de prácticas recomendadas

En los siguientes ejemplos, se muestra cómo implementar Kueue y Hotswap, y cómo combinarlos para las prácticas recomendadas que se describen en la sección anterior.

Kueue

En el siguiente manifiesto de ejemplo, se muestra una configuración de 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

Este manifiesto hace lo siguiente:

  • Define un ResourceFlavor llamado tpu-v6e-slice que especifica las etiquetas de nodo para las porciones de TPU v6e.
  • Define un ClusterQueue llamado tpu-training-cq que administra la cuota para los recursos de TPU.
  • Define un LocalQueue llamado default-queue que permite que las cargas de trabajo en el espacio de nombres default usen la cola de clúster tpu-training-cq.

Hotswap

En el siguiente ejemplo, se muestra una configuración de Hotswap que define dos clases de prioridad, low-priority-job y high-priority-job. Esta configuración de Hotswap crea una carga de trabajo de JobSet de prioridad alta y usa 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

Según esta configuración, Hotswap realiza las siguientes acciones:

  • Si una falla de infraestructura interrumpe la carga de trabajo de prioridad alta, el JobSet la reinicia. Hotswap interrumpe la carga de trabajo de prioridad baja para reprogramar la carga de trabajo de prioridad alta antes de que se recupere la infraestructura. La carga de trabajo de prioridad baja permanece en estado de error. Este proceso reduce significativamente el tiempo de inactividad de la carga de trabajo.
  • Cuando se recupera la infraestructura, Hotswap reprograma la carga de trabajo de prioridad baja en el grupo de nodos que se recuperó.

Kueue y Hotswap

Combina Kueue y Hotswap cuando operas en un entorno complejo con recursos limitados. Esta combinación proporciona un sistema sólido que prioriza las cargas de trabajo fundamentales durante la programación inicial y durante el tiempo de ejecución.

En el siguiente ejemplo, se muestra una configuración combinada de Kueue y Hotswap. En este ejemplo, se usa 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

Según esta configuración, Kueue se combina con Hotswap y realiza las siguientes acciones:

  • Kueue administra la admisión de JobSets low-jax-trillium y high-jax-trillium en la cola del clúster según sus prioridades definidas y los recursos disponibles.
  • Si el JobSet high-jax-trillium se interrumpe debido a una falla de infraestructura, Hotswap interrumpe el JobSet low-jax-trillium para reprogramar el JobSet de prioridad alta.
  • Hotswap garantiza que el JobSet de prioridad alta se reinicie rápidamente, lo que minimiza su tiempo de inactividad.
  • Cuando se recupera la infraestructura, Hotswap reprograma el JobSet de prioridad baja en el grupo de nodos recuperado.

¿Qué sigue?