Dynamisches Aufteilen mit einem benutzerdefinierten Scheduler verwenden

In diesem Dokument wird beschrieben, wie Sie dynamisches Slicing verwenden, indem Sie direkt mit benutzerdefinierten Slice-Ressourcen interagieren. Sie können Slices erstellen, den Status von Partitionen überwachen und den Zustand von Slices prüfen.

Bevor Sie diese Anleitung durcharbeiten, sollten Sie sich mit den Konzepten des dynamischen Aufteilens vertraut machen.

Vorteile von dynamischem Slicing mit einem benutzerdefinierten Scheduler

Verwenden Sie Ihren eigenen Scheduler, um Slice-Ressourcen zu verwalten, wenn Sie komplexe Planungsanforderungen haben oder dynamisches Slicing in Ihre vorhandene Planungsinfrastruktur einbinden möchten.

Wenn Sie lieber einen Scheduler verwenden möchten, anstatt Slice-benutzerdefinierte Ressourcen direkt zu verwalten, bietet GKE die Integration mit Kueue und Topology Aware Scheduling (TAS). Weitere Informationen finden Sie unter Dynamische Slices mit Kueue und TAS planen.

Workflowübersicht

Wenn Sie dynamisches Slicing mit einem benutzerdefinierten Scheduler verwenden möchten, führen Sie die folgenden Aufgaben in diesem Dokument aus:

  1. Slice-Controller aktivieren
  2. Knotenpools mit inkrementeller Bereitstellung erstellen
  3. Benutzerdefinierte Slice-Ressourcen erstellen, die auf Ihren Arbeitslastanforderungen basieren. Wenden Sie die benutzerdefinierte Slice-Ressource auf Ihren Cluster an.
  4. Partitionsstatus und Slice-Integrität überwachen
  5. Löschen Sie das Segment, wenn Sie fertig sind.

Weitere Informationen zu den Feldern und dem Status der benutzerdefinierten Slice-Ressource finden Sie in der Referenz zur benutzerdefinierten Slice-Ressource.

Voraussetzungen

Damit Sie dynamisches Slicing in GKE verwenden können, müssen die folgenden Anforderungen erfüllt sein:

  • Verwenden Sie einen Standardcluster in Version 1.35.2-gke.1842000 oder höher im Rapid Channel.
  • Verwenden Sie die Ironwood-Version (TPU7x).
  • Verwenden Sie das Container-Optimized OS-Image für Ihre Knoten.
  • Wenn Sie die inkrementelle Bereitstellung verwenden möchten, müssen Sie Reservierungen im Modus „Alle Kapazitäten“ verwenden. „Alle“ ist ein Kapazitätsmodus, der von TPU Cluster Director aktiviert wird.

Hinweis

Führen Sie die folgenden Aufgaben aus, bevor Sie beginnen:

  • Aktivieren Sie die Google Kubernetes Engine API.
  • Google Kubernetes Engine API aktivieren
  • Wenn Sie die Google Cloud CLI für diesen Task verwenden möchten, müssen Sie die gcloud CLI installieren und dann initialisieren. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit dem Befehl gcloud components update ab. In früheren gcloud CLI-Versionen werden die Befehle in diesem Dokument möglicherweise nicht unterstützt.

Slice-Controller aktivieren

Wenn Sie dynamisches Slicing verwenden möchten, aktivieren Sie den Slice-Controller in Ihrem Cluster.

  1. Cluster aktualisieren:

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

    Ersetzen Sie Folgendes:

  2. Rufen Sie Anmeldedaten ab, damit Sie mit kubectl-Befehlen mit Ihrem Cluster kommunizieren können:

    gcloud config set container/cluster CLUSTER_NAME
    gcloud container clusters get-credentials CLUSTER_NAME \
        --location=LOCATION
    
  3. Prüfen Sie in der Ausgabe des folgenden Befehls, ob der Wert slices.accelerator.gke.io vorhanden ist:

    kubectl get crd slices.accelerator.gke.io
    

    Die Ausgabe sieht etwa so aus:

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

Knotenpools mit inkrementeller Bereitstellung erstellen

In diesem Abschnitt wird beschrieben, wie Sie die TPU-Knotenpools mit inkrementeller Bereitstellung erstellen. GKE konvertiert Ihre gesamte TPU-Kapazität in Knotenpools mit einer Gruppe von 16 TPU-VMs oder Unterblöcken. GKE stellt diese Knotenpools auch dann bereit, wenn nicht alle 16 fehlerfreien VMs gefunden werden können. Dazu werden Knoten auf fehlerfreien Teilen des Hostcomputers platziert und fehlerhafte Maschinen inkrementell bereitgestellt, während sie repariert werden.

