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

Neste documento, explicamos como provisionar pools de nós de TPU e programar fatiamento dinâmico no Google Kubernetes Engine (GKE) usando o Kueue e o Topology Aware Scheduling (TAS).

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

Antes de seguir estas instruções, entenda os conceitos de fatiamento dinâmico.

Requisitos

Para usar o fatiamento 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 a configuração de super-slicing dinâmico (topologias iguais ou maiores que 4x4x4), use a versão 1.35.2-gke.1842000 ou mais recente.
    • Para a configuração de subfatiamento dinâmico (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 os nós.
  • Para usar o provisionamento incremental, use as reservas do modo de capacidade total. O modo de capacidade total é um recurso ativado pelo TPU Cluster Director.
  • Para a subfração dinâmica, verifique se os nós têm eventos de manutenção pendentes. Monitore as instâncias para verificar se há eventos de manutenção pendentes. Se algum dos 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, será necessário acionar manualmente o evento de manutenção do host nesses nós antes de usar o sub-slicing.

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
  • Para usar a Google Cloud CLI nesta tarefa, instale e 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 fatiamento dinâmico no GKE com o Kueue

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

  1. Confira a topologia e o status de integridade de todas as reservas do modo de capacidade.
  2. Ative o controlador de fração no cluster.
  3. Instale o Kueue, o JobSet e o LWS.
  4. Criar pools de nós de TPU.
  5. Configure o Kueue para criar um recurso personalizado de Slice.
  6. Executar cargas de trabalho no fatiamento dinâmico com o Kueue.
  7. Limpeza.

Ativar o controlador de fatias

Para usar o fatiamento dinâmico, ative o controlador de fatias no cluster.

  1. Atualize o cluster:

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

    Substitua:

  2. Receba credenciais para se comunicar com o cluster usando os 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

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 subfatiamento dinâmico, use o Kueue v0.18.2 ou mais recente. Para o super-slicing dinâmico, use o Kueue v0.16.6 ou mais recente.

Instalar o 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 sub-slicing dinâmico, use o JobSet v0.12.0 ou mais recente. Para o super-slicing dinâmico, use o JobSet v0.11.1 ou mais recente.

Instalar o LWS

O LeaderWorkerSet (LWS) é necessário apenas para a subfração dinâmica.

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 os pools de nós de TPU com provisionamento incremental. O GKE converte toda a capacidade de TPU em pools de nós compostos por 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 de forma incremental as máquinas não íntegras enquanto elas são reparadas.

É possível segmentar o pool de nós para pertencer a qualquer um dos seguintes:

  • Um bloco específico de TPUs, que é exposto em reservas do modo de capacidade total. A segmentação de blocos 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 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 o 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 sub-bloco.
  • 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 de ICI ou OCS.

Definir o pool de nós para pertencer a um bloco ou sub-bloco

É possível segmentar sub-blocos ou blocos específicos na 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 sub-bloco disponível nesse bloco. É necessário criar tantos pools de nós quantos forem os sub-blocos no bloco que você quer usar.

  • Segmentar um sub-bloco:cada pool de nós é mapeado para um sub-bloco específico e disponível. Ao usar a segmentação de sub-blocos, 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, conclua as seguintes 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 que 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, conclua as seguintes 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.

Nesse estágio, 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 de ICI necessários para formar o fatiamento e permitir que as cargas de trabalho sejam escalonadas, crie um fatiamento dinâmico 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 é ativada pelo controlador de fração.
  • Programe cargas de trabalho do GKE com o Kueue e o TAS. O Kueue processa automaticamente a criação e a exclusão de recursos personalizados de Slice. Evite modificar manualmente os recursos personalizados de Slice criados pelo Kueue.

Criar um fatiamento dinâmico com o Kueue e o TAS

Nesta seção, você vai programar 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: Role
metadata:
  labels:
    app.kubernetes.io/managed-by: kustomize
    app.kubernetes.io/name: slice-controller
    control-plane: controller-manager
  name: slice-controller-webhook-secrets-role
  namespace: slice-controller-system
rules:
- apiGroups:
  - ""
  resources:
  - secrets
  verbs:
  - get
  - list
  - watch
  - update
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
  labels:
    control-plane: controller-manager
  name: slice-controller-manager-role
rules:
- apiGroups:
  - ""
  resources:
  - nodes
  - pods
  verbs:
  - get
  - list
  - 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:
  - accelerator.gke.io
  resources:
  - slices/status
  verbs:
  - get
  - patch
  - update
- apiGroups:
  - admissionregistration.k8s.io
  resources:
  - mutatingwebhookconfigurations
  verbs:
  - get
  - list
  - update
  - watch
- apiGroups:
  - batch
  resources:
  - jobs
  verbs:
  - get
  - list
  - patch
  - update
  - watch
- apiGroups:
  - events.k8s.io
  resources:
  - events
  verbs:
  - create
  - patch
  - update
  - watch
- apiGroups:
  - jobset.x-k8s.io
  resources:
  - jobsets
  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
- apiGroups:
  - leaderworkerset.x-k8s.io
  resources:
  - leaderworkersets
  verbs:
  - get
  - list
  - patch
  - update
  - watch
- apiGroups:
  - ray.io
  resources:
  - rayclusters
  - rayjobs
  verbs:
  - 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: RoleBinding
metadata:
  labels:
    app.kubernetes.io/managed-by: kustomize
    app.kubernetes.io/name: slice-controller
    control-plane: controller-manager
  name: slice-controller-webhook-secrets-rolebinding
  namespace: slice-controller-system
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: Role
  name: slice-controller-webhook-secrets-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=60m
        - --kube-api-qps=100
        - --kube-api-burst=200
        - --default-slice-health-node-affinity=HEALTHY_AND_DEGRADED
        - --reuse-slices=true
        - --add-preferred-node-affinity=true
        command:
        - /manager
        image: tpuongke/kueue-slice-controller:v0.10
        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: 12
            memory: 32Gi
          requests:
            cpu: 8
            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
- admissionReviewVersions:
  - v1
  clientConfig:
    service:
      name: slice-controller-webhook-service
      namespace: slice-controller-system
      path: /mutate-ray-io-v1-raycluster
  failurePolicy: Fail
  name: mraycluster.kb.io
  rules:
  - apiGroups:
    - ray.io
    apiVersions:
    - v1
    operations:
    - CREATE
    resources:
    - rayclusters
  sideEffects: None
- admissionReviewVersions:
  - v1
  clientConfig:
    service:
      name: slice-controller-webhook-service
      namespace: slice-controller-system
      path: /mutate-ray-io-v1-rayjob
  failurePolicy: Fail
  name: mrayjob.kb.io
  rules:
  - apiGroups:
    - ray.io
    apiVersions:
    - v1
    operations:
    - CREATE
    resources:
    - rayjobs
  sideEffects: None

  2. Aplique o manifesto slice-controller.yaml:

    kubectl apply -f slice-controller.yaml

  3. Para configurar o Kueue para fatiamento dinâmico, salve o seguinte manifesto 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

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

    • Topologia de fatiamento dinâmico do Ironwood (TPU7x) (superslice-topology): a topologia define os níveis que o Kueue considera ao programar cargas de trabalho de fatiamento dinâmico. Esses níveis são os seguintes:
      • 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 fração.
      • 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 deles.
    • ResourceFlavor do SuperSlice do Ironwood (TPU7x) (superslice-rf): o tipo de recurso para sub-blocos do Ironwood (TPU7x) inclui o rótulo cloud.google.com/gke-tpu-accelerator: tpu7x para corresponder 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ção do GKE confirme que a fração se tornou ativa. A verificação de admissão é definida primeiro e depois adicionada ao ClusterQueue que processa as cargas de trabalho de fatiamento dinâmico.
    • ClusterQueue (cq) e LocalQueue (lq): esses campos gerenciam recursos do 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 cotas.
      • Cota ilimitada: defina o campo nominalQuota como um valor muito alto, como "999999", para modelar uma cota ilimitada. Para se concentrar no TAS e no fatiamento dinâmico, essa configuração ignora a funcionalidade de gerenciamento de cotas do Kueue.

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

Além da integridade e prontidão padrão do nó, o GKE expõe o estado específico de cada formato de partição usando o rótulo cloud.google.com/gke-tpu-partition-[shape]-state (em que [shape] corresponde ao formato 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 de TPU. A configuração de subfatiamento dinâmico (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 fração, mas o desempenho 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 é indefinido devido à falha na inicialização 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 permanecerá inalterado.

  • Se não houver essa configuração de rótulo, o webhook vai injetar a seguinte afinidade de nó padrão para garantir que apenas as 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 sub-bloco.

Executar cargas de trabalho de teste no fatiamento dinâmico com o Kueue

Nesta seção, descrevemos como implantar cargas de trabalho no fatiamento dinâmico com o Kueue e o TAS. Ele inclui exemplos que mostram como criar uma carga de trabalho de fatiamento dinâmico e uma carga de trabalho que consiste em várias fatias. As cargas de trabalho são enviadas como JobSets.

Exemplo 1: uma única carga de trabalho usa um único fatiamento dinâmico

O exemplo a seguir descreve como criar uma carga de trabalho usando uma fração com uma topologia 4x12x16, 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 fatiamento dinâmico. Os requisitos para a topologia do acelerador tpu7x incluem as seguintes regras:
      • Para subfatiamento dinâmico:é 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 super-slicing dinâmico:é possível especificar topologias iguais ou maiores que 4x4x4. Para a configuração dinâmica de super-slicing, 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 fraçã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 um único fatiamento dinâmico com uma topologia 4x12x16. Depois que a fração estiver ativa, o Kueue vai admitir a carga de trabalho, e os 192 pods serão programados nos nós para formar a fração dinâmica que executa as 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 dois fatiamentos dinâmicos, cada um composto por quatro sub-blocos que segmentam 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, você define o campo replicas como 2 na seção replicatedJobs. Depois de aplicar o manifesto, o Kueue tenta formar duas frações separadas com uma topologia 4x8x8. O Kueue cria um fatiamento dinâmico para cada réplica definida em jobset.spec.replicatedJobs[].replicas. Se n réplicas forem especificadas, o Kueue vai criar n fatiamento dinâmico para a carga de trabalho e aguardar que todas as fatias fiquem ativas antes de admitir a carga de trabalho.

Otimizar a programação de cargas de trabalho com a reutilização de frações

Por padrão, o Kueue instrui o controlador de fração a provisionar uma nova fração de TPU para cada execução de carga de trabalho e excluí-la quando a carga de trabalho for concluída. Esse comportamento pode causar rotatividade e sobrecarga na infraestrutura subjacente.

Com a reutilização de frações, novas cargas de trabalho podem reutilizar frações de TPU ociosas que têm topologia e limites de hardware idênticos. Quando você ativa a reutilização de frações, o Kueue marca a fração como inativa no final de um job, em vez de excluí-la. Um novo job que solicita a mesma partição reivindica e reutiliza instantaneamente a fração, o que ignora o atraso no provisionamento de hardware.

Ativar a reutilização de frações

Para otimizar totalmente a reutilização de fatias, configure o controlador de Slice e o Kueue para trabalharem juntos.

  1. Configurar o controlador de fração:atualize os argumentos de implantação do controlador de fração do Kueue para incluir a flag --reuse-slices=true. Para garantir que as cargas de trabalho sejam programadas de maneira ideal nas frações atuais, adicione também a flag --add-preferred-node-affinity=true.

  2. Configurar o Kueue (recomendado): para permitir que o Kueue use as sugestões de afinidade de nó preferenciais do controlador, ative o feature gate TASRespectNodeAffinityPreferred na configuração do Kueue.

    Por exemplo, edite o ConfigMap kueue-manager-config para adicionar o feature gate:

    apiVersion: v1
kind: ConfigMap
metadata:
  name: kueue-manager-config
  namespace: kueue-system
data:
  controller_manager_config.yaml: |
    apiVersion: config.kueue.x-k8s.io/v1beta1
    kind: Configuration
    featureGates:
      TASRespectNodeAffinityPreferred: true
    # ...

    Depois de mudar o ConfigMap, reinicie a implantação kueue-controller-manager para aplicar a mudança:

    kubectl rollout restart deploy kueue-controller-manager -n kueue-system

Aplicar a exclusão de frações por carga de trabalho

Para que uma carga de trabalho específica limpe as frações e evite a reutilização delas, siga estas etapas:

  1. Aplique o seguinte rótulo à definição de Job ou JobSet:

    kueue.gke.io/tpu-enforce-slice-deletion: "true"

    Depois de aplicar esse rótulo, a carga de trabalho ainda poderá ser programada em uma fração inativa. No entanto, depois que a carga de trabalho for concluída, a fração será excluída permanentemente, em vez de ser devolvida ao pool reutilizável.

  2. Configure o Kueue para transmitir esse rótulo do job do framework para o objeto de carga de trabalho do Kueue adicionando o rótulo à lista labelKeysToCopy na configuração do Kueue:

    labelKeysToCopy:
- "kueue.gke.io/tpu-enforce-slice-deletion"

Monitorar a fração

É possível ver 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

Quando uma fração é criada ou usada ativamente por uma carga de trabalho, o controlador de fração anexa rótulos para identificar a carga de trabalho que a está usando. Como os nomes das frações são gerados dinamicamente (e são dissociados do nome da carga de trabalho quando a reutilização está ativada), use esses rótulos para encontrar a fração atribuída à sua carga de trabalho.

  1. Para encontrar as frações atribuídas a uma carga de trabalho, execute o seguinte comando:

    kubectl get slice -l kueue.gke.io/tpu-slice-used-by-namespace=NAMESPACE,kueue.gke.io/tpu-slice-used-by-workload=WORKLOAD_NAME

    Substitua:

    • NAMESPACE: o namespace da carga de trabalho.
    • WORKLOAD_NAME: o nome da carga de trabalho do Kueue, que geralmente é o mesmo do job ou do JobSet.

  2. Copie o nome da fração da saída do comando anterior.

  3. Use o nome do fatiamento para verificar o status detalhado do fatiamento dinâmico:

    kubectl describe slice SLICE_NAME

    Substitua SLICE_NAME pelo nome da fração.

O resultado será o seguinte:

Name:         s-4x4x4-74d3fb0b
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>

Monitorar as métricas da fração

É possível monitorar as seguintes métricas do sistema do GKE que mostram 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.

Métricas do controlador de fração do Kueue

Também é possível monitorar as seguintes métricas do Prometheus expostas pelo controlador de fatias do Kueue para observar o ciclo de vida, a reutilização e as verificações de admissão das fatias:

Métricas do ciclo de vida do Slice

As métricas de ciclo de vida de Slice monitoram a criação e a exclusão de frações:

  • tpu_ac_slice_create_time: histograma que mede o tempo de criação da fração.
  • tpu_ac_slices_created: contador que rastreia as frações criadas (inclui um rótulo is_reusable).
  • tpu_ac_slices_status: medidor que rastreia o número atual de frações por status, incluindo os rótulos is_reusable e in_use.
  • tpu_ac_slice_delete_time: histograma que mede o tempo que levou para excluir a fração.
  • tpu_ac_slices_deleted: contador que rastreia as frações excluídas.
  • tpu_ac_slices_gc: o rastreamento de contadores limpou as frações automaticamente.

Métricas de reutilização de frações

As métricas de reutilização de Slice monitoram a reutilização bem-sucedida e os conflitos de frações de TPU:

  • tpu_ac_slices_reused: o rastreamento de contadores reutilizou as frações com sucesso.
  • tpu_ac_slices_conflict: contador que monitora as frações excluídas devido a incompatibilidades de topologia ou formato.

Métricas de carga de trabalho e admissão

As métricas de carga de trabalho e admissão monitoram a duração e os resultados das verificações de admissão:

  • tpu_ac_admission_time: histograma que mede a duração do processo de verificação de admissão. Observação: isso mede a duração da verificação de admissão do controlador de fração e não inclui o tempo geral de admissão da carga de trabalho do Kueue.
  • tpu_ac_workload_processed: contador de decisões de verificação de admissão, contendo um rótulo target_state.

Métricas gerais

As métricas gerais rastreiam problemas e erros operacionais:

  • tpu_ac_problem: um contador que rastreia problemas ou erros operacionais. Essa métrica inclui um rótulo problem que indica o tipo de erro específico.

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 aciona o Kueue para excluir 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 fatiamento 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 fatiamento dinâmico com seu próprio programador personalizado. Se você optar por usar um programador diferente, siga as informações de referência do recurso personalizado de fração.

A seguir