Entraîner un modèle avec PyTorch, Ray et Google Kubernetes Engine (GKE) sur des GPU

Ce guide explique comment entraîner un modèle sur Google Kubernetes Engine (GKE) à l'aide de Ray, PyTorch et du module complémentaire Ray Operator.

À propos de Ray

Ray est un framework de calcul évolutif Open Source pour les applications d'IA/ML. L'Ray Train est un composant de Ray conçu pour l'entraînement et le réglage de modèles distribués. Vous pouvez utiliser l'API Ray Train pour mettre à l'échelle l'entraînement sur plusieurs machines et pour l'intégrer à des bibliothèques de machine learning telles que PyTorch.

Vous pouvez déployer des jobs d'entraînement Ray à l'aide de la ressource RayCluster ou RayJob. Vous devez utiliser une ressource RayJob lorsque vous déployez des jobs Ray en production pour les raisons suivantes :

  • La ressource RayJob crée un cluster Ray éphémère qui peut être supprimé automatiquement lorsqu'un job est terminé.
  • La ressource RayJob est compatible avec les stratégies de nouvelle tentative pour une exécution de job résiliente.
  • Vous pouvez gérer les jobs Ray à l'aide de modèles d'API Kubernetes connus.

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 consoleGoogle 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-central1
export COMPUTE_ZONE=us-central1-c
export CLUSTER_VERSION=CLUSTER_VERSION
export TUTORIAL_HOME=`pwd`

    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/raytrain/pytorch-mnist

  5. Créez un environnement virtuel Python :

    python -m venv myenv && \
source myenv/bin/activate

  6. Installez Ray.

Créer un cluster GKE

Créez un cluster GKE Autopilot ou Standard :

Autopilot

Créez un cluster Autopilot :

gcloud container clusters create-auto ${CLUSTER_NAME}  \
    --enable-ray-operator \
    --cluster-version=${CLUSTER_VERSION} \
    --location=${COMPUTE_REGION}

Standard

Créez un cluster standard :

gcloud container clusters create ${CLUSTER_NAME} \
    --addons=RayOperator \
    --cluster-version=${CLUSTER_VERSION}  \
    --machine-type=e2-standard-8 \
    --location=${COMPUTE_ZONE} \
    --num-nodes=4

Déployer une ressource RayCluster

Déployez une ressource RayCluster sur votre cluster :

  1. Examinez le fichier manifeste suivant :

    apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: pytorch-mnist-cluster
spec:
  rayVersion: '2.37.0'
  headGroupSpec:
    rayStartParams:
      dashboard-host: '0.0.0.0'
    template:
      metadata:
      spec:
        containers:
        - name: ray-head
          image: rayproject/ray:2.37.0
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "2"
              ephemeral-storage: "9Gi"
              memory: "4Gi"
            requests:
              cpu: "2"
              ephemeral-storage: "9Gi"
              memory: "4Gi"
  workerGroupSpecs:
  - replicas: 4
    minReplicas: 1
    maxReplicas: 5
    groupName: worker-group
    rayStartParams: {}
    template:
      spec:
        containers:
        - name: ray-worker
          image: rayproject/ray:2.37.0
          resources:
            limits:
              cpu: "4"
              ephemeral-storage: "9Gi"
              memory: "8Gi"
            requests:
              cpu: "4"
              ephemeral-storage: "9Gi"
              memory: "8Gi"

    Ce fichier manifeste décrit une ressource personnalisée RayCluster.

  2. Appliquez le fichier manifeste à votre cluster GKE :

    kubectl apply -f ray-cluster.yaml

  3. Vérifiez que la ressource RayCluster est prête :

    kubectl get raycluster

    Le résultat ressemble à ce qui suit :

    NAME                    DESIRED WORKERS   AVAILABLE WORKERS   CPUS   MEMORY   GPUS   STATUS   AGE
pytorch-mnist-cluster   2                 2                   6      20Gi     0      ready    63s

    Dans ce résultat, ready dans la colonne STATUS indique que la ressource RayCluster est prête.

Se connecter à la ressource RayCluster

Connectez-vous à la ressource RayCluster pour envoyer un job Ray.

  1. Vérifiez que GKE a créé le service RayCluster :

    kubectl get svc pytorch-mnist-cluster-head-svc

    Le résultat ressemble à ce qui suit :

    NAME                             TYPE        CLUSTER-IP       EXTERNAL-IP   PORT(S)                                AGE
pytorch-mnist-cluster-head-svc   ClusterIP   34.118.238.247   <none>        10001/TCP,8265/TCP,6379/TCP,8080/TCP   109s

  2. Établissez une session de transfert de port vers le nœud principal Ray :

    kubectl port-forward svc/pytorch-mnist-cluster-head-svc 8265:8265 2>&1 >/dev/null &

  3. Vérifiez que le client Ray peut se connecter au cluster Ray à l'aide de localhost :

    ray list nodes --address http://localhost:8265

    Le résultat ressemble à ce qui suit :

    Stats:
------------------------------
Total: 3

Table:
------------------------------
    NODE_ID                                                   NODE_IP     IS_HEAD_NODE    STATE    NODE_NAME    RESOURCES_TOTAL                 LABELS
0  1d07447d7d124db641052a3443ed882f913510dbe866719ac36667d2  10.28.1.21  False           ALIVE    10.28.1.21   CPU: 2.0                        ray.io/node_id: 1d07447d7d124db641052a3443ed882f913510dbe866719ac36667d2
# Several lines of output omitted

Entraîner un modèle

