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 Vertex AI roles/aiplatform.user dans 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 du projet, la région et l'ID de cluster 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 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 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 contenant 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 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 nécessaire, car tout se fait via l'API Agent Platform.

Chaque appel renvoie directement la réponse de Slurm, sans opération à interroger. 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 de votre 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 correspond à une duré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 du 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 jusqu'à ce que l'état atteigne un état 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 de votre projet

• CLUSTER_ID : ID de votre cluster.

• JOB_ID : ID du job Slurm.

Répertorier des tâches

Pour lister les jobs en cours d'exécution, utilisez la commande suivante :

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 de votre projet

• CLUSTER_ID : ID de votre cluster.

Annuler une mission

L'annulation est acceptée immédiatement, mais l'arrêt de la tâche prend un certain temps. Vérifiez son état pour vous assurer qu'il est CANCELLED. Vous ne pouvez annuler que vos propres tâches. L'annulation d'un job terminé est sans risque et renvoie une réponse positive.

Pour annuler une tâche en cours d'exécution, procédez comme suit :

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 de votre 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 encore dans la file d'attente en direct. 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 de votre 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 à partir de 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'aide renvoie l'état de Slurm et 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 forme d'objet JSON. Vous pouvez donc lire directement ses champs, tels que 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

Limites

La taille totale de la requête ne doit pas dépasser 1 Mo. Pour utiliser des scripts volumineux, placez-les dans le stockage partagé du cluster, puis appelez-les à partir d'un script wrapper court.