使用 Slurm Job API 管理訓練叢集

透過 Slurm Job API，您可以使用 Agent Platform API、指令碼、筆記本或 CI/CD 管道，管理 Gemini Enterprise Agent Platform 訓練叢集。您可以使用 Slurm Job API 管理 Slurm 叢集工作，不必透過 SSH 連線至叢集。

必要條件

使用 Slurm Job API 管理叢集前，請確認您符合下列需求：

  • 您有權存取訓練叢集，通常是透過叢集專案的 Vertex AI 使用者角色 roles/aiplatform.user，或是任何包含 aiplatform.googleapis.com/modelDevelopmentClusters.run 權限的角色。

  • 已安裝 Google Cloud CLI，且您已設定應用程式預設憑證 (gcloud auth application-default login)。對於無法執行互動式登入的自動化工具，請以具備 aiplatform.googleapis.com/modelDevelopmentClusters.run 權限的服務帳戶進行驗證。

  • 您知道叢集的專案 ID、地區和叢集 ID。

此外，建議您設定便利別名：

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

總覽

每次呼叫都是對叢集 callSlurmRestApi 端點發出的 POST 要求。在要求主體中，您會提供下列欄位：

  • 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 物件。HTTP_METHOD_POST 請求必須包含主體。HTTP_METHOD_GET 和 HTTP_METHOD_DELETE 要求則可省略。

回應包含下列欄位：

  • status：Slurm Job API 傳回的 HTTP 狀態。例如：200。

  • body：Slurm Job API 的回應，以 JSON 字串形式呈現。剖析主體欄位，找出所需欄位，例如工作 ID 或工作清單。

您會在訓練叢集上，以自己的 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：專案 ID。

  • CLUSTER_ID：叢集 ID。

  • PARTITION：您要連線的 Slurm 分區。

  • USERNAME：Slurm 分割區的使用者名稱。

time_limit 值是以分鐘為單位的時間。script 欄位會保留指令碼內容，而非檔案路徑。

Slurm Job API 會在回應主體中傳回新工作的 ID。您可以使用 jq 指令來隔離工作 ID：

gcurl -sS -X POST "https://REGION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/REGION/modelDevelopmentClusters/CLUSTER_ID:callSlurmRestApi" \
  -d '{ ... }' | jq -r '.body | fromjson | .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：專案 ID。

  • CLUSTER_ID：叢集 ID。

  • JOB_ID：Slurm 工作的 ID。

列出工作

使用下列指令列出執行中的工作：

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：專案 ID。

  • CLUSTER_ID：叢集 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：專案 ID。

  • CLUSTER_ID：叢集 ID。

  • JOB_ID：Slurm 工作的 ID。

顯示已完成的工作

只有在工作仍在即時佇列中時，才能使用 /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：專案 ID。

  • CLUSTER_ID：叢集 ID。

  • JOB_ID：Slurm 工作的 ID。

從 Python 使用 Slurm Job API

如果您偏好使用 Python 而非 curl，可以使用下列輔助指令碼提交工作、檢查工作狀態及取消工作。輔助指令碼會使用 google-auth 程式庫取得應用程式預設憑證。

如要透過 Python 使用 Slurm Job API，請按照下列步驟操作：

  1. 安裝應用程式預設憑證：

    pip install google-auth requests

  2. 請使用下列 Python 指令碼來使用 API：

    import json
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()
    # body is a JSON string; parse it to read Slurm's fields.
    return envelope["status"], json.loads(envelope["body"])

# 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 的狀態和已剖析的主體，因此您可以直接讀取 job_id 或 job_state 等欄位。如要在自動化環境中使用指令碼，請以服務帳戶身分進行驗證，而非使用互動式登入。

API 回應

以下是典型的 Slurm Job API 回應：

{
  "status": 200,
  "body": "{ ... raw Slurm JSON ... }"
}

status 欄位是 Slurm 的 HTTP 狀態，body 欄位則是 JSON 字串。您可以使用 jq 剖析 body 欄位，例如：jq -r '.body | fromjson'。

建議您檢查 body 欄位，因為主體欄位包含 Slurm 報告的任何錯誤或警告。

如果 CallSlurmRestApi 要求本身格式錯誤，您可能會收到如下的 error 結果：

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

常見回覆

顯示內容 說明
狀態 200，並提供工作 ID 已成功提交工作。
狀態 200，沒有錯誤 要求成功。
狀態 200，但附帶警告 要求順利完成，但 Slurm 忽略或調整了某些項目。閱讀警告訊息。
工作狀態 404 工作不在即時佇列中，可能已完成。請嘗試使用歷史記錄 路徑。
狀態 400 或 500 Slurm 拒絕了這項要求。遭拒的原因可能包括主體格式錯誤或欄位類型錯誤。Read the error in the body.
PERMISSION_DENIED 錯誤 您無權存取這個叢集。
NOT_FOUND 錯誤 叢集名稱有誤或叢集不存在。
INVALID_ARGUMENT 錯誤 路徑不是有效的 Slurm 路徑，或要求過大。
UNAVAILABLE 錯誤 叢集暫時無法連線。請稍候片刻再試。

限制

整個要求的大小不得超過 1 MB。如要使用大型指令碼，請將指令碼放在叢集的共用儲存空間，然後從簡短的包裝函式指令碼呼叫該指令碼。