Criar um cluster do Cluster Director com o Pathways

Nesta página, mostramos aos administradores de plataforma e engenheiros de machine learning (ML) como criar um cluster do Cluster Director orquestrado pelo Slurm e configurá-lo para cargas de trabalho do Pathways. Ao concluir essas etapas, você provisiona um cluster com frações de TPU v6e, um nó controlador de CPU, um sistema de arquivos compartilhado do Google Cloud Managed Lustre e o ambiente de execução /home/pathways. O Pathways não é pré-instalado em um cluster do Cluster Director: depois de criar o cluster, faça o download dos scripts do Pathways para o nó de login e execute um script de configuração uma vez.

Antes de começar

Antes de começar

Antes de começar, conclua os seguintes pré-requisitos:

  • Instale a CLI do Google Cloud.
  • Ative a API Hypercompute Cluster.
  • Verifique se a conta de usuário e a conta de serviço padrão do Compute Engine têm os papéis do Identity and Access Management (IAM) necessários para criar e se conectar a um cluster do Cluster Director. Para ver a lista de papéis, consulte Antes de começar no guia de início rápido do Cluster Director.
  • Crie uma reserva de TPU v6e na zona em que você planeja criar o cluster (por exemplo, pelo menos quatro VMs ct6e-standard-4t para duas frações 2x4).
  • Configure o sudo sem senha nos nós do cluster para recuperação automática de frações. Peça ao administrador para conceder a você o papel de Login de administrador no SO do Compute (roles/compute.osAdminLogin) no projeto.
  • Crie um bucket do Cloud Storage para dados temporários e saídas de carga de trabalho do Pathways e conceda o papel de Usuário de objetos do Storage (roles/storage.objectUser) no projeto à sua conta de usuário e à conta de serviço padrão do Compute Engine.

Configurar seu ambiente local

Faça login com suas credenciais do Google Cloud :

gcloud auth login

Defina as seguintes variáveis de ambiente com valores adequados à sua carga de trabalho:

export PROJECT=PROJECT_ID
export REGION=REGION
export ZONE=ZONE
export CLUSTER_ID=CLUSTER_ID
export RESERVATION_NAME=RESERVATION_NAME

Substitua:

  • PROJECT_ID: o ID do projeto doGoogle Cloud .
  • REGION: a região em que você cria recursos.
  • ZONE: a zona em que você cria recursos.
  • CLUSTER_ID: o nome do cluster. O nome pode ter até 10 caracteres e usar apenas números e letras minúsculas.
  • RESERVATION_NAME: o nome da sua reserva de TPU v6e em ZONE.

Criar um cluster do Cluster Director

No exemplo a seguir, você cria um cluster com duas frações de TPU v6e 2x4 (quatro VMs ct6e-standard-4t) da sua reserva e um nó de CPU para o controlador do Pathways. Todos os nós montam um sistema de arquivos compartilhado do Managed Lustre em /home. Para mais informações sobre os campos na definição do cluster, consulte Criar um cluster personalizado.

  1. Crie o cluster usando a API REST:

    Crie o corpo da solicitação:

    cat <<EOF > "${CLUSTER_ID}_create.json"
{
  "networkResources": {
    "net": {
      "config": {
        "newNetwork": {
          "network": "projects/${PROJECT}/global/networks/${CLUSTER_ID}-net"
        }
      }
    }
  },
  "storageResources": {
    "home": {
      "config": {
        "newLustre": {
          "lustre": "projects/${PROJECT}/locations/${ZONE}/instances/${CLUSTER_ID}-lustre",
          "filesystem": "lustre",
          "capacityGb": "18000",
          "perUnitStorageThroughput": "500"
        }
      }
    }
  },
  "computeResources": {
    "crtpu": {
      "config": {
        "newReservedInstances": {
          "reservation": "projects/${PROJECT}/zones/${ZONE}/reservations/${RESERVATION_NAME}"
        }
      }
    },
    "crcpu": {
      "config": {
        "newOnDemandInstances": {
          "zone": "${ZONE}",
          "machineType": "n2-highmem-32"
        }
      }
    }
  },
  "orchestrator": {
    "slurm": {
      "loginNodes": {
        "count": "1",
        "zone": "${ZONE}",
        "machineType": "n2-standard-16",
        "enableOsLogin": true,
        "enablePublicIps": true,
        "bootDisk": {
          "type": "pd-balanced",
          "sizeGb": "100"
        },
        "storageConfigs": [
          {
            "id": "home",
            "localMount": "/home"
          }
        ]
      },
      "nodeSets": [
        {
          "id": "nstpu",
          "computeId": "crtpu",
          "staticNodeCount": "4",
          "maxDynamicNodeCount": "0",
          "enablePublicIps": true,
          "resourcePolicyConfig": {
            "acceleratorTopology": "2x4"
          },
          "computeInstance": {
            "bootDisk": {
              "type": "hyperdisk-balanced",
              "sizeGb": "100"
            }
          },
          "storageConfigs": [
            {
              "id": "home",
              "localMount": "/home"
            }
          ]
        },
        {
          "id": "nscpu",
          "computeId": "crcpu",
          "staticNodeCount": "1",
          "maxDynamicNodeCount": "0",
          "enablePublicIps": true,
          "computeInstance": {
            "bootDisk": {
              "type": "pd-balanced",
              "sizeGb": "100",
              "image": "projects/clusterdirector-public-images/global/images/family/tpu-v6e-v5p-ubuntu2204-slurm2605-v20260902"
            }
          },
          "storageConfigs": [
            {
              "id": "home",
              "localMount": "/home"
            }
          ]
        }
      ],
      "partitions": [
        {
          "id": "tpu",
          "nodeSetIds": ["nstpu"]
        },
        {
          "id": "cpu",
          "nodeSetIds": ["nscpu"]
        }
      ],
      "defaultPartition": "tpu",
      "config": {
        "requeueExitCodes": [1, 2, 134, 137, 139, 143]
      }
    }
  }
}
EOF

    Em seguida, envie a solicitação:

    curl -s -X POST \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -d @"${CLUSTER_ID}_create.json" \
    "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters?clusterId=${CLUSTER_ID}"

    O cluster inclui o seguinte:

    • Partição tpu: um conjunto de nós estático de VMs de TPU v6e reservadas. staticNodeCount é o número de VMs e acceleratorTopology é a topologia de cada fração. Cada VM ct6e-standard-4t tem quatro chips, então uma fração 2x4 tem duas VMs. Para mais frações, multiplique: duas frações 2x4 precisam de um staticNodeCount de 4.
    • Partição cpu: uma VM sob demanda que executa o gerenciador de recursos do Pathways, o servidor proxy do IFRT e o cliente do MaxText. Os scripts do Pathways enviam jobs para uma partição chamada cpu. A VM usa o tipo de máquina n2-highmem-32 porque o serviço Pathways e o MaxText precisam de até 240 GB de memória enquanto salvam checkpoints. Ele também usa uma imagem com o Slurm 26.05, que é compatível com frações de TPU, porque o serviço Pathways envia os jobs de worker da TPU dessa VM.
    • requeueExitCodes: o Slurm coloca em fila novamente um job que sai com um destes códigos. Essa configuração permite que o Slurm reinicie um job de worker da TPU quando qualquer VM na fração for interrompida, não apenas a primeira.

  2. Aguarde até que o cluster esteja pronto. Acompanhe o progresso na página do Cluster Director no console do Google Cloud ou verifique se reconciling é false na saída do seguinte comando:

    curl -s \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters/${CLUSTER_ID}"

