Gerenciar clusters de treinamento usando a API Slurm Job

A API Slurm Job permite gerenciar seus clusters de treinamento da Gemini Enterprise Agent Platform usando a API Agent Platform, um script, um notebook ou um pipeline de CI/CD. A API Slurm Job permite gerenciar o job do cluster do Slurm sem se conectar a ele com SSH.

Pré-requisitos

Antes de usar a API Slurm Job para gerenciar seu cluster, verifique se você atende aos seguintes requisitos:

• Você tem acesso aos clusters de treinamento, geralmente pela função Usuário da Vertex AI roles/aiplatform.user no projeto do cluster ou qualquer função que inclua a permissão aiplatform.googleapis.com/modelDevelopmentClusters.run.

• A Google Cloud CLI está instalada e você configurou as Application Default Credentials (gcloud auth application-default login). Para ferramentas automatizadas que não podem executar um login interativo, autentique-se como uma conta de serviço com a permissão aiplatform.googleapis.com/modelDevelopmentClusters.run.

• Você conhece o ID do projeto, a região e o ID do cluster.

Além disso, recomendamos que você configure um alias de conveniência:

alias gcurl='curl -H "Authorization: Bearer $(gcloud auth application-default print-access-token)" -H "Content-Type: application/json"'

Visão geral

Cada chamada é uma solicitação POST para o endpoint callSlurmRestApi do seu cluster. No corpo da solicitação, informe os seguintes campos:

• method: o método HTTP REST do Slurm, como HTTP_METHOD_GET, HTTP_METHOD_POST ou HTTP_METHOD_DELETE.

• path: o caminho REST do Slurm. Por exemplo, /slurm/v0.0.42/job/submit. Use a versão do Slurm que seu cluster está executando. Os exemplos neste guia usam a v0.0.42.

• body: um objeto JSON com o payload da solicitação. O corpo é obrigatório para solicitações HTTP_METHOD_POST. Omita para solicitações HTTP_METHOD_GET e HTTP_METHOD_DELETE.

A resposta contém os seguintes campos:

• status: o status HTTP retornado pela API Slurm Job. Por exemplo, 200.

• structuredBody: a resposta da API Slurm Job, como um objeto JSON. Leia os campos necessários, como o ID do job ou a lista de jobs, diretamente de structuredBody.

Você executa jobs na sua própria conta do Linux nos clusters de treinamento automaticamente. Não há configuração adicional, e não é possível executar o job como outro usuário. Não é necessário usar SSH, porque tudo é feito pela API da Agent Platform.

Cada chamada retorna a resposta do Slurm diretamente, sem uma operação para sondar. Você pode acompanhar o status de um job ao longo do tempo. Para mais informações, consulte Verificar o status do job.

Chamar a API

As seções a seguir mostram como usar a API Slurm Job.

Envie um job

gcurl -X POST \
  "https://REGION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/REGION/modelDevelopmentClusters/CLUSTER_ID:callSlurmRestApi" \
  -d '{
    "method": "HTTP_METHOD_POST",
    "path": "/slurm/v0.0.42/job/submit",
    "body": {
      "job": {
        "name": "my-run",
        "partition": "PARTITION",
        "current_working_directory": "/home/USERNAME",
        "minimum_nodes": 1,
        "tasks_per_node": 1,
        "environment": ["PATH=/bin:/usr/bin"],
        "time_limit": {"set": true, "number": 120}
      },
      "script": "#!/bin/bash\necho hello\nsleep 5\necho done\n"
    }
  }'

Substitua:

• REGION: a região em que o cluster está.

• PROJECT_ID: o ID do projeto.

• CLUSTER_ID: o ID do cluster.

• PARTITION: a partição do Slurm a que você está se conectando.

• USERNAME: seu nome de usuário na partição do Slurm.

O valor time_limit é o tempo em minutos. O campo script contém o conteúdo do script, não um caminho para um arquivo.

A API Slurm Job retorna o ID do novo job em structuredBody. Use o comando jq para isolar o ID do job:

gcurl -sS -X POST "https://REGION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/REGION/modelDevelopmentClusters/CLUSTER_ID:callSlurmRestApi" \
  -d '{ ... }' | jq -r '.structuredBody.job_id'

Verificar o status do job

Use o seguinte para monitorar um job em execução. Pesquise até que o status atinja um estado interrompido, como COMPLETED, FAILED ou CANCELLED:

gcurl -X POST \
  "https://REGION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/REGION/modelDevelopmentClusters/CLUSTER_ID:callSlurmRestApi" \
  -d '{
    "method": "HTTP_METHOD_GET",
    "path": "/slurm/v0.0.42/job/JOB_ID"
  }'

Substitua:

• REGION: a região em que o cluster está.

• PROJECT_ID: o ID do projeto.

• CLUSTER_ID: o ID do cluster.

• JOB_ID: o ID do job do Slurm.

Listar jobs

Use o seguinte para listar os jobs em execução:

gcurl -X POST \
  "https://REGION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/REGION/modelDevelopmentClusters/CLUSTER_ID:callSlurmRestApi" \
  -d '{
    "method": "HTTP_METHOD_GET",
    "path": "/slurm/v0.0.42/jobs/"
  }'

Substitua:

• REGION: a região em que o cluster está.

• PROJECT_ID: o ID do projeto.

• CLUSTER_ID: o ID do cluster.

Cancelar um job

