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 Standard na versão 1.35.2-gke.1842000 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 do modo de capacidade total. 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:

  • 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. 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 pools de nós de 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 subblocos. O GKE provisiona esses pools de nós mesmo quando não consegue encontrar todas as 16 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 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.

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 neste 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 sub-bloco 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 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.11.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.6 \
        --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. Para instalar o controlador de fração do Kueue, salve o seguinte manifesto como slice-controller.yaml:

    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
  verbs:
  - get
  - list
  - watch
- apiGroups:
  - kueue.x-k8s.io
  resources:
  - 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

  4. Aplique o manifesto slice-controller.yaml:

    kubectl apply -f slice-controller.yaml

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

  6. 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. O ClusterQueue cq 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 gerenciamento de cotas e compartilhamento justo.
      • Cota ilimitada: defina nominalQuota como 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 do sub-bloco

Além da integridade e prontidão padrão dos nós, o GKE expõe o estado específico de cada subbloco 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 do OCS. O sub-bloco ainda pode formar uma fração, mas o desempenho geral pode ser menor em comparação com sub-blocos íntegros.
  • UNHEALTHY: o sub-bloco está em estado não í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:

    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:

    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
              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.
    • 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:

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