Mengelola cluster pelatihan menggunakan Slurm Job API

Slurm Job API memungkinkan Anda mengelola cluster pelatihan Gemini Enterprise Agent Platform melalui Agent Platform API, skrip, notebook, atau pipeline CI/CD. Slurm Job API memungkinkan Anda mengelola tugas cluster Slurm tanpa terhubung ke cluster dengan SSH.

Prasyarat

Sebelum menggunakan Slurm Job API untuk mengelola cluster, pastikan Anda memenuhi persyaratan berikut:

• Anda memiliki akses ke cluster pelatihan, biasanya melalui peran Vertex AI User roles/aiplatform.user di project cluster, atau peran apa pun yang mencakup izin aiplatform.googleapis.com/modelDevelopmentClusters.run.

• Google Cloud CLI telah diinstal dan Anda telah mengonfigurasi Kredensial Default Aplikasi (gcloud auth application-default login). Untuk alat otomatis yang tidak dapat menjalankan login interaktif, lakukan autentikasi sebagai akun layanan dengan izin aiplatform.googleapis.com/modelDevelopmentClusters.run.

• Anda mengetahui project ID, region, dan ID cluster Anda.

Selain itu, sebaiknya siapkan alias praktis:

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

Ringkasan

Setiap panggilan adalah permintaan POST ke endpoint callSlurmRestApi untuk cluster Anda. Dalam isi permintaan, Anda memberikan kolom berikut:

• method: Metode HTTP REST Slurm, seperti: HTTP_METHOD_GET, HTTP_METHOD_POST, atau HTTP_METHOD_DELETE.

• path: Jalur REST Slurm. Misalnya, /slurm/v0.0.42/job/submit. Gunakan versi Slurm yang dijalankan cluster Anda. Contoh dalam panduan ini menggunakan v0.0.42.

• body: Objek JSON dengan payload permintaan. Isi diperlukan untuk permintaan HTTP_METHOD_POST. Jangan sertakan untuk permintaan HTTP_METHOD_GET dan HTTP_METHOD_DELETE.

Respons berisi kolom berikut:

• status: Status HTTP yang ditampilkan oleh Slurm Job API. Contohnya, 200.

• structuredBody: Respons dari Slurm Job API, sebagai objek JSON. Baca kolom yang Anda butuhkan, seperti ID tugas atau daftar tugas, langsung dari structuredBody.

Anda menjalankan tugas di akun Linux Anda sendiri di cluster pelatihan secara otomatis. Tidak ada konfigurasi tambahan, dan Anda tidak dapat menjalankan tugas sebagai pengguna lain. SSH tidak diperlukan, karena semuanya dilakukan melalui Agent Platform API.

Setiap panggilan menampilkan respons Slurm secara langsung, tanpa operasi untuk melakukan polling. Anda dapat memantau status tugas dari waktu ke waktu. Untuk mengetahui informasi selengkapnya, lihat Periksa status tugas.

Memanggil API

Bagian berikut menunjukkan cara menggunakan Slurm Job API.

Mengirim tugas

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

Ganti kode berikut:

• REGION: Region tempat cluster Anda berada.

• PROJECT_ID: Project ID Anda.

• CLUSTER_ID: ID cluster Anda.

• PARTITION: Partisi Slurm yang Anda hubungkan.

• USERNAME: Nama pengguna Anda di partisi Slurm.

Nilai time_limit adalah waktu dalam menit. Kolom script berisi konten skrip, bukan jalur ke file.

Slurm Job API menampilkan ID tugas baru dalam structuredBody. Anda dapat menggunakan perintah jq untuk mengisolasi ID tugas:

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'

Memeriksa status tugas

Gunakan perintah berikut untuk memantau tugas yang sedang berjalan. Polling hingga status mencapai status berhenti, seperti COMPLETED, FAILED, atau 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"
  }'

Ganti kode berikut:

• REGION: Region tempat cluster Anda berada.

• PROJECT_ID: Project ID Anda.

• CLUSTER_ID: ID cluster Anda.

• JOB_ID: ID tugas Slurm.

Membuat daftar tugas

Gunakan perintah berikut untuk mencantumkan tugas yang sedang berjalan:

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

Ganti kode berikut:

• REGION: Region tempat cluster Anda berada.

• PROJECT_ID: Project ID Anda.

• CLUSTER_ID: ID cluster Anda.

Membatalkan tugas

Pembatalan akan segera diterima, tetapi tugas memerlukan waktu beberapa saat untuk berhenti. Periksa statusnya untuk memverifikasi bahwa statusnya adalah CANCELLED. Anda hanya dapat membatalkan tugas Anda sendiri. Membatalkan tugas yang telah selesai aman dan akan menampilkan keberhasilan.

Gunakan perintah berikut untuk membatalkan tugas yang sedang berjalan:

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

Ganti kode berikut:

• REGION: Region tempat cluster Anda berada.

• PROJECT_ID: Project ID Anda.

• CLUSTER_ID: ID cluster Anda.

• JOB_ID: ID tugas Slurm.

Menampilkan tugas yang selesai

Memeriksa tugas dengan jalur /slurm hanya berfungsi saat tugas masih dalam antrean aktif. Untuk tugas yang telah selesai, gunakan jalur 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"
  }'

Ganti kode berikut:

• REGION: Region tempat cluster Anda berada.

• PROJECT_ID: Project ID Anda.

• CLUSTER_ID: ID cluster Anda.

• JOB_ID: ID tugas Slurm.

Menggunakan Slurm Job API dari Python

Jika lebih suka menggunakan Python daripada curl, Anda dapat menggunakan skrip helper berikut untuk mengirimkan tugas, memeriksa status tugas, dan membatalkan tugas. Skrip helper menggunakan library google-auth untuk Kredensial Default Aplikasi.

Lakukan hal berikut untuk menggunakan Slurm Job API dari Python:

1. Instal Kredensial Default Aplikasi:

   pip install google-auth requests
   

2. Gunakan skrip Python berikut untuk menggunakan 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"])
   
   

Skrip helper menampilkan status dari Slurm dan structuredBody, sehingga Anda dapat membaca kolom seperti job_id atau job_state secara langsung. Untuk menggunakan skrip di lingkungan otomatis, lakukan autentikasi sebagai akun layanan, bukan menggunakan login interaktif.

Respons API

Berikut adalah respons Slurm Job API yang umum:

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

Kolom status adalah status HTTP dari Slurm dan structuredBody berisi respons Slurm sebagai objek JSON, sehingga Anda dapat membaca kolomnya secara langsung, seperti: jq -r '.structuredBody.job_id'.

Selalu periksa structuredBody terlepas dari nilai status, karena structuredBody mencakup semua error atau peringatan yang dilaporkan Slurm.

Jika permintaan CallSlurmRestApi salah format, Anda mungkin mendapatkan hasil error seperti contoh berikut:

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

Respons umum

Batas

Seluruh permintaan harus tetap berukuran di bawah 1 MB. Untuk menggunakan skrip besar, tempatkan skrip di penyimpanan bersama cluster, lalu panggil dari skrip wrapper pendek.