Coverage for sparkle/tools/slurm_parsing.py: 42%

1"""This module helps to extract and structure information from Slurm I/O.""" 

2from __future__ import annotations 

3import re 

4from pathlib import Path 

5 

6 

7class SlurmBatch: 

8 """Class to parse a Slurm batch file and get structured information. 

9 

10 Attributes 

11 ---------- 

12 sbatch_options: list[str] 

13 The SBATCH options. Ex.: ["--array=-22%250", "--mem-per-cpu=3000"] 

14 cmd_params: list[str] 

15 The parameters to pass to the command 

16 cmd: str 

17 The command to execute 

18 srun_options: list[str] 

19 A list of arguments to pass to srun. Ex.: ["-n1", "--nodes=1"] 

20 file: Path 

21 The loaded file Path 

22 """ 

23 # Precompiled regex 

24 re_sbatch = re.compile("#SBATCH (.*)") 

25 re_params_all = re.compile(r"params=\( \\\n(?:(.*))\)", re.DOTALL) 

26 re_params_items = re.compile(r"'(.*)'") 

27 re_srun_all = re.compile(r"srun (.*)") 

28 re_srun_split = re.compile(r" (?!-)") 

29 

30 def __init__(self: SlurmBatch, srcfile: Path) -> None: 

31 """Parse the data contained in srcfile and localy store the information.""" 

32 self.file = Path(srcfile) 

33 

34 with Path(self.file).open() as f: 

35 filestr = f.read() 

36 

37 self.sbatch_options = SlurmBatch.re_sbatch.findall(filestr) 

38 

39 # First find the cmd_params block ... 

40 cmd_block = "" 

41 if len(SlurmBatch.re_params_all.findall(filestr)) > 0: 

42 cmd_block = SlurmBatch.re_params_all.findall(filestr)[0] 

43 

44 # ... then parse it 

45 self.cmd_params = SlurmBatch.re_params_items.findall(cmd_block) 

46 

47 srun = SlurmBatch.re_srun_all.findall(filestr)[0] 

48 srun_args, cmd = SlurmBatch.re_srun_split.split(srun, maxsplit=1) 

49 

50 self.srun_options = srun_args.split() 

51 

52 self.cmd = cmd.replace("${params[$SLURM_ARRAY_TASK_ID]}", "").strip() 

53 self.cmd = self.cmd.replace("${output[$SLURM_ARRAY_TASK_ID]}", "").strip() 

54 self.cmd = self.cmd.replace(">", "").strip()