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 Google Cloud-Dynamic Workload Scheduler (DWS).

Mit der Konfiguration, die Sie in diesem Dokument definieren, werden ein primärer Knotenpool und ein Backup-Pool von Knoten 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 zuverlässigen, unterbrechungsfreien Zugriff auf diese Maschinen haben, sobald sie ausgeführt werden.
  • DWS flex-start (Plan B): Dies ist Ihr Backup-Knotenpool. Wenn Maschinen vom Typ „Plan A“ nicht verfügbar sind, wird Ihr Job vom Kueue-Planungsprogramm automatisch diesem Pool vom Typ „Plan B“ zugewiesen. DWS sucht dann nach der Plan B-Hardware, kann aber keinen sofortigen Zugriff garantieren, 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 Script 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 Arbeitsabläufe 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 müssen.

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. Für Ihren Job wird immer zuerst versucht, diesen Pool zu verwenden.
  • DWS flex-start (Plan B)-Knotenpool:Der Backup-Knotenpool. Wenn keine Maschinen im primären Pool 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.

    Rollen, die zum Auswählen oder Erstellen eines Projekts erforderlich sind

    • Projekt auswählen: Für die Auswahl eines Projekts ist keine bestimmte IAM-Rolle erforderlich. Sie können jedes Projekt auswählen, für das Ihnen eine Rolle zugewiesen wurde.
    • Projekt erstellen: Zum Erstellen eines Projekts benötigen Sie die Rolle „Projektersteller“ (roles/resourcemanager.projectCreator), die die Berechtigung resourcemanager.projects.create enthält. Weitere Informationen zum Zuweisen von Rollen

    Zur Projektauswahl

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

  3. Aktivieren Sie die Google Kubernetes Engine API und die Cloud TPU API.

    Rollen, die zum Aktivieren von APIs erforderlich sind

    Zum Aktivieren von APIs benötigen Sie die IAM-Rolle „Service Usage-Administrator“ (roles/serviceusage.serviceUsageAdmin), die die Berechtigung serviceusage.services.enable enthält. Weitere 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 und Cloud TPU-Umgebung einrichten.

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 Werte jedes Mal neu einzugeben oder zu ersetzen. Dieser Ansatz erleichtert die Nachvollziehbarkeit des Prozesses und verringert das Fehlerrisiko.

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 finden Sie eine Erläuterung 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-Knotenpools mit flexiblem Start. (Das ist Ihr Plan B-Knotenpool.)

Infrastruktur konfigurieren (Clusteradministrator)

Als Clusteradministrator konfigurieren Sie den GKE-Cluster und die Knotenpools, um den Fallback-Mechanismus zu unterstützen.

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. Dieses Feature verkürzt die Zeit, die zum Starten großer KI-Container-Images benötigt wird, erheblich.

  2. Rufen Sie die Anmeldedaten des Clusters ab, damit die kubectl-Befehlszeile eine Verbindung zum Cluster herstellen kann. Mit diesem Befehl wird Ihre 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 Backup-Knotenpools 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, die Sie für Ihr KI-Modell benötigen, z. B. GPUs oder verschiedene TPUs.

  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. Wenn Sie diese spezielle Hardware unbedingt haben möchten, bedeutet das, dass Sie zu einer anderen Methode wechseln, um sie zu erwerben, 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 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --enable-queued-provisioning \
  --flex-start \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    Diese Befehle verwenden die folgenden Schlüsselflags:

    • --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: Damit wird sichergestellt, dass der Knotenpool nicht Ihre vorreservierte Hardware nutzt.Mit Google Cloud können Sie bestimmte Maschinen reservieren, um die Verfügbarkeit zu gewährleisten. Wenn Sie dieses Flag auf none setzen, werden diese Reservierungen umgangen und stattdessen dynamisch nicht reservierte Maschinen angefordert.
    • --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, wenn diese verfügbar wird.

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 folgendem Inhalt. In dieser Datei werden zwei Varianten von Ressourcen (On-Demand und DWS Flex-Start) und eine Clusterwarteschlange definiert, in der sie priorisiert werden. 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, wird stattdessen versucht, 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-Flex-Start-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 einen Trainings-Workload, 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 für die benötigten Maschinen.

Das Trainingsskript

Kopieren Sie das folgende Python-Skript und fügen Sie es in eine Datei namens 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 numerisches Hochleistungs-Computing, um ein lineares Regressionsmodell zu trainieren. Dieses Skript 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 4x4-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

Die RayJob, die 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 folgendem Inhalt. Dieses Manifest definiert Ihren Trainingsjob und gibt dem System an, 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:

  • Spezielle Hardware anfordern:Im Abschnitt nodeSelector wird die Hardware angegeben, die für das Skript erforderlich ist (in diesem Beispiel 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-Knoten mit flexiblem Start von GKE speziell markiert (tainted) 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. Beim ersten Job wird die On-Demand-Kapazität von Plan A genutzt, wodurch der zweite Job auf die Flex-Start-Kapazität von Plan B mit DWS zurückgreifen muss.

Führen Sie den folgenden Befehl aus, um die beiden Jobs zu senden. Im Befehl werden eine for-Schleife und envsubst verwendet, 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 setzt sie vorübergehend aus.
  2. Entscheidung:Kueue bewertet die Ressourcenverfügbarkeit anhand der 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 für Job 1 die Ressourcen von Plan A verwendet werden, 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 dem Ausführen dieses Befehls einen Webbrowser und rufen Sie http://localhost:8265 auf. Im Ray-Dashboard können Sie den Jobstatus und die gemeldeten Messwerte aufrufen, 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. Gegen Ende der Ausgabe sollten die Meldungen Training Complete! und Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded 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 Backup-DWS-Pool 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.

    Damit das funktioniert, haben Sie die folgenden Schritte ausgeführt:

    1. GKE-Cluster erstellt:Die Umgebung zum Hosten der Knotenpools und Planungstools wurde eingerichtet.
    2. Knotenpools konfiguriert:Ein On-Demand-Knotenpool (Plan A) und ein DWS-Knotenpool (Plan B) wurden erstellt.
    3. Kueue installiert und konfiguriert:Der Kueue-Controller wurde bereitgestellt und es wurden Prioritätsregeln angewendet, die das System anweisen, zuerst Plan A auszuführen und dann 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 wurde so konfiguriert, dass er bestimmte Hardware anfordert, an den Kueue-Controller weitergeleitet wird und DWS-Knoten toleriert.
    6. Arbeitslast gesendet:Es wurden zwei Jobs gesendet, um Kueue zu zwingen, den zweiten Job automatisch an Plan B weiterzuleiten, wenn die Ressourcen von Plan A aufgebraucht sind.
    7. Ergebnisse überprüft:Portweiterleitung verwendet, um eine Verbindung zum Ray-Dashboard herzustellen und zu bestätigen, dass beide Jobs erfolgreich ausgeführt wurden.

    Nächste Schritte