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

Você também pode usar o fatiamento dinâmico interagindo diretamente com Recursos personalizados de fatiamento. Para mais informações, consulte Usar fatiamento dinâmico com um escalonador personalizado.

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

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 de capacidade total é um recurso ativado pelo TPU Cluster Director.
  • Para o subfatiamento 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 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 de TPU.
  5. Configure o Kueue para criar um recurso Slice fração.
  6. Execute cargas de trabalho em fatiamento dinâmico 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 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 Kueue, JobSet e 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 comando a seguir:

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 subsegmentação dinâmica, 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 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 do JobSet necessária com base nos requisitos de topologia. Para subsegmentação dinâmica, 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) é necessário apenas 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 LWS v0.8.0 ou mais recente.

Criar pools de nós com provisionamento incremental

Esta seção descreve como criar os 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 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 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 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 do ICI ou do OCS.
.

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

Você pode 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 houver no bloco que você 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 Conferir a topologia e o status de integridade de todas as reservas no modo de capacidade:

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

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

  2. Defina o caminho da reserva:

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

    Substitua:

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

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

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

    Substitua:

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

Sub-bloco

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

    1. Para identificar o nome do bloco, liste todos os blocos de 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 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 de interconexão entre chips (ICI) ainda não estão ativos. Portanto, não é possível executar cargas de trabalho diretamente nesses pools de nós.

Para ativar todos os links do ICI necessários para formar a fração e permitir 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 lida automaticamente com a criação e exclusão de recursos personalizados do tipo "Slice". Evite modificar manualmente os recursos personalizados de Slice criados pelo Kueue.

Criar uma fração dinâmica 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 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 a qual bloco cada sub-bloco pertence, 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 ResourceFlavor para sub-blocos do Ironwood (TPU7x) inclui o rótulo cloud.google.com/gke-tpu-accelerator: tpu7x para corresponder a nós com máquinas 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 fatiamento dinâmico.
    • ClusterQueue (cq) e LocalQueue (lq): esses campos gerenciam recursos google.com/tpu. O cq ClusterQueue inclui a verificação de admissão superslice-ac. O 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 no TAS e no particionamento 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 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 a versão 1.36.0-gke.3712000 ou mais recente do GKE.

É 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 OCS. A partição ainda pode formar uma fatia, mas o desempenho geral pode ser menor em comparação com partições íntegras. Esse estado se aplica apenas à topologia de nível superior 4x4x4. 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 em fatiamento dinâmico com o Kueue

Esta seção descreve como implantar cargas de trabalho em fatiamento 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 4x12x16, que é composta de 12 sub-blocos. O número de pods foi calculado como: (4 * 12 * 16) / 4 chips por nó = 192 pods.

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

    # 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 para a 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 fração aceitas 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 admite 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 subblocos 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

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

Otimizar o agendamento de cargas de trabalho com reutilização de intervalos

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 a excluir quando a carga de trabalho for concluída. Esse comportamento pode causar rotatividade e sobrecarga na infraestrutura subjacente.

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

Ativar a reutilização de intervalos

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

  1. Configure 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 ajudar a 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. Configure o Kueue (recomendado): para permitir que o Kueue use as sugestões de afinidade de nó preferidas 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ção por carga de trabalho

Para fazer com que uma carga de trabalho específica limpe as fatias 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 ociosa. No entanto, depois que a carga de trabalho for concluída, a fração será excluída permanentemente em vez de ser retornada 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. Para isso, adicione o rótulo à lista labelKeysToCopy na configuração do Kueue:

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

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

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 usa. Como os nomes dos intervalos são gerados dinamicamente (e são separados do nome da carga de trabalho quando a reutilização está ativada), use esses rótulos para encontrar o intervalo atribuído à 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 nome do seu job ou JobSet.

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

  3. Use o nome do período para verificar o status detalhado do período dinâmico:

    kubectl describe slice SLICE_NAME

    Substitua SLICE_NAME pelo nome da sua 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 de 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.

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

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

Métricas de ciclo de vida de segmentação

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

  • tpu_ac_slice_create_time: histograma que mede quanto tempo levou para criar a fração.
  • tpu_ac_slices_created: contagem de 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 rótulos is_reusable e in_use.
  • tpu_ac_slice_delete_time: histograma que mede quanto tempo levou a exclusão de uma fração.
  • tpu_ac_slices_deleted: contador que rastreia intervalos excluídos.
  • tpu_ac_slices_gc: o rastreamento de contadores limpou automaticamente as divisões.

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

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

  • tpu_ac_slices_reused: o rastreamento do contador reutilizou frações com sucesso.
  • tpu_ac_slices_conflict: frações de acompanhamento de contador excluídas devido a incompatibilidades de topologia ou forma.

Métricas de carga de trabalho e admissão

As métricas de carga de trabalho e admissão rastreiam 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 faz com que o Kueue exclua os recursos personalizados de Slice associados.

    kubectl delete jobset JOBSET_NAME

    Substitua JOBSET_NAME pelo nome do seu 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 escalonador

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 agendador diferente, siga as informações de referência do recurso personalizado de fração.

A seguir