srun replacement

run

#!/usr/bin/python
from argparse import ArgumentParser, REMAINDER
from subprocess import check_output, call
from threading import Thread
from time import sleep
import os
import re

POLL_INTERVAL=2

def poll_job(job_id):
  ignore_states = [
    "RUNNING",
    "CONFIGURING",
    "COMPLETING",
    "PENDING",
    "RESV_DEL_HOLD",
    "REQUEUE_FED",
    "REQUEUE_HOLD",
    "REQUEUED",
    "RESIZING"
  ]

  while True:
    stdout = check_output(["scontrol", "show", "job", str(job_id)])
    job_state_search = re.search('JobState=([a-zA-Z_]*)', stdout)
    if job_state_search:
      job_state = job_state_search.group(1)
      if job_state not in ignore_states:
        os._exit(0)

    sleep(POLL_INTERVAL)

def watch_file(filename):
  while not os.path.isfile(filename):
    sleep(POLL_INTERVAL)
  call(["tail", "-f", filename])

def main():
  parser = ArgumentParser()
  parser.add_argument("--mem", required=True)
  parser.add_argument("--time", required=True)
  parser.add_argument("rest", nargs=REMAINDER)
  args = parser.parse_args()

  job_id = check_output(["sbatch", "--parsable", "--mem", args.mem, "--time", args.time, os.path.join(os.environ["HOME"], ".bin", "srun.sh")] + args.rest).strip()
  poll_thread = Thread(target=poll_job, args=(job_id,))
  out_thread = Thread(target=watch_file, args=("slurm-%s.out" % job_id,))
  poll_thread.start()
  out_thread.start()

  try:
    while poll_thread.isAlive():
      poll_thread.join(1)
  except (KeyboardInterrupt, SystemExit):
    check_output(["scancel", job_id])
    os._exit(1)

if __name__ == "__main__":
  main()

srun.sh

#!/bin/bash
"$@"