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 mencionamos no conteúdo do Trusted Cloud by S3NS , consulte Funções e tarefas de usuário comuns do GKE Enterprise.
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, veja se você realizou as seguintes tarefas:
- Ative 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, em seguida,
inicialize a
CLI gcloud. Se você instalou a gcloud CLI anteriormente, instale a versão
mais recente executando
gcloud components update
.
- 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
Para criar um pool de nós com o início flexível ativado em um cluster padrão, use a CLI gcloud.
Se você usar um cluster no modo Autopilot, pule esta seção e acesse a Executar uma carga de trabalho em lote.
É 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:
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, usect6e-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
como4
porque uma topologia4x4
consiste em 16 chips e cada VMct6e-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.
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 comoTrue
.
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.No Trusted Cloud console, inicie uma sessão do Cloud Shell clicando em
Ativar o Cloud Shell. Uma sessão é aberta no painel inferior do console Trusted Cloud .
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: python:3.10 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 - | pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html 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
outpu-v5-lite-podslice
.TPU_TOPOLOGY
: a topologia física da fração de TPU. Por exemplo, o valor pode ser4x4x4
ou2x2
, 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.
Aplique o manifesto
dws-flex-start.yaml
:kubectl apply -f dws-flex-start.yaml
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>
Limpar
Para evitar cobranças na sua conta do Trusted Cloud by S3NS 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
- In the Trusted Cloud console, go to the Manage resources page.
- In the project list, select the project that you want to delete, and then click Delete.
- In the dialog, type the project ID, and then click Shut down to delete the project.
Excluir o recurso individual
Exclua os jobs:
kubectl delete job -l "job-name in (job-1,job-2)"
Exclua o pool de nós:
gcloud container node-pools delete NODE_POOL_NAME \ --location LOCATION_NAME
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.