Treinar um modelo com PyTorch, Ray e Google Kubernetes Engine (GKE) em GPUs

Neste guia, demonstramos como treinar um modelo no Google Kubernetes Engine (GKE) usando Ray, PyTorch e o complemento Ray Operator.

Sobre o Ray

O Ray é um framework de computação escalonável e de código aberto para aplicativos de IA/ML. O Ray Train é um componente no Ray desenvolvido para treinamento e ajuste de modelos distribuídos. Você pode usar a API Ray Train para escalonar o treinamento em várias máquinas e fazer a integração com bibliotecas de machine learning, como o PyTorch.

É possível implantar jobs de treinamento do Ray usando os recursos RayCluster ou RayJob. Use um recurso do RayJob ao implantar jobs do Ray na produção pelos seguintes motivos:

  • O recurso RayJob cria um cluster do Ray temporário que pode ser excluído automaticamente quando um job é concluído.
  • O recurso RayJob oferece suporte a políticas de repetição para execução de jobs resilientes.
  • É possível gerenciar jobs do Ray usando os padrões conhecidos da API Kubernetes.

Preparar o ambiente

Para preparar o ambiente, siga estas etapas:

  1. Inicie uma sessão do Cloud Shell no console do Google Cloud clicando em Ícone de ativação do Cloud Shell Ativar o Cloud Shell no console doGoogle Cloud . Isso inicia uma sessão no painel inferior do console Google Cloud .

  2. Defina as variáveis de ambiente:

    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`
    

    Substitua:

    • PROJECT_ID: o Google Cloud ID do projeto.
    • CLUSTER_VERSION: a versão do GKE a ser usada. Precisa ser 1.30.1 ou mais recente.
  3. Clone o repositório do GitHub:

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
    
  4. Mude para o diretório de trabalho:

    cd kubernetes-engine-samples/ai-ml/gke-ray/raytrain/pytorch-mnist
    
  5. Crie um ambiente virtual em Python:

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

Criar um cluster do GKE

Crie um cluster do GKE Autopilot ou Padrão:

Piloto automático

Criar um cluster do Autopilot:

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

Padrão

Criar um cluster padrão

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

Implantar um recurso do RayCluster

Implante um recurso do RayCluster no cluster:

  1. Analise o seguinte manifesto:

    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"

    Esse manifesto descreve um recurso personalizado do RayCluster.

  2. Aplique o manifesto ao seu cluster do GKE:

    kubectl apply -f ray-cluster.yaml
    
  3. Verifique se o recurso RayCluster está pronto:

    kubectl get raycluster
    

    O resultado será assim:

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

    Nesta saída, ready na coluna STATUS indica que o recurso do RayCluster está pronto.

Conectar-se ao recurso do RayCluster

Conecte-se ao recurso do RayCluster para enviar um job do Ray.

  1. Verifique se o GKE criou o serviço RayCluster:

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

    O resultado será assim:

    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. Estabeleça uma sessão de encaminhamento de portas para o Ray head:

    kubectl port-forward svc/pytorch-mnist-cluster-head-svc 8265:8265 2>&1 >/dev/null &
    
  3. Verifique se o cliente Ray pode se conectar ao cluster usando o localhost:

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

    O resultado será assim:

    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
    

Treinar um modelo

Treinar um modelo PyTorch usando o Conjunto de dados Fashion MNIST:

  1. Envie um job do Ray e aguarde a conclusão:

    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
    

    O resultado será assim:

    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. Verifique o status do job:

    ray job status pytorch-mnist
    

    O resultado será assim:

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

    Aguarde até que Status for job seja COMPLETE. Isso pode levar 15 minutos ou mais.

  3. Veja os registros de jobs do Ray:

    ray job logs pytorch-mnist
    

    O resultado será assim:

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

Implantar um RayJob

O recurso RayJob personalizado gerencia o ciclo de vida de um recurso do RayCluster durante a execução de um único job do Ray.

  1. Analise o seguinte manifesto:

    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"

    Esse manifesto descreve um recurso personalizado do RayJob.

  2. Aplique o manifesto ao seu cluster do GKE:

    kubectl apply -f ray-job.yaml
    
  3. Verifique se o recurso RayJob está em execução:

    kubectl get rayjob
    

    O resultado será assim:

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

    Nesta saída, a coluna DEPLOYMENT STATUS indica que o recurso do RayJob é Running.

  4. Veja o status do recurso RayJob:

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

    O resultado será assim:

    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. Verifique se o job do Ray foi concluído:

    kubectl get rayjob
    

    O resultado será assim:

    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
    

    Nesta saída, a coluna DEPLOYMENT STATUS indica que o recurso do RayJob é Complete.