Slurm Job API를 사용하여 학습 클러스터 관리

Slurm Job API를 사용하면 Agent Platform API, 스크립트, 노트북 또는 CI/CD 파이프라인을 통해 Gemini Enterprise Agent Platform 학습 클러스터를 관리할 수 있습니다. Slurm 작업 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 객체). structuredBody에서 작업 ID 또는 작업 목록과 같은 필요한 필드를 직접 읽습니다.

학습 클러스터의 자체 Linux 계정에서 작업을 자동으로 실행합니다. 추가 구성이 없으며 다른 사용자로 작업을 실행할 수 없습니다. 모든 작업이 Agent Platform API를 통해 이루어지므로 SSH는 필요하지 않습니다.

각 호출은 폴링할 작업 없이 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는 structuredBody에서 새 작업의 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 '.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 작업 API 사용

curl 대신 Python을 사용하려면 다음 도우미 스크립트를 사용하여 작업을 제출하고, 작업 상태를 확인하고, 작업을 취소하면 됩니다. 헬퍼 스크립트는 애플리케이션 기본 사용자 인증 정보에 google-auth 라이브러리를 사용합니다.

Python에서 Slurm 작업 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 작업 API 응답입니다.

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

status 필드는 Slurm의 HTTP 상태이고 structuredBody에는 Slurm의 응답이 JSON 객체로 포함되어 있으므로 jq -r '.structuredBody.job_id'와 같은 필드를 직접 읽을 수 있습니다.

status 값과 관계없이 항상 structuredBody를 확인하세요. structuredBody에는 Slurm에서 보고하는 오류 또는 경고가 포함되어 있기 때문입니다.

CallSlurmRestApi 요청의 형식이 잘못된 경우 다음 예와 같이 error 결과가 표시될 수 있습니다.

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

일반적인 대답

한도

전체 요청의 크기는 1MB 미만이어야 합니다. 큰 스크립트를 사용하려면 클러스터의 공유 스토리지에 스크립트를 배치한 다음 짧은 래퍼 스크립트에서 호출합니다.