Coverage for src/shephex/experiment/experiment.py: 99%
141 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
1"""
2Experiment class definition.
3"""
5import json
6import pickle as pkl
7from pathlib import Path
8from time import time
9from typing import Any, Callable, Optional, TypeAlias, Union
11import shortuuid
13from shephex.experiment.context import ExperimentContext
14from shephex.experiment.meta import Meta
15from shephex.experiment.options import Options
16from shephex.experiment.procedure import PickleProcedure, ScriptProcedure
17from shephex.experiment.result import ExperimentResult
18from shephex.experiment.status import Status
20ProcedureType: TypeAlias = Union[PickleProcedure, ScriptProcedure]
23class Experiment:
24 """
25 Experiment class definition, shephex's central object.
26 """
28 extension = 'exp'
29 shep_dir = 'shephex'
31 def __init__(
32 self,
33 *args,
34 function: Optional[Union[Callable, str, Path]] = None,
35 procedure: Optional[ProcedureType] = None,
36 root_path: Optional[Union[Path, str]] = None,
37 identifier: Optional[str] = None,
38 status: Optional[Status] = None,
39 meta: Optional[Meta] = None,
40 **kwargs,
41 ) -> None:
42 """
43 An experiment object containing the procedure, options, and metadata.
45 Parameters
46 ----------
47 *args
48 Positional arguments to be passed to the function or script.
49 function : Optional[Union[Callable, str, Path]], optional
50 A callable function or path to a script, by default None
51 procedure : Optional[ProcedureType], optional
52 A procedure object, by default None
53 root_path : Union[Path, str], optional
54 The root path for the experiment, by default None
55 identifier : Optional[str], optional
56 The identifier for the experiment, by default None. When None
57 a random identifier is generated using shortuuid.
58 status : Optional[str], optional
59 The status of the experiment, by default None (pending).
60 meta : Optional[Meta], optional
61 A shephex.experiment.meta object, by default None. If supplied,
62 identifier and status are ignored.
63 """
65 # Root path
66 self.root_path = Path(root_path).resolve() if root_path else Path.cwd()
68 # Set the procedure
69 if function is not None:
70 self.procedure = function
71 elif procedure is not None:
72 self._procedure = procedure
73 else:
74 raise ValueError("Either 'func' or 'procedure' must be provided.")
76 # Set the options
77 args = args if args else []
78 kwargs = kwargs if kwargs else {}
79 self.options = Options(*args, **kwargs)
81 # Meta
82 if meta is None:
83 identifier = identifier if identifier is not None else shortuuid.uuid()
84 status = status if status is not None else Status.pending()
86 self.meta = Meta(
87 status=status,
88 identifier=identifier,
89 procedure=self.procedure.get_metadata(),
90 options_path=f'{self.shep_dir}/{self.options.name}',
91 time_stamp=time(),
92 )
93 else:
94 self.meta = meta
96 ############################################################################
97 # Properties
98 ############################################################################
100 @property
101 def root_path(self) -> Path:
102 """
103 The root path for the experiment.
105 Returns
106 -------
107 Path
108 The root path for the experiment
109 """
110 return self._root_path
112 @root_path.setter
113 def root_path(self, root_path: Path) -> None:
114 self._root_path = Path(root_path)
116 @property
117 def identifier(self) -> str:
118 """
119 The identifier for the experiment.
121 Returns
122 -------
123 str
124 The identifier for the experiment.
125 """
126 return self.meta['identifier']
128 @property
129 def procedure(self) -> ProcedureType:
130 """
131 The procedure for the experiment.
133 Returns
134 -------
135 ProcedureType
136 The procedure for the experiment.
137 """
138 return self._procedure
140 @procedure.setter
141 def procedure(self, procedure: ProcedureType):
142 """
143 Set the procedure for the experiment.
145 Parameters
146 ----------
147 procedure : ProcedureType
148 The procedure for the experiment. This can be a callable function,
149 a path to a script, or a Procedure object.
150 If path or str is provided, a ScriptProcedure object is created.
151 If callable is provided, a PickleProcedure object is created.
153 Raises
154 ------
155 ValueError
156 If the procedure type is not a valid type.
157 """
159 procedure_type = type(procedure)
160 if isinstance(procedure, (ScriptProcedure, PickleProcedure)):
161 pass
162 elif procedure_type is str or procedure_type is Path:
163 procedure = ScriptProcedure(procedure)
164 elif callable(procedure):
165 procedure = PickleProcedure(procedure)
166 else:
167 raise ValueError(f'Invalid procedure type: {procedure_type}')
169 self._procedure = procedure
171 @property
172 def status(self) -> Status:
173 """
174 The status of the experiment.
176 Returns
177 -------
178 str
179 The status of the experiment.
180 """
181 return self.meta['status']
183 @status.setter
184 def status(self, status: Union[str, Status]) -> None:
185 """
186 Set the status of the experiment.
188 Parameters
189 ----------
190 status : str
191 The status of the experiment. Valid statuses are:
192 'pending', 'submitted', 'running', 'completed', 'failed'.
194 Raises
195 ------
196 ValueError
197 If the status is not a valid status.
198 """
199 if isinstance(status, str):
200 status = Status(status)
201 self.meta['status'] = status
203 @property
204 def directory(self) -> Path:
205 """
206 The directory for the experiment.
208 Created as root_path/identifier-extension/
210 Returns
211 -------
212 Path
213 The directory for the experiment.
214 """
215 if not self.root_path.exists():
216 self.root_path.mkdir(parents=True)
218 directory = self.root_path / Path(f'{self.identifier}-{self.extension}/')
219 return directory
221 @property
222 def shephex_directory(self) -> Path:
223 return self.directory / self.shep_dir
225 ############################################################################
226 # Methods
227 ############################################################################
229 def dump(self) -> None:
230 """
231 Dump all the experiment data to the experiment directory, including
232 the options, meta, and procedure.
233 """
234 path = self.directory
235 if not path.exists():
236 path.mkdir(parents=True)
238 self.shephex_directory.mkdir(parents=True, exist_ok=True)
240 self._dump_options()
241 self._dump_meta()
242 self._dump_procedure()
244 def _dump_options(self) -> None:
245 self.options.dump(self.shephex_directory)
247 def _dump_procedure(self) -> None:
248 self.procedure.dump(self.shephex_directory)
250 def _dump_meta(self) -> None:
251 self.meta.dump(self.shephex_directory)
253 @classmethod
254 def load(
255 cls,
256 path: Union[str, Path],
257 override_procedure: Optional[ProcedureType] = None,
258 load_procedure: bool = True,
259 ) -> 'Experiment':
260 """
261 Load an experiment from a directory.
263 Parameters
264 ----------
265 path : Union[str, Path]
266 The path to the experiment directory.
267 override_procedure : Optional[ProcedureType], optional
268 Override the procedure object, by default None
269 load_procedure : bool, optional
270 Load the procedure object, by default True
272 Returns
273 -------
274 Experiment
275 An experiment object loaded from the directory.
276 """
278 path = Path(path)
280 # Load the meta file
281 meta = Meta.from_file(path / cls.shep_dir)
282 meta['status'] = Status(meta['status'])
284 # Load the options
285 with open(path / meta['options_path'], 'rb') as f:
286 options = json.load(f)
288 # Load the procedure
289 procedure = cls.load_procedure(path / cls.shep_dir, meta['procedure'], override_procedure, load_procedure)
291 # Create the experiment
292 experiment = cls(
293 *options['args'],
294 **options['kwargs'],
295 procedure=procedure,
296 root_path=path.parent,
297 meta=meta,
298 )
300 return experiment
302 @staticmethod
303 def load_procedure(
304 path: Path,
305 meta: dict,
306 override_procedure: Optional[ProcedureType] = None,
307 load_procedure: bool = True,
308 ) -> ProcedureType:
309 # Load the procedure
310 procedure_path = path / meta['name']
311 procedure_type = meta['type']
313 if override_procedure:
314 return override_procedure
315 elif not load_procedure:
316 return PickleProcedure(lambda: None)
318 if procedure_type == 'ScriptProcedure':
319 meta['path'] = str(procedure_path)
320 procedure = ScriptProcedure.from_metadata(meta)
322 elif procedure_type == 'PickleProcedure':
323 with open(procedure_path, 'rb') as f:
324 procedure = pkl.load(f)
326 return procedure
328 def _execute(
329 self, execution_directory: Optional[Union[Path, str]] = None
330 ) -> ExperimentResult:
331 """
332 Execute the experiment procedure.
334 Parameters
335 ----------
336 execution_directory : Optional[Union[Path, str]], optional
337 The directory where the experiment will be executed, defaults to
338 the experiment directory.
339 """
340 self.update_status(Status.running())
341 if self.procedure.context:
342 context = ExperimentContext(self.shephex_directory.resolve())
343 else:
344 context = None
346 if execution_directory is None:
347 execution_directory = self.directory
349 result = self.procedure._execute(
350 options=self.options,
351 directory=execution_directory,
352 shephex_directory=self.shephex_directory,
353 context=context,
354 )
355 self.meta.load(self.shephex_directory) # Reload the meta file
356 self.update_status(result.status)
357 return result
359 def update_status(self, status: Union[Status, str]) -> None:
360 """
361 Update the status of the experiment.
362 """
363 self.status = status
364 self._dump_meta()
366 def to_dict(self) -> dict:
367 """
368 Return a dictionary representation of the experiment. Used for printing
369 not for saving or comparing experiments.
371 Returns
372 -------
373 dict
374 A dictionary representation of the experiment.
375 """
376 experiment_dict = self.meta.get_dict()
377 experiment_dict.update(self.options.to_dict())
378 return experiment_dict
380 ############################################################################
381 # Magic Methods
382 ############################################################################
384 def __eq__(self, experiment: Any) -> bool:
385 """
386 Compare two experiments based on their options.
388 Parameters
389 ----------
390 experiment : Any
391 The experiment to compare.
393 Returns
394 -------
395 bool
396 True if the experiments have the same options, False otherwise.
397 """
398 if not isinstance(experiment, Experiment):
399 return False
401 if self.procedure != experiment.procedure:
402 return False
404 return experiment.options == self.options
406 def __repr__(self) -> str:
407 """
408 Return a string representation of the experiment.
410 Returns
411 -------
412 str
413 A string representation of the experiment.
414 """
415 rep_str = f'Experiment {self.identifier}'
416 for key, value in self.options.items():
417 rep_str += f'\n\t{key}: {value}'
418 rep_str += f'\n\tStatus: {self.status}'
419 rep_str += f'\n\tProcedure: {self.procedure.name}'
420 return rep_str