Skip to content
Snippets Groups Projects
Commit e82bd3ed authored by Mattia Mancini's avatar Mattia Mancini
Browse files

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

parent cf458076
No related branches found
No related tags found
No related merge requests found
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
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment