Optimiza el uso de recursos de GKE para cargas de trabajo mixtas de entrenamiento e inferencia de IA/AA

En este instructivo, se muestra cómo compartir de manera eficiente los recursos del acelerador entre las cargas de trabajo de entrenamiento y de entrega de inferencias dentro de un solo clúster de Google Kubernetes Engine (GKE). Si distribuyes tus cargas de trabajo mixtas en un solo clúster, mejorarás el uso de recursos, simplificarás la administración del clúster, reducirás los problemas derivados de las limitaciones en la cantidad de aceleradores y mejorarás la rentabilidad general.

En este instructivo, crearás una Deployment de entrega de alta prioridad con el modelo de lenguaje grande (LLM) Gemma 2 para la inferencia y el framework de entrega Hugging Face TGI (Interfaz de generación de texto), junto con un trabajo de ajuste del LLM de baja prioridad. Ambas cargas de trabajo se ejecutan en un solo clúster que usa GPUs NVIDIA L4. Usas Kueue, un sistema de cola de trabajos nativo de Kubernetes de código abierto, para administrar y programar tus cargas de trabajo. Kueue te permite priorizar las tareas de entrega y detener los trabajos de entrenamiento de menor prioridad para optimizar el uso de recursos. A medida que disminuye la demanda de la publicación, reasigna los aceleradores liberados para reanudar los trabajos de entrenamiento. Usas Kueue y las clases de prioridad para administrar las cuotas de recursos durante todo el proceso.

Este instructivo está dirigido a ingenieros de aprendizaje automático (AA), administradores y operadores de plataformas, y especialistas en IA y datos que deseen entrenar y alojar un modelo de aprendizaje automático (AA) en un clúster de GKE, y que también deseen reducir los costos y la sobrecarga de administración, en especial cuando se trabaja con una cantidad limitada de aceleradores. Para obtener más información sobre los roles comunes y las tareas de ejemplo a las que hacemos referencia en el contenido de Google Cloud , consulta Roles de usuario y tareas comunes de GKE.

Antes de leer esta página, asegúrate de estar familiarizado con lo siguiente:

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

Prepare el entorno

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

Obtén acceso al modelo

Para obtener acceso a los modelos de Gemma para la implementación en GKE, primero debes firmar el contrato de consentimiento de licencia y, luego, generar un token de acceso de Hugging Face.

1. Firma el acuerdo de consentimiento de licencia. Accede a la página de consentimiento del modelo, verifica el consentimiento con tu cuenta de Hugging Face y acepta las condiciones 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 nuevo si aún no tienes uno:

   1. Haz clic en Tu perfil > Configuración > Tokens de acceso.
   2. Selecciona Token nuevo.
   3. Especifica el nombre que desees y un rol de al menos Read.
   4. Selecciona Genera un token.
   5. Copia el token generado al portapapeles.

Inicia Cloud Shell

En este instructivo, usarás Cloud Shell para administrar recursos alojados enGoogle Cloud. Cloud Shell tiene preinstalado el software que necesitas para este instructivo, incluidos kubectl, la CLI de gcloud y Terraform.

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

1. En la Google Cloud consola, haz clic en Activar Cloud Shell en la Google Cloud consola para iniciar una sesión de Cloud Shell. Esto inicia una sesión en el panel inferior de la consola de Google Cloud .

2. Configura las variables de entorno predeterminadas:

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

   Reemplaza PROJECT_ID por tu Google Cloud ID del proyecto.

3. Clone el código de muestra desde 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)
   

Crea un clúster de GKE

Puedes usar un clúster de Autopilot o Standard para tus cargas de trabajo mixtas. Te recomendamos que uses un clúster de Autopilot para una experiencia de Kubernetes completamente administrada. Para elegir el modo de operación de GKE que se adapte mejor a tus cargas de trabajo, consulta Elige un modo de operación de GKE.

Autopilot

