Coverage for python/lsst/pipe/base/quantum_graph/aggregator/_ingester.py: 26%

1# This file is part of pipe_base. 

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__ = ("Ingester",) 

31 

32import dataclasses 

33import logging 

34import time 

35import uuid 

36from collections import defaultdict 

37 

38from lsst.daf.butler import Butler, CollectionType, DatasetRef, DimensionGroup 

39from lsst.daf.butler.datastore.record_data import DatastoreRecordData 

40from lsst.daf.butler.registry import ConflictingDefinitionError 

41 

42from ...pipeline_graph import TaskImportMode 

43from .._common import DatastoreName 

44from .._predicted import PredictedDatasetModel, PredictedQuantumGraphComponents, PredictedQuantumGraphReader 

45from ._communicators import IngesterCommunicator 

46 

47 

48@dataclasses.dataclass 

49class Ingester: 

50 """A helper class for the provenance aggregator that handles ingestion into 

51 the central butler repository. 

52 """ 

53 

54 predicted_path: str 

55 """Path to the predicted quantum graph.""" 

56 

57 butler_path: str 

58 """Path or alias to the central butler repository.""" 

59 

60 comms: IngesterCommunicator 

61 """Communicator object for this worker.""" 

62 

63 predicted: PredictedQuantumGraphComponents = dataclasses.field(init=False) 

64 """Components of the predicted graph.""" 

65 

66 butler: Butler = dataclasses.field(init=False) 

67 """Client for the central butler repository.""" 

68 

69 n_datasets_ingested: int = 0 

70 """Total number of datasets ingested by this invocation.""" 

71 

72 n_datasets_skipped: int = 0 

73 """Total number of datasets skipped because they were already present.""" 

74 

75 n_producers_pending: int = 0 

76 """Number of quanta whose outputs are currently pending ingest.""" 

77 

78 refs_pending: defaultdict[DimensionGroup, list[DatasetRef]] = dataclasses.field( 

79 default_factory=lambda: defaultdict(list) 

80 ) 

81 """Dataset references pending ingest, grouped by their dimensions.""" 

82 

83 records_pending: dict[DatastoreName, DatastoreRecordData] = dataclasses.field(default_factory=dict) 

84 """Datastore records pending ingest, grouped by datastore name.""" 

85 

86 already_ingested: set[uuid.UUID] | None = None 

87 """A set of all dataset IDs already present in the output RUN 

88 collection. 

89 

90 If this is not `None`, the ingester is in defensive ingest mode, either 

91 because it was configured to query for these dataset IDs up front, or 

92 because a transaction failed due to a dataset already being present. 

93 """ 

94 

95 last_ingest_time: float = dataclasses.field(default_factory=time.time) 

96 """POSIX timestamp since the last ingest transaction concluded.""" 

97 

98 def __post_init__(self) -> None: 

99 self.comms.log.verbose("Reading from predicted quantum graph.") 

100 with PredictedQuantumGraphReader.open( 

101 self.predicted_path, import_mode=TaskImportMode.DO_NOT_IMPORT 

102 ) as reader: 

103 # We only need the header and pipeline graph. 

104 self.predicted = reader.components 

105 if self.comms.config.mock_storage_classes: 

106 import lsst.pipe.base.tests.mocks # noqa: F401 

107 self.comms.log.verbose("Initializing butler.") 

108 self.butler = Butler.from_config(self.butler_path, writeable=not self.comms.config.dry_run) 

109 

110 @property 

111 def n_datasets_pending(self) -> int: 

112 """The number of butler datasets currently pending.""" 

113 return sum(len(v) for v in self.refs_pending.values()) 

114 

115 @staticmethod 

116 def run(predicted_path: str, butler_path: str, comms: IngesterCommunicator) -> None: 

117 """Run the ingester. 

118 

119 Parameters 

120 ---------- 

121 predicted_path : `str` 

122 Path to the predicted quantum graph. 

123 butler_path : `str` 

124 Path or alias to the central butler repository. 

125 comms : `IngesterCommunicator` 

126 Communicator for the ingester. 

127 

128 Notes 

129 ----- 

130 This method is designed to run as the ``target`` in 

131 `WorkerContext.make_worker`. 

132 """ 

133 with comms: 

134 ingester = Ingester(predicted_path, butler_path, comms) 

135 ingester.loop() 

136 

137 def loop(self) -> None: 

138 """Run the main loop for the ingester.""" 

139 self.comms.log.verbose("Registering collections and dataset types.") 

140 if not self.comms.config.dry_run: 

141 if self.comms.config.register_dataset_types: 

142 self.predicted.pipeline_graph.register_dataset_types( 

143 self.butler, 

144 include_inputs=False, 

145 include_packages=True, 

146 include_configs=True, 

147 include_logs=True, 

148 ) 

149 self.butler.collections.register(self.predicted.header.output_run) 

150 # Updating the output chain cannot happen inside the caching 

151 # context. 

152 if self.comms.config.update_output_chain: 

153 self.update_output_chain() 

154 with self.butler.registry.caching_context(): 

155 if self.comms.config.defensive_ingest: 

156 self.fetch_already_ingested() 

157 self.comms.log.info("Startup completed in %ss.", time.time() - self.last_ingest_time) 

158 self.last_ingest_time = time.time() 

159 for ingest_request in self.comms.poll(): 

160 self.n_producers_pending += 1 

161 self.comms.log.debug(f"Got ingest request for producer {ingest_request.producer_id}.") 

162 self.update_pending(ingest_request.datasets, ingest_request.records) 

163 if self.n_datasets_pending > self.comms.config.ingest_batch_size: 

164 self.ingest() 

165 self.comms.log.info("All ingest requests received.") 

166 # We use 'while' in case this fails with a conflict and we switch 

167 # to defensive mode (should be at most two iterations). 

168 ingest_start_time = time.time() 

169 while self.n_datasets_pending: 

170 n_datasets = self.n_datasets_pending 

171 self.ingest() 

172 self.comms.log.verbose( 

173 "Gathered %d final datasets in %ss and ingested them in %ss.", 

174 n_datasets, 

175 ingest_start_time - self.last_ingest_time, 

176 time.time() - ingest_start_time, 

177 ) 

178 if self.n_producers_pending: 

179 # We can finish with returns pending if we filtered out all of 

180 # the datasets we started with as already existing. 

181 self.report() 

182 self.comms.log_progress( 

183 logging.INFO, 

184 f"Ingested {self.n_datasets_ingested} dataset(s); " 

185 f"skipped {self.n_datasets_skipped} already present.", 

186 ) 

187 

188 def ingest(self) -> None: 

189 """Ingest all pending datasets and report success to the supervisor.""" 

190 ingest_start_time = time.time() 

191 self.comms.log.verbose( 

192 "Gathered %d datasets from %d quanta in %ss.", 

193 self.n_datasets_pending, 

194 self.n_producers_pending, 

195 ingest_start_time - self.last_ingest_time, 

196 ) 

197 try: 

198 if not self.comms.config.dry_run: 

199 with self.butler.registry.transaction(): 

200 for refs in self.refs_pending.values(): 

201 self.butler.registry._importDatasets(refs, expand=False, assume_new=True) 

202 self.butler._datastore.import_records(self.records_pending) 

203 self.last_ingest_time = time.time() 

204 self.comms.log.verbose( 

205 "Ingested %d datasets from %d quanta in %ss.", 

206 self.n_datasets_pending, 

207 self.n_producers_pending, 

208 self.last_ingest_time - ingest_start_time, 

209 ) 

210 self.n_datasets_ingested += self.n_datasets_pending 

211 except ConflictingDefinitionError: 

212 if self.already_ingested is None: 

213 self.comms.log_progress( 

214 logging.INFO, 

215 "Some outputs seem to have already been ingested; querying for existing datasets and " 

216 "switching to defensive ingest mode.", 

217 ) 

218 self.fetch_already_ingested() 

219 # We just return instead of trying again immediately because we 

