Execute uma carga de trabalho de pequeno lote com TPUs e o modo de aprovisionamento de início flexível

Autopilot Padrão

---

Pré-visualização — flex-start

Esta funcionalidade está sujeita aos "Termos das Ofertas de Pré-DG" na secção Termos Gerais de Serviço dos Termos Específicos do Serviço. As funcionalidades de pré-DG estão disponíveis "tal como estão" e podem ter apoio técnico limitado. Para mais informações, consulte as descrições das fases de lançamento.

Este guia mostra como otimizar o aprovisionamento de TPUs para cargas de trabalho de treino de pequena e média escala através do modo de aprovisionamento de início flexível. Neste guia, vai usar o início flexível para implementar uma carga de trabalho que consiste num conjunto de nós de uma fatia de TPU.

Este guia destina-se a engenheiros de aprendizagem automática (ML), administradores e operadores de plataformas, bem como a especialistas em dados e IA interessados em usar capacidades de orquestração de contentores do Kubernetes para executar cargas de trabalho em lote. Para mais informações sobre as funções comuns e exemplos de tarefas que referimos no Google Cloud conteúdo, consulte Funções e tarefas comuns de utilizadores do GKE.

Preços de início flexível

Recomendamos o início flexível se a sua carga de trabalho exigir recursos aprovisionados dinamicamente conforme necessário, durante um máximo de sete dias com reservas de curto prazo, sem gestão complexa de quotas e acesso rentável. O início flexível é alimentado pelo programador de carga de trabalho dinâmico e é faturado através dos preços do programador de carga de trabalho dinâmico:

• Com desconto (até 53%) para vCPUs, GPUs e TPUs.
• Paga à medida que usa.

Antes de começar

Antes de começar, certifique-se de que realizou as seguintes tarefas:

• Ative a API Google Kubernetes Engine.
Ative a API Google Kubernetes Engine
• Se quiser usar a CLI gcloud para esta tarefa, instale-a e, em seguida, inicialize-a. Se instalou anteriormente a CLI gcloud, execute gcloud components update para obter a versão mais recente.
  Nota: para instalações existentes da CLI gcloud, certifique-se de que define a compute/region propriedade. Se usar principalmente clusters zonais, defina o compute/zone. Ao definir uma localização predefinida, pode evitar erros na CLI gcloud, como os seguintes: One of [--zone, --region] must be supplied: Please specify location. Pode ter de especificar a localização em determinados comandos se a localização do seu cluster for diferente da predefinição que definiu.

• Verifique se tem um cluster do Autopilot ou um cluster padrão com a versão 1.33.0-gke.1712000 ou posterior.
• Verifique se conhece as limitações do início flexível.
• Quando usar um cluster padrão, verifique se mantém, pelo menos, um conjunto de nós sem o início flexível ativado para que o cluster funcione corretamente.
• Verifique se tem quota para TPUs preemptíveis nas localizações dos nós.

Crie um node pool com início flexível

Se usar um cluster no modo de piloto automático, ignore esta secção e aceda à secção Execute uma carga de trabalho em lote.

Para criar um node pool com o início flexível ativado num cluster padrão existente, use a CLI gcloud.

Pode criar um conjunto de nós de fatia de TPU de host único ou múltiplo com início flexível:

1. Crie um node pool com início flexível:

   Anfitrião único

   gcloud container node-pools create NODE_POOL_NAME \
       --cluster=CLUSTER_NAME \
       --location=LOCATION_NAME \
       --node-locations=NODE_ZONES \
       --machine-type=MACHINE_TYPE \
       --reservation-affinity=none \
       --enable-autoscaling \
       --flex-start \
       --num-nodes 0 \
       --min-nodes=0 \
       --max-nodes=1
   

   Substitua o seguinte:

   • NODE_POOL_NAME: o nome que escolher para o seu node pool.
   • CLUSTER_NAME: o nome do cluster.
   • LOCATION_NAME: a região de computação para o plano de controlo do cluster.
   • NODE_ZONES: a lista separada por vírgulas de uma ou mais zonas onde o GKE cria o conjunto de nós.
   • MACHINE_TYPE: o tipo de máquina a usar para os nós. Para mais informações sobre os tipos de máquinas compatíveis com a TPU, use a tabela em Escolha a versão da TPU.

   Vários anfitriões

   gcloud container node-pools create NODE_POOL_NAME \
       --cluster=CLUSTER_NAME \
       --location=LOCATION_NAME \
       --node-locations=NODE_ZONES \
       --machine-type=MACHINE_TYPE \
       --tpu-topology=TPU_TOPOLOGY \
       --flex-start \
       --enable-autoscaling \
       --num-nodes=0 \
       --max-nodes=2 \
       --reservation-affinity=none \
       --no-enable-autorepair
   

   Substitua o seguinte:

   • NODE_POOL_NAME: o nome que escolher para o seu node pool.
   • CLUSTER_NAME: o nome do cluster.
   • LOCATION_NAME: a região de computação para o plano de controlo do cluster.
   • NODE_ZONES: a lista separada por vírgulas de uma ou mais zonas onde o GKE cria o conjunto de nós.
   • MACHINE_TYPE: o tipo de máquina a usar para os nós. Por exemplo, pode usar ct6e-standard-4t para a TPU Trillium. Para saber mais sobre os tipos de máquinas disponíveis, consulte o artigo Escolha a versão da TPU.
   • TPU_TOPOLOGY: a topologia física do segmento da TPU. O formato da topologia depende da versão da TPU. Para saber mais sobre as topologias de TPUs, use a tabela em Escolha uma topologia.

   O comando anterior usa as seguintes flags obrigatórias quando cria um conjunto de nós com início flexível:

   • --enable-autoscaling: o início flexível disponibiliza apenas os recursos de computação necessários quando a sua carga de trabalho é executada. Tem de definir os seguintes parâmetros:

     • --num-nodes=0
     • --min-nodes=0

     • --max-nodes definido para o número de máquinas virtuais que a sua fatia de TPU requer.

       Por exemplo, o comando de criação do conjunto de nós pode incluir os seguintes parâmetros:

       ...
       --machine-type=ct6e-standard-4t \
       --tpu-topology=4x4 \
       --enable-autoscaling \
       --num-nodes=0 \
       --max-nodes=4 \
       

       Este comando define o campo --max-nodes como 4 porque uma topologia 4x4 consiste em 16 chips e cada VM ct6e-standard-4t tem 4 chips.

     O redimensionador automático de clusters é dimensionado até ao número de nós que a sua carga de trabalho requer. Depois de a carga de trabalho ser concluída, o escalador automático do cluster é reduzido para zero nós.

   • --reservation-affinity=none: o início flexível não usa nem requer reservas.

2. Valide o estado do início flexível no conjunto de nós:

   gcloud container node-pools describe NODE_POOL_NAME \
       --cluster CLUSTER_NAME \
       --location LOCATION_NAME \
       --format="get(config.flexStart)"
   

   Se o início flexível estiver ativado no conjunto de nós, o campo flexStart é definido como True.

Execute uma carga de trabalho em lote

Nesta secção, cria um trabalho que agenda um nó da TPU com início flexível. Um controlador de tarefas no Kubernetes cria um ou mais pods e garante que executam com êxito uma tarefa específica.

1. Na Google Cloud consola, inicie uma sessão do Cloud Shell clicando em Ativar Cloud Shell. É aberta uma sessão no painel inferior da consola Google Cloud .