1. Configura 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"
   

   Reemplaza los siguientes valores:

   • HF_TOKEN: Es el token de Hugging Face que generaste antes.
   • REGION: Una región que admita el tipo de acelerador que deseas usar, por ejemplo, us-central1 para la GPU L4.

   Puedes ajustar la variable MODEL_BUCKET, que representa el bucket 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 bucket de Cloud Storage para el trabajo de ajuste:

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

4. Para otorgar acceso al bucket 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. Configura 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"
   

   Reemplaza los siguientes valores:

   • HF_TOKEN: Es el token de Hugging Face que generaste antes.
   • REGION: Es la región que admite el tipo de acelerador que deseas usar, por ejemplo, us-central1 para la GPU L4.

   Puedes ajustar estas variables:

   • GPU_POOL_MACHINE_TYPE: Es la serie de máquinas del grupo de nodos que deseas usar en la región seleccionada. Este valor depende del tipo de acelerador que seleccionaste. Para obtener más información, consulta Limitaciones del uso de GPUs en GKE. Por ejemplo, en este instructivo, se usa g2-standard-24 con dos GPUs conectadas por nodo. Para obtener la lista más actualizada de las GPUs disponibles, consulta GPUs para cargas de trabajo de procesamiento.
   • GPU_POOL_ACCELERATOR_TYPE: Es el tipo de acelerador compatible con la región que seleccionaste. Por ejemplo, en este instructivo, se usa nvidia-l4. Para obtener la lista más reciente de las GPUs disponibles, consulta GPUs para cargas de trabajo de procesamiento.
   • MODEL_BUCKET: Es el bucket de Cloud Storage en el que almacenas 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:

   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 bucket de Cloud Storage para el trabajo de ajuste:

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

5. Para otorgar acceso al bucket 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
   

Crea un secreto de Kubernetes para las credenciales de Hugging Face

Para crear un Secret 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=-

Configura Kueue

En este instructivo, Kueue es el administrador de recursos central, lo que permite compartir de manera eficiente las GPUs entre tus cargas de trabajo de entrenamiento y de entrega. Kueue logra esto definiendo los requisitos de recursos ("sabores"), priorizando las cargas de trabajo a través de colas (con tareas de servicio priorizadas sobre el entrenamiento) y asignando recursos de forma dinámica según la demanda y la prioridad. En este instructivo, se usa el tipo de recurso Workload para agrupar las cargas de trabajo de inferencia y ajuste, respectivamente.

La función de interrupción de Kueue garantiza que las cargas de trabajo de procesamiento de prioridad alta siempre tengan los recursos necesarios, ya que pausa o expulsa los trabajos de entrenamiento de prioridad más baja cuando los recursos son escasos.

Para controlar la Deployment 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 administrador de recursos de 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 administració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
   

   El resultado debería ser similar al 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 Kueue administra las cuotas de recursos:

   ResourceFlavor

   En este manifiesto, se define un ResourceFlavor predeterminado en Kueue para la administración de recursos.

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

   ClusterQueue

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

   En este instructivo, se usan nodos con dos GPU Nvidia L4 conectadas, con el tipo de nodo correspondiente de g2-standard-24, que ofrece 24 CPU virtuales 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 en la configuración de ClusterQueue hace referencia a las PriorityClasses para determinar qué Pods se pueden interrumpir 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. Ejecuta estos comandos para aplicar los manifiestos correspondientes y crear los objetos de Kueue y Kubernetes.

   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
   

Implementa el servidor de inferencia de TGI

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

1. En el directorio /workloads, consulta el archivo tgi-gemma-2-9b-it-hp.yaml. Este manifiesto define una implementación de Kubernetes para implementar el tiempo de ejecución de entrega 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 que se distribuyen entre los nodos de un clúster.

   La implementación prioriza las tareas de inferencia y usa dos GPUs para el modelo. Usa el paralelismo de tensor configurando la variable de entorno NUM_SHARD para ajustar el modelo a 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 verificar si GKE creó correctamente la Deployment, ejecuta el siguiente comando:

   kubectl --namespace=llm get deployment
   

   El resultado debería ser similar al siguiente:

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

Verifica la administración de cuotas de Kueue

