Optimizar el uso de recursos de GKE para cargas de trabajo mixtas de entrenamiento e inferencia de IA y aprendizaje automático

En este tutorial se explica cómo compartir de forma eficiente recursos de acelerador entre cargas de trabajo de entrenamiento y de servicio de inferencias en un mismo clúster de Google Kubernetes Engine (GKE). Al distribuir tus cargas de trabajo mixtas en un solo clúster, mejoras el uso de los recursos, simplificas la gestión de clústeres, reduces los problemas derivados de las limitaciones de cantidad de aceleradores y aumentas la rentabilidad general.

En este tutorial, crearás una implementación de servicio de alta prioridad con el modelo de lenguaje extenso (LLM) Gemma 2 para la inferencia y el framework de servicio Hugging Face TGI (Text Generation Interface), junto con un trabajo de ajuste fino de LLM de baja prioridad. Ambas cargas de trabajo se ejecutan en un solo clúster que usa GPUs NVIDIA L4. Utilizas Kueue, un sistema de colas de trabajos nativo de Kubernetes y de código abierto, para gestionar y programar tus cargas de trabajo. Kueue te permite priorizar las tareas de servicio y anticipar los trabajos de entrenamiento de menor prioridad para optimizar el uso de los recursos. A medida que disminuye la demanda de servicio, reasignas los aceleradores liberados para reanudar los trabajos de entrenamiento. Utilizas Kueue y las clases de prioridad para gestionar las cuotas de recursos durante todo el proceso.

Este tutorial está dirigido a ingenieros de aprendizaje automático (ML), administradores y operadores de plataformas, y especialistas en datos e IA que quieran entrenar y alojar un modelo de aprendizaje automático en un clúster de GKE, así como reducir los costes y la sobrecarga de gestión, sobre todo cuando se trata de un número limitado de aceleradores. Para obtener más información sobre los roles habituales y las tareas de ejemplo a las que hacemos referencia en el contenido de Google Cloud , consulta Roles y tareas de usuario habituales de GKE.

Antes de leer esta página, asegúrese de que conoce los siguientes conceptos:

• Disponibilidad de la versión actual de la GPU
• GPUs en GKE

Preparar el entorno

En esta sección, aprovisionarás los recursos que necesitas para desplegar TGI y el modelo de tus cargas de trabajo de inferencia y entrenamiento.

Acceder al modelo

Para acceder a los modelos de Gemma e implementarlos en GKE, primero debes firmar el acuerdo de consentimiento de licencia y, a continuación, generar un token de acceso de Hugging Face.

1. Firma el contrato de consentimiento de licencia. Accede a la página de consentimiento del modelo, verifica el consentimiento con tu cuenta de Hugging Face y acepta los términos del modelo.

2. Genera un token de acceso. Para acceder al modelo a través de Hugging Face, necesitas un token de Hugging Face. Sigue estos pasos para generar un token si aún no tienes uno:

   1. Haz clic en Tu perfil > Configuración > Tokens de acceso.
   2. Selecciona New Token (Nuevo token).
   3. Especifica el nombre que quieras y un rol de al menos Read.
   4. Selecciona Generar un token.
   5. Copia el token generado en el portapapeles.

Abrir Cloud Shell

En este tutorial, usarás Cloud Shell para gestionar los recursos alojados enGoogle Cloud. Cloud Shell tiene preinstalado el software que necesitas para este tutorial, como kubectl, gcloud CLI y Terraform.

Para configurar tu entorno con Cloud Shell, sigue estos pasos:

1. En la Google Cloud consola, inicia una sesión de Cloud Shell haciendo clic en Activar Cloud Shell en la Google Cloud consola. Se iniciará una sesión en el panel inferior de la consola Google Cloud .

2. Define las variables de entorno predeterminadas:

   gcloud config set project PROJECT_ID
   export PROJECT_ID=$(gcloud config get project)
   

   Sustituye PROJECT_ID por el Google Cloud ID de tu proyecto.

3. Clona el código de ejemplo de GitHub. En Cloud Shell, ejecuta los siguientes comandos:

   git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/
   cd kubernetes-engine-samples/ai-ml/mix-train-and-inference
   export EXAMPLE_HOME=$(pwd)
   

Crear un clúster de GKE

Puedes usar un clúster de Autopilot o Estándar para tus cargas de trabajo mixtas. Te recomendamos que uses un clúster de Autopilot para disfrutar de una experiencia de Kubernetes totalmente gestionada. Para elegir el modo de funcionamiento de GKE que mejor se adapte a tus cargas de trabajo, consulta Elegir un modo de funcionamiento de GKE.