O cancelamento é aceito imediatamente, mas o job leva um tempo para ser interrompido. Confira o status para verificar se ele é CANCELLED. Você só pode cancelar seus próprios jobs. O cancelamento de um job concluído é seguro e retorna sucesso.

Use o seguinte para cancelar um job em execução:

gcurl -X POST \
  "https://REGION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/REGION/modelDevelopmentClusters/CLUSTER_ID:callSlurmRestApi" \
  -d '{
    "method": "HTTP_METHOD_DELETE",
    "path": "/slurm/v0.0.42/job/JOB_ID"
  }'

Substitua:

• REGION: a região em que o cluster está.

• PROJECT_ID: o ID do projeto.

• CLUSTER_ID: o ID do cluster.

• JOB_ID: o ID do job do Slurm.

Mostrar um job concluído

A verificação de um job com o caminho /slurm só funciona enquanto ele ainda está na fila ativa. Para jobs concluídos, use o caminho slurmdb:

gcurl -X POST \
  "https://REGION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/REGION/modelDevelopmentClusters/CLUSTER_ID:callSlurmRestApi" \
  -d '{
    "method": "HTTP_METHOD_GET",
    "path": "/slurmdb/v0.0.42/job/JOB_ID"
  }'

Substitua:

• REGION: a região em que o cluster está.

• PROJECT_ID: o ID do projeto.

• CLUSTER_ID: o ID do cluster.

• JOB_ID: o ID do job do Slurm.

Como usar a API Slurm Job com Python

Se você preferir usar Python em vez de curl, use o script auxiliar a seguir para enviar, verificar o status e cancelar jobs. O script auxiliar usa a biblioteca google-auth para Application Default Credentials.

Faça o seguinte para usar a API Slurm Job com Python:

1. Instale o Application Default Credentials:

   pip install google-auth requests
   

2. Use o seguinte script Python para usar a API:

   import google.auth
   import google.auth.transport.requests
   
   # Fill in your cluster's values.
   PROJECT_ID = "PROJECT_ID"
   REGION = "REGION"
   CLUSTER_ID = "CLUSTER_ID"
   PARTITION = "PARTITION"
   USERNAME = "USERNAME"
   
   # Application Default Credentials: gcloud auth application-default login, or a
   # service account for automated callers.
   credentials, _ = google.auth.default(
       scopes=["https://www.googleapis.com/auth/cloud-platform"])
   session = google.auth.transport.requests.AuthorizedSession(credentials)
   
   url = (
       f"https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}"
       f"/locations/{REGION}/modelDevelopmentClusters/{CLUSTER_ID}:callSlurmRestApi"
   )
   
   def call_slurm(method, path, body=None):
       request = {"method": method, "path": path}
       if body is not None:
           request["body"] = body
       response = session.post(url, json=request)
       response.raise_for_status()
       envelope = response.json()
       return envelope["status"], envelope.get("structuredBody", {})
   
   # Submit a job.
   status, result = call_slurm(
       "HTTP_METHOD_POST",
       "/slurm/v0.0.42/job/submit",
       {
           "job": {
               "name": "my-run",
               "partition": f"{PARTITION}",
               "current_working_directory": f"/home/{USERNAME}",
               "minimum_nodes": 1,
               "tasks_per_node": 1,
               "environment": ["PATH=/bin:/usr/bin"],
               "time_limit": {"set": True, "number": 120},
           },
           "script": "#!/bin/bash\necho hello\nsleep 5\necho done\n",
       },
   )
   job_id = result["job_id"]
   print("submitted job", job_id)
   
   # Check its status.
   _, pending_result = call_slurm("HTTP_METHOD_GET", f"/slurm/v0.0.42/job/{job_id}")
   print("state:", pending_result["jobs"][0]["job_state"])
   
   # Cancel it.
   call_slurm("HTTP_METHOD_DELETE", f"/slurm/v0.0.42/job/{job_id}")
   print("cancelled job", job_id)
   
   # Confirm the job is cancelled.
   _, cancelled_result  = call_slurm("HTTP_METHOD_GET", f"/slurm/v0.0.42/job/{job_id}")
   print("state", cancelled_result["jobs"][0]["job_state"])
   
   

O script auxiliar retorna o status do Slurm e do structuredBody para que você possa ler campos como job_id ou job_state diretamente. Para usar o script em ambientes automatizados, faça a autenticação como uma conta de serviço em vez de usar um login interativo.

Resposta da API

Confira a seguir uma resposta típica da API Slurm Job:

{
  "status": 200,
  "structuredBody": {
    "job_id": 1234,
    "errors": [],
    "warnings": []
  }
}

O campo status é o status HTTP do Slurm, e structuredBody contém a resposta do Slurm como um objeto JSON. Assim, é possível ler os campos diretamente, como jq -r '.structuredBody.job_id'.

Sempre verifique structuredBody, seja qual for o valor de status, porque structuredBody inclui todos os erros ou avisos informados pelo Slurm.

Se a solicitação CallSlurmRestApi estiver malformatada, você poderá receber um resultado error, como no exemplo a seguir:

{
  "error": {
    "code": 400,
    "message": "{ ... some error ... }",
    "status": "INVALID_ARGUMENT"
  }
}

Respostas comuns

Limites

A solicitação inteira precisa ter menos de 1 MB. Para usar scripts grandes, coloque-os no armazenamento compartilhado do cluster e chame-os de um script wrapper curto.