Sie können Ihren Knotenpool auf eine der folgenden Optionen ausrichten:

  • Ein bestimmter TPU-Block, der in Reservierungen im Modus „Alle Kapazitäten“ verfügbar ist. Durch die Blockausrichtung kann GKE den Knotenpool in einem beliebigen verfügbaren Unterblock innerhalb des angegebenen Blocks erstellen.
  • Ein bestimmter Unterblock oder eine bestimmte Gruppe von 16 TPU-VMs für eine detailliertere Steuerung.

Arbeitslastrichtlinie erstellen

Wenn Sie einen TPU-Slice-Knotenpool mit Ironwood (TPU7x) erstellen möchten, müssen Sie zuerst eine Arbeitslastrichtlinie erstellen, bei der das Feld accelerator-topology-mode auf provision_only festgelegt ist. Diese Einstellung löst den inkrementellen Bereitstellungsprozess aus.

Arbeitslastrichtlinie erstellen:

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

Ersetzen Sie Folgendes:

  • WORKLOAD_POLICY_NAME: ein Name für Ihre Arbeitslastrichtlinie.
  • PROJECT_ID: Projekt-ID in Google Cloud .
  • REGION: die Region für die Arbeitslastrichtlinie.

Führen Sie in diesem Befehl folgende Schritte aus:

  • Legen Sie das Feld accelerator-topology immer auf 4x4x4 fest, damit es der Gesamtzahl der Chips in einem einzelnen Unterblock entspricht.
  • Legen Sie das Feld accelerator-topology-mode immer auf provision_only fest, damit der inkrementelle Bereitstellungsprozess ausgelöst wird. Wenn das Feld provision_only festgelegt ist, werden im Knotenpool TPU-Knoten bereitgestellt, ohne ICI- oder OCS-Verbindungen zu erstellen.
können Sie die Liste der Arbeitslastrichtlinien aufrufen.

Knotenpool auf einen Block oder Unterblock ausrichten

Sie können bestimmte Unterblöcke oder Blöcke in Ihrer Reservierung im Modus „Gesamte Kapazität“ targetieren.

  • Auf einen Block ausrichten:Jeder Knotenpool verwendet Kapazität aus einem bestimmten Block. GKE platziert den Knotenpool in einem verfügbaren Unterblock in diesem Block. Sie müssen so viele Knotenpools erstellen, wie es Unterblöcke im Block gibt, den Sie verwenden möchten.
  • Auf einen Unterblock ausrichten:Jeder Knotenpool wird einem bestimmten verfügbaren Unterblock zugeordnet. Wenn Sie die Ausrichtung auf Unterblöcke verwenden, erstellt GKE den Knotenpool, wenn mindestens eine VM fehlerfrei ist. Die inkrementelle Bereitstellung trägt dazu bei, dass alle Knoten innerhalb des angegebenen Unterblocks platziert werden.

Blockieren

  1. Wenn Sie den Namen des Blocks in einer Reservierung und die Anzahl der verfügbaren Unterblöcke im Block abrufen möchten, führen Sie die folgenden Schritte im Dokument Topologie und Status von Reservierungen im Modus „Alle Kapazitäten“ ansehen aus:

    1. Ermitteln Sie den Namen des Blocks, indem Sie alle Reservierungsblöcke auflisten und den Wert im Feld name: kopieren. Dieser Wert ist der Name des Blocks oder BLOCK_NAME in diesem Dokument.

    2. Ermitteln Sie, wie viele Knotenpools erstellt werden sollen, indem Sie einen Reservierungsblock beschreiben und den Wert im Feld reservationSubBlockCount ermitteln. Dieser Wert gibt die Anzahl der verfügbaren Unterblöcke an. Der Wert reservationSubBlockCount: 4 gibt beispielsweise an, dass für den Block vier Unterblöcke verfügbar sind und Sie vier separate Knotenpools erstellen müssen.

  2. Reservierungspfad festlegen:

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

    Ersetzen Sie Folgendes:

    • RESERVATION_NAME: Der Name Ihrer TPU-Reservierung.
    • BLOCK_NAME: der Name des Blocks.
  3. Erstellen Sie für jeden im vorherigen Schritt ermittelten Unterblock einen Knotenpool. Wenn die Anzahl beispielsweise 4 ist, führen Sie diesen Befehl viermal aus. Verwenden Sie für jeden Knotenpool einen eindeutigen Namen.

    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}
    

    Ersetzen Sie Folgendes:

    • NODE_POOL_NAME ist der Name des neuen Knotenpools.
    • CLUSTER_NAME: der Name Ihres GKE-Clusters.
    • WORKLOAD_POLICY_NAME: der Name der von Ihnen erstellten Workload-Richtlinie.
    • ZONE: die Zone für den Knotenpool, z. B. us-central1-a.

