Administra clústeres de entrenamiento con la API de Slurm Job

La API de Slurm Job te permite administrar tus clústeres de entrenamiento de Gemini Enterprise Agent Platform a través de la API de Agent Platform, una secuencia de comandos, un notebook o una canalización de CI/CD. La API de Slurm Job te permite administrar el trabajo del clúster de Slurm sin conectarte al clúster con SSH.

Requisitos previos

Antes de usar la API de Slurm Job para administrar tu clúster, asegúrate de cumplir con los siguientes requisitos:

  • Tienes acceso a los clústeres de entrenamiento, por lo general, a través del rol Usuario de Vertex AI roles/aiplatform.user en el proyecto del clúster o cualquier rol que incluya el permiso aiplatform.googleapis.com/modelDevelopmentClusters.run.

  • Google Cloud CLI está instalada y configuraste las credenciales predeterminadas de la aplicación (gcloud auth application-default login). Para las herramientas automatizadas que no pueden ejecutar un acceso interactivo, autentícate como una cuenta de servicio con el permiso aiplatform.googleapis.com/modelDevelopmentClusters.run.

  • Conoces el ID del proyecto, la región y el ID del clúster.

Además, te recomendamos que configures un alias de conveniencia:

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

Descripción general

Cada llamada es una solicitud POST al extremo callSlurmRestApi de tu clúster. En el cuerpo de la solicitud, proporciona los siguientes campos:

  • method: Es el método HTTP de la API de REST de Slurm, como HTTP_METHOD_GET, HTTP_METHOD_POST o HTTP_METHOD_DELETE.

  • path: Es la ruta de acceso a la API de REST de Slurm. Por ejemplo, /slurm/v0.0.42/job/submit. Usa la versión de Slurm que ejecuta tu clúster. En los ejemplos de esta guía, se usa la versión 0.0.42.

  • body: Es un objeto JSON con la carga útil de la solicitud. El cuerpo es obligatorio para las solicitudes HTTP_METHOD_POST. Omítelo para las solicitudes HTTP_METHOD_GET y HTTP_METHOD_DELETE.

La respuesta contiene los siguientes campos:

  • status: Es el estado HTTP que devuelve la API de Slurm Job. Por ejemplo, 200.

  • structuredBody: Es la respuesta de la API de Slurm Job, como un objeto JSON. Lee los campos que necesitas, como el ID del trabajo o la lista de trabajos, directamente desde structuredBody.

Los trabajos se ejecutan automáticamente en tu propia cuenta de Linux en los clústeres de entrenamiento. No hay configuración adicional y no puedes ejecutar el trabajo como otro usuario. No se requiere SSH, ya que todo se realiza a través de la API de Agent Platform.

Cada llamada devuelve la respuesta de Slurm directamente, sin una operación para sondear. Puedes hacer un seguimiento del estado de un trabajo a lo largo del tiempo. Para obtener más información, consulta Cómo verificar el estado del trabajo.

Llama a la API

En las siguientes secciones, se muestra cómo usar la API de Slurm Job.

Envía un trabajo

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"
    }
  }'

Reemplaza lo siguiente:

  • REGION: Es la región en la que se encuentra el clúster.

  • PROJECT_ID: ID del proyecto

  • CLUSTER_ID: Es el ID de tu clúster.

  • PARTITION: Es la partición de Slurm a la que te conectas.

  • USERNAME: Tu nombre de usuario en la partición de Slurm.

El valor de time_limit es el tiempo en minutos. El campo script contiene el contenido del script, no una ruta de acceso a un archivo.

La API de Slurm Job devuelve el ID del trabajo nuevo en structuredBody. Puedes usar el comando jq para aislar el ID del trabajo:

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'

Verifica el estado del trabajo

Usa lo siguiente para supervisar un trabajo en ejecución. Consulta hasta que el estado alcance un estado detenido, como COMPLETED, FAILED o 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"
  }'

Reemplaza lo siguiente:

  • REGION: Es la región en la que se encuentra el clúster.

  • PROJECT_ID: ID del proyecto

  • CLUSTER_ID: Es el ID de tu clúster.

  • JOB_ID: Es el ID del trabajo de Slurm.

Mostrar trabajos

