Coverage for src/mockslurm/mock_sbatch.py: 53%
147 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"""Implement a mock of the sbatch command of slurm.
6The implementation mimics the sbatch API, while using subprocess to
7start the processes on the host machine without slurm daemons.
9Double fork is used to detach the subprocess, allowing the sbatch
10mock to return while the created process is still running.
11"""
13import argparse
14import logging
15import os
16import subprocess as sp
17import sys
18import time
19from enum import Enum
21from mockslurm.process_db import (
22 JobReason,
23 JobState,
24 append_job,
25 find_db_file,
26 get_db,
27 get_db_file_handle,
28 update_db_value,
29)
32class DependencyActions(Enum):
33 """Scheduler behaviour with respect to job dependencies."""
35 OK = 1
36 WAIT = 2
37 NEVER = 3
40def parse_dependency(job_dependency: str) -> tuple[bool, list[dict[str, list[int]]]]:
41 """Parse a job's dependency specification string in sbatch format.
43 Parameters
44 ----------
45 job_dependency : str
46 sbatch job's dependency string. Examples: afterok:20:21,afterany:23
48 Returns
49 -------
50 Tuple[bool, List[Dict[str, List[int]]]]
51 Tuple 1st element is whether all dependencies must be satisfied for a job to start (True, "," in slurm str),
52 or if a single dependency is enough (False, "?" in slurm str). When a single dependency is present: True.
53 The list values are dictionaries with a single key been the dependency type ("afterok", "afterany", etc.)
54 and the values been the job IDs for these dependency
55 Example
56 -------
57 >>> parse_dependency("afterok:20:21?afterany:23")
58 (False, [{'afterok': [20, 21]}, {'afterany': [23]}])
59 """
60 if "," in job_dependency and "?" in job_dependency:
61 raise ValueError('Multiple separator "?", ","in job dependency not supported.')
62 separator = "?" if "?" in job_dependency else "," # default is "," if only 1 dep
63 dependencies = []
64 while job_dependency:
65 if job_dependency[0] == separator:
66 job_dependency = job_dependency[1:]
68 dep_str = job_dependency.split(separator, maxsplit=1)[0]
69 job_dependency = job_dependency[len(dep_str) :]
70 split = dep_str.split(":")
71 dependencies.append({split[0]: [int(v) for v in split[1:]]})
73 return (separator == ",", dependencies)
76def check_dependency(job_dependency: list[tuple[bool, dict[str, list[int]]]]) -> str:
77 """Check if a job's dependency are satisfied.
79 Parameters
80 ----------
81 job_dependency : Tuple[bool, List[Dict[str, List[int]]]]
82 Job dependencies. Each tuple in the list corresponds to a dependency specification.
83 See `mockslurm.mock_sbatch.parse_dependency`
85 Warning
86 -------
87 This function only supports the "afterok" dependency specification, anything else causes
88 a `NotImplementedError` to be raised.
90 Raise
91 -----
92 NotImplementedError
93 If the dependency type is not "afterok"
95 Returns
96 -------
97 DependencyActions
98 The action to follow based on the dependency evaluation, for instance to wait, to start job, etc.
99 """
100 if not job_dependency[1]:
101 return DependencyActions.OK
103 deps_value = []
104 with get_db_file_handle(find_db_file()) as db_file:
105 db = get_db(db_file)
106 for dep in job_dependency[1]:
107 internal_deps_ok = True
108 for dep_type, job_idx in dep.items():
109 if dep_type != "afterok":
110 raise NotImplementedError(f"Dependency type {dep_type} is not implemented.")
111 dep_state = db[job_idx if isinstance(job_idx, list) else [job_idx]]["STATE"]
112 if any(dep_state == JobState.FAILED) or any(dep_state == JobState.CANCELLED):
113 return DependencyActions.NEVER
114 internal_deps_ok &= all(dep_state == JobState.COMPLETED)
116 deps_value.append(internal_deps_ok)
118 combined_dep = all(deps_value) if job_dependency[0] else any(deps_value)
119 return DependencyActions.OK if combined_dep else DependencyActions.WAIT
122def launch_job(
123 cmd: str, stdout: str, stderr: str, job_idx: int, job_dependency: list[tuple[bool, dict[str, list[int]]]]
124):
125 """Run `cmd` as a detached subprocess.
127 The return code of `cmd` is retrieved by the detached subprocess and updated in the DB.
129 Parameters
130 ----------
131 cmd : str
132 Command to run, similar to the content of the --wrap= argument of sbatch
133 stdout : str
134 Path to a file where the output of `cmd` will be piped
135 stderr : str
136 Path to a file where the errors of `cmd` will be piped
137 job_idx : int
138 Index of the job in the DB
139 job_dependency : List[Tuple[bool, Dict[str, List[int]]]]
140 Job dependencies. Each tuple in the list corresponds to a dependency specification.
141 See `mockslurm.mock_sbatch.parse_dependency`
142 """
143 logging.debug(
144 f"lauch_job with args cmd: {cmd}, stdout: {stdout}, stderr: {stderr}, job_idx: {job_idx}, job_dependency: {job_dependency}"
145 )
146 # Wait for dependencies to be ready
147 dependency_check = check_dependency(job_dependency)
148 if dependency_check == DependencyActions.WAIT:
149 logging.debug(f"Job {job_idx} dependency {job_dependency} 1st check is WAIT, updating db Reason DEPENDENCY")
150 with get_db_file_handle(find_db_file()) as db_file:
151 db = get_db(db_file)
152 update_db_value(db_file, job_idx, key="REASON", value=JobReason.Dependency)
153 while dependency_check == DependencyActions.WAIT:
154 time.sleep(0.25)
155 logging.debug(f"Job {job_idx} dependency {job_dependency} check remains WAIT, sleeping")
156 dependency_check = check_dependency(job_dependency)
158 # If not ok: do not start job and mark its state as FAILED
159 if dependency_check != DependencyActions.OK:
160 logging.debug(
161 f"Job {job_idx} dependency check is {dependency_check}, updating DB with REASON DependencyNeverSatisfied and STATE FAILED"
162 )
163 with get_db_file_handle(find_db_file()) as db_file:
164 db = get_db(db_file)
165 update_db_value(db_file, job_idx, key="REASON", value=JobReason.DependencyNeverSatisfied)
166 update_db_value(db_file, job_idx, key="STATE", value=JobState.FAILED)
167 else:
168 logging.debug(f"Job {job_idx} dependency check is OK, checking job STATE")
169 # lock the DB here, so STATE can not change before we start the process
170 with get_db_file_handle(find_db_file()) as db_file:
171 db = get_db(db_file)
172 if db[job_idx]["STATE"] == JobState.PENDING:
173 logging.debug(f"Job {job_idx} STATE is still pending, starting job")
174 try:
175 stdout_f = open(stdout, "a")
176 stderr_f = sp.STDOUT if stderr == stdout else open(stderr, "a")
177 # Can not use shlex here to split command, otherwise we would split bash commands at each
178 # space so we use a single string with shell=True
179 p = sp.Popen(cmd, stdout=stdout_f, stderr=stderr_f, start_new_session=True, shell=True)
180 logging.debug(
181 f"Job {job_idx} started with PID {p.pid}, cmd {cmd}, stdout {stdout_f}, stderr {stderr_f}"
182 )
183 update_db_value(db_file, job_idx, key="PID", value=p.pid)
184 update_db_value(db_file, job_idx, key="STATE", value=JobState.RUNNING)
185 update_db_value(db_file, job_idx, key="REASON", value=JobReason.NOREASON)
186 logging.debug(f"Job {job_idx} updated db with PID {p.pid}, STATE RUNNING, REASON NOREASON")
187 except:
188 logging.debug("Job {} failed to start. Updating DB with STATE FAILED and REASON JobLaunchFailure")
189 update_db_value(db_file, job_idx, key="STATE", value=JobState.FAILED)
190 update_db_value(db_file, job_idx, key="REASON", value=JobReason.JobLaunchFailure)
191 logging.debug("Job {} starting error", exc_info=True)
192 raise
193 else:
194 logging.debug(
195 "Job {} STATE is {}. Do not start start job and exit".format(job_idx, db[job_idx]["STATE"])
196 )
197 # Process is not pending anymore, it must have been killed already, so we won't start it
198 return
200 logging.debug(f"Job {job_idx} is running, waiting its completion")
201 # wait for process to be done
202 exit_code = p.wait()
203 logging.debug(f"Job {job_idx} completed with exit code {exit_code}")
204 # closing stdout and stderr file of job
205 stdout_f.close()
206 if stderr != stdout:
207 stderr_f.close()
208 logging.debug(f"Job {job_idx} closed stdout {stdout} and stderr {stderr}")
209 # Update job state in db
210 with get_db_file_handle(find_db_file()) as db_file:
211 job_state = JobState.COMPLETED if exit_code == 0 else JobState.FAILED
212 job_reason = JobReason.NOREASON if exit_code == 0 else JobReason.NonZeroExitCode
213 update_db_value(db_file, job_idx, key="EXIT_CODE", value=exit_code)
214 update_db_value(db_file, job_idx, key="STATE", value=job_state)
215 update_db_value(db_file, job_idx, key="REASON", value=job_reason)
216 logging.debug(
217 f"Job {job_idx} Updated db with EXIT_CODE {exit_code}, STATE {JobState(job_state).name}, REASON {JobReason(job_reason).name}"
218 )
221def main():
222 parser = argparse.ArgumentParser(
223 description="Slurm sbtach mock.", formatter_class=argparse.ArgumentDefaultsHelpFormatter
224 )
225 parser.add_argument("--account", "-A", type=str, dest="ACCOUNT", help="user account", required=False)
226 parser.add_argument(
227 "--dependency",
228 "-d",
229 type=str,
230 default="",
231 dest="dependency",
232 help="Defer the start of this job until the specified dependencies have been satisfied",
233 )
234 parser.add_argument(
235 "--error", "-e", type=str, dest="stderr", help="error file of slurm job", default="slurm-%j.out"
236 )
237 parser.add_argument("--job-name", "-J", type=str, dest="NAME", help="job name", default="wrap")
238 parser.add_argument(
239 "--nodelist",
240 "-w",
241 type=str,
242 dest="NODELIST",
243 help="Request a specific list of hosts",
244 nargs="*",
245 )
246 parser.add_argument(
247 "--output", "-o", type=str, dest="stdout", help="output file of slurm job", default="slurm-%j.out"
248 )
249 parser.add_argument("--partition", "-p", type=str, dest="PARTITION", help="job partition", required=False)
250 parser.add_argument(
251 "--parsable",
252 action="store_true",
253 dest="parsable",
254 help="Outputs only the job id number. Errors will still be displayed",
255 )
256 parser.add_argument("--reservation", type=str, dest="RESERVATION", help="job reservation", required=False)
257 parser.add_argument(
258 "--wrap",
259 type=str,
260 dest="CMD",
261 help="Command to be executed",
262 required=True,
263 )
264 parser.add_argument(
265 "--mock_sbatch_debug",
266 action="store_true",
267 dest="debug_mode",
268 help="If provided, logs all actions of the sbatch mock in a mock_sbatch.log file",
269 default=False,
270 )
272 args, _ = parser.parse_known_args()
273 defined_args = {arg: value for arg, value in vars(args).items() if value is not None}
274 defined_args.pop("stdout")
275 defined_args.pop("stderr")
276 defined_args.pop("parsable")
277 defined_args.pop("dependency")
278 defined_args.pop("debug_mode")
279 # TODO: raise error if arguments to sbatch aren't valid (reservation, partition, nodelist ? etc.)
281 if args.debug_mode: 281 ↛ 288line 281 didn't jump to line 288 because the condition on line 281 was always true
282 logging.basicConfig(
283 filename="mock_sbatch.log",
284 level=logging.DEBUG,
285 format="%(asctime)s %(levelname)s mock_sbatch %(pathname)s:%(lineno)s:%(funcName)s %(message)s",
286 )
288 logging.debug(
289 "\n".join(
290 ["mock_sbatch called with args:"] + [f"--{arg}: {arg_value}" for arg, arg_value in vars(args).items()]
291 )
292 )
294 with get_db_file_handle(find_db_file()) as db_file:
295 job_idx = append_job(db_file, **defined_args)
296 logging.debug(f"Appended job {job_idx} in DB with args {defined_args}")
298 # if parsable is set: print jobID
299 if args.parsable: 299 ↛ 300line 299 didn't jump to line 300 because the condition on line 299 was never true
300 print(job_idx, end="", flush=True)
301 logging.debug(f"Printing job ID {job_idx} to stdout")
303 # In order to create a process to run the sbatch command and exit sbatch without killing the child process or
304 # making it a zombie, we need to use the double fork technique used to create deamons:
305 # see https://stackoverflow.com/questions/473620/how-do-you-create-a-daemon-in-python
306 # https://stackoverflow.com/questions/881388/what-is-the-reason-for-performing-a-doube-fork-when-creating-a-daemon
307 # Double forking is used to create a new process B (that is attached to process group of A, its parent), then
308 # create a new process C and dettach it from A. The new process can then run and be cleaned up by process 1
309 # instead of A.
310 # There are a number of other things to consider to completely detach a process, such as closing all file handles,
311 # here we only do the minimum for the mock to work.
312 # See https://pagure.io/python-daemon/blob/main/f/src/daemon/daemon.py for a reference implementation of a daemon
313 # in python (old one, there is a pep associated that was never merged due to lack of interest)
315 logging.debug(f"Job {job_idx} Parent process will fork child")
316 # fork a first time
317 pid_1st_fork = os.fork()
318 if pid_1st_fork == 0: # we are in the child process 318 ↛ 319line 318 didn't jump to line 319 because the condition on line 318 was never true
319 logging.debug(f"Job {job_idx} Child process, creating new process session and becoming group leader")
320 # create new process session for the child
321 # the 2nd fork process will be in this session, detached from the initial process
322 os.setsid()
324 logging.debug(f"Job {job_idx} Child process will fork grandchild")
325 # second fork
326 pid_2nd_fork = os.fork()
327 if pid_2nd_fork == 0: # we are in the grandchild: run the sbatch job, then update db
328 logging.debug(f"Job {job_idx} Grandchild process re-derecting stdin, stdout and stderr")
329 # redirect stdin, stdout and stderr to /dev/null (so we separate from child's file handles)
330 # (child's file handles are the same than the parent child handles...)
331 os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdin.fileno())
332 os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdout.fileno())
333 os.dup2(os.open(os.devnull, os.O_RDWR), sys.stderr.fileno())
334 logging.debug(f"Job {job_idx} Grandchild process calls launch_job")
335 launch_job(
336 args.CMD,
337 args.stdout.replace("%j", str(job_idx).strip()),
338 args.stderr.replace("%j", str(job_idx).strip()),
339 job_idx,
340 parse_dependency(args.dependency),
341 )
342 # exit without clean up.
343 # it is REQUIRED to have the grandchild exit in this way if it is tested with pytest, otherwise pytest
344 # will run the program several times see https://github.com/pytest-dev/pytest/issues/12028
345 # Side effect: the stdio buffers will not be flushed (detached process can't output to stdio anyway)
346 os._exit(0)
347 else:
348 # exit the child without cleaning up file handlers or flushing buffers, because those are shared with parent
349 logging.debug(f"Job {job_idx} Child proces exiting without cleaning up.")
350 os._exit(0)
351 else:
352 os.wait() # this waits for the child (not the grandchild!) to finish to clean it up
353 logging.debug(f"Job {job_idx} Parent process exiting after waiting on child process.")
356if __name__ == "__main__": 356 ↛ 357line 356 didn't jump to line 357 because the condition on line 356 was never true
357 main()