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

Você pode treinar e implantar modelos de machine learning (ML) com eficiência em um ambiente robusto e escalonável. Isso é possível combinando Ray, PyTorch e o complemento Ray Operator no Google Kubernetes Engine (GKE). O GKE oferece o controle e o desempenho necessários para cargas de trabalho exigentes de IA/ML. Essas cargas de trabalho incluem o treinamento de grandes modelos de fundação e o atendimento de solicitações de inferência em escala. Neste tutorial, você vai usar o escalonamento automático do GKE e a computação distribuída do Ray para simplificar seus fluxos de trabalho de machine learning.

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.

Objetivos

Este guia é destinado a clientes de IA generativa, usuários novos ou atuais do GKE, engenheiros de ML, engenheiros de MLOps (DevOps) ou administradores de plataformas interessados em usar os recursos de orquestração de contêineres do Kubernetes para disponibilizar modelos usando o Ray.

  • Crie um cluster do GKE.
  • Crie um cluster do Ray usando o recurso personalizado do RayCluster.
  • Treine um modelo usando um job do Ray.
  • Implantar um job do Ray usando o recurso personalizado RayJob.

Custos

Neste documento, você vai usar os seguintes componentes faturáveis do Google Cloud:

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços.

Novos usuários do Google Cloud podem estar qualificados para um teste sem custo financeiro.

Ao concluir as tarefas descritas neste documento, é possível evitar o faturamento contínuo excluindo os recursos criados. Para mais informações, consulte Limpeza.

Antes de começar

O Cloud Shell vem pré-instalado com o software necessário para este tutorial, incluindo kubectl e a CLI gcloud. Se você não usa o Cloud Shell, é necessário instalar a CLI gcloud.

  1. Faça login na sua conta do Google Cloud . Se você começou a usar o Google Cloudagora, crie uma conta para avaliar o desempenho dos nossos produtos em situações reais. Clientes novos também recebem US $300 em créditos para executar, testar e implantar cargas de trabalho.
  2. Instale a CLI do Google Cloud.

  3. Ao usar um provedor de identidade (IdP) externo, primeiro faça login na CLI gcloud com sua identidade federada.

  4. Para inicializar a CLI gcloud, execute o seguinte comando:

    gcloud init
  5. Crie ou selecione um Google Cloud projeto.

    Papéis necessários para selecionar ou criar um projeto

    • Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
    • Criar um projeto: para criar um projeto, é preciso ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.
    • Crie um projeto do Google Cloud :

      gcloud projects create PROJECT_ID

      Substitua PROJECT_ID por um nome para o projeto do Google Cloud que você está criando.

    • Selecione o projeto Google Cloud que você criou:

      gcloud config set project PROJECT_ID

      Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

  6. Verifique se o faturamento está ativado para o projeto do Google Cloud .

  7. Ative a API GKE, se ainda não tiver feito isso:

    Funções necessárias para ativar APIs

    Para ativar APIs, você precisa da permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão pelo papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão pelo papel de Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin). Saiba como conceder papéis.

    gcloud services enable container.googleapis.com
  8. Instale a CLI do Google Cloud.

  9. Ao usar um provedor de identidade (IdP) externo, primeiro faça login na CLI gcloud com sua identidade federada.

  10. Para inicializar a CLI gcloud, execute o seguinte comando:

    gcloud init
  11. Crie ou selecione um Google Cloud projeto.

    Papéis necessários para selecionar ou criar um projeto

    • Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
    • Criar um projeto: para criar um projeto, é preciso ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.
    • Crie um projeto do Google Cloud :

      gcloud projects create PROJECT_ID

      Substitua PROJECT_ID por um nome para o projeto do Google Cloud que você está criando.

    • Selecione o projeto Google Cloud que você criou:

      gcloud config set project PROJECT_ID

      Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

  12. Verifique se o faturamento está ativado para o projeto do Google Cloud .

  13. Ative a API GKE, se ainda não tiver feito isso:

    Funções necessárias para ativar APIs

    Para ativar APIs, você precisa da permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão pelo papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão pelo papel de Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin). Saiba como conceder papéis.

    gcloud services enable container.googleapis.com
  14. Conceda papéis à sua conta de usuário. Execute uma vez o seguinte comando para cada um dos seguintes papéis do IAM: roles/container.clusterAdmin, roles/container.admin

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    Substitua:

    • PROJECT_ID: o ID do projeto.
    • USER_IDENTIFIER: o identificador da sua conta de usuário . Por exemplo, myemail@example.com.
    • ROLE: o papel do IAM que você concede à sua conta de usuário.

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 Google Cloud console. 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 do Ray 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.

Observar suas cargas de trabalho do Ray

Para conferir detalhes dos RayJobs, acesse a seção Kubernetes Engine > IA/ML > Jobs no Google Cloud console.

Ver RayJobs no console Google Cloud

Limpar

Excluir o projeto

    Excluir um projeto do Google Cloud :

    gcloud projects delete PROJECT_ID

Excluir recursos individuais

Se você usou um projeto existente e não quer excluí-lo, exclua os recursos individuais. Para excluir o cluster, digite:

gcloud container clusters delete ${CLUSTER_NAME}

A seguir