Otimizar a utilização de recursos do GKE para cargas de trabalho de treinamento e inferência mistas de IA/ML


Neste tutorial, mostramos como compartilhar recursos de acelerador de maneira eficiente entre cargas de trabalho de treinamento e de serviço de inferência em um único cluster do Google Kubernetes Engine (GKE). Ao distribuir suas cargas de trabalho mistas em um único cluster, você melhora a utilização de recursos, simplifica o gerenciamento de clusters, reduz problemas causados por limitações na quantidade de aceleradores e aumenta a relação custo-benefício geral.

Neste tutorial, você vai criar uma implantação de disponibilização de alta prioridade usando o modelo de linguagem grande (LLM) Gemma 2 para inferência e o framework de disponibilização TGI do Hugging Face (interface de geração de texto), além de um job de ajuste fino de LLM de baixa prioridade. As duas cargas de trabalho são executadas em um único cluster que usa GPUs NVIDIA L4. Você usa o Kueue, um sistema de enfileiramento de jobs nativo do Kubernetes de código aberto, para gerenciar e programar suas cargas de trabalho. Com o Kueue, é possível priorizar tarefas de veiculação e interromper jobs de treinamento de baixa prioridade para otimizar a utilização de recursos. À medida que as demandas de veiculação diminuem, você realoca os aceleradores liberados para retomar os jobs de treinamento. Você usa o Kueue e as classes de prioridade para gerenciar cotas de recursos durante todo o processo.

Este tutorial é destinado a engenheiros de machine learning (ML), administradores e operadores de plataforma e especialistas em dados e IA que querem treinar e hospedar um modelo de machine learning (ML) em um cluster do GKE, além de reduzir custos e sobrecarga de gerenciamento, principalmente ao lidar com um número limitado de aceleradores. Para saber mais sobre papéis comuns e tarefas de exemplo referenciados no conteúdo do Google Cloud , consulte Tarefas e funções de usuário comuns do GKE Enterprise.

Antes de ler esta página, confira se você conhece os seguintes conceitos:

Objetivos

Ao final deste guia, você será capaz de realizar as seguintes etapas:

  • Configure uma implantação de exibição de alta prioridade.
  • Configure jobs de treinamento de menor prioridade.
  • Implemente estratégias de substituição para lidar com a demanda variável.
  • Gerenciar a alocação de recursos entre tarefas de treinamento e veiculação usando o Kueue.

Antes de começar

  • 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.

      Acessar o IAM
    2. Selecionar um projeto.
    3. Clique em CONCEDER ACESSO.
    4. No campo Novos principais, insira seu identificador de usuário. Normalmente, é o endereço de e-mail de uma Conta do Google.

    5. Na lista Selecionar um papel, escolha um.
    6. Para conceder outros papéis, clique em Adicionar outro papel e adicione cada papel adicional.
    7. Clique em Salvar.

Prepare o ambiente

Nesta seção, você provisiona os recursos necessários para implantar o TGI e o modelo para suas cargas de trabalho de inferência e treinamento.

Receber acesso ao modelo

Para ter acesso aos modelos Gemma para implantação no GKE, primeiro assine o contrato de consentimento de licença e gere um token de acesso do Hugging Face.

  1. Assine o contrato de consentimento de licença. Acesse a página de consentimento do modelo, confirme seu consentimento usando sua conta do Hugging Face e aceite os termos do modelo.
  2. Gere um token de acesso. Para acessar o modelo pelo Hugging Face, você precisa de um token do Hugging Face. Siga as etapas abaixo para gerar um novo token, caso ainda não tenha um:

    1. Clique em Seu perfil > Configurações > Tokens de acesso.
    2. Selecione Novo token.
    3. Especifique um Nome de sua escolha e um Papel de pelo menos Read.
    4. Selecione Gerar um token.
    5. Copie o token gerado para a área de transferência.

Iniciar o Cloud Shell

Neste tutorial, você vai usar o Cloud Shell para gerenciar recursos hospedados no Google Cloud. O Cloud Shell vem pré-instalado com o software necessário para este tutorial, incluindo kubectl, a gcloud CLI e o Terraform.

Para configurar o ambiente com o Cloud Shell, siga estas etapas:

  1. No console Google Cloud , inicie uma sessão do Cloud Shell clicando em Ícone de ativação do Cloud Shell Ativar o Cloud Shell no consoleGoogle Cloud . Isso inicia uma sessão no painel inferior do console Google Cloud .

  2. Defina as variáveis de ambiente padrão:

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

    Substitua PROJECT_ID pelo Google Cloud ID do projeto.

  3. Clone o exemplo de código do GitHub. No Cloud Shell, execute estes comandos:

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

