Coverage for python/lsst/ctrl/mpexec/cli/butler_factory.py: 21%
170 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-04 01:49 -0800
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__ = [
31 "ButlerFactory",
32 "OutputChainedCollectionInfo",
33 "OutputRunCollectionInfo",
34]
36import atexit
37import shutil
38from collections.abc import Iterable, Sequence
40from lsst.daf.butler import Butler, CollectionType
41from lsst.daf.butler.datastore.cache_manager import DatastoreCacheManager
42from lsst.daf.butler.registry import MissingCollectionError, RegistryDefaults
43from lsst.daf.butler.registry.wildcards import CollectionWildcard
44from lsst.pipe.base import Instrument, PipelineGraph
45from lsst.pipe.base.pipeline_graph import NodeType
46from lsst.resources import ResourcePathExpression
47from lsst.utils.logging import getLogger
49_LOG = getLogger(__name__)
52class OutputChainedCollectionInfo:
53 """A helper class for handling command-line arguments related to an output
54 `~lsst.daf.butler.CollectionType.CHAINED` collection.
56 Parameters
57 ----------
58 butler : `lsst.daf.butler.Butler`
59 Butler that collections will be added to and/or queried from.
60 name : `str`
61 Name of the collection given on the command line.
62 """
64 def __init__(self, butler: Butler, name: str):
65 self.name = name
66 try:
67 self.chain = tuple(butler.collections.get_info(name).children)
68 self.exists = True
69 except MissingCollectionError:
70 self.chain = ()
71 self.exists = False
73 def __str__(self) -> str:
74 return self.name
76 name: str
77 """Name of the collection provided on the command line (`str`).
78 """
80 exists: bool
81 """Whether this collection already exists in the butler (`bool`).
82 """
84 chain: tuple[str, ...]
85 """The definition of the collection, if it already exists (`tuple`[`str`]).
87 Empty if the collection does not already exist.
88 """
91class OutputRunCollectionInfo:
92 """A helper class for handling command-line arguments related to an output
93 `~lsst.daf.butler.CollectionType.RUN` collection.
95 Parameters
96 ----------
97 butler : `lsst.daf.butler.Butler`
98 Butler that collections will be added to and/or queried from.
99 name : `str`
100 Name of the collection given on the command line.
101 """
103 def __init__(self, butler: Butler, name: str):
104 self.name = name
105 try:
106 actual_type = butler.collections.get_info(name).type
107 if actual_type is not CollectionType.RUN:
108 raise TypeError(f"Collection '{name}' exists but has type {actual_type.name}, not RUN.")
109 self.exists = True
110 except MissingCollectionError:
111 self.exists = False
113 name: str
114 """Name of the collection provided on the command line (`str`).
115 """
117 exists: bool
118 """Whether this collection already exists in the butler (`bool`).
119 """
122class ButlerFactory:
123 """A helper class for processing command-line arguments related to input
124 and output collections.
126 Parameters
127 ----------
128 butler : `lsst.daf.butler.Butler`
129 Butler that collections will be added to and/or queried from.
130 output : `str` or `None`
131 The name of a `~lsst.daf.butler.CollectionType.CHAINED` input/output
132 collection.
133 output_run : `str` or `None`
134 The name of a `~lsst.daf.butler.CollectionType.RUN` input/output
135 collection.
136 inputs : `str` or `~collections.abc.Iterable` [`str`]
137 Input collection name or iterable of collection names.
138 extend_run : `bool`
139 A boolean indicating whether ``output_run`` should already exist and be
140 extended.
141 rebase : `bool`
142 A boolean indicating whether to force the ``output`` collection to be
143 consistent with ``inputs`` and ``output`` run such that the ``output``
144 collection has output run collections first (i.e. those that start with
145 the same prefix), then the new inputs, then any original inputs not
146 included in the new inputs.
147 writeable : `bool`
148 If `True`, a `~lsst.daf.butler.Butler` is being initialized in a
149 context where actual writes should happens, and hence no output run
150 is necessary.
152 Raises
153 ------
154 ValueError
155 Raised if ``writeable is True`` but there are no output collections.
156 """
158 def __init__(
159 self,
160 butler: Butler,
161 *,
162 output: str | None,
163 output_run: str | None,
164 inputs: str | Iterable[str],
165 extend_run: bool = False,
166 rebase: bool = False,
167 writeable: bool,
168 ):
169 if output is not None:
170 self.output = OutputChainedCollectionInfo(butler, output)
171 else:
172 self.output = None
173 if output_run is not None:
174 if rebase and self.output and not output_run.startswith(self.output.name):
175 raise ValueError("Cannot rebase if output run does not start with output collection name.")
176 self.output_run = OutputRunCollectionInfo(butler, output_run)
177 elif self.output is not None:
178 if extend_run:
179 if not self.output.chain:
180 raise ValueError("Cannot use --extend-run option with non-existing or empty output chain")
181 run_name = self.output.chain[0]
182 else:
183 run_name = f"{self.output}/{Instrument.makeCollectionTimestamp()}"
184 self.output_run = OutputRunCollectionInfo(butler, run_name)
185 elif not writeable:
186 # If we're not writing yet, ok to have no output run.
187 self.output_run = None
188 else:
189 raise ValueError("Cannot write without at least one of (--output, --output-run).")
190 # Recursively flatten any input CHAINED collections. We do this up
191 # front so we can tell if the user passes the same inputs on subsequent
192 # calls, even though we also flatten when we define the output CHAINED
193 # collection.
194 self.inputs = tuple(butler.collections.query(inputs, flatten_chains=True)) if inputs else ()
196 # If things are inconsistent and user has asked for a rebase then
197 # construct the new output chain.
198 if rebase and self._check_output_input_consistency():
199 assert self.output is not None
200 newOutputChain = [item for item in self.output.chain if item.startswith(self.output.name)]
201 newOutputChain.extend([item for item in self.inputs if item not in newOutputChain])
202 newOutputChain.extend([item for item in self.output.chain if item not in newOutputChain])
203 self.output.chain = tuple(newOutputChain)
205 def check(self, *, extend_run: bool, replace_run: bool, prune_replaced: str | None = None) -> None:
206 """Check command-line options for consistency with each other and the
207 data repository.
209 Parameters
210 ----------
211 extend_run : `bool`
212 Whether the ``output_run`` should already exist and be extended.
213 replace_run : `bool`
214 Whether the ``output_run`` should be replaced in the ``output``
215 chain.
216 prune_replaced : `str` or `None`
217 If ``replace_run=True``, whether/how datasets in the old run should
218 be removed. Options are ``"purge"``, ``"unstore"``, and `None`.
219 """
220 assert not (extend_run and replace_run), "In mutually-exclusive group in ArgumentParser."
221 if consistencyError := self._check_output_input_consistency():
222 raise ValueError(consistencyError)
224 if extend_run:
225 if self.output_run is None:
226 raise ValueError("Cannot --extend-run when no output collection is given.")
227 elif not self.output_run.exists:
228 raise ValueError(
229 f"Cannot --extend-run; output collection '{self.output_run.name}' does not exist."
230 )
231 if not extend_run and self.output_run is not None and self.output_run.exists:
232 raise ValueError(
233 f"Output run '{self.output_run.name}' already exists, but --extend-run was not given."
234 )
235 if prune_replaced and not replace_run:
236 raise ValueError("--prune-replaced requires --replace-run.")
237 if replace_run and (self.output is None or not self.output.exists):
238 raise ValueError("--output must point to an existing CHAINED collection for --replace-run.")
240 def _check_output_input_consistency(self) -> str | None:
241 if self.inputs and self.output is not None and self.output.exists:
242 # Passing the same inputs that were used to initialize the output
243 # collection is allowed; this means the inputs must appear as a
244 # contiguous subsequence of outputs (normally they're also at the
245 # end, but --rebase will in general put them in the middle).
246 for n in reversed(range(1 + len(self.output.chain) - len(self.inputs))):
247 if self.inputs == self.output.chain[n : n + len(self.inputs)]:
248 return None
249 return (
250 f"Output CHAINED collection {self.output.name!r} exists and does not include the "
251 f"same sequence of (flattened) input collections {self.inputs} as a contiguous "
252 "subsequence. "
253 "Use --rebase to ignore this problem and reset the output collection, but note that "
254 "this may obfuscate what inputs were actually used to produce these outputs."
255 )
256 return None
258 @classmethod
259 def _make_read_parts(
260 cls,
261 butler_config: ResourcePathExpression,
262 *,
263 output: str | None,
264 output_run: str | None,
265 inputs: str | Iterable[str],
266 extend_run: bool = False,
267 rebase: bool = False,
268 replace_run: bool,
269 prune_replaced: str | None = None,
270 ) -> tuple[Butler, Sequence[str], ButlerFactory]:
271 """Parse arguments to support implementations of `make_read_butler` and
272 `make_butler_and_collections`.
274 Parameters
275 ----------
276 butler_config : convertible to `lsst.resources.ResourcePath`
277 Path to configuration for the butler.
278 output : `str` or `None`
279 The name of a `~lsst.daf.butler.CollectionType.CHAINED`
280 input/output collection.
281 output_run : `str` or `None`
282 The name of a `~lsst.daf.butler.CollectionType.RUN` input/output
283 collection.
284 inputs : `str` or `~collections.abc.Iterable` [`str`]
285 Input collection name or iterable of collection names.
286 extend_run : `bool`
287 A boolean indicating whether ``output_run`` should already exist
288 and be extended.
289 rebase : `bool`
290 A boolean indicating whether to force the ``output`` collection to
291 be consistent with ``inputs`` and ``output`` run such that the
292 ``output`` collection has output run collections first (i.e. those
293 that start with the same prefix), then the new inputs, then any
294 original inputs not included in the new inputs.
295 replace_run : `bool`
296 Whether the ``output_run`` should be replaced in the ``output``
297 chain.
298 prune_replaced : `str` or `None`
299 If ``replace_run=True``, whether/how datasets in the old run should
300 be removed. Options are ``"purge"``, ``"unstore"``, and `None`.
302 Returns
303 -------
304 butler : `lsst.daf.butler.Butler`
305 A read-only butler constructed from the repo at
306 ``args.butler_config``, but with no default collections.
307 inputs : `~collections.abc.Sequence` [ `str` ]
308 A collection search path constructed according to ``args``.
309 self : `ButlerFactory`
310 A new `ButlerFactory` instance representing the processed version
311 of ``args``.
312 """
313 butler = Butler.from_config(butler_config, writeable=False)
314 self = cls(
315 butler,
316 output=output,
317 output_run=output_run,
318 inputs=inputs,
319 extend_run=extend_run,
320 rebase=rebase,
321 writeable=False,
322 )
323 self.check(extend_run=extend_run, replace_run=replace_run, prune_replaced=prune_replaced)
324 if self.output and self.output.exists:
325 if replace_run:
326 replaced = self.output.chain[0]
327 inputs = list(self.output.chain[1:])
328 _LOG.debug(
329 "Simulating collection search in '%s' after removing '%s'.", self.output.name, replaced
330 )
331 else:
332 inputs = [self.output.name]
333 else:
334 inputs = list(self.inputs)
335 if extend_run:
336 assert self.output_run is not None, "Output collection has to be specified."
337 inputs.insert(0, self.output_run.name)
338 collSearch = CollectionWildcard.from_expression(inputs).require_ordered()
339 return butler, collSearch, self
341 @classmethod
342 def make_read_butler(
343 cls,
344 butler_config: ResourcePathExpression,
345 *,
346 output: str | None,
347 output_run: str | None,
348 inputs: str | Iterable[str],
349 extend_run: bool = False,
350 rebase: bool = False,
351 replace_run: bool,
352 prune_replaced: str | None = None,
353 ) -> Butler:
354 """Construct a read-only butler according to the given command-line
355 arguments.
357 Parameters
358 ----------
359 butler_config : convertible to `lsst.resources.ResourcePath`
360 Path to configuration for the butler.
361 output : `str` or `None`
362 The name of a `~lsst.daf.butler.CollectionType.CHAINED`
363 input/output collection.
364 output_run : `str` or `None`
365 The name of a `~lsst.daf.butler.CollectionType.RUN` input/output
366 collection.
367 inputs : `str` or `~collections.abc.Iterable` [`str`]
368 Input collection name or iterable of collection names.
369 extend_run : `bool`
370 A boolean indicating whether ``output_run`` should already exist
371 and be extended.
372 rebase : `bool`
373 A boolean indicating whether to force the ``output`` collection to
374 be consistent with ``inputs`` and ``output`` run such that the
375 ``output`` collection has output run collections first (i.e. those
376 that start with the same prefix), then the new inputs, then any
377 original inputs not included in the new inputs.
378 replace_run : `bool`
379 Whether the ``output_run`` should be replaced in the ``output``
380 chain.
381 prune_replaced : `str` or `None`
382 If ``replace_run=True``, whether/how datasets in the old run should
383 be removed. Options are ``"purge"``, ``"unstore"``, and `None`.
385 Returns
386 -------
387 butler : `lsst.daf.butler.Butler`
388 A read-only butler initialized with the given collections.
389 """
390 cls.define_datastore_cache() # Ensure that this butler can use a shared cache.
391 butler, inputs, _ = cls._make_read_parts(
392 butler_config,
393 output=output,
394 output_run=output_run,
395 inputs=inputs,
396 extend_run=extend_run,
397 rebase=rebase,
398 replace_run=replace_run,
399 prune_replaced=prune_replaced,
400 )
401 _LOG.debug("Preparing butler to read from %s.", inputs)
402 return Butler.from_config(butler=butler, collections=inputs)
404 @classmethod
405 def make_butler_and_collections(
406 cls,
407 butler_config: ResourcePathExpression,
408 *,
409 output: str | None,
410 output_run: str | None,
411 inputs: str | Iterable[str],
412 extend_run: bool = False,
413 rebase: bool = False,
414 replace_run: bool,
415 prune_replaced: str | None = None,
416 ) -> tuple[Butler, Sequence[str], str | None]:
417 """Return a read-only butler, a collection search path, and the name
418 of the run to be used for future writes.
420 Parameters
421 ----------
422 butler_config : convertible to `lsst.resources.ResourcePath`
423 Path to configuration for the butler.
424 output : `str` or `None`
425 The name of a `~lsst.daf.butler.CollectionType.CHAINED`
426 input/output collection.
427 output_run : `str` or `None`
428 The name of a `~lsst.daf.butler.CollectionType.RUN` input/output
429 collection.
430 inputs : `str` or `~collections.abc.Iterable` [`str`]
431 Input collection name or iterable of collection names.
432 extend_run : `bool`
433 A boolean indicating whether ``output_run`` should already exist
434 and be extended.
435 rebase : `bool`
436 A boolean indicating whether to force the ``output`` collection to
437 be consistent with ``inputs`` and ``output`` run such that the
438 ``output`` collection has output run collections first (i.e. those
439 that start with the same prefix), then the new inputs, then any
440 original inputs not included in the new inputs.
441 replace_run : `bool`
442 Whether the ``output_run`` should be replaced in the ``output``
443 chain.
444 prune_replaced : `str` or `None`
445 If ``replace_run=True``, whether/how datasets in the old run should
446 be removed. Options are ``"purge"``, ``"unstore"``, and `None`.
448 Returns
449 -------
450 butler : `lsst.daf.butler.Butler`
451 A read-only butler that collections will be added to and/or queried
452 from.
453 inputs : `Sequence` [ `str` ]
454 Collections to search for datasets.
455 run : `str` or `None`
456 Name of the output `~lsst.daf.butler.CollectionType.RUN` collection
457 if it already exists, or `None` if it does not.
458 """
459 butler, inputs, self = cls._make_read_parts(
460 butler_config,
461 output=output,
462 output_run=output_run,
463 inputs=inputs,
464 extend_run=extend_run,
465 rebase=rebase,
466 replace_run=replace_run,
467 prune_replaced=prune_replaced,
468 )
469 run: str | None = None
470 if extend_run:
471 assert self.output_run is not None, "Output collection has to be specified."
472 if self.output_run is not None:
473 run = self.output_run.name
474 _LOG.debug("Preparing butler to read from %s and expect future writes to '%s'.", inputs, run)
475 return butler, inputs, run
477 @staticmethod
478 def define_datastore_cache() -> None:
479 """Define where datastore cache directories should be found.
481 Notes
482 -----
483 All the jobs should share a datastore cache if applicable. This
484 method asks for a shared fallback cache to be defined and then
485 configures an exit handler to clean it up.
486 """
487 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset()
488 if defined:
489 atexit.register(shutil.rmtree, cache_dir, ignore_errors=True)
490 _LOG.debug("Defining shared datastore cache directory to %s", cache_dir)
492 @classmethod
493 def make_write_butler(
494 cls,
495 butler_config: ResourcePathExpression,
496 pipeline_graph: PipelineGraph,
497 *,
498 output: str | None,
499 output_run: str | None,
500 inputs: str | Iterable[str],
501 extend_run: bool = False,
502 rebase: bool = False,
503 replace_run: bool,
504 prune_replaced: str | None = None,
505 ) -> Butler:
506 """Return a read-write butler initialized to write to and read from
507 the collections specified by the given command-line arguments.
509 Parameters
510 ----------
511 butler_config : convertible to `lsst.resources.ResourcePath`
512 Path to configuration for the butler.
513 pipeline_graph : `lsst.pipe.base.PipelineGraph`
514 Definitions for tasks in a pipeline.
515 output : `str` or `None`
516 The name of a `~lsst.daf.butler.CollectionType.CHAINED`
517 input/output collection.
518 output_run : `str` or `None`
519 The name of a `~lsst.daf.butler.CollectionType.RUN` input/output
520 collection.
521 inputs : `str` or `~collections.abc.Iterable` [`str`]
522 Input collection name or iterable of collection names.
523 extend_run : `bool`
524 A boolean indicating whether ``output_run`` should already exist
525 and be extended.
526 rebase : `bool`
527 A boolean indicating whether to force the ``output`` collection to
528 be consistent with ``inputs`` and ``output`` run such that the
529 ``output`` collection has output run collections first (i.e. those
530 that start with the same prefix), then the new inputs, then any
531 original inputs not included in the new inputs.
532 replace_run : `bool`
533 Whether the ``output_run`` should be replaced in the ``output``
534 chain.
535 prune_replaced : `str` or `None`
536 If ``replace_run=True``, whether/how datasets in the old run should
537 be removed. Options are ``"purge"``, ``"unstore"``, and `None`.
539 Returns
540 -------
541 butler : `lsst.daf.butler.Butler`
542 A read-write butler initialized according to the given arguments.
543 """
544 cls.define_datastore_cache() # Ensure that this butler can use a shared cache.
545 butler = Butler.from_config(butler_config, writeable=True)
546 self = cls(
547 butler,
548 output=output,
549 output_run=output_run,
550 inputs=inputs,
551 extend_run=extend_run,
552 rebase=rebase,
553 writeable=True,
554 )
555 self.check(extend_run=extend_run, replace_run=replace_run, prune_replaced=prune_replaced)
556 assert self.output_run is not None, "Output collection has to be specified." # for mypy
557 if self.output is not None:
558 chain_definition = list(
559 butler.collections.query(
560 self.output.chain if self.output.exists else self.inputs,
561 flatten_chains=True,
562 include_chains=False,
563 )
564 )
565 if replace_run:
566 replaced = chain_definition.pop(0)
567 if prune_replaced == "unstore":
568 # Remove datasets from datastore
569 with butler.transaction():
570 # we want to remove regular outputs from this pipeline,
571 # but keep initOutputs, configs, and versions.
572 refs = [
573 ref
574 for ref in butler.registry.queryDatasets(..., collections=replaced)
575 if (
576 (producer := pipeline_graph.producer_of(ref.datasetType.name)) is not None
577 and producer.key.node_type is NodeType.TASK # i.e. not TASK_INIT
578 )
579 ]
580 butler.pruneDatasets(refs, unstore=True, disassociate=False)
581 elif prune_replaced == "purge":
582 # Erase entire collection and all datasets, need to remove
583 # collection from its chain collection first.
584 with butler.transaction():
585 butler.collections.redefine_chain(self.output.name, chain_definition)
586 butler.removeRuns([replaced], unstore=True)
587 elif prune_replaced is not None:
588 raise NotImplementedError(f"Unsupported --prune-replaced option '{prune_replaced}'.")
589 if not self.output.exists:
590 butler.collections.register(self.output.name, CollectionType.CHAINED)
591 if not extend_run:
592 butler.collections.register(self.output_run.name, CollectionType.RUN)
593 chain_definition.insert(0, self.output_run.name)
594 butler.collections.redefine_chain(self.output.name, chain_definition)
595 _LOG.debug(
596 "Preparing butler to write to '%s' and read from '%s'=%s",
597 self.output_run.name,
598 self.output.name,
599 chain_definition,
600 )
601 butler.registry.defaults = RegistryDefaults(
602 run=self.output_run.name, collections=self.output.name
603 )
604 else:
605 inputs = (self.output_run.name,) + self.inputs
606 _LOG.debug("Preparing butler to write to '%s' and read from %s.", self.output_run.name, inputs)
607 butler.registry.defaults = RegistryDefaults(run=self.output_run.name, collections=inputs)
608 return butler
610 output: OutputChainedCollectionInfo | None
611 """Information about the output chained collection, if there is or will be
612 one (`OutputChainedCollectionInfo` or `None`).
613 """
615 output_run: OutputRunCollectionInfo | None
616 """Information about the output run collection, if there is or will be
617 one (`OutputRunCollectionInfo` or `None`).
618 """
620 inputs: tuple[str, ...]
621 """Input collections provided directly by the user (`tuple` [ `str` ]).
622 """