Programar frações dinâmicas com o Kueue e o TAS

Este documento explica como provisionar pools de nós de TPU e programar intervalos dinâmicos no Google Kubernetes Engine (GKE) usando o Kueue e o agendamento com reconhecimento de topologia (TAS).

Também é possível usar o corte dinâmico interagindo diretamente com Recursos personalizados de corte. Para mais informações, consulte Usar o corte dinâmico com um programador personalizado.

Antes de seguir estas instruções, entenda os conceitos de segmentação dinâmica.

Requisitos

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

  • Use um cluster padrão no canal rápido em uma das seguintes versões:
    • Para configuração dinâmica de supersegmentação (topologias iguais ou maiores que 4x4x4), use a versão 1.35.2-gke.1842000 ou mais recente.
    • Para configuração dinâmica de subsegmentação (topologias menores que 4x4x4), use a versão 1.36.0-gke.3712000 ou mais recente.
  • 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 de capacidade total. O modo "Toda a capacidade" é um recurso ativado pelo TPU Cluster Director.
  • Para o subdimensionamento dinâmico, verifique se os nós têm eventos de manutenção pendentes. Monitore suas instâncias para eventos de manutenção pendentes. Se algum dos seus nós tiver um evento de manutenção pendente com um horário de término entre 18 de setembro de 2026 e 30 de setembro de 2026, você precisará acionar manualmente o evento de manutenção do host nesses nós antes de usar a subsegmentação.

Antes de começar

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

  • Ative a API Google Kubernetes Engine.
  • Ativar a API Google Kubernetes Engine
  • Se você quiser usar a Google Cloud CLI para essa tarefa, instale e, em seguida, inicialize a CLI gcloud. Se você instalou a CLI gcloud anteriormente, instale a versão mais recente executando o comando gcloud components update. Talvez as versões anteriores da CLI gcloud não sejam compatíveis com a execução dos comandos neste documento.

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. Instale o Kueue, o JobSet e o LWS.
  4. Criar pools de nós da TPU.
  5. Configure o Kueue para criar um recurso Slice fração.
  6. Execute cargas de trabalho em segmentação dinâmica com o Kueue.
  7. 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 do 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 de 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
    

Instalar o Kueue, o JobSet e o LWS

Se você já instalou o Kueue, o JobSet e o LWS, pule esta seção.

Instalar o Kueue

Siga as instruções da documentação do Kueue ou execute o seguinte comando:

kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/KUEUE_VERSION/manifests.yaml

Substitua KUEUE_VERSION pela versão necessária do Kueue com base nos requisitos de topologia. Para o subsegmento dinâmico, use o Kueue v0.18.2 ou mais recente. Para o super-fatiamento dinâmico, use o Kueue v0.16.6 ou mais recente.

Instalar JobSet

Siga as instruções da documentação do JobSet ou execute o seguinte comando:

kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml

Substitua JOBSET_VERSION pela versão necessária do JobSet com base nos requisitos de topologia. Para o subsegmento dinâmico, use o JobSet v0.12.0 ou mais recente. Para o supersegmento dinâmico, use o JobSet v0.11.1 ou mais recente.

Instalar o LWS

O LeaderWorkerSet (LWS) só é necessário para o subfatiamento dinâmico.

Siga as instruções da documentação do LWS ou execute o seguinte comando:

kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/LWS_VERSION/manifests.yaml

Substitua LWS_VERSION pela versão necessária do LWS. Use o LWS v0.8.0 ou mais recente.

Criar pools de nós com provisionamento incremental

Nesta seção, descrevemos como criar pools de nós de TPU com provisionamento incremental. O GKE converte toda a capacidade da TPU em pools de nós que compreendem grupos de 16 nós de VMs Ironwood (TPU7x) ou sub-blocos. O GKE provisiona esses pools de nós mesmo quando não consegue encontrar todas as VMs íntegras, colocando nós em partes íntegras da máquina host e provisionando incrementalmente máquinas não íntegras enquanto elas são reparadas.

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

  • 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 subbloco específico ou um grupo específico de 16 nós de VMs Ironwood (TPU7x) de TPUs 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.

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

Faça com que seu pool de nós pertença a um bloco ou sub-bloco