Entraînez un modèle PyTorch à l'aide de l'ensemble de données Fashion MNIST :

  1. Envoyez un job Ray et attendez qu'il se termine :

    ray job submit --submission-id pytorch-mnist-job --working-dir . --runtime-env-json='{"pip": ["torch", "torchvision"], "excludes": ["myenv"]}' --address http://localhost:8265 -- python train.py

    Le résultat ressemble à ce qui suit :

    Job submission server address: http://localhost:8265

--------------------------------------------
Job 'pytorch-mnist-job' submitted successfully
--------------------------------------------

Next steps
  Query the logs of the job:
    ray job logs pytorch-mnist-job
  Query the status of the job:
    ray job status pytorch-mnist-job
  Request the job to be stopped:
    ray job stop pytorch-mnist-job

Handling connection for 8265
Tailing logs until the job exits (disable with --no-wait):
...
...

  2. Vérifiez l'état du job :

    ray job status pytorch-mnist

    Le résultat ressemble à ce qui suit :

    Job submission server address: http://localhost:8265
Status for job 'pytorch-mnist-job': RUNNING
Status message: Job is currently running.

    Attendez que Status for job soit défini sur COMPLETE. Cette opération peut prendre 15 minutes, voire plus.

  3. Affichez les journaux des jobs Ray :

    ray job logs pytorch-mnist

    Le résultat ressemble à ce qui suit :

    Training started with configuration:
╭─────────────────────────────────────────────────╮
│ Training config                                  │
├──────────────────────────────────────────────────┤
│ train_loop_config/batch_size_per_worker       8  │
│ train_loop_config/epochs                     10  │
│ train_loop_config/lr                      0.001  │
╰─────────────────────────────────────────────────╯

# Several lines omitted

Training finished iteration 10 at 2024-06-19 08:29:36. Total running time: 9min 18s
╭───────────────────────────────╮
│ Training result                │
├────────────────────────────────┤
│ checkpoint_dir_name            │
│ time_this_iter_s      25.7394  │
│ time_total_s          351.233  │
│ training_iteration         10  │
│ accuracy               0.8656  │
│ loss                  0.37827  │
╰───────────────────────────────╯

# Several lines omitted
-------------------------------
Job 'pytorch-mnist' succeeded
-------------------------------

Déployer une ressource RayJob

La ressource personnalisée RayJob gère le cycle de vie d'une ressource RayCluster lors de l'exécution d'un seul job Ray.

  1. Examinez le fichier manifeste suivant :

    apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: pytorch-mnist-job
spec:
  shutdownAfterJobFinishes: true
  entrypoint: python ai-ml/gke-ray/raytrain/pytorch-mnist/train.py
  runtimeEnvYAML: |
    pip:
      - torch
      - torchvision
    working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/main.zip"
    env_vars:
      NUM_WORKERS: "4"
      CPUS_PER_WORKER: "2"
  rayClusterSpec:
    rayVersion: '2.37.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
            - name: ray-head
              image: rayproject/ray:2.37.0
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "2"
                  ephemeral-storage: "9Gi"
                  memory: "4Gi"
                requests:
                  cpu: "2"
                  ephemeral-storage: "9Gi"
                  memory: "4Gi"
    workerGroupSpecs:
      - replicas: 4
        minReplicas: 1
        maxReplicas: 5
        groupName: small-group
        rayStartParams: {}
        template:
          spec:
            containers:
              - name: ray-worker
                image: rayproject/ray:2.37.0
                resources:
                  limits:
                    cpu: "4"
                    ephemeral-storage: "9Gi"
                    memory: "8Gi"
                  requests:
                    cpu: "4"
                    ephemeral-storage: "9Gi"
                    memory: "8Gi"

    Ce fichier manifeste décrit une ressource personnalisée RayJob.

  2. Appliquez le fichier manifeste à votre cluster GKE :

    kubectl apply -f ray-job.yaml

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

    kubectl get rayjob

    Le résultat ressemble à ce qui suit :

    NAME                JOB STATUS   DEPLOYMENT STATUS   START TIME             END TIME   AGE
pytorch-mnist-job   RUNNING      Running             2024-06-19T15:43:32Z              2m29s

    Dans ce résultat, la colonne DEPLOYMENT STATUS indique que la ressource RayJob est Running.

  4. Affichez l'état de la ressource RayJob :

    kubectl logs -f -l job-name=pytorch-mnist-job

    Le résultat ressemble à ce qui suit :

    Training started with configuration:
╭─────────────────────────────────────────────────╮
│ Training config                                  │
├──────────────────────────────────────────────────┤
│ train_loop_config/batch_size_per_worker       8  │
│ train_loop_config/epochs                     10  │
│ train_loop_config/lr                      0.001  │
╰─────────────────────────────────────────────────╯

# Several lines omitted

Training finished iteration 10 at 2024-06-19 08:29:36. Total running time: 9min 18s
╭───────────────────────────────╮
│ Training result                │
├────────────────────────────────┤
│ checkpoint_dir_name            │
│ time_this_iter_s      25.7394  │
│ time_total_s          351.233  │
│ training_iteration         10  │
│ accuracy               0.8656  │
│ loss                  0.37827  │
╰───────────────────────────────╯

# Several lines omitted
-------------------------------
Job 'pytorch-mnist' succeeded
-------------------------------

  5. Vérifiez que le job Ray est terminé :

    kubectl get rayjob

    Le résultat ressemble à ce qui suit :

    NAME                JOB STATUS   DEPLOYMENT STATUS   START TIME             END TIME               AGE
pytorch-mnist-job   SUCCEEDED    Complete            2024-06-19T15:43:32Z   2024-06-19T15:51:12Z   9m6s

    Dans ce résultat, la colonne DEPLOYMENT STATUS indique que la ressource RayJob est Complete.