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

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

2 

3from __future__ import annotations 

4import re 

5from pathlib import Path 

6 

7 

8class SlurmBatch: 

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

10 

11 Attributes 

12 ---------- 

13 sbatch_options: list[str] 

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

15 cmd_params: list[str] 

16 The parameters to pass to the command 

17 cmd: str 

18 The command to execute 

19 srun_options: list[str] 

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

21 file: Path 

22 The loaded file Path 

23 """ 

24 

25 # Precompiled regex 

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

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

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

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

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

31 

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

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

34 self.file = Path(srcfile) 

35 

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

37 filestr = f.read() 

38 

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

40 

41 # First find the cmd_params block ... 

42 cmd_block = "" 

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

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

45 

46 # ... then parse it 

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

48 

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

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

51 

52 self.srun_options = srun_args.split() 

53 

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

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

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