Configurar TPUs com o particionamento dinâmico do GKE

Neste documento, explicamos como usar o fragmentação dinâmica no Google Kubernetes Engine (GKE). O particionamento dinâmico permite configurar sub-blocos de TPU provisionada em diferentes topologias. Esse recurso reduz a necessidade de recriar pools de nós, aumenta a tolerância a falhas permitindo a recuperação automática quando ocorre uma falha e otimiza a utilização de recursos.

Este documento é destinado a engenheiros de IA/ML e administradores de plataforma que querem otimizar a utilização de TPUs, reduzir o tempo de provisionamento e melhorar a tolerância a falhas para cargas de trabalho de treinamento e inferência em grande escala.

Antes de ler este documento, confira se você conhece os seguintes tópicos:

O que é segmentação dinâmica?

O particionamento dinâmico oferece flexibilidade no gerenciamento da capacidade do Cloud TPU, permitindo desacoplar o provisionamento da TPU. O processo de segmentação dinâmica envolve o seguinte:

  1. Você provisiona recursos como unidades menores chamadas de sub-blocos. Um subbloco é a unidade de criação lógica fundamental da capacidade do Ironwood (TPU7x). Para Ironwood (TPU7x), ele representa um grupo de 16 nós de VMs de TPU com uma topologia 4x4x4 de chips de TPU interconectados. No contexto do modo de capacidade total da TPU e do particionamento dinâmico, um pool de nós é mapeado diretamente para um subbloco.
  2. Em seguida, a segmentação dinâmica junta esses sub-blocos em segmentos maiores.

Benefícios da segmentação dinâmica

A segmentação dinâmica ajuda você a:

  • Reduzir o tempo de provisionamento: o provisionamento individual de sub-blocos leva a um provisionamento geral mais rápido porque minimiza o impacto de qualquer falha única.
  • Reduzir o tempo de recuperação: se ocorrer uma falha no chip da TPU, a menor unidade de falha será um subbloco. O corte dinâmico isola sub-blocos com falhas para que as cargas de trabalho possam ser reagendadas em sub-blocos íntegros mais rapidamente do que o reprovisionamento de uma grande fatia inteira.
  • Redefinir capacidade: se você tiver requisitos de carga de trabalho diversos, não será necessário excluir e recriar pools de nós para mudanças de topologia, o que seria necessário sem o particionamento dinâmico. Em vez disso, é possível reconfigurar dinamicamente os pools de nós provisionados para corresponder aos formatos especificados.

Principais elementos da segmentação dinâmica

O corte dinâmico apresenta os seguintes conceitos principais:

  • Provisionamento incremental de pools de nós:o segmentação dinâmica usa o provisionamento incremental, que é um modelo de provisionamento tolerante a falhas de pools de nós. Esse modelo converte toda a capacidade de TPU em pools de nós de um grupo de 16 nós de VMs de TPU.
  • Controlador de fração:um controlador de recursos personalizados do Kubernetes em execução no plano de controle do GKE que gerencia o fracionamento dinâmico. O controlador de fração gerencia o ciclo de vida de um recurso personalizado de fração, que representa uma fração dinâmica. O controlador de fração processa a criação, o monitoramento contínuo e a exclusão da Slice. Quando você usa um programador, ele direciona a criação e exclusão do recurso personalizado Slice.
  • Recurso personalizado de fração: une dinamicamente esses sub-blocos com base na topologia de TPU solicitada. Esse processo depende da reconfiguração dinâmica da rede OCS para conectar os pools de nós da TPU, o que ajuda a garantir um desempenho otimizado. Para inspecionar o progresso ou a integridade da formação de intervalos dinâmicos, verifique os campos de status do recurso personalizado de Slice.

Requisitos

Para usar o particionamento dinâmico no GKE, você precisa atender aos seguintes requisitos:

  • Use um cluster Standard na versão 1.35.0-gke.274500 ou mais recente no canal rápido.
  • Use a versão Ironwood (TPU7x).
  • Use a imagem do Container-Optimized OS para seus nós.
  • Para usar o provisionamento incremental, use reservas no modo Toda a capacidade. O modo de capacidade é um recurso ativado pelo TPU Cluster Director.

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 o comando gcloud components update. Talvez as versões anteriores da gcloud CLI não sejam compatíveis com a execução dos comandos neste documento.

