KubeRay mit TPU Trillium konfigurieren

Autopilot Standard

In dieser Anleitung erfahren Sie, wie Sie KubeRay mit TPU Trillium in Google Kubernetes Engine (GKE) konfigurieren. Hier erfahren Sie, wie Sie TPU-Konfigurationen mit einem einzelnen Host und mit mehreren Hosts einrichten, einschließlich der erforderlichen Umgebungsvariablen und Pod-Spezifikationen für TPU Trillium.

Diese Anleitung richtet sich an Plattformadministratoren und ‑operatoren sowie an Daten- und KI-Spezialisten, die erfahren möchten, wie die TPU-Trillium-Initialisierung mit KubeRay für Knotenpools mit einem und mehreren Hosts konfiguriert wird. In diesem Tutorial wird gezeigt, wie Sie ein Skript mit Jax ausführen, das die erfolgreiche TPU-Initialisierung überprüft. In dieser Anleitung wird kein Modell bereitgestellt.

Bevor Sie KubeRay in GKE konfigurieren, sollten Sie sich mit Ray-Definitionen und ‑Begriffen in GKE vertraut machen.

Übersicht

In dieser Anleitung erfahren Sie, wie Sie ein Python-Skript mit Jax ausführen, das überprüft, ob die TPU Trillium-Initialisierung mit KubeRay erfolgreich war. JAX ist eine leistungsstarke numerische Berechnungsbibliothek, die Arbeitslasten für maschinelles Lernen unterstützt. KubeRay ist ein Kubernetes-Operator, der eine einheitliche Möglichkeit zum Bereitstellen, Verwalten und Überwachen von Ray-Anwendungen in Kubernetes bietet.

Für Trillium-TPUs (v6e) sind bestimmte Umgebungsvariablen und Pod-Spezifikationen erforderlich, die sich von früheren TPU-Generationen unterscheiden. In dieser Anleitung finden Sie die erforderlichen Konfigurationen, um eine Arbeitslast mit KubeRay auf Trillium-TPUs bereitzustellen.

Hinweise

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 gcloud components update ab.
  Hinweis: Legen Sie für vorhandene Installationen der gcloud CLI das compute/region-Attribut fest. Wenn Sie hauptsächlich zonale Cluster verwenden, legen Sie stattdessen compute/zone fest. Wenn Sie einen standardmäßigen Speicherort festlegen, können Sie in der gcloud CLI Fehler wie diesen vermeiden: One of [--zone, --region] must be supplied: Please specify location. Bei bestimmten Befehlen müssen Sie möglicherweise den Speicherort angeben, wenn er sich von dem von Ihnen festgelegten Standard unterscheidet.

• Prüfen Sie, ob die Ray CLI (Version 2.37.0) installiert ist.

Cloud Shell aktivieren

Die in dieser Anleitung verwendeten Befehlszeilentools gcloud, helm und kubectl sind in Cloud Shell vorinstalliert.

1. Rufen Sie die Google Cloud Console auf.

2. Klicken Sie oben im Google Cloud Console-Fenster auf die Schaltfläche Cloud Shell aktivieren .

   In einem neuen Frame in der Google Cloud Console wird eine Cloud Shell-Sitzung geöffnet und darin eine Eingabeaufforderung angezeigt.

   

GKE-Cluster und -Knotenpool erstellen

Sie können KubeRay auf TPUs in einem GKE-Cluster im Autopilot- oder Standardmodus konfigurieren. Für eine vollständig verwaltete Kubernetes-Umgebung empfehlen wir die Verwendung eines Autopilot-Clusters. Informationen zum Auswählen des GKE-Betriebsmodus, der für Ihre Arbeitslasten am besten geeignet ist, finden Sie unter GKE-Betriebsmodi.

Autopilot

