Coverage for src/shephex/cli/slurm/cancel.py: 100%

35 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-03-29 18:45 +0100

1from pathlib import Path 

2from typing import List 

3 

4import rich_click as click 

5 

6from shephex.cli.slurm.slurm import slurm 

7from shephex.experiment import Experiment, ExperimentContext, Meta 

8 

9 

10def get_job_id(directory: Path) -> str: 

11 context = ExperimentContext(directory / Experiment.shep_dir) 

12 job_id = context.meta.get('job-id', None) 

13 return job_id 

14 

15def get_job_status(directory: Path) -> str: 

16 meta = Meta.from_file(directory / Experiment.shep_dir) 

17 return meta.get('status') 

18 

19@slurm.command() 

20@click.argument("directories", type=click.Path(exists=True, file_okay=False, dir_okay=True), nargs=-1) 

21@click.option("-p", '--print_only', is_flag=True, help="Print the command without executing it", default=False) 

22def cancel(directories: List[click.Path], print_only: bool) -> None: 

23 """ 

24 Cancel a job 

25 """ 

26 job_ids = [] 

27 for directory in directories: 

28 directory = Path(directory) 

29 job_id = get_job_id(directory) 

30 status = get_job_status(directory) 

31 

32 print(job_id, status) 

33 if status != 'running': 

34 print(f"Job in {directory} is not running") 

35 continue 

36 if job_id is not None: 

37 job_ids.append(job_id) 

38 

39 

40 

41 if len(job_ids) == 0: 

42 print("No job-ids found") 

43 return 

44 

45 job_ids = list(set(job_ids)) 

46 command = ['scancel'] + job_ids 

47 command = " ".join(command) 

48 

49 if print_only: 

50 print(command)