Coverage for src/shephex/executor/slurm/slurm_script.py: 100%

26 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-03-29 18:45 +0100

1import re 

2import subprocess 

3from pathlib import Path 

4 

5from shephex.executor.slurm import SlurmBody, SlurmHeader 

6 

7 

8class SlurmScript: 

9 def __init__( 

10 self, 

11 header: SlurmHeader, 

12 body: SlurmBody, 

13 directory: Path, 

14 name: str = 'submit.sh', 

15 ) -> None: 

16 """ 

17 Slurm script object, consisting of a header and body. 

18 

19 Parameters: 

20 ----------- 

21 header : SlurmHeader 

22 Header portion of the script 

23 body : SlurmBody 

24 Body portion of the script 

25 directory : Path 

26 Directory to write the script to. 

27 """ 

28 

29 self.header = header 

30 self.body = body 

31 self.directory = directory 

32 self.name = name 

33 

34 def __repr__(self) -> str: 

35 return str(self.header) + 2 * '\n' + str(self.body) + '\n' 

36 

37 @property 

38 def path(self) -> Path: 

39 return self.directory / self.name 

40 

41 def write(self) -> bool: 

42 """ 

43 Write the script to the specified directory. 

44 """ 

45 self.directory.mkdir(parents=True, exist_ok=True) 

46 

47 with open(self.path, 'w') as f: 

48 f.write(str(self)) 

49 

50 def submit(self, command: str = 'sbatch') -> int: 

51 """ 

52 Submit the script to the Slurm scheduler. 

53 

54 Parameters: 

55 ----------- 

56 command : str 

57 Command to use to submit the script. Default is 'sbatch'. 

58 Returns: 

59 -------- 

60 job_id : int 

61 Job ID of the submitted 

62 """ 

63 

64 if not self.path.exists(): 

65 self.write() 

66 

67 command = [command, str(self.name)] 

68 result = subprocess.check_output(command, cwd=self.directory) 

69 job_id = int(re.search(r'\d+', str(result)).group()) 

70 

71 return job_id