Unterblock

  1. Führen Sie die folgenden Schritte im Dokument Topologie und Integritätsstatus von Reservierungen im Modus „Alle Kapazitäten“ ansehen aus, um den Namen des Blocks und die IDs der verfügbaren Unterblöcke abzurufen:

    1. Wenn Sie den Namen des Blocks ermitteln möchten, listen Sie alle Reservierungsblöcke auf und kopieren Sie den Wert im Feld name:. Dieser Wert ist der Name des Blocks oder der BLOCK_NAME in diesem Dokument.

    2. Um den Namen der Unterblöcke zu ermitteln, listen Sie alle Unterblöcke eines Blocks auf und kopieren Sie den Wert im Feld name: für jeden Eintrag unter reservationSubBlocks. Dieser Wert ist der Name des Unterblocks oder SUBBLOCK_NAME in diesem Dokument.

  2. Reservierungspfad festlegen:

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

    Ersetzen Sie Folgendes:

    • RESERVATION_NAME: Der Name Ihrer TPU-Reservierung.
    • BLOCK_NAME: der Name des Blocks.
    • SUBBLOCK_NAME: der Name des Unterblocks.
  3. Knotenpool erstellen:

    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}
    

    Ersetzen Sie Folgendes:

    • NODE_POOL_NAME: ein eindeutiger Name für Ihren neuen Knotenpool, z. B. sub-block-pool-1.
    • PROJECT_ID: Projekt-ID in Google Cloud .
    • CLUSTER_NAME: Name Ihres GKE-Clusters.
    • ZONE: Die Zone für den Knotenpool, z. B. us-central2-b.
    • WORKLOAD_POLICY_NAME: der Name der von Ihnen erstellten Arbeitslastrichtlinie.

In dieser Phase werden die Knoten erstellt, ihre Inter-Chip Interconnect-Verbindungen (ICI) sind jedoch noch nicht aktiv. Daher können Sie Arbeitslasten nicht direkt in diesen Knotenpools ausführen.

Damit alle erforderlichen ICI-Links zum Erstellen des Slices und zum Planen von Arbeitslasten verwendet werden können, erstellen Sie einen dynamischen Slice mit einer der folgenden Methoden:

  • Erstellen Sie eine benutzerdefinierte Slice-Ressource. Anstelle von Pods verwenden Sie eine benutzerdefinierte Slice-Ressource, um die angegebene Topologie zu definieren, die vom Slice-Controller aktiviert wird.
  • GKE-Arbeitslasten mit Kueue und TAS planen. Kueue übernimmt automatisch das Erstellen und Löschen von benutzerdefinierten Slice-Ressourcen. Sie sollten von Kueue erstellte benutzerdefinierte Slice-Ressourcen nicht manuell ändern.

Dynamischen Ausschnitt erstellen

Nachdem Sie die Knotenpools erstellt haben, können Sie ein größeres dynamisches Slice bilden, indem Sie eine benutzerdefinierte Slice-Ressource erstellen. Anstelle von Pods verwenden Sie eine benutzerdefinierte Slice-Ressource, um die angegebene Topologie zu definieren, die dann vom Slice-Controller aktiviert wird.

