Coverage for python/lsst/ctrl/mpexec/preExecInit.py: 18%

155 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-23 10:58 +0000

1# This file is part of ctrl_mpexec. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27 

28from __future__ import annotations 

29 

30__all__ = ["PreExecInit"] 

31 

32# ------------------------------- 

33# Imports of standard modules -- 

34# ------------------------------- 

35import abc 

36import logging 

37from collections.abc import Iterable, Iterator 

38from contextlib import contextmanager 

39from typing import TYPE_CHECKING, Any 

40 

41# ----------------------------- 

42# Imports for other modules -- 

43# ----------------------------- 

44from lsst.daf.butler import DatasetRef, DatasetType 

45from lsst.daf.butler.registry import ConflictingDefinitionError 

46from lsst.pipe.base import PipelineDatasetTypes 

47from lsst.utils.packages import Packages 

48 

49if TYPE_CHECKING: 

50 from lsst.daf.butler import Butler, LimitedButler 

51 from lsst.pipe.base import QuantumGraph, TaskDef, TaskFactory 

52 

53_LOG = logging.getLogger(__name__) 

54 

55 

56class MissingReferenceError(Exception): 

57 """Exception raised when resolved reference is missing from graph.""" 

58 

59 pass 

60 

61 

62def _compare_packages(old_packages: Packages, new_packages: Packages) -> None: 

63 """Compare two versions of Packages. 

64 

65 Parameters 

66 ---------- 

67 old_packages : `Packages` 

68 Previously recorded package versions. 

69 new_packages : `Packages` 

70 New set of package versions. 

71 

72 Raises 

73 ------ 

74 TypeError 

75 Raised if parameters are inconsistent. 

76 """ 

77 diff = new_packages.difference(old_packages) 

78 if diff: 

79 versions_str = "; ".join(f"{pkg}: {diff[pkg][1]} vs {diff[pkg][0]}" for pkg in diff) 

80 raise TypeError(f"Package versions mismatch: ({versions_str})") 

81 else: 

82 _LOG.debug("new packages are consistent with old") 

83 

84 

85class PreExecInitBase(abc.ABC): 

86 """Common part of the implementation of PreExecInit classes that does not 

87 depend on Butler type. 

88 """ 

89 

90 def __init__(self, butler: LimitedButler, taskFactory: TaskFactory, extendRun: bool): 

91 self.butler = butler 

92 self.taskFactory = taskFactory 

93 self.extendRun = extendRun 

94 

95 def initialize( 

96 self, 

97 graph: QuantumGraph, 

98 saveInitOutputs: bool = True, 

99 registerDatasetTypes: bool = False, 

100 saveVersions: bool = True, 

101 ) -> None: 

102 """Perform all initialization steps. 

103 

104 Convenience method to execute all initialization steps. Instead of 

105 calling this method and providing all options it is also possible to 

106 call methods individually. 

107 

108 Parameters 

109 ---------- 

110 graph : `~lsst.pipe.base.QuantumGraph` 

111 Execution graph. 

112 saveInitOutputs : `bool`, optional 

113 If ``True`` (default) then save "init outputs", configurations, 

114 and package versions to butler. 

115 registerDatasetTypes : `bool`, optional 

116 If ``True`` then register dataset types in registry, otherwise 

117 they must be already registered. 

118 saveVersions : `bool`, optional 

119 If ``False`` then do not save package versions even if 

120 ``saveInitOutputs`` is set to ``True``. 

121 """ 

122 # register dataset types or check consistency 

123 self.initializeDatasetTypes(graph, registerDatasetTypes) 

124 

125 # Save task initialization data or check that saved data 

126 # is consistent with what tasks would save 

127 if saveInitOutputs: 

128 self.saveInitOutputs(graph) 

129 self.saveConfigs(graph) 

130 if saveVersions: 

131 self.savePackageVersions(graph) 

132 

133 @abc.abstractmethod 

134 def initializeDatasetTypes(self, graph: QuantumGraph, registerDatasetTypes: bool = False) -> None: 

135 """Save or check DatasetTypes output by the tasks in a graph. 

136 

137 Iterates over all DatasetTypes for all tasks in a graph and either 

138 tries to add them to registry or compares them to existing ones. 

139 

140 Parameters 

141 ---------- 

142 graph : `~lsst.pipe.base.QuantumGraph` 

143 Execution graph. 

144 registerDatasetTypes : `bool`, optional 

145 If ``True`` then register dataset types in registry, otherwise 

146 they must be already registered. 

147 

148 Raises 

149 ------ 

150 ValueError 

151 Raised if existing DatasetType is different from DatasetType 

152 in a graph. 

153 KeyError 

154 Raised if ``registerDatasetTypes`` is ``False`` and DatasetType 

155 does not exist in registry. 

156 """ 

157 raise NotImplementedError() 

158 

159 def saveInitOutputs(self, graph: QuantumGraph) -> None: 

160 """Write any datasets produced by initializing tasks in a graph. 

161 

162 Parameters 

163 ---------- 

164 graph : `~lsst.pipe.base.QuantumGraph` 

165 Execution graph. 

166 

167 Raises 

168 ------ 

169 TypeError 

170 Raised if the type of existing object in butler is different from 

171 new data. 

172 """ 

173 _LOG.debug("Will save InitOutputs for all tasks") 

174 for taskDef in self._task_iter(graph): 

175 init_input_refs = graph.initInputRefs(taskDef) or [] 

176 task = self.taskFactory.makeTask(taskDef, self.butler, init_input_refs) 

177 for name in taskDef.connections.initOutputs: 

178 attribute = getattr(taskDef.connections, name) 

179 init_output_refs = graph.initOutputRefs(taskDef) or [] 

180 init_output_ref, obj_from_store = self._find_dataset(init_output_refs, attribute.name) 

181 if init_output_ref is None: 

182 raise ValueError(f"Cannot find dataset reference for init output {name} in a graph") 

183 init_output_var = getattr(task, name) 

184 

185 if obj_from_store is not None: 

186 _LOG.debug( 

187 "Retrieving InitOutputs for task=%s key=%s dsTypeName=%s", task, name, attribute.name 

188 ) 

189 obj_from_store = self.butler.get(init_output_ref) 

190 # Types are supposed to be identical. 

191 # TODO: Check that object contents is identical too. 

192 if type(obj_from_store) is not type(init_output_var): 

193 raise TypeError( 

194 f"Stored initOutput object type {type(obj_from_store)} " 

195 "is different from task-generated type " 

196 f"{type(init_output_var)} for task {taskDef}" 

197 ) 

198 else: 

199 _LOG.debug("Saving InitOutputs for task=%s key=%s", taskDef.label, name) 

200 # This can still raise if there is a concurrent write. 

201 self.butler.put(init_output_var, init_output_ref) 

202 

203 def saveConfigs(self, graph: QuantumGraph) -> None: 

204 """Write configurations for pipeline tasks to butler or check that 

205 existing configurations are equal to the new ones. 

206 

207 Parameters 

208 ---------- 

209 graph : `~lsst.pipe.base.QuantumGraph` 

210 Execution graph. 

211 

212 Raises 

213 ------ 

214 TypeError 

215 Raised if existing object in butler is different from new data. 

216 Exception 

217 Raised if ``extendRun`` is `False` and datasets already exists. 

218 Content of a butler collection should not be changed if exception 

219 is raised. 

220 """ 

221 

222 def logConfigMismatch(msg: str) -> None: 

223 """Log messages about configuration mismatch.""" 

224 _LOG.fatal("Comparing configuration: %s", msg) 

225 

226 _LOG.debug("Will save Configs for all tasks") 

227 # start transaction to rollback any changes on exceptions 

228 with self.transaction(): 

229 for taskDef in self._task_iter(graph): 

230 # Config dataset ref is stored in task init outputs, but it 

231 # may be also be missing. 

232 task_output_refs = graph.initOutputRefs(taskDef) 

233 if task_output_refs is None: 

234 continue 

235 

236 config_ref, old_config = self._find_dataset(task_output_refs, taskDef.configDatasetName) 

237 if config_ref is None: 

238 continue 

239 

240 if old_config is not None: 

241 if not taskDef.config.compare(old_config, shortcut=False, output=logConfigMismatch): 

242 raise TypeError( 

243 f"Config does not match existing task config {taskDef.configDatasetName!r} in " 

244 "butler; tasks configurations must be consistent within the same run collection" 

245 ) 

246 else: 

247 _LOG.debug( 

248 "Saving Config for task=%s dataset type=%s", taskDef.label, taskDef.configDatasetName 

249 ) 

250 self.butler.put(taskDef.config, config_ref) 

251 

252 def savePackageVersions(self, graph: QuantumGraph) -> None: 