Autopilot

1. Define las variables de entorno predeterminadas en Cloud Shell:

   export HF_TOKEN=HF_TOKEN
   export REGION=REGION
   export CLUSTER_NAME="llm-cluster"
   export PROJECT_NUMBER=$(gcloud projects list \
       --filter="$(gcloud config get-value project)" \
       --format="value(PROJECT_NUMBER)")
   export MODEL_BUCKET="model-bucket-$PROJECT_ID"
   

   Sustituye los siguientes valores:

   • HF_TOKEN: el token de Hugging Face que has generado anteriormente.
   • REGION: una región que admita el tipo de acelerador que quieras usar. Por ejemplo, us-central1 para la GPU L4.

   Puedes ajustar la variable MODEL_BUCKET, que representa el segmento de Cloud Storage en el que almacenas los pesos del modelo entrenado.

2. Crea un clúster de Autopilot:

   gcloud container clusters create-auto ${CLUSTER_NAME} \
       --project=${PROJECT_ID} \
       --location=${REGION} \
       --release-channel=rapid
   

3. Crea el segmento de Cloud Storage para el trabajo de ajuste fino:

   gcloud storage buckets create gs://${MODEL_BUCKET} \
       --location ${REGION} \
       --uniform-bucket-level-access
   

4. Para conceder acceso al segmento de Cloud Storage, ejecuta este comando:

   gcloud storage buckets add-iam-policy-binding "gs://$MODEL_BUCKET" \
       --role=roles/storage.objectAdmin \
       --member=principal://iam.googleapis.com/projects/$PROJECT_NUMBER/locations/global/workloadIdentityPools/$PROJECT_ID.svc.id.goog/subject/ns/llm/sa/default \
       --condition=None
   

5. Para obtener las credenciales de autenticación del clúster, ejecuta este comando:

   gcloud container clusters get-credentials llm-cluster \
       --location=$REGION \
       --project=$PROJECT_ID
   

6. Crea un espacio de nombres para tus implementaciones. En Cloud Shell, ejecuta el siguiente comando:

   kubectl create ns llm
   

Estándar

1. Define las variables de entorno predeterminadas en Cloud Shell:

   export HF_TOKEN=HF_TOKEN
   export REGION=REGION
   export CLUSTER_NAME="llm-cluster"
   export GPU_POOL_MACHINE_TYPE="g2-standard-24"
   export GPU_POOL_ACCELERATOR_TYPE="nvidia-l4"
   export PROJECT_NUMBER=$(gcloud projects list \
       --filter="$(gcloud config get-value project)" \
       --format="value(PROJECT_NUMBER)")
   export MODEL_BUCKET="model-bucket-$PROJECT_ID"
   

   Sustituye los siguientes valores:

   • HF_TOKEN: el token de Hugging Face que has generado anteriormente.
   • REGION: la región que admite el tipo de acelerador que quieres usar. Por ejemplo, us-central1 para la GPU L4.

   Puedes ajustar estas variables:

   • GPU_POOL_MACHINE_TYPE: la serie de máquinas del grupo de nodos que quieres usar en la región seleccionada. Este valor depende del tipo de acelerador que hayas seleccionado. Para obtener más información, consulta Limitaciones del uso de GPUs en GKE. Por ejemplo, en este tutorial se usa g2-standard-24 con dos GPUs conectadas por nodo. Para ver la lista más actualizada de GPUs disponibles, consulta GPUs para cargas de trabajo de Compute.
   • GPU_POOL_ACCELERATOR_TYPE: el tipo de acelerador que se admite en la región seleccionada. Por ejemplo, en este tutorial se usa nvidia-l4. Para ver la lista más reciente de GPUs disponibles, consulta GPUs para cargas de trabajo de Compute.
   • MODEL_BUCKET: el segmento de Cloud Storage donde se almacenan los pesos del modelo entrenado.

2. Crea un clúster estándar:

   gcloud container clusters create ${CLUSTER_NAME} \
       --project=${PROJECT_ID} \
       --location=${REGION} \
       --workload-pool=${PROJECT_ID}.svc.id.goog \
       --release-channel=rapid \
       --machine-type=e2-standard-4 \
       --addons GcsFuseCsiDriver \
       --num-nodes=1
   