1. Führen Sie in Cloud Shell den folgenden Befehl aus:

   gcloud container clusters create-auto CLUSTER_NAME \
       --enable-ray-operator \
       --release-channel=rapid \
       --location=LOCATION
   

   Ersetzen Sie Folgendes:

   • CLUSTER_NAME ist der Name des neuen Clusters.
   • LOCATION: Die Region, in der Ihre TPU-Trillium-Kapazität verfügbar ist. Weitere Informationen finden Sie unter TPU-Verfügbarkeit in GKE.

   GKE erstellt einen Autopilot-Cluster mit aktiviertem Ray-Operator-Add-on. Das Add-on installiert den Ray TPU-Webhook automatisch in der Clustersteuerungsebene.

2. Konfigurieren Sie kubectl für die Kommunikation mit Ihrem Cluster :

   gcloud container clusters get-credentials CLUSTER_NAME --location=LOCATION
   

Standard

1. Erstellen Sie in Cloud Shell einen Standardcluster, der das Ray-Operator-Add-on aktiviert. Führen Sie dazu den folgenden Befehl aus :

   gcloud container clusters create CLUSTER_NAME \
     --location LOCATION \
     --addons=RayOperator \
     --cluster-version=1.33 \
     --machine-type=n1-standard-16
   

   Ersetzen Sie Folgendes:

   • CLUSTER_NAME ist der Name des neuen Clusters.
   • LOCATION: Die Region, in der Ihre TPU-Trillium-Kapazität verfügbar ist. Weitere Informationen finden Sie unter TPU-Verfügbarkeit in GKE.

   Die Erstellung eines Clusters kann einige Minuten dauern.

2. Konfigurieren Sie kubectl für die Kommunikation mit Ihrem Cluster :

   gcloud container clusters get-credentials CLUSTER_NAME --location=LOCATION
   

3. Sie können einen TPU-Slice-Knotenpool mit einem einzelnen Host oder einen TPU-Slice-Knotenpool mit mehreren Hosts erstellen:

Einzelner Host

Führen Sie in Cloud Shell den folgenden Befehl aus:

gcloud container node-pools create v6e-4 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=1 \
    --threads-per-core=1 \
    --tpu-topology=2x2

Mehrere Hosts

Führen Sie in Cloud Shell den folgenden Befehl aus:

gcloud container node-pools create v6e-16 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=4 \
    --threads-per-core=1 \
    --tpu-topology=4x4

Benutzerdefinierte RayJob-Ressource ausführen

Wenn Sie ein RayJob-Manifest definieren, weisen Sie KubeRay an, Folgendes zu tun:

• RayCluster erstellen:Die RayJob-Spezifikation enthält ein rayClusterSpec, das die gewünschte Ray-Clusterkonfiguration (Head- und Worker-Gruppen) definiert.
• Bestimmten Job ausführen:Im Feld entrypoint von RayJob wird der Befehl oder das Skript angegeben, das im erstellten Ray-Cluster ausgeführt werden soll. In dieser Anleitung ist entrypoint ein Python-Skript (tpu_list_devices.py), mit dem die TPU-Trillium-Initialisierung überprüft wird.

So erstellen Sie eine benutzerdefinierte RayJob-Ressource:

Einzelner Host

1. Erstellen Sie das folgende ray-job.tpu-v6e-singlehost.yaml-Manifest:

   apiVersion: ray.io/v1
   kind: RayJob
   metadata:
     name: v6e-4-job
   spec:
     entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
     runtimeEnvYAML: |
       working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
       pip:
         - jax[tpu]==0.4.33
         - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
     rayClusterSpec:
       rayVersion: '2.43.0'
       headGroupSpec:
         rayStartParams: {}
         template:
           spec:
             containers:
             -   name: ray-head
                 image: rayproject/ray:2.43.0-py310
                 ports:
                   - containerPort: 6379
                     name: gcs-server
                   - containerPort: 8265
                     name: dashboard
                   - containerPort: 10001
                     name: client
                 resources:
                   limits:
                     cpu: "8"
                     memory: 40G
                   requests:
                     cpu: "8"
                     memory: 40G
       workerGroupSpecs:
       -   replicas: 1
           minReplicas: 1
           maxReplicas: 1
           numOfHosts: 1
           groupName: tpu-group
           rayStartParams: {}
           template:
             spec:
               containers:
               -   name: ray-worker
                   image: rayproject/ray:2.43.0-py310
                   resources:
                     limits:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
                     requests:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
               nodeSelector:
                 cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                 cloud.google.com/gke-tpu-topology: 2x2

