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

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

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

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

  • יש לכם גישה לאשכולות האימון, בדרך כלל באמצעות התפקיד משתמש בפלטפורמת הסוכן 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 לנקודת הקצה (endpoint) של callSlurmRestApi באשכול. בגוף הבקשה, מציינים את השדות הבאים:

  • ‫method: ה-method של Slurm REST HTTP, למשל: 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. אתם יכולים לבטל רק את המשימות שלכם. אפשר לבטל משימה שהושלמה, והיא תחזיר את הערך success.

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

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

הדוגמה הבאה היא תגובה אופיינית מה-API של Slurm Job:

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

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

מה רואים מה המשמעות של
סטטוס 200 עם מזהה משימה המשימה נשלחה.
סטטוס 200, אין שגיאות הבקשה אושרה.
סטטוס 200 עם אזהרות הבקשה בוצעה, אבל Slurm התעלם ממשהו או שינה משהו. קוראים את האזהרה.
סטטוס 404 למשימה המשימה לא נמצאת בתור הפעיל, יכול להיות שהיא הסתיימה. אפשר לנסות את נתיב ההיסטוריה.
סטטוס 400 או 500 ‫Slurm דחה את הבקשה. הסיבות לדחייה יכולות להיות גוף בקשה בפורמט שגוי או שדה מסוג שגוי. קוראים את השגיאה ב-structuredBody.
שגיאת PERMISSION_DENIED אין לך גישה לאשכול הזה.
שגיאת NOT_FOUND שם האשכול שגוי או שהאשכול לא קיים.
שגיאת INVALID_ARGUMENT הנתיב לא תקין ב-Slurm, או שהבקשה גדולה מדי.
שגיאת UNAVAILABLE אי אפשר לגשת לאשכול באופן זמני. נסו שוב אחרי המתנה קצרה.

מגבלות

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