3. Crea el grupo de nodos de GPU para las cargas de trabajo de inferencia y ajuste fino:

   gcloud container node-pools create gpupool \
       --accelerator type=${GPU_POOL_ACCELERATOR_TYPE},count=2,gpu-driver-version=latest \
       --project=${PROJECT_ID} \
       --location=${REGION} \
       --node-locations=${REGION}-a \
       --cluster=${CLUSTER_NAME} \
       --machine-type=${GPU_POOL_MACHINE_TYPE} \
       --num-nodes=3
   

4. Crea el segmento de Cloud Storage para el trabajo de ajuste fino:

   gcloud storage buckets create gs://${MODEL_BUCKET} \
       --location ${REGION} \
       --uniform-bucket-level-access
   

5. Para conceder acceso al segmento de Cloud Storage, ejecuta este comando:

   gcloud storage buckets add-iam-policy-binding "gs://$MODEL_BUCKET" \
       --role=roles/storage.objectAdmin \
       --member=principal://iam.googleapis.com/projects/$PROJECT_NUMBER/locations/global/workloadIdentityPools/$PROJECT_ID.svc.id.goog/subject/ns/llm/sa/default \
       --condition=None
   

6. Para obtener las credenciales de autenticación del clúster, ejecuta este comando:

   gcloud container clusters get-credentials llm-cluster \
       --location=$REGION \
       --project=$PROJECT_ID
   

7. Crea un espacio de nombres para tus implementaciones. En Cloud Shell, ejecuta el siguiente comando:

   kubectl create ns llm
   

Crear un secreto de Kubernetes para las credenciales de Hugging Face

Para crear un secreto de Kubernetes que contenga el token de Hugging Face, ejecuta el siguiente comando:

kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=$HF_TOKEN \
    --dry-run=client -o yaml | kubectl apply --namespace=llm --filename=-

Configurar Kueue

En este tutorial, Kueue es el gestor de recursos central, que permite compartir GPUs de forma eficiente entre tus cargas de trabajo de entrenamiento y de servicio. Kueue lo consigue definiendo los requisitos de recursos ("sabores"), priorizando las cargas de trabajo mediante colas (con tareas de servicio priorizadas sobre el entrenamiento) y asignando recursos de forma dinámica en función de la demanda y la prioridad. En este tutorial se usa el tipo de recurso Workload para agrupar las cargas de trabajo de inferencia y ajuste, respectivamente.

La función de preferencia de Kueue asegura que las cargas de trabajo de servicio de alta prioridad siempre tengan los recursos necesarios pausando o desalojando los trabajos de entrenamiento de menor prioridad cuando los recursos sean escasos.

Para controlar la implementación del servidor de inferencia con Kueue, habilita la integración de pod y configura managedJobsNamespaceSelector para excluir los espacios de nombres kube-system y kueue-system.

1. En el directorio /kueue, consulta el código en kustomization.yaml. Este manifiesto instala el gestor de recursos Kueue con configuraciones personalizadas.

   apiVersion: kustomize.config.k8s.io/v1beta1
   kind: Kustomization
   resources:
   - https://github.com/kubernetes-sigs/kueue/releases/download/v0.12.3/manifests.yaml
   patches:
   - path: patch.yaml
     target:
       version: v1
       kind: ConfigMap
       name: kueue-manager-config
   

