Coverage for python / lsst / pipe / base / separable_pipeline_executor.py: 36%
64 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/>.
29from __future__ import annotations
31__all__ = [
32 "SeparablePipelineExecutor",
33]
36import datetime
37import getpass
38import logging
39from collections.abc import Iterable
40from typing import Any
42import lsst.resources
43from lsst.daf.butler import Butler
45from ._quantumContext import ExecutionResources
46from .all_dimensions_quantum_graph_builder import AllDimensionsQuantumGraphBuilder
47from .graph import QuantumGraph
48from .mp_graph_executor import MPGraphExecutor
49from .pipeline import Pipeline
50from .quantum_graph import PredictedQuantumGraph
51from .quantum_graph_builder import QuantumGraphBuilder
52from .quantum_graph_executor import QuantumGraphExecutor
53from .single_quantum_executor import SingleQuantumExecutor
54from .taskFactory import TaskFactory
56_LOG = logging.getLogger(__name__)
59class SeparablePipelineExecutor:
60 """An executor that allows each step of pipeline execution to be
61 run independently.
63 The executor can run any or all of the following steps:
65 * pre-execution initialization
66 * pipeline building
67 * quantum graph generation
68 * quantum graph execution
70 Any of these steps can also be handed off to external code without
71 compromising the remaining ones.
73 Parameters
74 ----------
75 butler : `lsst.daf.butler.Butler`
76 A Butler whose ``collections`` and ``run`` attributes contain the input
77 and output collections to use for processing.
78 clobber_output : `bool`, optional
79 If set, the pipeline execution overwrites existing output files.
80 Otherwise, any conflict between existing and new outputs is an error.
81 skip_existing_in : iterable [`str`], optional
82 If not empty, the pipeline execution searches the listed collections
83 for existing outputs, and skips any quanta that have run to completion
84 (or have no work to do). Otherwise, all tasks are attempted (subject to
85 ``clobber_output``).
86 task_factory : `.TaskFactory`, optional
87 A custom task factory for use in pre-execution and execution. By
88 default, a new instance of `.TaskFactory` is used.
89 resources : `.ExecutionResources`
90 The resources available to each quantum being executed.
91 raise_on_partial_outputs : `bool`, optional
92 If `True` raise exceptions chained by
93 `.AnnotatedPartialOutputsError` immediately, instead of
94 considering the partial result a success and continuing to run
95 downstream tasks.
96 """
98 def __init__(
99 self,
100 butler: Butler,
101 clobber_output: bool = False,
102 skip_existing_in: Iterable[str] | None = None,
103 task_factory: TaskFactory | None = None,
104 resources: ExecutionResources | None = None,
105 raise_on_partial_outputs: bool = True,
106 ):
107 self._butler = Butler.from_config(
108 butler=butler, collections=butler.collections.defaults, run=butler.run
109 )
110 if not self._butler.collections.defaults:
111 raise ValueError("Butler must specify input collections for pipeline.")
112 if not self._butler.run:
113 raise ValueError("Butler must specify output run for pipeline.")
115 self._clobber_output = clobber_output
116 self._skip_existing_in = list(skip_existing_in) if skip_existing_in else []
118 self._task_factory = task_factory if task_factory else TaskFactory()
119 self.resources = resources
120 self.raise_on_partial_outputs = raise_on_partial_outputs
122 def pre_execute_qgraph(
123 self,
124 graph: QuantumGraph | PredictedQuantumGraph,
125 register_dataset_types: bool = False,
126 save_init_outputs: bool = True,
127 save_versions: bool = True,
128 ) -> None:
129 """Run pre-execution initialization.
131 This method will be deprecated after DM-38041, to be replaced with a
132 method that takes either a `.Pipeline` or a
133 resolved `.pipeline_graph.PipelineGraph` instead of a `.QuantumGraph`.
135 Parameters
136 ----------
137 graph : `.QuantumGraph` or `.quantum_graph.PredictedQuantumGraph`
138 The quantum graph defining the pipeline and datasets to
139 be initialized.
140 register_dataset_types : `bool`, optional
141 If `True`, register all output dataset types from the pipeline
142 represented by ``graph``.
143 save_init_outputs : `bool`, optional
144 If `True`, create init-output datasets in this object's output run.
145 save_versions : `bool`, optional
146 If `True`, save a package versions dataset.
147 """
148 if register_dataset_types:
149 graph.pipeline_graph.register_dataset_types(self._butler, include_packages=save_versions)
150 if save_init_outputs:
151 graph.write_init_outputs(self._butler, skip_existing=(self._butler.run in self._skip_existing_in))
152 graph.write_configs(self._butler)
153 if save_versions:
154 graph.write_packages(self._butler)
156 def make_pipeline(self, pipeline_uri: str | lsst.resources.ResourcePath) -> Pipeline:
157 """Build a pipeline from pipeline and configuration information.
159 Parameters
160 ----------
161 pipeline_uri : `str` or `lsst.resources.ResourcePath`
162 URI to a file containing a pipeline definition. A URI fragment may
163 be used to specify a subset of the pipeline, as described in
164 :ref:`pipeline-running-intro`.
166 Returns
167 -------
168 pipeline : `.Pipeline`
169 The fully-built pipeline.
170 """
171 return Pipeline.from_uri(pipeline_uri)
173 def make_quantum_graph_builder(
174 self,
175 pipeline: Pipeline,
176 where: str = "",
177 *,
178 builder_class: type[QuantumGraphBuilder] = AllDimensionsQuantumGraphBuilder,
179 **kwargs: Any,
180 ) -> QuantumGraphBuilder:
181 """Initialize a quantum graph builder from a pipeline and input
182 datasets.
184 Parameters
185 ----------
186 pipeline : `.Pipeline`
187 The pipeline for which to generate a quantum graph.
188 where : `str`, optional
189 A data ID query that constrains the quanta generated. Must not be
190 provided if a custom ``builder_class`` is given and that class does
191 not accept ``where`` as a construction argument.
192 builder_class : `type` [ \
193 `.quantum_graph_builder.QuantumGraphBuilder` ], optional
194 Quantum graph builder implementation. Ignored if ``builder`` is
195 provided.
196 **kwargs
197 Additional keyword arguments are forwarded to ``builder_class``
198 when a quantum graph builder instance is constructed. All
199 arguments accepted by the
200 `~.quantum_graph_builder.QuantumGraphBuilder` base
201 class are provided automatically (from explicit arguments to this
202 method and executor attributes) and do not need to be included
203 as keyword arguments.
205 Returns
206 -------
207 builder : `.quantum_graph_builder.QuantumGraphBuilder`
208 A quantum graph builder.
209 """
210 if where:
211 # Only pass 'where' if it's actually provided, since some
212 # QuantumGraphBuilder subclasses may not accept it.
213 kwargs["where"] = where
214 return builder_class(
215 pipeline.to_graph(),
216 self._butler,
217 skip_existing_in=self._skip_existing_in,
218 clobber=self._clobber_output,
219 **kwargs,
220 )
222 def make_quantum_graph(
223 self,
224 pipeline: Pipeline,
225 where: str = "",
226 *,
227 builder_class: type[QuantumGraphBuilder] = AllDimensionsQuantumGraphBuilder,
228 attach_datastore_records: bool = False,
229 **kwargs: Any,
230 ) -> QuantumGraph:
231 """Build a quantum graph from a pipeline and input datasets.
233 This returns an instance of the old `.QuantumGraph` class. Use
234 `build_quantum_graph` to construct a
235 `.quantum_graph.PredictedQuantumGraph`.
237 Parameters
238 ----------
239 pipeline : `.Pipeline`
240 The pipeline for which to generate a quantum graph.
241 where : `str`, optional
242 A data ID query that constrains the quanta generated. Must not be
243 provided if a custom ``builder_class`` is given and that class does
244 not accept ``where`` as a construction argument.
245 builder_class : `type` [ \
246 `.quantum_graph_builder.QuantumGraphBuilder` ], optional
247 Quantum graph builder implementation. Ignored if ``builder`` is
248 provided.
249 attach_datastore_records : `bool`, optional
250 Whether to attach datastore records. These are currently used only
251 by `lsst.daf.butler.QuantumBackedButler`, which is not used by
252 `SeparablePipelineExecutor` for execution.
253 **kwargs
254 Additional keyword arguments are forwarded to ``builder_class``
255 when a quantum graph builder instance is constructed. All
256 arguments accepted by the
257 `~.quantum_graph_builder.QuantumGraphBuilder` base
258 class are provided automatically (from explicit arguments to this
259 method and executor attributes) and do not need to be included
260 as keyword arguments.
262 Returns
263 -------
264 graph : `.QuantumGraph`
265 The quantum graph for ``.Pipeline`` as run on the datasets
266 identified by ``where``.
268 Notes
269 -----
270 This method does no special handling of empty quantum graphs. If
271 needed, clients can use `len` to test if the returned graph is empty.
272 """
273 metadata = {
274 "input": self._butler.collections.defaults,
275 "output_run": self._butler.run,
276 "skip_existing_in": self._skip_existing_in,
277 "skip_existing": bool(self._skip_existing_in),
278 "data_query": where,
279 "user": getpass.getuser(),
280 "time": str(datetime.datetime.now()),
281 }
282 qg_builder = self.make_quantum_graph_builder(pipeline, where, builder_class=builder_class, **kwargs)
283 graph = qg_builder.build(metadata=metadata, attach_datastore_records=attach_datastore_records)
284 _LOG.info(
285 "QuantumGraph contains %d quanta for %d tasks, graph ID: %r",
286 len(graph),
287 len(graph.taskGraph),
288 graph.graphID,
289 )
290 return graph
292 def build_quantum_graph(
293 self,
294 pipeline: Pipeline,
295 where: str = "",
296 *,
297 builder_class: type[QuantumGraphBuilder] = AllDimensionsQuantumGraphBuilder,
298 attach_datastore_records: bool = False,
299 **kwargs: Any,
300 ) -> PredictedQuantumGraph:
301 """Build a quantum graph from a pipeline and input datasets.
303 This returns an instance of the new
304 `.quantum_graph.PredictedQuantumGraph` class. Use `make_quantum_graph`
305 to construct a `.QuantumGraph`.
307 Parameters
308 ----------
309 pipeline : `.Pipeline`
310 The pipeline for which to generate a quantum graph.
311 where : `str`, optional
312 A data ID query that constrains the quanta generated. Must not be
313 provided if a custom ``builder_class`` is given and that class does
314 not accept ``where`` as a construction argument.
315 builder_class : `type` [ \
316 `.quantum_graph_builder.QuantumGraphBuilder` ], optional
317 Quantum graph builder implementation. Ignored if ``builder`` is
318 provided.
319 attach_datastore_records : `bool`, optional
320 Whether to attach datastore records. These are currently used only
321 by `lsst.daf.butler.QuantumBackedButler`, which is not used by
322 `SeparablePipelineExecutor` for execution.
323 **kwargs
324 Additional keyword arguments are forwarded to ``builder_class``
325 when a quantum graph builder instance is constructed. All
326 arguments accepted by the
327 `~.quantum_graph_builder.QuantumGraphBuilder` base
328 class are provided automatically (from explicit arguments to this
329 method and executor attributes) and do not need to be included
330 as keyword arguments.
332 Returns
333 -------
334 graph : `.QuantumGraph`
335 The quantum graph for ``.Pipeline`` as run on the datasets
336 identified by ``where``.
338 Notes
339 -----
340 This method does no special handling of empty quantum graphs. If
341 needed, clients can use `len` to test if the returned graph is empty.
342 """
343 metadata = {
344 "skip_existing_in": self._skip_existing_in,
345 "skip_existing": bool(self._skip_existing_in),
346 "data_query": where,
347 }
348 qg_builder = self.make_quantum_graph_builder(pipeline, where, builder_class=builder_class, **kwargs)
349 graph = qg_builder.finish(
350 metadata=metadata, attach_datastore_records=attach_datastore_records
351 ).assemble()
352 _LOG.info(
353 "PredictedQuantumGraph contains %d quanta for %d tasks.",
354 len(graph),
355 len(graph.quanta_by_task),
356 )
357 return graph
359 def run_pipeline(
360 self,
361 graph: QuantumGraph | PredictedQuantumGraph,
362 fail_fast: bool = False,
363 graph_executor: QuantumGraphExecutor | None = None,
364 num_proc: int = 1,
365 ) -> None:
366 """Run a pipeline in the form of a prepared quantum graph.
368 Pre-execution initialization must have already been run;
369 see `pre_execute_qgraph`.
371 Parameters
372 ----------
373 graph : `.QuantumGraph` or `.quantum_graph.PredictedQuantumGraph`
374 The pipeline and datasets to execute.
375 fail_fast : `bool`, optional
376 If `True`, abort all execution if any task fails when
377 running with multiple processes. Only used with the default graph
378 executor).
379 graph_executor : `.quantum_graph_executor.QuantumGraphExecutor`,\
380 optional
381 A custom graph executor. By default, a new instance of
382 `.mp_graph_executor.MPGraphExecutor` is used.
383 num_proc : `int`, optional
384 The number of processes that can be used to run the pipeline. The
385 default value ensures that no subprocess is created. Only used with
386 the default graph executor.
387 """
388 if not graph_executor:
389 quantum_executor = SingleQuantumExecutor(
390 butler=self._butler,
391 task_factory=self._task_factory,
392 skip_existing_in=self._skip_existing_in,
393 clobber_outputs=self._clobber_output,
394 resources=self.resources,
395 raise_on_partial_outputs=self.raise_on_partial_outputs,
396 )
397 graph_executor = MPGraphExecutor(
398 num_proc=num_proc,
399 timeout=2_592_000.0, # In practice, timeout is never helpful; set to 30 days.
400 quantum_executor=quantum_executor,
401 fail_fast=fail_fast,
402 )
403 # Have to reset connection pool to avoid sharing connections with
404 # forked processes.
405 self._butler.registry.resetConnectionPool()
407 graph_executor.execute(graph)