Cette page a été traduite par l'API Cloud Translation.

Configurer KubeRay avec TPU Trillium

Ce tutoriel explique comment configurer KubeRay avec TPU Trillium sur Google Kubernetes Engine (GKE). Apprenez à configurer les configurations TPU à hôte unique et multi-hôte, y compris les variables d'environnement et les spécifications de pod nécessaires pour TPU Trillium.

Ce tutoriel s'adresse aux administrateurs et opérateurs de plate-forme, ainsi qu'aux spécialistes des données et de l'IA qui souhaitent apprendre à configurer l'initialisation de TPU Trillium avec KubeRay pour les pools de nœuds à hôte unique et multi-hôtes. Ce tutoriel montre comment exécuter un script avec Jax qui vérifie l'initialisation réussie de la TPU. Ce tutoriel ne déploie pas de modèle.

Avant de configurer KubeRay dans GKE, assurez-vous de bien connaître les définitions et la terminologie de Ray dans GKE.

Présentation

Ce tutoriel explique comment exécuter un script Python avec Jax qui vérifie que l'initialisation de TPU Trillium avec KubeRay a réussi. Jax est une bibliothèque de calcul numérique hautes performances qui prend en charge les charges de travail de machine learning. KubeRay est un opérateur Kubernetes qui fournit une méthode unifiée pour déployer, gérer et surveiller les applications Ray sur Kubernetes.

Les TPU Trillium (v6e) nécessitent des variables d'environnement et des spécifications de pod spécifiques, qui diffèrent des générations de TPU précédentes. Ce tutoriel fournit les configurations nécessaires pour déployer correctement une charge de travail avec KubeRay sur des TPU Trillium.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

Activez l'API Google Kubernetes Engine.

Activer l'API Google Kubernetes Engine

Si vous souhaitez utiliser Google Cloud CLI pour cette tâche, installez puis initialisez gcloud CLI. Si vous avez précédemment installé la gcloud CLI, obtenez la dernière version en exécutant la commande gcloud components update.
Remarque : Pour les installations de la gcloud CLI existantes, veillez à définir la propriété compute/region. Si vous utilisez principalement des clusters zonaux, définissez plutôt compute/zone. En définissant un emplacement par défaut, vous pouvez éviter les erreurs gcloud CLI de ce type : One of [--zone, --region] must be supplied: Please specify location. Vous devrez peut-être spécifier l'emplacement dans certaines commandes si celui de votre cluster diffère de l'emplacement par défaut que vous avez défini.

Assurez-vous d'avoir installé la CLI Ray (version 2.37.0).

Activer Cloud Shell

Cloud Shell est préinstallé avec les outils de ligne de commande gcloud, helm et kubectl utilisés dans ce tutoriel.

Accédez à la consoleGoogle Cloud .
En haut de la fenêtre de la console Google Cloud , cliquez sur le bouton Activer Cloud Shell .

Une session Cloud Shell s'ouvre dans un nouveau cadre dans la console Google Cloud et affiche une invite de ligne de commande.

Créer un cluster GKE et un pool de nœuds

Vous pouvez configurer KubeRay sur des TPU dans un cluster GKE Autopilot ou GKE Standard. Nous vous recommandons d'utiliser un cluster GKE Autopilot pour une expérience Kubernetes entièrement gérée. Pour choisir le mode de fonctionnement GKE le mieux adapté à vos charges de travail, consultez À propos des modes de fonctionnement GKE.

Autopilot

Dans Cloud Shell, exécutez la commande suivante :
```
gcloud container clusters create-auto CLUSTER_NAME \
    --enable-ray-operator \
    --release-channel=rapid \
    --location=LOCATION
```
Remplacez les éléments suivants :
- CLUSTER_NAME : nom du nouveau cluster
- LOCATION : région dans laquelle votre capacité de TPU Trillium est disponible. Pour en savoir plus, consultez la section Disponibilité des TPU dans GKE.
GKE crée un cluster Autopilot avec le module complémentaire de l'opérateur Ray activé. Le module complémentaire installe automatiquement le webhook Ray TPU dans le plan de contrôle du cluster.

Pour communiquer avec votre cluster, configurez kubectl :

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

Standard

Dans Cloud Shell, créez un cluster Standard qui active le module complémentaire de l'opérateur Ray en exécutant la commande suivante :
```
gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --addons=RayOperator \
  --cluster-version=1.33 \
  --machine-type=n1-standard-16
```
Remplacez les éléments suivants :
- CLUSTER_NAME : nom du nouveau cluster
- LOCATION : région dans laquelle votre capacité de TPU Trillium est disponible. Pour en savoir plus, consultez la section Disponibilité des TPU dans GKE.
La création du cluster peut prendre plusieurs minutes.

