diff --git a/lib/slurm_cli/slurm_control.py b/lib/slurm_cli/slurm_control.py index 2cf42cbed71f1ead744a6cc93c36b1ed67042c30..d04c031920c5b5cc8be24c3598b88688f96214f2 100644 --- a/lib/slurm_cli/slurm_control.py +++ b/lib/slurm_cli/slurm_control.py @@ -2,8 +2,16 @@ from typing import List, Tuple, Union, Dict from .jobs import SlurmJobStatus from subprocess import run as run_process -__SQUEUE_PATH='squeue' -__SRUN_PATH='srun' +__SQUEUE_PATH = 'squeue' +__SRUN_PATH = 'srun' + + +def execute_on_shell(cmd, args): + process_status = run_process([cmd] + args) + if process_status.returncode > 0: + raise SlurmCallError() + output = process_status.stdout + return output class EmptyListException(Exception): @@ -24,24 +32,24 @@ def __list_contains_valid_ids(ids_list): raise EmptyListException() -def __compose_get_processes_status_cmd(job_ids: Union[List, Tuple] = ()): +def __compose_get_processes_status_cmd(job_ids: Union[List, Tuple] = (), job_name: Union[List, Tuple] = ()): cmd = ['--states=all', '-h'] fmt = '%i;%j;%t;%T;%r' cmd += ['--format=%s' % fmt] + if job_ids: - cmd += ','.join(job_ids) + cmd += [','.join(job_ids)] else: cmd += ['-a'] + if job_name: + cmd += ['-n', ','.join(job_name)] + return cmd def __execute_squeue(args): - process_status = run_process([__SQUEUE_PATH] + args) - if process_status.returncode > 0: - raise SlurmCallError() - output = process_status.stdout - return output + return execute_on_shell(__SQUEUE_PATH, args) def __parse_squeue_output(squeue_output) -> List[SlurmJobStatus]: @@ -71,8 +79,34 @@ 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] = ()): - args = __compose_get_processes_status_cmd(job_ids) +def get_jobs_status(job_ids: Union[List, Tuple] = (), job_name: Union[List, Tuple] = ()): + args = __compose_get_processes_status_cmd(job_ids, job_name) output = __execute_squeue(args) parsed_output = __parse_squeue_output(output) - return __map_job_status_per_jobid(parsed_output) \ No newline at end of file + return __map_job_status_per_jobid(parsed_output) + + +def __execute_srun(args): + return execute_on_shell(__SRUN_PATH, args) + + +def __compose_run_job_arguments(cmd, queue=None, executor_config=None, job_name=None): + if executor_config: + raise NotImplementedError() + + args = [] + if queue: + args += ['-p', queue] + if job_name: + args += ['-J', job_name] + else: + args += ['-J', cmd] + args += cmd.split(' ') + return args + + +# TODO collect output +def run_job(cmd, queue=None, executor_config=None, task_name=None): + args = __compose_run_job_arguments(cmd, queue, executor_config, task_name) + + _ = __execute_srun(args)