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 testUtil import (AddTaskFactoryMock, makeSimpleQGraph)
50logging.basicConfig(level=logging.INFO)
52# Have to monkey-patch Instrument.fromName() to not retrieve non-existing
53# instrument from registry, these tests can run fine without actual instrument
54# and implementing full mock for Instrument is too complicated.
55Instrument.fromName = lambda name, reg: None 55 ↛ exitline 55 didn't run the lambda on line 55
58@contextlib.contextmanager
59def makeTmpFile(contents=None):
60 """Context manager for generating temporary file name.
62 Temporary file is deleted on exiting context.
64 Parameters
65 ----------
66 contents : `bytes`
67 Data to write into a file.
68 """
69 fd, tmpname = tempfile.mkstemp()
70 if contents:
71 os.write(fd, contents)
72 os.close(fd)
73 yield tmpname
74 with contextlib.suppress(OSError):
75 os.remove(tmpname)
78@contextlib.contextmanager
79def temporaryDirectory():
80 """Context manager that creates and destroys temporary directory.
82 Difference from `tempfile.TemporaryDirectory` is that it ignores errors
83 when deleting a directory, which may happen with some filesystems.
84 """
85 tmpdir = tempfile.mkdtemp()
86 yield tmpdir
87 shutil.rmtree(tmpdir, ignore_errors=True)
90@contextlib.contextmanager
91def makeSQLiteRegistry(create=True):
92 """Context manager to create new empty registry database.
94 Yields
95 ------
96 config : `RegistryConfig`
97 Registry configuration for initialized registry database.
98 """
99 with temporaryDirectory() as tmpdir:
100 uri = f"sqlite:///{tmpdir}/gen3.sqlite"
101 config = RegistryConfig()
102 config["db"] = uri
103 if create:
104 Registry.fromConfig(config, create=True)
105 yield config
108class SimpleConnections(PipelineTaskConnections, dimensions=(),
109 defaultTemplates={"template": "simple"}):
110 schema = cT.InitInput(doc="Schema",
111 name="{template}schema",
112 storageClass="SourceCatalog")
115class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections):
116 field = pexConfig.Field(dtype=str, doc="arbitrary string")
118 def setDefaults(self):
119 PipelineTaskConfig.setDefaults(self)
122class TaskOne(PipelineTask):
123 ConfigClass = SimpleConfig
124 _DefaultName = "taskOne"
127class TaskTwo(PipelineTask):
128 ConfigClass = SimpleConfig
129 _DefaultName = "taskTwo"
132class TaskFactoryMock(TaskFactory):
133 def loadTaskClass(self, taskName):
134 if taskName == "TaskOne":
135 return TaskOne, "TaskOne"
136 elif taskName == "TaskTwo":
137 return TaskTwo, "TaskTwo"
139 def makeTask(self, taskClass, config, overrides, butler):
140 if config is None:
141 config = taskClass.ConfigClass()
142 if overrides:
143 overrides.applyTo(config)
144 return taskClass(config=config, butler=butler)
147def _makeArgs(cmd="run", registryConfig=None, **kwargs):
148 """Return parsed command line arguments.
150 By default butler_config is set to `Config` populated with some defaults,
151 it can be overriden completely by keyword argument.
153 Parameters
154 ----------
155 cmd : `str`, optional
156 Produce arguments for this pipetask command.
157 registryConfig : `RegistryConfig`, optional
158 Override for registry configuration.
159 **kwargs
160 Overrides for other arguments.
161 """
162 # call parser for "run" command to set defaults for all arguments
163 parser = parser_mod.makeParser()
164 args = parser.parse_args([cmd])
165 # override butler_config with our defaults
166 args.butler_config = Config()
167 if registryConfig:
168 args.butler_config["registry"] = registryConfig
169 # The default datastore has a relocatable root, so we need to specify
170 # some root here for it to use
171 args.butler_config.configFile = "."
172 # override arguments from keyword parameters
173 for key, value in kwargs.items():
174 setattr(args, key, value)
175 return args
178def _makeQGraph():
179 """Make a trivial QuantumGraph with one quantum.
181 The only thing that we need to do with this quantum graph is to pickle
182 it, the quanta in this graph are not usable for anything else.
184 Returns
185 -------
186 qgraph : `~lsst.pipe.base.QuantumGraph`
187 """
188 taskDef = TaskDef(taskName="taskOne", config=SimpleConfig())
189 quanta = [Quantum()]
190 taskNodes = QuantumGraphTaskNodes(taskDef=taskDef, quanta=quanta, initInputs={}, initOutputs={})
191 qgraph = QuantumGraph([taskNodes])
192 return qgraph
195class CmdLineFwkTestCase(unittest.TestCase):
196 """A test case for CmdLineFwk
197 """
199 def testMakePipeline(self):
200 """Tests for CmdLineFwk.makePipeline method
201 """
202 fwk = CmdLineFwk()
204 # make empty pipeline
205 args = _makeArgs()
206 pipeline = fwk.makePipeline(args)
207 self.assertIsInstance(pipeline, Pipeline)
208 self.assertEqual(len(pipeline), 0)
210 # few tests with serialization
211 with makeTmpFile() as tmpname:
212 # make empty pipeline and store it in a file
213 args = _makeArgs(save_pipeline=tmpname)
214 pipeline = fwk.makePipeline(args)
215 self.assertIsInstance(pipeline, Pipeline)
217 # read pipeline from a file
218 args = _makeArgs(pipeline=tmpname)
219 pipeline = fwk.makePipeline(args)
220 self.assertIsInstance(pipeline, Pipeline)
221 self.assertEqual(len(pipeline), 0)
223 # single task pipeline
224 actions = [
225 _ACTION_ADD_TASK("TaskOne:task1")
226 ]
227 args = _makeArgs(pipeline_actions=actions)
228 pipeline = fwk.makePipeline(args)
229 self.assertIsInstance(pipeline, Pipeline)
230 self.assertEqual(len(pipeline), 1)
232 # many task pipeline
233 actions = [
234 _ACTION_ADD_TASK("TaskOne:task1a"),
235 _ACTION_ADD_TASK("TaskTwo:task2"),
236 _ACTION_ADD_TASK("TaskOne:task1b")
237 ]
238 args = _makeArgs(pipeline_actions=actions)
239 pipeline = fwk.makePipeline(args)
240 self.assertIsInstance(pipeline, Pipeline)
241 self.assertEqual(len(pipeline), 3)
243 # single task pipeline with config overrides, cannot use TaskOne, need
244 # something that can be imported with `doImport()`
245 actions = [
246 _ACTION_ADD_TASK("testUtil.AddTask:task"),
247 _ACTION_CONFIG("task:addend=100")
248 ]
249 args = _makeArgs(pipeline_actions=actions)
250 pipeline = fwk.makePipeline(args)
251 taskDefs = list(pipeline.toExpandedPipeline())
252 self.assertEqual(len(taskDefs), 1)
253 self.assertEqual(taskDefs[0].config.addend, 100)
255 overrides = b"config.addend = 1000\n"
256 with makeTmpFile(overrides) as tmpname:
257 actions = [
258 _ACTION_ADD_TASK("testUtil.AddTask:task"),
259 _ACTION_CONFIG_FILE("task:" + tmpname)
260 ]
261 args = _makeArgs(pipeline_actions=actions)
262 pipeline = fwk.makePipeline(args)
263 taskDefs = list(pipeline.toExpandedPipeline())
264 self.assertEqual(len(taskDefs), 1)
265 self.assertEqual(taskDefs[0].config.addend, 1000)
267 # Check --instrument option, for now it only checks that it does not crash
268 actions = [
269 _ACTION_ADD_TASK("testUtil.AddTask:task"),
270 _ACTION_ADD_INSTRUMENT("Instrument")
271 ]
272 args = _makeArgs(pipeline_actions=actions)
273 pipeline = fwk.makePipeline(args)
275 def testMakeGraphFromPickle(self):
276 """Tests for CmdLineFwk.makeGraph method.
278 Only most trivial case is tested that does not do actual graph
279 building.
280 """
281 fwk = CmdLineFwk()
283 with makeTmpFile() as tmpname, makeSQLiteRegistry() as registryConfig:
285 # make non-empty graph and store it in a file
286 qgraph = _makeQGraph()
287 with open(tmpname, "wb") as pickleFile:
288 qgraph.save(pickleFile)
289 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
290 qgraph = fwk.makeGraph(None, args)
291 self.assertIsInstance(qgraph, QuantumGraph)
292 self.assertEqual(len(qgraph), 1)
294 # pickle with wrong object type
295 with open(tmpname, "wb") as pickleFile:
296 pickle.dump({}, pickleFile)
297 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
298 with self.assertRaises(TypeError):
299 fwk.makeGraph(None, args)
301 # reading empty graph from pickle should work but makeGraph()
302 # will return None and make a warning
303 qgraph = QuantumGraph()
304 with open(tmpname, "wb") as pickleFile:
305 qgraph.save(pickleFile)
306 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
307 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"):
308 # this also tests that warning is generated for empty graph
309 qgraph = fwk.makeGraph(None, args)
310 self.assertIs(qgraph, None)
312 def testShowPipeline(self):
313 """Test for --show options for pipeline.
314 """
315 fwk = CmdLineFwk()
317 actions = [
318 _ACTION_ADD_TASK("testUtil.AddTask:task"),
319 _ACTION_CONFIG("task:addend=100")
320 ]
321 args = _makeArgs(pipeline_actions=actions)
322 pipeline = fwk.makePipeline(args)
324 args.show = ["pipeline"]
325 fwk.showInfo(args, pipeline)
326 args.show = ["config"]
327 fwk.showInfo(args, pipeline)
328 args.show = ["history=task::addend"]
329 fwk.showInfo(args, pipeline)
330 args.show = ["tasks"]
331 fwk.showInfo(args, pipeline)
334class CmdLineFwkTestCaseWithButler(unittest.TestCase):
335 """A test case for CmdLineFwk
336 """
338 def setUp(self):
339 super().setUpClass()
340 self.root = tempfile.mkdtemp()
342 def tearDown(self):
343 shutil.rmtree(self.root, ignore_errors=True)
344 super().tearDownClass()
346 def testSimpleQGraph(self):
347 """Test successfull execution of trivial quantum graph.
348 """
350 nQuanta = 5
351 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
353 # should have one task and number of quanta
354 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
356 args = _makeArgs()
357 fwk = CmdLineFwk()
358 taskFactory = AddTaskFactoryMock()
360 # run whole thing
361 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
362 self.assertEqual(taskFactory.countExec, nQuanta)
364 def testSimpleQGraphSkipExisting(self):
365 """Test continuing execution of trivial quantum graph with --skip-existing.
366 """
368 nQuanta = 5
369 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
371 # should have one task and number of quanta
372 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
374 args = _makeArgs()
375 fwk = CmdLineFwk()
376 taskFactory = AddTaskFactoryMock(stopAt=3)
378 # run first three quanta
379 with self.assertRaises(RuntimeError):
380 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
381 self.assertEqual(taskFactory.countExec, 3)
383 # run remaining ones
384 taskFactory.stopAt = -1
385 args.skip_existing = True
386 args.no_versions = True
387 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
388 self.assertEqual(taskFactory.countExec, nQuanta)
390 def testSimpleQGraphPartialOutputsFail(self):
391 """Test continuing execution of trivial quantum graph with partial
392 outputs.
393 """
395 nQuanta = 5
396 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
398 # should have one task and number of quanta
399 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
401 args = _makeArgs()
402 fwk = CmdLineFwk()
403 taskFactory = AddTaskFactoryMock(stopAt=3)
405 # run first three quanta
406 with self.assertRaises(RuntimeError):
407 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
408 self.assertEqual(taskFactory.countExec, 3)
410 # drop one of the two outputs from one task
411 ref = butler._findDatasetRef("add2_dataset2", instrument="INSTR", detector=0)
412 self.assertIsNotNone(ref)
413 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True)
415 taskFactory.stopAt = -1
416 args.skip_existing = True
417 args.no_versions = True
418 excRe = "Registry inconsistency while checking for existing outputs.*"
419 with self.assertRaisesRegex(RuntimeError, excRe):
420 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
422 def testSimpleQGraphClobberPartialOutputs(self):
423 """Test continuing execution of trivial quantum graph with
424 --clobber-partial-outputs.
425 """
427 nQuanta = 5
428 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
430 # should have one task and number of quanta
431 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
433 args = _makeArgs()
434 fwk = CmdLineFwk()
435 taskFactory = AddTaskFactoryMock(stopAt=3)
437 # run first three quanta
438 with self.assertRaises(RuntimeError):
439 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
440 self.assertEqual(taskFactory.countExec, 3)
442 # drop one of the two outputs from one task
443 ref = butler._findDatasetRef("add2_dataset2", instrument="INSTR", detector=0)
444 self.assertIsNotNone(ref)
445 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True)
447 taskFactory.stopAt = -1
448 args.skip_existing = True
449 args.clobber_partial_outputs = True
450 args.no_versions = True
451 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
452 # number of executed quanta is incremented
453 self.assertEqual(taskFactory.countExec, nQuanta + 1)
455 def testSimpleQGraphReplaceRun(self):
456 """Test repeated execution of trivial quantum graph with
457 --replace-run.
458 """
460 # need non-memory registry in this case
461 nQuanta = 5
462 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root, inMemory=False)
464 # should have one task and number of quanta
465 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
467 fwk = CmdLineFwk()
468 taskFactory = AddTaskFactoryMock()
470 # run whole thing
471 args = _makeArgs(
472 butler_config=self.root,
473 input="test",
474 output="output",
475 output_run="output/run1")
476 # deep copy is needed because quanta are updated in place
477 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
478 self.assertEqual(taskFactory.countExec, nQuanta)
480 # need to refresh collections explicitly (or make new butler/registry)
481 butler.registry._collections.refresh()
482 collections = set(butler.registry.queryCollections(...))
483 self.assertEqual(collections, {"test", "output", "output/run1"})
485 # number of datasets written by pipeline:
486 # - nQuanta of init_outputs
487 # - nQuanta of configs
488 # - packages (single dataset)
489 # - nQuanta * two output datasets
490 # - nQuanta of metadata
491 n_outputs = nQuanta * 5 + 1
492 refs = butler.registry.queryDatasets(..., collections="output/run1")
493 self.assertEqual(len(list(refs)), n_outputs)
495 # re-run with --replace-run (--inputs is not compatible)
496 args.input = None
497 args.replace_run = True
498 args.output_run = "output/run2"
499 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
501 butler.registry._collections.refresh()
502 collections = set(butler.registry.queryCollections(...))
503 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2"})
505 # new output collection
506 refs = butler.registry.queryDatasets(..., collections="output/run2")
507 self.assertEqual(len(list(refs)), n_outputs)
509 # old output collection is still there
510 refs = butler.registry.queryDatasets(..., collections="output/run1")
511 self.assertEqual(len(list(refs)), n_outputs)
513 # re-run with --replace-run and --prune-replaced=unstore
514 args.input = None
515 args.replace_run = True
516 args.prune_replaced = "unstore"
517 args.output_run = "output/run3"
518 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
520 butler.registry._collections.refresh()
521 collections = set(butler.registry.queryCollections(...))
522 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run3"})
524 # new output collection
525 refs = butler.registry.queryDatasets(..., collections="output/run3")
526 self.assertEqual(len(list(refs)), n_outputs)
528 # old output collection is still there, and it has all datasets but
529 # they are not in datastore
530 refs = butler.registry.queryDatasets(..., collections="output/run2")
531 refs = list(refs)
532 self.assertEqual(len(refs), n_outputs)
533 with self.assertRaises(FileNotFoundError):
534 butler.get(refs[0], collections="output/run2")
536 # re-run with --replace-run and --prune-replaced=purge
537 args.input = None
538 args.replace_run = True
539 args.prune_replaced = "purge"
540 args.output_run = "output/run4"
541 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args)
543 butler.registry._collections.refresh()
544 collections = set(butler.registry.queryCollections(...))
545 # output/run3 should disappear now
546 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"})
548 # new output collection
549 refs = butler.registry.queryDatasets(..., collections="output/run4")
550 self.assertEqual(len(list(refs)), n_outputs)
552 def testShowGraph(self):
553 """Test for --show options for quantum graph.
554 """
555 fwk = CmdLineFwk()
557 nQuanta = 2
558 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
560 args = _makeArgs(show=["graph"])
561 fwk.showInfo(args, pipeline=None, graph=qgraph)
562 # TODO: cannot test "workflow" option presently, it instanciates
563 # butler from command line options and there is no way to pass butler
564 # mock to that code.
567class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
568 pass
571def setup_module(module):
572 lsst.utils.tests.init()
575if __name__ == "__main__": 575 ↛ 576line 575 didn't jump to line 576, because the condition on line 575 was never true
576 lsst.utils.tests.init()
577 unittest.main()