Déployer une application Ray Serve avec un modèle Stable Diffusion sur Google Kubernetes Engine (GKE) avec des TPU

Ce guide explique comment déployer et diffuser un modèle Stable Diffusion sur Google Kubernetes Engine (GKE) à l'aide de TPU, de Ray Serve et du Module complémentaire Ray Operator.

Ce guide est destiné aux clients d'IA générative, aux utilisateurs nouveaux ou existants de GKE, aux ingénieurs en ML, aux ingénieurs MLOps (DevOps) ou aux administrateurs de plate-forme qui s'intéressent à l'utilisation des fonctionnalités d'orchestration de conteneurs Kubernetes pour diffuser des modèles utilisant Ray.

À propos de Ray et Ray Serve

Ray est un framework de calcul évolutif Open Source pour les applications d'IA/ML. Ray Serve est une bibliothèque de diffusion de modèles pour Ray, utilisée pour le scaling et la diffusion de modèles dans un environnement distribué. Pour en savoir plus, consultez la section Ray Serve dans la documentation Ray.

À propos des TPU

Les TPU (Tensor Processing Units) sont des accélérateurs matériels spécialisés conçus pour accélérer considérablement l'entraînement et l'inférence des modèles de machine learning à grande échelle. L'utilisation de Ray avec des TPU vous permet de procéder au scaling des applications de ML hautes performances en toute simplicité. Pour en savoir plus sur les TPU, consultez la Présentation de Cloud TPU dans la documentation de Cloud TPU.

À propos du webhook d'initialisation TPU KubeRay

Au sein du module complémentaire Ray Operator, GKE fournit des webhooks de validation et de mutation qui gèrent la planification des pods TPU et certaines variables d'environnement TPU requises par des frameworks tels que JAX pour l'initialisation du conteneur. Le webhook de TPU KubeRay modifie les pods ayant l'étiquette app.kubernetes.io/name: kuberay qui demandent des TPU avec les propriétés suivantes :

  • TPU_WORKER_ID : Entier unique pour chaque pod de nœud de calcul dans la tranche TPU.
  • TPU_WORKER_HOSTNAMES : liste de noms d'hôtes DNS pour tous les travailleurs TPU qui doivent communiquer entre eux au sein de la tranche. Cette variable n'est injectée que pour les pods TPU dans un groupe multi-hôtes.
  • replicaIndex : étiquette de pod qui contient un identifiant unique pour l'instance répliquée du groupe de nœuds de calcul auquel le pod appartient. Cela est utile pour les groupes de nœuds de calcul multi-hôtes, où plusieurs pods de nœuds de calcul peuvent appartenir à la même instance répliquée, et est utilisé par Ray pour permettre l'autoscaling multi-hôte.
  • TPU_NAME : chaîne représentant la tranche de pod TPU GKE à laquelle ce pod appartient, définie sur la même valeur que l'étiquette replicaIndex.
  • podAffinity : garantit que GKE planifie les pods TPU avec les étiquettes replicaIndex correspondantes sur le même pool de nœuds. Cela permet à GKE d'effectuer un scaling atomique des TPU multi-hôtes par pool de nœuds, plutôt que par nœuds individuels.

Objectifs

  • Créer un cluster GKE avec un pool de nœuds TPU.
  • Déployer un cluster Ray avec des TPU.
  • Déployer une ressource personnalisée RayService.
  • Interagir avec le serveur de modèle Stable Diffusion.

Coûts

Dans ce document, vous utilisez les composants facturables suivants de Google Cloud :

Obtenez une estimation des coûts en fonction de votre utilisation prévue à l'aide du simulateur de coût. Les nouveaux utilisateurs de Google Cloud peuvent bénéficier d'un essai gratuit.

Une fois que vous avez terminé les tâches décrites dans ce document, vous pouvez éviter de continuer à payer des frais en supprimant les ressources que vous avez créées. Pour en savoir plus, consultez la section Effectuer un nettoyage.

Avant de commencer

