Slurm Job API を使用してトレーニング クラスタを管理する

Slurm Job API を使用すると、Agent Platform API、スクリプト、ノートブック、CI/CD パイプラインを介して Gemini Enterprise Agent Platform トレーニング クラスタを管理できます。Slurm Job API を使用すると、SSH でクラスタに接続せずに Slurm クラスタ ジョブを管理できます。

前提条件

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

  • structuredBody: Slurm Job API からのレスポンス（JSON オブジェクト）。ジョブ ID やジョブリストなど、必要なフィールドを structuredBody から直接読み取ります。

トレーニング クラスタの独自の Linux アカウントでジョブを自動的に実行します。追加の構成はなく、別のユーザーとしてジョブを実行することはできません。すべてが Agent Platform API を介して行われるため、SSH は必要ありません。

各呼び出しは、ポーリングするオペレーションなしで、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 を structuredBody で返します。jq コマンドを使用して、ジョブ ID を分離できます。

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: プロジェクト 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 を使用する

curl ではなく Python を使用する場合は、次のヘルパースクリプトを使用して、ジョブの送信、ジョブのステータスの確認、ジョブのキャンセルを行うことができます。ヘルパースクリプトは、アプリケーションのデフォルト認証情報に google-auth ライブラリを使用します。

Python から Slurm Job API を使用するには、次の操作を行います。

  1. アプリケーションのデフォルト認証情報をインストールします。

    pip install google-auth requests

  2. 次の Python スクリプトを使用して 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"])

ヘルパー スクリプトは Slurm と structuredBody からステータスを返すため、job_id や job_state などのフィールドを直接読み取ることができます。自動化された環境でスクリプトを使用する場合は、インタラクティブ ログインを使用するのではなく、サービス アカウントとして認証します。

API レスポンス

一般的な Slurm Job API レスポンスは次のとおりです。

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

status フィールドは Slurm の HTTP ステータスで、structuredBody には Slurm のレスポンスが JSON オブジェクトとして含まれています。そのため、jq -r '.structuredBody.job_id' などのフィールドを直接読み取ることができます。

structuredBody には Slurm が報告するエラーや警告が含まれているため、status の値に関係なく、常に structuredBody を確認してください。

CallSlurmRestApi リクエストの形式が正しくない場合、次の例のように error 結果が返されることがあります。

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

一般的な回答

表示される内容 意味
ステータス 200 とジョブ ID ジョブが正常に送信されました。
ステータス 200、エラーなし リクエストが成功しました。
ステータス 200（警告あり） リクエストは成功しましたが、Slurm が何かを無視または調整しました。警告を読みます。
ジョブのステータス 404 ジョブはライブキューにありません。完了した可能性があります。履歴パスを試してください。
ステータス 400 または 500 Slurm がリクエストを拒否しました。拒否の理由としては、不正な形式の本文や、型が間違っているフィールドなどが考えられます。structuredBody のエラーを読み取ります。
PERMISSION_DENIED エラー このクラスタへのアクセス権がありません。
NOT_FOUND エラー クラスタ名が間違っているか、クラスタが存在しません。
INVALID_ARGUMENT エラー パスが有効な Slurm パスではないか、リクエストが大きすぎます。
UNAVAILABLE エラー クラスタに一時的にアクセスできません。しばらく待ってから再試行します。

上限

リクエスト全体のサイズは 1 MB 未満にする必要があります。大きなスクリプトを使用するには、スクリプトをクラスタの共有ストレージに配置し、短いラッパー スクリプトから呼び出します。