KubeRay mit TPU Trillium konfigurieren

Autopilot Standard

In dieser Anleitung erfahren Sie, wie Sie KubeRay mit TPU Trillium in Google Kubernetes Engine (GKE) konfigurieren. Sie lernen, TPU Konfigurationen mit einem oder mehreren Hosts einzurichten, 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-Experten, die erfahren möchten, wie sie die TPU Trillium-Initialisierung mit KubeRay für Knotenpools mit einem oder mehreren Hosts konfigurieren. In dieser Anleitung wird gezeigt, wie Sie ein Skript mit Jax ausführen, mit dem die erfolgreiche TPU-Initialisierung überprüft wird. In dieser Anleitung wird kein Modell bereitgestellt.

Bevor Sie KubeRay in GKE konfigurieren, sollten Sie mit den Ray-Definitionen und der Terminologie in GKE vertraut sein.

Übersicht

In dieser Anleitung wird gezeigt, wie Sie ein Python-Skript mit Jax ausführen, mit dem überprüft wird, 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 bietet, Ray-Anwendungen in Kubernetes bereitzustellen, zu verwalten und zu überwachen.

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.

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 diese Aufgabe verwenden möchten, installieren und dann initialisieren Sie die gcloud CLI. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit dem gcloud components update Befehl ab. Ältere gcloud CLI-Versionen unterstützen möglicherweise nicht die Ausführung der Befehle in diesem Dokument.
  Hinweis: Legen Sie für vorhandene Installationen der gcloud CLI das Attribut compute/region 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

Cloud Shell ist mit den in dieser Anleitung verwendeten Befehlszeilentools gcloud, helm und kubectl 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 ist 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, um mit Ihrem Cluster zu kommunizieren :

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

Standard

1. Erstellen Sie in Cloud Shell einen Standardcluster, in dem das Ray-Operator-Add-on aktiviert ist. 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 ist 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, um mit Ihrem Cluster zu kommunizieren :

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

3. Sie können einen TPU-Slice-Knotenpool mit einem oder 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 eine rayClusterSpec, die die gewünschte Ray-Clusterkonfiguration (Head- und Worker-Gruppen) definiert.
• Bestimmten Job ausführen:Das Feld entrypoint in der RayJob-Spezifikation gibt den Befehl oder das Skript an, das im erstellten Ray-Cluster ausgeführt werden soll. In dieser Anleitung ist das 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 v6e-16-RayJob 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 v6e-16-RayJob 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 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 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 Wert TPU_WORKER_ID=0 entspricht. In der TPU-Gruppe mit mehreren Hosts ist für den Ray-Knoten mit TPU_WORKER_ID=0 in den Ressourcen TPU-v6e-16-head: 1.0 festgelegt. Diese Umgebungsvariable TPU_WORKER_ID wird von einem mutierenden GKE-Webhook für KubeRay festgelegt.

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

Bereinigen

Löschen Sie nach Abschluss der Anleitung 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

• Learn about Ray in Kubernetes
• vLLM in GKE mit TPUs bereitstellen
• SDXL in GKE mit TPUs bereitstellen
• Learn more TPUs in GKE