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

Autopilot Standard

---

Pré-lançamento — flex-start

Este recurso está sujeito aos "Termos das Soluções Pré-GA" na seção Termos Gerais de Serviço dos Termos Específicos de Serviço. Os recursos pré-GA estão disponíveis "no estado em que se encontram" e podem ter suporte limitado. Para mais informações, consulte as descrições das fases de lançamento.

Este guia mostra como otimizar o provisionamento de TPU para cargas de trabalho de treinamento de média e pequena escala usando o modo de provisionamento de início flexível. Neste guia, você usa o flex-start para implantar uma carga de trabalho que consiste em um pool de nós de fração da TPU.

Este guia é destinado a engenheiros de machine learning (ML), administradores e operadores de plataforma e especialistas em dados e IA interessados em usar os recursos de orquestração de contêineres do Kubernetes para executar cargas de trabalho em lote. Para mais informações sobre papéis comuns e tarefas de exemplo que referenciamos no conteúdo do Google Cloud , consulte Funções e tarefas comuns do usuário do GKE.

Preços de início flexível

O início flexível é recomendado se a carga de trabalho exigir recursos provisionados dinamicamente conforme necessário, por até sete dias com reservas de curto prazo, sem gerenciamento complexo de cota e acesso econômico. O início flexível é alimentado pelo Programador dinâmico de cargas de trabalho e é faturado usando os preços do Programador dinâmico de cargas de trabalho:

• Desconto de até 53% para vCPUs, GPUs e TPUs.
• Você paga conforme a utilização.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

• Ativar a API Google Kubernetes Engine.
Ativar a API Google Kubernetes Engine
• Se você quiser usar a CLI do Google Cloud para essa tarefa, instale e inicialize a gcloud CLI. Se você instalou a gcloud CLI anteriormente, instale a versão mais recente executando gcloud components update.
  Observação: em instalações prévias da gcloud CLI, defina a propriedade compute/region. Se você usa principalmente clusters zonais, defina compute/zone. Ao definir um local padrão, é possível evitar erros na gcloud CLI como: One of [--zone, --region] must be supplied: Please specify location. Talvez seja necessário especificar o local em determinados comandos se o local do cluster for diferente do padrão definido.

• Verifique se você tem um cluster do Autopilot ou um cluster Standard que esteja executando a versão 1.33.0-gke.1712000 ou mais recente.
• Confira se você conhece as limitações do flex-start.
• Ao usar um cluster padrão, verifique se você mantém pelo menos um pool de nós sem o início flexível ativado para que o cluster funcione corretamente.
• Verifique se você tem cota para TPUs preemptivas nos locais dos nós.

Criar um pool de nós com início flexível

Se você usa um cluster no modo Autopilot, pule esta seção e vá para Executar uma carga de trabalho em lote.

Para criar um pool de nós com o início flexível ativado em um cluster padrão, use a CLI gcloud.

É possível criar um pool de nós de fração de TPU de host único ou vários hosts com início flexível:

1. Crie um pool de nós com flex-start:

   Host ú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:

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

   Vários hosts

   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:

   • NODE_POOL_NAME: o nome escolhido para o pool de nós.
   • CLUSTER_NAME: o nome do cluster.
   • LOCATION_NAME: a região de computação do plano de controle do cluster.
   • NODE_ZONES: a lista separada por vírgulas de uma ou mais zonas em que o GKE cria o pool de nós.
   • MACHINE_TYPE: o tipo de máquina a ser usado para nós. Por exemplo, use ct6e-standard-4t para TPU Trillium. Para saber mais sobre os tipos de máquina disponíveis, consulte Escolher a versão da TPU.
   • TPU_TOPOLOGY: a topologia física da fração de TPU. O formato da topologia depende da versão da TPU. Para saber mais sobre topologias de TPU, use a tabela em Escolher uma topologia.

   O comando anterior usa as seguintes flags obrigatórias ao criar um pool de nós com flex-start:

   • --enable-autoscaling:o flex-start provisiona apenas os recursos de computação necessários quando a carga de trabalho é executada. Você precisa definir os seguintes parâmetros:

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

     • --max-nodes definido como o número de máquinas virtuais necessárias para sua fração de TPU.

       Por exemplo, o comando de criação do pool 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 \
       

       Esse 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 escalonador automático de cluster aumenta a escala até o número de nós que sua carga de trabalho exige. Depois que a carga de trabalho for concluída, o escalonador automático de cluster vai reduzir a escala vertical para zero nós.

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

2. Verifique o status do flex-start no pool 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 pool de nós, o campo flexStart será definido como True.

Executar uma carga de trabalho em lote

Nesta seção, você cria um job que programa um nó de TPU com flex-start. Um controlador de job no Kubernetes cria um ou mais pods e garante que eles executem uma tarefa específica com sucesso.

1. No Google Cloud console, inicie uma sessão do Cloud Shell clicando em Ativar o Cloud Shell. Uma sessão é aberta no painel inferior do console Google Cloud .

2. Crie um arquivo chamado dws-flex-start.yaml:

   Host único

   Use o seguinte manifesto para o arquivo 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 hosts

   Use o seguinte manifesto para o arquivo 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:

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

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

   kubectl apply -f dws-flex-start.yaml
   

4. Verifique se os jobs estão sendo executados no mesmo nó:

   kubectl get pods
   

   O resultado será assim:

   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>
   

   Sucesso :você implantou uma carga de trabalho em lote usando flex-start.

Limpar

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

Excluir o projeto

Cuidado: excluir um projeto tem os seguintes efeitos:
• Tudo no projeto é excluído. Se você tiver usado um projeto atual para as tarefas neste documento, a exclusão dele vai resultar na exclusão de outros trabalhos feitos nele.
• Os IDs do projeto personalizados são perdidos. É possível que, ao criar o projeto, você tenha criado um ID de projeto personalizado para ser usado no futuro. Para preservar os URLs que usam o ID do projeto, como um URL appspot.com, exclua recursos específicos do projeto, em vez de excluir 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.

Excluir o recurso individual

1. Exclua os jobs:

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

2. Exclua o pool de nós:

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

3. Exclua o cluster:

   gcloud container clusters delete CLUSTER_NAME
   

A seguir

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