Configurar o ambiente do Pathways

O ambiente do Pathways fica em /home/pathways no sistema de arquivos compartilhado. Portanto, você o configura uma vez por cluster. Ele contém as imagens de contêiner do Pathways, um ambiente Python com JAX e pathwaysutils e os scripts do Pathways.

  1. Conecte-se ao nó de login. Para outras maneiras de se conectar, consulte Conectar-se a um cluster.

    gcloud compute ssh ${CLUSTER_ID}-login-001 \
    --project=${PROJECT} \
    --zone=${ZONE} \
    --tunnel-through-iap

  2. Verifique se as partições cpu e tpu existem e estão no estado idle:

    sinfo -o "%P %a %D %T %N %f"

    O resultado será o seguinte:

    PARTITION AVAIL NODES STATE NODELIST FEATURES
tpu* up 4 idle ${CLUSTER_ID}-nstpu-[0-3]
cpu up 1 idle ${CLUSTER_ID}-nscpu-0

  3. Baixe os scripts do Pathways para o nó de login:

    mkdir -p ~/pathways-scripts && cd ~/pathways-scripts
for f in setup_env.sh start_pathways.sh run_worker_slice.sh run_maxtext.sh; do
  curl -fsSLO "https://raw.githubusercontent.com/GoogleCloudPlatform/slurm-gcp/tpu-2605/pathways/scripts/${f}"
done

    Esse comando faz o download dos seguintes scripts:

    • setup_env.sh: realiza uma configuração única do ambiente do Pathways.
    • start_pathways.sh: inicia o serviço Pathways.
    • run_worker_slice.sh: inicia o worker do Pathways para uma fração de TPU (enviada por start_pathways.sh).
    • run_maxtext.sh: executa um exemplo de carga de trabalho de treinamento do MaxText.

  4. Execute o script de configuração. Isso leva cerca de 5 minutos:

    chmod +x *.sh
./setup_env.sh

    O script extrai as imagens do servidor e do proxy do Pathways, cria um ambiente Python com JAX, pathwaysutils e MaxText, copia os scripts para /home/pathways e vincula ~/pathways a /home/pathways. Ele termina com uma saída semelhante a esta:

    Environment ready: JAX 0.11.1

  5. Opcional: para dar a outros usuários do cluster acesso ao ambiente do Pathways, peça que eles vinculem /home/pathways ao diretório inicial:

    ln -sfn /home/pathways ~/pathways

Nesse momento, o cluster está pronto para o Pathways. Inicie o serviço Pathways como parte da execução de uma carga de trabalho, porque as cargas de trabalho em lote e interativas o iniciam com configurações diferentes. Para mais informações, consulte Executar uma carga de trabalho em lote com o Pathways e Executar uma carga de trabalho interativa com o Pathways.

Se você tiver problemas, consulte Resolver problemas do Pathways no Cloud.

Limpar

Para evitar cobranças na conta do Google Cloud pelos recursos usados nesta página, exclua o cluster do Cluster Director quando terminar:

curl -s -X DELETE \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters/${CLUSTER_ID}"

A seguir