Trainingscluster mit der Slurm Job API verwalten

Mit der Slurm Job API können Sie Ihre Trainingscluster der Gemini Enterprise Agent Platform über die Agent Platform API, ein Skript, ein Notebook oder eine CI/CD-Pipeline verwalten. Mit der Slurm Job API können Sie den Slurm-Clusterjob verwalten, ohne eine SSH-Verbindung zum Cluster herzustellen.

Vorbereitung

Bevor Sie die Slurm Job API zum Verwalten Ihres Clusters verwenden, müssen die folgenden Voraussetzungen erfüllt sein:

• Sie haben Zugriff auf die Trainingscluster, in der Regel über die Rolle Vertex AI-Nutzer roles/aiplatform.user im Projekt des Clusters oder über eine beliebige Rolle, die die Berechtigung aiplatform.googleapis.com/modelDevelopmentClusters.run enthält.

• Die Google Cloud CLI ist installiert und Sie haben Standardanmeldedaten für Anwendungen (gcloud auth application-default login) konfiguriert. Für automatisierte Tools, die keine interaktive Anmeldung ausführen können, authentifizieren Sie sich als Dienstkonto mit der Berechtigung aiplatform.googleapis.com/modelDevelopmentClusters.run.

• Sie kennen die Projekt-ID, die Region und die Cluster-ID Ihres Clusters.

Außerdem empfehlen wir, einen Alias einzurichten:

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

Übersicht

Jeder Aufruf ist eine POST-Anfrage an den callSlurmRestApi-Endpunkt für Ihren Cluster. Geben Sie im Anfragetext die folgenden Felder an:

• method: Die Slurm REST-HTTP-Methode, z. B. HTTP_METHOD_GET, HTTP_METHOD_POST oder HTTP_METHOD_DELETE.

• path: Der Slurm-REST-Pfad. Beispiel: /slurm/v0.0.42/job/submit Verwenden Sie die Slurm-Version, die auf Ihrem Cluster ausgeführt wird. In den Beispielen in diesem Leitfaden wird Version 0.0.42 verwendet.

• body: Ein JSON-Objekt mit der Anfrage-Nutzlast. Der Text ist für HTTP_METHOD_POST-Anfragen erforderlich. Lassen Sie es bei HTTP_METHOD_GET- und HTTP_METHOD_DELETE-Anfragen weg.

Die Antwort umfasst die folgenden Felder:

• status: Der von der Slurm Job API zurückgegebene HTTP-Status. Beispiel: 200.

• structuredBody: Die Antwort der Slurm Job API als JSON-Objekt. Lesen Sie die benötigten Felder wie die Job-ID oder die Jobliste direkt aus structuredBody.

Sie führen Jobs automatisch in Ihrem eigenen Linux-Konto in den Trainingsclustern aus. Es ist keine zusätzliche Konfiguration erforderlich und Sie können den Job nicht als anderer Nutzer ausführen. SSH ist nicht erforderlich, da alles über die Agent Platform API erfolgt.

Jeder Aufruf gibt die Antwort von Slurm direkt zurück, ohne dass ein Vorgang abgefragt werden muss. Sie können den Status eines Jobs im Zeitverlauf verfolgen. Weitere Informationen finden Sie unter Jobstatus prüfen.

API aufrufen

In den folgenden Abschnitten wird beschrieben, wie Sie die Slurm Job API verwenden.

Job senden

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

Ersetzen Sie Folgendes:

• REGION: Die Region, in der sich der Cluster befindet.

• PROJECT_ID: Ihre Projekt-ID.

• CLUSTER_ID: Ihre Cluster-ID.

• PARTITION: Die Slurm-Partition, mit der Sie eine Verbindung herstellen.

• USERNAME: Ihr Nutzername in der Slurm-Partition.

Der time_limit-Wert ist die Zeit in Minuten. Das Feld script enthält den Skriptinhalt und keinen Pfad zu einer Datei.

Die Slurm Job API gibt die ID des neuen Jobs in structuredBody zurück. Mit dem Befehl jq können Sie die Job-ID isolieren:

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'