Cloud Shell est préinstallé avec les logiciels dont vous avez besoin dans ce tutoriel, y compris kubectl et gcloud CLI. Si vous n'utilisez pas Cloud Shell, installez gcloud CLI.

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  2. Install the Google Cloud CLI.

  3. To initialize the gcloud CLI, run the following command: 

    gcloud init

  4. Create or select a Google Cloud project.

    • Create a Google Cloud project:

      gcloud projects create PROJECT_ID

      Replace PROJECT_ID with a name for the Google Cloud project you are creating.

    • Select the Google Cloud project that you created:

      gcloud config set project PROJECT_ID

      Replace PROJECT_ID with your Google Cloud project name.

  5. Make sure that billing is enabled for your Google Cloud project.

  6. Enable the GKE API: 

    gcloud services enable container.googleapis.com
  7. Install the Google Cloud CLI.

  8. To initialize the gcloud CLI, run the following command: 

    gcloud init

  9. Create or select a Google Cloud project.

    • Create a Google Cloud project:

      gcloud projects create PROJECT_ID

      Replace PROJECT_ID with a name for the Google Cloud project you are creating.

    • Select the Google Cloud project that you created:

      gcloud config set project PROJECT_ID

      Replace PROJECT_ID with your Google Cloud project name.

  10. Make sure that billing is enabled for your Google Cloud project.

  11. Enable the GKE API: 

    gcloud services enable container.googleapis.com

  12. Grant roles to your user account. Run the following command once for each of the following IAM roles: roles/container.clusterAdmin, roles/container.admin 

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE
    • Replace PROJECT_ID with your project ID.

    • Replace USER_IDENTIFIER with the identifier for your user account. For example, user:myemail@example.com.

    • Replace ROLE with each individual role.

Garantir un quota suffisant

Assurez-vous que votre projet Google Cloud dispose d'un quota TPU suffisant dans votre région ou zone Compute Engine. Pour en savoir plus, consultez la section Garantir un quota de TPU et GKE suffisant dans la documentation Cloud TPU. Vous devrez peut-être également augmenter vos quotas pour les éléments suivants :

  • Disque persistant SSD (Go)
  • Adresses IP en cours d'utilisation

Préparer votre environnement

Pour préparer votre environnement, procédez comme suit :

  1. Lancez une session Cloud Shell depuis la console Google Cloud en cliquant sur Icône d'activation Cloud Shell Activer Cloud Shell dans la console Google Cloud. Une session s'ouvre dans le volet inférieur de la console Google Cloud.

  2. Définissez les variables d'environnement :

    export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=ray-cluster
export COMPUTE_REGION=us-central2-b
export CLUSTER_VERSION=CLUSTER_VERSION

    Remplacez les éléments suivants :

    • PROJECT_ID : ID de votre projet Google Cloud.
    • CLUSTER_VERSION : version de GKE à utiliser. Doit être 1.30.1 ou une version ultérieure.

  3. Clonez le dépôt GitHub.

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples

  4. Accédez au répertoire de travail :

    cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/stable-diffusion

Créer un cluster avec un pool de nœuds TPU

Créez un cluster GKE Standard avec un pool de nœuds TPU :

  1. Créez un cluster en mode Standard avec l'opérateur Ray activé :

    gcloud container clusters create ${CLUSTER_NAME} \
    --addons=RayOperator \
    --machine-type=n1-standard-8 \
    --cluster-version=${CLUSTER_VERSION} \
    --location=${COMPUTE_REGION}

  2. Créez un pool de nœuds TPU à hôte unique :

    gcloud container node-pools create tpu-pool \
    --location=${COMPUTE_REGION} \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct4p-hightpu-4t \
    --num-nodes=1 \
    --tpu-topology=2x2x1

Pour utiliser des TPU avec le mode Standard, vous devez sélectionner les éléments suivants :

  • Un emplacement Compute Engine avec de la capacité pour les accélérateurs TPU
  • Un type de machine compatible pour le TPU
  • La topologie physique de la tranche de pod TPU

Configurer une ressource RayCluster avec des TPU

Configurez votre fichier manifeste RayCluster pour préparer votre charge de travail TPU :

Configurer le nodeSelector TPU

GKE utilise les nodeSelectors Kubernetes pour s'assurer que les charges de travail TPU sont planifiées selon la topologie TPU et l'accélérateur appropriés. Pour en savoir plus sur la sélection des nodeSelectors TPU, consultez la section Déployer des charges de travail TPU dans GKE Standard.

Mettez à jour le fichier manifeste ray-cluster.yaml pour planifier votre pod sur une tranche de pod TPU v4 avec une topologie de 2x2x1 :

nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
  cloud.google.com/gke-tpu-topology: 2x2x1

Configurer une ressource de conteneur TPU

Pour utiliser un accélérateur TPU, vous devez spécifier le nombre de puces TPU que GKE doit allouer à chaque pod en configurant les éléments limits et requests de ressource google.com/tpu dans le champ du conteneur TPU de votre fichier manifeste RayCluster workerGroupSpecs.

Mettez à jour le fichier manifeste ray-cluster.yaml avec les limites et demandes de ressources :

resources:
  limits:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"
   requests:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"

Configurer le numOfHosts du groupe de nœuds de calcul

