ניהול אשכולות אימון באמצעות Slurm Job API

‫Slurm Job API מאפשר לכם לנהל את אשכולות האימון של Gemini Enterprise Agent Platform באמצעות Agent Platform API, סקריפט, מחברת או צינור CI/CD. בעזרת Slurm Job API אפשר לנהל את העבודה של אשכול Slurm בלי להתחבר לאשכול באמצעות SSH.

דרישות מוקדמות

לפני שמשתמשים ב-Slurm Job API כדי לנהל את האשכול, צריך לוודא שמתקיימות הדרישות הבאות:

• יש לכם גישה לאשכולות האימון, בדרך כלל דרך התפקיד משתמש Vertex AI roles/aiplatform.user בפרויקט של האשכול, או כל תפקיד שכולל את ההרשאה aiplatform.googleapis.com/modelDevelopmentClusters.run.

• ‫Google Cloud CLI מותקן והגדרתם את פרטי הכניסה ב-Application Default Credentials ‏ (gcloud auth application-default login). כדי להשתמש בכלים אוטומטיים שלא יכולים להריץ התחברות אינטראקטיבית, צריך לבצע אימות כחשבון שירות עם הרשאה aiplatform.googleapis.com/modelDevelopmentClusters.run.

• אתם יודעים מה מזהה הפרויקט, האזור ומזהה האשכול שלכם.

בנוסף, מומלץ להגדיר כתובת אימייל נוחה:

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

סקירה כללית

כל שיחה היא בקשת POST לנקודת הקצה callSlurmRestApi של האשכול. בגוף הבקשה מציינים את השדות הבאים:

• ‫method: ה-method של HTTP ב-Slurm REST, כמו: HTTP_METHOD_GET,‏ HTTP_METHOD_POST או HTTP_METHOD_DELETE.

• ‫path: הנתיב של Slurm REST. לדוגמה, /slurm/v0.0.42/job/submit. משתמשים בגרסת Slurm שמופעלת באשכול. בדוגמאות במדריך הזה נעשה שימוש בגרסה v0.0.42.

• ‫body: אובייקט JSON עם המטען הייעודי (payload) של הבקשה. גוף הבקשה נדרש לבקשות HTTP_METHOD_POST. לא לכלול אותו בבקשות של HTTP_METHOD_GET ושל HTTP_METHOD_DELETE.

התשובה מכילה את השדות הבאים:

• ‫status: סטטוס ה-HTTP שהוחזר על ידי Slurm Job API. לדוגמה: 200.

• ‫structuredBody: התגובה מ-Slurm Job API כאובייקט JSON. לקרוא את השדות שאתם צריכים, כמו מזהה המשרה או רשימת המשרות, ישירות מ-structuredBody.

מריצים משימות בחשבון Linux שלכם באשכולות האימון באופן אוטומטי. אין הגדרות נוספות, ואי אפשר להריץ את העבודה כמשתמש אחר. לא צריך SSH, כי הכול מתבצע דרך Agent Platform API.

כל קריאה מחזירה את התגובה של Slurm ישירות, בלי פעולה של שליחת שאילתה. אתם יכולים לעקוב אחרי הסטטוס של משרה לאורך זמן. מידע נוסף זמין במאמר בנושא בדיקת הסטטוס של משימות.

שליחת קריאה ל-API

בקטעים הבאים מוסבר איך משתמשים ב-Slurm Job API.

שליחת משרה

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

מחליפים את מה שכתוב בשדות הבאים:

• ‫REGION: האזור שבו נמצא האשכול.

• ‫PROJECT_ID: מזהה הפרויקט.

• ‫CLUSTER_ID: מזהה האשכול.

• ‫PARTITION: המחיצה ב-Slurm שאליה מתחברים.

• ‫USERNAME: שם המשתמש במחיצת Slurm.

הערך time_limit הוא זמן בדקות. השדה script מכיל את התוכן של הסקריפט, ולא נתיב לקובץ.

‫Slurm Job API מחזיר את המזהה של המשימה החדשה ב-structuredBody. אפשר להשתמש בפקודה jq כדי לבודד את מזהה העבודה:

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'