Status der Knoten und Partitionen prüfen

  1. Führen Sie den folgenden Befehl aus, um die Knotennamen aus dem Knotenpool abzurufen:

    kubectl get nodes -l cloud.google.com/gke-nodepool=${NODE_POOL_NAME}
    

    Die Ausgabe sieht in etwa so aus:

    NAME                                 STATUS   ROLES    AGE    VERSION
    gke-np-status-update-7b4c890c-0jhp   Ready    <none>   2d1h   v1.35.1-gke.1396002
    gke-np-status-update-7b4c890c-377r   Ready    <none>   2d1h   v1.35.1-gke.1396002
    gke-np-status-update-7b4c890c-gb51   Ready    <none>   2d1h   v1.35.1-gke.1396002
    
  2. Prüfen Sie das Bereitstellungsmodell des Knotens:

    kubectl describe node NODE_NAME | grep "cloud.google.com/gke-accelerator-topology-mode"
    

    Die Ausgabe sieht in etwa so aus:

    cloud.google.com/gke-accelerator-topology-mode: PROVISION_ONLY
    

    Dieser Wert entspricht der Einstellung accelerator-topology-mode=provision_only, die Sie beim Erstellen der Arbeitslastrichtlinie definiert haben.

  3. Rufen Sie die Informationen zum Knotenlabel ab:

    kubectl describe node NODE_NAME | grep "cloud.google.com/gke-tpu-partition-4x4x4-id"
    

    Ersetzen Sie NODE_NAME durch den Namen eines der Knoten im Knotenpool.

    Die Ausgabe sieht in etwa so aus:

    cloud.google.com/gke-tpu-partition-4x4x4-id=fba785f80d18552357dcdef6d3d16c27
    

    Die Annotation cloud.google.com/gke-tpu-partition-4x4x4-state gibt an, ob der Knoten für die Bildung eines dynamischen Segments verfügbar ist. Dieses Label unterstützt die folgenden Werte:

    • HEALTHY: Der Knoten ist fehlerfrei und voll funktionsfähig.
    • DEGRADED: Der Knoten ist beeinträchtigt, kann aber weiterhin für die dynamische Segmentierung verwendet werden.
    • UNHEALTHY: Der Knoten funktioniert nicht richtig und kann nicht zum Erstellen eines Slice verwendet werden.
    • UNSET: Der Status ist aufgrund unzureichender Knoten im Knotenpool nicht definiert.
    • INCOMPLETE: Nicht alle Knoten in der Partition werden bereitgestellt.
  4. Prüfen Sie, ob der Knoten die Annotation node.gke.io/created-by-mig enthält:

    kubectl describe node NODE_NAME | grep "node.gke.io/created-by-mig"
    

    Ersetzen Sie NODE_NAME durch den Namen eines der Knoten im Knotenpool.

    Die Ausgabe sieht in etwa so aus:

    node.gke.io/created-by-mig: projects/735972712744/zones/us-central1-ai1a/team/string
    

    Die Ausgabe enthält das Label node.gke.io/created-by-mig, mit dem die GKE-Steuerungsebene Kubernetes-Knoten mit den zugrunde liegenden Compute Engine-Ressourcen verknüpfen kann.