É 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 ajuda a garantir 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 do 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. Determine quantos pools de nós criar descrevendo um bloco de reserva e identificando 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 este 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 reserva 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 Inter-Chip Interconnect (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 que as cargas de trabalho sejam programadas, 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 o controlador de fração do Kueue

  1. Para instalar o controlador de fração do Kueue, salve o seguinte manifesto como slice-controller.yaml:

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    apiVersion: v1
    kind: Namespace
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-system
    ---
    apiVersion: v1
    kind: ServiceAccount
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: Role
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-leader-election-role
      namespace: slice-controller-system
    rules:
    - apiGroups:
      - ""
      resources:
      - configmaps
      verbs:
      - get
      - list
      - watch
      - create
      - update
      - patch
      - delete
    - apiGroups:
      - coordination.k8s.io
      resources:
      - leases
      verbs:
      - get
      - list
      - watch
      - create
      - update
      - patch
      - delete
    - apiGroups:
      - ""
      resources:
      - events
      verbs:
      - create
      - patch
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: ClusterRole
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-manager-role
    rules:
    - apiGroups:
      - ""
      resources:
      - events
      verbs:
      - create
      - patch
      - update
      - watch
    - apiGroups:
      - ""
      resources:
      - nodes
      - pods
      verbs:
      - get
      - list
      - watch
    - apiGroups:
      - ""
      resources:
      - secrets
      verbs:
      - get
      - list
      - update
      - watch
    - apiGroups:
      - accelerator.gke.io
      resources:
      - slices
      verbs:
      - create
      - delete
      - get
      - list
      - patch
      - update
      - watch
    - apiGroups:
      - accelerator.gke.io
      resources:
      - slices/finalizers
      verbs:
      - update
    - apiGroups:
      - admissionregistration.k8s.io
      resources:
      - mutatingwebhookconfigurations
      verbs:
      - get
      - list
      - update
      - watch
    - apiGroups:
      - batch
      resources:
      - jobs
      verbs:
      - get
      - list
      - patch
      - update
      - watch
    - apiGroups:
      - jobset.x-k8s.io
      resources:
      - jobsets
      verbs:
      - get
      - list
      - patch
      - update
      - watch
    - apiGroups:
      - leaderworkerset.x-k8s.io
      resources:
      - leaderworkersets
      verbs:
      - get
      - list
      - patch
      - update
      - watch
    - apiGroups:
      - kueue.x-k8s.io
      resources:
      - admissionchecks
      - admissionchecks/status
      - workloads/status
      verbs:
      - get
      - patch
      - update
    - apiGroups:
      - kueue.x-k8s.io
      resources:
      - workloads
      verbs:
      - create
      - get
      - list
      - patch
      - update
      - watch
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: ClusterRole
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-metrics-auth-role
    rules:
    - apiGroups:
      - authentication.k8s.io
      resources:
      - tokenreviews
      verbs:
      - create
    - apiGroups:
      - authorization.k8s.io
      resources:
      - subjectaccessreviews
      verbs:
      - create
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: ClusterRole
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-metrics-reader
    rules:
    - nonResourceURLs:
      - /metrics
      verbs:
      - get
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: RoleBinding
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-leader-election-rolebinding
      namespace: slice-controller-system
    roleRef:
      apiGroup: rbac.authorization.k8s.io
      kind: Role
      name: slice-controller-leader-election-role
    subjects:
    - kind: ServiceAccount
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: ClusterRoleBinding
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-manager-rolebinding
    roleRef:
      apiGroup: rbac.authorization.k8s.io
      kind: ClusterRole
      name: slice-controller-manager-role
    subjects:
    - kind: ServiceAccount
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: ClusterRoleBinding
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-metrics-auth-rolebinding
    roleRef:
      apiGroup: rbac.authorization.k8s.io
      kind: ClusterRole
      name: slice-controller-metrics-auth-role
    subjects:
    - kind: ServiceAccount
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    ---
    apiVersion: v1
    kind: Secret
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-webhook-server-cert
      namespace: slice-controller-system
    ---
    apiVersion: v1
    kind: Service
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-controller-manager-metrics-service
      namespace: slice-controller-system
    spec:
      ports:
      - name: https
        port: 8443
        protocol: TCP
        targetPort: 8443
      selector:
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
    ---
    apiVersion: v1
    kind: Service
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-webhook-service
      namespace: slice-controller-system
    spec:
      ports:
      - port: 443
        protocol: TCP
        targetPort: 9443
      selector:
        control-plane: controller-manager
    ---
    apiVersion: apps/v1
    kind: Deployment
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    spec:
      replicas: 1
      selector:
        matchLabels:
          app.kubernetes.io/name: slice-controller
          control-plane: controller-manager
      template:
        metadata:
          annotations:
            kubectl.kubernetes.io/default-container: manager
          labels:
            app.kubernetes.io/name: slice-controller
            control-plane: controller-manager
        spec:
          containers:
          - args:
            - --metrics-bind-address=:8443
            - --leader-elect
            - --health-probe-bind-address=:8081
            - --zap-log-level=3
            - --feature-gates=UseRetryMechanismForSliceCreation=true
            - --activation-timeout=6m
            command:
            - /manager
            image: tpuongke/kueue-slice-controller:latest
            livenessProbe:
              httpGet:
                path: /healthz
                port: 8081
              initialDelaySeconds: 15
              periodSeconds: 20
            name: manager
            ports:
            - containerPort: 9443
              name: webhook-server
              protocol: TCP
            readinessProbe:
              httpGet:
                path: /readyz
                port: 8081
              initialDelaySeconds: 5
              periodSeconds: 10
            resources:
              limits:
                cpu: 12000m
                memory: 32Gi
              requests:
                cpu: 8000m
                memory: 16Gi
            securityContext:
              allowPrivilegeEscalation: false
              capabilities:
                drop:
                - ALL
            volumeMounts:
            - mountPath: /tmp/k8s-webhook-server/serving-certs
              name: cert
              readOnly: true
          securityContext:
            runAsNonRoot: true
            seccompProfile:
              type: RuntimeDefault
          serviceAccountName: slice-controller-controller-manager
          terminationGracePeriodSeconds: 10
          volumes:
          - name: cert
            secret:
              defaultMode: 420
              secretName: slice-controller-webhook-server-cert
    ---
    apiVersion: admissionregistration.k8s.io/v1
    kind: MutatingWebhookConfiguration
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-mutating-webhook-configuration
    webhooks:
    - admissionReviewVersions:
      - v1
      clientConfig:
        service:
          name: slice-controller-webhook-service
          namespace: slice-controller-system
          path: /mutate-batch-v1-job
      failurePolicy: Fail
      name: mjob.kb.io
      rules:
      - apiGroups:
        - batch
        apiVersions:
        - v1
        operations:
        - CREATE
        resources:
        - jobs
      sideEffects: None
    - admissionReviewVersions:
      - v1
      clientConfig:
        service:
          name: slice-controller-webhook-service
          namespace: slice-controller-system
          path: /mutate-jobset-x-k8s-io-v1alpha2-jobset
      failurePolicy: Fail
      name: mjobset.kb.io
      rules:
      - apiGroups:
        - jobset.x-k8s.io
        apiVersions:
        - v1alpha2
        operations:
        - CREATE
        resources:
        - jobsets
      sideEffects: None
    - admissionReviewVersions:
      - v1
      clientConfig:
        service:
          name: slice-controller-webhook-service
          namespace: slice-controller-system
          path: /mutate-leaderworkerset-x-k8s-io-v1-leaderworkerset
      failurePolicy: Fail
      name: mleaderworkerset.kb.io
      rules:
      - apiGroups:
        - leaderworkerset.x-k8s.io
        apiVersions:
        - v1
        operations:
        - CREATE
        resources:
        - leaderworkersets
      sideEffects: None
    
  2. Aplique o manifesto slice-controller.yaml:

    kubectl apply -f slice-controller.yaml
    
  3. Para configurar o Kueue para o particionamento dinâmico, salve o manifesto a seguir como dynamic-slice-topology.yaml:

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    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
    
  4. 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: este 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 sabor do recurso 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. A ClusterQueue cq inclui a verificação de admissão superslice-ac. O campo nominalQuota para google.com/tpu pode ser configurado de duas maneiras:
      • Cota específica: defina o campo nominalQuota para corresponder à capacidade atual para compartilhamento justo e gerenciamento de cota.
      • Cota ilimitada: defina o campo nominalQuota com um valor muito alto, como "999999", para modelar uma cota ilimitada. Para focar na TAS e no fracionamento dinâmico, essa configuração ignora a funcionalidade de gerenciamento de cota do Kueue.

Definir a seleção de integridade da partição

Além da integridade e da prontidão padrão dos nós, o GKE expõe o estado específico de cada forma de partição usando o rótulo cloud.google.com/gke-tpu-partition-[shape]-state (em que [shape] corresponde à forma do ID da partição, como 2x2x1, 2x2x2, 2x2x4, 2x4x4 ou 4x4x4). Esse rótulo permite que o GKE considere fatores que influenciam a formação de frações, como o estado dos links da TPU. A configuração dinâmica de subsegmentação (topologias menores que 4x4x4) requer o GKE versão 1.36.0-gke.3712000 ou posterior.

É possível definir o valor do rótulo de estado da partição da seguinte maneira:

  • HEALTHY: a partição está íntegra e totalmente funcional.
  • DEGRADED: a infraestrutura da partição está em um estado degradado, por exemplo, devido à degradação do link do OCS. A partição ainda pode formar uma fatia, mas a performance geral pode ser menor em comparação com partições íntegras. Esse estado se aplica apenas à topologia 4x4x4 de nível superior. Topologias menores não têm um estado degradado.
  • UNHEALTHY: a partição não está íntegra e não pode formar uma fração.
  • UNSET: o estado não está definido devido a uma inicialização sem êxito do controlador de fração do GKE.
  • INCOMPLETE: nem todos os nós na partição são provisionados.

O webhook do controlador de fração do Kueue valida se uma carga de trabalho inclui um requisito específico de integridade da partição. 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-[shape]-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 partições disponíveis sejam usadas:

    nodeAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
        nodeSelectorTerms:
        - matchExpressions:
          - key: cloud.google.com/gke-tpu-partition-4x4x4-state
            operator: In
            values:
            - "HEALTHY"
            - "DEGRADED"
    

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 no particionamento dinâmico com o Kueue

Nesta seção, descrevemos como implantar cargas de trabalho em um corte dinâmico com o Kueue e o TAS. Ele inclui 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 por 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:

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: big-super-slice
      labels:
        kueue.x-k8s.io/queue-name: lq
    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
                  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:
      • Para subsegmentação dinâmica:é possível especificar topologias menores que 4x4x4, como 2x2x1, 2x2x2, 2x2x4 ou 2x4x4. Essas topologias menores exigem o GKE versão 1.36.0-gke.3712000 ou posterior.
      • Para supersegmentação dinâmica:é possível especificar topologias iguais ou maiores que 4x4x4. Para a configuração dinâmica de supersegmentação, cada dimensão da topologia solicitada precisa ser um múltiplo de quatro, por exemplo, 4A x 4B x 4C.
      • A topologia precisa ser uma string tridimensional no formato AxBxC, por exemplo, 4x8x8.
      • 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.
    • O webhook injeta a afinidade de nó padrão para garantir que os nós HEALTHY e DEGRADED sejam usados.
  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 segmentando apenas nós HEALTHY.

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

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: two-super-slices
      labels:
        kueue.x-k8s.io/queue-name: lq
    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 o campo replicas como 2 na seção replicatedJobs. Depois de aplicar o manifesto, o Kueue tentará formar duas fatias separadas com uma topologia 4x8x8. O Kueue cria uma fatia dinâmica para cada réplica definida em jobset.spec.replicatedJobs[].replicas. Se n réplicas forem especificadas, 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.

Monitorar a fração

É possível conferir o status da fração e monitorar as métricas dela com as métricas do sistema do GKE.

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 49 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).

Monitorar as métricas da fração

É possível monitorar as seguintes métricas do sistema do GKE que expõem a condição de uma fração:

  • kubernetes.io/accelerator/slice/state
  • kubernetes.io/accelerator/partition/state
  • kubernetes.io/accelerator/slice/deformation_durations
  • kubernetes.io/accelerator/slice/formation_durations

Para mais informações sobre as métricas, consulte Métricas do sistema do GKE.

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 a divisão dinâmica 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