Criar um cluster do GKE

É possível usar um cluster do Autopilot ou Standard para suas cargas de trabalho mistas. Recomendamos que você use um cluster do Autopilot para ter uma experiência totalmente gerenciada do Kubernetes. Para escolher o modo de operação do GKE mais adequado para suas cargas de trabalho, consulte Escolher um modo de operação do GKE.

Piloto automático

  1. Defina as variáveis de ambiente padrão no 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"
    

    Substitua os seguintes valores:

    • HF_TOKEN: o token do Hugging Face que você gerou antes.
    • REGION: uma região compatível com o tipo de acelerador que você quer usar, por exemplo, us-central1 para a GPU L4.

    É possível ajustar a variável MODEL_BUCKET, que representa o bucket do Cloud Storage em que você armazena os pesos do modelo treinado.

  2. Criar um cluster do Autopilot:

    gcloud container clusters create-auto ${CLUSTER_NAME} \
        --project=${PROJECT_ID} \
        --location=${REGION} \
        --release-channel=rapid
    
  3. Crie o bucket do Cloud Storage para o job de ajuste refinado:

    gcloud storage buckets create gs://${MODEL_BUCKET} \
        --location ${REGION} \
        --uniform-bucket-level-access
    
  4. Para conceder acesso ao bucket do Cloud Storage, execute 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 receber as credenciais de autenticação do cluster, execute este comando:

    gcloud container clusters get-credentials llm-cluster \
        --location=$REGION \
        --project=$PROJECT_ID
    
  6. Crie um namespace para suas implantações. No Cloud Shell, execute o seguinte comando:

    kubectl create ns llm
    

Padrão

  1. Defina as variáveis de ambiente padrão no 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"
    

    Substitua os seguintes valores:

    • HF_TOKEN: o token do Hugging Face que você gerou antes.
    • REGION: a região compatível com o tipo de acelerador que você quer usar, por exemplo, us-central1 para a GPU L4.

    É possível ajustar estas variáveis:

    • GPU_POOL_MACHINE_TYPE: a série de máquinas do pool de nós que você quer usar na região selecionada. Esse valor depende do tipo de acelerador selecionado. Para saber mais, consulte Limitações do uso de GPUs no GKE. Por exemplo, este tutorial usa g2-standard-24 com duas GPUs anexadas por nó. Para conferir a lista mais atualizada de GPUs disponíveis, consulte GPUs para cargas de trabalho de computação.
    • GPU_POOL_ACCELERATOR_TYPE: o tipo de acelerador compatível com a região selecionada. Por exemplo, este tutorial usa nvidia-l4. Para conferir a lista mais recente de GPUs disponíveis, consulte GPUs para cargas de trabalho de computação.
    • MODEL_BUCKET: o bucket do Cloud Storage em que você armazena os pesos do modelo treinado.
  2. Criar um cluster padrão

    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. Crie o pool de nós de GPU para cargas de trabalho de inferência e ajuste refinado:

    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. Crie o bucket do Cloud Storage para o job de ajuste refinado:

    gcloud storage buckets create gs://${MODEL_BUCKET} \
        --location ${REGION} \
        --uniform-bucket-level-access
    
  5. Para conceder acesso ao bucket do Cloud Storage, execute 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 receber as credenciais de autenticação do cluster, execute este comando:

    gcloud container clusters get-credentials llm-cluster \
        --location=$REGION \
        --project=$PROJECT_ID
    
  7. Crie um namespace para suas implantações. No Cloud Shell, execute o seguinte comando:

    kubectl create ns llm
    

Criar um Secret do Kubernetes para as credenciais do Hugging Face

Para criar um secret do Kubernetes que contenha o token do Hugging Face, execute o seguinte 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 o Kueue

Neste tutorial, o Kueue é o gerenciador de recursos central, permitindo o compartilhamento eficiente de GPUs entre suas cargas de trabalho de treinamento e disponibilização. O Kueue faz isso definindo requisitos de recursos ("sabores"), priorizando cargas de trabalho por filas (com tarefas de veiculação priorizadas em relação ao treinamento) e alocando dinamicamente recursos com base na demanda e na prioridade. Este tutorial usa o tipo de recurso Workload para agrupar as cargas de trabalho de inferência e ajuste refinado, respectivamente.

O recurso de remoção preventiva do Kueue garante que as cargas de trabalho de veiculação de alta prioridade sempre tenham os recursos necessários, pausando ou removendo jobs de treinamento de baixa prioridade quando os recursos estão escassos.

