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