Limitações

  • Uma única fração precisa usar sub-blocos no mesmo bloco de TPU em uma reserva. Para usar sub-blocos em blocos de TPU, use TPU Multislices.
  • O corte dinâmico não é compatível com topologias menores que 4x4x4.

Usar o particionamento dinâmico no GKE com o Kueue

Nesta seção, descrevemos o fluxo de trabalho para usar o particionamento dinâmico no GKE.

  1. Confira a topologia e o status de integridade de todas as reservas no modo de capacidade total.
  2. Ative o controlador de fração no cluster.
  3. Criar pools de nós da TPU.
  4. Configure o Kueue para criar um recurso Slice fração.
  5. Execute cargas de trabalho em segmentação dinâmica com o Kueue.
  6. Limpeza.

Ativar o controlador de fração

Para usar o particionamento dinâmico, ative o controlador de fração no cluster.

  1. Atualize o cluster:

    gcloud container clusters update CLUSTER_NAME \
    --location=LOCATION \
    --enable-slice-controller

    Substitua:

  2. Receba as credenciais para se comunicar com o cluster usando comandos kubectl:

    gcloud config set container/cluster CLUSTER_NAME
gcloud container clusters get-credentials CLUSTER_NAME \
    --location=LOCATION

  3. Na saída do comando a seguir, verifique se o valor slices.accelerator.gke.io está presente:

    kubectl get crd slices.accelerator.gke.io

    O resultado será o seguinte:

    slices.accelerator.gke.io                2026-01-09T23:58:02Z

Criar pools de nós com provisionamento incremental

Nesta seção, descrevemos como criar os pools de nós da TPU com provisionamento incremental. O GKE converte toda a capacidade de TPU em pools de nós de um grupo de 16 nós de VMs de TPU ou sub-blocos. O GKE provisiona esses pools de nós mesmo quando não consegue encontrar todas as 16 VMs íntegras. Para isso, ele coloca nós em partes íntegras da máquina host e provisiona incrementalmente máquinas com problemas enquanto elas são reparadas.

É possível segmentar seu pool de nós para que ele pertença a qualquer um dos seguintes elementos:

  • Um bloco específico de TPUs, que é exposto em reservas no modo "Toda a capacidade". O bloqueio de segmentação permite que o GKE crie o pool de nós em qualquer sub-bloco disponível dentro do bloco especificado.
  • Um sub-bloco específico ou um grupo específico de 16 nós de VMs de TPU para um controle mais granular.

Criar uma política de carga de trabalho

Para criar um pool de nós de fração de TPU com Ironwood (TPU7x), primeiro crie uma política de carga de trabalho com o campo accelerator-topology-mode definido como provision_only. Essa configuração aciona o processo de provisionamento incremental.

Crie uma política de carga de trabalho:

gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
        --project=PROJECT_ID \
        --region=REGION  \
        --type=HIGH_THROUGHPUT \
        --accelerator-topology=4x4x4 \
        --accelerator-topology-mode=provision_only

Substitua:

  • WORKLOAD_POLICY_NAME: um nome para a política de carga de trabalho.
  • PROJECT_ID: o ID do projeto Google Cloud .
  • REGION: a região da política de carga de trabalho.

Neste comando, faça o seguinte:

  • Sempre defina o campo accelerator-topology como 4x4x4 para corresponder ao número total de chips em um único subbloco.
  • Sempre defina o campo accelerator-topology-mode como provision_only para garantir que o processo de provisionamento incremental seja acionado. Quando o campo provision_only é definido, o pool de nós provisiona nós de TPU sem formar links do ICI ou do OCS.

Segmentar seu pool de nós para que ele pertença a um bloco ou subbloco

É possível segmentar sub-blocos ou blocos específicos na sua reserva do modo "Toda a capacidade".

  • Segmentar um bloco:cada pool de nós usa a capacidade de um bloco especificado. O GKE coloca o pool de nós em um subbloco disponível nesse bloco. Você precisa criar tantos pools de nós quanto sub-blocos no bloco que quer usar.

  • Segmentar um subbloco:cada pool de nós é mapeado para um subbloco específico e disponível. Ao usar o direcionamento por sub-bloco, o GKE cria o pool de nós se pelo menos uma VM estiver íntegra. O provisionamento incremental garante que todos os nós sejam colocados no sub-bloco especificado.

