Optimizar la priorización de cargas de trabajo de IA/ML 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 de IA/ML heterogéneas en Google Kubernetes Engine (GKE), especialmente cuando no hay capacidad en las reservas o a través de recursos bajo demanda. Las cargas de trabajo heterogéneas hacen referencia a diferentes tipos de cargas de trabajo de IA y aprendizaje automático que se ejecutan simultáneamente en el mismo clúster de GKE. Por ejemplo, puedes ejecutar un servicio de inferencia online sensible a la latencia junto con una serie de trabajos de entrenamiento por lotes interrumpibles.

En esta guía se ofrecen recomendaciones para administradores y operadores de la plataforma, así como para especialistas en datos e IA.

Ventajas de priorizar las cargas de trabajo de IA y aprendizaje automático

Las cargas de trabajo heterogéneas tienen prioridades diferentes y comparten una capacidad y unos recursos limitados. En esta página se describen las prácticas recomendadas para configurar GKE y las herramientas de código abierto, que te ayudarán a obtener las siguientes ventajas:

  • Minimiza el tiempo de inactividad de las cargas de trabajo de alta prioridad.
  • Ejecuta rápidamente las 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 los recursos.

  • Kueue: 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 ampliar para gestionar otros tipos de cargas de trabajo, como las definidas por definiciones de recursos personalizadas, como leaderworkerset. Kueue gestiona las cuotas y cómo las consumen las cargas de trabajo en un clúster de Kubernetes. Kueue toma decisiones sobre cuándo espera una carga de trabajo, cuándo se inicia (por ejemplo, creando el pod) y cuándo se expropia un pod perteneciente a una carga de trabajo.

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

  • Intercambio en caliente: técnica que reduce el tiempo medio de recuperación (MTTR). El intercambio en caliente permite la expropiación en función de la prioridad de la carga de trabajo cuando los recursos del clúster están totalmente utilizados y no hay capacidad adicional disponible, ya sea de instancias bajo demanda o de reservas.

    • Cuando un nodo que aloja una carga de trabajo deja de estar en buen estado, la carga de trabajo se vuelve a programar 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 a la carga de trabajo que se está recuperando.
    • Si configuras tus pods con PriorityClass, la carga de trabajo configurada con mayor prioridad desalojará una carga de trabajo de baja prioridad en ejecución para adquirir sus recursos. Este proceso de desalojo se conoce como "prelación".

Casos prácticos

Consulta la siguiente tabla para conocer las prácticas recomendadas de cada caso práctico:

Caso práctico 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 en función de su importancia. Kueue puede gestionar las cuotas para que determinados equipos o proyectos tengan acceso a una cantidad determinada de recursos.

Kueue te permite aplicar las siguientes configuraciones:

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

Para probar la configuración de las prácticas recomendadas, consulta el ejemplo de Kueue de este documento.
Debes reducir el MTTR actual. Usa Hotswap para reprogramar cargas de trabajo en recursos en buen estado cuando se produzca una interrupción y para interrumpir cargas de trabajo de baja prioridad en favor de cargas de trabajo de alta prioridad.

La función de intercambio en caliente te permite aplicar las siguientes configuraciones:

  • Configura PriorityClasses para definir los niveles de prioridad de 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 correctos cuando se produzcan interrupciones.

Para probar la configuración de práctica recomendada, consulta el ejemplo de intercambio en caliente de 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 las siguientes configuraciones:

  • Usa Kueue para gestionar la programación inicial y la admisión de cargas de trabajo en función de la prioridad.
  • Usa Hotswap para gestionar las interrupciones de las cargas de trabajo y permitir una 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 las prácticas recomendadas, consulta el ejemplo de Kueue y Hotswap de 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 seguir las prácticas recomendadas descritas en la sección anterior.

Kueue

En el siguiente ejemplo de manifiesto 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 archivo de manifiesto hace lo siguiente:

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

Intercambio en caliente

En el ejemplo siguiente se muestra una configuración de intercambio en caliente 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 un fallo de la infraestructura interrumpe la carga de trabajo de alta prioridad, 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 baja prioridad sigue teniendo el estado de error. Este proceso reduce significativamente el tiempo de inactividad de la carga de trabajo.
  • Cuando se recupera la infraestructura, Hotswap vuelve a programar la carga de trabajo de baja prioridad en el grupo de nodos que se ha recuperado.

Kueue y Hotswap

Combina Kueue y Hotswap cuando trabajes en un entorno complejo con recursos limitados. Esta combinación proporciona un sistema robusto 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 gestiona la admisión de low-jax-trillium y high-jax-trillium JobSets en la cola del clúster en función de sus prioridades definidas y los recursos disponibles.
  • Si la infraestructura interrumpe el high-jax-trillium JobSet, Hotswap lo sustituye para reprogramar el low-jax-trillium JobSet de alta prioridad.
  • El intercambio en caliente asegura 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.

Siguientes pasos