Coverage for python / lsst / pipe / base / quantum_graph / aggregator / _writer.py: 24%
219 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:14 -0800
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/>.
28from __future__ import annotations
30__all__ = ("Writer",)
32import dataclasses
33import itertools
34import logging
35import operator
36import uuid
37from typing import TypeVar
39import networkx
40import zstandard
42from lsst.utils.packages import Packages
44from ... import automatic_connection_constants as acc
45from ...pipeline_graph import TaskImportMode
46from .._common import BaseQuantumGraphWriter
47from .._multiblock import Compressor, MultiblockWriter
48from .._predicted import PredictedDatasetModel, PredictedQuantumGraphComponents, PredictedQuantumGraphReader
49from .._provenance import (
50 DATASET_ADDRESS_INDEX,
51 DATASET_MB_NAME,
52 LOG_ADDRESS_INDEX,
53 LOG_MB_NAME,
54 METADATA_ADDRESS_INDEX,
55 METADATA_MB_NAME,
56 QUANTUM_ADDRESS_INDEX,
57 QUANTUM_MB_NAME,
58 ProvenanceDatasetModel,
59 ProvenanceInitQuantaModel,
60 ProvenanceInitQuantumModel,
61 ProvenanceQuantumModel,
62)
63from ._communicators import WriterCommunicator
64from ._structs import WriteRequest
67@dataclasses.dataclass
68class _DataWriters:
69 """A struct of low-level writer objects for the main components of a
70 provenance quantum graph.
72 Parameters
73 ----------
74 comms : `WriterCommunicator`
75 Communicator helper object for the writer.
76 predicted : `.PredictedQuantumGraphComponents`
77 Components of the predicted graph.
78 indices : `dict` [ `uuid.UUID`, `int` ]
79 Mapping from UUID to internal integer ID, including both quanta and
80 datasets.
81 compressor : `Compressor`
82 Object that can compress `bytes`.
83 cdict_data : `bytes` or `None`, optional
84 Bytes representation of the compression dictionary used by the
85 compressor.
86 """
88 def __init__(
89 self,
90 comms: WriterCommunicator,
91 predicted: PredictedQuantumGraphComponents,
92 indices: dict[uuid.UUID, int],
93 compressor: Compressor,
94 cdict_data: bytes | None = None,
95 ) -> None:
96 assert comms.config.output_path is not None
97 header = predicted.header.model_copy()
98 header.graph_type = "provenance"
99 self.graph = comms.enter(
100 BaseQuantumGraphWriter.open(
101 comms.config.output_path,
102 header,
103 predicted.pipeline_graph,
104 indices,
105 address_filename="nodes",
106 compressor=compressor,
107 cdict_data=cdict_data,
108 ),
109 on_close="Finishing writing provenance quantum graph.",
110 is_progress_log=True,
111 )
112 self.graph.address_writer.addresses = [{}, {}, {}, {}]
113 self.logs = comms.enter(
114 MultiblockWriter.open_in_zip(self.graph.zf, LOG_MB_NAME, header.int_size, use_tempfile=True),
115 on_close="Copying logs into zip archive.",
116 is_progress_log=True,
117 )
118 self.graph.address_writer.addresses[LOG_ADDRESS_INDEX] = self.logs.addresses
119 self.metadata = comms.enter(
120 MultiblockWriter.open_in_zip(self.graph.zf, METADATA_MB_NAME, header.int_size, use_tempfile=True),
121 on_close="Copying metadata into zip archive.",
122 is_progress_log=True,
123 )
124 self.graph.address_writer.addresses[METADATA_ADDRESS_INDEX] = self.metadata.addresses
125 self.datasets = comms.enter(
126 MultiblockWriter.open_in_zip(self.graph.zf, DATASET_MB_NAME, header.int_size, use_tempfile=True),
127 on_close="Copying dataset provenance into zip archive.",
128 is_progress_log=True,
129 )
130 self.graph.address_writer.addresses[DATASET_ADDRESS_INDEX] = self.datasets.addresses
131 self.quanta = comms.enter(
132 MultiblockWriter.open_in_zip(self.graph.zf, QUANTUM_MB_NAME, header.int_size, use_tempfile=True),
133 on_close="Copying quantum provenance into zip archive.",
134 is_progress_log=True,
135 )
136 self.graph.address_writer.addresses[QUANTUM_ADDRESS_INDEX] = self.quanta.addresses
138 graph: BaseQuantumGraphWriter
139 """The parent graph writer."""
141 datasets: MultiblockWriter
142 """A writer for dataset provenance."""
144 quanta: MultiblockWriter
145 """A writer for quantum provenance."""
147 metadata: MultiblockWriter
148 """A writer for metadata content."""
150 logs: MultiblockWriter
151 """A writer for log content."""
153 @property
154 def compressor(self) -> Compressor:
155 """Object that should be used to compress all JSON blocks."""
156 return self.graph.compressor
159@dataclasses.dataclass
160class Writer:
161 """A helper class for the provenance aggregator actually writes the
162 provenance quantum graph file.
163 """
165 predicted_path: str
166 """Path to the predicted quantum graph."""
168 comms: WriterCommunicator
169 """Communicator object for this worker."""
171 predicted: PredictedQuantumGraphComponents = dataclasses.field(init=False)
172 """Components of the predicted quantum graph."""
174 existing_init_outputs: dict[uuid.UUID, set[uuid.UUID]] = dataclasses.field(default_factory=dict)
175 """Mapping that tracks which init-outputs exist.
177 This mapping is updated as scanners inform the writer about init-output
178 existence, since we want to write that provenance information out only at
179 the end.
180 """
182 indices: dict[uuid.UUID, int] = dataclasses.field(default_factory=dict)
183 """Mapping from UUID to internal integer ID, including both quanta and
184 datasets.
186 This is fully initialized at construction.
187 """
189 output_dataset_ids: set[uuid.UUID] = dataclasses.field(default_factory=set)
190 """The IDs of all datasets that are produced by this graph.
192 This is fully initialized at construction.
193 """
195 overall_inputs: dict[uuid.UUID, PredictedDatasetModel] = dataclasses.field(default_factory=dict)
196 """All datasets that are not produced by any quantum in this graph."""
198 xgraph: networkx.DiGraph = dataclasses.field(default_factory=networkx.DiGraph)
199 """A bipartite NetworkX graph linking datasets to quanta and quanta to
200 datasets.
202 This is fully initialized at construction. There are no node or edge
203 attributes in this graph; we only need it to store adjacency information
204 with datasets as well as with quanta.
205 """
207 pending_compression_training: list[WriteRequest] = dataclasses.field(default_factory=list)
208 """Unprocessed quantum scans that are being accumulated in order to
209 build a compression dictionary.
210 """
212 def __post_init__(self) -> None:
213 assert self.comms.config.output_path is not None, "Writer should not be used if writing is disabled."
214 self.comms.log.info("Reading predicted quantum graph.")
215 with PredictedQuantumGraphReader.open(
216 self.predicted_path, import_mode=TaskImportMode.DO_NOT_IMPORT
217 ) as reader:
218 self.comms.check_for_cancel()
219 reader.read_init_quanta()
220 self.comms.check_for_cancel()
221 reader.read_quantum_datasets()
222 self.predicted = reader.components
223 for predicted_init_quantum in self.predicted.init_quanta.root:
224 self.existing_init_outputs[predicted_init_quantum.quantum_id] = set()
225 self.comms.check_for_cancel()
226 self.comms.log.info("Generating integer indexes and identifying outputs.")
227 self._populate_indices_and_outputs()
228 self.comms.check_for_cancel()
229 self._populate_xgraph_and_inputs()
230 self.comms.check_for_cancel()
231 self.comms.log_progress(
232 # We add one here for 'packages', which we do ingest but don't
233 # record provenance for.
234 logging.INFO,
235 f"Graph has {len(self.output_dataset_ids) + 1} predicted output dataset(s).",
236 )
238 def _populate_indices_and_outputs(self) -> None:
239 all_uuids = set(self.predicted.quantum_datasets.keys())
240 for quantum in self.comms.periodically_check_for_cancel(
241 itertools.chain(
242 self.predicted.init_quanta.root,
243 self.predicted.quantum_datasets.values(),
244 )
245 ):
246 if not quantum.task_label:
247 # Skip the 'packages' producer quantum.
248 continue
249 all_uuids.update(quantum.iter_input_dataset_ids())
250 self.output_dataset_ids.update(quantum.iter_output_dataset_ids())
251 all_uuids.update(self.output_dataset_ids)
252 self.indices = {
253 node_id: node_index
254 for node_index, node_id in self.comms.periodically_check_for_cancel(
255 enumerate(sorted(all_uuids, key=operator.attrgetter("int")))
256 )
257 }
259 def _populate_xgraph_and_inputs(self) -> None:
260 for predicted_quantum in self.comms.periodically_check_for_cancel(
261 itertools.chain(
262 self.predicted.init_quanta.root,
263 self.predicted.quantum_datasets.values(),
264 )
265 ):
266 if not predicted_quantum.task_label:
267 # Skip the 'packages' producer quantum.
268 continue
269 for predicted_input in itertools.chain.from_iterable(predicted_quantum.inputs.values()):
270 self.xgraph.add_edge(predicted_input.dataset_id, predicted_quantum.quantum_id)
271 if predicted_input.dataset_id not in self.output_dataset_ids:
272 self.overall_inputs.setdefault(predicted_input.dataset_id, predicted_input)
273 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
274 self.xgraph.add_edge(predicted_quantum.quantum_id, predicted_output.dataset_id)
276 @staticmethod
277 def run(predicted_path: str, comms: WriterCommunicator) -> None:
278 """Run the writer.
280 Parameters
281 ----------
282 predicted_path : `str`
283 Path to the predicted quantum graph.
284 comms : `WriterCommunicator`
285 Communicator for the writer.
287 Notes
288 -----
289 This method is designed to run as the ``target`` in
290 `WorkerContext.make_worker`.
291 """
292 with comms:
293 writer = Writer(predicted_path, comms)
294 writer.loop()
296 def loop(self) -> None:
297 """Run the main loop for the writer."""
298 data_writers: _DataWriters | None = None
299 if not self.comms.config.zstd_dict_size:
300 data_writers = self.make_data_writers()
301 self.comms.log.info("Polling for write requests from scanners.")
302 for request in self.comms.poll():
303 if data_writers is None:
304 self.pending_compression_training.append(request)
305 if len(self.pending_compression_training) >= self.comms.config.zstd_dict_n_inputs:
306 data_writers = self.make_data_writers()
307 else:
308 self.process_request(request, data_writers)
309 if data_writers is None:
310 data_writers = self.make_data_writers()
311 self.write_init_outputs(data_writers)
313 def make_data_writers(self) -> _DataWriters:
314 """Make a compression dictionary, open the low-level writers, and
315 write any accumulated scans that were needed to make the compression
316 dictionary.
318 Returns
319 -------
320 data_writers : `_DataWriters`
321 Low-level writers struct.
322 """
323 cdict = self.make_compression_dictionary()
324 self.comms.send_compression_dict(cdict.as_bytes())
325 assert self.comms.config.output_path is not None
326 self.comms.log.info("Opening output files.")
327 data_writers = _DataWriters(
328 self.comms,
329 self.predicted,
330 self.indices,
331 compressor=zstandard.ZstdCompressor(self.comms.config.zstd_level, cdict),
332 cdict_data=cdict.as_bytes(),
333 )
334 self.comms.check_for_cancel()
335 self.comms.log.info("Compressing and writing queued scan requests.")
336 for request in self.pending_compression_training:
337 self.process_request(request, data_writers)
338 del self.pending_compression_training
339 self.comms.check_for_cancel()
340 self.write_overall_inputs(data_writers)
341 self.write_packages(data_writers)
342 self.comms.log.info("Returning to write request loop.")
343 return data_writers
345 def make_compression_dictionary(self) -> zstandard.ZstdCompressionDict:
346 """Make the compression dictionary.
348 Returns
349 -------
350 cdict : `zstandard.ZstdCompressionDict`
351 The compression dictionary.
352 """
353 if (
354 not self.comms.config.zstd_dict_size
355 or len(self.pending_compression_training) < self.comms.config.zstd_dict_n_inputs
356 ):
357 self.comms.log.info("Making compressor with no dictionary.")
358 return zstandard.ZstdCompressionDict(b"")
359 self.comms.log.info("Training compression dictionary.")
360 training_inputs: list[bytes] = []
361 # We start the dictionary training with *predicted* quantum dataset
362 # models, since those have almost all of the same attributes as the
363 # provenance quantum and dataset models, and we can get a nice random
364 # sample from just the first N, since they're ordered by UUID. We
365 # chop out the datastore records since those don't appear in the
366 # provenance graph.
367 for predicted_quantum in self.predicted.quantum_datasets.values():
368 if len(training_inputs) == self.comms.config.zstd_dict_n_inputs:
369 break
370 predicted_quantum.datastore_records.clear()
371 training_inputs.append(predicted_quantum.model_dump_json().encode())
372 # Add the provenance quanta, metadata, and logs we've accumulated.
373 for write_request in self.pending_compression_training:
374 assert not write_request.is_compressed, "We can't compress without the compression dictionary."
375 training_inputs.append(write_request.quantum)
376 training_inputs.append(write_request.metadata)
377 training_inputs.append(write_request.logs)
378 return zstandard.train_dictionary(self.comms.config.zstd_dict_size, training_inputs)
380 def write_init_outputs(self, data_writers: _DataWriters) -> None:
381 """Write provenance for init-output datasets and init-quanta.
383 Parameters
384 ----------
385 data_writers : `_DataWriters`
386 Low-level writers struct.
387 """
388 self.comms.log.info("Writing init outputs.")
389 init_quanta = ProvenanceInitQuantaModel()
390 for predicted_init_quantum in self.predicted.init_quanta.root:
391 if not predicted_init_quantum.task_label:
392 # Skip the 'packages' producer quantum.
393 continue
394 existing_outputs = self.existing_init_outputs[predicted_init_quantum.quantum_id]
395 for predicted_output in itertools.chain.from_iterable(predicted_init_quantum.outputs.values()):
396 provenance_output = ProvenanceDatasetModel.from_predicted(
397 predicted_output,
398 producer=predicted_init_quantum.quantum_id,
399 consumers=self.xgraph.successors(predicted_output.dataset_id),
400 )
401 provenance_output.produced = predicted_output.dataset_id in existing_outputs
402 data_writers.datasets.write_model(
403 provenance_output.dataset_id, provenance_output, data_writers.compressor
404 )
405 init_quanta.root.append(ProvenanceInitQuantumModel.from_predicted(predicted_init_quantum))
406 data_writers.graph.write_single_model("init_quanta", init_quanta)
408 def write_overall_inputs(self, data_writers: _DataWriters) -> None:
409 """Write provenance for overall-input datasets.
411 Parameters
412 ----------
413 data_writers : `_DataWriters`
414 Low-level writers struct.
415 """
416 self.comms.log.info("Writing overall inputs.")
417 for predicted_input in self.comms.periodically_check_for_cancel(self.overall_inputs.values()):
418 if predicted_input.dataset_id not in data_writers.datasets.addresses:
419 data_writers.datasets.write_model(
420 predicted_input.dataset_id,
421 ProvenanceDatasetModel.from_predicted(
422 predicted_input,
423 producer=None,
424 consumers=self.xgraph.successors(predicted_input.dataset_id),
425 ),
426 data_writers.compressor,
427 )
428 del self.overall_inputs
430 @staticmethod
431 def write_packages(data_writers: _DataWriters) -> None:
432 """Write package version information to the provenance graph.
434 Parameters
435 ----------
436 data_writers : `_DataWriters`
437 Low-level writers struct.
438 """
439 packages = Packages.fromSystem(include_all=True)
440 data = packages.toBytes("json")
441 data_writers.graph.write_single_block("packages", data)
443 def process_request(self, request: WriteRequest, data_writers: _DataWriters) -> None:
444 """Process a `WriteRequest` into `_ScanData`.
446 Parameters
447 ----------
448 request : `WriteRequest`
449 Result of a quantum scan.
450 data_writers : `_DataWriters`
451 Low-level writers struct.
452 """
453 if (existing_init_outputs := self.existing_init_outputs.get(request.quantum_id)) is not None:
454 self.comms.log.debug("Handling init-output scan for %s.", request.quantum_id)
455 existing_init_outputs.update(request.existing_outputs)
456 self.comms.report_write()
457 return
458 self.comms.log.debug("Handling quantum scan for %s.", request.quantum_id)
459 predicted_quantum = self.predicted.quantum_datasets[request.quantum_id]
460 outputs: dict[uuid.UUID, bytes] = {}
461 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
462 provenance_output = ProvenanceDatasetModel.from_predicted(
463 predicted_output,
464 producer=predicted_quantum.quantum_id,
465 consumers=self.xgraph.successors(predicted_output.dataset_id),
466 )
467 provenance_output.produced = provenance_output.dataset_id in request.existing_outputs
468 outputs[provenance_output.dataset_id] = data_writers.compressor.compress(
469 provenance_output.model_dump_json().encode()
470 )
471 if not request.quantum:
472 request.quantum = (
473 ProvenanceQuantumModel.from_predicted(predicted_quantum).model_dump_json().encode()
474 )
475 if request.is_compressed:
476 request.quantum = data_writers.compressor.compress(request.quantum)
477 if not request.is_compressed:
478 request.quantum = data_writers.compressor.compress(request.quantum)
479 if request.metadata:
480 request.metadata = data_writers.compressor.compress(request.metadata)
481 if request.logs:
482 request.logs = data_writers.compressor.compress(request.logs)
483 self.comms.log.debug("Writing quantum %s.", request.quantum_id)
484 data_writers.quanta.write_bytes(request.quantum_id, request.quantum)
485 for dataset_id, dataset_data in outputs.items():
486 data_writers.datasets.write_bytes(dataset_id, dataset_data)
487 if request.metadata:
488 (metadata_output,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
489 address = data_writers.metadata.write_bytes(request.quantum_id, request.metadata)
490 data_writers.metadata.addresses[metadata_output.dataset_id] = address
491 if request.logs:
492 (log_output,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
493 address = data_writers.logs.write_bytes(request.quantum_id, request.logs)
494 data_writers.logs.addresses[log_output.dataset_id] = address
495 # We shouldn't need this predicted quantum anymore; delete it in the
496 # hopes that'll free up some memory.
497 del self.predicted.quantum_datasets[request.quantum_id]
498 self.comms.report_write()
501_T = TypeVar("_T")