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.pipe.base import (
55 Instrument,
56 Pipeline,
57 PipelineTaskConfig,
58 PipelineTaskConnections,
59 QuantumGraph,
60 TaskDef,
61)
62from lsst.pipe.base.graphBuilder import DatasetQueryConstraintVariant as DQCVariant
63from lsst.pipe.base.tests.simpleQGraph import (
64 AddTask,
65 AddTaskFactoryMock,
66 makeSimpleButler,
67 makeSimplePipeline,
68 makeSimpleQGraph,
69 populateButler,
70)
71from lsst.utils.tests import temporaryDirectory
73logging.basicConfig(level=getattr(logging, os.environ.get("UNIT_TEST_LOGGING_LEVEL", "INFO"), logging.INFO))
75# Have to monkey-patch Instrument.fromName() to not retrieve non-existing
76# instrument from registry, these tests can run fine without actual instrument
77# and implementing full mock for Instrument is too complicated.
78Instrument.fromName = lambda name, reg: None 78 ↛ exitline 78 didn't run the lambda on line 78
81@contextlib.contextmanager
82def makeTmpFile(contents=None, suffix=None):
83 """Context manager for generating temporary file name.
85 Temporary file is deleted on exiting context.
87 Parameters
88 ----------
89 contents : `bytes`
90 Data to write into a file.
91 """
92 fd, tmpname = tempfile.mkstemp(suffix=suffix)
93 if contents:
94 os.write(fd, contents)
95 os.close(fd)
96 yield tmpname
97 with contextlib.suppress(OSError):
98 os.remove(tmpname)
101@contextlib.contextmanager
102def makeSQLiteRegistry(create=True):
103 """Context manager to create new empty registry database.
105 Yields
106 ------
107 config : `RegistryConfig`
108 Registry configuration for initialized registry database.
109 """
110 with temporaryDirectory() as tmpdir:
111 uri = f"sqlite:///{tmpdir}/gen3.sqlite"
112 config = RegistryConfig()
113 config["db"] = uri
114 if create:
115 Registry.createFromConfig(config)
116 yield config
119class SimpleConnections(PipelineTaskConnections, dimensions=(), defaultTemplates={"template": "simple"}):
120 schema = cT.InitInput(doc="Schema", name="{template}schema", storageClass="SourceCatalog")
123class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections):
124 field = pexConfig.Field(dtype=str, doc="arbitrary string")
126 def setDefaults(self):
127 PipelineTaskConfig.setDefaults(self)
130def _makeArgs(registryConfig=None, **kwargs):
131 """Return parsed command line arguments.
133 By default butler_config is set to `Config` populated with some defaults,
134 it can be overridden completely by keyword argument.
136 Parameters
137 ----------
138 cmd : `str`, optional
139 Produce arguments for this pipetask command.
140 registryConfig : `RegistryConfig`, optional
141 Override for registry configuration.
142 **kwargs
143 Overrides for other arguments.
144 """
145 # Use a mock to get the default value of arguments to 'run'.
147 mock = unittest.mock.Mock()
149 @click.command(cls=PipetaskCommand)
150 @run_options()
151 def fake_run(ctx, **kwargs):
152 """Fake "pipetask run" command for gathering input arguments.
154 The arguments & options should always match the arguments & options in
155 the "real" command function `lsst.ctrl.mpexec.cli.cmd.run`.
156 """
157 mock(**kwargs)
159 runner = click.testing.CliRunner()
160 # --butler-config is the only required option
161 result = runner.invoke(fake_run, "--butler-config /")
162 if result.exit_code != 0:
163 raise RuntimeError(f"Failure getting default args from 'fake_run': {result}")
164 mock.assert_called_once()
165 args = mock.call_args[1]
166 args["enableLsstDebug"] = args.pop("debug")
167 args["execution_butler_location"] = args.pop("save_execution_butler")
168 if "pipeline_actions" not in args:
169 args["pipeline_actions"] = []
170 if "mock_configs" not in args:
171 args["mock_configs"] = []
172 args = SimpleNamespace(**args)
174 # override butler_config with our defaults
175 if "butler_config" not in kwargs:
176 args.butler_config = Config()
177 if registryConfig:
178 args.butler_config["registry"] = registryConfig
179 # The default datastore has a relocatable root, so we need to specify
180 # some root here for it to use
181 args.butler_config.configFile = "."
183 # override arguments from keyword parameters
184 for key, value in kwargs.items():
185 setattr(args, key, value)
186 args.dataset_query_constraint = DQCVariant.fromExpression(args.dataset_query_constraint)
187 return args
190class FakeDSType(NamedTuple):
191 name: str
194@dataclass(frozen=True)
195class FakeDSRef:
196 datasetType: str
197 dataId: tuple
199 def isComponent(self):
200 return False
203# Task class name used by tests, needs to be importable
204_TASK_CLASS = "lsst.pipe.base.tests.simpleQGraph.AddTask"
207def _makeQGraph():
208 """Make a trivial QuantumGraph with one quantum.
210 The only thing that we need to do with this quantum graph is to pickle
211 it, the quanta in this graph are not usable for anything else.
213 Returns
214 -------
215 qgraph : `~lsst.pipe.base.QuantumGraph`
216 """
217 config = Config(
218 {
219 "version": 1,
220 "skypix": {
221 "common": "htm7",
222 "htm": {
223 "class": "lsst.sphgeom.HtmPixelization",
224 "max_level": 24,
225 },
226 },
227 "elements": {
228 "A": {
229 "keys": [
230 {
231 "name": "id",
232 "type": "int",
233 }
234 ],
235 "storage": {
236 "cls": "lsst.daf.butler.registry.dimensions.table.TableDimensionRecordStorage",
237 },
238 },
239 "B": {
240 "keys": [
241 {
242 "name": "id",
243 "type": "int",
244 }
245 ],
246 "storage": {
247 "cls": "lsst.daf.butler.registry.dimensions.table.TableDimensionRecordStorage",
248 },
249 },
250 },
251 "packers": {},
252 }
253 )
254 universe = DimensionUniverse(config=config)
255 fakeDSType = DatasetType("A", tuple(), storageClass="ExposureF", universe=universe)
256 taskDef = TaskDef(taskName=_TASK_CLASS, config=AddTask.ConfigClass(), taskClass=AddTask)
257 quanta = [
258 Quantum(
259 taskName=_TASK_CLASS,
260 inputs={
261 fakeDSType: [
262 DatasetRef(fakeDSType, DataCoordinate.standardize({"A": 1, "B": 2}, universe=universe))
263 ]
264 },
265 )
266 ] # type: ignore
267 qgraph = QuantumGraph({taskDef: set(quanta)})
268 return qgraph
271class CmdLineFwkTestCase(unittest.TestCase):
272 """A test case for CmdLineFwk"""
274 def testMakePipeline(self):
275 """Tests for CmdLineFwk.makePipeline method"""
276 fwk = CmdLineFwk()
278 # make empty pipeline
279 args = _makeArgs()
280 pipeline = fwk.makePipeline(args)
281 self.assertIsInstance(pipeline, Pipeline)
282 self.assertEqual(len(pipeline), 0)
284 # few tests with serialization
285 with makeTmpFile() as tmpname:
286 # make empty pipeline and store it in a file
287 args = _makeArgs(save_pipeline=tmpname)
288 pipeline = fwk.makePipeline(args)
289 self.assertIsInstance(pipeline, Pipeline)
291 # read pipeline from a file
292 args = _makeArgs(pipeline=tmpname)
293 pipeline = fwk.makePipeline(args)
294 self.assertIsInstance(pipeline, Pipeline)
295 self.assertEqual(len(pipeline), 0)
297 # single task pipeline, task name can be anything here
298 actions = [_ACTION_ADD_TASK("TaskOne:task1")]
299 args = _makeArgs(pipeline_actions=actions)
300 pipeline = fwk.makePipeline(args)
301 self.assertIsInstance(pipeline, Pipeline)
302 self.assertEqual(len(pipeline), 1)
304 # many task pipeline
305 actions = [
306 _ACTION_ADD_TASK("TaskOne:task1a"),
307 _ACTION_ADD_TASK("TaskTwo:task2"),
308 _ACTION_ADD_TASK("TaskOne:task1b"),
309 ]
310 args = _makeArgs(pipeline_actions=actions)
311 pipeline = fwk.makePipeline(args)
312 self.assertIsInstance(pipeline, Pipeline)
313 self.assertEqual(len(pipeline), 3)
315 # single task pipeline with config overrides, need real task class
316 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG("task:addend=100")]
317 args = _makeArgs(pipeline_actions=actions)
318 pipeline = fwk.makePipeline(args)
319 taskDefs = list(pipeline.toExpandedPipeline())
320 self.assertEqual(len(taskDefs), 1)
321 self.assertEqual(taskDefs[0].config.addend, 100)
323 overrides = b"config.addend = 1000\n"
324 with makeTmpFile(overrides) as tmpname:
325 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG_FILE("task:" + tmpname)]
326 args = _makeArgs(pipeline_actions=actions)
327 pipeline = fwk.makePipeline(args)
328 taskDefs = list(pipeline.toExpandedPipeline())
329 self.assertEqual(len(taskDefs), 1)
330 self.assertEqual(taskDefs[0].config.addend, 1000)
332 # Check --instrument option, for now it only checks that it does not
333 # crash.
334 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_ADD_INSTRUMENT("Instrument")]
335 args = _makeArgs(pipeline_actions=actions)
336 pipeline = fwk.makePipeline(args)
338 def testMakeGraphFromSave(self):
339 """Tests for CmdLineFwk.makeGraph method.
341 Only most trivial case is tested that does not do actual graph
342 building.
343 """
344 fwk = CmdLineFwk()
346 with makeTmpFile(suffix=".qgraph") as tmpname, makeSQLiteRegistry() as registryConfig:
348 # make non-empty graph and store it in a file
349 qgraph = _makeQGraph()
350 with open(tmpname, "wb") as saveFile:
351 qgraph.save(saveFile)
352 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
353 qgraph = fwk.makeGraph(None, args)
354 self.assertIsInstance(qgraph, QuantumGraph)
355 self.assertEqual(len(qgraph), 1)
357 # will fail if graph id does not match
358 args = _makeArgs(
359 qgraph=tmpname,
360 qgraph_id="R2-D2 is that you?",
361 registryConfig=registryConfig,
362 execution_butler_location=None,
363 )
364 with self.assertRaisesRegex(ValueError, "graphID does not match"):
365 fwk.makeGraph(None, args)
367 # save with wrong object type
368 with open(tmpname, "wb") as saveFile:
369 pickle.dump({}, saveFile)
370 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
371 with self.assertRaises(ValueError):
372 fwk.makeGraph(None, args)
374 # reading empty graph from pickle should work but makeGraph()
375 # will return None and make a warning
376 qgraph = QuantumGraph(dict())
377 with open(tmpname, "wb") as saveFile:
378 qgraph.save(saveFile)
379 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
380 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"):
381 # this also tests that warning is generated for empty graph
382 qgraph = fwk.makeGraph(None, args)
383 self.assertIs(qgraph, None)
385 def testShowPipeline(self):
386 """Test for --show options for pipeline."""
387 fwk = CmdLineFwk()
389 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG("task:addend=100")]
390 args = _makeArgs(pipeline_actions=actions)
391 pipeline = fwk.makePipeline(args)
393 args.show = ["pipeline"]
394 fwk.showInfo(args, pipeline)
395 args.show = ["config"]
396 fwk.showInfo(args, pipeline)
397 args.show = ["history=task::addend"]
398 fwk.showInfo(args, pipeline)
399 args.show = ["tasks"]
400 fwk.showInfo(args, pipeline)
403class CmdLineFwkTestCaseWithButler(unittest.TestCase):
404 """A test case for CmdLineFwk"""
406 def setUp(self):
407 super().setUpClass()
408 self.root = tempfile.mkdtemp()
409 self.nQuanta = 5
410 self.pipeline = makeSimplePipeline(nQuanta=self.nQuanta)
412 def tearDown(self):
413 shutil.rmtree(self.root, ignore_errors=True)
414 super().tearDownClass()
416 def testSimpleQGraph(self):
417 """Test successfull execution of trivial quantum graph."""
418 args = _makeArgs(butler_config=self.root, input="test", output="output")
419 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
420 populateButler(self.pipeline, butler)
422 fwk = CmdLineFwk()
423 taskFactory = AddTaskFactoryMock()
425 qgraph = fwk.makeGraph(self.pipeline, args)
426 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
427 self.assertEqual(len(qgraph), self.nQuanta)
429 # run whole thing
430 fwk.runPipeline(qgraph, taskFactory, args)
431 self.assertEqual(taskFactory.countExec, self.nQuanta)
433 def testSimpleQGraphNoSkipExisting_inputs(self):
434 """Test for case when output data for one task already appears in
435 _input_ collection, but no ``--extend-run`` or ``-skip-existing``
436 option is present.
437 """
438 args = _makeArgs(
439 butler_config=self.root,
440 input="test",
441 output="output",
442 )
443 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
444 populateButler(
445 self.pipeline,
446 butler,
447 datasetTypes={
448 args.input: [
449 "add_dataset0",
450 "add_dataset1",
451 "add2_dataset1",
452 "add_init_output1",
453 "task0_config",
454 "task0_metadata",
455 "task0_log",
456 ]
457 },
458 )
460 fwk = CmdLineFwk()
461 taskFactory = AddTaskFactoryMock()
463 qgraph = fwk.makeGraph(self.pipeline, args)
464 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
465 # With current implementation graph has all nQuanta quanta, but when
466 # executing one quantum is skipped.
467 self.assertEqual(len(qgraph), self.nQuanta)
469 # run whole thing
470 fwk.runPipeline(qgraph, taskFactory, args)
471 self.assertEqual(taskFactory.countExec, self.nQuanta)
473 def testSimpleQGraphSkipExisting_inputs(self):
474 """Test for ``--skip-existing`` with output data for one task already
475 appears in _input_ collection. No ``--extend-run`` option is needed
476 for this case.
477 """
478 args = _makeArgs(
479 butler_config=self.root,
480 input="test",
481 output="output",
482 skip_existing_in=("test",),
483 )
484 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
485 populateButler(
486 self.pipeline,
487 butler,
488 datasetTypes={
489 args.input: [
490 "add_dataset0",
491 "add_dataset1",
492 "add2_dataset1",
493 "add_init_output1",
494 "task0_config",
495 "task0_metadata",
496 "task0_log",
497 ]
498 },
499 )
501 fwk = CmdLineFwk()
502 taskFactory = AddTaskFactoryMock()
504 qgraph = fwk.makeGraph(self.pipeline, args)
505 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
506 self.assertEqual(len(qgraph), self.nQuanta - 1)
508 # run whole thing
509 fwk.runPipeline(qgraph, taskFactory, args)
510 self.assertEqual(taskFactory.countExec, self.nQuanta - 1)
512 def testSimpleQGraphSkipExisting_outputs(self):
513 """Test for ``--skip-existing`` with output data for one task already
514 appears in _output_ collection. The ``--extend-run`` option is needed
515 for this case.
516 """
517 args = _makeArgs(
518 butler_config=self.root,
519 input="test",
520 output_run="output/run",
521 skip_existing_in=("output/run",),
522 )
523 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
524 populateButler(
525 self.pipeline,
526 butler,
527 datasetTypes={
528 args.input: ["add_dataset0"],
529 args.output_run: [
530 "add_dataset1",
531 "add2_dataset1",
532 "add_init_output1",
533 "task0_metadata",
534 "task0_log",
535 ],
536 },
537 )
539 fwk = CmdLineFwk()
540 taskFactory = AddTaskFactoryMock()
542 # fails without --extend-run
543 with self.assertRaisesRegex(ValueError, "--extend-run was not given"):
544 qgraph = fwk.makeGraph(self.pipeline, args)
546 # retry with --extend-run
547 args.extend_run = True
548 qgraph = fwk.makeGraph(self.pipeline, args)
550 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
551 # Graph does not include quantum for first task
552 self.assertEqual(len(qgraph), self.nQuanta - 1)
554 # run whole thing
555 fwk.runPipeline(qgraph, taskFactory, args)
556 self.assertEqual(taskFactory.countExec, self.nQuanta - 1)
558 def testSimpleQGraphOutputsFail(self):
559 """Test continuing execution of trivial quantum graph with partial
560 outputs.
561 """
562 args = _makeArgs(butler_config=self.root, input="test", output="output")
563 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
564 populateButler(self.pipeline, butler)
566 fwk = CmdLineFwk()
567 taskFactory = AddTaskFactoryMock(stopAt=3)
569 qgraph = fwk.makeGraph(self.pipeline, args)
570 self.assertEqual(len(qgraph), self.nQuanta)
572 # run first three quanta
573 with self.assertRaises(MPGraphExecutorError):
574 fwk.runPipeline(qgraph, taskFactory, args)
575 self.assertEqual(taskFactory.countExec, 3)
577 butler.registry.refresh()
579 # drop one of the two outputs from one task
580 ref1 = butler.registry.findDataset(
581 "add2_dataset2", collections=args.output, instrument="INSTR", detector=0
582 )
583 self.assertIsNotNone(ref1)
584 # also drop the metadata output
585 ref2 = butler.registry.findDataset(
586 "task1_metadata", collections=args.output, instrument="INSTR", detector=0
587 )
588 self.assertIsNotNone(ref2)
589 butler.pruneDatasets([ref1, ref2], disassociate=True, unstore=True, purge=True)
591 taskFactory.stopAt = -1
592 args.skip_existing_in = (args.output,)
593 args.extend_run = True
594 args.no_versions = True
595 with self.assertRaises(MPGraphExecutorError):
596 fwk.runPipeline(qgraph, taskFactory, args)
598 def testSimpleQGraphClobberOutputs(self):
599 """Test continuing execution of trivial quantum graph with
600 --clobber-outputs.
601 """
602 args = _makeArgs(butler_config=self.root, input="test", output="output")
603 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
604 populateButler(self.pipeline, butler)
606 fwk = CmdLineFwk()
607 taskFactory = AddTaskFactoryMock(stopAt=3)
609 qgraph = fwk.makeGraph(self.pipeline, args)
611 # should have one task and number of quanta
612 self.assertEqual(len(qgraph), self.nQuanta)
614 # run first three quanta
615 with self.assertRaises(MPGraphExecutorError):
616 fwk.runPipeline(qgraph, taskFactory, args)
617 self.assertEqual(taskFactory.countExec, 3)
619 butler.registry.refresh()
621 # drop one of the two outputs from one task
622 ref1 = butler.registry.findDataset(
623 "add2_dataset2", collections=args.output, dataId=dict(instrument="INSTR", detector=0)
624 )
625 self.assertIsNotNone(ref1)
626 # also drop the metadata output
627 ref2 = butler.registry.findDataset(
628 "task1_metadata", collections=args.output, dataId=dict(instrument="INSTR", detector=0)
629 )
630 self.assertIsNotNone(ref2)
631 butler.pruneDatasets([ref1, ref2], disassociate=True, unstore=True, purge=True)
633 taskFactory.stopAt = -1
634 args.skip_existing = True
635 args.extend_run = True
636 args.clobber_outputs = True
637 args.no_versions = True
638 fwk.runPipeline(qgraph, taskFactory, args)
639 # number of executed quanta is incremented
640 self.assertEqual(taskFactory.countExec, self.nQuanta + 1)
642 def testSimpleQGraphReplaceRun(self):
643 """Test repeated execution of trivial quantum graph with
644 --replace-run.
645 """
646 args = _makeArgs(butler_config=self.root, input="test", output="output", output_run="output/run1")
647 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
648 populateButler(self.pipeline, butler)
650 fwk = CmdLineFwk()
651 taskFactory = AddTaskFactoryMock()
653 qgraph = fwk.makeGraph(self.pipeline, args)
655 # should have one task and number of quanta
656 self.assertEqual(len(qgraph), self.nQuanta)
658 # deep copy is needed because quanta are updated in place
659 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
660 self.assertEqual(taskFactory.countExec, self.nQuanta)
662 # need to refresh collections explicitly (or make new butler/registry)
663 butler.registry.refresh()
664 collections = set(butler.registry.queryCollections(...))
665 self.assertEqual(collections, {"test", "output", "output/run1"})
667 # number of datasets written by pipeline:
668 # - nQuanta of init_outputs
669 # - nQuanta of configs
670 # - packages (single dataset)
671 # - nQuanta * two output datasets
672 # - nQuanta of metadata
673 # - nQuanta of log output
674 n_outputs = self.nQuanta * 6 + 1
675 refs = butler.registry.queryDatasets(..., collections="output/run1")
676 self.assertEqual(len(list(refs)), n_outputs)
678 # re-run with --replace-run (--inputs is ignored, as long as it hasn't
679 # changed)
680 args.replace_run = True
681 args.output_run = "output/run2"
682 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
684 butler.registry.refresh()
685 collections = set(butler.registry.queryCollections(...))
686 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2"})
688 # new output collection
689 refs = butler.registry.queryDatasets(..., collections="output/run2")
690 self.assertEqual(len(list(refs)), n_outputs)
692 # old output collection is still there
693 refs = butler.registry.queryDatasets(..., collections="output/run1")
694 self.assertEqual(len(list(refs)), n_outputs)
696 # re-run with --replace-run and --prune-replaced=unstore
697 args.replace_run = True
698 args.prune_replaced = "unstore"
699 args.output_run = "output/run3"
700 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
702 butler.registry.refresh()
703 collections = set(butler.registry.queryCollections(...))
704 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run3"})
706 # new output collection
707 refs = butler.registry.queryDatasets(..., collections="output/run3")
708 self.assertEqual(len(list(refs)), n_outputs)
710 # old output collection is still there, and it has all datasets but
711 # non-InitOutputs are not in datastore
712 refs = butler.registry.queryDatasets(..., collections="output/run2")
713 refs = list(refs)
714 self.assertEqual(len(refs), n_outputs)
715 initOutNameRe = re.compile("packages|task.*_config|add_init_output.*")
716 for ref in refs:
717 if initOutNameRe.fullmatch(ref.datasetType.name):
718 butler.get(ref, collections="output/run2")
719 else:
720 with self.assertRaises(FileNotFoundError):
721 butler.get(ref, collections="output/run2")
723 # re-run with --replace-run and --prune-replaced=purge
724 # This time also remove --input; passing the same inputs that we
725 # started with and not passing inputs at all should be equivalent.
726 args.input = None
727 args.replace_run = True
728 args.prune_replaced = "purge"
729 args.output_run = "output/run4"
730 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
732 butler.registry.refresh()
733 collections = set(butler.registry.queryCollections(...))
734 # output/run3 should disappear now
735 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
737 # new output collection
738 refs = butler.registry.queryDatasets(..., collections="output/run4")
739 self.assertEqual(len(list(refs)), n_outputs)
741 # Trying to run again with inputs that aren't exactly what we started
742 # with is an error, and the kind that should not modify the data repo.
743 with self.assertRaises(ValueError):
744 args.input = ["test", "output/run2"]
745 args.prune_replaced = None
746 args.replace_run = True
747 args.output_run = "output/run5"
748 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
749 butler.registry.refresh()
750 collections = set(butler.registry.queryCollections(...))
751 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
752 with self.assertRaises(ValueError):
753 args.input = ["output/run2", "test"]
754 args.prune_replaced = None
755 args.replace_run = True
756 args.output_run = "output/run6"
757 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
758 butler.registry.refresh()
759 collections = set(butler.registry.queryCollections(...))
760 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
762 def testMockTask(self):
763 """Test --mock option."""
764 args = _makeArgs(
765 butler_config=self.root, input="test", output="output", mock=True, register_dataset_types=True
766 )
767 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
768 populateButler(self.pipeline, butler)
770 fwk = CmdLineFwk()
771 taskFactory = AddTaskFactoryMock()
773 qgraph = fwk.makeGraph(self.pipeline, args)
774 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
775 self.assertEqual(len(qgraph), self.nQuanta)
777 # run whole thing
778 fwk.runPipeline(qgraph, taskFactory, args)
779 # None of the actual tasks is executed
780 self.assertEqual(taskFactory.countExec, 0)
782 # check dataset types
783 butler.registry.refresh()
784 datasetTypes = list(butler.registry.queryDatasetTypes(re.compile("^_mock_.*")))
785 self.assertEqual(len(datasetTypes), self.nQuanta * 2)
787 def testMockTaskFailure(self):
788 """Test --mock option and configure one of the tasks to fail."""
789 args = _makeArgs(
790 butler_config=self.root,
791 input="test",
792 output="output",
793 mock=True,
794 register_dataset_types=True,
795 mock_configs=[
796 _ACTION_CONFIG("task3-mock:failCondition='detector = 0'"),
797 ],
798 fail_fast=True,
799 )
800 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
801 populateButler(self.pipeline, butler)
803 fwk = CmdLineFwk()
804 taskFactory = AddTaskFactoryMock()
806 qgraph = fwk.makeGraph(self.pipeline, args)
807 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
808 self.assertEqual(len(qgraph), self.nQuanta)
810 with self.assertRaises(MPGraphExecutorError) as cm:
811 fwk.runPipeline(qgraph, taskFactory, args)
813 self.assertIsNotNone(cm.exception.__cause__)
814 self.assertRegex(str(cm.exception.__cause__), "Simulated failure: task=task3")
816 def testSubgraph(self):
817 """Test successfull execution of trivial quantum graph."""
818 args = _makeArgs(butler_config=self.root, input="test", output="output")
819 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
820 populateButler(self.pipeline, butler)
822 fwk = CmdLineFwk()
823 qgraph = fwk.makeGraph(self.pipeline, args)
825 # Select first two nodes for execution. This depends on node ordering
826 # which I assume is the same as execution order.
827 nNodes = 2
828 nodeIds = [node.nodeId for node in qgraph]
829 nodeIds = nodeIds[:nNodes]
831 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
832 self.assertEqual(len(qgraph), self.nQuanta)
834 with makeTmpFile(suffix=".qgraph") as tmpname, makeSQLiteRegistry() as registryConfig:
835 with open(tmpname, "wb") as saveFile:
836 qgraph.save(saveFile)
838 args = _makeArgs(
839 qgraph=tmpname,
840 qgraph_node_id=nodeIds,
841 registryConfig=registryConfig,
842 execution_butler_location=None,
843 )
844 fwk = CmdLineFwk()
846 # load graph, should only read a subset
847 qgraph = fwk.makeGraph(pipeline=None, args=args)
848 self.assertEqual(len(qgraph), nNodes)
850 def testShowGraph(self):
851 """Test for --show options for quantum graph."""
852 fwk = CmdLineFwk()
854 nQuanta = 2
855 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
857 args = _makeArgs(show=["graph"])
858 fwk.showInfo(args, pipeline=None, graph=qgraph)
860 def testShowGraphWorkflow(self):
861 fwk = CmdLineFwk()
863 nQuanta = 2
864 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
866 args = _makeArgs(show=["workflow"])
867 fwk.showInfo(args, pipeline=None, graph=qgraph)
869 # TODO: cannot test "uri" option presently, it instanciates
870 # butler from command line options and there is no way to pass butler
871 # mock to that code.
874class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
875 pass
878def setup_module(module):
879 lsst.utils.tests.init()
882if __name__ == "__main__": 882 ↛ 883line 882 didn't jump to line 883, because the condition on line 882 was never true
883 lsst.utils.tests.init()
884 unittest.main()