2. En el directorio /kueue, consulta el código en patch.yaml. Este ConfigMap personaliza Kueue para excluir la gestión de pods en los espacios de nombres kube-system y kueue-system.

   apiVersion: v1
   kind: ConfigMap
   metadata:
     name: kueue-manager-config
   data:
     controller_manager_config.yaml: |
       apiVersion: config.kueue.x-k8s.io/v1beta1
       kind: Configuration
       health:
         healthProbeBindAddress: :8081
       metrics:
         bindAddress: :8080
       # enableClusterQueueResources: true
       webhook:
         port: 9443
       leaderElection:
         leaderElect: true
         resourceName: c1f6bfd2.kueue.x-k8s.io
       controller:
         groupKindConcurrency:
           Job.batch: 5
           Pod: 5
           Workload.kueue.x-k8s.io: 5
           LocalQueue.kueue.x-k8s.io: 1
           ClusterQueue.kueue.x-k8s.io: 1
           ResourceFlavor.kueue.x-k8s.io: 1
       clientConnection:
         qps: 50
         burst: 100
       #pprofBindAddress: :8083
       #waitForPodsReady:
       #  enable: false
       #  timeout: 5m
       #  blockAdmission: false
       #  requeuingStrategy:
       #    timestamp: Eviction
       #    backoffLimitCount: null # null indicates infinite requeuing
       #    backoffBaseSeconds: 60
       #    backoffMaxSeconds: 3600
       #manageJobsWithoutQueueName: true
       managedJobsNamespaceSelector:
         matchExpressions:
           - key: kubernetes.io/metadata.name
             operator: NotIn
             values: [ kube-system, kueue-system ]
       #internalCertManagement:
       #  enable: false
       #  webhookServiceName: ""
       #  webhookSecretName: ""
       integrations:
         frameworks:
         - "batch/job"
         - "kubeflow.org/mpijob"
         - "ray.io/rayjob"
         - "ray.io/raycluster"
         - "jobset.x-k8s.io/jobset"
         - "kubeflow.org/paddlejob"
         - "kubeflow.org/pytorchjob"
         - "kubeflow.org/tfjob"
         - "kubeflow.org/xgboostjob"
         - "kubeflow.org/jaxjob"
         - "workload.codeflare.dev/appwrapper"
         - "pod"
       #  - "deployment" # requires enabling pod integration
       #  - "statefulset" # requires enabling pod integration
       #  - "leaderworkerset.x-k8s.io/leaderworkerset" # requires enabling pod integration
       #  externalFrameworks:
       #  - "Foo.v1.example.com"
       #fairSharing:
       #  enable: true
       #  preemptionStrategies: [LessThanOrEqualToFinalShare, LessThanInitialShare]
       #admissionFairSharing:
       #  usageHalfLifeTime: "168h" # 7 days
       #  usageSamplingInterval: "5m"
       #  resourceWeights: # optional, defaults to 1 for all resources if not specified
       #    cpu: 0    # if you want to completely ignore cpu usage
       #    memory: 0 # ignore completely memory usage
       #    example.com/gpu: 100 # and you care only about GPUs usage
       #resources:
       #  excludeResourcePrefixes: []
       #  transformations:
       #  - input: nvidia.com/mig-4g.5gb
       #    strategy: Replace | Retain
       #    outputs:
       #      example.com/accelerator-memory: 5Gi
       #      example.com/accelerator-gpc: 4
       #objectRetentionPolicies:
       #  workloads:
       #    afterFinished: null # null indicates infinite retention, 0s means no retention at all
       #    afterDeactivatedByKueue: null # null indicates infinite retention, 0s means no retention at all
   

3. En Cloud Shell, ejecuta el siguiente comando para instalar Kueue:

   cd ${EXAMPLE_HOME}
   kubectl kustomize kueue |kubectl apply --server-side --filename=-
   

   Espera hasta que los pods de Kueue estén listos:

   watch kubectl --namespace=kueue-system get pods
   

   La salida debería ser similar a la siguiente:

   NAME                                        READY   STATUS    RESTARTS   AGE
   kueue-controller-manager-bdc956fc4-vhcmx    1/1     Running   0          3m15s
   

4. En el directorio /workloads, consulta los archivos flavors.yaml, cluster-queue.yaml y local-queue.yaml. Estos manifiestos especifican cómo gestiona Kueue las cuotas de recursos:

   ResourceFlavor

   Este manifiesto define un ResourceFlavor predeterminado en Kueue para la gestión de recursos.

   apiVersion: kueue.x-k8s.io/v1beta1
   kind: ResourceFlavor
   metadata:
     name: default-flavor
   

   ClusterQueue

   Este manifiesto configura una ClusterQueue de Kueue con límites de recursos para CPU, memoria y GPU.

   En este tutorial se usan nodos con dos GPUs Nvidia L4 conectadas, con el tipo de nodo g2-standard-24 correspondiente, que ofrece 24 vCPUs y 96 GB de RAM. El código de ejemplo muestra cómo limitar el uso de recursos de tu carga de trabajo a un máximo de seis GPUs.

   El campo preemption de la configuración de ClusterQueue hace referencia a las PriorityClasses para determinar qué pods se pueden desalojar cuando los recursos son escasos.

   apiVersion: kueue.x-k8s.io/v1beta1
   kind: ClusterQueue
   metadata:
     name: "cluster-queue"
   spec:
     namespaceSelector: {} # match all.
     preemption:
       reclaimWithinCohort: LowerPriority
       withinClusterQueue: LowerPriority
     resourceGroups:
     - coveredResources: [ "cpu", "memory", "nvidia.com/gpu", "ephemeral-storage" ]
       flavors:
       - name: default-flavor
         resources:
         - name: "cpu"
           nominalQuota: 72
         - name: "memory"
           nominalQuota: 288Gi
         - name: "nvidia.com/gpu"
           nominalQuota: 6
         - name: "ephemeral-storage"
           nominalQuota: 200Gi
   

   LocalQueue

   Este manifiesto crea una LocalQueue de Kueue llamada lq en el espacio de nombres llm.

   apiVersion: kueue.x-k8s.io/v1beta1
   kind: LocalQueue
   metadata:
     namespace: llm # LocalQueue under llm namespace 
     name: lq
   spec:
     clusterQueue: cluster-queue # Point to the ClusterQueue
   