KubeRay v1.1.0 ajoute un champ numOfHosts à la ressource personnalisée RayCluster, qui spécifie le nombre d'hôtes TPU à créer par instance répliquée du groupe de nœuds de calcul. Pour les groupes de nœuds de calcul à hôtes multiples, les instances répliquées sont traitées comme des tranches de pods plutôt que comme des nœuds de calcul individuels, les nœuds de calcul numOfHosts étant créés par instance répliquée.

Mettez à jour le fichier manifeste ray-cluster.yaml avec les éléments suivants :

workerGroupSpecs:
  # Several lines omitted
  numOfHosts: 1 # the number of "hosts" or workers per replica

Créer une ressource personnalisée RayService

Créez une ressource personnalisée RayService :

  1. Examinez le fichier manifeste suivant :

    apiVersion: ray.io/v1
kind: RayService
metadata:
  name: stable-diffusion-tpu
spec:
  serveConfigV2: |
    applications:
      - name: stable_diffusion
        import_path: ai-ml.gke-ray.rayserve.stable-diffusion.stable_diffusion_tpu:deployment
        runtime_env:
          working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
          pip:
            - diffusers==0.7.2
            - flax
            - jax[tpu]==0.4.11
            - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
            - fastapi
  rayClusterConfig:
    rayVersion: '2.9.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
          - name: ray-head
            image: rayproject/ray-ml:2.9.0-py310
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "2"
                memory: "8G"
              requests:
                cpu: "2"
                memory: "8G"
    workerGroupSpecs:
    - replicas: 1
      minReplicas: 1
      maxReplicas: 10
      numOfHosts: 1
      groupName: tpu-group
      rayStartParams: {}
      template:
        spec:
          containers:
          - name: ray-worker
            image: rayproject/ray-ml:2.9.0-py310
            resources:
              limits:
                cpu: "100"
                ephemeral-storage: 20Gi
                google.com/tpu: "4"
                memory: 200G
              requests:
                cpu: "100"
                ephemeral-storage: 20Gi
                google.com/tpu: "4"
                memory: 200G
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
            cloud.google.com/gke-tpu-topology: 2x2x1

    Ce fichier manifeste décrit une ressource personnalisée RayService qui crée une ressource RayCluster avec un nœud principal et un groupe de nœuds de calcul TPU avec une topologie 2x2x1, ce qui signifie que chaque nœud de calcul aura quatre puces TPU v4.

    Le nœud TPU appartient à une seule tranche de pod TPU v4 avec une topologie 2x2x1. Pour créer un groupe de nœuds de calcul multi-hôtes, remplacez les valeurs gke-tpu nodeSelector, les limites et les requêtes de conteneur google.com/tpu, ainsi que les valeurs numOfHosts par votre configuration multi-hôte. Pour en savoir plus sur les topologies multi-hôtes TPU, consultez la section Architecture du système dans la documentation Cloud TPU.

  2. Appliquez le fichier manifeste à votre cluster :

    kubectl apply -f ray-service-tpu.yaml

  3. Vérifiez que la ressource RayService est en cours d'exécution :

    kubectl get rayservices

    Le résultat ressemble à ce qui suit :

    NAME                   SERVICE STATUS   NUM SERVE ENDPOINTS
stable-diffusion-tpu   Running          2

    Dans ce résultat, Running dans la colonne SERVICE STATUS indique que la ressource RayService est prête.

(Facultatif) Afficher le tableau de bord Ray

Vous pouvez afficher votre déploiement Ray Serve et les journaux pertinents à partir du tableau de bord Ray.

  1. Établissez une session de transfert de port vers le tableau de bord Ray à partir du service principal Ray :

    kubectl port-forward svc/stable-diffusion-tpu-head-svc 8265:8265

  2. Dans un navigateur Web, accédez à http://localhost:8265/.

  3. Cliquez sur l'onglet Diffuser.

Envoyer des requêtes au serveur de modèles

  1. Établissez une session de transfert de port vers le point de terminaison Serve à partir du service principal Ray :

    kubectl port-forward svc/stable-diffusion-tpu-serve-svc 8000

  2. Ouvrez une nouvelle session Cloud Shell.

  3. Envoyez une requête texte-vers-image au serveur de modèle Stable Diffusion :

    python stable_diffusion_tpu_req.py  --save_pictures

    Les résultats de l'inférence de diffusion stable sont enregistrés dans un fichier nommé diffusion_results.png.

    Image générée par Stable Diffusion comprenant 8 sections : une chaise verte, un homme debout à l'extérieur d'une maison, un robot dans la rue, une famille assise à une table, un chien marchant dans un parc, un dragon volant, un portrait d'ours de style japonais et une cascade.

Effectuer un nettoyage

Supprimer le projet

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

Supprimer des ressources individuelles

Pour supprimer le cluster, saisissez ce qui suit :

gcloud container clusters delete ${CLUSTER_NAME}

Étape suivante