Coverage for src/mockslurm/mock_squeue.py: 85%
83 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-07 13:34 +0000
1# Copyright 2026 CNRS
2# This software is distributed under the terms of the CeCILL-C free software license.
4"""Mock of the squeue command of slurm.
6The squeue -o, --format argument is supported except for the %all directive
7"""
9import argparse
10import datetime
11from collections.abc import Callable
13import numpy as np
15from mockslurm.process_db import (
16 JobReason,
17 JobState,
18 find_db_file,
19 get_db,
20 get_db_file_handle,
21 get_filtered_DB_mask,
22)
23from mockslurm.utils import filter_dict_from_args
25_SQUEUE_DEFAULT_SHORT_FORMAT = "%.18i %.9P %.8j %.8u %.2t %.10M %.6D %R"
26_SQUEUE_DEFAULT_LONG_FORMAT = "%.18i %.9P %.8j %.8u %.8T %.10M %.9l %.6D %R"
28_SQUEUE_STATE_CODE_LONG_TO_SHORT = {
29 JobState.PENDING: "PD",
30 JobState.RUNNING: "R",
31 JobState.FAILED: "F",
32 JobState.COMPLETED: "CD",
33 JobState.CANCELLED: "CA",
34}
37def count_nodes(nodelist: str) -> int:
38 """Count the number of nodes requested by users in the nodelist.
40 Parameters
41 ----------
42 nodelist : str
43 Nodelist string in slurm format: comma separated list of nodes or node ranges
44 example: "cp12,cp18", "cp[10-14],cp28"
46 Returns
47 -------
48 int
49 Number of nodes concerned by the nodelist
50 """
51 if not nodelist: # mock considers no nodelist jobs get allocated 1 node 51 ↛ 52line 51 didn't jump to line 52 because the condition on line 51 was never true
52 return 1
54 nodecount = 0
55 nlist = nodelist.split(",")
56 for node_specs in nlist:
57 if "[" in node_specs:
58 _, node_specs = node_specs.split("[") # ignore node name
59 if "-" in node_specs: # count nodes in range
60 first, last = node_specs.split("-")
61 last = last.split("]")[0] # ignore closing "]" if any
62 nodecount += int(last) - int(first) + 1
63 else: # not a range: single node
64 nodecount += 1
66 return nodecount
69SQUEUE_FORMAT_LETTER_TO_DB_FIELD = {
70 # "": lambda job_info, _: job_info["PID"],
71 "i": ("JOBID", lambda _, job_idx: str(job_idx)),
72 "A": ("JOBID", lambda _, job_idx: str(job_idx)),
73 "j": ("NAME", lambda job_info, _: job_info["NAME"].decode()),
74 "u": ("USER", lambda job_info, _: job_info["USER"].decode()),
75 "a": ("ACCOUNT", lambda job_info, _: job_info["ACCOUNT"].decode()),
76 "P": ("PARTITION", lambda job_info, _: job_info["PARTITION"].decode()),
77 "v": ("RESERVATION", lambda job_info, _: job_info["RESERVATION"].decode()),
78 "M": (
79 "TIME",
80 lambda job_info, _: ( # split is to remove the floating point part, that contains microseconds that should not be displayed
81 str(datetime.datetime.now() - datetime.datetime.fromtimestamp(job_info["START_TIME"])).split(".")[0]
82 if job_info["STATE"] == JobState.RUNNING
83 else "0:00:00"
84 ),
85 ),
86 "l": ("TIME_LIMIT", lambda job_info, _: "UNLIMITED"), # no time limit in mock
87 "n": ("REQ_NODES", lambda job_info, _: job_info["NODELIST"].decode()),
88 "N": ("NODELIST", lambda job_info, _: job_info["NODELIST"].decode()),
89 "D": ("NODES", lambda job_info, _: str(count_nodes(job_info["NODELIST"].decode()))),
90 "S": ("START_TIME", lambda job_info, _: job_info["START_TIME"]),
91 "V": (
92 "SUBMIT_TIME",
93 lambda job_info, _: job_info["START_TIME"].decode(),
94 ), # equal to start time for mock)
95 "o": ("COMMAND", lambda job_info, _: job_info["CMD"].decode()),
96 "r": ("REASON", lambda job_info, _: JobReason(job_info["REASON"]).name),
97 "t": ("ST", lambda job_info, _: _SQUEUE_STATE_CODE_LONG_TO_SHORT[job_info["STATE"]]),
98 "T": ("STATE", lambda job_info, _: JobState(job_info["STATE"]).name),
99 "R": (
100 "NODELIST(REASON)",
101 lambda job_info, _: (
102 "(" + JobReason(job_info["REASON"]).name + ")"
103 if JobReason(job_info["STATE"]) == JobState.PENDING
104 or JobReason(job_info["REASON"]) == JobReason.NonZeroExitCode
105 else job_info["NODELIST"].decode()
106 ),
107 ),
108 # "": lambda job_info, _: job_info["EXIT_CODE"],
109}
112def parse_squeue_format(squeue_format_str: str) -> tuple[str, list[str], list[Callable]]:
113 """Convert squeue format argument (-o format) to a python format string and
114 list of function filling the values of the format string for each job DB row.
116 Parameters
117 ----------
118 squeue_format_str : str
119 Squeue format string, eg "%.18i %.9P %.8j %.8u %.2t %.10M %.6D %R".
121 Warning
122 -------
123 Does not support squeue's "%all" formatting string.
125 Returns
126 -------
127 tuple[str, list[str], list[Callable]]
128 python format string, and a list of callable that should be used to fill the values
129 in the formatted string with a job DB row and job_index.
131 Examples
132 --------
133 >>> job_DB_row = {
134 ... "PID": -1,
135 ... "NAME": b"jobname",
136 ... "USER": b"bob",
137 ... "ACCOUNT": b"bob_name",
138 ... "PARTITION": b"",
139 ... "RESERVATION": b"",
140 ... "NODELIST": b"mocknode1",
141 ... "TIME": 0,
142 ... "START_TIME": datetime.datetime.now().timestamp(),
143 ... "CMD": "",
144 ... "STATE": JobState.PENDING,
145 ... "REASON": JobReason.WaitingForScheduling,
146 ... "EXIT_CODE": np.iinfo(np.int16).max,
147 ... }
148 >>> row_idx = 0
149 >>> format_str, fields_header, fields_filler_fcts = parse_squeue_format("%.18i %.9P %.8j %.8u %.2t %.10M %.6D %R")
150 >>> squeue_output = format_str.format(*[fct(job_DB_row, row_idx) for fct in fields_filler_fcts])
151 """
152 # Iterate on the format string and for each found %
153 # add a "{:ndigit}" to the python string, with optional ">" if "." follows "%"
154 # ndigit is found by converting the characters following %
155 # the function filling the value is taken from the map of squeue format letters to DB function
156 python_format_string = ""
157 fields_header = []
158 fields_filler_functions = []
159 # split 1st value is empty if string starts with delimiter
160 for field_format_str in squeue_format_str.split("%")[1:]:
161 python_format_string += "{:"
162 if field_format_str[0] == ".":
163 python_format_string += ">"
164 field_format_str = field_format_str[1:]
165 idx = 0
166 while field_format_str[idx].isdigit():
167 python_format_string += field_format_str[idx]
168 idx += 1
169 header, fct = SQUEUE_FORMAT_LETTER_TO_DB_FIELD[field_format_str[idx]]
170 fields_header.append(header)
171 fields_filler_functions.append(fct)
172 python_format_string += "}" + field_format_str[idx + 1 :]
174 return python_format_string, fields_header, fields_filler_functions
177def main():
178 parser = argparse.ArgumentParser(
179 description="Slurm scancel mock", formatter_class=argparse.ArgumentDefaultsHelpFormatter, add_help=False
180 )
181 user_group = parser.add_mutually_exclusive_group()
182 parser.add_argument(
183 "--account",
184 "-A",
185 type=str,
186 dest="ACCOUNT",
187 help="Specify the accounts of the jobs to view. Accepts a comma separated list of account names",
188 )
189 parser.add_argument(
190 "--name",
191 "-n",
192 type=str,
193 dest="NAME",
194 help="Request jobs having one of the specified names. The list consists of a comma separated list of job names.",
195 )
196 user_group.add_argument(
197 "--me",
198 action="store_true",
199 dest="me",
200 help="Equivalent to --user=<my username>",
201 )
202 parser.add_argument(
203 "--nodelist",
204 "-w",
205 type=str,
206 dest="NODELIST",
207 help="Report only on jobs allocated to the specified node or list of nodes",
208 )
209 parser.add_argument(
210 "--format",
211 "-o",
212 type=str,
213 default=_SQUEUE_DEFAULT_SHORT_FORMAT,
214 dest="format_str",
215 help="Specify the information to be displayed, its size and position",
216 )
217 parser.add_argument(
218 "--noheader", "-h", action="store_true", dest="no_header", help="Do not print a header on the output"
219 )
220 parser.add_argument(
221 "--long",
222 "-l",
223 action="store_true",
224 dest="long",
225 help="Report more of the available information for the selected jobs or job steps, subject to any constraints specified",
226 )
227 parser.add_argument(
228 "--partition",
229 "-p",
230 type=str,
231 dest="PARTITION",
232 help="Specify the partitions of the jobs or steps to view. Accepts a comma separated list of partition names",
233 )
234 parser.add_argument(
235 "--reservation",
236 "-R",
237 type=str,
238 dest="RESERVATION",
239 help="Specify the reservation of the jobs to view",
240 )
241 parser.add_argument(
242 "--usage",
243 action="store_true",
244 dest="print_help",
245 help="Print a brief help message listing the squeue options",
246 )
247 parser.add_argument(
248 "--jobs",
249 "-j",
250 type=str,
251 dest="jobs",
252 help="Specify a comma separated list of job IDs to display, Defaults to all jobs",
253 )
254 args = parser.parse_args()
256 if args.print_help: 256 ↛ 257line 256 didn't jump to line 257 because the condition on line 256 was never true
257 parser.print_help()
259 if args.long and not args.format_str: 259 ↛ 260line 259 didn't jump to line 260 because the condition on line 259 was never true
260 args.format_str = _SQUEUE_DEFAULT_LONG_FORMAT
262 # Split args that take list in comma separated string to list!
263 args.ACCOUNT = args.ACCOUNT.split(",") if args.ACCOUNT is not None else None
264 args.NAME = args.NAME.split(",") if args.NAME is not None else None
265 args.PARTITION = args.PARTITION.split(",") if args.PARTITION is not None else None
266 args.jobs = [int(job_id) for job_id in args.jobs.split(",")] if args.jobs is not None else None
268 field_filter_values = filter_dict_from_args(args)
269 # filter out finished jobs
270 field_filter_values["STATE"] = [JobState.RUNNING, JobState.PENDING]
272 with get_db_file_handle(find_db_file()) as db_file:
273 # Get mask to select DB rows
274 mask = get_filtered_DB_mask(db_file, field_filter_values)
275 # filter job IDs if some were specified
276 if args.jobs and len(mask) > 0: # if mask is empty (no jobs in DB), skip 276 ↛ 279line 276 didn't jump to line 279 because the condition on line 276 was never true
277 # job IDs are just the index of the jobs in the DB!
278 # remove (silently like real squeue...) job IDs that do not exist
279 args.jobs = [job_ID for job_ID in args.jobs if 0 <= job_ID < len(mask)]
280 if not args.jobs: # no jobs remaining: exit with error like squeue
281 print("slurm_load_jobs error: Invalid job id specified")
282 exit(1)
283 mask[args.jobs] = True
285 # Get format string and function formating field value based on squeue --format
286 format_str, fields_header, fields_filler_fct = parse_squeue_format(args.format_str)
288 # Print header
289 if not args.no_header: 289 ↛ 293line 289 didn't jump to line 293 because the condition on line 289 was always true
290 print(format_str.format(*fields_header))
292 # Print the jobs found in filtered DB
293 job_indices = np.nonzero(mask)[0]
294 for idx, job in zip(job_indices, get_db(db_file)[mask]):
295 print(format_str.format(*[fct(job, idx) for fct in fields_filler_fct]))
298if __name__ == "__main__": 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true
299 main()