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

Autopilot Estándar

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

Información acerca de Ray

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

Puedes implementar trabajos de entrenamiento de Ray con el recurso de RayCluster o RayJob. Debes usar un recurso de RayJob cuando implementes trabajos de Ray en producción por los siguientes motivos

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

Objetivos

Esta guía está dirigida a clientes de IA generativa, usuarios nuevos o existentes de GKE, ingenieros de AA, ingenieros de MLOps (DevOps) o administradores de plataformas que estén interesados en el uso de las capacidades de organización de contenedores de Kubernetes para entregar 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.
Implementa un trabajo de Ray con el recurso personalizado RayJob.

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

Para generar una estimación de costos en función del uso previsto, usa la calculadora de precios. Es posible que los usuarios nuevos de Google Cloud califiquen para obtener una prueba gratuita.

Cuando finalices las tareas que se describen en este documento, puedes borrar los recursos que creaste para evitar que continúe la facturación. Para obtener más información, consulta Cómo realizar una limpieza.

Antes de comenzar

Cloud Shell tiene preinstalado el software que necesitas para este instructivo, incluidos kubectl y gcloud CLI. Si no usas Cloud Shell, debes instalar gcloud CLI.

Accede a tu cuenta de Google Cloud. Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.

Install the Google Cloud CLI.

To initialize the gcloud CLI, run the following command:

gcloud init

Note: If you installed the gcloud CLI previously, make sure you have the latest version by running

gcloud components
      update

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.

Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.

Enable the GKE API:

gcloud services enable container.googleapis.com

Install the Google Cloud CLI.

To initialize the gcloud CLI, run the following command:

gcloud init

Note: If you installed the gcloud CLI previously, make sure you have the latest version by running

gcloud components
      update

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.

Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.

Enable the GKE API:

gcloud services enable container.googleapis.com

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_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.
Instala Ray.

Prepare el entorno

Para preparar tu entorno, sigue estos pasos:

Para iniciar una sesión de Cloud Shell desde la consola de Google Cloud, haz clic en Activar Cloud Shell en la consola de Google Cloud. Esto inicia una sesión en el panel inferior de la consola de Google Cloud.

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

Reemplaza lo siguiente:

PROJECT_ID: El ID del proyecto de Google Cloud.
CLUSTER_VERSION: la versión de GKE que se usará. Debe ser 1.30.1 o una versión posterior.

Clona el repositorio de GitHub:

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

Cambia al directorio de trabajo:

cd kubernetes-engine-samples/ai-ml/gke-ray/raytrain/pytorch-mnist

Cree un clúster de GKE

Crea un cluster 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

Implementa un recurso de RayCluster

Implementa un recurso de RayCluster en tu clúster:

Revisa el siguiente manifiesto:

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

En este manifiesto, se describe un recurso personalizado de RayCluster.

Aplica el manifiesto a tu clúster de GKE:
```
kubectl apply -f ray-cluster.yaml
```

Verifica que el recurso de RayCluster esté listo:

kubectl get raycluster

El resultado es similar a este:

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 de RayCluster está listo.

Conéctate al recurso de RayCluster

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

Verifica que GKE haya creado el Service de RayCluster:

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

El resultado es similar a este:

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

Establece una sesión de redirección de puertos al encabezado de Ray:

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

Verifica que el cliente de Ray pueda conectarse al clúster de Ray con localhost:

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

El resultado es similar a este:

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

Entrenar un modelo

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

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"]}' --address http://localhost:8265 -- python train.py

El resultado es similar a este:

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):
...
...

Verifica el estado del trabajo:

ray job status pytorch-mnist

El resultado es similar a este:

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.

Visualiza los registros de trabajos de Ray:

ray job logs pytorch-mnist

El resultado es similar a este:

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

Implementa un RayJob

El recurso personalizado RayJob administra el ciclo de vida de un recurso de RayCluster durante la ejecución de un solo trabajo de Ray.

Revisa el siguiente 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.9.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
            - name: ray-head
              image: rayproject/ray:2.9.0
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "2"
                  memory: "4Gi"
                requests:
                  cpu: "2"
                  memory: "4Gi"
    workerGroupSpecs:
      - replicas: 4
        minReplicas: 1
        maxReplicas: 5
        groupName: small-group
        rayStartParams: {}
        template:
          spec:
            containers:
              - name: ray-worker
                image: rayproject/ray:2.9.0
                resources:
                  limits:
                    cpu: "2"
                    memory: "8Gi"
                  requests:
                    cpu: "2"
                    memory: "8Gi"

En este manifiesto, se describe un recurso personalizado de RayJob.

Aplica el manifiesto a tu clúster de GKE:
```
kubectl apply -f ray-job.yaml
```

Verifica que el recurso de RayJob esté en ejecución:

kubectl get rayjob

El resultado es similar a este:

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 de RayJob es Running.

Visualiza el estado de los recursos de RayJob:

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

El resultado es similar a este:

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

Verifica que el trabajo de Ray esté completo:

kubectl get rayjob

El resultado es similar a este:

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 de RayJob es Complete.

Limpia

Borra el proyecto

Precaución: Borrar un proyecto tiene estas consecuencias:

Se borra todo en el proyecto. Si usaste un proyecto existente para las tareas de este documento, cuando lo borres, también se borrará cualquier otro trabajo que hayas realizado en el proyecto.
Se pierden los ID personalizados de proyectos. Cuando creaste este proyecto, es posible que hayas creado un ID del proyecto personalizado que desees utilizar en el futuro. Para conservar las URL que utilizan el ID del proyecto, como una URL appspot.com, borra los recursos seleccionados dentro del proyecto en lugar de borrar todo el proyecto.

Si planeas explorar varias infraestructuras, instructivos y guías de inicio rápido la reutilización de proyectos puede ayudarte a evitar exceder los límites de las cuotas del proyecto.

Borra un proyecto de Google Cloud:

gcloud projects delete PROJECT_ID

Borra los recursos individuales

Si usaste un proyecto existente y no quieres borrarlo, puedes borrar los recursos individuales. Para borrar el clúster, escribe lo siguiente:

gcloud container clusters delete ${CLUSTER_NAME}

¿Qué sigue?

Explora arquitecturas de referencia, diagramas y prácticas recomendadas sobre Google Cloud. Consulta nuestro Cloud Architecture Center.