2. Wenden Sie das Manifest an:

   kubectl apply -f ray-job.tpu-v6e-singlehost.yaml
   

3. Prüfen Sie, ob der RayJob erstellt wurde und ausgeführt wird:

   kubectl get rayjobs v6e-4-job
   

   Die Ausgabe sieht etwa so aus:

   NAME      JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME       START TIME  END TIME   AGE
   v6e-4-job PENDING      Running             v6e-4-job-raycluster   2024-10-15T23:15:22Z  20s
   

4. Geben Sie die Ausgabe des RayJob aus.

   kubectl logs -l=job-name=v6e-4-job
   

   Die Ausgabe sieht etwa so aus:

   2024-10-15 16:15:40,222 INFO cli.py:300 -- ray job stop v6e-4-job-hzq5q
   2024-10-15 16:15:40,246 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
   2024-10-15 16:15:40,112 INFO job_manager.py:528 -- Runtime env is setting up.
   2024-10-15 16:15:50,181 INFO worker.py:1461 -- Using address 10.84.1.25:6379 set in the environment variable RAY_ADDRESS
   2024-10-15 16:15:50,181 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.25:6379...
   2024-10-15 16:15:50,186 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.25:8265
   ['TPU cores:4']
   2024-10-15 16:16:12,349 SUCC cli.py:63 -- -------------------------------------
   2024-10-15 16:16:12,349 SUCC cli.py:64 -- Job 'v6e-4-job-hzq5q' succeeded
   2024-10-15 16:16:12,349 SUCC cli.py:65 -- -------------------------------------
   

Mehrere Hosts

1. Erstellen Sie das folgende ray-job.tpu-v6e-multihost.yaml-Manifest:

   apiVersion: ray.io/v1
   kind: RayJob
   metadata:
     name: v6e-16-job
   spec:
     entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
     runtimeEnvYAML: |
       working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
       pip:
         - jax[tpu]==0.4.33
         - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
     rayClusterSpec:
       rayVersion: '2.43.0'
       headGroupSpec:
         rayStartParams: {}
         template:
           spec:
             containers:
             -   name: ray-head
                 image: rayproject/ray:2.43.0-py310
                 ports:
                   - containerPort: 6379
                     name: gcs-server
                   - containerPort: 8265
                     name: dashboard
                   - containerPort: 10001
                     name: client
                 resources:
                   limits:
                     cpu: "8"
                     memory: 40G
                   requests:
                     cpu: "8"
                     memory: 40G
       workerGroupSpecs:
         - replicas: 1
           minReplicas: 1
           maxReplicas: 1
           numOfHosts: 4
           groupName: tpu-group
           rayStartParams: {}
           template:
             spec:
               containers:
                 - name: ray-worker
                   image: rayproject/ray:2.43.0-py310
                   resources:
                     limits:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
                     requests:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
                   env:
                   - name: NODE_IP
                     valueFrom:
                       fieldRef:
                         fieldPath: status.hostIP
                   - name: VBAR_CONTROL_SERVICE_URL
                     value: $(NODE_IP):8353
                   - name: JAX_PLATFORMS
                     value: tpu,cpu
                   - name: ENABLE_PJRT_COMPATIBILITY
                     value: "true"
                   ports:
                   - containerPort: 8081
                     name: mxla
               nodeSelector:
                 cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                 cloud.google.com/gke-tpu-topology: 4x4

2. Wenden Sie das Manifest an:

   kubectl apply -f ray-job.tpu-v6e-multihost.yaml
   

3. Prüfen Sie, ob der RayJob „v6e-16“ erstellt wurde und ausgeführt wird:

   kubectl get rayjobs v6e-16-job
   

   Die Ausgabe sieht etwa so aus:

   NAME         JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME              START TIME             END TIME   AGE
   v6e-16-job                Running             v6e-16-job-raycluster-qr6vk   2024-10-16T19:28:19Z              66s
   

