Mengonfigurasi KubeRay dengan TPU Trillium

Autopilot Standard

Tutorial ini menunjukkan cara mengonfigurasi KubeRay dengan TPU Trillium di Google Kubernetes Engine (GKE). Pelajari cara menyiapkan konfigurasi TPU host tunggal dan multi-host, termasuk variabel lingkungan dan spesifikasi Pod yang diperlukan untuk TPU Trillium.

Tutorial ini ditujukan bagi admin dan operator Platform serta spesialis Data dan AI yang ingin mempelajari cara mengonfigurasi inisialisasi TPU Trillium dengan KubeRay untuk node pool host tunggal dan multi-host. Tutorial ini menunjukkan cara menjalankan skrip dengan Jax yang memverifikasi keberhasilan inisialisasi TPU. Tutorial ini tidak men-deploy model.

Sebelum mengonfigurasi KubeRay di GKE, pastikan Anda memahami definisi dan terminologi Ray di GKE.

Ringkasan

Tutorial ini menunjukkan cara menjalankan skrip Python dengan Jax yang memverifikasi bahwa inisialisasi TPU Trillium dengan KubeRay berhasil. Jax adalah library komputasi numerik berperforma tinggi yang mendukung beban kerja machine learning. KubeRay adalah operator Kubernetes yang menyediakan cara terpadu untuk men-deploy, mengelola, dan memantau aplikasi Ray di Kubernetes.

TPU Trillium (v6e) memerlukan variabel lingkungan dan spesifikasi Pod tertentu yang berbeda dari generasi TPU sebelumnya. Tutorial ini memberikan konfigurasi yang diperlukan untuk berhasil men-deploy beban kerja dengan KubeRay di TPU Trillium.

Sebelum memulai

Sebelum memulai, pastikan Anda telah melakukan tugas berikut:

• Aktifkan Google Kubernetes Engine API.
Aktifkan Google Kubernetes Engine API
• Jika ingin menggunakan Google Cloud CLI untuk tugas ini, instal lalu lakukan inisialisasi gcloud CLI. Jika sebelumnya Anda telah menginstal gcloud CLI, dapatkan versi terbaru dengan menjalankan gcloud components update.
  Catatan: Untuk penginstalan gcloud CLI yang ada, pastikan untuk menyetel properti compute/region. Jika Anda terutama menggunakan cluster zona, tetapkan compute/zone. Dengan menyetel lokasi default, Anda dapat menghindari error di gcloud CLI yang seperti ini: One of [--zone, --region] must be supplied: Please specify location. Anda mungkin perlu menentukan lokasi dalam perintah tertentu jika lokasi cluster Anda berbeda dengan lokasi default yang Anda tetapkan.

• Pastikan Anda telah menginstal Ray CLI (versi 2.37.0).

Mengaktifkan Cloud Shell

Cloud Shell telah diinstal dengan alat command line gcloud, helm, dan kubectl yang digunakan dalam tutorial ini.

1. Buka Google Cloud console.

2. Di bagian atas jendela konsol Google Cloud , klik tombol Activate Cloud Shell .

   Sesi Cloud Shell akan terbuka di dalam frame baru di konsol Google Cloud dan menampilkan perintah command line.

   

Membuat cluster dan node pool GKE

Anda dapat mengonfigurasi KubeRay di TPU dalam cluster GKE Autopilot atau Standard. Sebaiknya gunakan cluster Autopilot untuk pengalaman Kubernetes yang terkelola sepenuhnya. Untuk memilih mode operasi GKE yang paling sesuai untuk workload Anda, lihat Tentang mode operasi GKE.

Autopilot

1. Jalankan perintah berikut di Cloud Shell:

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

   Ganti kode berikut:

   • CLUSTER_NAME: nama cluster baru.
   • LOCATION: region tempat kapasitas TPU Trillium Anda tersedia. Untuk mengetahui informasi selengkapnya, lihat Ketersediaan TPU di GKE.

   GKE membuat cluster Autopilot dengan add-on operator Ray yang diaktifkan. Add-on ini otomatis menginstal webhook TPU Ray di bidang kontrol cluster.

2. Untuk berkomunikasi dengan cluster Anda, konfigurasi kubectl :

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

Standar

1. Di Cloud Shell, buat cluster Standard yang mengaktifkan add-on operator Ray dengan menjalankan perintah berikut :

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

   Ganti kode berikut:

   • CLUSTER_NAME: nama cluster baru.
   • LOCATION: region tempat kapasitas TPU Trillium Anda tersedia. Untuk mengetahui informasi selengkapnya, lihat Ketersediaan TPU di GKE.

   Pembuatan cluster mungkin memerlukan waktu beberapa menit.