En esta sección, confirmarás que Kueue aplique correctamente la cuota de GPU para tu Deployment.

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

   kubectl --namespace=llm get workloads
   

   El resultado debería ser similar al 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 Deployment a cuatro réplicas:

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

3. Ejecuta el siguiente comando para ver la cantidad de réplicas que implementa GKE:

   kubectl get workloads --namespace=llm
   

   El resultado debería ser similar al 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 el siguiente comando para mostrar los Pods en el espacio de nombres llm:

   kubectl get pod --namespace=llm
   

   El resultado debería ser similar al 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 la escala del Deployment a 1. Este paso es obligatorio antes de implementar el trabajo de ajuste, ya que, de lo contrario, no se admitirá debido a que el trabajo de inferencia tiene prioridad.

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

Explicación del comportamiento

El ejemplo de ajuste de escala da como resultado solo tres réplicas (a pesar de que se ajustó a cuatro) debido al límite de cuota de GPU que estableciste en la configuración de ClusterQueue. La sección spec.resourceGroups de ClusterQueue define un nominalQuota de "6" para nvidia.com/gpu. La implementación especifica que cada Pod requiere "2" GPUs. Por lo tanto, ClusterQueue solo puede admitir un máximo de tres réplicas de la 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 excederí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 la cuota de recursos de Kueue.

Implementa el trabajo de entrenamiento

En esta sección, implementarás un trabajo de ajuste de baja prioridad para un modelo de Gemma 2 que requiere cuatro GPUs en dos Pods. Un controlador de Job en Kubernetes crea uno o más Pods y garantiza que ejecuten correctamente una tarea específica.

Este trabajo usará la cuota de GPU restante en ClusterQueue. El trabajo usa una imagen compilada previamente y guarda puntos de control para permitir el reinicio a partir de resultados intermedios.

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

1. Consulta el archivo fine-tune-l4.yaml. Este 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 manifiesto para crear el trabajo de ajuste:

   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 Deployments se estén ejecutando. Para verificar el estado de los objetos Workload, ejecuta el siguiente comando:

   kubectl get workloads --namespace=llm
   

   El resultado debería ser similar al 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, ejecuta este comando para ver los Pods en el espacio de nombres llm:

   kubectl get pod --namespace=llm
   

   El resultado debería ser similar al 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 admite la ejecución de tu trabajo de ajuste y de los Pods del servidor de inferencia, y reserva los recursos correctos según los límites de cuota especificados.

4. Visualiza los registros de salida para verificar que tu trabajo de ajuste fino guarde puntos de control en el bucket de Cloud Storage. El trabajo de ajuste tarda alrededor de 10 minutos antes de comenzar a guardar el primer punto de control.

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

   El resultado del primer punto de control guardado es 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}
   

Prueba la interrupción y la asignación dinámica de Kueue en tu 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 verticalmente. En este caso, se muestra cómo Kueue prioriza el servidor de inferencia de alta prioridad suspendiendo y deteniendo el trabajo de ajuste de baja prioridad cuando los recursos son limitados.

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

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

2. Verifica el estado de los objetos de 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 ya no se admite el trabajo de ajuste porque las réplicas aumentadas del servidor de inferencia están usando la cuota de GPU disponible.

3. Verifica el estado del trabajo de ajuste:

   kubectl get job --namespace=llm
   

   El resultado es similar al siguiente, lo que indica que el estado del trabajo de ajuste ahora está suspendido:

   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 es similar al siguiente, lo que indica que Kueue finalizó los Pods del Job de ajuste para liberar recursos para la Deployment 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 la situación en la que disminuye la carga del servidor de inferencia y se reduce la escala de sus Pods. Ejecuta el siguiente comando:

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

   Ejecuta el siguiente comando para mostrar los objetos de Workload:

   kubectl get workloads --namespace=llm
   

   El resultado es similar al siguiente, lo que indica que se finalizó una de las implementaciones del servidor de inferencia y se volvió 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 está en ejecución nuevamente y se reanuda desde el último punto de control disponible.

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