使用 Slurm 作业 API 管理训练集群

借助 Slurm 作业 API，您可以通过 Agent Platform API、脚本、笔记本或 CI/CD 流水线来管理 Gemini Enterprise Agent Platform 训练集群。借助 Slurm 作业 API，您无需通过 SSH 连接到集群即可管理 Slurm 集群作业。

前提条件

在使用 Slurm 作业 API 管理集群之前，请确保满足以下要求：

  • 您有权访问训练集群，通常是通过集群项目中的 Vertex AI User 角色 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 作业 API 返回的 HTTP 状态。例如 200。

  • body：来自 Slurm 作业 API 的响应，以 JSON 字符串的形式呈现。解析正文字段，找到所需的字段，例如作业 ID 或作业列表。

您可以在训练集群上自动运行自己 Linux 账号中的作业。无需进行其他配置，并且您无法以其他用户身份运行作业。无需使用 SSH，因为所有操作都是通过 Agent Platform API 完成的。

每次调用都会直接返回 Slurm 的响应，而无需轮询操作。您可以随时跟踪作业的状态。如需了解详情，请参阅检查作业状态。

调用 API

以下部分介绍了如何使用 Slurm 作业 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 作业 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 作业 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 作业 API 响应：

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

status 字段是来自 Slurm 的 HTTP 状态，body 字段是 JSON 字符串。您可以使用 jq 解析 body 字段，例如：jq -r '.body | fromjson'。

我们建议您检查 body 字段，无论 Slurm 状态如何，因为该正文字段包含 Slurm 报告的任何错误或警告。

如果 CallSlurmRestApi 请求本身格式有误，您可能会获得如下所示的 error 结果：

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

常见回答

看到的内容 具体含义
状态 200，并包含作业 ID 作业已成功提交。
状态 200，无错误 请求成功。
状态 200，但有警告 请求成功，但 Slurm 忽略或调整了某些内容。阅读警告。
作业的状态 404 作业不在实时队列中；可能已完成。尝试历史记录路径。
状态 400 或 500 Slurm 拒绝了请求。拒绝的原因可能包括正文格式错误或字段类型错误。读取正文中的错误。
PERMISSION_DENIED 错误 您无权访问此集群。
NOT_FOUND 错误 集群名称有误或集群不存在。
INVALID_ARGUMENT 错误 路径不是有效的 Slurm 路径，或者请求过大。
UNAVAILABLE 错误 集群暂时无法访问。稍等片刻后重试。

限制

整个请求的大小必须小于 1 MB。如需使用大型脚本，请将脚本放在集群的共享存储空间中，然后通过简短的封装脚本调用该脚本。