Usa el siguiente comando para enumerar los trabajos en ejecución:

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/"
  }'

Reemplaza lo siguiente:

  • REGION: Es la región en la que se encuentra el clúster.

  • PROJECT_ID: ID del proyecto

  • CLUSTER_ID: Es el ID de tu clúster.

Cancelar un trabajo

La cancelación se acepta de inmediato, pero el trabajo tarda un momento en detenerse. Verifica su estado para confirmar que sea CANCELLED. Solo puedes cancelar tus propios trabajos. Cancelar un trabajo completado es seguro y devuelve un mensaje de éxito.

Usa el siguiente comando para cancelar un trabajo en ejecución:

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"
  }'

Reemplaza lo siguiente:

  • REGION: Es la región en la que se encuentra el clúster.

  • PROJECT_ID: ID del proyecto

  • CLUSTER_ID: Es el ID de tu clúster.

  • JOB_ID: Es el ID del trabajo de Slurm.

Cómo mostrar un trabajo finalizado

Verificar un trabajo con la ruta de acceso /slurm solo funciona mientras el trabajo aún está en la cola en vivo. Para los trabajos completados, usa la ruta de acceso 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"
  }'

Reemplaza lo siguiente:

  • REGION: Es la región en la que se encuentra el clúster.

  • PROJECT_ID: ID del proyecto

  • CLUSTER_ID: Es el ID de tu clúster.

  • JOB_ID: Es el ID del trabajo de Slurm.

Usa la API de Slurm Job desde Python

Si prefieres usar Python en lugar de curl, puedes usar la siguiente secuencia de comandos auxiliar para enviar trabajos, verificar el estado de los trabajos y cancelarlos. La secuencia de comandos auxiliar usa la biblioteca google-auth para las credenciales predeterminadas de la aplicación.

Para usar la API de Slurm Job desde Python, haz lo siguiente:

  1. Instala las credenciales predeterminadas de la aplicación:

    pip install google-auth requests

  2. Usa la siguiente secuencia de comandos de Python para usar la 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"])

La secuencia de comandos auxiliar devuelve el estado de Slurm y structuredBody, de modo que puedes leer campos como job_id o job_state directamente. Para usar la secuencia de comandos en entornos automatizados, autentícate como una cuenta de servicio en lugar de usar un acceso interactivo.

Respuesta de la API

A continuación, se muestra una respuesta típica de la API de Slurm Job:

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

El campo status es el estado HTTP de Slurm y structuredBody contiene la respuesta de Slurm como un objeto JSON, por lo que puedes leer sus campos directamente, como jq -r '.structuredBody.job_id'.

Siempre verifica structuredBody, independientemente del valor de status, ya que structuredBody incluye los errores o las advertencias que informa Slurm.

Si la solicitud CallSlurmRestApi tiene un formato incorrecto, es posible que obtengas un resultado error, como en el siguiente ejemplo:

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

Respuestas habituales

Lo que se muestra Qué significa
Estado 200 con un ID de trabajo El trabajo se envió correctamente.
Estado 200, sin errores La solicitud se realizó correctamente.
Estado 200 con advertencias La solicitud se realizó correctamente, pero Slurm ignoró o ajustó algo. Lee la advertencia.
Estado 404 para un trabajo El trabajo no está en la cola en vivo; es posible que haya finalizado. Prueba la ruta de historial.
Estado 400 o 500 Slurm rechazó la solicitud. Los motivos del rechazo pueden incluir un cuerpo con formato incorrecto o un campo del tipo incorrecto. Lee el error en structuredBody.
Error PERMISSION_DENIED No tienes acceso a este clúster.
Error NOT_FOUND El nombre del clúster es incorrecto o el clúster no existe.
Error de INVALID_ARGUMENT La ruta no es una ruta de Slurm válida o la solicitud es demasiado grande.
Error UNAVAILABLE No se puede acceder al clúster temporalmente. Vuelve a intentarlo después de una breve espera.

Límites

El tamaño de la solicitud completa debe ser inferior a 1 MB. Para usar secuencias de comandos grandes, colócalas en el almacenamiento compartido del clúster y, luego, llámalas desde una secuencia de comandos de wrapper corta.