Gestire i cluster di addestramento utilizzando l'API Slurm Job

L'API Slurm Job consente di gestire i cluster di addestramento di Gemini Enterprise Agent Platform tramite l'API Agent Platform, uno script, un blocco note o una pipeline CI/CD. L'API Slurm Job ti consente di gestire il job del cluster Slurm senza connetterti al cluster con SSH.

Prerequisiti

Prima di utilizzare l'API Slurm Job per gestire il cluster, assicurati di soddisfare i seguenti requisiti:

  • Hai accesso ai cluster di addestramento, in genere tramite il ruolo Vertex AI User roles/aiplatform.user nel progetto del cluster o qualsiasi ruolo che includa l'autorizzazione aiplatform.googleapis.com/modelDevelopmentClusters.run.

  • Google Cloud CLI è installata e hai configurato le Credenziali predefinite dell'applicazione (gcloud auth application-default login). Per gli strumenti automatizzati che non possono eseguire un accesso interattivo, esegui l'autenticazione come account di servizio con l'autorizzazione aiplatform.googleapis.com/modelDevelopmentClusters.run.

  • Conosci l'ID progetto, la regione e l'ID cluster del tuo cluster.

Inoltre, ti consigliamo di configurare un alias pratico:

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

Panoramica

Ogni chiamata è una richiesta POST all'endpoint callSlurmRestApi per il tuo cluster. Nel corpo della richiesta fornisci i seguenti campi:

  • method: il metodo HTTP REST di Slurm, ad esempio HTTP_METHOD_GET, HTTP_METHOD_POST o HTTP_METHOD_DELETE.

  • path: il percorso REST di Slurm. Ad esempio: /slurm/v0.0.42/job/submit. Utilizza la versione di Slurm in esecuzione sul cluster. Gli esempi di questa guida utilizzano la versione 0.0.42.

  • body: un oggetto JSON con il payload della richiesta. Il corpo è obbligatorio per le richieste HTTP_METHOD_POST. Omettilo per le richieste HTTP_METHOD_GET e HTTP_METHOD_DELETE.

La risposta contiene i seguenti campi:

  • status: lo stato HTTP restituito dall'API Slurm Job. Ad esempio 200.

  • structuredBody: la risposta dell'API Slurm Job, come oggetto JSON. Leggi direttamente i campi che ti servono, ad esempio l'ID job o l'elenco dei job, da structuredBody.

I job vengono eseguiti automaticamente sul tuo account Linux nei cluster di addestramento. Non è necessaria alcuna configurazione aggiuntiva e non puoi eseguire il job come un altro utente. SSH non è obbligatorio, perché tutto viene eseguito tramite l'API Agent Platform.

Ogni chiamata restituisce direttamente la risposta di Slurm, senza un'operazione da eseguire. Puoi monitorare lo stato di un job nel tempo. Per saperne di più, consulta Controllare lo stato del job.

Chiama l'API

Le sezioni seguenti mostrano come utilizzare l'API Slurm Job.

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

Sostituisci quanto segue:

  • REGION: la regione in cui si trova il cluster.

  • PROJECT_ID: il tuo ID progetto.

  • CLUSTER_ID: l'ID cluster.

  • PARTITION: la partizione Slurm a cui ti stai connettendo.

  • USERNAME: Il tuo nome utente sulla partizione Slurm.

Il valore time_limit è il tempo in minuti. Il campo script contiene i contenuti dello script, non un percorso a un file.

L'API Slurm Job restituisce l'ID del nuovo job in structuredBody. Puoi utilizzare il comando jq per isolare l'ID 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'

Verifica dello stato di un job

Utilizza quanto segue per monitorare un job in esecuzione. Esegui il polling finché lo stato non raggiunge uno stato di arresto, ad esempio 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"
  }'

Sostituisci quanto segue:

  • REGION: la regione in cui si trova il cluster.

  • PROJECT_ID: il tuo ID progetto.

  • CLUSTER_ID: l'ID cluster.

  • JOB_ID: l'ID job del job Slurm.

