Coverage for tests/test_cmdLineFwk.py : 21%

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 argparse
26import contextlib
27import logging
28import os
29import pickle
30import shutil
31import tempfile
32import unittest
34from lsst.ctrl.mpexec.cmdLineFwk import CmdLineFwk
35from lsst.ctrl.mpexec.cmdLineParser import (_ACTION_ADD_TASK, _ACTION_CONFIG,
36 _ACTION_CONFIG_FILE, _ACTION_ADD_INSTRUMENT)
37from lsst.daf.butler import Config, Quantum, Registry
38from lsst.daf.butler.registry import RegistryConfig
39from lsst.obs.base import Instrument
40import lsst.pex.config as pexConfig
41from lsst.pipe.base import (Pipeline, PipelineTask, PipelineTaskConfig,
42 QuantumGraph, QuantumGraphTaskNodes,
43 TaskDef, TaskFactory, PipelineTaskConnections)
44import lsst.pipe.base.connectionTypes as cT
45import lsst.utils.tests
46from testUtil import (AddTaskFactoryMock, makeSimpleQGraph)
49logging.basicConfig(level=logging.INFO)
51# Have to monkey-patch Instrument.fromName() to not retrieve non-existing
52# instrument from registry, these tests can run fine without actual instrument
53# and implementing full mock for Instrument is too complicated.
54Instrument.fromName = lambda name, reg: None 54 ↛ exitline 54 didn't run the lambda on line 54
57@contextlib.contextmanager
58def makeTmpFile(contents=None):
59 """Context manager for generating temporary file name.
61 Temporary file is deleted on exiting context.
63 Parameters
64 ----------
65 contents : `bytes`
66 Data to write into a file.
67 """
68 fd, tmpname = tempfile.mkstemp()
69 if contents:
70 os.write(fd, contents)
71 os.close(fd)
72 yield tmpname
73 with contextlib.suppress(OSError):
74 os.remove(tmpname)
77@contextlib.contextmanager
78def makeSQLiteRegistry():
79 """Context manager to create new empty registry database.
81 Yields
82 ------
83 config : `RegistryConfig`
84 Registry configuration for initialized registry database.
85 """
86 with makeTmpFile() as filename:
87 uri = f"sqlite:///{filename}"
88 config = RegistryConfig()
89 config["db"] = uri
90 Registry.fromConfig(config, create=True)
91 yield config
94class SimpleConnections(PipelineTaskConnections, dimensions=(),
95 defaultTemplates={"template": "simple"}):
96 schema = cT.InitInput(doc="Schema",
97 name="{template}schema",
98 storageClass="SourceCatalog")
101class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections):
102 field = pexConfig.Field(dtype=str, doc="arbitrary string")
104 def setDefaults(self):
105 PipelineTaskConfig.setDefaults(self)
108class TaskOne(PipelineTask):
109 ConfigClass = SimpleConfig
110 _DefaultName = "taskOne"
113class TaskTwo(PipelineTask):
114 ConfigClass = SimpleConfig
115 _DefaultName = "taskTwo"
118class TaskFactoryMock(TaskFactory):
119 def loadTaskClass(self, taskName):
120 if taskName == "TaskOne":
121 return TaskOne, "TaskOne"
122 elif taskName == "TaskTwo":
123 return TaskTwo, "TaskTwo"
125 def makeTask(self, taskClass, config, overrides, butler):
126 if config is None:
127 config = taskClass.ConfigClass()
128 if overrides:
129 overrides.applyTo(config)
130 return taskClass(config=config, butler=butler)
133def _makeArgs(pipeline=None, qgraph=None, pipeline_actions=(), order_pipeline=False,
134 save_pipeline="", save_qgraph="", save_single_quanta="",
135 pipeline_dot="", qgraph_dot="", registryConfig=None):
136 """Return parsed command line arguments.
138 Parameters
139 ----------
140 pipeline : `str`, optional
141 Name of the YAML file with pipeline.
142 qgraph : `str`, optional
143 Name of the pickle file with QGraph.
144 pipeline_actions : itrable of `cmdLinePArser._PipelineAction`, optional
145 order_pipeline : `bool`
146 save_pipeline : `str`
147 Name of the YAML file to store pipeline.
148 save_qgraph : `str`
149 Name of the pickle file to store QGraph.
150 save_single_quanta : `str`
151 Name of the pickle file pattern to store individual QGraph.
152 pipeline_dot : `str`
153 Name of the DOT file to write pipeline graph.
154 qgraph_dot : `str`
155 Name of the DOT file to write QGraph representation.
156 """
157 args = argparse.Namespace()
158 args.butler_config = Config()
159 if registryConfig:
160 args.butler_config["registry"] = registryConfig
161 # The default datastore has a relocatable root, so we need to specify
162 # some root here for it to use
163 args.butler_config.configFile = "."
164 args.pipeline = pipeline
165 args.qgraph = qgraph
166 args.pipeline_actions = pipeline_actions
167 args.order_pipeline = order_pipeline
168 args.save_pipeline = save_pipeline
169 args.save_qgraph = save_qgraph
170 args.save_single_quanta = save_single_quanta
171 args.pipeline_dot = pipeline_dot
172 args.qgraph_dot = qgraph_dot
173 args.input = ""
174 args.output = None
175 args.output_run = None
176 args.extend_run = False
177 args.replace_run = False
178 args.prune_replaced = False
179 args.register_dataset_types = False
180 args.skip_init_writes = False
181 args.no_versions = False
182 args.skip_existing = False
183 args.clobber_partial_outputs = False
184 args.init_only = False
185 args.processes = 1
186 args.profile = None
187 args.enableLsstDebug = False
188 args.graph_fixup = None
189 args.timeout = None
190 args.fail_fast = False
191 return args
194def _makeQGraph():
195 """Make a trivial QuantumGraph with one quantum.
197 The only thing that we need to do with this quantum graph is to pickle
198 it, the quanta in this graph are not usable for anything else.
200 Returns
201 -------
202 qgraph : `~lsst.pipe.base.QuantumGraph`
203 """
204 taskDef = TaskDef(taskName="taskOne", config=SimpleConfig())
205 quanta = [Quantum()]
206 taskNodes = QuantumGraphTaskNodes(taskDef=taskDef, quanta=quanta, initInputs={}, initOutputs={})
207 qgraph = QuantumGraph([taskNodes])
208 return qgraph
211class CmdLineFwkTestCase(unittest.TestCase):
212 """A test case for CmdLineFwk
213 """
215 def testMakePipeline(self):
216 """Tests for CmdLineFwk.makePipeline method
217 """
218 fwk = CmdLineFwk()
220 # make empty pipeline
221 args = _makeArgs()
222 pipeline = fwk.makePipeline(args)
223 self.assertIsInstance(pipeline, Pipeline)
224 self.assertEqual(len(pipeline), 0)
226 # few tests with serialization
227 with makeTmpFile() as tmpname:
228 # make empty pipeline and store it in a file
229 args = _makeArgs(save_pipeline=tmpname)
230 pipeline = fwk.makePipeline(args)
231 self.assertIsInstance(pipeline, Pipeline)
233 # read pipeline from a file
234 args = _makeArgs(pipeline=tmpname)
235 pipeline = fwk.makePipeline(args)
236 self.assertIsInstance(pipeline, Pipeline)
237 self.assertEqual(len(pipeline), 0)
239 # single task pipeline
240 actions = [
241 _ACTION_ADD_TASK("TaskOne:task1")
242 ]
243 args = _makeArgs(pipeline_actions=actions)
244 pipeline = fwk.makePipeline(args)
245 self.assertIsInstance(pipeline, Pipeline)
246 self.assertEqual(len(pipeline), 1)
248 # many task pipeline
249 actions = [
250 _ACTION_ADD_TASK("TaskOne:task1a"),
251 _ACTION_ADD_TASK("TaskTwo:task2"),
252 _ACTION_ADD_TASK("TaskOne:task1b")
253 ]
254 args = _makeArgs(pipeline_actions=actions)
255 pipeline = fwk.makePipeline(args)
256 self.assertIsInstance(pipeline, Pipeline)
257 self.assertEqual(len(pipeline), 3)
259 # single task pipeline with config overrides, cannot use TaskOne, need
260 # something that can be imported with `doImport()`
261 actions = [
262 _ACTION_ADD_TASK("testUtil.AddTask:task"),
263 _ACTION_CONFIG("task:addend=100")
264 ]
265 args = _makeArgs(pipeline_actions=actions)
266 pipeline = fwk.makePipeline(args)
267 taskDefs = list(pipeline.toExpandedPipeline())
268 self.assertEqual(len(taskDefs), 1)
269 self.assertEqual(taskDefs[0].config.addend, 100)
271 overrides = b"config.addend = 1000\n"
272 with makeTmpFile(overrides) as tmpname:
273 actions = [
274 _ACTION_ADD_TASK("testUtil.AddTask:task"),
275 _ACTION_CONFIG_FILE("task:" + tmpname)
276 ]
277 args = _makeArgs(pipeline_actions=actions)
278 pipeline = fwk.makePipeline(args)
279 taskDefs = list(pipeline.toExpandedPipeline())
280 self.assertEqual(len(taskDefs), 1)
281 self.assertEqual(taskDefs[0].config.addend, 1000)
283 # Check --instrument option, for now it only checks that it does not crash
284 actions = [
285 _ACTION_ADD_TASK("testUtil.AddTask:task"),
286 _ACTION_ADD_INSTRUMENT("Instrument")
287 ]
288 args = _makeArgs(pipeline_actions=actions)
289 pipeline = fwk.makePipeline(args)
291 def testMakeGraphFromPickle(self):
292 """Tests for CmdLineFwk.makeGraph method.
294 Only most trivial case is tested that does not do actual graph
295 building.
296 """
297 fwk = CmdLineFwk()
299 with makeTmpFile() as tmpname, makeSQLiteRegistry() as registryConfig:
301 # make non-empty graph and store it in a file
302 qgraph = _makeQGraph()
303 with open(tmpname, "wb") as pickleFile:
304 qgraph.save(pickleFile)
305 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
306 qgraph = fwk.makeGraph(None, args)
307 self.assertIsInstance(qgraph, QuantumGraph)
308 self.assertEqual(len(qgraph), 1)
310 # pickle with wrong object type
311 with open(tmpname, "wb") as pickleFile:
312 pickle.dump({}, pickleFile)
313 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
314 with self.assertRaises(TypeError):
315 fwk.makeGraph(None, args)
317 # reading empty graph from pickle should work but makeGraph()
318 # will return None and make a warning
319 qgraph = QuantumGraph()
320 with open(tmpname, "wb") as pickleFile:
321 qgraph.save(pickleFile)
322 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
323 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"):
324 # this also tests that warning is generated for empty graph
325 qgraph = fwk.makeGraph(None, args)
326 self.assertIs(qgraph, None)
328 def testShowPipeline(self):
329 """Test for --show options for pipeline.
330 """
331 fwk = CmdLineFwk()
333 actions = [
334 _ACTION_ADD_TASK("testUtil.AddTask:task"),
335 _ACTION_CONFIG("task:addend=100")
336 ]
337 args = _makeArgs(pipeline_actions=actions)
338 pipeline = fwk.makePipeline(args)
340 args.show = ["pipeline"]
341 fwk.showInfo(args, pipeline)
342 args.show = ["config"]
343 fwk.showInfo(args, pipeline)
344 args.show = ["history=task::addend"]
345 fwk.showInfo(args, pipeline)
346 args.show = ["tasks"]
347 fwk.showInfo(args, pipeline)
350class CmdLineFwkTestCaseWithButler(unittest.TestCase):
351 """A test case for CmdLineFwk
352 """
354 def setUp(self):
355 super().setUpClass()
356 self.root = tempfile.mkdtemp()
358 def tearDown(self):
359 shutil.rmtree(self.root, ignore_errors=True)
360 super().tearDownClass()
362 def testSimpleQGraph(self):
363 """Test successfull execution of trivial quantum graph.
364 """
366 nQuanta = 5
367 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
369 # should have one task and number of quanta
370 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
372 args = _makeArgs()
373 fwk = CmdLineFwk()
374 taskFactory = AddTaskFactoryMock()
376 # run whole thing
377 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
378 self.assertEqual(taskFactory.countExec, nQuanta)
380 def testSimpleQGraphSkipExisting(self):
381 """Test continuing execution of trivial quantum graph with --skip-existing.
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 # run remaining ones
400 taskFactory.stopAt = -1
401 args.skip_existing = True
402 args.no_versions = True
403 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
404 self.assertEqual(taskFactory.countExec, nQuanta)
406 def testSimpleQGraphPartialOutputsFail(self):
407 """Test continuing execution of trivial quantum graph with partial
408 outputs.
409 """
411 nQuanta = 5
412 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
414 # should have one task and number of quanta
415 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
417 args = _makeArgs()
418 fwk = CmdLineFwk()
419 taskFactory = AddTaskFactoryMock(stopAt=3)
421 # run first three quanta
422 with self.assertRaises(RuntimeError):
423 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
424 self.assertEqual(taskFactory.countExec, 3)
426 # drop one of the two outputs from one task
427 ref = butler._findDatasetRef("add2_dataset2", instrument="INSTR", detector=0)
428 self.assertIsNotNone(ref)
429 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True)
431 taskFactory.stopAt = -1
432 args.skip_existing = True
433 args.no_versions = True
434 excRe = "Registry inconsistency while checking for existing outputs.*"
435 with self.assertRaisesRegex(RuntimeError, excRe):
436 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
438 def testSimpleQGraphClobberPartialOutputs(self):
439 """Test continuing execution of trivial quantum graph with
440 --clobber-partial-outputs.
441 """
443 nQuanta = 5
444 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
446 # should have one task and number of quanta
447 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
449 args = _makeArgs()
450 fwk = CmdLineFwk()
451 taskFactory = AddTaskFactoryMock(stopAt=3)
453 # run first three quanta
454 with self.assertRaises(RuntimeError):
455 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
456 self.assertEqual(taskFactory.countExec, 3)
458 # drop one of the two outputs from one task
459 ref = butler._findDatasetRef("add2_dataset2", instrument="INSTR", detector=0)
460 self.assertIsNotNone(ref)
461 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True)
463 taskFactory.stopAt = -1
464 args.skip_existing = True
465 args.clobber_partial_outputs = True
466 args.no_versions = True
467 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
468 # number of executed quanta is incremented
469 self.assertEqual(taskFactory.countExec, nQuanta + 1)
471 def testShowGraph(self):
472 """Test for --show options for quantum graph.
473 """
474 fwk = CmdLineFwk()
476 nQuanta = 2
477 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root)
479 args = _makeArgs()
480 args.show = ["graph"]
481 fwk.showInfo(args, pipeline=None, graph=qgraph)
482 # TODO: cannot test "workflow" option presently, it instanciates
483 # butler from command line options and there is no way to pass butler
484 # mock to that code.
487class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
488 pass
491def setup_module(module):
492 lsst.utils.tests.init()
495if __name__ == "__main__": 495 ↛ 496line 495 didn't jump to line 496, because the condition on line 495 was never true
496 lsst.utils.tests.init()
497 unittest.main()