220 # might have just shrunk the number of pending datasets below 

221 # the batch threshold. 

222 return 

223 else: 

224 raise 

225 self.report() 

226 self.records_pending.clear() 

227 self.refs_pending.clear() 

228 

229 def report(self) -> None: 

230 """Report a successful ingest to the supervisor.""" 

231 self.comms.report_ingest(self.n_producers_pending) 

232 self.n_producers_pending = 0 

233 

234 def fetch_already_ingested(self) -> None: 

235 """Query for the UUIDs of all dataset already present in the output 

236 RUN collection, and filter and pending datasets accordingly. 

237 """ 

238 self.comms.log.info("Fetching all UUIDs in output collection %r.", self.predicted.header.output_run) 

239 self.already_ingested = set( 

240 self.butler.registry._fetch_run_dataset_ids(self.predicted.header.output_run) 

241 ) 

242 kept: set[uuid.UUID] = set() 

243 for dimensions, refs in self.refs_pending.items(): 

244 filtered_refs: list[DatasetRef] = [] 

245 for ref in refs: 

246 if ref.id not in self.already_ingested: 

247 kept.add(ref.id) 

248 filtered_refs.append(ref) 

249 else: 

250 self.n_datasets_skipped += 1 

251 self.refs_pending[dimensions] = filtered_refs 

252 for datastore_name, datastore_records in list(self.records_pending.items()): 

253 if (filtered_records := datastore_records.subset(kept)) is not None: 

254 self.records_pending[datastore_name] = filtered_records 

255 else: 

256 del self.records_pending[datastore_name] 

257 

258 def update_pending( 

259 self, datasets: list[PredictedDatasetModel], records: dict[DatastoreName, DatastoreRecordData] 

260 ) -> None: 

261 """Add an ingest request to the pending-ingest data structures. 

262 

263 Parameters 

264 ---------- 

265 datasets : `list` [ `PredictedDatasetModel` ] 

266 Registry information about the datasets. 

267 records : `dict` [ `str`, \ 

268 `lsst.daf.butler.datastore.record_data.DatastoreRecordData` ] 

269 Datastore information about the datasets. 

270 """ 

271 n_given = len(datasets) 

272 if self.already_ingested is not None: 

273 datasets = [d for d in datasets if d.dataset_id not in self.already_ingested] 

274 kept = {d.dataset_id for d in datasets} 

275 self.n_datasets_skipped += n_given - len(kept) 

276 records = { 

277 datastore_name: filtered_records 

278 for datastore_name, original_records in records.items() 

279 if (filtered_records := original_records.subset(kept)) is not None 

280 } 

281 for dataset in datasets: 

282 ref = self.predicted.make_dataset_ref(dataset) 

283 self.refs_pending[ref.datasetType.dimensions].append(ref) 

284 for datastore_name, datastore_records in records.items(): 

285 if (existing_records := self.records_pending.get(datastore_name)) is not None: 

286 existing_records.update(datastore_records) 

287 else: 

288 self.records_pending[datastore_name] = datastore_records 

289 

290 def update_output_chain(self) -> None: 

291 """Update the output CHAINED collection to include the output RUN 

292 collection (and the inputs, if the output CHAINED collection does not 

293 exist). 

294 

295 Notes 

296 ----- 

297 This method cannot be called inside the registry caching context. 

298 """ 

299 if self.predicted.header.output is None: 

300 return 

301 self.comms.log.info( 

302 "Updating output collection %s to include %s.", 

303 self.predicted.header.output, 

304 self.predicted.header.output_run, 

305 ) 

306 if self.butler.collections.register(self.predicted.header.output, CollectionType.CHAINED): 

307 # Chain is new; need to add inputs, but we want to flatten them 

308 # first. 

309 if self.predicted.header.inputs: 

310 flattened = self.butler.collections.query(self.predicted.header.inputs, flatten_chains=True) 

311 self.butler.collections.extend_chain(self.predicted.header.output, flattened) 

312 self.butler.collections.prepend_chain(self.predicted.header.output, self.predicted.header.output_run)