Entrenar un modelo con Ray y PyTorch en Google Kubernetes Engine (GKE)

En esta guía se muestra cómo entrenar un modelo en Google Kubernetes Engine (GKE) con Ray, PyTorch y el complemento Ray Operator.

Acerca de Ray

Ray es un framework de computación escalable de código abierto para aplicaciones de IA y aprendizaje automático. Ray Train es un componente de Ray diseñado para el entrenamiento y el ajuste de modelos distribuidos. Puedes usar la API Ray Train para escalar el entrenamiento en varias máquinas e integrarlo con bibliotecas de aprendizaje automático, como PyTorch.

Puedes desplegar tareas de entrenamiento de Ray con el recurso RayCluster o RayJob. Deberías usar un recurso RayJob al implementar trabajos de Ray en producción por los siguientes motivos:

  • El recurso RayJob crea un clúster de Ray efímero que se puede eliminar automáticamente cuando se completa un trabajo.
  • El recurso RayJob admite políticas de reintento para la ejecución de trabajos resilientes.
  • Puedes gestionar los trabajos de Ray con patrones de API de Kubernetes que ya conoces.

Objetivos

Esta guía está dirigida a clientes de IA generativa, usuarios nuevos o actuales de GKE, ingenieros de aprendizaje automático, ingenieros de MLOps (DevOps) o administradores de plataformas que estén interesados en usar las funciones de orquestación de contenedores de Kubernetes para servir modelos con Ray.

  • Crea un clúster de GKE.
  • Crea un clúster de Ray con el recurso personalizado RayCluster.
  • Entrena un modelo con un trabajo de Ray.
  • Despliega un trabajo de Ray con el recurso personalizado RayJob.

Costes

En este documento, se utilizan los siguientes componentes facturables de Google Cloud:

Para generar una estimación de costes basada en el uso previsto, utiliza la calculadora de precios.

Los usuarios nuevos Google Cloud pueden disfrutar de una prueba gratuita.

Cuando termines las tareas que se describen en este documento, puedes evitar que se te siga facturando eliminando los recursos que has creado. Para obtener más información, consulta la sección Limpiar.

Antes de empezar

Cloud Shell tiene preinstalado el software que necesitas para este tutorial, incluido kubectl y la CLI de gcloud. Si no usas Cloud Shell, debes instalar la CLI de gcloud.

  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. Si utilizas un proveedor de identidades (IdP) externo, primero debes iniciar sesión en la CLI de gcloud con tu identidad federada.

  4. Para inicializar gcloud CLI, ejecuta el siguiente comando:

    gcloud init
  5. Create or select a Google Cloud project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.
    • 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.

  6. Verify that billing is enabled for your Google Cloud project.

  7. Enable the GKE API:

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

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

  9. Si utilizas un proveedor de identidades (IdP) externo, primero debes iniciar sesión en la CLI de gcloud con tu identidad federada.

  10. Para inicializar gcloud CLI, ejecuta el siguiente comando:

    gcloud init
  11. Create or select a Google Cloud project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.
    • 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.

  12. Verify that billing is enabled for your Google Cloud project.

  13. Enable the GKE API:

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    gcloud services enable container.googleapis.com
  14. 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 the following:

    • PROJECT_ID: Your project ID.
    • USER_IDENTIFIER: The identifier for your user account. For example, myemail@example.com.
    • ROLE: The IAM role that you grant to your user account.
  15. Prepara tu entorno

    Para preparar tu entorno, sigue estos pasos:

    1. Inicia una sesión de Cloud Shell desde la Google Cloud consolaIcono de activación de Cloud Shell haciendo clic en Activar Cloud Shell en la Google Cloud consola. Se iniciará una sesión en el panel inferior de la consola Google Cloud .

    2. Define las variables de entorno:

      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`
      

      Haz los cambios siguientes:

      • PROJECT_ID: tu Google Cloud ID de proyecto.
      • CLUSTER_VERSION: la versión de GKE que se va a usar. Debe ser 1.30.1 o posterior.
    3. Clona el repositorio de GitHub:

      git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
      
    4. Cambia al directorio de trabajo:

      cd kubernetes-engine-samples/ai-ml/gke-ray/raytrain/pytorch-mnist
      
    5. Crea un entorno virtual de Python:

      python -m venv myenv && \
      source myenv/bin/activate
      
    6. Instala Ray.

    Crear un clúster de GKE

    Crea un clúster de GKE Autopilot o Standard:

    Autopilot

    Crea un clúster de Autopilot:

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

    Estándar

    Crea un clúster estándar:

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

    Desplegar un recurso RayCluster

    Despliega un recurso RayCluster en tu clúster:

    1. Revisa el siguiente archivo de manifiesto:

      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"

      Este manifiesto describe un recurso personalizado de RayCluster.

    2. Aplica el manifiesto a tu clúster de GKE:

      kubectl apply -f ray-cluster.yaml
      
    3. Verifica que el recurso RayCluster esté listo:

      kubectl get raycluster
      

      El resultado debería ser similar al siguiente:

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

      En este resultado, ready en la columna STATUS indica que el recurso RayCluster está listo.

    Conectarse al recurso RayCluster

    Conéctate al recurso RayCluster para enviar un trabajo de Ray.

    1. Verifica que GKE haya creado el servicio RayCluster:

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

      El resultado debería ser similar al siguiente:

      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. Establece una sesión de redirección de puertos al dispositivo Ray:

      kubectl port-forward svc/pytorch-mnist-cluster-head-svc 8265:8265 2>&1 >/dev/null &
      
    3. Verifica que el cliente de Ray se puede conectar al clúster de Ray mediante localhost:

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

      El resultado debería ser similar al siguiente:

      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
      

    Preparar un modelo

    Entrena un modelo de PyTorch con el conjunto de datos Fashion MNIST:

    1. Envía un trabajo de Ray y espera a que se complete:

      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
      

      El resultado debería ser similar al siguiente:

      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. Verifica el estado de la tarea:

      ray job status pytorch-mnist
      

      El resultado debería ser similar al siguiente:

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

      Espera a que Status for job sea COMPLETE. Este proceso puede tardar 15 minutos o más.

    3. Para ver los registros de tareas de Ray, haz lo siguiente:

      ray job logs pytorch-mnist
      

      El resultado debería ser similar al siguiente:

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

    Desplegar un RayJob

    El recurso personalizado RayJob gestiona el ciclo de vida de un recurso RayCluster durante la ejecución de un único trabajo de Ray.

    1. Revisa el siguiente archivo de manifiesto:

      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"

      Este manifiesto describe un recurso personalizado de RayJob.

    2. Aplica el manifiesto a tu clúster de GKE:

      kubectl apply -f ray-job.yaml
      
    3. Comprueba que el recurso RayJob se está ejecutando:

      kubectl get rayjob
      

      El resultado debería ser similar al siguiente:

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

      En este resultado, la columna DEPLOYMENT STATUS indica que el recurso RayJob es Running.

    4. Consulta el estado del recurso RayJob:

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

      El resultado debería ser similar al siguiente:

      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. Verifica que el trabajo de Ray se haya completado:

      kubectl get rayjob
      

      El resultado debería ser similar al siguiente:

      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
      

      En este resultado, la columna DEPLOYMENT STATUS indica que el recurso RayJob es Complete.

    Limpieza

    Eliminar el proyecto

      Delete a Google Cloud project:

      gcloud projects delete PROJECT_ID

    Eliminar recursos concretos

    Si has usado un proyecto que ya existía y no quieres eliminarlo, puedes eliminar los recursos concretos. Para eliminar el clúster, escribe lo siguiente:

    gcloud container clusters delete ${CLUSTER_NAME}
    

    Siguientes pasos