5. Consulta los archivos default-priorityclass.yaml, low-priorityclass.yaml y high-priorityclass.yaml. Estos manifiestos definen los objetos PriorityClass para la programación de Kubernetes.

   Prioridad predeterminada

   apiVersion: scheduling.k8s.io/v1
   kind: PriorityClass
   metadata:
     name: default-priority-nonpreempting
   value: 10
   preemptionPolicy: Never
   globalDefault: true
   description: "This priority class will not cause other pods to be preempted."
   

   Prioridad baja

   apiVersion: scheduling.k8s.io/v1
   kind: PriorityClass
   metadata:
     name: low-priority-preempting
   value: 20
   preemptionPolicy: PreemptLowerPriority
   globalDefault: false
   description: "This priority class will cause pods with lower priority to be preempted."
   

   Prioridad alta

   apiVersion: scheduling.k8s.io/v1
   kind: PriorityClass
   metadata:
     name: high-priority-preempting
   value: 30
   preemptionPolicy: PreemptLowerPriority
   globalDefault: false
   description: "This high priority class will cause other pods to be preempted."
   

6. Crea los objetos de Kueue y Kubernetes ejecutando estos comandos para aplicar los manifiestos correspondientes.

   cd ${EXAMPLE_HOME}/workloads
   kubectl apply --filename=flavors.yaml
   kubectl apply --filename=default-priorityclass.yaml
   kubectl apply --filename=high-priorityclass.yaml
   kubectl apply --filename=low-priorityclass.yaml
   kubectl apply --filename=cluster-queue.yaml
   kubectl apply --filename=local-queue.yaml --namespace=llm
   

Desplegar el servidor de inferencia de TGI

En esta sección, desplegarás el contenedor de TGI para servir el modelo Gemma 2.

1. En el directorio /workloads, consulta el archivo tgi-gemma-2-9b-it-hp.yaml. Este manifiesto define un despliegue de Kubernetes para desplegar el tiempo de ejecución de servicio de TGI y el modelo gemma-2-9B-it. Un Deployment es un objeto de la API de Kubernetes que te permite ejecutar varias réplicas de pods distribuidas entre los nodos de un clúster.

   El despliegue prioriza las tareas de inferencia y usa dos GPUs para el modelo. Usa el paralelismo de tensores. Para ello, define la variable de entorno NUM_SHARD para que el modelo quepa en la memoria de la GPU.

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: tgi-gemma-deployment
     labels:
       app: gemma-server
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: gemma-server
     template:
       metadata:
         labels:
           app: gemma-server
           ai.gke.io/model: gemma-2-9b-it
           ai.gke.io/inference-server: text-generation-inference
           examples.ai.gke.io/source: user-guide
           kueue.x-k8s.io/queue-name: lq
       spec:
         priorityClassName: high-priority-preempting
         containers:
         - name: inference-server
           image: us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310
           resources:
             requests:
               cpu: "4"
               memory: "30Gi"
               ephemeral-storage: "30Gi"
               nvidia.com/gpu: "2"
             limits:
               cpu: "4"
               memory: "30Gi"
               ephemeral-storage: "30Gi"
               nvidia.com/gpu: "2"
           env:
           - name: AIP_HTTP_PORT
             value: '8000'
           - name: NUM_SHARD
             value: '2'
           - name: MODEL_ID
             value: google/gemma-2-9b-it
           - name: HUGGING_FACE_HUB_TOKEN
             valueFrom:
               secretKeyRef:
                 name: hf-secret
                 key: hf_api_token
           volumeMounts:
           - mountPath: /dev/shm
             name: dshm
         volumes:
         - name: dshm
           emptyDir:
             medium: Memory
         nodeSelector:
           cloud.google.com/gke-accelerator: "nvidia-l4"
   ---
   apiVersion: v1
   kind: Service
   metadata:
     name: llm-service
   spec:
     selector:
       app: gemma-server
     type: ClusterIP
     ports:
     - protocol: TCP
       port: 8000
       targetPort: 8000
   

2. Aplica el manifiesto ejecutando el siguiente comando:

   kubectl apply --filename=tgi-gemma-2-9b-it-hp.yaml --namespace=llm
   

   La operación de implementación tardará unos minutos en completarse.

