Coverage for tests/test_cmdLineFwk.py: 17%
437 statements
« prev ^ index » next coverage.py v6.4.1, created at 2022-06-17 02:13 -0700
« prev ^ index » next coverage.py v6.4.1, created at 2022-06-17 02:13 -0700
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 "namespace": "ctrl_mpexec_test",
221 "skypix": {
222 "common": "htm7",
223 "htm": {
224 "class": "lsst.sphgeom.HtmPixelization",
225 "max_level": 24,
226 },
227 },
228 "elements": {
229 "A": {
230 "keys": [
231 {
232 "name": "id",
233 "type": "int",
234 }
235 ],
236 "storage": {
237 "cls": "lsst.daf.butler.registry.dimensions.table.TableDimensionRecordStorage",
238 },
239 },
240 "B": {
241 "keys": [
242 {
243 "name": "id",
244 "type": "int",
245 }
246 ],
247 "storage": {
248 "cls": "lsst.daf.butler.registry.dimensions.table.TableDimensionRecordStorage",
249 },
250 },
251 },
252 "packers": {},
253 }
254 )
255 universe = DimensionUniverse(config=config)
256 fakeDSType = DatasetType("A", tuple(), storageClass="ExposureF", universe=universe)
257 taskDef = TaskDef(taskName=_TASK_CLASS, config=AddTask.ConfigClass(), taskClass=AddTask)
258 quanta = [
259 Quantum(
260 taskName=_TASK_CLASS,
261 inputs={
262 fakeDSType: [
263 DatasetRef(fakeDSType, DataCoordinate.standardize({"A": 1, "B": 2}, universe=universe))
264 ]
265 },
266 )
267 ] # type: ignore
268 qgraph = QuantumGraph({taskDef: set(quanta)})
269 return qgraph
272class CmdLineFwkTestCase(unittest.TestCase):
273 """A test case for CmdLineFwk"""
275 def testMakePipeline(self):
276 """Tests for CmdLineFwk.makePipeline method"""
277 fwk = CmdLineFwk()
279 # make empty pipeline
280 args = _makeArgs()
281 pipeline = fwk.makePipeline(args)
282 self.assertIsInstance(pipeline, Pipeline)
283 self.assertEqual(len(pipeline), 0)
285 # few tests with serialization
286 with makeTmpFile() as tmpname:
287 # make empty pipeline and store it in a file
288 args = _makeArgs(save_pipeline=tmpname)
289 pipeline = fwk.makePipeline(args)
290 self.assertIsInstance(pipeline, Pipeline)
292 # read pipeline from a file
293 args = _makeArgs(pipeline=tmpname)
294 pipeline = fwk.makePipeline(args)
295 self.assertIsInstance(pipeline, Pipeline)
296 self.assertEqual(len(pipeline), 0)
298 # single task pipeline, task name can be anything here
299 actions = [_ACTION_ADD_TASK("TaskOne:task1")]
300 args = _makeArgs(pipeline_actions=actions)
301 pipeline = fwk.makePipeline(args)
302 self.assertIsInstance(pipeline, Pipeline)
303 self.assertEqual(len(pipeline), 1)
305 # many task pipeline
306 actions = [
307 _ACTION_ADD_TASK("TaskOne:task1a"),
308 _ACTION_ADD_TASK("TaskTwo:task2"),
309 _ACTION_ADD_TASK("TaskOne:task1b"),
310 ]
311 args = _makeArgs(pipeline_actions=actions)
312 pipeline = fwk.makePipeline(args)
313 self.assertIsInstance(pipeline, Pipeline)
314 self.assertEqual(len(pipeline), 3)
316 # single task pipeline with config overrides, need real task class
317 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG("task:addend=100")]
318 args = _makeArgs(pipeline_actions=actions)
319 pipeline = fwk.makePipeline(args)
320 taskDefs = list(pipeline.toExpandedPipeline())
321 self.assertEqual(len(taskDefs), 1)
322 self.assertEqual(taskDefs[0].config.addend, 100)
324 overrides = b"config.addend = 1000\n"
325 with makeTmpFile(overrides) as tmpname:
326 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG_FILE("task:" + tmpname)]
327 args = _makeArgs(pipeline_actions=actions)
328 pipeline = fwk.makePipeline(args)
329 taskDefs = list(pipeline.toExpandedPipeline())
330 self.assertEqual(len(taskDefs), 1)
331 self.assertEqual(taskDefs[0].config.addend, 1000)
333 # Check --instrument option, for now it only checks that it does not
334 # crash.
335 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_ADD_INSTRUMENT("Instrument")]
336 args = _makeArgs(pipeline_actions=actions)
337 pipeline = fwk.makePipeline(args)
339 def testMakeGraphFromSave(self):
340 """Tests for CmdLineFwk.makeGraph method.
342 Only most trivial case is tested that does not do actual graph
343 building.
344 """
345 fwk = CmdLineFwk()
347 with makeTmpFile(suffix=".qgraph") as tmpname, makeSQLiteRegistry() as registryConfig:
349 # make non-empty graph and store it in a file
350 qgraph = _makeQGraph()
351 with open(tmpname, "wb") as saveFile:
352 qgraph.save(saveFile)
353 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
354 qgraph = fwk.makeGraph(None, args)
355 self.assertIsInstance(qgraph, QuantumGraph)
356 self.assertEqual(len(qgraph), 1)
358 # will fail if graph id does not match
359 args = _makeArgs(
360 qgraph=tmpname,
361 qgraph_id="R2-D2 is that you?",
362 registryConfig=registryConfig,
363 execution_butler_location=None,
364 )
365 with self.assertRaisesRegex(ValueError, "graphID does not match"):
366 fwk.makeGraph(None, args)
368 # save with wrong object type
369 with open(tmpname, "wb") as saveFile:
370 pickle.dump({}, saveFile)
371 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
372 with self.assertRaises(ValueError):
373 fwk.makeGraph(None, args)
375 # reading empty graph from pickle should work but makeGraph()
376 # will return None and make a warning
377 qgraph = QuantumGraph(dict())
378 with open(tmpname, "wb") as saveFile:
379 qgraph.save(saveFile)
380 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
381 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"):
382 # this also tests that warning is generated for empty graph
383 qgraph = fwk.makeGraph(None, args)
384 self.assertIs(qgraph, None)
386 def testShowPipeline(self):
387 """Test for --show options for pipeline."""
388 fwk = CmdLineFwk()
390 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG("task:addend=100")]
391 args = _makeArgs(pipeline_actions=actions)
392 pipeline = fwk.makePipeline(args)
394 args.show = ["pipeline"]
395 fwk.showInfo(args, pipeline)
396 args.show = ["config"]
397 fwk.showInfo(args, pipeline)
398 args.show = ["history=task::addend"]
399 fwk.showInfo(args, pipeline)
400 args.show = ["tasks"]
401 fwk.showInfo(args, pipeline)
404class CmdLineFwkTestCaseWithButler(unittest.TestCase):
405 """A test case for CmdLineFwk"""
407 def setUp(self):
408 super().setUpClass()
409 self.root = tempfile.mkdtemp()
410 self.nQuanta = 5
411 self.pipeline = makeSimplePipeline(nQuanta=self.nQuanta)
413 def tearDown(self):
414 shutil.rmtree(self.root, ignore_errors=True)
415 super().tearDownClass()
417 def testSimpleQGraph(self):
418 """Test successfull execution of trivial quantum graph."""
419 args = _makeArgs(butler_config=self.root, input="test", output="output")
420 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
421 populateButler(self.pipeline, butler)
423 fwk = CmdLineFwk()
424 taskFactory = AddTaskFactoryMock()
426 qgraph = fwk.makeGraph(self.pipeline, args)
427 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
428 self.assertEqual(len(qgraph), self.nQuanta)
430 # run whole thing
431 fwk.runPipeline(qgraph, taskFactory, args)
432 self.assertEqual(taskFactory.countExec, self.nQuanta)
434 def testSimpleQGraphNoSkipExisting_inputs(self):
435 """Test for case when output data for one task already appears in
436 _input_ collection, but no ``--extend-run`` or ``-skip-existing``
437 option is present.
438 """
439 args = _makeArgs(
440 butler_config=self.root,
441 input="test",
442 output="output",
443 )
444 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
445 populateButler(
446 self.pipeline,
447 butler,
448 datasetTypes={
449 args.input: [
450 "add_dataset0",
451 "add_dataset1",
452 "add2_dataset1",
453 "add_init_output1",
454 "task0_config",
455 "task0_metadata",
456 "task0_log",
457 ]
458 },
459 )
461 fwk = CmdLineFwk()
462 taskFactory = AddTaskFactoryMock()
464 qgraph = fwk.makeGraph(self.pipeline, args)
465 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
466 # With current implementation graph has all nQuanta quanta, but when
467 # executing one quantum is skipped.
468 self.assertEqual(len(qgraph), self.nQuanta)
470 # run whole thing
471 fwk.runPipeline(qgraph, taskFactory, args)
472 self.assertEqual(taskFactory.countExec, self.nQuanta)
474 def testSimpleQGraphSkipExisting_inputs(self):
475 """Test for ``--skip-existing`` with output data for one task already
476 appears in _input_ collection. No ``--extend-run`` option is needed
477 for this case.
478 """
479 args = _makeArgs(
480 butler_config=self.root,
481 input="test",
482 output="output",
483 skip_existing_in=("test",),
484 )
485 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
486 populateButler(
487 self.pipeline,
488 butler,
489 datasetTypes={
490 args.input: [
491 "add_dataset0",
492 "add_dataset1",
493 "add2_dataset1",
494 "add_init_output1",
495 "task0_config",
496 "task0_metadata",
497 "task0_log",
498 ]
499 },
500 )
502 fwk = CmdLineFwk()
503 taskFactory = AddTaskFactoryMock()
505 qgraph = fwk.makeGraph(self.pipeline, args)
506 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
507 self.assertEqual(len(qgraph), self.nQuanta - 1)
509 # run whole thing
510 fwk.runPipeline(qgraph, taskFactory, args)
511 self.assertEqual(taskFactory.countExec, self.nQuanta - 1)
513 def testSimpleQGraphSkipExisting_outputs(self):
514 """Test for ``--skip-existing`` with output data for one task already
515 appears in _output_ collection. The ``--extend-run`` option is needed
516 for this case.
517 """
518 args = _makeArgs(
519 butler_config=self.root,
520 input="test",
521 output_run="output/run",
522 skip_existing_in=("output/run",),
523 )
524 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
525 populateButler(
526 self.pipeline,
527 butler,
528 datasetTypes={
529 args.input: ["add_dataset0"],
530 args.output_run: [
531 "add_dataset1",
532 "add2_dataset1",
533 "add_init_output1",
534 "task0_metadata",
535 "task0_log",
536 ],
537 },
538 )
540 fwk = CmdLineFwk()
541 taskFactory = AddTaskFactoryMock()
543 # fails without --extend-run
544 with self.assertRaisesRegex(ValueError, "--extend-run was not given"):
545 qgraph = fwk.makeGraph(self.pipeline, args)
547 # retry with --extend-run
548 args.extend_run = True
549 qgraph = fwk.makeGraph(self.pipeline, args)
551 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
552 # Graph does not include quantum for first task
553 self.assertEqual(len(qgraph), self.nQuanta - 1)
555 # run whole thing
556 fwk.runPipeline(qgraph, taskFactory, args)
557 self.assertEqual(taskFactory.countExec, self.nQuanta - 1)
559 def testSimpleQGraphOutputsFail(self):
560 """Test continuing execution of trivial quantum graph with partial
561 outputs.
562 """
563 args = _makeArgs(butler_config=self.root, input="test", output="output")
564 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
565 populateButler(self.pipeline, butler)
567 fwk = CmdLineFwk()
568 taskFactory = AddTaskFactoryMock(stopAt=3)
570 qgraph = fwk.makeGraph(self.pipeline, args)
571 self.assertEqual(len(qgraph), self.nQuanta)
573 # run first three quanta
574 with self.assertRaises(MPGraphExecutorError):
575 fwk.runPipeline(qgraph, taskFactory, args)
576 self.assertEqual(taskFactory.countExec, 3)
578 butler.registry.refresh()
580 # drop one of the two outputs from one task
581 ref1 = butler.registry.findDataset(
582 "add2_dataset2", collections=args.output, instrument="INSTR", detector=0
583 )
584 self.assertIsNotNone(ref1)
585 # also drop the metadata output
586 ref2 = butler.registry.findDataset(
587 "task1_metadata", collections=args.output, instrument="INSTR", detector=0
588 )
589 self.assertIsNotNone(ref2)
590 butler.pruneDatasets([ref1, ref2], disassociate=True, unstore=True, purge=True)
592 taskFactory.stopAt = -1
593 args.skip_existing_in = (args.output,)
594 args.extend_run = True
595 args.no_versions = True
596 with self.assertRaises(MPGraphExecutorError):
597 fwk.runPipeline(qgraph, taskFactory, args)
599 def testSimpleQGraphClobberOutputs(self):
600 """Test continuing execution of trivial quantum graph with
601 --clobber-outputs.
602 """
603 args = _makeArgs(butler_config=self.root, input="test", output="output")
604 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
605 populateButler(self.pipeline, butler)
607 fwk = CmdLineFwk()
608 taskFactory = AddTaskFactoryMock(stopAt=3)
610 qgraph = fwk.makeGraph(self.pipeline, args)
612 # should have one task and number of quanta
613 self.assertEqual(len(qgraph), self.nQuanta)
615 # run first three quanta
616 with self.assertRaises(MPGraphExecutorError):
617 fwk.runPipeline(qgraph, taskFactory, args)
618 self.assertEqual(taskFactory.countExec, 3)
620 butler.registry.refresh()
622 # drop one of the two outputs from one task
623 ref1 = butler.registry.findDataset(
624 "add2_dataset2", collections=args.output, dataId=dict(instrument="INSTR", detector=0)
625 )
626 self.assertIsNotNone(ref1)
627 # also drop the metadata output
628 ref2 = butler.registry.findDataset(
629 "task1_metadata", collections=args.output, dataId=dict(instrument="INSTR", detector=0)
630 )
631 self.assertIsNotNone(ref2)
632 butler.pruneDatasets([ref1, ref2], disassociate=True, unstore=True, purge=True)
634 taskFactory.stopAt = -1
635 args.skip_existing = True
636 args.extend_run = True
637 args.clobber_outputs = True
638 args.no_versions = True
639 fwk.runPipeline(qgraph, taskFactory, args)
640 # number of executed quanta is incremented
641 self.assertEqual(taskFactory.countExec, self.nQuanta + 1)
643 def testSimpleQGraphReplaceRun(self):
644 """Test repeated execution of trivial quantum graph with
645 --replace-run.
646 """
647 args = _makeArgs(butler_config=self.root, input="test", output="output", output_run="output/run1")
648 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
649 populateButler(self.pipeline, butler)
651 fwk = CmdLineFwk()
652 taskFactory = AddTaskFactoryMock()
654 qgraph = fwk.makeGraph(self.pipeline, args)
656 # should have one task and number of quanta
657 self.assertEqual(len(qgraph), self.nQuanta)
659 # deep copy is needed because quanta are updated in place
660 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
661 self.assertEqual(taskFactory.countExec, self.nQuanta)
663 # need to refresh collections explicitly (or make new butler/registry)
664 butler.registry.refresh()
665 collections = set(butler.registry.queryCollections(...))
666 self.assertEqual(collections, {"test", "output", "output/run1"})
668 # number of datasets written by pipeline:
669 # - nQuanta of init_outputs
670 # - nQuanta of configs
671 # - packages (single dataset)
672 # - nQuanta * two output datasets
673 # - nQuanta of metadata
674 # - nQuanta of log output
675 n_outputs = self.nQuanta * 6 + 1
676 refs = butler.registry.queryDatasets(..., collections="output/run1")
677 self.assertEqual(len(list(refs)), n_outputs)
679 # re-run with --replace-run (--inputs is ignored, as long as it hasn't
680 # changed)
681 args.replace_run = True
682 args.output_run = "output/run2"
683 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
685 butler.registry.refresh()
686 collections = set(butler.registry.queryCollections(...))
687 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2"})
689 # new output collection
690 refs = butler.registry.queryDatasets(..., collections="output/run2")
691 self.assertEqual(len(list(refs)), n_outputs)
693 # old output collection is still there
694 refs = butler.registry.queryDatasets(..., collections="output/run1")
695 self.assertEqual(len(list(refs)), n_outputs)
697 # re-run with --replace-run and --prune-replaced=unstore
698 args.replace_run = True
699 args.prune_replaced = "unstore"
700 args.output_run = "output/run3"
701 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
703 butler.registry.refresh()
704 collections = set(butler.registry.queryCollections(...))
705 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run3"})
707 # new output collection
708 refs = butler.registry.queryDatasets(..., collections="output/run3")
709 self.assertEqual(len(list(refs)), n_outputs)
711 # old output collection is still there, and it has all datasets but
712 # non-InitOutputs are not in datastore
713 refs = butler.registry.queryDatasets(..., collections="output/run2")
714 refs = list(refs)
715 self.assertEqual(len(refs), n_outputs)
716 initOutNameRe = re.compile("packages|task.*_config|add_init_output.*")
717 for ref in refs:
718 if initOutNameRe.fullmatch(ref.datasetType.name):
719 butler.get(ref, collections="output/run2")
720 else:
721 with self.assertRaises(FileNotFoundError):
722 butler.get(ref, collections="output/run2")
724 # re-run with --replace-run and --prune-replaced=purge
725 # This time also remove --input; passing the same inputs that we
726 # started with and not passing inputs at all should be equivalent.
727 args.input = None
728 args.replace_run = True
729 args.prune_replaced = "purge"
730 args.output_run = "output/run4"
731 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
733 butler.registry.refresh()
734 collections = set(butler.registry.queryCollections(...))
735 # output/run3 should disappear now
736 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
738 # new output collection
739 refs = butler.registry.queryDatasets(..., collections="output/run4")
740 self.assertEqual(len(list(refs)), n_outputs)
742 # Trying to run again with inputs that aren't exactly what we started
743 # with is an error, and the kind that should not modify the data repo.
744 with self.assertRaises(ValueError):
745 args.input = ["test", "output/run2"]
746 args.prune_replaced = None
747 args.replace_run = True
748 args.output_run = "output/run5"
749 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
750 butler.registry.refresh()
751 collections = set(butler.registry.queryCollections(...))
752 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
753 with self.assertRaises(ValueError):
754 args.input = ["output/run2", "test"]
755 args.prune_replaced = None
756 args.replace_run = True
757 args.output_run = "output/run6"
758 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
759 butler.registry.refresh()
760 collections = set(butler.registry.queryCollections(...))
761 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
763 def testMockTask(self):
764 """Test --mock option."""
765 args = _makeArgs(
766 butler_config=self.root, input="test", output="output", mock=True, register_dataset_types=True
767 )
768 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
769 populateButler(self.pipeline, butler)
771 fwk = CmdLineFwk()
772 taskFactory = AddTaskFactoryMock()
774 qgraph = fwk.makeGraph(self.pipeline, args)
775 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
776 self.assertEqual(len(qgraph), self.nQuanta)
778 # run whole thing
779 fwk.runPipeline(qgraph, taskFactory, args)
780 # None of the actual tasks is executed
781 self.assertEqual(taskFactory.countExec, 0)
783 # check dataset types
784 butler.registry.refresh()
785 datasetTypes = list(butler.registry.queryDatasetTypes(re.compile("^_mock_.*")))
786 self.assertEqual(len(datasetTypes), self.nQuanta * 2)
788 def testMockTaskFailure(self):
789 """Test --mock option and configure one of the tasks to fail."""
790 args = _makeArgs(
791 butler_config=self.root,
792 input="test",
793 output="output",
794 mock=True,
795 register_dataset_types=True,
796 mock_configs=[
797 _ACTION_CONFIG("task3-mock:failCondition='detector = 0'"),
798 ],
799 fail_fast=True,
800 )
801 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
802 populateButler(self.pipeline, butler)
804 fwk = CmdLineFwk()
805 taskFactory = AddTaskFactoryMock()
807 qgraph = fwk.makeGraph(self.pipeline, args)
808 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
809 self.assertEqual(len(qgraph), self.nQuanta)
811 with self.assertRaises(MPGraphExecutorError) as cm:
812 fwk.runPipeline(qgraph, taskFactory, args)
814 self.assertIsNotNone(cm.exception.__cause__)
815 self.assertRegex(str(cm.exception.__cause__), "Simulated failure: task=task3")
817 def testSubgraph(self):
818 """Test successfull execution of trivial quantum graph."""
819 args = _makeArgs(butler_config=self.root, input="test", output="output")
820 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
821 populateButler(self.pipeline, butler)
823 fwk = CmdLineFwk()
824 qgraph = fwk.makeGraph(self.pipeline, args)
826 # Select first two nodes for execution. This depends on node ordering
827 # which I assume is the same as execution order.
828 nNodes = 2
829 nodeIds = [node.nodeId for node in qgraph]
830 nodeIds = nodeIds[:nNodes]
832 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
833 self.assertEqual(len(qgraph), self.nQuanta)
835 with makeTmpFile(suffix=".qgraph") as tmpname, makeSQLiteRegistry() as registryConfig:
836 with open(tmpname, "wb") as saveFile:
837 qgraph.save(saveFile)
839 args = _makeArgs(
840 qgraph=tmpname,
841 qgraph_node_id=nodeIds,
842 registryConfig=registryConfig,
843 execution_butler_location=None,
844 )
845 fwk = CmdLineFwk()
847 # load graph, should only read a subset
848 qgraph = fwk.makeGraph(pipeline=None, args=args)
849 self.assertEqual(len(qgraph), nNodes)
851 def testShowGraph(self):
852 """Test for --show options for quantum graph."""
853 fwk = CmdLineFwk()
855 nQuanta = 2
856 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
858 args = _makeArgs(show=["graph"])
859 fwk.showInfo(args, pipeline=None, graph=qgraph)
861 def testShowGraphWorkflow(self):
862 fwk = CmdLineFwk()
864 nQuanta = 2
865 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
867 args = _makeArgs(show=["workflow"])
868 fwk.showInfo(args, pipeline=None, graph=qgraph)
870 # TODO: cannot test "uri" option presently, it instanciates
871 # butler from command line options and there is no way to pass butler
872 # mock to that code.
874 def testSimpleQGraphDatastoreRecords(self):
875 """Test quantum graph generation with --qgraph-datastore-records."""
876 args = _makeArgs(
877 butler_config=self.root, input="test", output="output", qgraph_datastore_records=True
878 )
879 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
880 populateButler(self.pipeline, butler)
882 fwk = CmdLineFwk()
883 qgraph = fwk.makeGraph(self.pipeline, args)
884 self.assertEqual(len(qgraph), self.nQuanta)
885 for i, qnode in enumerate(qgraph):
886 quantum = qnode.quantum
887 self.assertIsNotNone(quantum.datastore_records)
888 # only the first quantum has a pre-existing input
889 if i == 0:
890 datastore_name = "FileDatastore@<butlerRoot>"
891 self.assertEqual(set(quantum.datastore_records.keys()), {datastore_name})
892 records_data = quantum.datastore_records[datastore_name]
893 records = dict(records_data.records)
894 self.assertEqual(len(records), 1)
895 _, records = records.popitem()
896 records = records["file_datastore_records"]
897 self.assertEqual(
898 [record.path for record in records],
899 ["test/add_dataset0/add_dataset0_INSTR_det0_test.pickle"],
900 )
901 else:
902 self.assertEqual(quantum.datastore_records, {})
905class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
906 pass
909def setup_module(module):
910 lsst.utils.tests.init()
913if __name__ == "__main__": 913 ↛ 914line 913 didn't jump to line 914, because the condition on line 913 was never true
914 lsst.utils.tests.init()
915 unittest.main()