Gérer les clusters d'entraînement à l'aide de l'API Slurm Job

L'API Slurm Job vous permet de gérer vos clusters d'entraînement Gemini Enterprise Agent Platform via l'API Agent Platform, un script, un notebook ou un pipeline CI/CD. L'API Slurm Job vous permet de gérer le job du cluster Slurm sans vous connecter au cluster avec SSH.

Prérequis

Avant d'utiliser l'API Slurm Job pour gérer votre cluster, assurez-vous de remplir les conditions suivantes :

  • Vous avez accès aux clusters d'entraînement, généralement via le rôle Utilisateur Agent Platform roles/aiplatform.user sur le projet du cluster, ou tout rôle incluant l'autorisation aiplatform.googleapis.com/modelDevelopmentClusters.run.

  • Google Cloud CLI est installé et vous avez configuré les identifiants par défaut de l'application (gcloud auth application-default login). Pour les outils automatisés qui ne peuvent pas exécuter de connexion interactive, authentifiez-vous en tant que compte de service avec l'autorisation aiplatform.googleapis.com/modelDevelopmentClusters.run.

  • Vous connaissez l'ID de projet, la région et l'ID de votre cluster.

Nous vous recommandons également de configurer un alias pratique :

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

Présentation

Chaque appel est une requête POST envoyée au point de terminaison callSlurmRestApi de votre cluster. Dans le corps de la requête, vous fournissez les champs suivants :

  • method : méthode HTTP REST Slurm, telle que HTTP_METHOD_GET, HTTP_METHOD_POST ou HTTP_METHOD_DELETE.

  • path : chemin d'accès REST Slurm. Par exemple, /slurm/v0.0.42/job/submit. Utilisez la version de Slurm exécutée par votre cluster. Les exemples de ce guide utilisent la version 0.0.42.

  • body : objet JSON avec la charge utile de la requête. Le corps est obligatoire pour les requêtes HTTP_METHOD_POST. Omettez-le pour les requêtes HTTP_METHOD_GET et HTTP_METHOD_DELETE.

La réponse contient les champs suivants :

  • status : état HTTP renvoyé par l'API Slurm Job. Exemple : 200.

  • structuredBody : réponse de l'API Slurm Job, sous la forme d'un objet JSON. Lisez les champs dont vous avez besoin, tels que l'ID du job ou la liste des jobs, directement à partir de structuredBody.

Vous exécutez automatiquement des jobs sur votre propre compte Linux sur les clusters d'entraînement. Aucune configuration supplémentaire n'est requise, et vous ne pouvez pas exécuter le job en tant qu'autre utilisateur. SSH n'est pas requis, car tout est effectué via l'API Agent Platform.

Chaque appel renvoie directement la réponse de Slurm, sans qu'il soit nécessaire d'interroger une opération. Vous pouvez suivre l'état d'un job au fil du temps. Pour en savoir plus, consultez Vérifier l'état d'un job.

Appeler l'API

Les sections suivantes expliquent comment utiliser l'API Slurm Job.

Envoyer une tâche

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

Remplacez les éléments suivants :

  • REGION : région dans laquelle se trouve votre cluster.

  • PROJECT_ID : ID du projet.

  • CLUSTER_ID : ID de votre cluster.

  • PARTITION : partition Slurm à laquelle vous vous connectez.

  • USERNAME : votre nom d'utilisateur sur la partition Slurm.

La valeur time_limit est exprimée en minutes. Le champ script contient le contenu du script, et non le chemin d'accès à un fichier.

L'API Slurm Job renvoie l'ID du nouveau job dans structuredBody. Vous pouvez utiliser la commande jq pour isoler l'ID de 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'

Vérifier l'état d'une tâche

Utilisez les éléments suivants pour surveiller un job en cours d'exécution. Interrogez l'état jusqu'à ce qu'il atteigne un état d'arrêt, tel que 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"
  }'

Remplacez les éléments suivants :

  • REGION : région dans laquelle se trouve votre cluster.

  • PROJECT_ID : ID du projet.

  • CLUSTER_ID : ID de votre cluster.

  • JOB_ID : ID du job Slurm.