3. Para comprobar si GKE ha creado correctamente la implementación, ejecuta el siguiente comando:

   kubectl --namespace=llm get deployment
   

   La salida debería ser similar a la siguiente:

   NAME                   READY   UP-TO-DATE   AVAILABLE   AGE
   tgi-gemma-deployment   1/1     1            1           5m13s
   

Verificar la gestión de cuotas de Kueue

En esta sección, confirmará que Kueue está aplicando correctamente la cuota de GPU de su Deployment.

1. Para comprobar si Kueue conoce tu Deployment, ejecuta este comando para obtener el estado de los objetos Workload:

   kubectl --namespace=llm get workloads
   

   La salida debería ser similar a la siguiente:

   NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
   pod-tgi-gemma-deployment-6bf9ffdc9b-zcfrh-84f19   lq      cluster-queue   True                  8m23s
   

2. Para probar la anulación de los límites de cuota, escala la implementación a cuatro réplicas:

   kubectl scale --replicas=4 deployment/tgi-gemma-deployment --namespace=llm
   

3. Ejecuta el siguiente comando para ver el número de réplicas que implementa GKE:

   kubectl get workloads --namespace=llm
   

   La salida debería ser similar a la siguiente:

   NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
   pod-tgi-gemma-deployment-6cb95cc7f5-5thgr-3f7d4   lq      cluster-queue   True                  14s
   pod-tgi-gemma-deployment-6cb95cc7f5-cbxg2-d9fe7   lq      cluster-queue   True                  5m41s
   pod-tgi-gemma-deployment-6cb95cc7f5-tznkl-80f6b   lq                                            13s
   pod-tgi-gemma-deployment-6cb95cc7f5-wd4q9-e4302   lq      cluster-queue   True                  13s
   

   El resultado muestra que solo se admiten tres pods debido a la cuota de recursos que aplica Kueue.

4. Ejecuta lo siguiente para mostrar los pods en el espacio de nombres llm:

   kubectl get pod --namespace=llm
   

   La salida debería ser similar a la siguiente:

   NAME                                    READY   STATUS            RESTARTS   AGE
   tgi-gemma-deployment-7649884d64-6j256   1/1     Running           0          4m45s
   tgi-gemma-deployment-7649884d64-drpvc   0/1     SchedulingGated   0          7s
   tgi-gemma-deployment-7649884d64-thdkq   0/1     Pending           0          7s
   tgi-gemma-deployment-7649884d64-znvpb   0/1     Pending           0          7s
   

5. Ahora, reduce el escalado de la implementación a 1. Este paso es obligatorio antes de implementar el trabajo de ajuste, ya que, de lo contrario, no se admitirá porque el trabajo de inferencia tiene prioridad.

   kubectl scale --replicas=1 deployment/tgi-gemma-deployment --namespace=llm
   

Explicación del comportamiento

En el ejemplo de escalado solo se obtienen tres réplicas (aunque se haya escalado a cuatro) debido al límite de cuota de GPU que has definido en la configuración de ClusterQueue. La sección spec.resourceGroups de ClusterQueue define una nominalQuota de "6" para nvidia.com/gpu. La implementación especifica que cada Pod requiere 2 GPUs. Por lo tanto, ClusterQueue solo puede alojar un máximo de tres réplicas de Deployment a la vez (ya que 3 réplicas * 2 GPUs por réplica = 6 GPUs, que es la cuota total).

Cuando intentas escalar a cuatro réplicas, Kueue reconoce que esta acción superaría la cuota de GPU y evita que se programe la cuarta réplica. Esto se indica con el estado SchedulingGated del cuarto pod. Este comportamiento demuestra la aplicación de cuotas de recursos de Kueue.

Desplegar la tarea de entrenamiento

En esta sección, desplegarás un trabajo de ajuste fino de menor prioridad para un modelo de Gemma 2 que requiere cuatro GPUs en dos pods. Un controlador de trabajo de Kubernetes crea uno o varios pods y se asegura de que ejecuten correctamente una tarea específica.

Esta tarea usará la cuota de GPU restante en ClusterQueue. El trabajo usa una imagen precompilada y guarda puntos de control para permitir que se reinicie a partir de resultados intermedios.

El trabajo de ajuste fino usa el conjunto de datos b-mc2/sql-create-context. La fuente del trabajo de ajuste se puede encontrar en el repositorio.

