Optimiza 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 la utilización de recursos y minimizar el tiempo de inactividad de las cargas de trabajo heterogéneas de AA/ML 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 los administradores y operadores de plataformas, y los especialistas en IA y datos.

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 esta página, se describen las prácticas recomendadas para configurar GKE y las herramientas de código abierto, lo que te permitirá obtener los siguientes beneficios:

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

Fondo

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

  • Kueue: Es un sistema de filas 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 que se definen con 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 se inicia una carga de trabajo (por ejemplo, creando el Pod) y cuándo se expulsa un Pod perteneciente a una carga de trabajo.

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

  • Intercambio en caliente: Es una técnica que reduce el tiempo promedio de recuperación (MTTR). Hotswap permite la expulsión preventiva según la prioridad de la carga de trabajo cuando los recursos del clúster se utilizan por completo y no hay capacidad adicional disponible, ya sea de instancias on demand 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 nodos de reserva aptos. Si no hay nodos de repuesto disponibles, Hotswap puede interrumpir una carga de trabajo de menor prioridad para dejar espacio a la carga de trabajo que se está recuperando.
    • Si configuras tus Pods con PriorityClass, la carga de trabajo configurada con mayor prioridad expulsa una carga de trabajo en ejecución de baja prioridad para adquirir sus recursos. Este proceso de desalojo se conoce como expulsió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 valor de WorkloadPriority más alto en Kueue.
  • Habilita la cola de uso compartido equitativo de Kueue para que todas las cargas de trabajo reciban recursos, incluso las de baja prioridad.

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 cargas de trabajo en recursos en buen estado cuando se produce una interrupción y para interrumpir cargas de trabajo de baja prioridad en favor de cargas de trabajo de alta prioridad.

El intercambio en caliente te permite aplicar los siguientes parámetros de configuración:

  • Configura PriorityClasses para definir niveles de prioridad para tus cargas de trabajo.
  • Asigna un valor de PriorityClasses más alto a las cargas de trabajo críticas.
  • Reprogramar automáticamente las cargas de trabajo en nodos en buen estado cuando se producen interrupciones

Para probar la configuración de práctica recomendada, consulta el ejemplo de intercambio en caliente en este documento.

Varias cargas de trabajo de IA compiten por recursos limitados Combina Kueue y Hotswap. Esta combinación proporciona un sistema sólido que prioriza las cargas de trabajo críticas tanto durante la programación inicial como 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 y la admisión iniciales de las cargas de trabajo según la prioridad.
  • Usa Hotswap para controlar las interrupciones de la carga de trabajo y habilitar la recuperación rápida. El intercambio en caliente ayuda a reducir el tiempo de recuperación de una carga de trabajo de alta prioridad cuando se produce una interrupción.

Para probar la configuración de práctica recomendada, 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 aplicar 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 del clúster tpu-training-cq.

Intercambio en caliente

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 intercambio en caliente crea una carga de trabajo de JobSet de alta prioridad 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 alta prioridad, el JobSet la reinicia. El intercambio en caliente interrumpe la carga de trabajo de baja prioridad para reprogramar la carga de trabajo de alta prioridad 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 baja prioridad en el grupo de nodos que se recuperó.

Kueue y Hotswap

Combina Kueue y Hotswap cuando operes en un entorno complejo con recursos limitados. Esta combinación proporciona un sistema sólido que prioriza las cargas de trabajo críticas 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 los 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 anula el JobSet low-jax-trillium para reprogramar el JobSet de alta prioridad.
  • El intercambio en caliente garantiza que el JobSet de alta prioridad se reinicie rápidamente, lo que minimiza su tiempo de inactividad.
  • Cuando se recupera la infraestructura, Hotswap vuelve a programar el JobSet de baja prioridad en el grupo de nodos recuperado.

¿Qué sigue?