4. Geben Sie die Ausgabe des RayJob „v6e-16“ aus:

   kubectl logs -l=job-name=v6e-16-job
   

   Die Ausgabe sieht etwa so aus:

   2024-10-16 12:21:33,986 INFO cli.py:300 -- ray job stop v6e-16-job-z44s7
   2024-10-16 12:21:34,011 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
   2024-10-16 12:21:33,826 INFO job_manager.py:528 -- Runtime env is setting up.
   2024-10-16 12:21:46,327 INFO worker.py:1461 -- Using address 10.84.1.61:6379 set in the environment variable RAY_ADDRESS
   2024-10-16 12:21:46,327 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.61:6379...
   2024-10-16 12:21:46,333 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.61:8265
   ['TPU cores:16', 'TPU cores:16', 'TPU cores:16', 'TPU cores:16']
   2024-10-16 12:22:12,156 SUCC cli.py:63 -- ---------------------------------
   2024-10-16 12:22:12,156 SUCC cli.py:64 -- Job 'v6e-16-job-z44s7' succeeded
   2024-10-16 12:22:12,156 SUCC cli.py:65 -- ---------------------------------
   

RayJob im Ray-Dashboard ansehen

Prüfen Sie, ob GKE den RayCluster-Dienst erstellt hat, und stellen Sie eine Verbindung zur RayCluster-Instanz her.

Einzelner Host

1. Rufen Sie den Namen des generierten RayCluster für den RayJob ab:

   export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-4-job -o jsonpath='{.status.rayClusterName}')
   

2. Rufen Sie den Namen des RayCluster-Head-Dienstes ab:

   export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')
   

3. Stellen Sie eine Verbindung zum Ray-Dashboard her, indem Sie den Head-Dienst per Portweiterleitung weiterleiten:

   kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
   

4. Öffnen Sie einen Webbrowser und geben Sie die folgende URL ein:

   http://localhost:8265/#/jobs
   

5. Rufen Sie den RayJob-Status und relevante Logs auf.

Mehrere Hosts

1. Rufen Sie den Namen des generierten RayCluster für den RayJob ab:

   export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-16-job -o jsonpath='{.status.rayClusterName}')
   

2. Rufen Sie den Namen des RayCluster-Head-Dienstes ab:

   export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')
   

3. Stellen Sie eine Verbindung zum Ray-Dashboard her, indem Sie den Head-Dienst per Portweiterleitung weiterleiten:

   kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
   

4. Öffnen Sie einen Webbrowser und geben Sie die folgende URL ein:

   http://localhost:8265/#/jobs
   

5. Rufen Sie den RayJob-Status und relevante Logs auf.

Ray legt eine TPU-{accelerator}-Head-Ressource fest, um den Ray-Worker-Knoten zu identifizieren, der dem TPU_WORKER_ID=0-Wert entspricht. In der TPU-Gruppe mit mehreren Hosts hat der Ray-Knoten mit TPU_WORKER_ID=0 TPU-v6e-16-head: 1.0 in seinen Ressourcen festgelegt. Diese Umgebungsvariable TPU_WORKER_ID wird von einem mutierenden GKE-Webhook für KubeRay festgelegt.

Erfolg: Sie haben erfolgreich einen RayJob in einem TPU-Slice-Knotenpool ausgeführt, auf dem TPU Trillium ausgeführt wird. Der Abschluss des Ray-Jobs gibt an, dass Jax erfolgreich initialisiert wurde und die verteilten TPU-Knoten korrekt aktiviert sind.

Bereinigen

Wenn Sie die Anleitung abgeschlossen haben, löschen Sie den RayJob, um unerwünschte Kosten für Ihr Konto zu vermeiden:

Einzelner Host

kubectl delete rayjobs v6e-4-job

Mehrere Hosts

kubectl delete rayjobs v6e-16-job

Nächste Schritte

• Weitere Informationen zu Ray in Kubernetes
• Informationen zum Bereitstellen von vLLM in GKE mit TPUs
• Informationen zum Bereitstellen von SDXL in GKE mit TPUs
• Weitere Informationen zu TPUs in GKE