1. Consulta el archivo fine-tune-l4.yaml. Este archivo de manifiesto define el trabajo de ajuste.

   apiVersion: v1
   kind: Service
   metadata:
     name: headless-svc-l4
   spec:
     clusterIP: None # clusterIP must be None to create a headless service
     selector:
       job-name: finetune-gemma-l4 # must match Job name
   ---
   apiVersion: batch/v1
   kind: Job
   metadata:
     name: finetune-gemma-l4
     labels:
       kueue.x-k8s.io/queue-name: lq
   spec:
     backoffLimit: 4
     completions: 2
     parallelism: 2
     completionMode: Indexed
     suspend: true # Set to true to allow Kueue to control the Job when it starts
     template:
       metadata:
         labels:
           app: finetune-job
         annotations:
           gke-gcsfuse/volumes: "true"
           gke-gcsfuse/memory-limit: "35Gi"
       spec:
         priorityClassName: low-priority-preempting
         containers:
         - name: gpu-job
           imagePullPolicy: Always
           image: us-docker.pkg.dev/google-samples/containers/gke/gemma-fine-tuning:v1.0.0
           ports:
           - containerPort: 29500
           resources:
             requests:
               nvidia.com/gpu: "2"
             limits:
               nvidia.com/gpu: "2"
           command:
           - bash
           - -c
           - |
             accelerate launch \
             --config_file fsdp_config.yaml \
             --debug \
             --main_process_ip finetune-gemma-l4-0.headless-svc-l4 \
             --main_process_port 29500 \
             --machine_rank ${JOB_COMPLETION_INDEX} \
             --num_processes 4 \
             --num_machines 2 \
             fine_tune.py
           env:
           - name: "EXPERIMENT"
             value: "finetune-experiment"
           - name: MODEL_NAME
             value: "google/gemma-2-2b"
           - name: NEW_MODEL
             value: "gemma-ft"
           - name: MODEL_PATH
             value: "/model-data/model-gemma2/experiment"
           - name: DATASET_NAME
             value: "b-mc2/sql-create-context"
           - name: DATASET_LIMIT
             value: "5000"
           - name: EPOCHS
             value: "1"
           - name: GRADIENT_ACCUMULATION_STEPS
             value: "2"
           - name: CHECKPOINT_SAVE_STEPS
             value: "10"
           - name: HF_TOKEN
             valueFrom:
               secretKeyRef:
                 name: hf-secret
                 key: hf_api_token
           volumeMounts:
           - mountPath: /dev/shm
             name: dshm
           - name: gcs-fuse-csi-ephemeral
             mountPath: /model-data
             readOnly: false
         nodeSelector:
           cloud.google.com/gke-accelerator: nvidia-l4
         restartPolicy: OnFailure
         serviceAccountName: default
         subdomain: headless-svc-l4
         terminationGracePeriodSeconds: 60
         volumes:
         - name: dshm
           emptyDir:
             medium: Memory
         - name: gcs-fuse-csi-ephemeral
           csi:
             driver: gcsfuse.csi.storage.gke.io
             volumeAttributes:
               bucketName: <MODEL_BUCKET>
               mountOptions: "implicit-dirs"
               gcsfuseLoggingSeverity: warning
   

2. Aplica el archivo de manifiesto para crear el trabajo de ajuste fino:

   cd ${EXAMPLE_HOME}/workloads
   
   sed -e "s/<MODEL_BUCKET>/$MODEL_BUCKET/g" \
       -e "s/<PROJECT_ID>/$PROJECT_ID/g" \
       -e "s/<REGION>/$REGION/g" \
       fine-tune-l4.yaml |kubectl apply --filename=- --namespace=llm
   

3. Verifica que tus implementaciones se estén ejecutando. Para comprobar el estado de los objetos de carga de trabajo, ejecuta el siguiente comando:

   kubectl get workloads --namespace=llm
   

   La salida debería ser similar a la siguiente:

   NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
   job-finetune-gemma-l4-3316f                       lq      cluster-queue   True                  29m
   pod-tgi-gemma-deployment-6cb95cc7f5-cbxg2-d9fe7   lq      cluster-queue   True                  68m
   

   A continuación, consulta los pods del espacio de nombres llm ejecutando este comando:

   kubectl get pod --namespace=llm
   

   La salida debería ser similar a la siguiente:

   NAME                                    READY   STATUS    RESTARTS   AGE
   finetune-gemma-l4-0-vcxpz               2/2     Running   0          31m
   finetune-gemma-l4-1-9ppt9               2/2     Running   0          31m
   tgi-gemma-deployment-6cb95cc7f5-cbxg2   1/1     Running   0          70m
   

   El resultado muestra que Kueue permite que se ejecuten tanto tu trabajo de ajuste como los pods del servidor de inferencia, y reserva los recursos correctos en función de los límites de cuota que hayas especificado.