2. Crie um ficheiro com o nome dws-flex-start.yaml:

   Anfitrião único

   Use o seguinte manifesto para o ficheiro dws-flex-start.yaml:

   apiVersion: batch/v1
   kind: Job
   metadata:
     name: job-1
   spec:
     template:
       spec:
         nodeSelector:
           cloud.google.com/gke-flex-start: "true"
           cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
           cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
         containers:
         - name: container-1
           image: gcr.io/k8s-staging-perf-tests/sleep:latest
           args: ["3600s"] # Sleep for 1 hour
           resources:
             requests:
                 google.com/tpu: NUM_CHIPS
             limits:
                 google.com/tpu: NUM_CHIPS
         restartPolicy: OnFailure
   

   Vários anfitriões

   Use o seguinte manifesto para o ficheiro dws-flex-start.yaml:

   apiVersion: v1
   kind: Service
   metadata:
     name: headless-svc
   spec:
     clusterIP: None
     selector:
       job-name: job-1
   ---
   apiVersion: batch/v1
   kind: Job
   metadata:
     name: job-1
   spec:
     backoffLimit: 0
     completions: 2
     parallelism: 2
     completionMode: Indexed
     template:
       spec:
         subdomain: headless-svc
         restartPolicy: Never
         nodeSelector:
             cloud.google.com/gke-flex-start: "true"
             cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
             cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
         containers:
         - name: tpu-job
           image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
           ports:
           - containerPort: 8471 # Default port using which TPU VMs communicate
           - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
           securityContext:
             privileged: true
           command:
           - bash
           - -c
           - |
             python -c 'import jax; print("TPU cores:", jax.device_count())'
           resources:
             requests:
               google.com/tpu: NUM_CHIPS
             limits:
               google.com/tpu: NUM_CHIPS
   

   Substitua o seguinte:

   • ACCELERATOR_TYPE: o tipo de acelerador de TPU que usou quando criou os conjuntos de nós. Por exemplo, tpu-v4-podslice ou tpu-v5-lite-podslice.
   • TPU_TOPOLOGY: a topologia física da fatia de TPU. Por exemplo, o valor pode ser 4x4x4 ou 2x2, consoante a versão da TPU.
   • NUM_CHIPS: o número de chips de TPU em cada VM é um, quatro ou oito. Para saber mais, consulte o artigo Versões da TPU.

3. Aplique o manifesto dws-flex-start.yaml:

   kubectl apply -f dws-flex-start.yaml
   

4. Verifique se os trabalhos estão a ser executados no mesmo nó:

   kubectl get pods
   

   O resultado é semelhante ao seguinte:

   NAME    READY   STATUS      RESTARTS   AGE   IP       NODE               NOMINATED NODE   READINESS GATES
   job-1   0/1     Completed   0          19m   10.(...) gke-flex-zonal-a2  <none>           <none>
   

   Êxito: implementou com êxito uma carga de trabalho em lote usando o início flexível.

Limpar

Para evitar incorrer em cobranças na sua Google Cloud conta pelos recursos que usou nesta página, elimine o projeto que contém os recursos ou mantenha o projeto e elimine os recursos individuais.

Elimine o projeto

Atenção: a eliminação de um projeto tem os seguintes efeitos:
• Tudo no projeto é eliminado. Se usou um projeto existente para as tarefas neste documento, quando o elimina, também elimina qualquer outro trabalho que tenha feito no projeto.
• Os IDs dos projetos personalizados são perdidos. Quando criou este projeto, pode ter criado um ID do projeto personalizado que quer usar no futuro. Para preservar os URLs que usam o ID do projeto, como um URL appspot.com, elimine os recursos selecionados no projeto em vez de eliminar o projeto inteiro.
1. In the Google Cloud console, go to the Manage resources page.

   Go to Manage resources

2. In the project list, select the project that you want to delete, and then click Delete.
3. In the dialog, type the project ID, and then click Shut down to delete the project.

Elimine o recurso individual

1. Elimine os trabalhos:

   kubectl delete job -l "job-name in (job-1,job-2)"
   

2. Elimine o node pool:

   gcloud container node-pools delete NODE_POOL_NAME \
         --location LOCATION_NAME
   

3. Elimine o cluster:

   gcloud container clusters delete CLUSTER_NAME
   

O que se segue?

• Saiba mais sobre as TPUs no GKE.
• Saiba mais acerca do aprovisionamento automático de nós.
• Saiba mais acerca das práticas recomendadas para executar cargas de trabalho em lote no GKE.