Bloquear

  1. Para recuperar o nome do bloco em uma reserva e a contagem de sub-blocos disponíveis no bloco, siga estas etapas no documento Ver a topologia e o status de integridade de todas as reservas no modo de capacidade:

    1. Identifique o nome do bloco listando todos os blocos de reserva e copiando o valor no campo name:. Esse valor é o nome do bloco ou BLOCK_NAME neste documento.

    2. Para determinar quantos pools de nós criar, descreva um bloco de reserva e identifique o valor no campo reservationSubBlockCount. Esse valor é o número de sub-blocos disponíveis. Por exemplo, o valor reservationSubBlockCount: 4 indica que o bloco tem quatro sub-blocos disponíveis, e você precisa criar quatro pools de nós separados.

  2. Defina o caminho da reserva:

    export RESERVATION_PATH="projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME"

    Substitua:

    • RESERVATION_NAME: o nome da sua reserva de TPU.
    • BLOCK_NAME: o nome do bloco.

  3. Crie um pool de nós para cada sub-bloco identificado na etapa anterior. Por exemplo, se a contagem for 4, execute esse comando quatro vezes. Use um nome exclusivo para cada pool de nós.

    gcloud container node-pools create NODE_POOL_NAME \
      --cluster=CLUSTER_NAME \
      --node-locations=ZONE \
      --machine-type=tpu7x-standard-4t \
      --num-nodes=16 \
      --placement-policy=WORKLOAD_POLICY_NAME \
      --reservation-affinity=specific \
      --reservation=${RESERVATION_PATH}

    Substitua:

    • NODE_POOL_NAME: o nome do novo pool de nós.
    • CLUSTER_NAME: o nome do cluster do GKE.
    • WORKLOAD_POLICY_NAME: o nome da política de carga de trabalho que você criou.
    • ZONE: a zona do pool de nós, por exemplo, us-central1-a.

Sub-bloco

  1. Para recuperar o nome do bloco e os IDs dos sub-blocos disponíveis, siga estas etapas no documento Ver a topologia e o status de integridade de todas as reservas do modo de capacidade:

    1. Para identificar o nome do bloco, liste todos os blocos de reservas e copie o valor no campo name:. Esse valor é o nome do bloco ou BLOCK_NAME no documento.

    2. Para identificar o nome dos sub-blocos, liste todos os sub-blocos de um bloco e copie o valor no campo name: para cada entrada em reservationSubBlocks. Esse valor é o nome do subbloco ou SUBBLOCK_NAME neste documento.

  2. Defina o caminho da reserva:

    export RESERVATION_PATH="projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME/reservationSubBlocks/SUBBLOCK_NAME"

    Substitua:

    • RESERVATION_NAME: o nome da sua reserva de TPU.
    • BLOCK_NAME: o nome do bloco.
    • SUBBLOCK_NAME: o nome do sub-bloco.

  3. Crie o pool de nós:

    gcloud container node-pools create NODE_POOL_NAME \
        --project=PROJECT_ID \
        --cluster=CLUSTER_NAME \
        --node-locations=ZONE \
        --machine-type=tpu7x-standard-4t \
        --num-nodes=16 \
        --placement-policy=WORKLOAD_POLICY_NAME \
        --reservation-affinity=specific \
        --reservation=${RESERVATION_PATH}

    Substitua:

    • NODE_POOL_NAME: um nome exclusivo para o novo pool de nós, por exemplo, sub-block-pool-1.
    • PROJECT_ID: o ID do projeto Google Cloud .
    • CLUSTER_NAME: o nome do cluster do GKE.
    • ZONE: a zona do pool de nós, por exemplo, us-central2-b.
    • WORKLOAD_POLICY_NAME: o nome da política de carga de trabalho que você criou.

Nessa etapa, os nós são criados, mas os links de interconexão entre chips (ICI) ainda não estão ativos. Portanto, não é possível executar cargas de trabalho diretamente nesses pools de nós.

Para ativar todos os links do ICI necessários para formar a fração e permitir o agendamento de cargas de trabalho, crie uma fração dinâmica usando um dos seguintes métodos:

  • Crie um recurso personalizado de fração. Em vez de pods, use um recurso personalizado de Slice para definir a topologia especificada, que o controlador de fração ativa.
  • Agende cargas de trabalho do GKE com o Kueue e o TAS. O Kueue processa automaticamente a criação e exclusão de recursos personalizados Slice. Evite modificar manualmente os recursos personalizados Slice criados pelo Kueue.

