Coverage for python/lsst/ctrl/mpexec/cli/script/qgraph.py: 79%
92 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 02:05 -0700
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/>.
28from __future__ import annotations
30__all__ = ("qgraph",)
32import uuid
33from collections.abc import Iterable, Mapping, Sequence
34from typing import TYPE_CHECKING
36import yaml
37from astropy.table import Table
39from lsst.pipe.base import BuildId, QuantumGraph
40from lsst.pipe.base.all_dimensions_quantum_graph_builder import (
41 AllDimensionsQuantumGraphBuilder,
42 DatasetQueryConstraintVariant,
43)
44from lsst.pipe.base.dot_tools import graph2dot
45from lsst.pipe.base.mermaid_tools import graph2mermaid
46from lsst.pipe.base.pipeline_graph import TaskImportMode
47from lsst.pipe.base.quantum_graph import PredictedQuantumGraph, PredictedQuantumGraphComponents
48from lsst.resources import ResourcePath, ResourcePathExpression
49from lsst.utils.iteration import ensure_iterable
50from lsst.utils.logging import getLogger
52from ..._pipeline_graph_factory import PipelineGraphFactory
53from ...showInfo import ShowInfo
54from ..butler_factory import ButlerFactory
55from ..utils import summarize_quantum_graph
57if TYPE_CHECKING:
58 from lsst.pipe.base.tests.mocks import ForcedFailure # this monkey patches; only import for annotation!
60_LOG = getLogger(__name__)
63def qgraph(
64 pipeline_graph_factory: PipelineGraphFactory | None,
65 *,
66 qgraph: ResourcePathExpression | None,
67 qgraph_id: str | None,
68 qgraph_node_id: Iterable[uuid.UUID | str] | None,
69 qgraph_datastore_records: bool,
70 skip_existing_in: Iterable[str] | None,
71 skip_existing: bool,
72 retained_dataset_types: str | None,
73 prune_unanchored_quanta: tuple[str, str] | None = None,
74 save_qgraph: ResourcePathExpression | None,
75 qgraph_dot: str | None,
76 qgraph_mermaid: str | None,
77 butler_config: ResourcePathExpression,
78 input: Iterable[str] | str,
79 output: str | None,
80 output_run: str | None,
81 extend_run: bool,
82 replace_run: bool,
83 prune_replaced: str | None,
84 data_query: str | None,
85 data_id_table: Iterable[ResourcePathExpression],
86 show: ShowInfo,
87 clobber_outputs: bool,
88 dataset_query_constraint: str,
89 rebase: bool,
90 mock: bool = False,
91 unmocked_dataset_types: Sequence[str],
92 mock_failure: Mapping[str, ForcedFailure],
93 for_execution: bool = False,
94 for_init_output_run: bool = False,
95 **kwargs: object,
96) -> PredictedQuantumGraph | None:
97 """Implement the command line interface `pipetask qgraph` subcommand.
99 Should only be called by command line tools and unit test code that test
100 this function.
102 Parameters
103 ----------
104 pipeline_graph_factory : `..PipelineGraphFactory` or `None`
105 A factory that holds the pipeline and can produce a pipeline graph.
106 If this is not `None` then ``qgraph`` should be `None`.
107 qgraph : convertible to `lsst.resources.ResourcePath`, or `None`
108 URI location for a serialized quantum graph definition. If this option
109 is not `None` then ``pipeline_graph_factory`` should be `None`.
110 qgraph_id : `str` or `None`
111 Quantum graph identifier, if specified must match the identifier of the
112 graph loaded from a file. Ignored if graph is not loaded from a file.
113 qgraph_node_id : `~collections.abc.Iterable` [`str` | `uuid.UUID`] or \
114 `None`
115 Only load a specified set of nodes if graph is loaded from a file,
116 nodes are identified by integer IDs.
117 qgraph_datastore_records : `bool`
118 If `True` then include datastore records into generated quanta.
119 skip_existing_in : `~collections.abc.Iterable` [ `str` ] or `None`
120 Accepts list of collections, if all Quantum outputs already exist in
121 the specified list of collections then that Quantum will be excluded
122 from the QuantumGraph.
123 skip_existing : `bool`
124 Appends output RUN collection to the ``skip_existing_in`` list.
125 retained_dataset_types : `str` or `None`
126 Path to a YAML file listing dataset type names or glob-style wildcard
127 patterns that should exist in ``skip_existing_in`` when the producing
128 task ran successfully. When a quantum should run, the builder
129 propagates the must-run signal backward through non-retained input
130 datasets, forcing the upstream quanta that need to regenerate those
131 intermediates to also run. Has no effect without ``skip_existing_in``.
132 prune_unanchored_quanta : `tuple` [ `str`, `str` ] or `None`, optional
133 A ``(source_label, anchor_label)`` pair of task labels. If not `None`,
134 source quanta with no reachable anchor quantum downstream are removed,
135 along with their entire downstream chain.
136 save_qgraph : convertible to `lsst.resources.ResourcePath` or `None`
137 URI location for saving the quantum graph.
138 qgraph_dot : `str` or `None`
139 Path location for storing GraphViz DOT representation of a quantum
140 graph.
141 qgraph_mermaid : `str` or `None`
142 Path location for storing Mermaid representation of a quantum graph.
143 butler_config : convertible to `lsst.resources.ResourcePath`
144 Path to butler repository configuration.
145 input : `~collections.abc.Iterable` [ `str` ] or `None`
146 List of names of the input collection(s).
147 output : `str` or `None`
148 Name of the output CHAINED collection. This may either be an existing
149 CHAINED collection to use as both input and output (if `input` is
150 `None`), or a new CHAINED collection created to include all inputs
151 (if `input` is not `None`). In both cases, the collection's children
152 will start with an output RUN collection that directly holds all new
153 datasets (see `output_run`).
154 output_run : `str` or `None`
155 Name of the new output RUN collection. If not provided then `output`
156 must be provided and a new RUN collection will be created by appending
157 a timestamp to the value passed with `output`. If this collection
158 already exists then `extend_run` must be passed.
159 extend_run : `bool`
160 Instead of creating a new RUN collection, insert datasets into either
161 the one given by `output_run` (if provided) or the first child
162 collection of `output` (which must be of type RUN).
163 replace_run : `bool`
164 Before creating a new RUN collection in an existing CHAINED collection,
165 remove the first child collection (which must be of type RUN). This can
166 be used to repeatedly write to the same (parent) collection during
167 development, but it does not delete the datasets associated with the
168 replaced run unless `prune-replaced` is also True. Requires `output`,
169 and `extend_run` must be `None`.
170 prune_replaced : `str` or `None`
171 If not `None`, delete the datasets in the collection replaced by
172 `replace_run`, either just from the datastore ("unstore") or by
173 removing them and the RUN completely ("purge"). Requires
174 ``replace_run`` to be `True`.
175 data_query : `str`
176 User query selection expression.
177 data_id_table : `~collections.abc.Iterable` [convertible to \
178 `lsst.resources.ResourcePath`]
179 Paths to data ID tables to join in.
180 show : `lsst.ctrl.mpexec.showInfo.ShowInfo`
181 Descriptions of what to dump to stdout.
182 clobber_outputs : `bool`
183 Remove outputs from previous execution of the same quantum before new
184 execution. If ``skip_existing`` is also passed, then only failed
185 quanta will be clobbered.
186 dataset_query_constraint : `str`
187 Control constraining graph building using pre-existing dataset types.
188 Valid values are off, all, or a comma separated list of dataset type
189 names.
190 rebase : `bool`
191 If `True` then reset output collection chain if it is inconsistent with
192 the ``inputs``.
193 mock : `bool`
194 If True, use a mocked version of the pipeline.
195 unmocked_dataset_types : `collections.abc.Sequence` [ `str` ], optional
196 List of overall-input dataset types that should not be mocked.
197 mock_failure : `~collections.abc.Mapping`
198 Quanta that should raise exceptions.
199 for_execution : `bool`, optional
200 If `True`, the script is being used to feed another that will execute
201 the given quanta, and hence all information needed for execution must
202 be loaded.
203 for_init_output_run : `bool`, optional
204 If `True`, the script is being used to feed another that will
205 initialize the output run, and hence all information needed to do so
206 must be loaded.
207 **kwargs : `dict` [`str`, `str`]
208 Ignored; click commands may accept options for more than one script
209 function and pass all the option kwargs to each of the script functions
210 which ignore these unused kwargs.
212 Returns
213 -------
214 qg : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph`
215 The quantum graph object that was created or loaded.
216 """
217 # make sure that --extend-run always enables --skip-existing
218 if extend_run:
219 skip_existing = True
221 skip_existing_in = tuple(skip_existing_in) if skip_existing_in is not None else ()
222 retained_dataset_type_patterns: list[str] | None = None
223 if retained_dataset_types is not None:
224 with open(retained_dataset_types) as f:
225 retained_dataset_type_patterns = yaml.safe_load(f)
226 if not isinstance(retained_dataset_type_patterns, list) or not retained_dataset_type_patterns:
227 raise ValueError(
228 f"--retained-dataset-types file {retained_dataset_types!r} must contain "
229 "a non-empty YAML sequence of strings."
230 )
231 if data_query is None: 231 ↛ 233line 231 didn't jump to line 233 because the condition on line 231 was always true
232 data_query = ""
233 inputs = list(ensure_iterable(input)) if input else []
234 del input
236 butler, collections, run = ButlerFactory.make_butler_and_collections(
237 butler_config,
238 output=output,
239 output_run=output_run,
240 inputs=inputs,
241 extend_run=extend_run,
242 rebase=rebase,
243 replace_run=replace_run,
244 prune_replaced=prune_replaced,
245 )
246 with butler:
247 if skip_existing and run:
248 skip_existing_in += (run,)
250 qgc: PredictedQuantumGraphComponents
251 if qgraph is not None:
252 # click passes empty tuple as default value for qgraph_node_id
253 quantum_ids = (
254 {uuid.UUID(q) if not isinstance(q, uuid.UUID) else q for q in qgraph_node_id}
255 if qgraph_node_id
256 else None
257 )
258 qgraph = ResourcePath(qgraph)
259 match qgraph.getExtension():
260 case ".qgraph": 260 ↛ 261line 260 didn't jump to line 261 because the pattern on line 260 never matched
261 qgc = PredictedQuantumGraphComponents.from_old_quantum_graph(
262 QuantumGraph.loadUri(
263 qgraph,
264 butler.dimensions,
265 nodes=quantum_ids,
266 graphID=BuildId(qgraph_id) if qgraph_id is not None else None,
267 )
268 )
269 case ".qg": 269 ↛ 291line 269 didn't jump to line 291 because the pattern on line 269 always matched
270 if qgraph_id is not None: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 _LOG.warning("--qgraph-id is ignored when loading new '.qg' files.")
272 if for_execution or for_init_output_run or save_qgraph or show.needs_full_qg:
273 import_mode = TaskImportMode.ASSUME_CONSISTENT_EDGES
274 else:
275 import_mode = TaskImportMode.DO_NOT_IMPORT
276 with PredictedQuantumGraph.open(qgraph, import_mode=import_mode) as reader:
277 if (
278 for_execution
279 or qgraph_dot
280 or qgraph_mermaid
281 or show.needs_full_qg
282 or qgraph_node_id
283 ):
284 # This reads everything for the given quanta.
285 reader.read_execution_quanta(quantum_ids)
286 elif for_init_output_run: 286 ↛ 287line 286 didn't jump to line 287 because the condition on line 286 was never true
287 reader.read_init_quanta()
288 else:
289 reader.read_thin_graph()
290 qgc = reader.components
291 case ext:
292 raise ValueError(f"Unrecognized extension for quantum graph: {ext!r}")
294 # pipeline can not be provided in this case
295 if pipeline_graph_factory: 295 ↛ 296line 295 didn't jump to line 296 because the condition on line 295 was never true
296 raise ValueError(
297 "Pipeline must not be given when quantum graph is read from "
298 f"file: {bool(pipeline_graph_factory)}"
299 )
300 else:
301 if pipeline_graph_factory is None: 301 ↛ 302line 301 didn't jump to line 302 because the condition on line 301 was never true
302 raise ValueError("Pipeline must be given when quantum graph is not read from file.")
303 # We can't resolve the pipeline graph if we're mocking until after
304 # we've done the mocking (and the QG build will resolve on its own
305 # anyway).
306 pipeline_graph = pipeline_graph_factory(resolve=False)
307 if mock: 307 ↛ 308line 307 didn't jump to line 308 because the condition on line 307 was never true
308 from lsst.pipe.base.tests.mocks import mock_pipeline_graph
310 pipeline_graph = mock_pipeline_graph(
311 pipeline_graph,
312 unmocked_dataset_types=unmocked_dataset_types,
313 force_failures=mock_failure,
314 )
315 data_id_tables = []
316 for table_file in data_id_table: 316 ↛ 317line 316 didn't jump to line 317 because the loop on line 316 never started
317 with ResourcePath(table_file).as_local() as local_path:
318 table = Table.read(local_path.ospath)
319 # Add the filename to the metadata for more logging
320 # information down in the QG builder.
321 table.meta["filename"] = table_file
322 data_id_tables.append(table)
323 # make execution plan (a.k.a. DAG) for pipeline
324 graph_builder = AllDimensionsQuantumGraphBuilder(
325 pipeline_graph,
326 butler,
327 where=data_query,
328 skip_existing_in=skip_existing_in,
329 retained_dataset_types=retained_dataset_type_patterns,
330 prune_unanchored_quanta=prune_unanchored_quanta,
331 clobber=clobber_outputs,
332 dataset_query_constraint=DatasetQueryConstraintVariant.fromExpression(
333 dataset_query_constraint
334 ),
335 input_collections=collections,
336 output_run=run,
337 data_id_tables=data_id_tables,
338 )
339 # Accumulate metadata (QB builder adds some of its own).
340 metadata = {
341 "butler_argument": str(butler_config),
342 "extend_run": extend_run,
343 "skip_existing_in": skip_existing_in,
344 "skip_existing": skip_existing,
345 "retained_dataset_types": retained_dataset_types,
346 "data_query": data_query,
347 }
348 assert run is not None, "Butler output run collection must be defined"
349 qgc = graph_builder.finish(
350 output, metadata=metadata, attach_datastore_records=qgraph_datastore_records
351 )
353 if save_qgraph:
354 _LOG.verbose("Writing quantum graph to %r.", save_qgraph)
355 qgc.write(save_qgraph)
357 qg = qgc.assemble()
359 if not summarize_quantum_graph(qg):
360 return None
362 if qgraph_dot: 362 ↛ 363line 362 didn't jump to line 363 because the condition on line 362 was never true
363 _LOG.verbose("Writing quantum graph DOT visualization to %r.", qgraph_dot)
364 graph2dot(qg, qgraph_dot)
366 if qgraph_mermaid: 366 ↛ 367line 366 didn't jump to line 367 because the condition on line 366 was never true
367 _LOG.verbose("Writing quantum graph Mermaid visualization to %r.", qgraph_mermaid)
368 graph2mermaid(qg, qgraph_mermaid)
370 # optionally dump some info.
371 show.show_graph_info(qg, butler_config)
373 return qg