Benutzerdefinierte Slice-Ressource erstellen

  1. Definieren Sie die benutzerdefinierte Slice-Ressource:

    apiVersion: accelerator.gke.io/v1beta1
    kind: Slice
    metadata:
      # Name of the slice resource
      name: SLICE_NAME
    spec:
      # Specify the type of accelerator for this slice
      type: "tpu7x"
      # Define the desired topology for the accelerator slice
      topology: TOPOLOGY
      partitionIds:
        - PARTITION_ID # Example: a9476d1b02bd4f4e75ffffae3bd23c01
        - PARTITION_ID_2
        # ... add more partition IDs as needed
    

    Ersetzen Sie Folgendes:

    • SLICE_NAME: Ein Name für Ihr Segment. Der Name muss den metadata.name-Bedingungen entsprechen.
    • TOPOLOGY: Die Topologie für das dynamische Slice. Die Topologie muss die folgenden Bedingungen erfüllen:
      • Jede Dimension der angeforderten Topologie muss ein Vielfaches von vier sein, z. B. 4A x 4B x 4C.
      • Die drei Werte in den Topologiedimensionen, AxBxC, müssen in nicht absteigender Reihenfolge angegeben werden (A ≤ B ≤ C). Beispiel: 4x4x8 ist gültig, 4x8x4 jedoch nicht. Diese Reihenfolge sorgt für eine konsistente Segmentierung und vermeidet unerwartetes Verhalten.
      • Das Produkt der drei Werte in den Topologiedimensionen, A × B × C,darf 9.216 nicht überschreiten.
    • PARTITION_ID: Eine Liste von Strings, die die 4x4x4-Partitionen identifizieren, aus denen das Segment besteht. Berechnen Sie die Anzahl der Partitionen anhand der Gesamtzahl der Chips, wobei jede Partition aus 64 Chips besteht. Die Anzahl der Elemente in Ihrer spec.partitionIds-Liste muss genau mit der berechneten Anzahl der Partitionen ((A × B × C) / 64) übereinstimmen. Die partitionIds muss die folgenden Bedingungen erfüllen:
      • Jede Partition muss einem Reservierungsunterblock zugeordnet werden.
      • Alle zugehörigen Unterblöcke müssen zur selben Reservierung gehören.
      • Alle zugehörigen Blöcke müssen sich in derselben Reservierung befinden.
      • Die zugehörigen Knotenpools müssen alle Knoten im Status ready haben.
    • Der Wert des Felds type muss tpu7x sein.
    • Optional können Sie der benutzerdefinierten Slice-Ressource die Annotation slice.gke.io/retry-on-failure: "true" hinzufügen, damit der Slice-Controller während der Slice-Erstellung automatisch Wiederholungsversuche unternimmt. Wenn das Segment aufgrund des Statusgrunds SliceCreationFailed nicht erstellt wird, versucht der Controller es noch einmal, bis das Segment erfolgreich erstellt wurde.

    Wenn Sie beispielsweise einen 4x8x8-Slice erstellen möchten, müssen Sie vier eindeutige Partitions-IDs angeben.

    apiVersion: accelerator.gke.io/v1beta1
    kind: Slice
    metadata:
        name: test-slice-example
        annotations:
          slice.gke.io/retry-on-failure: "true" # Optional annotation to retry slice formation
    spec:
        type: "tpu7x"
        topology: "4x8x8" # (4*8*8)/64 = 4 partitions
        partitionIds:
            - "p0"
            - "p1"
            - "p2"
            - "p3"
    
  2. Wenden Sie die benutzerdefinierte Slice-Ressource an:

    kubectl apply -f test-slice-example.yaml
    

    An diesem Punkt versucht GKE, den Slice zu erstellen. Wenn eines der folgenden Probleme auftritt, schlägt die Erstellung des Slice fehl und der Statusgrund in der benutzerdefinierten Slice-Ressource wird auf SliceCreationFailed oder FAILED aktualisiert:

    • Wenn die ausgewählten Knoten in der benutzerdefinierten Ressource nicht vorhanden sind, lautet der Statusgrund SliceCreationFailed.
    • Wenn Knoten in der benutzerdefinierten Ressource von einem anderen Slice verwendet werden, lautet der Statusgrund SliceCreationFailed.
    • Wenn die Knoten in der benutzerdefinierten Ressource nicht Teil desselben Reservierungsblocks sind, lautet der Statusgrund FAILED.
    • Wenn sich die Knoten nicht in derselben Reservierung befinden, lautet der Statusgrund FAILED.
    • Wenn die Topologie nicht mit der Anzahl der Partitionen übereinstimmt, lautet der Statusgrund SliceCreationFailed.

    Weitere Informationen zum Status der benutzerdefinierten Ressource „Slice“ finden Sie unter Slice-Status.

Status der benutzerdefinierten Slice-Ressource prüfen

Führen Sie den folgenden Befehl aus, um den Status der benutzerdefinierten Slice-Ressource zu prüfen:

kubectl describe slice SLICE_NAME

Ersetzen Sie SLICE_NAME durch den Namen des Segments.

Die Ausgabe sieht etwa so aus:

Name:         test-slice
Namespace:
Labels:       <none>
Annotations:  <none>
API Version:  accelerator.gke.io/v1beta1
Kind:         Slice
Metadata:
  Creation Timestamp:  2026-01-11T23:45:15Z
  Finalizers:
    accelerator.gke.io/slice-finalizer
  Generation:        1
  Resource Version:  1768175347356335006
  UID:               d0b71e5c-be3f-4788-aead-930c7afec4f2
Spec:
  Partition Ids:
    2c79463990ff67c4e3c2648666bfedfa
    ba898ffcac0ad0946e8ff036d771ee53
    [more partition IDs]
  Topology:  8x16x16
  Type:      tpu7x
Status:
  Conditions:
    Last Transition Time:  2026-01-11T23:45:38Z
    Message:               ""
    
    Reason:                FAILED
    
    Status:                False
    Type:                  Ready
Events:

