Rescale のジョブを任意の経過時間後に自動停止する

この記事について

この記事では, Rescale REST API を使用して, Rescale のジョブを任意の経過時間後に停止させる作業を自動化します

事前準備: Embedded Job API Key

この記事中のジョブでは, Embedded Job API Key を使用します。この機能を有効化していない場合は, Rescale の人（営業、技術）に相談したり, Rescale Support にサポートリクエストを出したりして有効化します

背景

Rescale のジョブを, 途中で自動的に止めたいというケースがあるかと思います。例えば、ベンチマークテストをするときに、ジョブを最後まで実行せずに、所定の時間内でのシミュレーションの進捗度合いによってパフォーマンスを評価する場合などです。1時間単位で自動的に止めたい場合は Walltime で設定できますが、例えば10分や30分、90分などでジョブを自動的に止めたい場合は, Rescale に備わっている機能だけでは実現できません。そこで, Rescale REST API を使用して、ジョブのステータス確認および所定時間経過後の停止を自動化します

全体の流れ

この記事では、ジョブのステータス確認と停止を行う Python スクリプトを, Rescale の独立したジョブとして実行します

(1) ユーザーは、途中で止めたいジョブを設定し保存する（ Submit はしない）。途中で止めたいジョブの Job ID が発行される。この Job ID を FgHiJ とする
(2) ユーザーは、ジョブのステータス確認と停止を（行うスクリプトを実行）するジョブを submit する。このジョブの Job ID を AbCdE とする
(3) ジョブ AbCdE のクラスタ起動が完了すると、ジョブコマンド中で job_submit_stop.py が実行される
(4) job_submit_stop.py は, API でジョブ FgHiJ を submit し、ステータスを polling する
(5) ジョブ FgHiJ のコマンドが開始されると, job_submit_stop.py は設定された時間だけ待機したのち、ジョブ FgHiJ を stop させる
(6) job_submit_stop.py は終了し、ジョブ AbCdE は終了する

ジョブ設定

ジョブのステータス確認と停止をするジョブ（上記例ではジョブ AbCdE ）は、以下の設定で実行します

• Inputs
  • job_submit_stop.py
• Software Settings
  • Bring Your Own Software (CPU)
  • Command

python3 job_submit_stop.py <JobId> <JobDurationSec>

• Hardware Settings
  • Starlite 1 cores
  • Walltime: 充分長く設定

コマンドは、例えばジョブ FgHiJ を10分（600秒）で自動停止させたい場合は、以下のようになります

python3 job_submit_stop.py FgHiJ 600

job_submit_stop.py は以下の内容です

import json
import os
import requests
import sys
import time

api_key = os.environ["RESCALE_API_KEY"]
api_key_in_header = "Token " + api_key
platform = os.environ["RESCALE_API_BASE_URL"]

args = sys.argv
draft_jobid = args[1]
job_duration_sec = int(args[2])

print("Submit the Job " + draft_jobid)

url = platform + "/api/v2/jobs/" + draft_jobid + "/submit/"
requests.post(
  url,
  headers={"Authorization": api_key_in_header}
)

print("Poll the status of the Job " + draft_jobid)

cluster_started = False
while cluster_started == False:
  time.sleep(10)
  url = platform + "/api/v2/jobs/?state=not_completed"
  response = requests.get(
    url,
    headers={"Authorization": api_key_in_header}
  )
  not_completed_jobs = json.loads(response.text)

  for job in not_completed_jobs["results"]:
    if job["id"] == draft_jobid and job["jobStatus"]["content"] == "Executing":
      cluster_started = True
      break

time.sleep(job_duration_sec)

url = platform + "/api/v2/jobs/" + draft_jobid + "/stop/"
requests.post(
  url,
  headers={"Authorization": api_key_in_header}
)
print("Requested to stop the job " + draft_jobid)

上記の方法で, Rescale のジョブを任意の経過時間後（秒単位で設定可能）に自動停止することができます。ジョブ AbCdE の実行コストが発生しますが、ジョブ FgHiJ の時間あたりコストが高いケースにおいては、これを適切に自動停止できるメリットの方が大きくなります