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

En este instructivo, se muestra cómo compartir recursos de acelerador de manera eficiente entre cargas de trabajo de entrenamiento y entrega de inferencias en un solo clúster de Google Kubernetes Engine (GKE). Cuando distribuyes tus cargas de trabajo mixtas en un solo clúster, mejoras la utilización de recursos, simplificas la administración del clúster, reduces los problemas de las limitaciones de cantidad de aceleradores y mejoras 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 de TGI de Hugging Face (interfaz de generación de texto), 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. Usas Kueue, un sistema de colas de trabajo nativo de Kubernetes de código abierto, para administrar y programar tus cargas de trabajo. Kueue te permite priorizar las tareas de publicación y anular las tareas de entrenamiento de prioridad más baja para optimizar el uso de recursos. A medida que disminuyen las demandas de entrega, reasignas los aceleradores liberados para reanudar las tareas de entrenamiento. Usas Kueue y clases de prioridad para administrar las cuotas de recursos durante el proceso.

Este instructivo está dirigido a ingenieros de aprendizaje automático (AA), administradores y operadores de plataformas, y especialistas en datos y IA que desean entrenar y alojar un modelo de aprendizaje automático (AA) en un clúster de GKE, y que también desean reducir los costos y la sobrecarga administrativa, en especial cuando se trata de 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 Tareas y roles comunes de los usuarios de GKE Enterprise.

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

Objetivos

Al final de esta guía, deberías poder realizar los siguientes pasos:

  • Configura una implementación de publicación de alta prioridad.
  • Configura trabajos de entrenamiento de prioridad más baja.
  • Implementa estrategias de anulación para abordar la demanda variable.
  • Administra la asignación de recursos entre las tareas de entrenamiento y entrega con Kueue.

Antes de comenzar

  • Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required APIs.

    Enable the APIs

    •

  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required APIs.

    Enable the APIs

    •

  • Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin

    Check for the roles

    1. In the Google Cloud console, go to the IAM page.

      Go to IAM
    2. Select the project.

    3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

    4. For all rows that specify or include you, check the Role column to see whether the list of roles includes the required roles.

    Grant the roles

    1. In the Google Cloud console, go to the IAM page.

      Ir a IAM
    2. Selecciona el proyecto.
    3. Haz clic en Grant access.

    4. En el campo Principales nuevas, ingresa tu identificador de usuario. Esta suele ser la dirección de correo electrónico de una Cuenta de Google.

    5. En la lista Seleccionar un rol, elige un rol.
    6. Para otorgar funciones adicionales, haz clic en Agregar otro rol y agrega cada rol adicional.
    7. Haz clic en Guardar.

Prepare el entorno

En esta sección, aprovisionarás los recursos que necesitas para implementar la 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 la 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 en Google 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 consola de Google Cloud , haz clic en Ícono de activación de Cloud Shell Activar Cloud Shell en la consola de Google Cloud 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 el ID del proyecto de tu Google Cloud.

  3. Clone el código de muestra desde GitHub. En Cloud Shell, ejecute 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)

Cree un clúster de GKE