253 """Write versions of software packages to butler. 

254 

255 Parameters 

256 ---------- 

257 graph : `~lsst.pipe.base.QuantumGraph` 

258 Execution graph. 

259 

260 Raises 

261 ------ 

262 TypeError 

263 Raised if existing object in butler is incompatible with new data. 

264 """ 

265 packages = Packages.fromSystem() 

266 _LOG.debug("want to save packages: %s", packages) 

267 

268 # start transaction to rollback any changes on exceptions 

269 with self.transaction(): 

270 # Packages dataset ref is stored in graph's global init outputs, 

271 # but it may be also be missing. 

272 

273 packages_ref, old_packages = self._find_dataset( 

274 graph.globalInitOutputRefs(), PipelineDatasetTypes.packagesDatasetName 

275 ) 

276 if packages_ref is None: 

277 return 

278 

279 if old_packages is not None: 

280 # Note that because we can only detect python modules that have 

281 # been imported, the stored list of products may be more or 

282 # less complete than what we have now. What's important is 

283 # that the products that are in common have the same version. 

284 _compare_packages(old_packages, packages) 

285 # Update the old set of packages in case we have more packages 

286 # that haven't been persisted. 

287 extra = packages.extra(old_packages) 

288 if extra: 

289 _LOG.debug("extra packages: %s", extra) 

290 old_packages.update(packages) 

291 # have to remove existing dataset first, butler has no 

292 # replace option. 

293 self.butler.pruneDatasets([packages_ref], unstore=True, purge=True) 

294 self.butler.put(old_packages, packages_ref) 

295 else: 

296 self.butler.put(packages, packages_ref) 

297 

298 def _find_dataset( 

299 self, refs: Iterable[DatasetRef], dataset_type: str 

300 ) -> tuple[DatasetRef | None, Any | None]: 

301 """Find a ref with a given dataset type name in a list of references 

302 and try to retrieve its data from butler. 

303 

304 Parameters 

305 ---------- 

306 refs : `~collections.abc.Iterable` [ `~lsst.daf.butler.DatasetRef` ] 

307 References to check for matching dataset type. 

308 dataset_type : `str` 

309 Name of a dataset type to look for. 

310 

311 Returns 

312 ------- 

313 ref : `~lsst.daf.butler.DatasetRef` or `None` 

314 Dataset reference or `None` if there is no matching dataset type. 

315 data : `Any` 

316 An existing object extracted from butler, `None` if ``ref`` is 

317 `None` or if there is no existing object for that reference. 

318 """ 

319 ref: DatasetRef | None = None 

320 for ref in refs: 

321 if ref.datasetType.name == dataset_type: 

322 break 

323 else: 

324 return None, None 

325 

326 try: 

327 data = self.butler.get(ref) 

328 if data is not None and not self.extendRun: 

329 # It must not exist unless we are extending run. 

330 raise ConflictingDefinitionError(f"Dataset {ref} already exists in butler") 

331 except (LookupError, FileNotFoundError): 

332 data = None 

333 return ref, data 

334 

335 def _task_iter(self, graph: QuantumGraph) -> Iterator[TaskDef]: 

336 """Iterate over TaskDefs in a graph, return only tasks that have one or 

337 more associated quanta. 

338 """ 

339 for taskDef in graph.iterTaskGraph(): 

340 if graph.getNumberOfQuantaForTask(taskDef) > 0: 

341 yield taskDef 

342 

343 @contextmanager 

344 def transaction(self) -> Iterator[None]: 

345 """Context manager for transaction. 

346 

347 Default implementation has no transaction support. 

348 """ 

349 yield 

350 

351 

352class PreExecInit(PreExecInitBase): 

353 """Initialization of registry for QuantumGraph execution. 

354 

355 This class encapsulates all necessary operations that have to be performed 

356 on butler and registry to prepare them for QuantumGraph execution. 

357 

358 Parameters 

359 ---------- 

360 butler : `~lsst.daf.butler.Butler` 

361 Data butler instance. 

362 taskFactory : `~lsst.pipe.base.TaskFactory` 

363 Task factory. 

364 extendRun : `bool`, optional 

365 If `True` then do not try to overwrite any datasets that might exist 

366 in ``butler.run``; instead compare them when appropriate/possible. If 

367 `False`, then any existing conflicting dataset will cause a butler 

368 exception to be raised. 

369 """ 

370 

371 def __init__(self, butler: Butler, taskFactory: TaskFactory, extendRun: bool = False): 

372 super().__init__(butler, taskFactory, extendRun) 

373 self.full_butler = butler 

374 if self.extendRun and self.full_butler.run is None: 

