Coverage for src/shephex/experiment/experiment.py: 99%

141 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-03-30 12:30 +0200

1""" 

2Experiment class definition. 

3""" 

4 

5import json 

6import pickle as pkl 

7from pathlib import Path 

8from time import time 

9from typing import Any, Callable, Optional, TypeAlias, Union 

10 

11import shortuuid 

12 

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 

19 

20ProcedureType: TypeAlias = Union[PickleProcedure, ScriptProcedure] 

21 

22 

23class Experiment: 

24 """ 

25 Experiment class definition, shephex's central object. 

26 """ 

27 

28 extension = 'exp' 

29 shep_dir = 'shephex' 

30 

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. 

44 

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 """ 

64 

65 # Root path 

66 self.root_path = Path(root_path).resolve() if root_path else Path.cwd() 

67 

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.") 

75 

76 # Set the options 

77 args = args if args else [] 

78 kwargs = kwargs if kwargs else {} 

79 self.options = Options(*args, **kwargs) 

80 

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() 

85 

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 

95 

96 ############################################################################ 

97 # Properties 

98 ############################################################################ 

99 

100 @property 

101 def root_path(self) -> Path: 

102 """ 

103 The root path for the experiment. 

104 

105 Returns 

106 ------- 

107 Path 

108 The root path for the experiment 

109 """ 

110 return self._root_path 

111 

112 @root_path.setter 

113 def root_path(self, root_path: Path) -> None: 

114 self._root_path = Path(root_path) 

115 

116 @property 

117 def identifier(self) -> str: 

118 """ 

119 The identifier for the experiment. 

120 

121 Returns 

122 ------- 

123 str 

124 The identifier for the experiment. 

125 """ 

126 return self.meta['identifier'] 

127 

128 @property 

129 def procedure(self) -> ProcedureType: 

130 """ 

131 The procedure for the experiment. 

132 

133 Returns 

134 ------- 

135 ProcedureType 

136 The procedure for the experiment. 

137 """ 

138 return self._procedure 

139 

140 @procedure.setter 

141 def procedure(self, procedure: ProcedureType): 

142 """ 

143 Set the procedure for the experiment. 

144 

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. 

152 

153 Raises 

154 ------ 

155 ValueError 

156 If the procedure type is not a valid type. 

157 """ 

158 

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}') 

168 

169 self._procedure = procedure 

170 

171 @property 

172 def status(self) -> Status: 

173 """ 

174 The status of the experiment. 

175 

176 Returns 

177 ------- 

178 str 

179 The status of the experiment. 

180 """ 

181 return self.meta['status'] 

182 

183 @status.setter 

184 def status(self, status: Union[str, Status]) -> None: 

185 """ 

186 Set the status of the experiment. 

187 

188 Parameters 

189 ---------- 

190 status : str 

191 The status of the experiment. Valid statuses are: 

192 'pending', 'submitted', 'running', 'completed', 'failed'. 

193 

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 

202 

203 @property 

204 def directory(self) -> Path: 

205 """ 

206 The directory for the experiment. 

207 

208 Created as root_path/identifier-extension/ 

209 

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) 

217 

218 directory = self.root_path / Path(f'{self.identifier}-{self.extension}/') 

219 return directory 

220 

221 @property 

222 def shephex_directory(self) -> Path: 

223 return self.directory / self.shep_dir 

224 

225 ############################################################################ 

226 # Methods 

227 ############################################################################ 

228 

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) 

237 

238 self.shephex_directory.mkdir(parents=True, exist_ok=True) 

239 

240 self._dump_options() 

241 self._dump_meta() 

242 self._dump_procedure() 

243 

244 def _dump_options(self) -> None: 

245 self.options.dump(self.shephex_directory) 

246 

247 def _dump_procedure(self) -> None: 

248 self.procedure.dump(self.shephex_directory) 

249 

250 def _dump_meta(self) -> None: 

251 self.meta.dump(self.shephex_directory) 

252 

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. 

262 

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 

271 

272 Returns 

273 ------- 

274 Experiment 

275 An experiment object loaded from the directory. 

276 """ 

277 

278 path = Path(path) 

279 

280 # Load the meta file 

281 meta = Meta.from_file(path / cls.shep_dir) 

282 meta['status'] = Status(meta['status']) 

283 

284 # Load the options 

285 with open(path / meta['options_path'], 'rb') as f: 

286 options = json.load(f) 

287 

288 # Load the procedure 

289 procedure = cls.load_procedure(path / cls.shep_dir, meta['procedure'], override_procedure, load_procedure) 

290 

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 ) 

299 

300 return experiment 

301 

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'] 

312 

313 if override_procedure: 

314 return override_procedure 

315 elif not load_procedure: 

316 return PickleProcedure(lambda: None) 

317 

318 if procedure_type == 'ScriptProcedure': 

319 meta['path'] = str(procedure_path) 

320 procedure = ScriptProcedure.from_metadata(meta) 

321 

322 elif procedure_type == 'PickleProcedure': 

323 with open(procedure_path, 'rb') as f: 

324 procedure = pkl.load(f) 

325 

326 return procedure 

327 

328 def _execute( 

329 self, execution_directory: Optional[Union[Path, str]] = None 

330 ) -> ExperimentResult: 

331 """ 

332 Execute the experiment procedure. 

333 

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 

345 

346 if execution_directory is None: 

347 execution_directory = self.directory 

348 

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 

358 

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() 

365 

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. 

370 

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 

379 

380 ############################################################################ 

381 # Magic Methods 

382 ############################################################################ 

383 

384 def __eq__(self, experiment: Any) -> bool: 

385 """ 

386 Compare two experiments based on their options. 

387 

388 Parameters 

389 ---------- 

390 experiment : Any 

391 The experiment to compare. 

392 

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 

400 

401 if self.procedure != experiment.procedure: 

402 return False 

403 

404 return experiment.options == self.options 

405 

406 def __repr__(self) -> str: 

407 """ 

408 Return a string representation of the experiment. 

409 

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