Elenca job

Utilizza il seguente comando per elencare i job in esecuzione:

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

Sostituisci quanto segue:

  • REGION: la regione in cui si trova il cluster.

  • PROJECT_ID: il tuo ID progetto.

  • CLUSTER_ID: l'ID cluster.

Annullare un job

L'annullamento viene accettato immediatamente, ma l'arresto del job richiede un po' di tempo. Controlla il suo stato per verificare che sia CANCELLED. Puoi annullare solo i tuoi job. L'annullamento di un job completato è sicuro e restituisce esito positivo.

Utilizza quanto segue per annullare un job in esecuzione:

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

Sostituisci quanto segue:

  • REGION: la regione in cui si trova il cluster.

  • PROJECT_ID: il tuo ID progetto.

  • CLUSTER_ID: l'ID cluster.

  • JOB_ID: l'ID job del job Slurm.

Visualizzare un job completato

Il controllo di un job con il percorso /slurm funziona solo mentre il job è ancora nella coda live. Per i job completati, utilizza invece il percorso 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"
  }'

Sostituisci quanto segue:

  • REGION: la regione in cui si trova il cluster.

  • PROJECT_ID: il tuo ID progetto.

  • CLUSTER_ID: l'ID cluster.

  • JOB_ID: l'ID job del job Slurm.

Utilizzo dell'API Slurm Job da Python

Se preferisci utilizzare Python anziché curl, puoi utilizzare il seguente script helper per inviare job, controllare lo stato dei job e annullarli. Lo script helper utilizza la libreria google-auth perCredenziali predefinite dell'applicazionee.

Per utilizzare l'API Slurm Job da Python:

  1. Installa le Credenziali predefinite dell'applicazione:

    pip install google-auth requests

  2. Utilizza il seguente script Python per utilizzare l'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"])

Lo script helper restituisce lo stato da Slurm e structuredBody, in modo da poter leggere direttamente campi come job_id o job_state. Per utilizzare lo script in ambienti automatizzati, esegui l'autenticazione come service account anziché utilizzare un accesso interattivo.

Risposta dell'API

Di seguito è riportata una tipica risposta dell'API Slurm Job:

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

Il campo status è lo stato HTTP di Slurm e structuredBody contiene la risposta di Slurm come oggetto JSON, in modo da poter leggere direttamente i suoi campi, ad esempio jq -r '.structuredBody.job_id'.

Controlla sempre structuredBody indipendentemente dal valore di status, perché structuredBody include eventuali errori o avvisi segnalati da Slurm.

Se la richiesta CallSlurmRestApi non è formattata correttamente, potresti ottenere un risultato error come nell'esempio seguente:

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

Risposte comuni

Cosa vedi Che cosa significa
Stato 200 con un ID job Job inviato correttamente.
Stato 200, nessun errore Richiesta riuscita.
Stato 200 con avvisi Richiesta riuscita, ma Slurm ha ignorato o modificato qualcosa. Leggi l'avviso.
Stato 404 per un job Il job non è nella coda live, quindi potrebbe essere stato completato. Prova il percorso della cronologia.
Stato 400 o 500 Slurm ha rifiutato la richiesta. I motivi del rifiuto potrebbero includere un corpo non valido o un campo del tipo errato. Leggi l'errore in structuredBody.
Errore PERMISSION_DENIED Non hai accesso a questo cluster.
Errore NOT_FOUND Il nome del cluster non è corretto o il cluster non esiste.
Errore INVALID_ARGUMENT Il percorso non è un percorso Slurm valido o la richiesta è troppo grande.
Errore UNAVAILABLE Il cluster non è temporaneamente raggiungibile. Riprova dopo una breve attesa.

Limiti

Le dimensioni dell'intera richiesta devono essere inferiori a 1 MB. Per utilizzare script di grandi dimensioni, posiziona lo script nello spazio di archiviazione condiviso del cluster e poi chiamalo da un breve script wrapper.