Commit e82bd3ed authored by mancini's avatar mancini

capturing output of srun and check allow the specification of of a list of commands

parent cf458076
from typing import List, Tuple, Union, Dict
from .jobs import SlurmJobStatus
from subprocess import run as run_process
from subprocess import run as __run_process
__SQUEUE_PATH = 'squeue'
__SRUN_PATH = 'srun'
def execute_on_shell(cmd, args):
process_status = run_process([cmd] + args)
def __execute_on_shell(cmd, args):
process_status = __run_process([cmd] + args, capture_output=True)
if process_status.returncode > 0:
raise SlurmCallError()
output = process_status.stdout
output = process_status.stdout.decode()
return output
......@@ -49,7 +49,7 @@ def __compose_get_processes_status_cmd(job_ids: Union[List, Tuple] = (), job_nam
def __execute_squeue(args):
return execute_on_shell(__SQUEUE_PATH, args)
return __execute_on_shell(__SQUEUE_PATH, args)
def __parse_squeue_output(squeue_output) -> List[SlurmJobStatus]:
......@@ -79,7 +79,7 @@ def __map_job_status_per_jobid(job_status_list: List[SlurmJobStatus]) -> Dict[st
return {job_status.job_id: job_status for job_status in job_status_list}
def get_jobs_status(job_ids: Union[List, Tuple] = (), job_name: Union[List, Tuple] = ()):
def get_jobs_status(job_ids: Union[List, Tuple] = (), job_name: Union[List, Tuple] = ()) -> Dict[str, SlurmJobStatus]:
args = __compose_get_processes_status_cmd(job_ids, job_name)
output = __execute_squeue(args)
parsed_output = __parse_squeue_output(output)
......@@ -87,7 +87,7 @@ def get_jobs_status(job_ids: Union[List, Tuple] = (), job_name: Union[List, Tupl
def __execute_srun(args):
return execute_on_shell(__SRUN_PATH, args)
return __execute_on_shell(__SRUN_PATH, args)
def __compose_run_job_arguments(cmd, queue=None, executor_config=None, job_name=None):
......@@ -101,7 +101,7 @@ def __compose_run_job_arguments(cmd, queue=None, executor_config=None, job_name=
args += ['-J', job_name]
else:
args += ['-J', cmd]
args += cmd.split(' ')
args += cmd.split(' ') if isinstance(cmd, str) else cmd
return args
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment