srutch.py - dev - kotkov/eco

import re
import time
import subprocess

# TODO: move hardcoded user and host variables to config file
USER = "dev"
HOST = "irus17.cc.dvo.ru"


def get_job_status(sacct_output, job_id):
    lines = sacct_output.strip().split("\n")
    if len(lines) < 3:
        return None
    for line in lines[2:]:
        parts = line.split()
        if len(parts) >= 2 and parts[0].startswith(str(job_id)):
            return parts[1]


POLLING_DELAY = 1


def srutch(params: str, submitline: str, callback_func) -> None:
    """(srun + sbatch) over ssh."""

    # Batch a job (stdout example: "Submitted batch job 189142")
    sbatch_command = f'sbatch {params} --wrap "{submitline}"'
    ssh_wrapper = f'ssh {USER}@{HOST} "{sbatch_command}"'
    process = subprocess.run(ssh_wrapper, shell=True,
                             capture_output=True, text=True)
    job_id = re.search("\\d+", process.stdout).group()
    print(f"Batched job: {job_id}")

    # Monitor the batched job
    sacct_command = f'sacct  -j {job_id} -o "jobid,state"'
    prev_status = None
    while True:
        process = subprocess.run(sacct_command, shell=True,
                                 capture_output=True, text=True)
        sacct_output = process.stdout
        status = get_job_status(sacct_output, job_id)
        # On status change
        if status != prev_status:
            callback_func(job_id, status)
            if status == "COMPLETED":
                return
        prev_status = status
        time.sleep(POLLING_DELAY)


def example_callback(job_id, status):
    print(f"Job ID={job_id} changed its status to {status}")


if __name__ == "__main__":
    srutch("-n 1", "srun -N 1 -t 10 sleep 10", example_callback)