Répertorier des tâches

Utilisez la commande suivante pour lister les jobs en cours d'exécution :

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

Remplacez les éléments suivants :

  • REGION : région dans laquelle se trouve votre cluster.

  • PROJECT_ID : ID du projet.

  • CLUSTER_ID : ID de votre cluster.

Annuler une mission

L'annulation est acceptée immédiatement, mais l'arrêt du job prend un certain temps. Vérifiez son état pour vous assurer qu'il est défini sur CANCELLED. Vous ne pouvez annuler que vos propres jobs. L'annulation d'un job terminé est sans risque et renvoie un message de réussite.

Utilisez la commande suivante pour annuler un job en cours d'exécution :

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

Remplacez les éléments suivants :

  • REGION : région dans laquelle se trouve votre cluster.

  • PROJECT_ID : ID du projet.

  • CLUSTER_ID : ID de votre cluster.

  • JOB_ID : ID du job Slurm.

Afficher un job terminé

La vérification d'un job avec le chemin d'accès /slurm ne fonctionne que lorsque le job est toujours dans la file d'attente active. Pour les jobs terminés, utilisez plutôt le chemin d'accès 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"
  }'

Remplacez les éléments suivants :

  • REGION : région dans laquelle se trouve votre cluster.

  • PROJECT_ID : ID du projet.

  • CLUSTER_ID : ID de votre cluster.

  • JOB_ID : ID du job Slurm.

Utiliser l'API Slurm Job depuis Python

Si vous préférez utiliser Python plutôt que curl, vous pouvez utiliser le script d'assistance suivant pour envoyer des jobs, vérifier leur état et les annuler. Le script d'assistance utilise la bibliothèque google-auth pour les Identifiants par défaut de l'application.

Pour utiliser l'API Slurm Job depuis Python, procédez comme suit :

  1. Installez les identifiants par défaut de l'application :

    pip install google-auth requests

  2. Utilisez le script Python suivant pour utiliser 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"])

Le script d'assistance renvoie l'état de Slurm et de structuredBody, ce qui vous permet de lire directement des champs tels que job_id ou job_state. Pour utiliser le script dans des environnements automatisés, authentifiez-vous en tant que compte de service au lieu d'utiliser une connexion interactive.

Réponse de l'API

Voici une réponse typique de l'API Slurm Job :

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

Le champ status correspond à l'état HTTP de Slurm et structuredBody contient la réponse de Slurm sous la forme d'un objet JSON. Vous pouvez donc lire ses champs directement, par exemple : jq -r '.structuredBody.job_id'.

Vérifiez toujours structuredBody, quelle que soit la valeur de status, car structuredBody inclut toutes les erreurs ou tous les avertissements signalés par Slurm.

Si la requête CallSlurmRestApi est mal formatée, vous pouvez obtenir un résultat error, comme dans l'exemple suivant :

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

Réponses courantes

Ce que vous voyez Signification
État 200 avec un ID de job Le job a bien été envoyé.
État 200, aucune erreur La requête a abouti.
État 200 avec des avertissements La requête a abouti, mais Slurm a ignoré ou ajusté un élément. Lisez l'avertissement.
Code d'état 404 pour un job Le job n'est pas dans la file d'attente en direct. Il est peut-être terminé. Essayez le chemin d'accès de l'historique.
État 400 ou 500 Slurm a rejeté la requête. Le refus peut être dû à un corps incorrect ou à un champ de type incorrect. Lisez l'erreur dans structuredBody.
Erreur PERMISSION_DENIED Vous n'avez pas accès à ce cluster.
Erreur NOT_FOUND Le nom du cluster est incorrect ou le cluster n'existe pas.
Erreur INVALID_ARGUMENT Le chemin d'accès n'est pas un chemin d'accès Slurm valide, ou la requête est trop volumineuse.
Erreur UNAVAILABLE Le cluster est temporairement inaccessible. Réessayez après un court délai.

Limites

La taille de la requête entière doit être inférieure à 1 Mo. Pour utiliser des scripts volumineux, placez le script sur l'espace de stockage partagé du cluster, puis appelez-le à partir d'un script wrapper court.