KI-Training auf TPUs mit DWS und Kueue optimieren

In diesem Dokument erfahren Sie, wie Sie einen GKE-Cluster konfigurieren, um die Verfügbarkeit von KI-Training auf Tensor Processing Units (TPUs) zu maximieren. Sie konfigurieren ein automatisiertes Fallback-System mit dem Open-Source-Tool Kueue für die Jobwarteschlange und dem Dynamic Workload Scheduler (DWS) von Google Cloud.

Mit der in diesem Dokument definierten Konfiguration werden ein primärer Knotenpool und ein Sicherungsknotenpool eingerichtet:

  • On-Demand (Plan A): Dieser Knotenpool ist Ihre erste Wahl für Knoten. Das System versucht zuerst, Jobs auf diesen On-Demand-Maschinen zu planen. Der Begriff On-Demand bedeutet, dass Sie nach dem Start dieser Maschinen zuverlässigen, unterbrechungsfreien Zugriff auf sie haben.
  • DWS Flex-Start (Plan B): Dies ist Ihr Backup-Knotenpool. Wenn Plan-A-Maschinen nicht verfügbar sind, weist das Kueue-Planungsprogramm Ihren Job automatisch diesem Plan-B-Pool zu. DWS sucht dann nach der Plan-B-Hardware, garantiert aber keinen sofortigen Zugriff, da diese Hardware möglicherweise auch nicht verfügbar ist. DWS gibt jedoch nicht auf: Ihre Anfrage wird bis zu 7 Tage lang in der Warteschlange gehalten und die Maschinen werden automatisch bereitgestellt, sobald sie verfügbar sind.

So wird die Wartezeit von Jobs minimiert. Sie müssen also nicht manuell nach verfügbaren Ressourcen suchen oder Ihr Skript für verschiedene Computer neu schreiben.

Übersicht über die Konfigurationsschritte

Um das automatische Fallback-System zu konfigurieren, müssen Sie mehrere Konfigurationsschritte ausführen. Es ist hilfreich, diese Konfiguration in zwei Kategorien zu unterteilen:

  • Aufgaben des Clusteradministrators:Einmalige Infrastrukturkonfiguration, z. B. das Erstellen des GKE-Cluster, das Bereitstellen von Knotenpools und die Installation des Kueue-Planungscontrollers.
  • Aufgaben für KI-Entwickler:Wiederkehrende, tägliche Workflows wie das Definieren der Anforderungen für den Trainingsjob und das Einreichen der Arbeitslast.

Auch wenn Sie alle diese Schritte selbst ausführen, hilft es, diese Unterscheidung im Hinterkopf zu behalten, um den Gesamtprozess zu verdeutlichen.

Bevor Sie das System konfigurieren, sollten Sie sich die Konfigurationsschritte ansehen, die Sie ausführen werden.

Thema Aufgabe
Infrastruktur konfigurieren (Clusteradministrator) 1. GKE-Cluster erstellen
2. Knotenpools erstellen
3. Status von Flex-Start im Knotenpool prüfen
Kueue installieren und konfigurieren (Clusteradministrator) 1. Kueue installieren
2. Konfigurationsregeln definieren
Trainingsjob ausführen (AI Developer) 1. ConfigMap erstellen
2. RayJob-Manifest definieren
3. Arbeitslast einreichen
4. Verbindung zum RayJob herstellen
5. Logs prüfen

Wichtige Konzepte

  • On-Demand-Knotenpool (Plan A): Der primäre Knotenpool mit hoher Priorität. Bei Ihrem Job wird immer zuerst versucht, diesen Pool zu verwenden.
  • DWS Flex-Start-Knotenpool (Plan B):Der Backup-Knotenpool. Wenn Maschinen im primären Pool nicht verfügbar sind, verwendet das System automatisch diesen Pool, um nach verfügbarer Hardware zu suchen.
  • Kueue:Ein Planungsprogramm, das die Jobwarteschlange verwaltet. Sie fängt Ihre Jobanfrage ab und entscheidet, welcher Knotenpool verwendet werden soll (Plan A oder Plan B).
  • Job:Die KI-Trainingsarbeitslast, die Sie ausführen möchten. In diesem Dokument definieren Sie sie mit einem RayJob-Manifest.

