Coverage for python/lsst/pipe/base/separable_pipeline_executor.py: 97%
74 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-25 08:10 +0000
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, DatasetRef
44from lsst.daf.butler._rubin.temporary_for_ingest import TemporaryForIngest
46from ._quantumContext import ExecutionResources
47from .all_dimensions_quantum_graph_builder import AllDimensionsQuantumGraphBuilder
48from .graph import QuantumGraph
49from .mp_graph_executor import MPGraphExecutor, MPGraphExecutorError
50from .pipeline import Pipeline
51from .quantum_graph import PredictedQuantumGraph
52from .quantum_graph_builder import QuantumGraphBuilder
53from .quantum_graph_executor import QuantumGraphExecutor
54from .single_quantum_executor import SingleQuantumExecutor
55from .taskFactory import TaskFactory
57_LOG = logging.getLogger(__name__)
60class SeparablePipelineExecutor:
61 """An executor that allows each step of pipeline execution to be
62 run independently.
64 The executor can run any or all of the following steps:
66 * pre-execution initialization
67 * pipeline building
68 * quantum graph generation
69 * quantum graph execution
71 Any of these steps can also be handed off to external code without
72 compromising the remaining ones.
74 Parameters
75 ----------
76 butler : `lsst.daf.butler.Butler`
77 A Butler whose ``collections`` and ``run`` attributes contain the input
78 and output collections to use for processing.
79 clobber_output : `bool`, optional
80 If set, the pipeline execution overwrites existing output files.
81 Otherwise, any conflict between existing and new outputs is an error.
82 skip_existing_in : `~collections.abc.Iterable` [`str`], optional
83 If not empty, the pipeline execution searches the listed collections
84 for existing outputs, and skips any quanta that have run to completion
85 (or have no work to do). Otherwise, all tasks are attempted (subject to
86 ``clobber_output``).
87 retained_dataset_types : `~collections.abc.Iterable` [`str`], optional
88 Dataset type names or glob-style wildcard patterns for types that
89 should be present in ``skip_existing_in`` whenever the producing task
90 ran successfully. Dataset types not in this list are treated as not
91 retained: when a downstream quantum must run, the builder propagates
92 the must-run signal backward through non-retained input edges, forcing
93 the upstream quanta that need to regenerate those intermediates to also
94 run. Has no effect without ``skip_existing_in``. ``["*"]`` means
95 retaining all datasets, equivalent to not providing this option.
96 task_factory : `.TaskFactory`, optional
97 A custom task factory for use in pre-execution and execution. By
98 default, a new instance of `.TaskFactory` is used.
99 resources : `.ExecutionResources`
100 The resources available to each quantum being executed.
101 raise_on_partial_outputs : `bool`, optional
102 If `True` raise exceptions chained by
103 `.AnnotatedPartialOutputsError` immediately, instead of
104 considering the partial result a success and continuing to run
105 downstream tasks.
106 """
108 def __init__(
109 self,
110 butler: Butler,
111 clobber_output: bool = False,
112 skip_existing_in: Iterable[str] | None = None,
113 retained_dataset_types: Iterable[str] | None = None,
114 task_factory: TaskFactory | None = None,
115 resources: ExecutionResources | None = None,
116 raise_on_partial_outputs: bool = True,
117 ):
118 self._butler = Butler.from_config(
119 butler=butler, collections=butler.collections.defaults, run=butler.run
120 )
121 if not self._butler.collections.defaults:
122 raise ValueError("Butler must specify input collections for pipeline.")
123 if not self._butler.run:
124 raise ValueError("Butler must specify output run for pipeline.")
126 self._clobber_output = clobber_output
127 self._skip_existing_in = list(skip_existing_in) if skip_existing_in else []
128 self._retained_dataset_types = list(retained_dataset_types) if retained_dataset_types else None
130 self._task_factory = task_factory if task_factory else TaskFactory()
131 self.resources = resources
132 self.raise_on_partial_outputs = raise_on_partial_outputs
134 def pre_execute_qgraph(
135 self,
136 graph: QuantumGraph | PredictedQuantumGraph,
137 register_dataset_types: bool = False,
138 save_init_outputs: bool = True,
139 save_versions: bool = True,
140 ) -> None:
141 """Run pre-execution initialization.
143 This method will be deprecated after DM-38041, to be replaced with a
144 method that takes either a `.Pipeline` or a
145 resolved `.pipeline_graph.PipelineGraph` instead of a `.QuantumGraph`.
147 Parameters
148 ----------
149 graph : `.QuantumGraph` or `.quantum_graph.PredictedQuantumGraph`
150 The quantum graph defining the pipeline and datasets to
151 be initialized.
152 register_dataset_types : `bool`, optional
153 If `True`, register all output dataset types from the pipeline
154 represented by ``graph``.
155 save_init_outputs : `bool`, optional
156 If `True`, create init-output datasets in this object's output run.
157 save_versions : `bool`, optional
158 If `True`, save a package versions dataset.
159 """
160 if register_dataset_types:
161 graph.pipeline_graph.register_dataset_types(self._butler, include_packages=save_versions)
162 if save_init_outputs:
163 graph.write_init_outputs(self._butler, skip_existing=(self._butler.run in self._skip_existing_in))
164 graph.write_configs(self._butler)
165 if save_versions:
166 graph.write_packages(self._butler)
168 def make_pipeline(self, pipeline_uri: str | lsst.resources.ResourcePath) -> Pipeline:
169 """Build a pipeline from pipeline and configuration information.
171 Parameters
172 ----------
173 pipeline_uri : `str` or `lsst.resources.ResourcePath`
174 URI to a file containing a pipeline definition. A URI fragment may
175 be used to specify a subset of the pipeline, as described in
176 :ref:`pipeline-running-intro`.
178 Returns
179 -------
180 pipeline : `.Pipeline`
181 The fully-built pipeline.
182 """
183 return Pipeline.from_uri(pipeline_uri)
185 def make_quantum_graph_builder(
186 self,
187 pipeline: Pipeline,
188 where: str = "",
189 *,
190 builder_class: type[QuantumGraphBuilder] = AllDimensionsQuantumGraphBuilder,
191 **kwargs: Any,
192 ) -> QuantumGraphBuilder:
193 """Initialize a quantum graph builder from a pipeline and input
194 datasets.
196 Parameters
197 ----------
198 pipeline : `.Pipeline`
199 The pipeline for which to generate a quantum graph.
200 where : `str`, optional
201 A data ID query that constrains the quanta generated. Must not be
202 provided if a custom ``builder_class`` is given and that class does
203 not accept ``where`` as a construction argument.
204 builder_class : `type` [ \
205 `.quantum_graph_builder.QuantumGraphBuilder` ], optional
206 Quantum graph builder implementation. Ignored if ``builder`` is
207 provided.
208 **kwargs
209 Additional keyword arguments are forwarded to ``builder_class``
210 when a quantum graph builder instance is constructed. All
211 arguments accepted by the
212 `~.quantum_graph_builder.QuantumGraphBuilder` base
213 class are provided automatically (from explicit arguments to this
214 method and executor attributes) and do not need to be included
215 as keyword arguments.
217 Returns
218 -------
219 builder : `.quantum_graph_builder.QuantumGraphBuilder`
220 A quantum graph builder.
221 """
222 if where: 222 ↛ 225line 222 didn't jump to line 225 because the condition on line 222 was never true
223 # Only pass 'where' if it's actually provided, since some
224 # QuantumGraphBuilder subclasses may not accept it.
225 kwargs["where"] = where
226 return builder_class(
227 pipeline.to_graph(),
228 self._butler,
229 skip_existing_in=self._skip_existing_in,
230 retained_dataset_types=self._retained_dataset_types,
231 clobber=self._clobber_output,
232 **kwargs,
233 )
235 def make_quantum_graph(
236 self,
237 pipeline: Pipeline,
238 where: str = "",
239 *,
240 builder_class: type[QuantumGraphBuilder] = AllDimensionsQuantumGraphBuilder,
241 attach_datastore_records: bool = False,
242 **kwargs: Any,
243 ) -> QuantumGraph:
244 """Build a quantum graph from a pipeline and input datasets.
246 This returns an instance of the old `.QuantumGraph` class. Use
247 `build_quantum_graph` to construct a
248 `.quantum_graph.PredictedQuantumGraph`.
250 Parameters
251 ----------
252 pipeline : `.Pipeline`
253 The pipeline for which to generate a quantum graph.
254 where : `str`, optional
255 A data ID query that constrains the quanta generated. Must not be
256 provided if a custom ``builder_class`` is given and that class does
257 not accept ``where`` as a construction argument.
258 builder_class : `type` [ \
259 `.quantum_graph_builder.QuantumGraphBuilder` ], optional
260 Quantum graph builder implementation. Ignored if ``builder`` is
261 provided.
262 attach_datastore_records : `bool`, optional
263 Whether to attach datastore records. These are currently used only
264 by `lsst.daf.butler.QuantumBackedButler`, which is not used by
265 `SeparablePipelineExecutor` for execution.
266 **kwargs
267 Additional keyword arguments are forwarded to ``builder_class``
268 when a quantum graph builder instance is constructed. All
269 arguments accepted by the
270 `~.quantum_graph_builder.QuantumGraphBuilder` base
271 class are provided automatically (from explicit arguments to this
272 method and executor attributes) and do not need to be included
273 as keyword arguments.
275 Returns
276 -------
277 graph : `.QuantumGraph`
278 The quantum graph for ``.Pipeline`` as run on the datasets
279 identified by ``where``.
281 Notes
282 -----
283 This method does no special handling of empty quantum graphs. If
284 needed, clients can use `len` to test if the returned graph is empty.
285 """
286 metadata = {
287 "input": self._butler.collections.defaults,
288 "output_run": self._butler.run,
289 "skip_existing_in": self._skip_existing_in,
290 "skip_existing": bool(self._skip_existing_in),
291 "retained_dataset_types": self._retained_dataset_types,
292 "data_query": where,
293 "user": getpass.getuser(),
294 "time": str(datetime.datetime.now()),
295 }
296 qg_builder = self.make_quantum_graph_builder(pipeline, where, builder_class=builder_class, **kwargs)
297 graph = qg_builder.build(metadata=metadata, attach_datastore_records=attach_datastore_records)
298 _LOG.info(
299 "QuantumGraph contains %d quanta for %d tasks, graph ID: %r",
300 len(graph),
301 len(graph.taskGraph),
302 graph.graphID,
303 )
304 return graph
306 def build_quantum_graph(
307 self,
308 pipeline: Pipeline,
309 where: str = "",
310 *,
311 builder_class: type[QuantumGraphBuilder] = AllDimensionsQuantumGraphBuilder,
312 attach_datastore_records: bool = False,
313 **kwargs: Any,
314 ) -> PredictedQuantumGraph:
315 """Build a quantum graph from a pipeline and input datasets.
317 This returns an instance of the new
318 `.quantum_graph.PredictedQuantumGraph` class. Use `make_quantum_graph`
319 to construct a `.QuantumGraph`.
321 Parameters
322 ----------
323 pipeline : `.Pipeline`
324 The pipeline for which to generate a quantum graph.
325 where : `str`, optional
326 A data ID query that constrains the quanta generated. Must not be
327 provided if a custom ``builder_class`` is given and that class does
328 not accept ``where`` as a construction argument.
329 builder_class : `type` [ \
330 `.quantum_graph_builder.QuantumGraphBuilder` ], optional
331 Quantum graph builder implementation. Ignored if ``builder`` is
332 provided.
333 attach_datastore_records : `bool`, optional
334 Whether to attach datastore records. These are currently used only
335 by `lsst.daf.butler.QuantumBackedButler`, which is not used by
336 `SeparablePipelineExecutor` for execution.
337 **kwargs
338 Additional keyword arguments are forwarded to ``builder_class``
339 when a quantum graph builder instance is constructed. All
340 arguments accepted by the
341 `~.quantum_graph_builder.QuantumGraphBuilder` base
342 class are provided automatically (from explicit arguments to this
343 method and executor attributes) and do not need to be included
344 as keyword arguments.
346 Returns
347 -------
348 graph : `.QuantumGraph`
349 The quantum graph for ``.Pipeline`` as run on the datasets
350 identified by ``where``.
352 Notes
353 -----
354 This method does no special handling of empty quantum graphs. If
355 needed, clients can use `len` to test if the returned graph is empty.
356 """
357 metadata = {
358 "skip_existing_in": self._skip_existing_in,
359 "skip_existing": bool(self._skip_existing_in),
360 "retained_dataset_types": self._retained_dataset_types,
361 "data_query": where,
362 }
363 qg_builder = self.make_quantum_graph_builder(pipeline, where, builder_class=builder_class, **kwargs)
364 graph = qg_builder.finish(
365 metadata=metadata, attach_datastore_records=attach_datastore_records
366 ).assemble()
367 _LOG.info(
368 "PredictedQuantumGraph contains %d quanta for %d tasks.",
369 len(graph),
370 len(graph.quanta_by_task),
371 )
372 return graph
374 def run_pipeline(
375 self,
376 graph: QuantumGraph | PredictedQuantumGraph,
377 fail_fast: bool = False,
378 graph_executor: QuantumGraphExecutor | None = None,
379 num_proc: int = 1,
380 *,
381 provenance_dataset_ref: DatasetRef | None = None,
382 ) -> None:
383 """Run a pipeline in the form of a prepared quantum graph.
385 Pre-execution initialization must have already been run;
386 see `pre_execute_qgraph`.
388 Parameters
389 ----------
390 graph : `.QuantumGraph` or `.quantum_graph.PredictedQuantumGraph`
391 The pipeline and datasets to execute.
392 fail_fast : `bool`, optional
393 If `True`, abort all execution if any task fails when
394 running with multiple processes. Only used with the default graph
395 executor).
396 graph_executor : `.quantum_graph_executor.QuantumGraphExecutor`,\
397 optional
398 A custom graph executor. By default, a new instance of
399 `.mp_graph_executor.MPGraphExecutor` is used.
400 num_proc : `int`, optional
401 The number of processes that can be used to run the pipeline. The
402 default value ensures that no subprocess is created. Only used with
403 the default graph executor.
404 provenance_dataset_ref : `lsst.daf.butler.DatasetRef`, optional
405 Dataset that should be used to save provenance. Provenance is only
406 supported when running in a single process (at least for the
407 default quantum executor), and should not be enabled in contexts
408 where a quantum might be executed more than once (i.e. retried)
409 within the same `~lsst.daf.butler.CollectionType.RUN` collection.
410 The caller is responsible for registering the dataset type and for
411 ensuring that the dimensions of this dataset do not lead to
412 uniqueness conflicts.
413 """
414 if not graph_executor: 414 ↛ 433line 414 didn't jump to line 433 because the condition on line 414 was always true
415 quantum_executor = SingleQuantumExecutor(
416 butler=self._butler,
417 task_factory=self._task_factory,
418 skip_existing_in=self._skip_existing_in,
419 clobber_outputs=self._clobber_output,
420 resources=self.resources,
421 raise_on_partial_outputs=self.raise_on_partial_outputs,
422 )
423 graph_executor = MPGraphExecutor(
424 num_proc=num_proc,
425 timeout=2_592_000.0, # In practice, timeout is never helpful; set to 30 days.
426 quantum_executor=quantum_executor,
427 fail_fast=fail_fast,
428 )
429 # Have to reset connection pool to avoid sharing connections with
430 # forked processes.
431 self._butler.registry.resetConnectionPool()
433 if provenance_dataset_ref is not None:
434 with TemporaryForIngest(self._butler, provenance_dataset_ref) as temporary:
435 try:
436 graph_executor.execute(graph, provenance_graph_file=temporary.ospath)
437 temporary.ingest()
438 except MPGraphExecutorError:
439 # If the graph executor itself raised, it will have
440 # finished the provenance rewrite. In other cases the
441 # temporary file might be incomplete or corrupted and we
442 # can't roll the dice on ingesting it.
443 temporary.ingest()
444 raise
446 else:
447 graph_executor.execute(graph)