בדיקת סטטוס המשימה

כדי לעקוב אחרי עבודה שפועלת: מבצעים סקר עד שהסטטוס מגיע למצב עצירה, כמו COMPLETED, ‏FAILED או 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"
  }'

מחליפים את מה שכתוב בשדות הבאים:

• ‫REGION: האזור שבו נמצא האשכול.

• ‫PROJECT_ID: מזהה הפרויקט.

• ‫CLUSTER_ID: מזהה האשכול.

• ‫JOB_ID: מזהה המשימה של Slurm.

פרסום משרות

כדי לראות רשימה של משימות שפועלות:

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

מחליפים את מה שכתוב בשדות הבאים:

• ‫REGION: האזור שבו נמצא האשכול.

• ‫PROJECT_ID: מזהה הפרויקט.

• ‫CLUSTER_ID: מזהה האשכול.

ביטול משימה

הביטול מתקבל מיד, אבל לוקח רגע עד שהעבודה מפסיקה. בודקים את הסטטוס שלו כדי לוודא שהסטטוס הוא CANCELLED. אתם יכולים לבטל רק את העבודות שלכם. ביטול של עבודה שהושלמה הוא בטוח ומחזיר הצלחה.

כדי לבטל עבודה שפועלת:

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

מחליפים את מה שכתוב בשדות הבאים:

• ‫REGION: האזור שבו נמצא האשכול.

• ‫PROJECT_ID: מזהה הפרויקט.

• ‫CLUSTER_ID: מזהה האשכול.

• ‫JOB_ID: מזהה המשימה של Slurm.

הצגת עבודה שהסתיימה

בדיקת משימה באמצעות הנתיב /slurm פועלת רק בזמן שהמשימה עדיין בתור הפעיל. לגבי משימות שהושלמו, משתמשים בנתיב 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"
  }'

מחליפים את מה שכתוב בשדות הבאים:

• ‫REGION: האזור שבו נמצא האשכול.

• ‫PROJECT_ID: מזהה הפרויקט.

• ‫CLUSTER_ID: מזהה האשכול.

• ‫JOB_ID: מזהה המשימה של Slurm.

שימוש ב-Slurm Job API מ-Python

אם אתם מעדיפים להשתמש ב-Python במקום ב-curl, אתם יכולים להשתמש בסקריפט העזר הבא כדי לשלוח משימות, לבדוק את סטטוס המשימות ולבטל משימות. סקריפט העזר משתמש בספריית google-auth ל-Application Default Credentials.

כדי להשתמש ב-Slurm Job API מ-Python:

1. מתקינים Application Default Credentials:

   pip install google-auth requests
   

2. כדי להשתמש ב-API, אפשר להשתמש בסקריפט Python הבא:

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

סקריפט העזר מחזיר את הסטטוס מ-Slurm ואת structuredBody, כך שאפשר לקרוא שדות כמו job_id או job_state ישירות. כדי להשתמש בסקריפט בסביבות אוטומטיות, צריך לבצע אימות כחשבון שירות במקום להשתמש בהתחברות אינטראקטיבית.

תגובה מה-API

זוהי תשובה אופיינית של Slurm Job API:

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

השדה status הוא סטטוס ה-HTTP מ-Slurm, והשדה structuredBody מכיל את התגובה של Slurm כאובייקט JSON, כך שאפשר לקרוא את השדות שלו ישירות, כמו: jq -r '.structuredBody.job_id'.

תמיד בודקים את structuredBody בלי קשר לערך של status, כי structuredBody כולל את כל השגיאות או האזהרות ש-Slurm מדווח עליהן.

אם הבקשה CallSlurmRestApi לא מעוצבת כראוי, יכול להיות שתקבלו תוצאה error כמו בדוגמה הבאה:

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

תשובות נפוצות

מגבלות

הגודל של הבקשה כולה צריך להיות פחות מ-1MB. כדי להשתמש בסקריפטים גדולים, צריך למקם את הסקריפט באחסון המשותף של האשכול ואז לקרוא לו מסקריפט wrapper קצר.