4. Consulta los registros de salida para verificar que la tarea de ajuste fino guarda los puntos de control en el segmento de Cloud Storage. La tarea de ajuste fino tarda unos 10 minutos en empezar a guardar el primer punto de control.

   kubectl logs --namespace=llm --follow --selector=app=finetune-job
   

   El resultado del primer punto de control guardado será similar al siguiente:

   {"name": "finetune", "thread": 133763559483200, "threadName": "MainThread", "processName": "MainProcess", "process": 33, "message": "Fine tuning started", "timestamp": 1731002351.0016131, "level": "INFO", "runtime": 451579.89835739136}
   …
   {"name": "accelerate.utils.fsdp_utils", "thread": 136658669348672, "threadName": "MainThread", "processName": "MainProcess", "process": 32, "message": "Saving model to /model-data/model-gemma2/experiment/checkpoint-10/pytorch_model_fsdp_0", "timestamp": 1731002386.1763802, "level": "INFO", "runtime": 486753.8924217224}
   

Probar la apropiación y la asignación dinámica de Kueue en una carga de trabajo mixta

En esta sección, simularás una situación en la que aumenta la carga del servidor de inferencia, lo que requiere que se escale. En este caso, se muestra cómo Kueue prioriza el servidor de inferencia de alta prioridad suspendiendo y adelantando el trabajo de ajuste fino de menor prioridad cuando los recursos son limitados.

1. Ejecuta el siguiente comando para ampliar las réplicas del servidor de inferencia a dos:

   kubectl scale --replicas=2 deployment/tgi-gemma-deployment --namespace=llm
   

2. Comprueba el estado de los objetos Workload:

   kubectl get workloads --namespace=llm
   

   El resultado es similar al siguiente:

   NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
   job-finetune-gemma-l4-3316f                       lq                      False                 32m
   pod-tgi-gemma-deployment-6cb95cc7f5-cbxg2-d9fe7   lq      cluster-queue   True                  70m
   pod-tgi-gemma-deployment-6cb95cc7f5-p49sh-167de   lq      cluster-queue   True                  14s
   

   El resultado muestra que el trabajo de ajuste no se admite porque las réplicas del servidor de inferencia aumentadas están usando la cuota de GPU disponible.

3. Comprueba el estado de la tarea de ajuste fino:

   kubectl get job --namespace=llm
   

   La salida será similar a la siguiente, lo que indica que el estado de la tarea de ajuste fino ahora es "suspended":

   NAME                STATUS      COMPLETIONS   DURATION   AGE
   finetune-gemma-l4   Suspended   0/2                      33m
   

4. Ejecuta el siguiente comando para inspeccionar tus pods:

   kubectl get pod --namespace=llm
   

   El resultado será similar al siguiente, lo que indica que Kueue ha terminado los pods de trabajo de ajuste para liberar recursos para la implementación del servidor de inferencia de mayor prioridad.

   NAME                                    READY   STATUS              RESTARTS   AGE
   tgi-gemma-deployment-6cb95cc7f5-cbxg2   1/1     Running             0          72m
   tgi-gemma-deployment-6cb95cc7f5-p49sh   0/1     ContainerCreating   0          91s
   

5. A continuación, prueba el caso en el que la carga del servidor de inferencia disminuye y sus pods se reducen. Ejecuta el siguiente comando:

   kubectl scale --replicas=1 deployment/tgi-gemma-deployment --namespace=llm
   

   Ejecuta el siguiente comando para mostrar los objetos Workload:

   kubectl get workloads --namespace=llm
   

   El resultado es similar al siguiente, lo que indica que se ha terminado una de las implementaciones del servidor de inferencia y que se ha vuelto a admitir el trabajo de ajuste.

   NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
   job-finetune-gemma-l4-3316f                       lq      cluster-queue   True                  37m
   pod-tgi-gemma-deployment-6cb95cc7f5-cbxg2-d9fe7   lq      cluster-queue   True                  75m
   

6. Ejecuta este comando para mostrar los trabajos:

   kubectl get job --namespace=llm
   

   El resultado es similar al siguiente, lo que indica que el trabajo de ajuste fino se está ejecutando de nuevo y se reanuda desde el último punto de control disponible.

   NAME                STATUS    COMPLETIONS   DURATION   AGE
   finetune-gemma-l4   Running   0/2           2m11s      38m