Coverage for src/shephex/cli/slurm/open.py: 100%
13 statements
« prev ^ index » next coverage.py v7.6.1, created at 2025-06-20 14:13 +0200
« prev ^ index » next coverage.py v7.6.1, created at 2025-06-20 14:13 +0200
1import subprocess
2from pathlib import Path
3from typing import Union
5import rich_click as click
7from shephex.cli.slurm.slurm import slurm
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()
17 if len(output) == 0:
18 return None
20 return output[0]
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)
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()
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()
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')
52 output = output.split('\n')
54 found=False
55 for line in output:
57 if '#SBATCH --output' in line or 'SBATCH -o' in line:
58 file_name = line.split('--output=')[1]
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
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"
73 output_file = f"{workdir}/{file_name}"
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()
80 return output_file
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)
94 if print_path:
95 print(file_name)
96 else: # pragma: no cover
97 command = [command_call, file_name]
98 subprocess.call(command)