Coverage for tests/test_cmdLineFwk.py: 14%
491 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-10-22 01:58 -0700
« prev ^ index » next coverage.py v6.5.0, created at 2022-10-22 01:58 -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 io import StringIO
36from types import SimpleNamespace
37from typing import NamedTuple
39import click
40import lsst.pex.config as pexConfig
41import lsst.pipe.base.connectionTypes as cT
42import lsst.utils.tests
43from lsst.ctrl.mpexec import CmdLineFwk, MPGraphExecutorError
44from lsst.ctrl.mpexec.cli.opt import run_options
45from lsst.ctrl.mpexec.cli.utils import (
46 _ACTION_ADD_INSTRUMENT,
47 _ACTION_ADD_TASK,
48 _ACTION_CONFIG,
49 _ACTION_CONFIG_FILE,
50 PipetaskCommand,
51)
52from lsst.ctrl.mpexec.showInfo import ShowInfo
53from lsst.daf.butler import (
54 Config,
55 DataCoordinate,
56 DatasetRef,
57 DimensionConfig,
58 DimensionUniverse,
59 Quantum,
60 Registry,
61)
62from lsst.daf.butler.core.datasets.type import DatasetType
63from lsst.daf.butler.registry import RegistryConfig
64from lsst.pipe.base import (
65 Instrument,
66 Pipeline,
67 PipelineTaskConfig,
68 PipelineTaskConnections,
69 QuantumGraph,
70 TaskDef,
71)
72from lsst.pipe.base.graphBuilder import DatasetQueryConstraintVariant as DQCVariant
73from lsst.pipe.base.tests.simpleQGraph import (
74 AddTask,
75 AddTaskFactoryMock,
76 makeSimpleButler,
77 makeSimplePipeline,
78 makeSimpleQGraph,
79 populateButler,
80)
81from lsst.utils.tests import temporaryDirectory
83logging.basicConfig(level=getattr(logging, os.environ.get("UNIT_TEST_LOGGING_LEVEL", "INFO"), logging.INFO))
85# Have to monkey-patch Instrument.fromName() to not retrieve non-existing
86# instrument from registry, these tests can run fine without actual instrument
87# and implementing full mock for Instrument is too complicated.
88Instrument.fromName = lambda name, reg: None 88 ↛ exitline 88 didn't run the lambda on line 88
91@contextlib.contextmanager
92def makeTmpFile(contents=None, suffix=None):
93 """Context manager for generating temporary file name.
95 Temporary file is deleted on exiting context.
97 Parameters
98 ----------
99 contents : `bytes`
100 Data to write into a file.
101 """
102 fd, tmpname = tempfile.mkstemp(suffix=suffix)
103 if contents:
104 os.write(fd, contents)
105 os.close(fd)
106 yield tmpname
107 with contextlib.suppress(OSError):
108 os.remove(tmpname)
111@contextlib.contextmanager
112def makeSQLiteRegistry(create=True, universe=None):
113 """Context manager to create new empty registry database.
115 Yields
116 ------
117 config : `RegistryConfig`
118 Registry configuration for initialized registry database.
119 """
120 dimensionConfig = universe.dimensionConfig if universe is not None else _makeDimensionConfig()
121 with temporaryDirectory() as tmpdir:
122 uri = f"sqlite:///{tmpdir}/gen3.sqlite"
123 config = RegistryConfig()
124 config["db"] = uri
125 if create:
126 Registry.createFromConfig(config, dimensionConfig=dimensionConfig)
127 yield config
130class SimpleConnections(PipelineTaskConnections, dimensions=(), defaultTemplates={"template": "simple"}):
131 schema = cT.InitInput(doc="Schema", name="{template}schema", storageClass="SourceCatalog")
134class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections):
135 field = pexConfig.Field(dtype=str, doc="arbitrary string")
137 def setDefaults(self):
138 PipelineTaskConfig.setDefaults(self)
141def _makeArgs(registryConfig=None, **kwargs):
142 """Return parsed command line arguments.
144 By default butler_config is set to `Config` populated with some defaults,
145 it can be overridden completely by keyword argument.
147 Parameters
148 ----------
149 cmd : `str`, optional
150 Produce arguments for this pipetask command.
151 registryConfig : `RegistryConfig`, optional
152 Override for registry configuration.
153 **kwargs
154 Overrides for other arguments.
155 """
156 # Use a mock to get the default value of arguments to 'run'.
158 mock = unittest.mock.Mock()
160 @click.command(cls=PipetaskCommand)
161 @run_options()
162 def fake_run(ctx, **kwargs):
163 """Fake "pipetask run" command for gathering input arguments.
165 The arguments & options should always match the arguments & options in
166 the "real" command function `lsst.ctrl.mpexec.cli.cmd.run`.
167 """
168 mock(**kwargs)
170 runner = click.testing.CliRunner()
171 # --butler-config is the only required option
172 result = runner.invoke(fake_run, "--butler-config /")
173 if result.exit_code != 0:
174 raise RuntimeError(f"Failure getting default args from 'fake_run': {result}")
175 mock.assert_called_once()
176 args = mock.call_args[1]
177 args["enableLsstDebug"] = args.pop("debug")
178 args["execution_butler_location"] = args.pop("save_execution_butler")
179 if "pipeline_actions" not in args:
180 args["pipeline_actions"] = []
181 if "mock_configs" not in args:
182 args["mock_configs"] = []
183 args = SimpleNamespace(**args)
185 # override butler_config with our defaults
186 if "butler_config" not in kwargs:
187 args.butler_config = Config()
188 if registryConfig:
189 args.butler_config["registry"] = registryConfig
190 # The default datastore has a relocatable root, so we need to specify
191 # some root here for it to use
192 args.butler_config.configFile = "."
194 # override arguments from keyword parameters
195 for key, value in kwargs.items():
196 setattr(args, key, value)
197 args.dataset_query_constraint = DQCVariant.fromExpression(args.dataset_query_constraint)
198 return args
201class FakeDSType(NamedTuple):
202 name: str
205@dataclass(frozen=True)
206class FakeDSRef:
207 datasetType: str
208 dataId: tuple
210 def isComponent(self):
211 return False
214# Task class name used by tests, needs to be importable
215_TASK_CLASS = "lsst.pipe.base.tests.simpleQGraph.AddTask"
218def _makeDimensionConfig():
219 """Make a simple dimension universe configuration."""
220 return DimensionConfig(
221 {
222 "version": 1,
223 "namespace": "ctrl_mpexec_test",
224 "skypix": {
225 "common": "htm7",
226 "htm": {
227 "class": "lsst.sphgeom.HtmPixelization",
228 "max_level": 24,
229 },
230 },
231 "elements": {
232 "A": {
233 "keys": [
234 {
235 "name": "id",
236 "type": "int",
237 }
238 ],
239 "storage": {
240 "cls": "lsst.daf.butler.registry.dimensions.table.TableDimensionRecordStorage",
241 },
242 },
243 "B": {
244 "keys": [
245 {
246 "name": "id",
247 "type": "int",
248 }
249 ],
250 "storage": {
251 "cls": "lsst.daf.butler.registry.dimensions.table.TableDimensionRecordStorage",
252 },
253 },
254 },
255 "packers": {},
256 }
257 )
260def _makeQGraph():
261 """Make a trivial QuantumGraph with one quantum.
263 The only thing that we need to do with this quantum graph is to pickle
264 it, the quanta in this graph are not usable for anything else.
266 Returns
267 -------
268 qgraph : `~lsst.pipe.base.QuantumGraph`
269 """
270 universe = DimensionUniverse(config=_makeDimensionConfig())
271 fakeDSType = DatasetType("A", tuple(), storageClass="ExposureF", universe=universe)
272 taskDef = TaskDef(taskName=_TASK_CLASS, config=AddTask.ConfigClass(), taskClass=AddTask)
273 quanta = [
274 Quantum(
275 taskName=_TASK_CLASS,
276 inputs={
277 fakeDSType: [
278 DatasetRef(fakeDSType, DataCoordinate.standardize({"A": 1, "B": 2}, universe=universe))
279 ]
280 },
281 )
282 ] # type: ignore
283 qgraph = QuantumGraph({taskDef: set(quanta)}, universe=universe)
284 return qgraph
287class CmdLineFwkTestCase(unittest.TestCase):
288 """A test case for CmdLineFwk"""
290 def testMakePipeline(self):
291 """Tests for CmdLineFwk.makePipeline method"""
292 fwk = CmdLineFwk()
294 # make empty pipeline
295 args = _makeArgs()
296 pipeline = fwk.makePipeline(args)
297 self.assertIsInstance(pipeline, Pipeline)
298 self.assertEqual(len(pipeline), 0)
300 # few tests with serialization
301 with makeTmpFile() as tmpname:
302 # make empty pipeline and store it in a file
303 args = _makeArgs(save_pipeline=tmpname)
304 pipeline = fwk.makePipeline(args)
305 self.assertIsInstance(pipeline, Pipeline)
307 # read pipeline from a file
308 args = _makeArgs(pipeline=tmpname)
309 pipeline = fwk.makePipeline(args)
310 self.assertIsInstance(pipeline, Pipeline)
311 self.assertEqual(len(pipeline), 0)
313 # single task pipeline, task name can be anything here
314 actions = [_ACTION_ADD_TASK("TaskOne:task1")]
315 args = _makeArgs(pipeline_actions=actions)
316 pipeline = fwk.makePipeline(args)
317 self.assertIsInstance(pipeline, Pipeline)
318 self.assertEqual(len(pipeline), 1)
320 # many task pipeline
321 actions = [
322 _ACTION_ADD_TASK("TaskOne:task1a"),
323 _ACTION_ADD_TASK("TaskTwo:task2"),
324 _ACTION_ADD_TASK("TaskOne:task1b"),
325 ]
326 args = _makeArgs(pipeline_actions=actions)
327 pipeline = fwk.makePipeline(args)
328 self.assertIsInstance(pipeline, Pipeline)
329 self.assertEqual(len(pipeline), 3)
331 # single task pipeline with config overrides, need real task class
332 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG("task:addend=100")]
333 args = _makeArgs(pipeline_actions=actions)
334 pipeline = fwk.makePipeline(args)
335 taskDefs = list(pipeline.toExpandedPipeline())
336 self.assertEqual(len(taskDefs), 1)
337 self.assertEqual(taskDefs[0].config.addend, 100)
339 overrides = b"config.addend = 1000\n"
340 with makeTmpFile(overrides) as tmpname:
341 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG_FILE("task:" + tmpname)]
342 args = _makeArgs(pipeline_actions=actions)
343 pipeline = fwk.makePipeline(args)
344 taskDefs = list(pipeline.toExpandedPipeline())
345 self.assertEqual(len(taskDefs), 1)
346 self.assertEqual(taskDefs[0].config.addend, 1000)
348 # Check --instrument option, for now it only checks that it does not
349 # crash.
350 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_ADD_INSTRUMENT("Instrument")]
351 args = _makeArgs(pipeline_actions=actions)
352 pipeline = fwk.makePipeline(args)
354 def testMakeGraphFromSave(self):
355 """Tests for CmdLineFwk.makeGraph method.
357 Only most trivial case is tested that does not do actual graph
358 building.
359 """
360 fwk = CmdLineFwk()
362 with makeTmpFile(suffix=".qgraph") as tmpname, makeSQLiteRegistry() as registryConfig:
364 # make non-empty graph and store it in a file
365 qgraph = _makeQGraph()
366 with open(tmpname, "wb") as saveFile:
367 qgraph.save(saveFile)
368 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
369 qgraph = fwk.makeGraph(None, args)
370 self.assertIsInstance(qgraph, QuantumGraph)
371 self.assertEqual(len(qgraph), 1)
373 # will fail if graph id does not match
374 args = _makeArgs(
375 qgraph=tmpname,
376 qgraph_id="R2-D2 is that you?",
377 registryConfig=registryConfig,
378 execution_butler_location=None,
379 )
380 with self.assertRaisesRegex(ValueError, "graphID does not match"):
381 fwk.makeGraph(None, args)
383 # save with wrong object type
384 with open(tmpname, "wb") as saveFile:
385 pickle.dump({}, saveFile)
386 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
387 with self.assertRaises(ValueError):
388 fwk.makeGraph(None, args)
390 # reading empty graph from pickle should work but makeGraph()
391 # will return None.
392 qgraph = QuantumGraph(dict(), universe=DimensionUniverse(_makeDimensionConfig()))
393 with open(tmpname, "wb") as saveFile:
394 qgraph.save(saveFile)
395 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None)
396 qgraph = fwk.makeGraph(None, args)
397 self.assertIs(qgraph, None)
399 def testShowPipeline(self):
400 """Test for --show options for pipeline."""
401 fwk = CmdLineFwk()
403 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG("task:addend=100")]
404 args = _makeArgs(pipeline_actions=actions)
405 pipeline = fwk.makePipeline(args)
407 with self.assertRaises(ValueError):
408 ShowInfo(["unrecognized", "config"])
410 stream = StringIO()
411 show = ShowInfo(
412 ["pipeline", "config", "history=task::addend", "tasks", "dump-config", "config=task::add*"],
413 stream=stream,
414 )
415 show.show_pipeline_info(pipeline)
416 self.assertEqual(show.unhandled, frozenset({}))
417 stream.seek(0)
418 output = stream.read()
419 self.assertIn("config.addend=100", output) # config option
420 self.assertIn("addend\n3", output) # History output
421 self.assertIn("class: lsst.pipe.base.tests.simpleQGraph.AddTask", output) # pipeline
423 show = ShowInfo(["pipeline", "uri"], stream=stream)
424 show.show_pipeline_info(pipeline)
425 self.assertEqual(show.unhandled, frozenset({"uri"}))
426 self.assertEqual(show.handled, {"pipeline"})
428 stream = StringIO()
429 show = ShowInfo(["config=task::addend.missing"], stream=stream) # No match
430 show.show_pipeline_info(pipeline)
431 stream.seek(0)
432 output = stream.read().strip()
433 self.assertEqual("### Configuration for task `task'", output)
435 stream = StringIO()
436 show = ShowInfo(["config=task::addEnd:NOIGNORECASE"], stream=stream) # No match
437 show.show_pipeline_info(pipeline)
438 stream.seek(0)
439 output = stream.read().strip()
440 self.assertEqual("### Configuration for task `task'", output)
442 stream = StringIO()
443 show = ShowInfo(["config=task::addEnd"], stream=stream) # Match but warns
444 show.show_pipeline_info(pipeline)
445 stream.seek(0)
446 output = stream.read().strip()
447 self.assertIn("NOIGNORECASE", output)
449 show = ShowInfo(["dump-config=notask"])
450 with self.assertRaises(ValueError) as cm:
451 show.show_pipeline_info(pipeline)
452 self.assertIn("Pipeline has no tasks named notask", str(cm.exception))
454 show = ShowInfo(["history"])
455 with self.assertRaises(ValueError) as cm:
456 show.show_pipeline_info(pipeline)
457 self.assertIn("Please provide a value", str(cm.exception))
459 show = ShowInfo(["history=notask::param"])
460 with self.assertRaises(ValueError) as cm:
461 show.show_pipeline_info(pipeline)
462 self.assertIn("Pipeline has no tasks named notask", str(cm.exception))
465class CmdLineFwkTestCaseWithButler(unittest.TestCase):
466 """A test case for CmdLineFwk"""
468 def setUp(self):
469 super().setUpClass()
470 self.root = tempfile.mkdtemp()
471 self.nQuanta = 5
472 self.pipeline = makeSimplePipeline(nQuanta=self.nQuanta)
474 def tearDown(self):
475 shutil.rmtree(self.root, ignore_errors=True)
476 super().tearDownClass()
478 def testSimpleQGraph(self):
479 """Test successfull execution of trivial quantum graph."""
480 args = _makeArgs(butler_config=self.root, input="test", output="output")
481 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
482 populateButler(self.pipeline, butler)
484 fwk = CmdLineFwk()
485 taskFactory = AddTaskFactoryMock()
487 qgraph = fwk.makeGraph(self.pipeline, args)
488 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
489 self.assertEqual(len(qgraph), self.nQuanta)
491 # run whole thing
492 fwk.runPipeline(qgraph, taskFactory, args)
493 self.assertEqual(taskFactory.countExec, self.nQuanta)
495 def testEmptyQGraph(self):
496 """Test that making an empty QG produces the right error messages."""
497 # We make QG generation fail by populating one input collection in the
498 # butler while using a different one (that we only register, not
499 # populate) to make the QG.
500 args = _makeArgs(butler_config=self.root, input="bad_input", output="output")
501 butler = makeSimpleButler(self.root, run="good_input", inMemory=False)
502 butler.registry.registerCollection("bad_input")
503 populateButler(self.pipeline, butler)
505 fwk = CmdLineFwk()
506 with self.assertLogs(level=logging.CRITICAL) as cm:
507 qgraph = fwk.makeGraph(self.pipeline, args)
508 self.assertRegex(
509 cm.output[0], ".*Initial data ID query returned no rows, so QuantumGraph will be empty.*"
510 )
511 self.assertRegex(cm.output[1], ".*No datasets.*bad_input.*")
512 self.assertIsNone(qgraph)
514 def testSimpleQGraphNoSkipExisting_inputs(self):
515 """Test for case when output data for one task already appears in
516 _input_ collection, but no ``--extend-run`` or ``-skip-existing``
517 option is present.
518 """
519 args = _makeArgs(
520 butler_config=self.root,
521 input="test",
522 output="output",
523 )
524 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
525 populateButler(
526 self.pipeline,
527 butler,
528 datasetTypes={
529 args.input: [
530 "add_dataset0",
531 "add_dataset1",
532 "add2_dataset1",
533 "add_init_output1",
534 "task0_config",
535 "task0_metadata",
536 "task0_log",
537 ]
538 },
539 )
541 fwk = CmdLineFwk()
542 taskFactory = AddTaskFactoryMock()
544 qgraph = fwk.makeGraph(self.pipeline, args)
545 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
546 # With current implementation graph has all nQuanta quanta, but when
547 # executing one quantum is skipped.
548 self.assertEqual(len(qgraph), self.nQuanta)
550 # run whole thing
551 fwk.runPipeline(qgraph, taskFactory, args)
552 self.assertEqual(taskFactory.countExec, self.nQuanta)
554 def testSimpleQGraphSkipExisting_inputs(self):
555 """Test for ``--skip-existing`` with output data for one task already
556 appears in _input_ collection. No ``--extend-run`` option is needed
557 for this case.
558 """
559 args = _makeArgs(
560 butler_config=self.root,
561 input="test",
562 output="output",
563 skip_existing_in=("test",),
564 )
565 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
566 populateButler(
567 self.pipeline,
568 butler,
569 datasetTypes={
570 args.input: [
571 "add_dataset0",
572 "add_dataset1",
573 "add2_dataset1",
574 "add_init_output1",
575 "task0_config",
576 "task0_metadata",
577 "task0_log",
578 ]
579 },
580 )
582 fwk = CmdLineFwk()
583 taskFactory = AddTaskFactoryMock()
585 qgraph = fwk.makeGraph(self.pipeline, args)
586 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
587 self.assertEqual(len(qgraph), self.nQuanta - 1)
589 # run whole thing
590 fwk.runPipeline(qgraph, taskFactory, args)
591 self.assertEqual(taskFactory.countExec, self.nQuanta - 1)
593 def testSimpleQGraphSkipExisting_outputs(self):
594 """Test for ``--skip-existing`` with output data for one task already
595 appears in _output_ collection. The ``--extend-run`` option is needed
596 for this case.
597 """
598 args = _makeArgs(
599 butler_config=self.root,
600 input="test",
601 output_run="output/run",
602 skip_existing_in=("output/run",),
603 )
604 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
605 populateButler(
606 self.pipeline,
607 butler,
608 datasetTypes={
609 args.input: ["add_dataset0"],
610 args.output_run: [
611 "add_dataset1",
612 "add2_dataset1",
613 "add_init_output1",
614 "task0_metadata",
615 "task0_log",
616 ],
617 },
618 )
620 fwk = CmdLineFwk()
621 taskFactory = AddTaskFactoryMock()
623 # fails without --extend-run
624 with self.assertRaisesRegex(ValueError, "--extend-run was not given"):
625 qgraph = fwk.makeGraph(self.pipeline, args)
627 # retry with --extend-run
628 args.extend_run = True
629 qgraph = fwk.makeGraph(self.pipeline, args)
631 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
632 # Graph does not include quantum for first task
633 self.assertEqual(len(qgraph), self.nQuanta - 1)
635 # run whole thing
636 fwk.runPipeline(qgraph, taskFactory, args)
637 self.assertEqual(taskFactory.countExec, self.nQuanta - 1)
639 def testSimpleQGraphOutputsFail(self):
640 """Test continuing execution of trivial quantum graph with partial
641 outputs.
642 """
643 args = _makeArgs(butler_config=self.root, input="test", output="output")
644 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
645 populateButler(self.pipeline, butler)
647 fwk = CmdLineFwk()
648 taskFactory = AddTaskFactoryMock(stopAt=3)
650 qgraph = fwk.makeGraph(self.pipeline, args)
651 self.assertEqual(len(qgraph), self.nQuanta)
653 # run first three quanta
654 with self.assertRaises(MPGraphExecutorError):
655 fwk.runPipeline(qgraph, taskFactory, args)
656 self.assertEqual(taskFactory.countExec, 3)
658 butler.registry.refresh()
660 # drop one of the two outputs from one task
661 ref1 = butler.registry.findDataset(
662 "add2_dataset2", collections=args.output, instrument="INSTR", detector=0
663 )
664 self.assertIsNotNone(ref1)
665 # also drop the metadata output
666 ref2 = butler.registry.findDataset(
667 "task1_metadata", collections=args.output, instrument="INSTR", detector=0
668 )
669 self.assertIsNotNone(ref2)
670 butler.pruneDatasets([ref1, ref2], disassociate=True, unstore=True, purge=True)
672 taskFactory.stopAt = -1
673 args.skip_existing_in = (args.output,)
674 args.extend_run = True
675 args.no_versions = True
676 with self.assertRaises(MPGraphExecutorError):
677 fwk.runPipeline(qgraph, taskFactory, args)
679 def testSimpleQGraphClobberOutputs(self):
680 """Test continuing execution of trivial quantum graph with
681 --clobber-outputs.
682 """
683 args = _makeArgs(butler_config=self.root, input="test", output="output")
684 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
685 populateButler(self.pipeline, butler)
687 fwk = CmdLineFwk()
688 taskFactory = AddTaskFactoryMock(stopAt=3)
690 qgraph = fwk.makeGraph(self.pipeline, args)
692 # should have one task and number of quanta
693 self.assertEqual(len(qgraph), self.nQuanta)
695 # run first three quanta
696 with self.assertRaises(MPGraphExecutorError):
697 fwk.runPipeline(qgraph, taskFactory, args)
698 self.assertEqual(taskFactory.countExec, 3)
700 butler.registry.refresh()
702 # drop one of the two outputs from one task
703 ref1 = butler.registry.findDataset(
704 "add2_dataset2", collections=args.output, dataId=dict(instrument="INSTR", detector=0)
705 )
706 self.assertIsNotNone(ref1)
707 # also drop the metadata output
708 ref2 = butler.registry.findDataset(
709 "task1_metadata", collections=args.output, dataId=dict(instrument="INSTR", detector=0)
710 )
711 self.assertIsNotNone(ref2)
712 butler.pruneDatasets([ref1, ref2], disassociate=True, unstore=True, purge=True)
714 taskFactory.stopAt = -1
715 args.skip_existing = True
716 args.extend_run = True
717 args.clobber_outputs = True
718 args.no_versions = True
719 fwk.runPipeline(qgraph, taskFactory, args)
720 # number of executed quanta is incremented
721 self.assertEqual(taskFactory.countExec, self.nQuanta + 1)
723 def testSimpleQGraphReplaceRun(self):
724 """Test repeated execution of trivial quantum graph with
725 --replace-run.
726 """
727 args = _makeArgs(butler_config=self.root, input="test", output="output", output_run="output/run1")
728 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
729 populateButler(self.pipeline, butler)
731 fwk = CmdLineFwk()
732 taskFactory = AddTaskFactoryMock()
734 qgraph = fwk.makeGraph(self.pipeline, args)
736 # should have one task and number of quanta
737 self.assertEqual(len(qgraph), self.nQuanta)
739 # deep copy is needed because quanta are updated in place
740 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
741 self.assertEqual(taskFactory.countExec, self.nQuanta)
743 # need to refresh collections explicitly (or make new butler/registry)
744 butler.registry.refresh()
745 collections = set(butler.registry.queryCollections(...))
746 self.assertEqual(collections, {"test", "output", "output/run1"})
748 # number of datasets written by pipeline:
749 # - nQuanta of init_outputs
750 # - nQuanta of configs
751 # - packages (single dataset)
752 # - nQuanta * two output datasets
753 # - nQuanta of metadata
754 # - nQuanta of log output
755 n_outputs = self.nQuanta * 6 + 1
756 refs = butler.registry.queryDatasets(..., collections="output/run1")
757 self.assertEqual(len(list(refs)), n_outputs)
759 # re-run with --replace-run (--inputs is ignored, as long as it hasn't
760 # changed)
761 args.replace_run = True
762 args.output_run = "output/run2"
763 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
765 butler.registry.refresh()
766 collections = set(butler.registry.queryCollections(...))
767 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2"})
769 # new output collection
770 refs = butler.registry.queryDatasets(..., collections="output/run2")
771 self.assertEqual(len(list(refs)), n_outputs)
773 # old output collection is still there
774 refs = butler.registry.queryDatasets(..., collections="output/run1")
775 self.assertEqual(len(list(refs)), n_outputs)
777 # re-run with --replace-run and --prune-replaced=unstore
778 args.replace_run = True
779 args.prune_replaced = "unstore"
780 args.output_run = "output/run3"
781 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
783 butler.registry.refresh()
784 collections = set(butler.registry.queryCollections(...))
785 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run3"})
787 # new output collection
788 refs = butler.registry.queryDatasets(..., collections="output/run3")
789 self.assertEqual(len(list(refs)), n_outputs)
791 # old output collection is still there, and it has all datasets but
792 # non-InitOutputs are not in datastore
793 refs = butler.registry.queryDatasets(..., collections="output/run2")
794 refs = list(refs)
795 self.assertEqual(len(refs), n_outputs)
796 initOutNameRe = re.compile("packages|task.*_config|add_init_output.*")
797 for ref in refs:
798 if initOutNameRe.fullmatch(ref.datasetType.name):
799 butler.get(ref, collections="output/run2")
800 else:
801 with self.assertRaises(FileNotFoundError):
802 butler.get(ref, collections="output/run2")
804 # re-run with --replace-run and --prune-replaced=purge
805 # This time also remove --input; passing the same inputs that we
806 # started with and not passing inputs at all should be equivalent.
807 args.input = None
808 args.replace_run = True
809 args.prune_replaced = "purge"
810 args.output_run = "output/run4"
811 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
813 butler.registry.refresh()
814 collections = set(butler.registry.queryCollections(...))
815 # output/run3 should disappear now
816 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
818 # new output collection
819 refs = butler.registry.queryDatasets(..., collections="output/run4")
820 self.assertEqual(len(list(refs)), n_outputs)
822 # Trying to run again with inputs that aren't exactly what we started
823 # with is an error, and the kind that should not modify the data repo.
824 with self.assertRaises(ValueError):
825 args.input = ["test", "output/run2"]
826 args.prune_replaced = None
827 args.replace_run = True
828 args.output_run = "output/run5"
829 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
830 butler.registry.refresh()
831 collections = set(butler.registry.queryCollections(...))
832 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
833 with self.assertRaises(ValueError):
834 args.input = ["output/run2", "test"]
835 args.prune_replaced = None
836 args.replace_run = True
837 args.output_run = "output/run6"
838 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
839 butler.registry.refresh()
840 collections = set(butler.registry.queryCollections(...))
841 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
843 def testMockTask(self):
844 """Test --mock option."""
845 args = _makeArgs(
846 butler_config=self.root, input="test", output="output", mock=True, register_dataset_types=True
847 )
848 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
849 populateButler(self.pipeline, butler)
851 fwk = CmdLineFwk()
852 taskFactory = AddTaskFactoryMock()
854 qgraph = fwk.makeGraph(self.pipeline, args)
855 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
856 self.assertEqual(len(qgraph), self.nQuanta)
858 # run whole thing
859 fwk.runPipeline(qgraph, taskFactory, args)
860 # None of the actual tasks is executed
861 self.assertEqual(taskFactory.countExec, 0)
863 # check dataset types
864 butler.registry.refresh()
865 datasetTypes = list(butler.registry.queryDatasetTypes(re.compile("^_mock_.*")))
866 self.assertEqual(len(datasetTypes), self.nQuanta * 2)
868 def testMockTaskFailure(self):
869 """Test --mock option and configure one of the tasks to fail."""
870 args = _makeArgs(
871 butler_config=self.root,
872 input="test",
873 output="output",
874 mock=True,
875 register_dataset_types=True,
876 mock_configs=[
877 _ACTION_CONFIG("task3-mock:failCondition='detector = 0'"),
878 ],
879 fail_fast=True,
880 )
881 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
882 populateButler(self.pipeline, butler)
884 fwk = CmdLineFwk()
885 taskFactory = AddTaskFactoryMock()
887 qgraph = fwk.makeGraph(self.pipeline, args)
888 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
889 self.assertEqual(len(qgraph), self.nQuanta)
891 with self.assertRaises(MPGraphExecutorError) as cm:
892 fwk.runPipeline(qgraph, taskFactory, args)
894 self.assertIsNotNone(cm.exception.__cause__)
895 self.assertRegex(str(cm.exception.__cause__), "Simulated failure: task=task3")
897 def testSubgraph(self):
898 """Test successful execution of trivial quantum graph."""
899 args = _makeArgs(butler_config=self.root, input="test", output="output")
900 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
901 populateButler(self.pipeline, butler)
903 fwk = CmdLineFwk()
904 qgraph = fwk.makeGraph(self.pipeline, args)
906 # Select first two nodes for execution. This depends on node ordering
907 # which I assume is the same as execution order.
908 nNodes = 2
909 nodeIds = [node.nodeId for node in qgraph]
910 nodeIds = nodeIds[:nNodes]
912 self.assertEqual(len(qgraph.taskGraph), self.nQuanta)
913 self.assertEqual(len(qgraph), self.nQuanta)
915 with makeTmpFile(suffix=".qgraph") as tmpname, makeSQLiteRegistry(
916 universe=butler.registry.dimensions
917 ) as registryConfig:
918 with open(tmpname, "wb") as saveFile:
919 qgraph.save(saveFile)
921 args = _makeArgs(
922 qgraph=tmpname,
923 qgraph_node_id=nodeIds,
924 registryConfig=registryConfig,
925 execution_butler_location=None,
926 )
927 fwk = CmdLineFwk()
929 # load graph, should only read a subset
930 qgraph = fwk.makeGraph(pipeline=None, args=args)
931 self.assertEqual(len(qgraph), nNodes)
933 def testShowGraph(self):
934 """Test for --show options for quantum graph."""
935 nQuanta = 2
936 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
938 show = ShowInfo(["graph"])
939 show.show_graph_info(qgraph)
940 self.assertEqual(show.handled, {"graph"})
942 def testShowGraphWorkflow(self):
943 nQuanta = 2
944 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
946 show = ShowInfo(["workflow"])
947 show.show_graph_info(qgraph)
948 self.assertEqual(show.handled, {"workflow"})
950 # TODO: cannot test "uri" option presently, it instantiates
951 # butler from command line options and there is no way to pass butler
952 # mock to that code.
953 show = ShowInfo(["uri"])
954 with self.assertRaises(ValueError): # No args given
955 show.show_graph_info(qgraph)
957 def testSimpleQGraphDatastoreRecords(self):
958 """Test quantum graph generation with --qgraph-datastore-records."""
959 args = _makeArgs(
960 butler_config=self.root, input="test", output="output", qgraph_datastore_records=True
961 )
962 butler = makeSimpleButler(self.root, run=args.input, inMemory=False)
963 populateButler(self.pipeline, butler)
965 fwk = CmdLineFwk()
966 qgraph = fwk.makeGraph(self.pipeline, args)
967 self.assertEqual(len(qgraph), self.nQuanta)
968 for i, qnode in enumerate(qgraph):
969 quantum = qnode.quantum
970 self.assertIsNotNone(quantum.datastore_records)
971 # only the first quantum has a pre-existing input
972 if i == 0:
973 datastore_name = "FileDatastore@<butlerRoot>"
974 self.assertEqual(set(quantum.datastore_records.keys()), {datastore_name})
975 records_data = quantum.datastore_records[datastore_name]
976 records = dict(records_data.records)
977 self.assertEqual(len(records), 1)
978 _, records = records.popitem()
979 records = records["file_datastore_records"]
980 self.assertEqual(
981 [record.path for record in records],
982 ["test/add_dataset0/add_dataset0_INSTR_det0_test.pickle"],
983 )
984 else:
985 self.assertEqual(quantum.datastore_records, {})
988class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
989 pass
992def setup_module(module):
993 lsst.utils.tests.init()
996if __name__ == "__main__": 996 ↛ 997line 996 didn't jump to line 997, because the condition on line 996 was never true
997 lsst.utils.tests.init()
998 unittest.main()