Das Feld reason im Status der benutzerdefinierten Slice-Ressource gibt den aktuellen Status des Slice an. Der Lebenszyklus einer benutzerdefinierten Slice-Ressource sieht so aus:

  • SliceNotCreated: Der Controller führt die Initialisierung und Ressourcenprüfungen durch.
    • Wenn die Voraussetzungen nicht erfüllt sind, wechselt der Status zu SliceCreationFailed.
    • Wenn die Validierung erfolgreich ist, ändert sich der Status zu ACTIVATING.
  • ACTIVATING: GKE bildet den Slice.
    • Bei Erfolg wechselt der Status zu ACTIVE.
    • Wenn untergeordnete Blöcke beeinträchtigt sind, der Slice aber verwendet werden kann, ändert sich der Status zu ACTIVE_DEGRADED.
    • Wenn die Bildung fehlschlägt, ändert sich der Status zu FAILED.
  • DEACTIVATING: Wenn die benutzerdefinierte Slice-Ressource gelöscht wird oder ein kritischer Fehler in einem aktiven oder fehlgeschlagenen Status auftritt, wird das Slice abgebaut.
  • INCOMPLETE: Der letzte Schritt, bevor die Ressource vollständig gelöscht wird.

Weitere Informationen zum Status der benutzerdefinierten Ressource „Slice“ finden Sie unter Slice-Status.

Arbeitslasten mit dynamischem Slicing ausführen

Wenn sich die benutzerdefinierte Ressource „Slice“ im Status ACTIVE befindet, können Sie Arbeitslasten darauf ausführen. Im folgenden Abschnitt finden Sie Beispiele für Arbeitslasten, bei denen dynamisches Slicing verwendet wird. Die Arbeitslasten werden als Jobs oder JobSets eingereicht.

Beispiel 1: Eine einzelne Arbeitslast verwendet einen einzelnen Slice

Das folgende Beispiel zeigt eine Arbeitslast, die einen einzelnen Subblock-Slice verwendet.

  1. Speichern Sie das folgende Beispielmanifest als tpu-job-jax-v7x-64.yaml:

    apiVersion: v1
    kind: Service
    metadata:
    name: headless-svc
    spec:
    clusterIP: None
    selector:
        job-name: tpu-job-jax-v7x-64
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
    name: tpu-job-jax-v7x-64
    spec:
    backoffLimit: 0
    completions: 16
    parallelism: 16
    completionMode: Indexed
    template:
        metadata:
        annotations:
            cloud.google.com/gke-tpu-slice-topology: 4x4x4
        spec:
        nodeSelector:
            cloud.google.com/gke-tpu-topology: 4x4x4
            cloud.google.com/gke-tpu-accelerator: tpu7x
            cloud.google.com/gke-tpu-slice: test-slice
        subdomain: headless-svc
        restartPolicy: Never
        containers:
        - name: tpu-job-jax
            env:
            - name: TPU_ACCELERATOR_TYPE
              value: tpu7x-128
            image: python:3.12
            securityContext:
            privileged: false
            command:
            - bash
            - -c
            - |
            set -ex
            pip install -U --pre jax jaxlib libtpu requests -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
            pip list
            python -c 'import jax; print("Total TPU devices (cores):", jax.device_count())'
            resources:
            requests:
                google.com/tpu: 4
            limits:
                google.com/tpu: 4
    

    In diesem Manifest:

    • cloud.google.com/gke-tpu-slice-topology und cloud.google.com/gke-tpu-topology definieren die Topologie des dynamischen Slice.
    • env.value: tpu7x-128 ist der TPU-Beschleunigertyp und die Gesamtzahl der Kerne im Slice. Die Anzahl der Kerne wird berechnet, indem die Dimensionen der Topologie mit der Anzahl der Kerne pro Chip multipliziert werden. Bei einer 4x4x4-Topologie lautet die Berechnung beispielsweise 4 × 4 × 4 × 2 = 128, wobei 2 die Anzahl der Kerne pro Chip für tpu7x (Ironwood (TPU7x)) ist. Daher ist die TPU_ACCELERATOR_TYPE tpu7x-128.
  2. Wenden Sie das tpu-job-jax-v7x-64.yaml-Manifest an:

    kubectl apply -f tpu-job-jax-v7x-64.yaml
    

Beispiel 2: Arbeitslast mit JobSet in Multislice-Knotenpools bereitstellen

