Coverage for src/shephex/executor/slurm/slurm_script.py: 100%
26 statements
« prev ^ index » next coverage.py v7.6.1, created at 2025-06-20 14:13 +0200
« prev ^ index » next coverage.py v7.6.1, created at 2025-06-20 14:13 +0200
1import re
2import subprocess
3from pathlib import Path
5from shephex.executor.slurm import SlurmBody, SlurmHeader
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.
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 """
29 self.header = header
30 self.body = body
31 self.directory = directory
32 self.name = name
34 def __repr__(self) -> str:
35 return str(self.header) + 2 * '\n' + str(self.body) + '\n'
37 @property
38 def path(self) -> Path:
39 return self.directory / self.name
41 def write(self) -> bool:
42 """
43 Write the script to the specified directory.
44 """
45 self.directory.mkdir(parents=True, exist_ok=True)
47 with open(self.path, 'w') as f:
48 f.write(str(self))
50 def submit(self, command: str = 'sbatch') -> int:
51 """
52 Submit the script to the Slurm scheduler.
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 """
64 if not self.path.exists():
65 self.write()
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())
71 return job_id