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 tempfile
31import unittest
33from lsst.ctrl.mpexec.cmdLineFwk import CmdLineFwk
34from lsst.ctrl.mpexec.cmdLineParser import (_ACTION_ADD_TASK, _ACTION_CONFIG,
35 _ACTION_CONFIG_FILE, _ACTION_ADD_INSTRUMENT)
36from lsst.daf.butler import Config, Quantum, Registry
37from lsst.daf.butler.registry import RegistryConfig
38import lsst.pex.config as pexConfig
39from lsst.pipe.base import (Pipeline, PipelineTask, PipelineTaskConfig,
40 QuantumGraph, QuantumGraphTaskNodes,
41 TaskDef, TaskFactory, PipelineTaskConnections)
42import lsst.pipe.base.connectionTypes as cT
43import lsst.utils.tests
44from testUtil import (AddTask, AddTaskFactoryMock, makeSimpleQGraph)
47logging.basicConfig(level=logging.INFO)
50@contextlib.contextmanager
51def makeTmpFile(contents=None):
52 """Context manager for generating temporary file name.
54 Temporary file is deleted on exiting context.
56 Parameters
57 ----------
58 contents : `bytes`
59 Data to write into a file.
60 """
61 fd, tmpname = tempfile.mkstemp()
62 if contents:
63 os.write(fd, contents)
64 os.close(fd)
65 yield tmpname
66 with contextlib.suppress(OSError):
67 os.remove(tmpname)
70@contextlib.contextmanager
71def makeSQLiteRegistry():
72 """Context manager to create new empty registry database.
74 Yields
75 ------
76 config : `RegistryConfig`
77 Registry configuration for initialized registry database.
78 """
79 with makeTmpFile() as filename:
80 uri = f"sqlite:///{filename}"
81 config = RegistryConfig()
82 config["db"] = uri
83 Registry.fromConfig(config, create=True)
84 yield config
87class SimpleConnections(PipelineTaskConnections, dimensions=(),
88 defaultTemplates={"template": "simple"}):
89 schema = cT.InitInput(doc="Schema",
90 name="{template}schema",
91 storageClass="SourceCatalog")
94class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections):
95 field = pexConfig.Field(dtype=str, doc="arbitrary string")
97 def setDefaults(self):
98 PipelineTaskConfig.setDefaults(self)
101class TaskOne(PipelineTask):
102 ConfigClass = SimpleConfig
103 _DefaultName = "taskOne"
106class TaskTwo(PipelineTask):
107 ConfigClass = SimpleConfig
108 _DefaultName = "taskTwo"
111class TaskFactoryMock(TaskFactory):
112 def loadTaskClass(self, taskName):
113 if taskName == "TaskOne":
114 return TaskOne, "TaskOne"
115 elif taskName == "TaskTwo":
116 return TaskTwo, "TaskTwo"
118 def makeTask(self, taskClass, config, overrides, butler):
119 if config is None:
120 config = taskClass.ConfigClass()
121 if overrides:
122 overrides.applyTo(config)
123 return taskClass(config=config, butler=butler)
126def _makeArgs(pipeline=None, qgraph=None, pipeline_actions=(), order_pipeline=False,
127 save_pipeline="", save_qgraph="", save_single_quanta="",
128 pipeline_dot="", qgraph_dot="", registryConfig=None):
129 """Return parsed command line arguments.
131 Parameters
132 ----------
133 pipeline : `str`, optional
134 Name of the YAML file with pipeline.
135 qgraph : `str`, optional
136 Name of the pickle file with QGraph.
137 pipeline_actions : itrable of `cmdLinePArser._PipelineAction`, optional
138 order_pipeline : `bool`
139 save_pipeline : `str`
140 Name of the YAML file to store pipeline.
141 save_qgraph : `str`
142 Name of the pickle file to store QGraph.
143 save_single_quanta : `str`
144 Name of the pickle file pattern to store individual QGraph.
145 pipeline_dot : `str`
146 Name of the DOT file to write pipeline graph.
147 qgraph_dot : `str`
148 Name of the DOT file to write QGraph representation.
149 """
150 args = argparse.Namespace()
151 args.butler_config = Config()
152 if registryConfig:
153 args.butler_config["registry"] = registryConfig
154 # The default datastore has a relocatable root, so we need to specify
155 # some root here for it to use
156 args.butler_config.configFile = "."
157 args.pipeline = pipeline
158 args.qgraph = qgraph
159 args.pipeline_actions = pipeline_actions
160 args.order_pipeline = order_pipeline
161 args.save_pipeline = save_pipeline
162 args.save_qgraph = save_qgraph
163 args.save_single_quanta = save_single_quanta
164 args.pipeline_dot = pipeline_dot
165 args.qgraph_dot = qgraph_dot
166 args.input = ""
167 args.output = None
168 args.output_run = None
169 args.extend_run = False
170 args.replace_run = False
171 args.prune_replaced = False
172 args.register_dataset_types = False
173 args.skip_init_writes = False
174 args.no_versions = False
175 args.skip_existing = False
176 args.init_only = False
177 args.processes = 1
178 args.profile = None
179 args.enableLsstDebug = False
180 args.graph_fixup = None
181 return args
184def _makeQGraph():
185 """Make a trivial QuantumGraph with one quantum.
187 The only thing that we need to do with this quantum graph is to pickle
188 it, the quanta in this graph are not usable for anything else.
190 Returns
191 -------
192 qgraph : `~lsst.pipe.base.QuantumGraph`
193 """
194 taskDef = TaskDef(taskName="taskOne", config=SimpleConfig())
195 quanta = [Quantum()]
196 taskNodes = QuantumGraphTaskNodes(taskDef=taskDef, quanta=quanta, initInputs={}, initOutputs={})
197 qgraph = QuantumGraph([taskNodes])
198 return qgraph
201class CmdLineFwkTestCase(unittest.TestCase):
202 """A test case for CmdLineFwk
203 """
205 def testMakePipeline(self):
206 """Tests for CmdLineFwk.makePipeline method
207 """
208 fwk = CmdLineFwk()
210 # make empty pipeline
211 args = _makeArgs()
212 pipeline = fwk.makePipeline(args)
213 self.assertIsInstance(pipeline, Pipeline)
214 self.assertEqual(len(pipeline), 0)
216 # few tests with serialization
217 with makeTmpFile() as tmpname:
218 # make empty pipeline and store it in a file
219 args = _makeArgs(save_pipeline=tmpname)
220 pipeline = fwk.makePipeline(args)
221 self.assertIsInstance(pipeline, Pipeline)
223 # read pipeline from a file
224 args = _makeArgs(pipeline=tmpname)
225 pipeline = fwk.makePipeline(args)
226 self.assertIsInstance(pipeline, Pipeline)
227 self.assertEqual(len(pipeline), 0)
229 # single task pipeline
230 actions = [
231 _ACTION_ADD_TASK("TaskOne:task1")
232 ]
233 args = _makeArgs(pipeline_actions=actions)
234 pipeline = fwk.makePipeline(args)
235 self.assertIsInstance(pipeline, Pipeline)
236 self.assertEqual(len(pipeline), 1)
238 # many task pipeline
239 actions = [
240 _ACTION_ADD_TASK("TaskOne:task1a"),
241 _ACTION_ADD_TASK("TaskTwo:task2"),
242 _ACTION_ADD_TASK("TaskOne:task1b")
243 ]
244 args = _makeArgs(pipeline_actions=actions)
245 pipeline = fwk.makePipeline(args)
246 self.assertIsInstance(pipeline, Pipeline)
247 self.assertEqual(len(pipeline), 3)
249 # single task pipeline with config overrides, cannot use TaskOne, need
250 # something that can be imported with `doImport()`
251 actions = [
252 _ACTION_ADD_TASK("testUtil.AddTask:task"),
253 _ACTION_CONFIG("task:addend=100")
254 ]
255 args = _makeArgs(pipeline_actions=actions)
256 pipeline = fwk.makePipeline(args)
257 taskDefs = list(pipeline.toExpandedPipeline())
258 self.assertEqual(len(taskDefs), 1)
259 self.assertEqual(taskDefs[0].config.addend, 100)
261 overrides = b"config.addend = 1000\n"
262 with makeTmpFile(overrides) as tmpname:
263 actions = [
264 _ACTION_ADD_TASK("testUtil.AddTask:task"),
265 _ACTION_CONFIG_FILE("task:" + tmpname)
266 ]
267 args = _makeArgs(pipeline_actions=actions)
268 pipeline = fwk.makePipeline(args)
269 taskDefs = list(pipeline.toExpandedPipeline())
270 self.assertEqual(len(taskDefs), 1)
271 self.assertEqual(taskDefs[0].config.addend, 1000)
273 # Check --instrument option, for now it only checks that it does not crash
274 actions = [
275 _ACTION_ADD_TASK("testUtil.AddTask:task"),
276 _ACTION_ADD_INSTRUMENT("Instrument")
277 ]
278 args = _makeArgs(pipeline_actions=actions)
279 pipeline = fwk.makePipeline(args)
281 def testMakeGraphFromPickle(self):
282 """Tests for CmdLineFwk.makeGraph method.
284 Only most trivial case is tested that does not do actual graph
285 building.
286 """
287 fwk = CmdLineFwk()
289 with makeTmpFile() as tmpname, makeSQLiteRegistry() as registryConfig:
291 # make non-empty graph and store it in a file
292 qgraph = _makeQGraph()
293 with open(tmpname, "wb") as pickleFile:
294 qgraph.save(pickleFile)
295 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
296 qgraph = fwk.makeGraph(None, args)
297 self.assertIsInstance(qgraph, QuantumGraph)
298 self.assertEqual(len(qgraph), 1)
300 # pickle with wrong object type
301 with open(tmpname, "wb") as pickleFile:
302 pickle.dump({}, pickleFile)
303 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
304 with self.assertRaises(TypeError):
305 fwk.makeGraph(None, args)
307 # reading empty graph from pickle should work but makeGraph()
308 # will return None and make a warning
309 qgraph = QuantumGraph()
310 with open(tmpname, "wb") as pickleFile:
311 qgraph.save(pickleFile)
312 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
313 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"):
314 # this also tests that warning is generated for empty graph
315 qgraph = fwk.makeGraph(None, args)
316 self.assertIs(qgraph, None)
318 def testSimpleQGraph(self):
319 """Test successfull execution of trivial quantum graph.
320 """
322 nQuanta = 5
323 butler, qgraph = makeSimpleQGraph(nQuanta)
325 # should have one task and number of quanta
326 self.assertEqual(len(qgraph), 1)
327 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
329 args = _makeArgs()
330 fwk = CmdLineFwk()
331 taskFactory = AddTaskFactoryMock()
333 # run whole thing
334 AddTask.countExec = 0
335 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
336 self.assertEqual(AddTask.countExec, nQuanta)
338 def testSimpleQGraphSkipExisting(self):
339 """Test continuing execution of trivial quantum graph with --skip-existing.
340 """
342 nQuanta = 5
343 butler, qgraph = makeSimpleQGraph(nQuanta)
345 # should have one task and number of quanta
346 self.assertEqual(len(qgraph), 1)
347 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
349 args = _makeArgs()
350 fwk = CmdLineFwk()
351 taskFactory = AddTaskFactoryMock()
353 # run whole thing
354 AddTask.countExec = 0
355 AddTask.stopAt = 3
356 with self.assertRaises(RuntimeError):
357 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
358 self.assertEqual(AddTask.countExec, 3)
360 AddTask.stopAt = -1
361 args.skip_existing = True
362 args.no_versions = True
363 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
364 self.assertEqual(AddTask.countExec, nQuanta)
366 def testShowPipeline(self):
367 """Test for --show options for pipeline.
368 """
369 fwk = CmdLineFwk()
371 actions = [
372 _ACTION_ADD_TASK("testUtil.AddTask:task"),
373 _ACTION_CONFIG("task:addend=100")
374 ]
375 args = _makeArgs(pipeline_actions=actions)
376 pipeline = fwk.makePipeline(args)
378 args.show = ["pipeline"]
379 fwk.showInfo(args, pipeline)
380 args.show = ["config"]
381 fwk.showInfo(args, pipeline)
382 args.show = ["history=task::addend"]
383 fwk.showInfo(args, pipeline)
384 args.show = ["tasks"]
385 fwk.showInfo(args, pipeline)
387 def testShowGraph(self):
388 """Test for --show options for quantum graph.
389 """
390 fwk = CmdLineFwk()
392 nQuanta = 2
393 butler, qgraph = makeSimpleQGraph(nQuanta)
395 args = _makeArgs()
396 args.show = ["graph"]
397 fwk.showInfo(args, pipeline=None, graph=qgraph)
398 # TODO: cannot test "workflow" option presently, it instanciates
399 # butler from command line options and there is no way to pass butler
400 # mock to that code.
403class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
404 pass
407def setup_module(module):
408 lsst.utils.tests.init()
411if __name__ == "__main__": 411 ↛ 412line 411 didn't jump to line 412, because the condition on line 411 was never true
412 lsst.utils.tests.init()
413 unittest.main()