In diesem Beispiel wird gezeigt, wie Sie eine Arbeitslast mit JobSet in Multislice-Knotenpools bereitstellen.

  1. JobSet installieren:

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.10.1/manifests.yaml
    
  2. Speichern Sie das folgende Beispielmanifest als tpu-multislice-jax.yaml:

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: tpu-multislice-jax
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice
    spec:
      failurePolicy:
        maxRestarts: 3
      replicatedJobs:
      - name: slice-job
        replicas: 2
        template:
          spec:
            parallelism: 16
            completions: 16
            backoffLimit: 0
            completionMode: Indexed
            template:
              metadata:
                annotations:
                  # The shape of the slice
                  cloud.google.com/gke-tpu-slice-topology: 4x4x4
              spec:
                hostNetwork: true
                dnsPolicy: ClusterFirstWithHostNet
                nodeSelector:
                  cloud.google.com/gke-tpu-topology: 4x4x4
                  cloud.google.com/gke-tpu-accelerator: tpu7x
                  # IMPORTANT: Do NOT put 'cloud.google.com/gke-tpu-slice' here manually.
                  # The exclusive-topology annotation handles the slice assignment automatically.
                containers:
                - name: jax-worker
                  image: python:3.12
                  securityContext:
                    privileged: true
                  ports:
                  - containerPort: 8471
                  command:
                  - bash
                  - -c
                  - |
                    set -ex
                    pip install -U --pre jax jaxlib libtpu requests -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                    # Verify JobSet injected the specific slice ID for this worker
                    echo "JobSet Index: $JOB_COMPLETION_INDEX"
                    python -c 'import jax; print("Total TPU devices:", jax.device_count())'
                  resources:
                    requests:
                      google.com/tpu: 4
                    limits:
                      google.com/tpu: 4
    
  3. Wenden Sie das tpu-multislice-jax.yaml-Manifest an:

    kubectl apply -f tpu-multislice-jax.yaml
    

    In diesem Manifest:

    • Das Feld replicas: 2 unter replicatedJobs gibt an, dass JobSet zwei separate Jobs erstellt, die jeweils einem 4x4x4-TPU-Slice entsprechen.
    • Die Annotation alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice sorgt dafür, dass jeder Job einem eindeutigen TPU-Slice zugewiesen wird.
    • Die cloud.google.com/gke-tpu-slice-topology: 4x4x4-Annotation definiert die Topologie jedes dynamischen Slice.
    • Die Umgebungsvariable TPU_ACCELERATOR_TYPE wird in diesem Beispiel nicht explizit festgelegt, da die Zuweisung von Slices von JobSet übernommen wird. Im JAX-Code werden die verfügbaren TPU-Geräte im zugewiesenen Slice automatisch erkannt.

Segment löschen

  1. Löschen Sie das Segment:

    kubectl patch slice $SLICE_NAME --type json \
      -p='[{"op": "remove", "path": "/metadata/finalizers"}]'
    
  2. Prüfen Sie, ob das Segment gelöscht wurde:

    kubectl get slices
    

Slice Controller deaktivieren

Wenn Sie den Slice-Controller deaktivieren möchten, entfernen Sie ihn aus dem Cluster.

  1. Prüfen Sie, ob die benutzerdefinierten Slice-Ressourcen leer sind:

    kubectl get slice -A
    
  2. Aktualisieren Sie den Cluster, um den Slice-Controller zu deaktivieren:

    gcloud container clusters update ${CLUSTER_NAME} \
        --location=${REGION} \
        --no-enable-slice-controller
    
  3. Löschen Sie die benutzerdefinierte Slice-Ressource :

    kubectl delete crd slices.accelerator.gke.io
    
  4. Prüfen Sie, ob die benutzerdefinierte Ressource „Slice“ gelöscht wurde:

    kubectl get crd | grep slices.accelerator.gke.io
    
  5. Entfernen Sie die vom Slice-Controller hinzugefügten Labels. Diese Labels müssen entfernt werden:

    • cloud.google.com/gke-tpu-slice
    • cloud.google.com/gke-tpu-topology
    1. Wenn Sie ein Element aus einem bestimmten Knoten entfernen möchten, aktualisieren Sie den Knotennamen.
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
    kubectl label node $NODE_NAME cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    
    1. So entfernen Sie diese Labels von jedem Knoten in Ihrem Cluster:
    kubectl label nodes --all cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    
    1. Prüfen Sie die Knotenlabels und bestätigen Sie, dass sie leer sind:
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
    kubectl describe node $NODE_NAME | grep "cloud.google.com/gke-tpu-slice"
    

Nächste Schritte