Coverage for src/shephex/executor/slurm/slurm_executor.py: 86%
128 statements
« prev ^ index » next coverage.py v7.6.1, created at 2025-06-20 14:58 +0200
1import inspect
2import json
3from pathlib import Path
4from typing import List, Literal, Optional, Union
6from shephex.executor.executor import Executor
7from shephex.executor.slurm import (
8 SlurmBody,
9 SlurmHeader,
10 SlurmProfileManager,
11 SlurmScript,
12)
13from shephex.experiment import FutureResult
14from shephex.experiment.experiment import Experiment
17class SlurmSafetyError(Exception):
18 pass
20class SlurmExecutor(Executor):
21 """
22 Shephex SLURM executor for executing experiments on a SLURM cluster.
23 """
24 def __init__(
25 self, directory: Union[str, Path] = None,
26 scratch: bool = False,
27 ulimit: Union[int, Literal['default']] = 8000,
28 move_output_file: bool = True,
29 safety_check: bool = True,
30 array_limit: int | None = None,
31 **kwargs
32 ) -> None:
33 """
34 shephex SLURM executor.
36 Parameters
37 ----------
38 directory : Union[str, Path], optional
39 Directory where the SLURM script and output files will be stored,
40 defaults to /slurm.
41 scratch : bool, optional
42 If True, the executor will use the /scratch directory for the
43 execution of the experiments. Defaults to False. When true
44 files will automatically be copied back to the original directory
45 once the job is finished.
46 **kwargs
47 Additional keyword arguments to be passed to the SlurmHeader,
48 these are the SLURM parameters for the job. Supports all the
49 arguments for sbatch, see https://slurm.schedmd.com/sbatch.html.
50 """
51 if safety_check:
52 self.safety_check(frame_index=2)
54 self.header = SlurmHeader()
55 for key, value in kwargs.items():
56 self.header.add(key, value)
58 if directory is None:
59 directory = 'slurm'
60 self.directory = Path(directory)
62 # Containers for commands to be executed before and after the main execution
63 self._commands_pre_execution = []
64 self._commands_post_execution = []
66 # Special options
67 self.ulimit = ulimit
68 self.move_output_file = move_output_file
69 self.scratch = scratch
70 self.array_limit = array_limit
72 # To kepe track of the special options for saving the config
73 self.special_options = {
74 'scratch': scratch,
75 'ulimit': ulimit,
76 'move_output_file': move_output_file,
77 'array_limit': array_limit,
78 }
81 @classmethod
82 def from_config(cls, path: Path, safety_check: bool = True, **kwargs) -> 'SlurmExecutor':
84 if safety_check:
85 cls.safety_check(frame_index=1)
87 if not isinstance(path, Path):
88 path = Path(path)
90 if not path.exists():
91 raise FileNotFoundError(f'File {path} does not exist')
92 if not path.suffix == '.json':
93 raise ValueError(f'File {path} is not a json file')
95 with open(path) as f:
96 config = json.load(f)
97 config.update(kwargs)
99 pre_commands = config.pop('commands_pre_execution', list())
100 post_commands = config.pop('commands_post_execution', list())
102 instance = cls(**config, safety_check=False)
103 instance._commands_pre_execution = pre_commands
104 instance._commands_post_execution = post_commands
105 return instance
107 def to_config(self, path: Path | str) -> None:
108 if not isinstance(path, Path):
109 path = Path(path)
110 config = self.header.to_dict()
111 config.update(self.special_options)
113 config['commands_pre_execution'] = self._commands_pre_execution
114 config['commands_post_execution'] = self._commands_post_execution
116 with open(path, 'w') as f:
117 json.dump(config, f, indent=4)
119 @classmethod
120 def from_profile(cls, name: str, safety_check: bool = True, **kwargs) -> 'SlurmExecutor':
121 """
122 Create a new SlurmExecutor from a profile.
124 Parameters
125 ----------
126 name : str
127 Name of the profile.
128 safety_check : bool, optional
129 If True, a safety check will be performed to ensure that the executor
130 is not instantiated on a script that is not the main script. Defaults to True.
131 **kwargs
132 Additional keyword arguments to be passed to the SlurmExecutor.
133 """
134 if safety_check:
135 cls.safety_check(frame_index=2)
136 kwargs.pop("safety_check", None)
138 spm = SlurmProfileManager()
139 profile = spm.get_profile_path(name)
140 return cls.from_config(profile, **kwargs, safety_check=False)
142 def _single_execute(self) -> None:
143 raise NotImplementedError('Single execution is not supported for SLURM Executor, everything is executed with _sequence execute.')
145 def _sequence_execute(
146 self,
147 experiments: List[Experiment],
148 dry: bool = False,
149 execution_directory: Union[Path, str] = None,
150 ) -> List[FutureResult]:
151 """
152 Execute a sequence of experiments as an array job.
154 Parameters
155 ----------
156 experiments : List[Experiment]
157 List of experiments to be executed.
158 dry : bool, optional
159 If True, the script will be printed instead of executed.
160 execution_directory : Union[Path, str], optional
161 Directory where the experiments will be executed.
163 Returns
164 -------
165 List[FutureResult]
166 List of FutureResult objects.
167 """
169 if len(experiments) == 0:
170 return []
172 # Dump config:
173 self.directory.mkdir(parents=True, exist_ok=True)
174 index = len(list(self.directory.glob('config*.json')))
175 path = self.directory / f'config_{index}.json'
176 self.to_config(path)
178 header = self.header.copy()
179 if self.array_limit is not None:
180 array_limit = min(self.array_limit, len(experiments))
181 else:
182 array_limit = len(experiments)
184 header.add('array', f'0-{len(experiments)-1}%{array_limit}')
186 body = self._make_slurm_body(experiments)
188 count = len(list(self.directory.glob('submit*.sh')))
189 script = SlurmScript(header, body, directory=self.directory, name=f'submit_{count}.sh')
190 if dry:
191 print(script)
192 return [FutureResult() for _ in experiments]
194 script.write()
196 job_id = script.submit()
197 for experiment in experiments:
198 experiment.update_status('submitted')
200 return [FutureResult(info={'job_id': job_id}) for _ in experiments]
202 def _bash_array_str(self, strings: List[str]) -> str:
203 """
204 Convert a list of strings into a nicely formatted bash array of string.
206 Parameters
207 ----------
208 strings : List[str]
209 List of strings to be converted.
211 Returns
212 -------
213 str
214 A python string representing a bash array of strings.
215 """
216 bash_str = ' \n\t'.join(strings)
217 return f'(\n\t{bash_str}\n)'
219 def _body_add(self, command: str, when: Optional[Literal['pre', 'post']] = None) -> None:
220 """
221 Add a command to the body of the SLURM script.
223 Parameters
224 ----------
225 command : str
226 Command to be added to the body.
227 """
228 if when is None:
229 when = 'pre'
231 if when == 'pre':
232 self._commands_pre_execution.append(command)
234 elif when == 'post':
235 self._commands_post_execution.append(command)
237 def _make_slurm_body(self, experiments: List[Experiment]) -> SlurmBody:
238 """
239 Make a new SlurmBody object.
241 Returns
242 -------
243 SlurmBody
244 A new SlurmBody object.
245 """
247 identifiers = [str(experiment.identifier) for experiment in experiments]
248 directories = [str(experiment.directory.resolve()) for experiment in experiments]
250 body = SlurmBody()
252 body.add(f'directories={self._bash_array_str(directories)}')
253 body.add(f'identifiers={self._bash_array_str(identifiers)}')
255 if self.move_output_file:
256 self._body_add(r"mv slurm-${SLURM_ARRAY_JOB_ID}_${SLURM_ARRAY_TASK_ID}.out ${directories[$SLURM_ARRAY_TASK_ID]}", when='pre')
257 if self.ulimit != 'default':
258 self._body_add(f'ulimit -Su {self.ulimit}', when='pre')
260 for command in self._commands_pre_execution:
261 body.add(command)
263 # Slurm info command:
264 command = r'hex slurm add-info -d ${directories[$SLURM_ARRAY_TASK_ID]} -j "${SLURM_ARRAY_JOB_ID}_${SLURM_ARRAY_TASK_ID}"'
265 body.add(command)
267 # Execution command
268 command = r'hex execute ${directories[$SLURM_ARRAY_TASK_ID]}'
270 if self.scratch:
271 command += ' -e /scratch/$SLURM_JOB_ID'
273 body.add(command)
275 for command in self._commands_post_execution:
276 body.add(command)
278 if self.scratch:
279 body.add(
280 r'cp -r /scratch/$SLURM_JOB_ID/* ${directories[$SLURM_ARRAY_TASK_ID]}'
281 )
283 return body
285 @staticmethod
286 def safety_check(frame_index: int = 2) -> None:
287 """
288 Check if the executor is being called from the main script.
290 Parameters
291 ----------
292 frame_index : int, optional
293 Index of the frame to be checked. Defaults to 2.
295 Raises
296 ------
297 SlurmSafetyError
298 If the executor is not being called from the main script.
300 Frame index depends on which creation method is used:
301 - from_profile: 2
302 - from_config: 1
303 - __init__: 0
304 """
306 caller_frames = inspect.stack()
307 caller_frame = caller_frames[frame_index]
308 caller_module = inspect.getmodule(caller_frame[0])
310 if caller_module and caller_module.__name__ != "__main__" or caller_module is None:
311 raise SlurmSafetyError("""SlurmExecutor should only be called from the main script.
312 If the you really want, you can disable this check. This error may be caused by not having
313 a 'if __name__ == "__main__":' block in the main script.""")