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

13 statements  

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

1import subprocess 

2from pathlib import Path 

3from typing import Union 

4 

5import rich_click as click 

6 

7from shephex.cli.slurm.slurm import slurm 

8 

9 

10def check_job_status(job_id: int) -> Union[str, None]: # pragma: no cover 

11 """ 

12 Return the status of a given job id. 

13 """ 

14 output = subprocess.check_output(['sacct', '-j', job_id, '--format', 'state, JobName', '--noheader']).decode('utf-8') 

15 output = output.split() 

16 

17 if len(output) == 0: 

18 return None 

19 

20 return output[0] 

21 

22def find_output_file(job_id: int) -> str: # pragma: no cover 

23 """ 

24 Return the path of the output file for a given job id. 

25 """ 

26 status = check_job_status(job_id) 

27 

28 if status is None: 

29 print(f"Job id {job_id} not found.") 

30 exit() 

31 elif status == "RUNNING": 

32 command = ["scontrol", "show", "job", job_id] 

33 output = subprocess.check_output(command).decode('utf-8') 

34 output = output.split() 

35 

36 output_file = None 

37 for info in output: 

38 if 'StdOut' in info: 

39 output_file = info.split('=')[1] 

40 break 

41 else: 

42 # We need to use sacct to find the output file 

43 # First determine the working directory 

44 command = ['sacct', '-j', job_id, '--format', 'workdir%-1000', '--noheader'] 

45 output = subprocess.check_output(command).decode('utf-8') 

46 workdir = output.strip() 

47 

48 # Now we need to read the submission script to find the output file 

49 command = ['sacct', '-B', '-j', job_id] 

50 output = subprocess.check_output(command).decode('utf-8') 

51 

52 output = output.split('\n') 

53 

54 found=False 

55 for line in output: 

56 

57 if '#SBATCH --output' in line or 'SBATCH -o' in line: 

58 file_name = line.split('--output=')[1] 

59 

60 if '_' in job_id: 

61 job_id, arr_id = job_id.split('_') 

62 file_name = file_name.replace(r'%A', f'{job_id}') 

63 file_name = file_name.replace(r'%a', f'{arr_id}') 

64 else: 

65 file_name = file_name.replace(r'%A', f'{job_id}') 

66 found=True 

67 break 

68 

69 if not found: 

70 # Then we assume that thee job was submitted with the default output file 

71 file_name=f"slurm-{job_id}.out" 

72 

73 output_file = f"{workdir}/{file_name}" 

74 

75 # CHeck if the file exists 

76 if not Path(output_file).exists(): 

77 print(f"File {output_file} does not exist - Might be I wasn't able to find it correctly.") 

78 exit() 

79 

80 return output_file 

81 

82@slurm.command(name='open') 

83@click.argument("job_id", type=str) 

84@click.option("-c", "--command_call", type=str, help="Command to open file with", default="code") 

85@click.option("-p", "--print_path", help="Print the path of the file and exit.", is_flag=True) 

86def open_slurm(job_id: int, 

87 command_call: str, 

88 print_path: bool) -> None: 

89 """ 

90 Open slurm output file for a given job id. 

91 """ 

92 file_name = find_output_file(job_id) 

93 

94 if print_path: 

95 print(file_name) 

96 else: # pragma: no cover 

97 command = [command_call, file_name] 

98 subprocess.call(command)