Criar uma fração dinâmica com o Kueue e o TAS

Nesta seção, você programa cargas de trabalho do GKE com o Kueue e o TAS.

Instalar recursos do JobSet e do Kueue para segmentação dinâmica

  1. Instale o JobSet:

    helm install jobset oci://registry.k8s.io/jobset/charts/jobset \
        --version 0.10.1 \
        --namespace jobset-system \
        --create-namespace \
        --set controller.resources.requests.cpu=4 \
        --set controller.resources.requests.memory=16Gi

  2. Instale o Kueue:

    helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
        --version 0.16.1 \
        --namespace kueue-system \
        --create-namespace \
        --wait \
        --set controllerManager.replicas=3 \
        --set controllerManager.manager.resources.requests.cpu=16 \
        --set controllerManager.manager.resources.requests.memory=64Gi

  3. Instale o controlador de fração do Kueue:

    kubectl apply -f https://gist.githubusercontent.com/mwysokin/cd90010d0d375b3bf57c536905692547/raw/506c36dd070f4ac222ba8a5e58ba28bbfcfa8ed3/kueue-slice-controller-v0.8.0-130.yaml

  4. Para configurar o Kueue para o particionamento dinâmico, salve o manifesto a seguir como dynamic-slice-topology.yaml:

    apiVersion: kueue.x-k8s.io/v1beta1
kind: Topology
metadata:
  name: superslice-topology
spec:
  levels:
  # Label to identify the physical block a sub-block belongs to.
  # Only sub-blocks from the same block can form a slice.
  - nodeLabel: cloud.google.com/gce-topology-block
  # Label to identify individual TPU sub-blocks (4x4x4 topology).
  - nodeLabel: cloud.google.com/gke-tpu-partition-4x4x4-id
  # Standard Kubernetes label for individual nodes.
  # Required to assign Pods to specific VMs.
  - nodeLabel: kubernetes.io/hostname
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: superslice-rf
spec:
  nodeLabels:
    cloud.google.com/gke-tpu-accelerator: tpu7x
  topologyName: superslice-topology
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: AdmissionCheck
metadata:
  name: superslice-ac
spec:
  controllerName: accelerator.gke.io/slice
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: cq
spec:
  namespaceSelector: {}
  admissionChecks:
  - superslice-ac
  resourceGroups:
  - coveredResources:
    - google.com/tpu
    flavors:
    - name: superslice-rf
      resources:
      - name: google.com/tpu
        nominalQuota: "999999"  # modeling unlimited quota
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  name: lq
  namespace: default
spec:
  clusterQueue: cq

  5. Aplique o manifesto dynamic-slice-topology.yaml:

    kubectl apply -f dynamic-slice-topology.yaml

    Neste manifesto, você configura o Kueue para o particionamento dinâmico definindo os seguintes recursos:

    • Topologia de fração dinâmica do Ironwood (TPU7x) (superslice-topology): a topologia define os níveis que o Kueue considera ao programar cargas de trabalho de fração dinâmica. Estes são os níveis:
      • Rótulo cloud.google.com/gce-topology-block: esse nível é necessário para entender quais sub-blocos pertencem a quais blocos, porque apenas sub-blocos do mesmo bloco podem formar uma fatia.
      • Rótulo cloud.google.com/gke-tpu-partition-4x4x4-id: esse nível representa sub-blocos individuais do Ironwood (TPU7x) (topologia 4x4x4).
      • Rótulo kubernetes.io/hostname: esse nível é necessário para atribuir pods a VMs específicas e observar os rótulos e taints delas.
    • ResourceFlavor SuperSlice do Ironwood (TPU7x) (superslice-rf): o ResourceFlavor para sub-blocos do Ironwood (TPU7x) inclui o rótulo cloud.google.com/gke-tpu-accelerator: tpu7x para corresponder a nós com máquinas do Ironwood (TPU7x).
    • SuperSlice AdmissionCheck (superslice-ac): essa verificação de admissão informa ao Kueue para não programar uma carga de trabalho até que o controlador de frações do GKE confirme que a fração foi ativada. A verificação de admissão é definida primeiro e depois adicionada ao ClusterQueue que processa cargas de trabalho de segmentação dinâmica.
    • ClusterQueue (cq) e LocalQueue (lq): esses campos gerenciam recursos google.com/tpu. O cq ClusterQueue inclui a verificação de admissão superslice-ac. O nominalQuota para google.com/tpu pode ser configurado de duas maneiras:
      • Cota específica: defina nominalQuota para corresponder à capacidade atual para compartilhamento justo e gerenciamento de cotas.
      • Cota ilimitada: defina nominalQuota como um valor muito alto, como "999999", para modelar uma cota ilimitada. Para focar na TAS e no particionamento dinâmico, essa configuração ignora a funcionalidade de gerenciamento de cotas do Kueue.