Hinweis

  1. Wählen Sie in der Google Cloud Console auf der Seite für die Projektauswahl ein Google Cloud -Projekt aus oder erstellen Sie eines.

    Erforderliche Rollen zum Auswählen oder Erstellen eines Projekts

    Zur Projektauswahl

  2. Prüfen Sie, ob die Abrechnung für Ihr Google Cloud -Projekt aktiviert ist.

  3. Aktivieren Sie die Google Kubernetes Engine API und die Cloud TPU API, falls sie noch nicht aktiviert sind.

    Rollen, die zum Aktivieren von APIs erforderlich sind

    Zum Aktivieren von APIs benötigen Sie die Berechtigung serviceusage.services.enable. Wenn Sie das Projekt erstellt haben, haben Sie diese Berechtigung wahrscheinlich bereits über die Inhaberrolle (roles/owner). Andernfalls können Sie diese Berechtigung über die Rolle „Service Usage-Administrator“ (roles/serviceusage.serviceUsageAdmin) erhalten. Informationen zum Zuweisen von Rollen

    APIs aktivieren

  4. Aktivieren Sie Cloud Shell in der Google Cloud Console.

    Cloud Shell aktivieren

  5. Prüfen Sie, ob Sie für die Verwendung von TPU-Flex-Start-VMs genug auf Abruf verfügbares Kontingent haben. Wenn das Standardkontingent für Ihre Anforderungen nicht ausreicht, können Sie eine höhere Zuweisung anfordern. Weitere Informationen finden Sie unter Cloud TPU-Kontingente.

Umgebungsvariablen definieren

Um die Befehle, die Sie in diesem Dokument ausführen, zu vereinfachen, können Sie Umgebungsvariablen in Cloud Shell festlegen. In diesen Variablen werden Werte wie die ID Ihres Google Cloud -Projekts, die Namen Ihrer Knotenpools und der Standort Ihres GKE-Cluster gespeichert.

Nachdem Sie diese Variablen definiert haben, können Sie sie in mehreren Befehlen wiederverwenden, indem Sie auf den Variablennamen verweisen (z. B. $CLUSTER_NAME), anstatt die Werte jedes Mal neu einzugeben oder zu ersetzen. So ist der Prozess einfacher nachzuvollziehen und das Fehlerrisiko wird verringert.

Führen Sie die folgenden Befehle aus, um die folgenden nützlichen Umgebungsvariablen in Cloud Shell zu definieren:

export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export ZONE="us-east5-b"
export REGION="us-east5"
export CLUSTER_NAME="tpu-cluster"
export GKE_VERSION="1.34"
export ONDEMAND_NODEPOOL="on-demand-pool"
export DWS_NODEPOOL="dws-pool"

Hier eine Erklärung dieser Umgebungsvariablen:

  • PROJECT_ID: Die ID Ihres Google Cloud -Projekts.
  • PROJECT_NUMBER: Die eindeutige Kennnummer für Ihr Projekt (z. B. 123456789012).
  • ZONE: Die Computing-Zone für Ihren Cluster, z. B. us-east5-b. Wählen Sie eine Zone aus, in der der von Ihnen ausgewählte Beschleunigertyp verfügbar ist. Informationen zur Verfügbarkeit finden Sie unter Cloud TPU-Kontingente oder GPU-Kontingente.
  • REGION: Die Region, in der Sie die Clusterressourcen erstellen (z. B. us-east5).
  • CLUSTER_NAME: Der Name, den Sie für Ihren GKE-Cluster auswählen.
  • GKE_VERSION: Die GKE-Version Ihres Clusters. Verwenden Sie Version 1.34 oder höher.
  • ONDEMAND_NODEPOOL: Der Name Ihres Standard-On-Demand-Knotenpools. (Das ist Ihr Plan A-Knotenpool.)
  • DWS_NODEPOOL: Der Name des DWS-Flex-Start-Knotenpools. (Dies ist Ihr Plan-B-Knotenpool).