2. Untuk berkomunikasi dengan cluster Anda, konfigurasi kubectl :

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

3. Anda dapat membuat node pool slice TPU host tunggal atau multi-host:

Host tunggal

Jalankan perintah berikut di Cloud Shell:

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

Multi-host

Jalankan perintah berikut di Cloud Shell:

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

Menjalankan resource kustom RayJob

Dengan menentukan manifes RayJob, Anda menginstruksikan KubeRay untuk melakukan hal berikut:

• Buat RayCluster: spesifikasi RayJob mencakup rayClusterSpec yang menentukan konfigurasi cluster Ray (grup head dan worker) yang Anda inginkan.
• Menjalankan Job tertentu: kolom entrypoint dalam RayJob menentukan perintah atau skrip yang akan dieksekusi dalam cluster Ray yang dibuat. Dalam tutorial ini, entrypoint adalah skrip Python (tpu_list_devices.py) yang dirancang untuk memverifikasi inisialisasi TPU Trillium.

Untuk membuat resource kustom RayJob, selesaikan langkah-langkah berikut:

Host tunggal

1. Buat manifes ray-job.tpu-v6e-singlehost.yaml berikut:

   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. Terapkan manifes:

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

3. Pastikan RayJob dibuat dan berjalan:

   kubectl get rayjobs v6e-4-job
   

   Outputnya mirip dengan hal berikut ini:

   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. Cetak output RayJob.

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

   Outputnya mirip dengan hal berikut ini:

   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 -- -------------------------------------
   

Multi-host

1. Buat manifes ray-job.tpu-v6e-multihost.yaml berikut:

   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. Terapkan manifes:

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

3. Verifikasi bahwa v6e-16 RayJob dibuat dan berjalan:

   kubectl get rayjobs v6e-16-job
   

   Outputnya mirip dengan hal berikut ini:

   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. Cetak output v6e-16 RayJob:

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

   Outputnya mirip dengan hal berikut ini:

   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 -- ---------------------------------
   

Melihat RayJob di Dasbor Ray

Verifikasi bahwa GKE membuat layanan RayCluster, dan juga terhubung ke instance RayCluster.

Host tunggal

1. Ambil nama RayCluster yang dihasilkan untuk RayJob:

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

2. Ambil nama layanan head RayCluster:

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

3. Hubungkan ke Dasbor Ray dengan meneruskan port layanan head:

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

4. Buka browser web dan masukkan URL berikut:

   http://localhost:8265/#/jobs
   

5. Lihat status RayJob dan log yang relevan.

Multi-host

1. Ambil nama RayCluster yang dihasilkan untuk RayJob:

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

2. Ambil nama layanan head RayCluster:

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

3. Hubungkan ke Dasbor Ray dengan meneruskan port layanan head:

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

4. Buka browser web dan masukkan URL berikut:

   http://localhost:8265/#/jobs
   

5. Lihat status RayJob dan log yang relevan.

Ray menetapkan resource TPU-{accelerator}-Head untuk mengidentifikasi node pekerja Ray yang sesuai dengan nilai TPU_WORKER_ID=0. Dalam grup TPU multi-host, node Ray dengan TPU_WORKER_ID=0 memiliki TPU-v6e-16-head: 1.0 yang ditetapkan di resource-nya. Variabel lingkungan TPU_WORKER_ID ini ditetapkan oleh webhook GKE yang mengubah untuk KubeRay.

Berhasil: Anda telah berhasil menjalankan RayJob di node pool slice TPU yang menjalankan TPU Trillium. Penyelesaian RayJob menunjukkan bahwa Jax berhasil diinisialisasi dan node TPU terdistribusi diaktifkan dengan benar.

Pembersihan

Setelah Anda menyelesaikan tutorial, untuk mencegah timbulnya biaya yang tidak diinginkan pada akun Anda, hapus RayJob:

Host tunggal

kubectl delete rayjobs v6e-4-job

Multi-host

kubectl delete rayjobs v6e-16-job

Langkah berikutnya

• Pelajari Ray di Kubernetes.
• Pelajari cara Menyajikan vLLM di GKE dengan TPU.
• Pelajari cara Menyajikan SDXL di GKE dengan TPU.
• Pelajari lebih lanjut Tentang TPU di GKE.