Coverage for python / lsst / pipe / base / cli / cmd / commands.py: 83%
72 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-09 01:22 -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 typing import Any
30import click
32from lsst.daf.butler.cli.opt import (
33 dataset_type_option,
34 options_file_option,
35 register_dataset_types_option,
36 repo_argument,
37 transfer_dimensions_option,
38 transfer_option,
39)
40from lsst.daf.butler.cli.utils import ButlerCommand, split_commas, unwrap
42from ... import script
43from ...quantum_graph import aggregator
44from ..opt import instrument_argument, update_output_chain_option
47@click.command(short_help="Add an instrument definition to the repository", cls=ButlerCommand)
48@repo_argument(required=True)
49@instrument_argument(required=True, nargs=-1, help="The fully-qualified name of an Instrument subclass.")
50@click.option("--update", is_flag=True)
51def register_instrument(*args: Any, **kwargs: Any) -> None:
52 """Add an instrument to the data repository."""
53 script.register_instrument(*args, **kwargs)
56@click.command(short_help="Transfer datasets from a graph to a butler.", cls=ButlerCommand)
57@click.argument("graph", required=True)
58@click.argument("dest", required=True)
59@register_dataset_types_option()
60@transfer_dimensions_option(default=False)
61@update_output_chain_option()
62@click.option(
63 "--dry-run", is_flag=True, default=False, help="Run the transfer but do not update the destination butler"
64)
65@dataset_type_option(help="Subset of dataset types to transfer from graph.")
66@options_file_option()
67def transfer_from_graph(**kwargs: Any) -> None:
68 """Transfer datasets from a quantum graph to a destination butler.
70 GRAPH is a URI to the source quantum graph file.
72 DEST is a URI to the Butler repository that will receive copies of the
73 datasets.
74 """
75 number = script.transfer_from_graph(**kwargs)
76 print(f"Number of datasets transferred: {number}")
79@click.command(short_help="Make Zip archive from output files using graph.", cls=ButlerCommand)
80@click.argument("graph", required=True)
81@repo_argument(
82 required=True,
83 help="REPO is a URI to a butler configuration that is used to configure "
84 "the datastore of the quantum-backed butler.",
85)
86@click.argument("dest", required=True)
87@dataset_type_option(help="Dataset types to include in Zip archive.")
88@options_file_option()
89def zip_from_graph(**kwargs: Any) -> None:
90 """Transfer datasets from a quantum graph to a Zip archive.
92 GRAPH is a URI to the source quantum graph file to use when building the
93 Zip archive.
95 DEST is a directory to write the Zip archive.
96 """
97 zip = script.zip_from_graph(**kwargs)
98 print(f"Zip archive written to {zip}")
101@click.command(short_help="Retrieve artifacts from subset of graph.", cls=ButlerCommand)
102@click.argument("graph", required=True)
103@repo_argument(
104 required=True,
105 help="REPO is a URI to a butler configuration that is used to configure "
106 "the datastore of the quantum-backed butler.",
107)
108@click.argument("dest", required=True)
109@transfer_option()
110@click.option(
111 "--preserve-path/--no-preserve-path",
112 is_flag=True,
113 default=True,
114 help="Preserve the datastore path to the artifact at the destination.",
115)
116@click.option(
117 "--clobber/--no-clobber",
118 is_flag=True,
119 default=False,
120 help="If clobber, overwrite files if they exist locally.",
121)
122@click.option(
123 "--qgraph-node-id",
124 callback=split_commas,
125 multiple=True,
126 help=unwrap(
127 """Only load a specified set of nodes when graph is
128 loaded from a file, nodes are identified by UUID
129 values. One or more comma-separated strings are
130 accepted. By default all nodes are loaded. Ignored if
131 graph is not loaded from a file."""
132 ),
133)
134@click.option(
135 "--include-inputs/--no-include-inputs",
136 is_flag=True,
137 default=True,
138 help="Whether to include input datasets in retrieval.",
139)
140@click.option(
141 "--include-outputs/--no-include-outputs",
142 is_flag=True,
143 default=True,
144 help="Whether to include output datasets in retrieval.",
145)
146@options_file_option()
147def retrieve_artifacts_for_quanta(**kwargs: Any) -> None:
148 """Retrieve artifacts from given quanta defined in quantum graph.
150 GRAPH is a URI to the source quantum graph file to use when building the
151 Zip archive.
153 DEST is a directory to write the Zip archive.
154 """
155 artifacts = script.retrieve_artifacts_for_quanta(**kwargs)
156 print(f"Written {len(artifacts)} artifacts to {kwargs['dest']}.")
159_AGGREGATOR_DEFAULTS = aggregator.AggregatorConfig()
162@click.command(short_help="Scan for the outputs of an active or completed quantum graph.", cls=ButlerCommand)
163@click.argument("predicted_graph", required=True)
164@repo_argument(required=True, help="Path to the central butler repository.")
165@click.option(
166 "-o",
167 "--output",
168 "output_path",
169 default=_AGGREGATOR_DEFAULTS.output_path,
170 help=(
171 "Path to the output provenance quantum graph. THIS OPTION IS FOR "
172 "DEVELOPMENT AND DEBUGGING ONLY. IT MAY BE REMOVED IN THE FUTURE."
173 ),
174)
175@click.option(
176 "--processes",
177 "-j",
178 "n_processes",
179 default=_AGGREGATOR_DEFAULTS.n_processes,
180 type=click.IntRange(min=1),
181 help="Number of processes to use.",
182)
183@click.option(
184 "--complete/--incomplete",
185 "assume_complete",
186 default=_AGGREGATOR_DEFAULTS.assume_complete,
187 help="Whether execution has completed (and failures cannot be retried).",
188)
189@click.option(
190 "--dry-run",
191 is_flag=True,
192 default=_AGGREGATOR_DEFAULTS.dry_run,
193 help="Do not actually perform any central database ingests.",
194)
195@click.option(
196 "--interactive-status/--no-interactive-status",
197 "interactive_status",
198 default=_AGGREGATOR_DEFAULTS.interactive_status,
199 help="Use progress bars for status reporting instead of periodic logging.",
200)
201@click.option(
202 "--log-status-interval",
203 type=int,
204 default=_AGGREGATOR_DEFAULTS.log_status_interval,
205 help="Interval (in seconds) between periodic logger status updates.",
206)
207@click.option(
208 "--register-dataset-types/--no-register-dataset-types",
209 default=_AGGREGATOR_DEFAULTS.register_dataset_types,
210 help="Register output dataset types.",
211)
212@click.option(
213 "--update-output-chain/--no-update-output-chain",
214 default=_AGGREGATOR_DEFAULTS.update_output_chain,
215 help="Prepend the output RUN collection to the output CHAINED collection.",
216)
217@click.option(
218 "--worker-log-dir",
219 type=str,
220 default=_AGGREGATOR_DEFAULTS.worker_log_dir,
221 help="Path to a directory (POSIX only) for parallel worker logs.",
222)
223@click.option(
224 "--worker-log-level",
225 type=str,
226 default=_AGGREGATOR_DEFAULTS.worker_log_level,
227 help="Log level for worker processes/threads (use DEBUG for per-quantum messages).",
228)
229@click.option(
230 "--zstd-level",
231 type=int,
232 default=_AGGREGATOR_DEFAULTS.zstd_level,
233 help="Compression level for the provenance quantum graph file.",
234)
235@click.option(
236 "--zstd-dict-size",
237 type=int,
238 default=_AGGREGATOR_DEFAULTS.zstd_dict_size,
239 help="Size (in bytes) of the ZStandard compression dictionary.",
240)
241@click.option(
242 "--zstd-dict-n-inputs",
243 type=int,
244 default=_AGGREGATOR_DEFAULTS.zstd_dict_n_inputs,
245 help=("Number of samples of each type to include in ZStandard compression dictionary training."),
246)
247@click.option(
248 "--mock-storage-classes/--no-mock-storage-classes",
249 default=_AGGREGATOR_DEFAULTS.mock_storage_classes,
250 help="Enable support for storage classes created by the lsst.pipe.base.tests.mocks package.",
251)
252def aggregate_graph(predicted_graph: str, repo: str, **kwargs: Any) -> None:
253 """Scan for quantum graph's outputs to gather provenance, ingest datasets
254 into the central butler repository, and delete datasets that are no
255 longer needed.
256 """
257 # It'd be nice to allow to the user to provide a path to an
258 # AggregatorConfig JSON file for options that weren't provided, but Click
259 # 8.1 fundamentally cannot handle flag options that default to None rather
260 # than True or False (i.e. so they fall back to the config value when not
261 # set). It's not clear whether Click 8.2.x has actually fixed this; Click
262 # 8.2.0 tried but caused new problems.
264 config = aggregator.AggregatorConfig(**kwargs)
265 try:
266 aggregator.aggregate_graph(predicted_graph, repo, config)
267 except aggregator.FatalWorkerError as err:
268 # When this exception is raised, we'll have already logged the relevant
269 # traceback from a separate worker.
270 raise click.ClickException(str(err)) from None