Infrastruktur konfigurieren (Clusteradministrator)

Als Clusteradministrator konfigurieren Sie den GKE-Cluster und die Knotenpools so, dass der Fallback-Mechanismus unterstützt wird.

GKE-Cluster erstellen

Erstellen Sie zuerst den GKE-Cluster. In diesem Cluster installieren Sie den Kueue-Controller, konfigurieren Ihre Knotenpools und führen Ihre KI-Trainingsjobs aus. So erstellen Sie einen Cluster und stellen eine Verbindung zu ihm her:

  1. Erstellen Sie den Cluster:

    gcloud container clusters create ${CLUSTER_NAME} \
  --cluster-version=${GKE_VERSION} \
  --machine-type=n2-standard-16 \
  --location=${ZONE} \
  --enable-image-streaming \
  --addons=RayOperator \
  --project=${PROJECT_ID}

    Dieser Befehl verwendet die folgenden wichtigen Flags:

    • --addons=RayOperator: Installiert den Ray Operator in Ihrem Cluster. Sie benötigen diesen Operator, um die RayJob-Arbeitslast zu verwalten, die Sie später in diesem Dokument einreichen.
    • --enable-image-streaming: Ermöglicht es Ihrem Cluster, Container-Images schneller abzurufen. Mit diesem Feature wird die Zeit, die große KI-Container-Images zum Starten benötigen, erheblich verkürzt.

  2. Rufen Sie die Anmeldedaten des Clusters ab, damit die kubectl-CLI eine Verbindung zu ihm herstellen kann. Mit diesem Befehl wird die Kubernetes-Konfigurationsdatei aktualisiert, die standardmäßig im Verzeichnis ~/.kube/config gespeichert ist:

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --location=${ZONE} \
  --project=${PROJECT_ID}

Knotenpools erstellen