Puedes usar un clúster de Autopilot o Standard para tus cargas de trabajo combinadas. 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: Es una región que admite 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} \
    --region=${REGION} \
    --release-channel=rapid

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

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

  4. Para otorgar acceso al bucket de Cloud Storage, ejecuta el siguiente 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 \
    --region=$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 GPU para cargas de trabajo de procesamiento.
    • GPU_POOL_ACCELERATOR_TYPE: Es el tipo de acelerador que se admite en 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 GPU 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} \
    --region=${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 bucket de Cloud Storage para el trabajo de ajuste fino:

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

  5. Para otorgar acceso al bucket de Cloud Storage, ejecuta el siguiente 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 \
    --region=$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 GPUs de manera eficiente entre las cargas de trabajo de entrenamiento y publicación. Kueue logra esto definiendo los requisitos de recursos ("variantes"), priorizando las cargas de trabajo a través de colas (con tareas de entrega priorizadas sobre el entrenamiento) y asignando recursos de forma dinámica en función de la demanda y la prioridad. En este instructivo, se usa el tipo de recurso Workload para agrupar las cargas de trabajo de inferencia y de ajuste fino, respectivamente.

La función de prioridad de Kueue garantiza que las cargas de trabajo de publicación de alta prioridad siempre tengan los recursos necesarios, ya que pausa o desaloja 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 v1/pod. Para ello, aplica una configuración personalizada con Kustomize para garantizar que los Pods del servidor estén etiquetados con "kueue-job: true".

  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.10.0/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 administrar Pods con la etiqueta "kueue-job: true".

    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
    #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/mxjob"
      - "kubeflow.org/paddlejob"
      - "kubeflow.org/pytorchjob"
      - "kubeflow.org/tfjob"
      - "kubeflow.org/xgboostjob"
      - "pod"
    #  externalFrameworks:
    #  - "Foo.v1.example.com"
      podOptions:
        # You can change namespaceSelector to define in which 
        # namespaces kueue will manage the pods.
        namespaceSelector:
          matchExpressions:
          - key: kubernetes.io/metadata.name
            operator: NotIn
            values: [ kube-system, kueue-system ]
        # Kueue uses podSelector to manage pods with particular 
        # labels. The default podSelector will match all the pods. 
        podSelector:
          matchExpressions:
          - key: kueue-job
            operator: In
            values: [ "true", "True", "yes" ]

  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    2/2     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 una ClusterQueue de Kueue con límites de recursos para la CPU, la memoria y la GPU.

    En este instructivo, se usan nodos con dos GPUs Nvidia L4 conectadas, con el tipo de nodo correspondiente de g2-standard-24, que ofrece 24 vCPU y 96 GB de RAM. En el código de ejemplo, se 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 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 Gemma 2.

  1. En el directorio /workloads, consulta el archivo tgi-gemma-2-9b-it-hp.yaml. En este manifiesto, se define una implementación de Kubernetes para implementar el entorno de ejecución de entrega de TGI y el modelo gemma-2-9B-it.

    La implementación prioriza las tareas de inferencia y usa dos GPUs para el modelo. Usa paralelismo de tensores, ya que establece la variable de entorno NUM_SHARD para que el modelo se ajuste 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
        kueue-job: "true"
    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. Ejecuta el siguiente comando para aplicar el manifiesto:

    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ó la Deployment correctamente, 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 aplica correctamente la cuota de GPU para tu Deployment.

  1. Para verificar si Kueue está al tanto de tu Deployment, ejecuta este comando para recuperar el estado de los objetos de 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 GKE implementa:

    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, vuelve a reducir la escala de la Deployment a 1. Este paso es obligatorio antes de implementar la tarea de ajuste fino. De lo contrario, no se admitirá porque la tarea de inferencia tiene prioridad.

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

Explicación del comportamiento

El ejemplo de escalamiento genera solo tres réplicas (a pesar de escalar a cuatro) debido al límite de cuota de GPU que estableceste en la configuración de ClusterQueue. La sección spec.resourceGroups de ClusterQueue define un valor de 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 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 forzosa de cuotas de recursos de Kueue.

Implementa el trabajo de entrenamiento

En esta sección, implementarás un trabajo de ajuste fino de prioridad más baja para un modelo Gemma 2 que requiere cuatro GPUs en dos Pods. Esta tarea usará la cuota de GPU restante en ClusterQueue. La tarea usa una imagen precompilada y guarda puntos de control para permitir el reinicio desde resultados intermedios.

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

  1. Consulta el archivo fine-tune-l4.yaml. En este manifiesto, se define el trabajo de ajuste fino.

    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 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 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 el siguiente 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 que se ejecuten los Pods del servidor de inferencia y del trabajo de ajuste fino, y reserva los recursos correctos según los límites de cuota especificados.

  4. Consulta los registros de salida para verificar que tu trabajo de ajuste fino guarde los puntos de control en el bucket de Cloud Storage. La tarea de ajuste fino 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 agrande. En esta situación, se muestra cómo Kueue prioriza el servidor de inferencia de alta prioridad suspendiendo y anulando la tarea de ajuste fino de menor prioridad cuando los recursos están 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 la tarea de ajuste fino porque las réplicas del servidor de inferencia aumentadas usan la cuota de GPU disponible.

  3. Verifica el estado del trabajo de ajuste fino:

    kubectl get job --namespace=llm

    El resultado es similar al siguiente, lo que indica que el estado de la tarea de ajuste fino 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 de trabajo de ajuste fino para liberar recursos para la Deployment del servidor de inferencia de prioridad más alta.

    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 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 de Workload:

    kubectl get workloads --namespace=llm

    El resultado es similar al siguiente, lo que indica que se cerró una de las implementaciones del servidor de inferencia y que se volvió a admitir la tarea de ajuste fino.

    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 la tarea de ajuste fino se está ejecutando nuevamente y se reanuda desde el punto de control más reciente disponible.

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

Limpia

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

Borra los recursos implementados

Para evitar que se generen cargos en tu cuenta de Google Cloud por los recursos que creaste en esta guía, ejecuta los siguientes comandos:

gcloud storage rm --recursive gs://${MODEL_BUCKET}
gcloud container clusters delete ${CLUSTER_NAME} --location ${REGION}

¿Qué sigue?