Para controlar a implantação do servidor de inferência com o Kueue, ative a integração do pod e configure o managedJobsNamespaceSelector para excluir os namespaces kube-system e kueue-system.

  1. No diretório /kueue, confira o código em kustomization.yaml. Esse manifesto instala o gerenciador de recursos do Kueue com configurações 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. No diretório /kueue, confira o código em patch.yaml. Esse ConfigMap personaliza o Kueue para excluir o gerenciamento de pods nos namespaces kube-system e 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. No Cloud Shell, execute o comando a seguir para instalar o Kueue:

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

    Aguarde até que os pods do Kueue estejam prontos:

    watch kubectl --namespace=kueue-system get pods
    

    A saída será parecida com esta:

    NAME                                        READY   STATUS    RESTARTS   AGE
    kueue-controller-manager-bdc956fc4-vhcmx    1/1     Running   0          3m15s
    
  4. No diretório /workloads, confira os arquivos flavors.yaml, cluster-queue.yaml e local-queue.yaml. Esses manifestos especificam como o Kueue gerencia as cotas de recursos:

    ResourceFlavor

    Esse manifesto define um ResourceFlavor padrão no Kueue para gerenciamento de recursos.

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

    ClusterQueue

    Esse manifesto configura uma ClusterQueue do Kueue com limites de recursos para CPU, memória e GPU.

    Este tutorial usa nós com duas GPUs Nvidia L4 anexadas, com o tipo de nó correspondente g2-standard-24, oferecendo 24 vCPUs e 96 GB de RAM. O exemplo de código mostra como limitar o uso de recursos da sua carga de trabalho a um máximo de seis GPUs.

    O campo preemption na configuração do ClusterQueue faz referência às PriorityClasses para determinar quais pods podem ser preemptivos quando os recursos estão escassos.

    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

    Esse manifesto cria uma LocalQueue do Kueue chamada lq no namespace 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. Confira os arquivos default-priorityclass.yaml, low-priorityclass.yaml e high-priorityclass.yaml. Esses manifestos definem os objetos PriorityClass para o agendamento do Kubernetes.

    Prioridade padrão

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

    Prioridade baixa

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

    Prioridade 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. Crie os objetos do Kueue e do Kubernetes executando estes comandos para aplicar os manifestos correspondentes.

    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
    

Implantar o servidor de inferência TGI

Nesta seção, você implanta o contêiner do TGI para exibir o modelo Gemma 2.

  1. No diretório /workloads, confira o arquivo tgi-gemma-2-9b-it-hp.yaml. Esse manifesto define uma implantação do Kubernetes para implantar o ambiente de execução de serviço do TGI e o modelo gemma-2-9B-it. Uma implantação é um objeto da API Kubernetes que permite executar várias réplicas de pods distribuídos entre os nós de um cluster.

    A implantação prioriza tarefas de inferência e usa duas GPUs para o modelo. Ele usa o paralelismo de tensor, definindo a variável de ambiente NUM_SHARD, para ajustar o modelo à memória da 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. Execute este comando para aplicar o manifesto:

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

    A operação de implantação leva alguns minutos para ser concluída.

  3. Para verificar se o GKE criou a implantação, execute o seguinte comando:

    kubectl --namespace=llm get deployment
    

    A saída será parecida com esta:

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

Verificar o gerenciamento de cota do Kueue

Nesta seção, você confirma se o Kueue está aplicando corretamente a cota de GPU para sua implantação.

  1. Para verificar se o Kueue está ciente da sua implantação, execute este comando para recuperar o status dos objetos de carga de trabalho:

    kubectl --namespace=llm get workloads
    

    A saída será parecida com esta:

    NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
    pod-tgi-gemma-deployment-6bf9ffdc9b-zcfrh-84f19   lq      cluster-queue   True                  8m23s
    
  2. Para testar a substituição dos limites de cota, dimensione a implantação para quatro réplicas:

    kubectl scale --replicas=4 deployment/tgi-gemma-deployment --namespace=llm
    
  3. Execute o seguinte comando para ver o número de réplicas que o GKE implanta:

    kubectl get workloads --namespace=llm
    

    A saída será parecida com esta:

    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
    

    A saída mostra que apenas três pods são aceitos devido à cota de recursos imposta pelo Kueue.

  4. Execute o comando a seguir para mostrar os pods no namespace llm:

    kubectl get pod --namespace=llm
    

    A saída será parecida com esta:

    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. Agora, reduzir escala vertical a escala da implantação para 1. Essa etapa é necessária antes de implantar o job de ajuste refinado. Caso contrário, ele não será aceito porque o job de inferência tem prioridade.

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

Explicação do comportamento