Erstellen Sie die primären und Sicherungsknotenpools für Ihre Umgebung: den On-Demand-Knotenpool (Plan A) und den DWS-Flex-Start-Knotenpool (Plan B):

  1. On-Demand-Knotenpool erstellen: Dieser Pool dient als primäre Ressource für Trainingsjobs:

    gcloud container node-pools create ${ONDEMAND_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    In diesem Beispiel sind Ihre bevorzugten Maschinen TPU v6e-Beschleuniger. Sie geben diese Hardware mit dem Flag --machine-type=ct6e-standard-4t an. Sie können diesen Maschinentyp an die Hardware anpassen, z. B. GPUs oder andere TPUs, die Sie für Ihr KI-Modell verwenden möchten.

  2. DWS-Knotenpool mit Flex-Start erstellen: In diesem Beispiel wählen Sie denselben Maschinentyp (--machine-type=ct6e-standard-4t) aus, den Sie für den primären On-Demand-Pool ausgewählt haben. Für den Plan B-Knotenpool muss kein anderer Maschinentyp verwendet werden. Da Sie diese Hardware unbedingt haben möchten, bedeutet die Wahl als Plan B, dass Sie auf eine andere Methode zur Beschaffung umsteigen, wenn sie nicht sofort verfügbar ist. Bei dieser alternativen Methode wird DWS verwendet, das bis zu sieben Tage lang kontinuierlich nach der verfügbaren Hardware sucht:

    gcloud container node-pools create ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --enable-autoscaling \
  --enable-queued-provisioning \
  --flex-start \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    Für diese Befehle werden die folgenden wichtigen Flags verwendet:

    • --num-nodes=0, --min-nodes=0, --max-nodes=4 und --enable-autoscaling: Mit dieser Kombination können die Knotenpools bei Bedarf auf null Knoten skaliert und bei Inaktivität wieder heruntergefahren werden. So lassen sich Kosten sparen.
    • --tpu-topology: Definiert die physische Anordnung der TPU-Chips. Sie geben dieses Layout an, weil die physische Anordnung der Chips sich auf die Geschwindigkeit Ihres verteilten Trainingsjobs auswirkt.
    • --reservation-affinity=none: (Nur Plan A-Pool) Dieses Flag verhindert, dass der On-Demand-Knotenpool Ihre reservierte Kapazität nutzt. Der DWS-Knotenpool mit Flex-Start nutzt keine Reservierungen.
    • --enable-queued-provisioning und --flex-start: (nur Plan B-Pool) Mit diesen Flags kann DWS Knoten für Ihren Plan B-Pool aus flexibler Kapazität bereitstellen, sobald diese verfügbar ist.

Status von Flex-Start im Knotenpool prüfen

Prüfen Sie den DWS-Knotenpool mit flexiblem Start und vergewissern Sie sich, dass der flexible Start aktiviert ist:

gcloud container node-pools describe ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --format="get(config.flexStart)"

Wenn Flex Start aktiviert ist, lautet die Ausgabe True.

Kueue installieren und konfigurieren (Clusteradministrator)

In diesem Abschnitt installieren Sie den Kueue-Controller in Ihrem Cluster. Kueue ist ein Planungsprogramm, das die Jobwarteschlange verwaltet. Sie fängt Ihre Jobanfrage ab, entscheidet, welcher Knotenpool verwendet werden soll (On-Demand oder DWS Flex-Start), und weist den Job dann zu.

Kueue installieren

Führen Sie den folgenden Befehl aus, um Kueue zu installieren. Mit diesem Befehl werden die Installationsmanifeste aus dem offiziellen Repository heruntergeladen und auf Ihren Cluster angewendet:

helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
  --namespace kueue-system \
  --create-namespace \
  --set "controllerManager.featureGates[0].name=ElasticJobsViaWorkloadSlices" \
  --set "controllerManager.featureGates[0].enabled=true"

Konfigurationsregeln definieren

Erstellen Sie ein YAML-Manifest, in dem die Prioritätsregeln definiert sind. Diese Regeln weisen Kueue an, zuerst den On-Demand-Pool und dann den DWS-Flex-Start-Pool zu verwenden:

  1. Erstellen Sie eine Datei mit dem Namen dws-tpu-queue.yaml und mit folgendem Inhalt. In dieser Datei werden zwei Arten von Ressourcen (On-Demand und DWS Flex-Start) und eine Clusterwarteschlange definiert, die sie priorisiert. Diese Konfigurationsdatei definiert die Logik, die Kueue zum Verarbeiten Ihrer Jobs verwendet:

    • ResourceFlavor: Zu Beginn dieses Dokuments haben Sie zwei Knotenpools erstellt und ihnen mit den Umgebungsvariablen ${ONDEMAND_NODEPOOL} und ${DWS_NODEPOOL} Namen zugewiesen. Beim Erstellen dieser Knotenpools hat GKE automatisch jeden Knoten in diesen Pools mit dem Namen gekennzeichnet, den Sie für diese Umgebungsvariablen ausgewählt haben. Im Abschnitt ResourceFlavor wird Kueue angewiesen, nach Knoten mit diesen Labels zu suchen.
    • ClusterQueue: In diesem Abschnitt des Manifests wird die Prioritätsregel definiert. Die On-Demand-Variante wird zuerst aufgeführt, sodass Kueue zuerst versucht, On-Demand-Maschinen bereitzustellen. Wenn Kueue diese Maschinen nicht abrufen kann, versucht es stattdessen, DWS-Maschinen mit flexiblem Start bereitzustellen.
    • Quotas: In der Datei wird ein Kontingent festgelegt, das ein Limit für die Gesamtzahl der Ressourcen (z. B. CPU, Arbeitsspeicher und TPU-Chips) ist, die Ihre Jobs zu einem bestimmten Zeitpunkt im On-Demand-Knotenpool verwenden können. Wenn Ihre Jobs dieses Limit erreichen, versucht Kueue automatisch, DWS FlexStart-Maschinen (Ihre Plan B-Maschinen) bereitzustellen, die Sie in dws-tpu-queue.yaml mit einem viel höheren Kontingentlimit konfiguriert haben. 
      apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "default-cpu"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: default-pool
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "on-demand"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${ONDEMAND_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "dws"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${DWS_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  namespaceSelector: {}
  resourceGroups:
    - coveredResources: ["cpu", "memory", "google.com/tpu"]
      flavors:
        - name: "default-cpu" # Used for Ray Head Pod.
          resources:
            - name: "cpu"
              nominalQuota: 10
            - name: "memory"
              nominalQuota: 20Gi
            - name: "google.com/tpu"
              nominalQuota: 0
        - name: "on-demand" # First choice: on-demand node-pool.
          resources:
            - name: "cpu"
              nominalQuota: 40
            - name: "memory"
              nominalQuota: 75Gi
            - name: "google.com/tpu"
              nominalQuota: 16
        - name: "dws" # If on-demand is unavailable, fallback to DWS.
          resources:
            - name: "cpu"
              nominalQuota: 1000000000
            - name: "memory"
              nominalQuota: 1000000000Gi
            - name: "google.com/tpu"
              nominalQuota: 1000000000 # "Infinite" quota
  admissionChecksStrategy:
    admissionChecks:
      - name: "dws-prov"
        onFlavors: [dws]
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: "default"
  name: "user-queue"
spec:
  clusterQueue: "cluster-queue"
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: AdmissionCheck
metadata:
  name: dws-prov
spec:
  controllerName: kueue.x-k8s.io/provisioning-request
  parameters:
    apiGroup: kueue.x-k8s.io
    kind: ProvisioningRequestConfig
    name: dws-config
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ProvisioningRequestConfig
metadata:
  name: dws-config
spec:
  provisioningClassName: queued-provisioning.gke.io
  managedResources:
    - google.com/tpu

  2. Wenden Sie die Konfiguration auf Ihren Cluster an. Im folgenden Befehl wird ein Befehlszeilentool namens envsubst verwendet, um Platzhaltervariablen in der Datei dws-tpu-queue.yaml zu ersetzen. envsubst ersetzt die Platzhalter durch die Werte der Umgebungsvariablen, die Sie zuvor definiert haben:

    envsubst < dws-tpu-queue.yaml | kubectl apply -f -

Trainingsjob ausführen (AI Developer)

Als KI-Entwickler definieren und senden Sie eine Trainingsarbeitslast, indem Sie ein RayJob-Manifest erstellen. Sie geben Ihre Ressourcenanforderungen in diesem Manifest an. Das automatisierte Fallback-System, das der Clusteradministrator zuvor mit Kueue und DWS konfiguriert hat, kümmert sich um die zugrunde liegenden Knotenpools.

In diesem Abschnitt führen Sie die folgenden Schritte aus:

  • Erstellen Sie ein Python-Trainingsscript.
  • Speichern Sie das Skript in einer Kubernetes-ConfigMap.
  • Stellen Sie einen RayJob bereit, der die ConfigMap als Volume bereitstellt, damit das Trainingsskript auf den Knoten ausgeführt werden kann.

Nachdem Sie diese Schritte ausgeführt haben, verteilt Ray Train die JAX-Arbeitslast automatisch auf die Knoten und Kueue sorgt dafür, dass Sie die erforderlichen Maschinen erhalten.

Das Trainingsskript

Kopieren Sie das folgende Python-Skript und fügen Sie es in eine Datei mit dem Namen train.py ein:

import time
import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.v2.jax import JaxTrainer
import jax
import jax.numpy as jnp

def train_func():
    # JaxTrainer handles JAX distributed setup.
    print(f"Local Devices: {jax.local_devices()}")

    # Simple Linear Regression Training Loop.
    key = jax.random.PRNGKey(0)
    x = jax.random.normal(key, (1000, 10))
    w_true = jax.random.normal(key, (10, 1))
    y = jnp.dot(x, w_true)

    # Initialize weights
    w = jnp.zeros((10, 1))
    learning_rate = 0.1

    @jax.jit
    def update(w, x, y):
        y_pred = jnp.dot(x, w)
        loss = jnp.mean((y_pred - y) ** 2)
        grad = jax.grad(lambda w: jnp.mean((jnp.dot(x, w) - y) ** 2))(w)
        return w - learning_rate * grad, loss

    # Training loop
    print("Starting training...")
    for epoch in range(50):
        w, loss = update(w, x, y)
        if epoch % 10 == 0:
            train.report({"loss": loss.item(), "epoch": epoch})
            print(f"Epoch {epoch}: Loss {loss:.4f}")

    print("Training Complete!")
    # Allow metrics to sync before closing
    time.sleep(3)

def main():
    scaling_config = ScalingConfig(
        num_workers=4,
        resources_per_worker={"TPU": 4},
        use_tpu=True,
        topology="4x4",
        accelerator_type="TPU-V6E"
    )

    trainer = JaxTrainer(
        train_loop_per_worker=train_func,
        scaling_config=scaling_config
    )

    result = trainer.fit()
    print(f"Run Result: {result.metrics}")

if __name__ == "__main__":
    main()

Im Trainingsskript wird JAX verwendet, eine Python-Bibliothek für numerische Hochleistungsberechnungen, um ein lineares Regressionsmodell zu trainieren. Dieses Script ist ein vereinfachtes Beispiel, das zeigen soll, wie DWS und Kueue für das automatisierte Fallback verwendet werden. Es führt keine Datenparallelität oder Modellparallelität aus.

Im Abschnitt ScalingConfig des Trainingsskripts werden die Hardwareanforderungen für den Trainingsjob definiert. In diesem Abschnitt wird eine 4 × 4-TPU-Topologie angefordert, die dem physischen Layout der zuvor konfigurierten Knotenpools entspricht.

ConfigMap erstellen

Laden Sie den Inhalt Ihres train.py-Scripts in ein Kubernetes-ConfigMap-Objekt hoch. So kann das Skript im Cluster gespeichert und für Ihren RayJob verfügbar gemacht werden:

kubectl create configmap jax-train-script --from-file=train.py

Der RayJob, den Sie im nächsten Abschnitt definieren, hängt diese ConfigMap als Volume ein. Dadurch wird die Skriptdatei in den Ray-Containern angezeigt, sodass die Ray-Software sie finden und ausführen kann.

RayJob-Manifest anwenden

Erstellen Sie eine Datei mit dem Namen rayjob-tpu-v6e-dws.yaml und mit folgendem Inhalt. In diesem Manifest wird Ihr Trainingsjob definiert und dem System mitgeteilt, wie er weitergeleitet werden soll:

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: rayjob-tpu-v6e-dws-${JOB_ID}
  labels:
    kueue.x-k8s.io/queue-name: user-queue
  annotations:
    kueue.x-k8s.io/elastic-job: "true"
spec:
  shutdownAfterJobFinishes: true
  entrypoint: python /app/train.py
  runtimeEnvYAML: |
    pip:
      - jax[tpu]==0.8.2
      - pandas==2.3.3
  rayClusterSpec:
    enableInTreeAutoscaling: true
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-nodepool: default-pool
          containers:
            - name: ray-head
              image: rayproject/ray:2.53.0-py311
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "2"
                  memory: "4Gi"
                requests:
                  cpu: "2"
                  memory: "4Gi"
              volumeMounts:
                - mountPath: /app
                  name: train-script-volume
          volumes:
            - name: train-script-volume
              configMap:
                name: jax-train-script
    workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 2
        numOfHosts: 4
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            tolerations:
              - key: "google.com/tpu"
                operator: "Exists"
                effect: "NoSchedule"
              - key: "cloud.google.com/gke-queued"
                operator: "Exists"
                effect: "NoSchedule"
            containers:
              - name: ray-worker
                image: rayproject/ray:2.53.0-py311
                resources:
                  limits:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                  requests:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                volumeMounts:
                  - mountPath: /app
                    name: train-script-volume
            volumes:
              - name: train-script-volume
                configMap:
                  name: jax-train-script
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

Dieses Manifest enthält drei Konfigurationen, die das Fallback-System ermöglichen:

  • Spezifische Hardware anfordern:Im Abschnitt nodeSelector wird die Hardware angegeben, die für Ihr Skript erforderlich ist. In diesem Beispiel ist das ein tpu-v6e-slice mit einer 4x4-Topologie.
  • Wählt die Warteschlange aus:Das Label kueue.x-k8s.io/queue-name leitet Ihren Job direkt an Kueue weiter. Dadurch wird die automatische Fallback-Logik aktiviert.
  • Tolerates DWS flex-start nodes (Toleriert DWS-Flex-Start-Knoten): Im Abschnitt tolerations kann der Job in Ihrem Plan B-Knotenpool ausgeführt werden. Da DWS-Flex-Start-Knoten von GKE speziell markiert (mit einer Markierung versehen) werden, damit normale Arbeitslasten nicht versehentlich auf ihnen ausgeführt werden, muss Ihr Job die Markierung cloud.google.com/gke-queued explizit tolerieren.

Arbeitslast einreichen

Um nachzuweisen, dass das Fallback-System funktioniert, müssen Sie zwei Jobs einreichen. Der erste Job verbraucht die On-Demand-Kapazität von Plan A, sodass der zweite Job auf die DWS-Flex-Start-Kapazität von Plan B zurückgreifen muss.

Führen Sie den folgenden Befehl aus, um die beiden Jobs zu senden. Der Befehl verwendet eine for-Schleife und envsubst, um für jeden Lauf eine eindeutige Job-ID in das Manifest einzufügen:

for i in 1 2; do
  export JOB_ID=$i
  envsubst < rayjob-tpu-v6e-dws.yaml | kubectl apply -f -
  echo "Submitted Job $i"
  sleep 2
done

Nachdem Sie die Jobs gesendet haben, verarbeitet das System die Arbeitslast so:

  1. Abfangen:Kueue erkennt die Jobs anhand des Warteschlangenlabels und hält sie vorübergehend an.
  2. Entscheidung:Kueue vergleicht die Ressourcenverfügbarkeit mit den Regeln des Administrators. Zuerst wird der Pool für Plan A geprüft.
  3. Aufgabe:
    • Da Ressourcen für Plan A für den ersten Job verfügbar sind, weist Kueue Job 1 dort zu.
    • Da Job 1 die Ressourcen von Plan A nutzt, weist Kueue Job 2 automatisch dem Pool von Plan B (DWS-Flex-Start) zu.
  4. Starten:Kueue setzt die Jobs fort. Diese Aktion löst das GKE-Cluster-Autoscaling aus, um die Knoten bereitzustellen und die Trainingsskripts zu starten.

Verbindung zum RayJob herstellen

Als letzten Überprüfungsschritt können Sie mit dem Befehl kubectl port-forward eine Verbindung zum Ray-Dashboard herstellen und die Ausführung Ihrer Jobs beobachten.

Führen Sie den folgenden Befehl aus, um den Status Ihres ersten Jobs zu prüfen:

kubectl port-forward service/rayjob-tpu-v6e-dws-1-head-svc 8265:8265 &

Öffnen Sie nach der Ausführung dieses Befehls einen Webbrowser und rufen Sie http://localhost:8265 auf. Im Ray-Dashboard können Sie den Jobstatus und die gemeldeten Messwerte einsehen, um zu prüfen, ob beide Jobs in ihren jeweiligen Knotenpools erfolgreich abgeschlossen wurden.

Sie können die Logs für den ersten Job auch mit dem folgenden Befehl aufrufen:

kubectl logs job/rayjob-tpu-v6e-dws-1

Die gekürzte Ausgabe des Trainingsskripts sollte in etwa so aussehen: Die Meldungen Training Complete! und Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded sollten am Ende der Ausgabe angezeigt werden:

(pid=, ip=10.68.3.4) 5] XLA::TPU program HBM usage: 52.5K / 31.25G
(pid=, ip=10.68.9.4) :2152] XLA::TPU program VMEM usage: 141.0K / 128.00M [repeated 5x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722540     855 deepsea_compiler_backend.cc:2163] Total hbm usage >= 260.14M: [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167]     reserved           204B [repeated 19x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722542     855 deepsea_compiler_backend.cc:2163]     program           70.0K  [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     arguments            0B  [repeated 12x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777627     888 deepsea_compiler_backend.cc:2163] Output size 0B; shares 0B with arguments. [repeated 14x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777625     888 deepsea_compiler_backend.cc:2163] Total host usage >= 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     program         unknown size  [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167] Program sflag requirement 224B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167]     scoped              40B [repeated 21x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167] Program vmem requirement 141.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program smem requirement 40B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program host requirement 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program hbm requirement 70.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2167]     overlays          70.0K [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2175] XLA::TPU program SMEM usage: 1.9K / 1.00M (3 parameters) [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167]     HLO temp          76.0K (0.0% utilization: Unpadded (0B) Padded (0B), 100.0% fragmentation (76.0K)) [repeated 14x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Training Complete! [repeated 3x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Epoch 40: Loss 0.0000 [repeated 3x across cluster]
2026-03-20 03:59:51,008 SUCC cli.py:65 -- ------------------------------------------
2026-03-20 03:59:51,008 SUCC cli.py:66 -- Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded
2026-03-20 03:59:51,008 SUCC cli.py:67 -- ------------------------------------------

Bereinigen

Damit Ihrem Google Cloud -Konto die in diesem Dokument verwendeten Ressourcen nicht in Rechnung gestellt werden, können Sie entweder das Projekt löschen, das die Ressourcen enthält, oder das Projekt beibehalten und die einzelnen Ressourcen löschen.

Projekt löschen

  • Wechseln Sie in der Google Cloud -Console zur Seite Ressourcen verwalten.

    Zur Seite „Ressourcen verwalten“

  • Wählen Sie in der Projektliste das Projekt aus, das Sie löschen möchten, und klicken Sie dann auf Löschen.
  • Geben Sie im Dialogfeld die Projekt-ID ein und klicken Sie auf Shut down (Beenden), um das Projekt zu löschen.

    • Einzelne Ressourcen löschen

    Wenn Sie das in diesem Dokument verwendete GGoogle Cloud -Projekt beibehalten möchten, führen Sie den folgenden Befehl aus, um den Cluster zu löschen:

    gcloud container clusters delete ${CLUSTER_NAME} \
    --location=${ZONE} \
    --project=${PROJECT_ID} \
    --quiet

    Zusammenfassung

    In diesem Dokument haben Sie eine Ray-Trainingsumgebung konfiguriert und getestet. In dieser Umgebung werden ein primärer Knotenpool und ein DWS-Sicherungspool verwendet, um die Hardwareverfügbarkeit zu maximieren. Durch die automatische Umstellung auf DWS, wenn primäre Maschinen nicht verfügbar sind, haben Sie die Wartezeit Ihrer Trainingsjobs minimiert.

    Dazu haben Sie die folgenden Schritte ausgeführt:

    1. GKE-Cluster erstellt:Die Umgebung zum Hosten der Knotenpools und Planungstools wurde eingerichtet.
    2. Knotenpools konfiguriert: Sie haben einen On-Demand-Knotenpool (Plan A) und einen DWS-Knotenpool (Plan B) erstellt.
    3. Kueue installiert und konfiguriert: Den Kueue-Controller bereitgestellt und Prioritätsregeln angewendet, die das System anweisen, zuerst Plan A zu versuchen und auf Plan B zurückzugreifen.
    4. ConfigMap erstellt: Ein vereinfachtes JAX-Trainingsskript wurde im Cluster bereitgestellt, um als Testarbeitslast zu dienen.
    5. RayJob-Manifest definiert: Der Job ist so konfiguriert, dass er bestimmte Hardware anfordert, an den Kueue-Controller weitergeleitet wird und DWS-Knoten toleriert.
    6. Arbeitslast übermittelt:Es wurden zwei Jobs übermittelt, um Kueue zu zwingen, den zweiten Job automatisch an Plan B weiterzuleiten, wenn die Ressourcen von Plan A aufgebraucht sind.
    7. Ergebnisse überprüft: Sie haben die Portweiterleitung verwendet, um eine Verbindung zum Ray-Dashboard herzustellen, und bestätigt, dass beide Jobs erfolgreich ausgeführt wurden.

    Nächste Schritte