Definir a seleção de integridade do sub-bloco

Além da integridade e prontidão padrão do nó, o GKE expõe o estado específico de cada sub-bloco usando o rótulo cloud.google.com/gke-tpu-partition-4x4x4-state. Esse rótulo permite que o GKE considere fatores que influenciam a formação de frações, como o estado dos links da TPU.

É possível definir o valor do rótulo cloud.google.com/gke-tpu-partition-4x4x4-state da seguinte maneira:

  • HEALTHY: a infraestrutura está íntegra.
  • DEGRADED: a infraestrutura do sub-bloco está em um estado degradado, por exemplo, devido à degradação do link OCS. O sub-bloco ainda pode formar uma fração, mas o desempenho geral pode ser menor em comparação com sub-blocos íntegros. Se você puder tolerar uma possível degradação de desempenho, configure sua carga de trabalho para usar sub-blocos DEGRADED usando a afinidade de nó, conforme mostrado no Exemplo 3.
  • UNHEALTHY: o sub-bloco não está íntegro e não pode formar uma fração.

O webhook do controlador de fração do Kueue valida se uma carga de trabalho inclui um requisito de integridade de sub-bloco específico. Se nenhuma preferência for indicada, o webhook vai injetar uma afinidade de nó padrão.

O comportamento é o seguinte:

  • Se um nodeSelector ou nodeAffinity que segmenta o rótulo cloud.google.com/gke-tpu-partition-4x4x4-state estiver presente, ele vai permanecer inalterado.

  • Se não houver uma configuração de rótulo, o webhook vai injetar a seguinte afinidade de nó padrão para garantir que apenas sub-blocos disponíveis sejam usados:

    nodeSelector:
  cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"

A seção a seguir inclui exemplos em que o rótulo cloud.google.com/gke-tpu-partition-4x4x4-state é configurado para especificar as diferentes configurações de integridade do subbloco.

Executar cargas de trabalho de teste em segmentação dinâmica com o Kueue

Nesta seção, descrevemos como implantar cargas de trabalho em segmentação dinâmica com o Kueue e o TAS. Ele inclui três exemplos que mostram como criar uma carga de trabalho de fração dinâmica e uma carga de trabalho que consiste em várias frações. As cargas de trabalho são enviadas como JobSets.

Exemplo 1: uma única carga de trabalho usa uma única fração dinâmica

O exemplo a seguir descreve como criar uma carga de trabalho usando uma fração com uma topologia de 4x12x16, que é composta de 12 sub-blocos. O número de pods foi calculado como: (4 * 12 * 16) / 4 chips por nó = 192 pods.

  1. Salve o seguinte manifesto como big-super-slice.yaml:

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: big-super-slice
  labels:
    kueue.x-k8s.io/queue-name: lq
  annotations:
spec:
  replicatedJobs:
    - name: job-jax
      replicas: 1
      template:
        spec:
          parallelism: 192  # pods per slice calculation: 4*12*16 / 4 = 192
          completions: 192
          backoffLimit: 10
          template:
            metadata:
              annotations:
                cloud.google.com/gke-tpu-slice-topology: 4x12x16
            spec:
              tolerations:
                - key: "google.com/tpu"
                  operator: "Equal"
                  value: "present"
                  effect: "NoSchedule"
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu7x
                cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"
              containers:
                - name: jax
                  image: python:latest
                  command:
                    - bash
                    - -c
                    - |
                      printenv
                      pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                      python -c 'import jax; print("Global device count:", jax.device_count(), "Local device count:", jax.local_device_count())'
                  resources:
                    limits:
                      google.com/tpu: 4
              restartPolicy: Never

    Neste manifesto, as seguintes anotações informam ao Kueue as características e a topologia da fração para configurar o seguinte:

    • cloud.google.com/gke-tpu-slice-topology: especifica "4x12x16" como a topologia de fração dinâmica. Os requisitos da topologia do acelerador tpu7x incluem as seguintes regras:
      • A topologia mínima é 4x4x4.
      • A topologia precisa ser uma string tridimensional no formato AxBxC. Por exemplo, 4x8x8.
      • Cada dimensão (A, B e C) precisa ser um múltiplo de quatro.
      • As dimensões precisam ser classificadas em ordem não decrescente: A <= B <= C. Por exemplo, 4x8x4 é inválido. O correto é 4x4x8.
      • O produto das dimensões (ABC) não pode exceder 9.216.
      • As maiores topologias de segmentação compatíveis podem incluir até 32 sub-blocos. Por exemplo, 8x16x16 com 32 sub-blocos, 8x12x20 com 30 sub-blocos ou 12x12x12 com 27 sub-blocos estão dentro dos limites aceitos.
    • cloud.google.com/gke-tpu-accelerator: tpu7x: programa pods em VMs que executam o Ironwood (TPU7x).
    • kueue.x-k8s.io/queue-name: atribui o JobSet a uma LocalQueue do Kueue.

  2. Aplique o manifesto big-super-slice.yaml:

    kubectl apply -f big-super-slice.yaml

    Depois de aplicar o manifesto, o Kueue cria um JobSet chamado big-super-slice. Em seguida, o Kueue tenta formar uma única fração dinâmica com uma topologia 4x12x16. Depois que a fração é ativada, o Kueue aceita a carga de trabalho, e os 192 pods são programados nos nós para formar a fração dinâmica que executa suas cargas de trabalho.

Exemplo 2: carga de trabalho com mais de uma réplica

O exemplo a seguir demonstra como criar uma carga de trabalho que usa duas frações dinâmicas, cada uma composta de quatro sub-blocos.

  1. Salve o seguinte manifesto como two-super-slices.yaml:

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: two-super-slices
labels:
    kueue.x-k8s.io/queue-name: lq
annotations:
spec:
replicatedJobs:
    - name: job-jax
    replicas: 2
    template:
        spec:
        parallelism: 64  # Pods per slice calculation: (4*8*8) / 4 = 64
        completions: 64
        backoffLimit: 10
        template:
            metadata:
            annotations:
                cloud.google.com/gke-tpu-slice-topology: 4x8x8
            spec:
            tolerations:
                - key: "google.com/tpu"
                operator: "Equal"
                value: "present"
                effect: "NoSchedule"
            nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu7x
                cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"
            containers:
                - name: jax
                image: python:latest
                command:
                    - bash
                    - -c
                    - |
                    printenv
                    pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                    python -c 'import jax; print("Global device count:", jax.device_count(), "Local device count:", jax.local_device_count())'
                resources:
                    limits:
                    google.com/tpu: 4
            restartPolicy: Never

  2. Aplique o manifesto two-super-slices.yaml:

    kubectl apply -f two-super-slices.yaml

Nesse manifesto, defina replicas: 2 no campo replicatedJobs. Depois de aplicar o manifesto, o Kueue tentará formar duas partes separadas com uma topologia 4x8x8. O Kueue cria uma fatia dinâmica para cada réplica definida em jobset.spec.replicatedJobs[].replicas. Se forem especificadas réplicas n, o Kueue vai criar n intervalos dinâmicos para a carga de trabalho e aguardar que todos os intervalos sejam ativados antes de admitir a carga de trabalho.

Exemplo 3: carga de trabalho com uma única fração dinâmica e NodeAffinity

A partir do Kueue 0.15, o Kueue oferece suporte a NodeAffinity para seleção de nós do TAS. Essa funcionalidade pode ser usada para permitir que nós HEALTHY e DEGRADED façam parte de uma fração dinâmica. O exemplo a seguir mostra como configurar uma carga de trabalho com uma única fração dinâmica e NodeAffinity:

  1. Salve o seguinte manifesto como slice-8x8x8-na.yaml:

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: slice-8x8x8-na
  labels:
    kueue.x-k8s.io/queue-name: lq
spec:
  replicatedJobs:
    - name: rj1
      replicas: 1
      template:
        spec:
          parallelism: 128
          completions: 128
          backoffLimit: 10
          template:
            metadata:
              annotations:
                cloud.google.com/gke-tpu-slice-topology: 8x8x8
            spec:
              tolerations:
                - key: "google.com/tpu"
                  operator: "Equal"
                  value: "present"
                  effect: "NoSchedule"
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu7x
              affinity:
                nodeAffinity:
                  requiredDuringSchedulingIgnoredDuringExecution:
                    nodeSelectorTerms:
                      - matchExpressions:
                          - key: cloud.google.com/gke-tpu-partition-4x4x4-state
                            operator: In
                            values:
                              - "HEALTHY"
                              - "DEGRADED"
              containers:
                - name: jax
                  image: python:latest
                  command:
                    - bash
                    - -c
                    - |
                      printenv
                      pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                      python -c 'import jax; print("Global device count:", jax.device_count(), "Local device count:", jax.local_device_count())'
                  resources:
                    limits:
                      google.com/tpu: 4
              restartPolicy: Never

  2. Aplique o manifesto slice-8x8x8-na.yaml:

    kubectl apply -f slice-8x8x8-na.yaml

    Depois de aplicar o manifesto, o Kueue cria um JobSet chamado slice-8x8x8-na. Em seguida, o Kueue tenta formar uma única fração dinâmica com uma topologia 8x8x8, o que permite incluir nós HEALTHY e DEGRADED devido à NodeAffinity especificada. Depois que a fração é ativada, o Kueue aceita a carga de trabalho, e os 128 pods são programados nos nós que formam a fração dinâmica.

Monitorar o status da fração

Para verificar o status das suas divisões dinâmicas, execute o seguinte comando:

kubectl describe slice SLICE_NAME

Substitua SLICE_NAME pelo nome da sua fração. O nome do slice geralmente é derivado do nome do JobSet e do índice de réplica. No Exemplo 1, uma fração criada pelo Kueue teria um nome semelhante a default-jobset-big-super-slice-yyyyy-job-jax-0.

O resultado será o seguinte:

Name:         test-slice
Namespace:
Labels:       <none>
Annotations:  <none>
API Version:  accelerator.gke.io/v1beta1
Kind:         Slice
Metadata:
  Creation Timestamp:  2026-02-12T23:44:28Z
  Finalizers:
    accelerator.gke.io/slice-finalizer
  Generation:        1
  Resource Version:  1770939905695871008
  UID:               6dbbfe14-4486-4462-864d-e078d0ca8b5b
Spec:
  Partition Ids:
    5eae6a4f59d59cf30a9bf49de618eb2b
  Topology:  4x4x4
  Type:      tpu7x
Status:
  Conditions:
    Last Transition Time:  2026-02-12T23:45:05Z
    Message:
    Reason:                ACTIVE
    Status:                True
    Type:                  Ready
    Last Transition Time:  2026-02-12T23:45:05Z
    Message:               NodeLabelingCompleted
    Reason:                NodeLabelIsAdded
    Status:                True
    Type:                  NodeLabeled
Events:                    <none>

O nome da fração segue estas regras para garantir a compatibilidade com as convenções de nomenclatura de recursos do Compute Engine:

  • Modelo:{namespace}-jobset-{jobset.metadata.name}-kueueHash[5-character]-{jobset.spec.replicatedJobs[].name}-sliceIndex.
  • Comprimento: o nome tem 54 caracteres ou menos. O controlador acrescenta um hífen e um hash de cluster de oito caracteres para criar nomes de recursos do Compute Engine, que têm um limite de 63 caracteres.
  • Formato: o nome corresponde à expressão regular ^[a-z]([-a-z0-9]*[a-z0-9])?$. O nome tem as seguintes características:
    • Começa com uma letra minúscula.
    • Contém apenas letras minúsculas, números e hífens (-).
    • Termina com uma letra minúscula ou um número (não pode terminar com um hífen).

Limpar

Para evitar cobranças inesperadas, exclua as frações antes de excluir os pools de nós.

  1. Exclua o JobSet. Essa ação faz com que o Kueue exclua os recursos personalizados de Slice associados.

    kubectl delete jobset JOBSET_NAME

    Substitua JOBSET_NAME pelo nome do JobSet, por exemplo, big-super-slice.

  2. Exclua o pool de nós da TPU:

    gcloud container node-pools delete NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --location=LOCATION

(Opcional) Usar o corte dinâmico com seu próprio programador

Este documento se concentra no uso do Kueue e do TAS. No entanto, também é possível gerenciar o particionamento dinâmico com seu próprio programador personalizado. Se você escolher usar um programador diferente, siga as informações de referência do recurso personalizado de fração.

A seguir