O exemplo de escalonamento resulta em apenas três réplicas (apesar do escalonamento para quatro) devido ao limite de cota de GPU definido na configuração do ClusterQueue. A seção spec.resourceGroups de ClusterQueue define um nominalQuota de "6" para nvidia.com/gpu. A implantação especifica que cada pod requer "2" GPUs. Portanto, a ClusterQueue só pode acomodar no máximo três réplicas da implantação por vez (já que 3 réplicas * 2 GPUs por réplica = 6 GPUs, que é a cota total).

Quando você tenta escalonar para quatro réplicas, o Kueue reconhece que essa ação excederia a cota de GPU e impede que a quarta réplica seja programada. Isso é indicado pelo status SchedulingGated do quarto pod. Esse comportamento demonstra a aplicação de cotas de recursos do Kueue.

Implante o job de treinamento

Nesta seção, você implanta um job de ajuste refinado de prioridade mais baixa para um modelo Gemma 2 que exige quatro GPUs em dois pods. Um controlador de job no Kubernetes cria um ou mais pods e garante que eles executem uma tarefa específica com sucesso.

Esse job vai usar a cota restante de GPU na ClusterQueue. O job usa uma imagem pré-criada e salva pontos de verificação para permitir a reinicialização dos resultados intermediários.

O job de ajuste fino usa o conjunto de dados b-mc2/sql-create-context. A origem do job de ajuste detalhado pode ser encontrada no repositório.

  1. Visualize o arquivo fine-tune-l4.yaml. Esse manifesto define o job de ajuste refinado.

    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. Aplique o manifesto para criar o job 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. Verifique se as implantações estão em execução. Para verificar o status dos objetos de carga de trabalho, execute o seguinte comando:

    kubectl get workloads --namespace=llm
    

    A saída será parecida com esta:

    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
    

    Em seguida, veja os pods no namespace llm executando este comando:

    kubectl get pod --namespace=llm
    

    A saída será parecida com esta:

    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
    

    A saída mostra que o Kueue permite a execução do job de ajuste refinado e dos pods do servidor de inferência, reservando os recursos corretos com base nos limites de cota especificados.

  4. Confira os registros de saída para verificar se o job de ajuste refinado salva pontos de verificação no bucket do Cloud Storage. O job de ajuste refinado leva cerca de 10 minutos para começar a salvar o primeiro checkpoint.

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

    A saída do primeiro ponto de verificação salvo é semelhante a esta:

    {"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}
    

Testar a preempção e a alocação dinâmica do Kueue na sua carga de trabalho mista

Nesta seção, você simula um cenário em que a carga do servidor de inferência aumenta, exigindo o escalonamento vertical. Este cenário demonstra como o Kueue prioriza o servidor de inferência de alta prioridade suspendendo e substituindo o job de ajuste refinado de baixa prioridade quando os recursos estão limitados.

  1. Execute o comando a seguir para escalonar as réplicas do servidor de inferência para duas:

    kubectl scale --replicas=2 deployment/tgi-gemma-deployment --namespace=llm
    
  2. Verifique o status dos objetos de carga de trabalho:

    kubectl get workloads --namespace=llm
    

    A saída será assim:

    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
    

    A saída mostra que o job de ajuste refinado não é mais aceito porque o aumento das réplicas do servidor de inferência está usando a cota de GPU disponível.

  3. Verifique o status do job de ajuste refinado:

    kubectl get job --namespace=llm
    

    A saída será semelhante a esta, indicando que o status do job de ajuste refinado está suspenso:

    NAME                STATUS      COMPLETIONS   DURATION   AGE
    finetune-gemma-l4   Suspended   0/2                      33m
    
  4. Execute o seguinte comando para inspecionar seus pods:

    kubectl get pod --namespace=llm
    

    A saída é semelhante à seguinte, indicando que o Kueue encerrou os pods de job de ajuste fino para liberar recursos para a implantação do servidor de inferência de maior prioridade.

    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. Em seguida, teste o cenário em que a carga do servidor de inferência diminui e os pods reduzir escala vertical. Execute este comando:

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

    Execute o comando a seguir para mostrar os objetos de carga de trabalho:

    kubectl get workloads --namespace=llm
    

    A saída será semelhante a esta, indicando que uma das implantações do servidor de inferência foi encerrada e o job de ajuste refinado foi readmitido.

    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. Execute este comando para mostrar os jobs:

    kubectl get job --namespace=llm
    

    A saída será semelhante a esta, indicando que o trabalho de ajuste refinado está sendo executado novamente, retomando do último ponto de verificação disponível.

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

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

Excluir os recursos implantados

Para evitar cobranças na sua conta do Google Cloud pelos recursos criados neste guia, execute os comandos a seguir:

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

A seguir