375 raise RuntimeError( 

376 "Cannot perform extendRun logic unless butler is initialized " 

377 "with a default output RUN collection." 

378 ) 

379 

380 @contextmanager 

381 def transaction(self) -> Iterator[None]: 

382 # dosctring inherited 

383 with self.full_butler.transaction(): 

384 yield 

385 

386 def initializeDatasetTypes(self, graph: QuantumGraph, registerDatasetTypes: bool = False) -> None: 

387 # docstring inherited 

388 pipeline = graph.taskGraph 

389 pipelineDatasetTypes = PipelineDatasetTypes.fromPipeline( 

390 pipeline, registry=self.full_butler.registry, include_configs=True, include_packages=True 

391 ) 

392 

393 for datasetTypes, is_input in ( 

394 (pipelineDatasetTypes.initIntermediates, True), 

395 (pipelineDatasetTypes.initOutputs, False), 

396 (pipelineDatasetTypes.intermediates, True), 

397 (pipelineDatasetTypes.outputs, False), 

398 ): 

399 self._register_output_dataset_types(registerDatasetTypes, datasetTypes, is_input) 

400 

401 def _register_output_dataset_types( 

402 self, registerDatasetTypes: bool, datasetTypes: Iterable[DatasetType], is_input: bool 

403 ) -> None: 

404 def _check_compatibility(datasetType: DatasetType, expected: DatasetType, is_input: bool) -> bool: 

405 # These are output dataset types so check for compatibility on put. 

406 is_compatible = expected.is_compatible_with(datasetType) 

407 

408 if is_input: 

409 # This dataset type is also used for input so must be 

410 # compatible on get as ell. 

411 is_compatible = is_compatible and datasetType.is_compatible_with(expected) 

412 

413 if is_compatible: 

414 _LOG.debug( 

415 "The dataset type configurations differ (%s from task != %s from registry) " 

416 "but the storage classes are compatible. Can continue.", 

417 datasetType, 

418 expected, 

419 ) 

420 return is_compatible 

421 

422 missing_datasetTypes = set() 

423 for datasetType in datasetTypes: 

424 # Only composites are registered, no components, and by this point 

425 # the composite should already exist. 

426 if registerDatasetTypes and not datasetType.isComponent(): 

427 _LOG.debug("Registering DatasetType %s with registry", datasetType) 

428 # this is a no-op if it already exists and is consistent, 

429 # and it raises if it is inconsistent. 

430 try: 

431 self.full_butler.registry.registerDatasetType(datasetType) 

432 except ConflictingDefinitionError: 

433 if not _check_compatibility( 

434 datasetType, self.full_butler.get_dataset_type(datasetType.name), is_input 

435 ): 

436 raise 

437 else: 

438 _LOG.debug("Checking DatasetType %s against registry", datasetType) 

439 try: 

440 expected = self.full_butler.get_dataset_type(datasetType.name) 

441 except KeyError: 

442 # Likely means that --register-dataset-types is forgotten. 

443 missing_datasetTypes.add(datasetType.name) 

444 continue 

445 if expected != datasetType: 

446 if not _check_compatibility(datasetType, expected, is_input): 

447 raise ValueError( 

448 f"DatasetType configuration does not match Registry: {datasetType} != {expected}" 

449 ) 

450 

451 if missing_datasetTypes: 

452 plural = "s" if len(missing_datasetTypes) != 1 else "" 

453 raise KeyError( 

454 f"Missing dataset type definition{plural}: {', '.join(missing_datasetTypes)}. " 

455 "Dataset types have to be registered with either `butler register-dataset-type` or " 

456 "passing `--register-dataset-types` option to `pipetask run`." 

457 ) 

458 

459 

460class PreExecInitLimited(PreExecInitBase): 

461 """Initialization of registry for QuantumGraph execution. 

462 

463 This class works with LimitedButler and expects that all references in 

464 QuantumGraph are resolved. 

465 

466 Parameters 

467 ---------- 

468 butler : `~lsst.daf.butler.LimitedButler` 

469 Limited data butler instance. 

470 taskFactory : `~lsst.pipe.base.TaskFactory` 

471 Task factory. 

472 """ 

473 

474 def __init__(self, butler: LimitedButler, taskFactory: TaskFactory): 

475 super().__init__(butler, taskFactory, False) 

476 

477 def initializeDatasetTypes(self, graph: QuantumGraph, registerDatasetTypes: bool = False) -> None: 

478 # docstring inherited 

479 # With LimitedButler we never create or check dataset types. 

480 pass