Coverage for tests/test_cmdLineFwk.py: 19%
Shortcuts on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
Shortcuts on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
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# (https://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 program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
22"""Simple unit test for cmdLineFwk module.
23"""
25import contextlib
26import copy
27import logging
28import os
29import pickle
30import re
31import shutil
32import tempfile
33import unittest
34from dataclasses import dataclass
35from types import SimpleNamespace
36from typing import NamedTuple
38import click
39import lsst.pex.config as pexConfig
40import lsst.pipe.base.connectionTypes as cT
41import lsst.utils.tests
42from lsst.ctrl.mpexec import CmdLineFwk, MPGraphExecutorError
43from lsst.ctrl.mpexec.cli.opt import run_options
44from lsst.ctrl.mpexec.cli.utils import (
45 _ACTION_ADD_INSTRUMENT,
46 _ACTION_ADD_TASK,
47 _ACTION_CONFIG,
48 _ACTION_CONFIG_FILE,
49 PipetaskCommand,
50)
51from lsst.daf.butler import Config, DataCoordinate, DatasetRef, DimensionUniverse, Quantum, Registry
52from lsst.daf.butler.core.datasets.type import DatasetType
53from lsst.daf.butler.registry import RegistryConfig
54from lsst.obs.base import Instrument
55from lsst.pipe.base import Pipeline, PipelineTaskConfig, PipelineTaskConnections, QuantumGraph, TaskDef
56from lsst.pipe.base.graphBuilder import DatasetQueryConstraintVariant as DQCVariant
57from lsst.pipe.base.tests.simpleQGraph import (
58 AddTask,
59 AddTaskFactoryMock,
60 makeSimpleButler,
61 makeSimplePipeline,
62 makeSimpleQGraph,
63 populateButler,
64)
65from lsst.utils.tests import temporaryDirectory
67logging.basicConfig(level=getattr(logging, os.environ.get("UNIT_TEST_LOGGING_LEVEL", "INFO"), logging.INFO))
69# Have to monkey-patch Instrument.fromName() to not retrieve non-existing
70# instrument from registry, these tests can run fine without actual instrument
71# and implementing full mock for Instrument is too complicated.
72Instrument.fromName = lambda name, reg: None 72 ↛ exitline 72 didn't run the lambda on line 72
75@contextlib.contextmanager
76def makeTmpFile(contents=None, suffix=None):
77 """Context manager for generating temporary file name.
79 Temporary file is deleted on exiting context.
81 Parameters
82 ----------
83 contents : `bytes`
84 Data to write into a file.
85 """
86 fd, tmpname = tempfile.mkstemp(suffix=suffix)
87 if contents:
88 os.write(fd, contents)
89 os.close(fd)
90 yield tmpname
91 with contextlib.suppress(OSError):
92 os.remove(tmpname)
95@contextlib.contextmanager
96def makeSQLiteRegistry(create=True):
97 """Context manager to create new empty registry database.
99 Yields
100 ------
101 config : `RegistryConfig`
102 Registry configuration for initialized registry database.
103 """
104 with temporaryDirectory() as tmpdir:
105 uri = f"sqlite:///{tmpdir}/gen3.sqlite"
106 config = RegistryConfig()
107 config["db"] = uri
108 if create:
109 Registry.createFromConfig(config)
110 yield config
113class SimpleConnections(PipelineTaskConnections, dimensions=(), defaultTemplates={"template": "simple"}):
114 schema = cT.InitInput(doc="Schema", name="{template}schema", storageClass="SourceCatalog")
117class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections):
118 field = pexConfig.Field(dtype=str, doc="arbitrary string")
120 def setDefaults(self):
121 PipelineTaskConfig.setDefaults(self)
124def _makeArgs(registryConfig=None, **kwargs):
125 """Return parsed command line arguments.
127 By default butler_config is set to `Config` populated with some defaults,
128 it can be overridden completely by keyword argument.
130 Parameters
131 ----------
132 cmd : `str`, optional
133 Produce arguments for this pipetask command.
134 registryConfig : `RegistryConfig`, optional
135 Override for registry configuration.
136 **kwargs
137 Overrides for other arguments.
138 """
139 # Use a mock to get the default value of arguments to 'run'.
141 mock = unittest.mock.Mock()
143 @click.command(cls=PipetaskCommand)
144 @run_options()
145 def fake_run(ctx, **kwargs):
146 """Fake "pipetask run" command for gathering input arguments.
148 The arguments & options should always match the arguments & options in
149 the "real" command function `lsst.ctrl.mpexec.cli.cmd.run`.
150 """
151 mock(**kwargs)
153 runner = click.testing.CliRunner()
154 # --butler-config is the only required option
155 result = runner.invoke(fake_run, "--butler-config /")
156 if result.exit_code != 0:
157 raise RuntimeError(f"Failure getting default args from 'fake_run': {result}")
158 mock.assert_called_once()
159 args = mock.call_args[1]
160 args["enableLsstDebug"] = args.pop("debug")
161 args["execution_butler_location"] = args.pop("save_execution_butler")
162 if "pipeline_actions" not in args:
163 args["pipeline_actions"] = []
164 if "mock_configs" not in args:
165 args["mock_configs"] = []
166 args = SimpleNamespace(**args)
168 # override butler_config with our defaults
169 if "butler_config" not in kwargs:
170 args.butler_config = Config()
171 if registryConfig:
172 args.butler_config["registry"] = registryConfig
173 # The default datastore has a relocatable root, so we need to specify
174 # some root here for it to use
175 args.butler_config.configFile = "."
177 # override arguments from keyword parameters
178 for key, value in kwargs.items():
179 setattr(args, key, value)
180 args.dataset_query_constraint = DQCVariant.fromExpression(args.dataset_query_constraint)
181 return args
184class FakeDSType(NamedTuple):
185 name: str
188@dataclass(frozen=True)
189class FakeDSRef:
190 datasetType: str
191 dataId: tuple
193 def isComponent(self):
194 return False
197# Task class name used by tests, needs to be importable
198_TASK_CLASS = "lsst.pipe.base.tests.simpleQGraph.AddTask"
201def _makeQGraph():
202 """Make a trivial QuantumGraph with one quantum.
204 The only thing that we need to do with this quantum graph is to pickle
205 it, the quanta in this graph are not usable for anything else.
207 Returns
208 -------
209 qgraph : `~lsst.pipe.base.QuantumGraph`
210 """
211 config = Config(
212 {
213 "version": 1,
214 "skypix": {
215 "common": "htm7",
216 "htm": {
217 "class": "lsst.sphgeom.HtmPixelization",
218 "max_level": 24,
219 },
220 },
221 "elements": {
222 "A": {
223 "keys": [
224 {
225 "name": "id",
226 "type": "int",
227 }
228 ],
229 "storage": {
230 "cls": "lsst.daf.butler.registry.dimensions.table.TableDimensionRecordStorage",
231 },
232 },
233 "B": {
234 "keys": [
235 {
236 "name": "id",
237 "type": "int",
238 }
239 ],
240 "storage": {
241 "cls": "lsst.daf.butler.registry.dimensions.table.TableDimensionRecordStorage",
242 },
243 },
244 },
245 "packers": {},
246 }
247 )
248 universe = DimensionUniverse(config=config)
249 fakeDSType = DatasetType("A", tuple(), storageClass="ExposureF", universe=universe)
250 taskDef = TaskDef(taskName=_TASK_CLASS, config=AddTask.ConfigClass(), taskClass=AddTask)
251 quanta = [
252 Quantum(
253 taskName=_TASK_CLASS,
254 inputs={
255 fakeDSType: [
256 DatasetRef(fakeDSType, DataCoordinate.standardize({"A": 1, "B": 2}, universe=universe))
257 ]
258 },
259 )
260 ] # type: ignore
261 qgraph = QuantumGraph({taskDef: set(quanta)})
262 return qgraph
265class CmdLineFwkTestCase(unittest.TestCase):
266 """A test case for CmdLineFwk"""
268 def testMakePipeline(self):
269 """Tests for CmdLineFwk.makePipeline method"""
270 fwk = CmdLineFwk()
272 # make empty pipeline
273 args = _makeArgs()
274 pipeline = fwk.makePipeline(args)
275 self.assertIsInstance(pipeline, Pipeline)
276 self.assertEqual(len(pipeline), 0)
278 # few tests with serialization
279 with makeTmpFile() as tmpname:
280 # make empty pipeline and store it in a file
281 args = _makeArgs(save_pipeline=tmpname)
282 pipeline = fwk.makePipeline(args)
283 self.assertIsInstance(pipeline, Pipeline)
285 # read pipeline from a file
286 args = _makeArgs(pipeline=tmpname)
287 pipeline = fwk.makePipeline(args)
288 self.assertIsInstance(pipeline, Pipeline)
289 self.assertEqual(len(pipeline), 0)
291 # single task pipeline, task name can be anything here
292 actions = [_ACTION_ADD_TASK("TaskOne:task1")]
293 args = _makeArgs(pipeline_actions=actions)
294 pipeline = fwk.makePipeline(args)
295 self.assertIsInstance(pipeline, Pipeline)
296 self.assertEqual(len(pipeline), 1)
298 # many task pipeline
299 actions = [
300 _ACTION_ADD_TASK("TaskOne:task1a"),
301 _ACTION_ADD_TASK("TaskTwo:task2"),
302 _ACTION_ADD_TASK("TaskOne:task1b"),
303 ]
304 args = _makeArgs(pipeline_actions=actions)
305 pipeline = fwk.makePipeline(args)
306 self.assertIsInstance(pipeline, Pipeline)
307 self.assertEqual(len(pipeline), 3)
309 # single task pipeline with config overrides, need real task class
310 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG("task:addend=100")]
311 args = _makeArgs(pipeline_actions=actions)
312 pipeline = fwk.makePipeline(args)
313 taskDefs = list(pipeline.toExpandedPipeline())
314 self.assertEqual(len(taskDefs), 1)
315 self.assertEqual(taskDefs[0].config.addend, 100)
317 overrides = b"config.addend = 1000\n"
318 with makeTmpFile(overrides) as tmpname:
319 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG_FILE("task:" + tmpname)]
320 args = _makeArgs(pipeline_actions=actions)
321 pipeline = fwk.makePipeline(args)
322 taskDefs = list(pipeline.toExpandedPipeline())
323 self.assertEqual(len(taskDefs), 1)
324 self.assertEqual(taskDefs[0].config.addend, 1000)
326 # Check --instrument option, for now it only checks that it does not
327 # crash.
328 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_ADD_INSTRUMENT("Instrument")]
329 args = _makeArgs(pipeline_actions=actions)
330 pipeline = fwk.makePipeline(args)
332 def testMakeGraphFromSave(self):
333 """Tests for CmdLineFwk.makeGraph method.
335 Only most trivial case is tested that does not do actual graph
336 building.
337 """
338 fwk = CmdLineFwk()
340 with makeTmpFile(suffix=".qgraph") as tmpname, makeSQLiteRegistry() as registryConfig:
342 # make non-empty graph and store it in a file
343 qgraph = _makeQGraph()
344 with open(tmpname, "wb") as saveFile:
345 qgraph.save(saveFile)
346 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
347 qgraph = fwk.makeGraph(None, args)
348 self.assertIsInstance(qgraph, QuantumGraph)
349 self.assertEqual(len(qgraph), 1)
351 # will fail if graph id does not match
352 args = _makeArgs(
353 qgraph=tmpname,
354 qgraph_id="R2-D2 is that you?",
355 registryConfig=registryConfig,
356 execution_butler_location=None,
357 )
358 with self.assertRaisesRegex(ValueError, "graphID does not match"):
359 fwk.makeGraph(None, args)
361 # save with wrong object type
362 with open(tmpname, "wb") as saveFile:
363 pickle.dump({}, saveFile)
364 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
365 with self.assertRaises(ValueError):
366 fwk.makeGraph(None, args)
368 # reading empty graph from pickle should work but makeGraph()
369 # will return None and make a warning
370 qgraph = QuantumGraph(dict())
371 with open(tmpname, "wb") as saveFile:
372 qgraph.save(saveFile)
373 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
374 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"):
375 # this also tests that warning is generated for empty graph
376 qgraph = fwk.makeGraph(None, args)
377 self.assertIs(qgraph, None)
379 def testShowPipeline(self):
380 """Test for --show options for pipeline."""
381 fwk = CmdLineFwk()
383 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG("task:addend=100")]
384 args = _makeArgs(pipeline_actions=actions)
385 pipeline = fwk.makePipeline(args)
387 args.show = ["pipeline"]
388 fwk.showInfo(args, pipeline)
389 args.show = ["config"]
390 fwk.showInfo(args, pipeline)
391 args.show = ["history=task::addend"]
392 fwk.showInfo(args, pipeline)
393 args.show = ["tasks"]
394 fwk.showInfo(args, pipeline)
397class CmdLineFwkTestCaseWithButler(unittest.TestCase):
398 """A test case for CmdLineFwk"""
400 def setUp(self):
401 super().setUpClass()
402 self.root = tempfile.mkdtemp()
403 self.nQuanta = 5
404 self.pipeline = makeSimplePipeline(nQuanta=self.nQuanta)
406 def tearDown(self):
407 shutil.rmtree(self.root, ignore_errors=True)
408 super().tearDownClass()
410 def testSimpleQGraph(self):
411 """Test successfull execution of trivial quantum graph."""
412 args = _makeArgs(butler_config=self.root, input="test", output="output")
413 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
414 populateButler(self.pipeline, butler)
416 fwk = CmdLineFwk()
417 taskFactory = AddTaskFactoryMock()
419 qgraph = fwk.makeGraph(self.pipeline, args)
420 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
421 self.assertEqual(len(qgraph), self.nQuanta)
423 # run whole thing
424 fwk.runPipeline(qgraph, taskFactory, args)
425 self.assertEqual(taskFactory.countExec, self.nQuanta)
427 def testSimpleQGraphNoSkipExisting_inputs(self):
428 """Test for case when output data for one task already appears in
429 _input_ collection, but no ``--extend-run`` or ``-skip-existing``
430 option is present.
431 """
432 args = _makeArgs(
433 butler_config=self.root,
434 input="test",
435 output="output",
436 )
437 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
438 populateButler(
439 self.pipeline,
440 butler,
441 datasetTypes={
442 args.input: [
443 "add_dataset0",
444 "add_dataset1",
445 "add2_dataset1",
446 "add_init_output1",
447 "task0_config",
448 "task0_metadata",
449 "task0_log",
450 ]
451 },
452 )
454 fwk = CmdLineFwk()
455 taskFactory = AddTaskFactoryMock()
457 qgraph = fwk.makeGraph(self.pipeline, args)
458 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
459 # With current implementation graph has all nQuanta quanta, but when
460 # executing one quantum is skipped.
461 self.assertEqual(len(qgraph), self.nQuanta)
463 # run whole thing
464 fwk.runPipeline(qgraph, taskFactory, args)
465 self.assertEqual(taskFactory.countExec, self.nQuanta)
467 def testSimpleQGraphSkipExisting_inputs(self):
468 """Test for ``--skip-existing`` with output data for one task already
469 appears in _input_ collection. No ``--extend-run`` option is needed
470 for this case.
471 """
472 args = _makeArgs(
473 butler_config=self.root,
474 input="test",
475 output="output",
476 skip_existing_in=("test",),
477 )
478 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
479 populateButler(
480 self.pipeline,
481 butler,
482 datasetTypes={
483 args.input: [
484 "add_dataset0",
485 "add_dataset1",
486 "add2_dataset1",
487 "add_init_output1",
488 "task0_config",
489 "task0_metadata",
490 "task0_log",
491 ]
492 },
493 )
495 fwk = CmdLineFwk()
496 taskFactory = AddTaskFactoryMock()
498 qgraph = fwk.makeGraph(self.pipeline, args)
499 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
500 self.assertEqual(len(qgraph), self.nQuanta - 1)
502 # run whole thing
503 fwk.runPipeline(qgraph, taskFactory, args)
504 self.assertEqual(taskFactory.countExec, self.nQuanta - 1)
506 def testSimpleQGraphSkipExisting_outputs(self):
507 """Test for ``--skip-existing`` with output data for one task already
508 appears in _output_ collection. The ``--extend-run`` option is needed
509 for this case.
510 """
511 args = _makeArgs(
512 butler_config=self.root,
513 input="test",
514 output_run="output/run",
515 skip_existing_in=("output/run",),
516 )
517 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
518 populateButler(
519 self.pipeline,
520 butler,
521 datasetTypes={
522 args.input: ["add_dataset0"],
523 args.output_run: [
524 "add_dataset1",
525 "add2_dataset1",
526 "add_init_output1",
527 "task0_metadata",
528 "task0_log",
529 ],
530 },
531 )
533 fwk = CmdLineFwk()
534 taskFactory = AddTaskFactoryMock()
536 # fails without --extend-run
537 with self.assertRaisesRegex(ValueError, "--extend-run was not given"):
538 qgraph = fwk.makeGraph(self.pipeline, args)
540 # retry with --extend-run
541 args.extend_run = True
542 qgraph = fwk.makeGraph(self.pipeline, args)
544 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
545 # Graph does not include quantum for first task
546 self.assertEqual(len(qgraph), self.nQuanta - 1)
548 # run whole thing
549 fwk.runPipeline(qgraph, taskFactory, args)
550 self.assertEqual(taskFactory.countExec, self.nQuanta - 1)
552 def testSimpleQGraphOutputsFail(self):
553 """Test continuing execution of trivial quantum graph with partial
554 outputs.
555 """
556 args = _makeArgs(butler_config=self.root, input="test", output="output")
557 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
558 populateButler(self.pipeline, butler)
560 fwk = CmdLineFwk()
561 taskFactory = AddTaskFactoryMock(stopAt=3)
563 qgraph = fwk.makeGraph(self.pipeline, args)
564 self.assertEqual(len(qgraph), self.nQuanta)
566 # run first three quanta
567 with self.assertRaises(MPGraphExecutorError):
568 fwk.runPipeline(qgraph, taskFactory, args)
569 self.assertEqual(taskFactory.countExec, 3)
571 butler.registry.refresh()
573 # drop one of the two outputs from one task
574 ref1 = butler.registry.findDataset(
575 "add2_dataset2", collections=args.output, instrument="INSTR", detector=0
576 )
577 self.assertIsNotNone(ref1)
578 # also drop the metadata output
579 ref2 = butler.registry.findDataset(
580 "task1_metadata", collections=args.output, instrument="INSTR", detector=0
581 )
582 self.assertIsNotNone(ref2)
583 butler.pruneDatasets([ref1, ref2], disassociate=True, unstore=True, purge=True)
585 taskFactory.stopAt = -1
586 args.skip_existing_in = (args.output,)
587 args.extend_run = True
588 args.no_versions = True
589 with self.assertRaises(MPGraphExecutorError):
590 fwk.runPipeline(qgraph, taskFactory, args)
592 def testSimpleQGraphClobberOutputs(self):
593 """Test continuing execution of trivial quantum graph with
594 --clobber-outputs.
595 """
596 args = _makeArgs(butler_config=self.root, input="test", output="output")
597 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
598 populateButler(self.pipeline, butler)
600 fwk = CmdLineFwk()
601 taskFactory = AddTaskFactoryMock(stopAt=3)
603 qgraph = fwk.makeGraph(self.pipeline, args)
605 # should have one task and number of quanta
606 self.assertEqual(len(qgraph), self.nQuanta)
608 # run first three quanta
609 with self.assertRaises(MPGraphExecutorError):
610 fwk.runPipeline(qgraph, taskFactory, args)
611 self.assertEqual(taskFactory.countExec, 3)
613 butler.registry.refresh()
615 # drop one of the two outputs from one task
616 ref1 = butler.registry.findDataset(
617 "add2_dataset2", collections=args.output, dataId=dict(instrument="INSTR", detector=0)
618 )
619 self.assertIsNotNone(ref1)
620 # also drop the metadata output
621 ref2 = butler.registry.findDataset(
622 "task1_metadata", collections=args.output, dataId=dict(instrument="INSTR", detector=0)
623 )
624 self.assertIsNotNone(ref2)
625 butler.pruneDatasets([ref1, ref2], disassociate=True, unstore=True, purge=True)
627 taskFactory.stopAt = -1
628 args.skip_existing = True
629 args.extend_run = True
630 args.clobber_outputs = True
631 args.no_versions = True
632 fwk.runPipeline(qgraph, taskFactory, args)
633 # number of executed quanta is incremented
634 self.assertEqual(taskFactory.countExec, self.nQuanta + 1)
636 def testSimpleQGraphReplaceRun(self):
637 """Test repeated execution of trivial quantum graph with
638 --replace-run.
639 """
640 args = _makeArgs(butler_config=self.root, input="test", output="output", output_run="output/run1")
641 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
642 populateButler(self.pipeline, butler)
644 fwk = CmdLineFwk()
645 taskFactory = AddTaskFactoryMock()
647 qgraph = fwk.makeGraph(self.pipeline, args)
649 # should have one task and number of quanta
650 self.assertEqual(len(qgraph), self.nQuanta)
652 # deep copy is needed because quanta are updated in place
653 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
654 self.assertEqual(taskFactory.countExec, self.nQuanta)
656 # need to refresh collections explicitly (or make new butler/registry)
657 butler.registry.refresh()
658 collections = set(butler.registry.queryCollections(...))
659 self.assertEqual(collections, {"test", "output", "output/run1"})
661 # number of datasets written by pipeline:
662 # - nQuanta of init_outputs
663 # - nQuanta of configs
664 # - packages (single dataset)
665 # - nQuanta * two output datasets
666 # - nQuanta of metadata
667 # - nQuanta of log output
668 n_outputs = self.nQuanta * 6 + 1
669 refs = butler.registry.queryDatasets(..., collections="output/run1")
670 self.assertEqual(len(list(refs)), n_outputs)
672 # re-run with --replace-run (--inputs is ignored, as long as it hasn't
673 # changed)
674 args.replace_run = True
675 args.output_run = "output/run2"
676 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
678 butler.registry.refresh()
679 collections = set(butler.registry.queryCollections(...))
680 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2"})
682 # new output collection
683 refs = butler.registry.queryDatasets(..., collections="output/run2")
684 self.assertEqual(len(list(refs)), n_outputs)
686 # old output collection is still there
687 refs = butler.registry.queryDatasets(..., collections="output/run1")
688 self.assertEqual(len(list(refs)), n_outputs)
690 # re-run with --replace-run and --prune-replaced=unstore
691 args.replace_run = True
692 args.prune_replaced = "unstore"
693 args.output_run = "output/run3"
694 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
696 butler.registry.refresh()
697 collections = set(butler.registry.queryCollections(...))
698 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run3"})
700 # new output collection
701 refs = butler.registry.queryDatasets(..., collections="output/run3")
702 self.assertEqual(len(list(refs)), n_outputs)
704 # old output collection is still there, and it has all datasets but
705 # non-InitOutputs are not in datastore
706 refs = butler.registry.queryDatasets(..., collections="output/run2")
707 refs = list(refs)
708 self.assertEqual(len(refs), n_outputs)
709 initOutNameRe = re.compile("packages|task.*_config|add_init_output.*")
710 for ref in refs:
711 if initOutNameRe.fullmatch(ref.datasetType.name):
712 butler.get(ref, collections="output/run2")
713 else:
714 with self.assertRaises(FileNotFoundError):
715 butler.get(ref, collections="output/run2")
717 # re-run with --replace-run and --prune-replaced=purge
718 # This time also remove --input; passing the same inputs that we
719 # started with and not passing inputs at all should be equivalent.
720 args.input = None
721 args.replace_run = True
722 args.prune_replaced = "purge"
723 args.output_run = "output/run4"
724 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
726 butler.registry.refresh()
727 collections = set(butler.registry.queryCollections(...))
728 # output/run3 should disappear now
729 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
731 # new output collection
732 refs = butler.registry.queryDatasets(..., collections="output/run4")
733 self.assertEqual(len(list(refs)), n_outputs)
735 # Trying to run again with inputs that aren't exactly what we started
736 # with is an error, and the kind that should not modify the data repo.
737 with self.assertRaises(ValueError):
738 args.input = ["test", "output/run2"]
739 args.prune_replaced = None
740 args.replace_run = True
741 args.output_run = "output/run5"
742 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
743 butler.registry.refresh()
744 collections = set(butler.registry.queryCollections(...))
745 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
746 with self.assertRaises(ValueError):
747 args.input = ["output/run2", "test"]
748 args.prune_replaced = None
749 args.replace_run = True
750 args.output_run = "output/run6"
751 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
752 butler.registry.refresh()
753 collections = set(butler.registry.queryCollections(...))
754 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
756 def testMockTask(self):
757 """Test --mock option."""
758 args = _makeArgs(
759 butler_config=self.root, input="test", output="output", mock=True, register_dataset_types=True
760 )
761 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
762 populateButler(self.pipeline, butler)
764 fwk = CmdLineFwk()
765 taskFactory = AddTaskFactoryMock()
767 qgraph = fwk.makeGraph(self.pipeline, args)
768 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
769 self.assertEqual(len(qgraph), self.nQuanta)
771 # run whole thing
772 fwk.runPipeline(qgraph, taskFactory, args)
773 # None of the actual tasks is executed
774 self.assertEqual(taskFactory.countExec, 0)
776 # check dataset types
777 butler.registry.refresh()
778 datasetTypes = list(butler.registry.queryDatasetTypes(re.compile("^_mock_.*")))
779 self.assertEqual(len(datasetTypes), self.nQuanta * 2)
781 def testMockTaskFailure(self):
782 """Test --mock option and configure one of the tasks to fail."""
783 args = _makeArgs(
784 butler_config=self.root,
785 input="test",
786 output="output",
787 mock=True,
788 register_dataset_types=True,
789 mock_configs=[
790 _ACTION_CONFIG("task3-mock:failCondition='detector = 0'"),
791 ],
792 fail_fast=True,
793 )
794 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
795 populateButler(self.pipeline, butler)
797 fwk = CmdLineFwk()
798 taskFactory = AddTaskFactoryMock()
800 qgraph = fwk.makeGraph(self.pipeline, args)
801 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
802 self.assertEqual(len(qgraph), self.nQuanta)
804 with self.assertRaises(MPGraphExecutorError) as cm:
805 fwk.runPipeline(qgraph, taskFactory, args)
807 self.assertIsNotNone(cm.exception.__cause__)
808 self.assertRegex(str(cm.exception.__cause__), "Simulated failure: task=task3")
810 def testSubgraph(self):
811 """Test successfull execution of trivial quantum graph."""
812 args = _makeArgs(butler_config=self.root, input="test", output="output")
813 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
814 populateButler(self.pipeline, butler)
816 fwk = CmdLineFwk()
817 qgraph = fwk.makeGraph(self.pipeline, args)
819 # Select first two nodes for execution. This depends on node ordering
820 # which I assume is the same as execution order.
821 nNodes = 2
822 nodeIds = [node.nodeId for node in qgraph]
823 nodeIds = nodeIds[:nNodes]
825 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
826 self.assertEqual(len(qgraph), self.nQuanta)
828 with makeTmpFile(suffix=".qgraph") as tmpname, makeSQLiteRegistry() as registryConfig:
829 with open(tmpname, "wb") as saveFile:
830 qgraph.save(saveFile)
832 args = _makeArgs(
833 qgraph=tmpname,
834 qgraph_node_id=nodeIds,
835 registryConfig=registryConfig,
836 execution_butler_location=None,
837 )
838 fwk = CmdLineFwk()
840 # load graph, should only read a subset
841 qgraph = fwk.makeGraph(pipeline=None, args=args)
842 self.assertEqual(len(qgraph), nNodes)
844 def testShowGraph(self):
845 """Test for --show options for quantum graph."""
846 fwk = CmdLineFwk()
848 nQuanta = 2
849 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
851 args = _makeArgs(show=["graph"])
852 fwk.showInfo(args, pipeline=None, graph=qgraph)
854 def testShowGraphWorkflow(self):
855 fwk = CmdLineFwk()
857 nQuanta = 2
858 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
860 args = _makeArgs(show=["workflow"])
861 fwk.showInfo(args, pipeline=None, graph=qgraph)
863 # TODO: cannot test "uri" option presently, it instanciates
864 # butler from command line options and there is no way to pass butler
865 # mock to that code.
868class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
869 pass
872def setup_module(module):
873 lsst.utils.tests.init()
876if __name__ == "__main__": 876 ↛ 877line 876 didn't jump to line 877, because the condition on line 876 was never true
877 lsst.utils.tests.init()
878 unittest.main()