Coverage for tests/test_cmdLineFwk.py : 20%

Hot-keys 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 shutil
31import tempfile
32import unittest
34from lsst.ctrl.mpexec.cmdLineFwk import CmdLineFwk
35import lsst.ctrl.mpexec.cmdLineParser as parser_mod
36from lsst.ctrl.mpexec.cmdLineParser import (_ACTION_ADD_TASK, _ACTION_CONFIG,
37 _ACTION_CONFIG_FILE, _ACTION_ADD_INSTRUMENT)
38from lsst.daf.butler import Config, Quantum, Registry
39from lsst.daf.butler.registry import RegistryConfig
40from lsst.obs.base import Instrument
41import lsst.pex.config as pexConfig
42from lsst.pipe.base import (Pipeline, PipelineTask, PipelineTaskConfig,
43 QuantumGraph, QuantumGraphTaskNodes,
44 TaskDef, TaskFactory, PipelineTaskConnections)
45import lsst.pipe.base.connectionTypes as cT
46import lsst.utils.tests
47from lsst.pipe.base.tests.simpleQGraph import (AddTaskFactoryMock, makeSimpleQGraph)
48from lsst.utils.tests import temporaryDirectory
51logging.basicConfig(level=logging.INFO)
53# Have to monkey-patch Instrument.fromName() to not retrieve non-existing
54# instrument from registry, these tests can run fine without actual instrument
55# and implementing full mock for Instrument is too complicated.
56Instrument.fromName = lambda name, reg: None 56 ↛ exitline 56 didn't run the lambda on line 56
59@contextlib.contextmanager
60def makeTmpFile(contents=None):
61 """Context manager for generating temporary file name.
63 Temporary file is deleted on exiting context.
65 Parameters
66 ----------
67 contents : `bytes`
68 Data to write into a file.
69 """
70 fd, tmpname = tempfile.mkstemp()
71 if contents:
72 os.write(fd, contents)
73 os.close(fd)
74 yield tmpname
75 with contextlib.suppress(OSError):
76 os.remove(tmpname)
79@contextlib.contextmanager
80def makeSQLiteRegistry(create=True):
81 """Context manager to create new empty registry database.
83 Yields
84 ------
85 config : `RegistryConfig`
86 Registry configuration for initialized registry database.
87 """
88 with temporaryDirectory() as tmpdir:
89 uri = f"sqlite:///{tmpdir}/gen3.sqlite"
90 config = RegistryConfig()
91 config["db"] = uri
92 if create:
93 Registry.fromConfig(config, create=True)
94 yield config
97class SimpleConnections(PipelineTaskConnections, dimensions=(),
98 defaultTemplates={"template": "simple"}):
99 schema = cT.InitInput(doc="Schema",
100 name="{template}schema",
101 storageClass="SourceCatalog")
104class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections):
105 field = pexConfig.Field(dtype=str, doc="arbitrary string")
107 def setDefaults(self):
108 PipelineTaskConfig.setDefaults(self)
111class TaskOne(PipelineTask):
112 ConfigClass = SimpleConfig
113 _DefaultName = "taskOne"
116class TaskTwo(PipelineTask):
117 ConfigClass = SimpleConfig
118 _DefaultName = "taskTwo"
121class TaskFactoryMock(TaskFactory):
122 def loadTaskClass(self, taskName):
123 if taskName == "TaskOne":
124 return TaskOne, "TaskOne"
125 elif taskName == "TaskTwo":
126 return TaskTwo, "TaskTwo"
128 def makeTask(self, taskClass, config, overrides, butler):
129 if config is None:
130 config = taskClass.ConfigClass()
131 if overrides:
132 overrides.applyTo(config)
133 return taskClass(config=config, butler=butler)
136def _makeArgs(cmd="run", registryConfig=None, **kwargs):
137 """Return parsed command line arguments.
139 By default butler_config is set to `Config` populated with some defaults,
140 it can be overriden completely by keyword argument.
142 Parameters
143 ----------
144 cmd : `str`, optional
145 Produce arguments for this pipetask command.
146 registryConfig : `RegistryConfig`, optional
147 Override for registry configuration.
148 **kwargs
149 Overrides for other arguments.
150 """
151 # call parser for "run" command to set defaults for all arguments
152 parser = parser_mod.makeParser()
153 args = parser.parse_args([cmd])
154 # override butler_config with our defaults
155 args.butler_config = Config()
156 if registryConfig:
157 args.butler_config["registry"] = registryConfig
158 # The default datastore has a relocatable root, so we need to specify
159 # some root here for it to use
160 args.butler_config.configFile = "."
161 # override arguments from keyword parameters
162 for key, value in kwargs.items():
163 setattr(args, key, value)
164 return args
167def _makeQGraph():
168 """Make a trivial QuantumGraph with one quantum.
170 The only thing that we need to do with this quantum graph is to pickle
171 it, the quanta in this graph are not usable for anything else.
173 Returns
174 -------
175 qgraph : `~lsst.pipe.base.QuantumGraph`
176 """
177 taskDef = TaskDef(taskName="taskOne", config=SimpleConfig())
178 quanta = [Quantum()]
179 taskNodes = QuantumGraphTaskNodes(taskDef=taskDef, quanta=quanta, initInputs={}, initOutputs={})
180 qgraph = QuantumGraph([taskNodes])
181 return qgraph
184class CmdLineFwkTestCase(unittest.TestCase):
185 """A test case for CmdLineFwk
186 """
188 def testMakePipeline(self):
189 """Tests for CmdLineFwk.makePipeline method
190 """
191 fwk = CmdLineFwk()
193 # make empty pipeline
194 args = _makeArgs()
195 pipeline = fwk.makePipeline(args)
196 self.assertIsInstance(pipeline, Pipeline)
197 self.assertEqual(len(pipeline), 0)
199 # few tests with serialization
200 with makeTmpFile() as tmpname:
201 # make empty pipeline and store it in a file
202 args = _makeArgs(save_pipeline=tmpname)
203 pipeline = fwk.makePipeline(args)
204 self.assertIsInstance(pipeline, Pipeline)
206 # read pipeline from a file
207 args = _makeArgs(pipeline=tmpname)
208 pipeline = fwk.makePipeline(args)
209 self.assertIsInstance(pipeline, Pipeline)
210 self.assertEqual(len(pipeline), 0)
212 # single task pipeline
213 actions = [
214 _ACTION_ADD_TASK("TaskOne:task1")
215 ]
216 args = _makeArgs(pipeline_actions=actions)
217 pipeline = fwk.makePipeline(args)
218 self.assertIsInstance(pipeline, Pipeline)
219 self.assertEqual(len(pipeline), 1)
221 # many task pipeline
222 actions = [
223 _ACTION_ADD_TASK("TaskOne:task1a"),
224 _ACTION_ADD_TASK("TaskTwo:task2"),
225 _ACTION_ADD_TASK("TaskOne:task1b")
226 ]
227 args = _makeArgs(pipeline_actions=actions)
228 pipeline = fwk.makePipeline(args)
229 self.assertIsInstance(pipeline, Pipeline)
230 self.assertEqual(len(pipeline), 3)
232 # single task pipeline with config overrides, cannot use TaskOne, need
233 # something that can be imported with `doImport()`
234 actions = [
235 _ACTION_ADD_TASK("lsst.pipe.base.tests.simpleQGraph.AddTask:task"),
236 _ACTION_CONFIG("task:addend=100")
237 ]
238 args = _makeArgs(pipeline_actions=actions)
239 pipeline = fwk.makePipeline(args)
240 taskDefs = list(pipeline.toExpandedPipeline())
241 self.assertEqual(len(taskDefs), 1)
242 self.assertEqual(taskDefs[0].config.addend, 100)
244 overrides = b"config.addend = 1000\n"
245 with makeTmpFile(overrides) as tmpname:
246 actions = [
247 _ACTION_ADD_TASK("lsst.pipe.base.tests.simpleQGraph.AddTask:task"),
248 _ACTION_CONFIG_FILE("task:" + tmpname)
249 ]
250 args = _makeArgs(pipeline_actions=actions)
251 pipeline = fwk.makePipeline(args)
252 taskDefs = list(pipeline.toExpandedPipeline())
253 self.assertEqual(len(taskDefs), 1)
254 self.assertEqual(taskDefs[0].config.addend, 1000)
256 # Check --instrument option, for now it only checks that it does not crash
257 actions = [
258 _ACTION_ADD_TASK("lsst.pipe.base.tests.simpleQGraph.AddTask:task"),
259 _ACTION_ADD_INSTRUMENT("Instrument")
260 ]
261 args = _makeArgs(pipeline_actions=actions)
262 pipeline = fwk.makePipeline(args)
264 def testMakeGraphFromPickle(self):
265 """Tests for CmdLineFwk.makeGraph method.
267 Only most trivial case is tested that does not do actual graph
268 building.
269 """
270 fwk = CmdLineFwk()
272 with makeTmpFile() as tmpname, makeSQLiteRegistry() as registryConfig:
274 # make non-empty graph and store it in a file
275 qgraph = _makeQGraph()
276 with open(tmpname, "wb") as pickleFile:
277 qgraph.save(pickleFile)
278 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
279 qgraph = fwk.makeGraph(None, args)
280 self.assertIsInstance(qgraph, QuantumGraph)
281 self.assertEqual(len(qgraph), 1)
283 # pickle with wrong object type
284 with open(tmpname, "wb") as pickleFile:
285 pickle.dump({}, pickleFile)
286 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
287 with self.assertRaises(TypeError):
288 fwk.makeGraph(None, args)
290 # reading empty graph from pickle should work but makeGraph()
291 # will return None and make a warning
292 qgraph = QuantumGraph()
293 with open(tmpname, "wb") as pickleFile:
294 qgraph.save(pickleFile)
295 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
296 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"):
297 # this also tests that warning is generated for empty graph
298 qgraph = fwk.makeGraph(None, args)
299 self.assertIs(qgraph, None)
301 def testShowPipeline(self):
302 """Test for --show options for pipeline.
303 """
304 fwk = CmdLineFwk()
306 actions = [
307 _ACTION_ADD_TASK("lsst.pipe.base.tests.simpleQGraph.AddTask:task"),
308 _ACTION_CONFIG("task:addend=100")
309 ]
310 args = _makeArgs(pipeline_actions=actions)
311 pipeline = fwk.makePipeline(args)
313 args.show = ["pipeline"]
314 fwk.showInfo(args, pipeline)
315 args.show = ["config"]
316 fwk.showInfo(args, pipeline)
317 args.show = ["history=task::addend"]
318 fwk.showInfo(args, pipeline)
319 args.show = ["tasks"]
320 fwk.showInfo(args, pipeline)
323class CmdLineFwkTestCaseWithButler(unittest.TestCase):
324 """A test case for CmdLineFwk
325 """
327 def setUp(self):
328 super().setUpClass()
329 self.root = tempfile.mkdtemp()
331 def tearDown(self):
332 shutil.rmtree(self.root, ignore_errors=True)
333 super().tearDownClass()
335 def testSimpleQGraph(self):
336 """Test successfull execution of trivial quantum graph.
337 """
339 nQuanta = 5
340 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
342 # should have one task and number of quanta
343 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
345 args = _makeArgs()
346 fwk = CmdLineFwk()
347 taskFactory = AddTaskFactoryMock()
349 # run whole thing
350 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
351 self.assertEqual(taskFactory.countExec, nQuanta)
353 def testSimpleQGraphSkipExisting(self):
354 """Test continuing execution of trivial quantum graph with --skip-existing.
355 """
357 nQuanta = 5
358 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
360 # should have one task and number of quanta
361 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
363 args = _makeArgs()
364 fwk = CmdLineFwk()
365 taskFactory = AddTaskFactoryMock(stopAt=3)
367 # run first three quanta
368 with self.assertRaises(RuntimeError):
369 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
370 self.assertEqual(taskFactory.countExec, 3)
372 # run remaining ones
373 taskFactory.stopAt = -1
374 args.skip_existing = True
375 args.no_versions = True
376 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
377 self.assertEqual(taskFactory.countExec, nQuanta)
379 def testSimpleQGraphPartialOutputsFail(self):
380 """Test continuing execution of trivial quantum graph with partial
381 outputs.
382 """
384 nQuanta = 5
385 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
387 # should have one task and number of quanta
388 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
390 args = _makeArgs()
391 fwk = CmdLineFwk()
392 taskFactory = AddTaskFactoryMock(stopAt=3)
394 # run first three quanta
395 with self.assertRaises(RuntimeError):
396 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
397 self.assertEqual(taskFactory.countExec, 3)
399 # drop one of the two outputs from one task
400 ref = butler._findDatasetRef("add2_dataset2", instrument="INSTR", detector=0)
401 self.assertIsNotNone(ref)
402 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True)
404 taskFactory.stopAt = -1
405 args.skip_existing = True
406 args.no_versions = True
407 excRe = "Registry inconsistency while checking for existing outputs.*"
408 with self.assertRaisesRegex(RuntimeError, excRe):
409 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
411 def testSimpleQGraphClobberPartialOutputs(self):
412 """Test continuing execution of trivial quantum graph with
413 --clobber-partial-outputs.
414 """
416 nQuanta = 5
417 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
419 # should have one task and number of quanta
420 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
422 args = _makeArgs()
423 fwk = CmdLineFwk()
424 taskFactory = AddTaskFactoryMock(stopAt=3)
426 # run first three quanta
427 with self.assertRaises(RuntimeError):
428 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
429 self.assertEqual(taskFactory.countExec, 3)
431 # drop one of the two outputs from one task
432 ref = butler._findDatasetRef("add2_dataset2", instrument="INSTR", detector=0)
433 self.assertIsNotNone(ref)
434 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True)
436 taskFactory.stopAt = -1
437 args.skip_existing = True
438 args.clobber_partial_outputs = True
439 args.no_versions = True
440 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
441 # number of executed quanta is incremented
442 self.assertEqual(taskFactory.countExec, nQuanta + 1)
444 def testSimpleQGraphReplaceRun(self):
445 """Test repeated execution of trivial quantum graph with
446 --replace-run.
447 """
449 # need non-memory registry in this case
450 nQuanta = 5
451 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root, inMemory=False)
453 # should have one task and number of quanta
454 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
456 fwk = CmdLineFwk()
457 taskFactory = AddTaskFactoryMock()
459 # run whole thing
460 args = _makeArgs(
461 butler_config=self.root,
462 input="test",
463 output="output",
464 output_run="output/run1")
465 # deep copy is needed because quanta are updated in place
466 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
467 self.assertEqual(taskFactory.countExec, nQuanta)
469 # need to refresh collections explicitly (or make new butler/registry)
470 butler.registry._collections.refresh()
471 collections = set(butler.registry.queryCollections(...))
472 self.assertEqual(collections, {"test", "output", "output/run1"})
474 # number of datasets written by pipeline:
475 # - nQuanta of init_outputs
476 # - nQuanta of configs
477 # - packages (single dataset)
478 # - nQuanta * two output datasets
479 # - nQuanta of metadata
480 n_outputs = nQuanta * 5 + 1
481 refs = butler.registry.queryDatasets(..., collections="output/run1")
482 self.assertEqual(len(list(refs)), n_outputs)
484 # re-run with --replace-run (--inputs is not compatible)
485 args.input = None
486 args.replace_run = True
487 args.output_run = "output/run2"
488 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
490 butler.registry._collections.refresh()
491 collections = set(butler.registry.queryCollections(...))
492 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2"})
494 # new output collection
495 refs = butler.registry.queryDatasets(..., collections="output/run2")
496 self.assertEqual(len(list(refs)), n_outputs)
498 # old output collection is still there
499 refs = butler.registry.queryDatasets(..., collections="output/run1")
500 self.assertEqual(len(list(refs)), n_outputs)
502 # re-run with --replace-run and --prune-replaced=unstore
503 args.input = None
504 args.replace_run = True
505 args.prune_replaced = "unstore"
506 args.output_run = "output/run3"
507 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
509 butler.registry._collections.refresh()
510 collections = set(butler.registry.queryCollections(...))
511 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run3"})
513 # new output collection
514 refs = butler.registry.queryDatasets(..., collections="output/run3")
515 self.assertEqual(len(list(refs)), n_outputs)
517 # old output collection is still there, and it has all datasets but
518 # they are not in datastore
519 refs = butler.registry.queryDatasets(..., collections="output/run2")
520 refs = list(refs)
521 self.assertEqual(len(refs), n_outputs)
522 with self.assertRaises(FileNotFoundError):
523 butler.get(refs[0], collections="output/run2")
525 # re-run with --replace-run and --prune-replaced=purge
526 args.input = None
527 args.replace_run = True
528 args.prune_replaced = "purge"
529 args.output_run = "output/run4"
530 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
532 butler.registry._collections.refresh()
533 collections = set(butler.registry.queryCollections(...))
534 # output/run3 should disappear now
535 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
537 # new output collection
538 refs = butler.registry.queryDatasets(..., collections="output/run4")
539 self.assertEqual(len(list(refs)), n_outputs)
541 def testShowGraph(self):
542 """Test for --show options for quantum graph.
543 """
544 fwk = CmdLineFwk()
546 nQuanta = 2
547 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
549 args = _makeArgs(show=["graph"])
550 fwk.showInfo(args, pipeline=None, graph=qgraph)
551 # TODO: cannot test "workflow" option presently, it instanciates
552 # butler from command line options and there is no way to pass butler
553 # mock to that code.
556class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
557 pass
560def setup_module(module):
561 lsst.utils.tests.init()
564if __name__ == "__main__": 564 ↛ 565line 564 didn't jump to line 565, because the condition on line 564 was never true
565 lsst.utils.tests.init()
566 unittest.main()