Pour communiquer avec votre cluster, configurez kubectl :

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

Vous pouvez créer un pool de nœuds de tranche TPU à hôte unique ou multi-hôte :

Un seul hôte

Dans Cloud Shell, exécutez la commande suivante :

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

Hôte multiple

Dans Cloud Shell, exécutez la commande suivante :

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

Exécuter une ressource personnalisée RayJob

En définissant un fichier manifeste RayJob, vous demandez à KubeRay d'effectuer les opérations suivantes :

Créez un RayCluster : la spécification RayJob inclut un rayClusterSpec qui définit la configuration du cluster Ray (groupes de nœuds principaux et de nœuds de calcul) de votre choix.
Exécuter un job spécifique : le champ entrypoint de RayJob spécifie la commande ou le script à exécuter dans le cluster Ray créé. Dans ce tutoriel, entrypoint est un script Python (tpu_list_devices.py) conçu pour vérifier l'initialisation de la TPU Trillium.

Pour créer une ressource personnalisée RayJob, procédez comme suit :

Un seul hôte

Créez le fichier manifeste ray-job.tpu-v6e-singlehost.yaml suivant :

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

Appliquez le fichier manifeste :

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

Vérifiez que RayJob a été créé et est en cours d'exécution :

kubectl get rayjobs v6e-4-job

Le résultat ressemble à ce qui suit :

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

Imprimez la sortie du RayJob.

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

Le résultat ressemble à ce qui suit :

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

Hôte multiple

Créez le fichier manifeste ray-job.tpu-v6e-multihost.yaml suivant :

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

Appliquez le fichier manifeste :

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

Vérifiez que le RayJob v6e-16 est créé et en cours d'exécution :

kubectl get rayjobs v6e-16-job

Le résultat ressemble à ce qui suit :

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

Imprimez le résultat du RayJob v6e-16 :

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

Le résultat ressemble à ce qui suit :

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

Afficher le RayJob dans le tableau de bord Ray

Vérifiez que GKE a créé le service RayCluster, puis connectez-vous à l'instance RayCluster.

Un seul hôte

Récupérez le nom du RayCluster généré pour le RayJob :

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

Récupérez le nom du service principal RayCluster :

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

Connectez-vous au tableau de bord Ray en transférant le port du service principal :
```
kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
```
Ouvrez un navigateur Web et saisissez l'URL suivante :
```
http://localhost:8265/#/jobs
```
Affichez l'état de RayJob et les journaux pertinents.

Hôte multiple

Récupérez le nom du RayCluster généré pour le RayJob :

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

Récupérez le nom du service principal RayCluster :

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

Connectez-vous au tableau de bord Ray en transférant le port du service principal :
```
kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
```
Ouvrez un navigateur Web et saisissez l'URL suivante :
```
http://localhost:8265/#/jobs
```
Affichez l'état de RayJob et les journaux pertinents.

Ray définit une ressource TPU-{accelerator}-Head pour identifier le nœud de calcul Ray qui correspond à la valeur TPU_WORKER_ID=0. Dans le groupe de TPU à hôtes multiples, le nœud Ray avec TPU_WORKER_ID=0 a TPU-v6e-16-head: 1.0 défini dans ses ressources. Cette variable d'environnement TPU_WORKER_ID est définie par un webhook GKE de mutation pour KubeRay.

Effectuer un nettoyage

Une fois le tutoriel terminé, supprimez le RayJob pour éviter que des frais inutiles ne vous soient facturés :

Un seul hôte

kubectl delete rayjobs v6e-4-job

Hôte multiple

kubectl delete rayjobs v6e-16-job

Étapes suivantes

Découvrez Ray sur Kubernetes.
Découvrez comment diffuser vLLM sur GKE avec des TPU.
Découvrez comment diffuser SDXL sur GKE avec des TPU.
En savoir plus sur les TPU dans GKE

Configurer KubeRay avec TPU Trillium Restez organisé à l'aide des collections Enregistrez et classez les contenus selon vos préférences.

Présentation

Avant de commencer

Activer Cloud Shell

Créer un cluster GKE et un pool de nœuds

Autopilot

Standard

Un seul hôte

Hôte multiple

Exécuter une ressource personnalisée RayJob

Un seul hôte

Hôte multiple

Afficher le RayJob dans le tableau de bord Ray

Un seul hôte

Hôte multiple

Effectuer un nettoyage

Un seul hôte

Hôte multiple

Étapes suivantes

Configurer KubeRay avec TPU Trillium