Jobstatus prüfen

Verwenden Sie die folgenden Optionen, um einen laufenden Job zu überwachen. Führen Sie den Vorgang so lange aus, bis der Status einen beendeten Status erreicht, z. B. COMPLETED, FAILED oder 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"
  }'

Ersetzen Sie Folgendes:

• REGION: Die Region, in der sich der Cluster befindet.

• PROJECT_ID: Ihre Projekt-ID.

• CLUSTER_ID: Ihre Cluster-ID.

• JOB_ID: Die Job-ID des Slurm-Jobs.

Jobs auflisten

Verwenden Sie Folgendes, um laufende Jobs aufzulisten:

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

Ersetzen Sie Folgendes:

• REGION: Die Region, in der sich der Cluster befindet.

• PROJECT_ID: Ihre Projekt-ID.

• CLUSTER_ID: Ihre Cluster-ID.

Job abbrechen

Der Abbruch wird sofort akzeptiert, es dauert aber einen Moment, bis der Job beendet wird. Prüfen Sie den Status, um festzustellen, ob er CANCELLED lautet. Sie können nur Ihre eigenen Jobs abbrechen. Das Abbrechen eines abgeschlossenen Jobs ist sicher und gibt „success“ zurück.

So brechen Sie einen laufenden Job ab:

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

Ersetzen Sie Folgendes:

• REGION: Die Region, in der sich der Cluster befindet.

• PROJECT_ID: Ihre Projekt-ID.

• CLUSTER_ID: Ihre Cluster-ID.

• JOB_ID: Die Job-ID des Slurm-Jobs.

Abgeschlossenen Job anzeigen

Das Prüfen eines Jobs mit dem Pfad /slurm funktioniert nur, solange sich der Job noch in der Live-Warteschlange befindet. Verwenden Sie für abgeschlossene Jobs stattdessen den Pfad 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"
  }'

Ersetzen Sie Folgendes:

• REGION: Die Region, in der sich der Cluster befindet.

• PROJECT_ID: Ihre Projekt-ID.

• CLUSTER_ID: Ihre Cluster-ID.

• JOB_ID: Die Job-ID des Slurm-Jobs.

Slurm Job API über Python verwenden

Wenn Sie Python anstelle von curl bevorzugen, können Sie das folgende Hilfsskript verwenden, um Jobs zu senden, den Jobstatus zu prüfen und Jobs abzubrechen. Das Hilfsskript verwendet die google-auth-Bibliothek für Standardanmeldedaten für Anwendungen.

So verwenden Sie die Slurm Job API in Python:

1. Standardanmeldedaten für Anwendungen installieren:

   pip install google-auth requests
   

2. Verwenden Sie das folgende Python-Skript, um die API zu nutzen:

   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"])
   
   

Das Hilfsskript gibt den Status von Slurm und structuredBody zurück, sodass Sie Felder wie job_id oder job_state direkt lesen können. Wenn Sie das Skript in automatisierten Umgebungen verwenden möchten, authentifizieren Sie sich als Dienstkonto anstatt sich interaktiv anzumelden.

API-Antwort

Im Folgenden sehen Sie eine typische Antwort der Slurm Job API:

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

Das Feld status ist der HTTP-Status von Slurm und structuredBody enthält die Antwort von Slurm als JSON-Objekt, sodass Sie die Felder direkt lesen können, z. B. jq -r '.structuredBody.job_id'.

Prüfen Sie immer structuredBody, unabhängig vom status-Wert, da structuredBody alle Fehler oder Warnungen enthält, die von Slurm gemeldet werden.

Wenn die CallSlurmRestApi-Anfrage falsch formatiert ist, erhalten Sie möglicherweise ein error-Ergebnis wie im folgenden Beispiel:

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

Häufige Antworten

Limits

Die gesamte Anfrage darf nicht größer als 1 MB sein. Wenn Sie große Skripts verwenden möchten, legen Sie das Skript im freigegebenen Speicher des Clusters ab und rufen Sie es dann über ein kurzes Wrapper-Skript auf.