Hide keyboard shortcuts

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/>. 

21 

22"""Simple unit test for cmdLineFwk module. 

23""" 

24 

25import argparse 

26import contextlib 

27import logging 

28import os 

29import pickle 

30import tempfile 

31import unittest 

32 

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 Quantum, Config 

37import lsst.pex.config as pexConfig 

38from lsst.pipe.base import (Pipeline, PipelineTask, PipelineTaskConfig, 

39 QuantumGraph, QuantumGraphTaskNodes, 

40 TaskDef, TaskFactory, PipelineTaskConnections) 

41import lsst.pipe.base.connectionTypes as cT 

42import lsst.utils.tests 

43from testUtil import (AddTask, AddTaskFactoryMock, makeSimpleQGraph) 

44 

45 

46logging.basicConfig(level=logging.INFO) 

47 

48 

49@contextlib.contextmanager 

50def makeTmpFile(contents=None): 

51 """Context manager for generating temporary file name. 

52 

53 Temporary file is deleted on exiting context. 

54 

55 Parameters 

56 ---------- 

57 contents : `bytes` 

58 Data to write into a file. 

59 """ 

60 fd, tmpname = tempfile.mkstemp() 

61 if contents: 

62 os.write(fd, contents) 

63 os.close(fd) 

64 yield tmpname 

65 with contextlib.suppress(OSError): 

66 os.remove(tmpname) 

67 

68 

69class SimpleConnections(PipelineTaskConnections, dimensions=(), 

70 defaultTemplates={"template": "simple"}): 

71 schema = cT.InitInput(doc="Schema", 

72 name="{template}schema", 

73 storageClass="SourceCatalog") 

74 

75 

76class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections): 

77 field = pexConfig.Field(dtype=str, doc="arbitrary string") 

78 

79 def setDefaults(self): 

80 PipelineTaskConfig.setDefaults(self) 

81 

82 

83class TaskOne(PipelineTask): 

84 ConfigClass = SimpleConfig 

85 _DefaultName = "taskOne" 

86 

87 

88class TaskTwo(PipelineTask): 

89 ConfigClass = SimpleConfig 

90 _DefaultName = "taskTwo" 

91 

92 

93class TaskFactoryMock(TaskFactory): 

94 def loadTaskClass(self, taskName): 

95 if taskName == "TaskOne": 

96 return TaskOne, "TaskOne" 

97 elif taskName == "TaskTwo": 

98 return TaskTwo, "TaskTwo" 

99 

100 def makeTask(self, taskClass, config, overrides, butler): 

101 if config is None: 

102 config = taskClass.ConfigClass() 

103 if overrides: 

104 overrides.applyTo(config) 

105 return taskClass(config=config, butler=butler) 

106 

107 

108def _makeArgs(pipeline=None, qgraph=None, pipeline_actions=(), order_pipeline=False, 

109 save_pipeline="", save_qgraph="", save_single_quanta="", pipeline_dot="", qgraph_dot=""): 

110 """Return parsed command line arguments. 

111 

112 Parameters 

113 ---------- 

114 pipeline : `str`, optional 

115 Name of the YAML file with pipeline. 

116 qgraph : `str`, optional 

117 Name of the pickle file with QGraph. 

118 pipeline_actions : itrable of `cmdLinePArser._PipelineAction`, optional 

119 order_pipeline : `bool` 

120 save_pipeline : `str` 

121 Name of the YAML file to store pipeline. 

122 save_qgraph : `str` 

123 Name of the pickle file to store QGraph. 

124 save_single_quanta : `str` 

125 Name of the pickle file pattern to store individual QGraph. 

126 pipeline_dot : `str` 

127 Name of the DOT file to write pipeline graph. 

128 qgraph_dot : `str` 

129 Name of the DOT file to write QGraph representation. 

130 """ 

131 args = argparse.Namespace() 

132 args.butler_config = Config() 

133 # The default datastore has a relocatable root, so we need to specify 

134 # some root here for it to use 

135 args.butler_config.configFile = "." 

136 args.pipeline = pipeline 

137 args.qgraph = qgraph 

138 args.pipeline_actions = pipeline_actions 

139 args.order_pipeline = order_pipeline 

140 args.save_pipeline = save_pipeline 

141 args.save_qgraph = save_qgraph 

142 args.save_single_quanta = save_single_quanta 

143 args.pipeline_dot = pipeline_dot 

144 args.qgraph_dot = qgraph_dot 

145 args.output = {} 

146 args.register_dataset_types = False 

147 args.skip_init_writes = False 

148 args.skip_existing = False 

149 args.init_only = False 

150 args.processes = 1 

151 args.profile = None 

152 args.enableLsstDebug = False 

153 return args 

154 

155 

156def _makeQGraph(): 

157 """Make a trivial QuantumGraph with one quantum. 

158 

159 The only thing that we need to do with this quantum graph is to pickle 

160 it, the quanta in this graph are not usable for anything else. 

161 

162 Returns 

163 ------- 

164 qgraph : `~lsst.pipe.base.QuantumGraph` 

165 """ 

166 taskDef = TaskDef(taskName="taskOne", config=SimpleConfig()) 

167 quanta = [Quantum()] 

168 taskNodes = QuantumGraphTaskNodes(taskDef=taskDef, quanta=quanta, initInputs={}, initOutputs={}) 

169 qgraph = QuantumGraph([taskNodes]) 

170 return qgraph 

171 

172 

173class CmdLineFwkTestCase(unittest.TestCase): 

174 """A test case for CmdLineFwk 

175 """ 

176 

177 def testMakePipeline(self): 

178 """Tests for CmdLineFwk.makePipeline method 

179 """ 

180 fwk = CmdLineFwk() 

181 

182 # make empty pipeline 

183 args = _makeArgs() 

184 pipeline = fwk.makePipeline(args) 

185 self.assertIsInstance(pipeline, Pipeline) 

186 self.assertEqual(len(pipeline), 0) 

187 

188 # few tests with serialization 

189 with makeTmpFile() as tmpname: 

190 # make empty pipeline and store it in a file 

191 args = _makeArgs(save_pipeline=tmpname) 

192 pipeline = fwk.makePipeline(args) 

193 self.assertIsInstance(pipeline, Pipeline) 

194 

195 # read pipeline from a file 

196 args = _makeArgs(pipeline=tmpname) 

197 pipeline = fwk.makePipeline(args) 

198 self.assertIsInstance(pipeline, Pipeline) 

199 self.assertEqual(len(pipeline), 0) 

200 

201 # single task pipeline 

202 actions = [ 

203 _ACTION_ADD_TASK("TaskOne:task1") 

204 ] 

205 args = _makeArgs(pipeline_actions=actions) 

206 pipeline = fwk.makePipeline(args) 

207 self.assertIsInstance(pipeline, Pipeline) 

208 self.assertEqual(len(pipeline), 1) 

209 

210 # many task pipeline 

211 actions = [ 

212 _ACTION_ADD_TASK("TaskOne:task1a"), 

213 _ACTION_ADD_TASK("TaskTwo:task2"), 

214 _ACTION_ADD_TASK("TaskOne:task1b") 

215 ] 

216 args = _makeArgs(pipeline_actions=actions) 

217 pipeline = fwk.makePipeline(args) 

218 self.assertIsInstance(pipeline, Pipeline) 

219 self.assertEqual(len(pipeline), 3) 

220 

221 # single task pipeline with config overrides, cannot use TaskOne, need 

222 # something that can be imported with `doImport()` 

223 actions = [ 

224 _ACTION_ADD_TASK("testUtil.AddTask:task"), 

225 _ACTION_CONFIG("task:addend=100") 

226 ] 

227 args = _makeArgs(pipeline_actions=actions) 

228 pipeline = fwk.makePipeline(args) 

229 taskDefs = list(pipeline.toExpandedPipeline()) 

230 self.assertEqual(len(taskDefs), 1) 

231 self.assertEqual(taskDefs[0].config.addend, 100) 

232 

233 overrides = b"config.addend = 1000\n" 

234 with makeTmpFile(overrides) as tmpname: 

235 actions = [ 

236 _ACTION_ADD_TASK("testUtil.AddTask:task"), 

237 _ACTION_CONFIG_FILE("task:" + tmpname) 

238 ] 

239 args = _makeArgs(pipeline_actions=actions) 

240 pipeline = fwk.makePipeline(args) 

241 taskDefs = list(pipeline.toExpandedPipeline()) 

242 self.assertEqual(len(taskDefs), 1) 

243 self.assertEqual(taskDefs[0].config.addend, 1000) 

244 

245 # Check --instrument option, for now it only checks that it does not crash 

246 actions = [ 

247 _ACTION_ADD_TASK("testUtil.AddTask:task"), 

248 _ACTION_ADD_INSTRUMENT("Instrument") 

249 ] 

250 args = _makeArgs(pipeline_actions=actions) 

251 pipeline = fwk.makePipeline(args) 

252 

253 def testMakeGraphFromPickle(self): 

254 """Tests for CmdLineFwk.makeGraph method. 

255 

256 Only most trivial case is tested that does not do actual graph 

257 building. 

258 """ 

259 fwk = CmdLineFwk() 

260 

261 with makeTmpFile() as tmpname: 

262 

263 # make non-empty graph and store it in a file 

264 qgraph = _makeQGraph() 

265 with open(tmpname, "wb") as pickleFile: 

266 pickle.dump(qgraph, pickleFile) 

267 args = _makeArgs(qgraph=tmpname) 

268 qgraph = fwk.makeGraph(None, args) 

269 self.assertIsInstance(qgraph, QuantumGraph) 

270 self.assertEqual(len(qgraph), 1) 

271 

272 # pickle with wrong object type 

273 with open(tmpname, "wb") as pickleFile: 

274 pickle.dump({}, pickleFile) 

275 args = _makeArgs(qgraph=tmpname) 

276 with self.assertRaises(TypeError): 

277 fwk.makeGraph(None, args) 

278 

279 # reading empty graph from pickle should return None 

280 qgraph = QuantumGraph() 

281 with open(tmpname, "wb") as pickleFile: 

282 pickle.dump(qgraph, pickleFile) 

283 args = _makeArgs(qgraph=tmpname) 

284 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"): 

285 # this also tests that warning is generated for empty graph 

286 qgraph = fwk.makeGraph(None, args) 

287 self.assertIs(qgraph, None) 

288 

289 def testSimpleQGraph(self): 

290 """Test successfull execution of trivial quantum graph. 

291 """ 

292 

293 nQuanta = 5 

294 butler, qgraph = makeSimpleQGraph(nQuanta) 

295 

296 # should have one task and number of quanta 

297 self.assertEqual(len(qgraph), 1) 

298 self.assertEqual(len(list(qgraph.quanta())), nQuanta) 

299 

300 args = _makeArgs() 

301 fwk = CmdLineFwk() 

302 taskFactory = AddTaskFactoryMock() 

303 

304 # run whole thing 

305 AddTask.countExec = 0 

306 fwk.runPipeline(qgraph, taskFactory, args, butler=butler) 

307 self.assertEqual(AddTask.countExec, nQuanta) 

308 

309 def testSimpleQGraphSkipExisting(self): 

310 """Test continuing execution of trivial quantum graph with --skip-existing. 

311 """ 

312 

313 nQuanta = 5 

314 butler, qgraph = makeSimpleQGraph(nQuanta) 

315 

316 # should have one task and number of quanta 

317 self.assertEqual(len(qgraph), 1) 

318 self.assertEqual(len(list(qgraph.quanta())), nQuanta) 

319 

320 args = _makeArgs() 

321 fwk = CmdLineFwk() 

322 taskFactory = AddTaskFactoryMock() 

323 

324 # run whole thing 

325 AddTask.countExec = 0 

326 AddTask.stopAt = 3 

327 with self.assertRaises(RuntimeError): 

328 fwk.runPipeline(qgraph, taskFactory, args, butler=butler) 

329 self.assertEqual(AddTask.countExec, 3) 

330 

331 AddTask.stopAt = -1 

332 args.skip_existing = True 

333 fwk.runPipeline(qgraph, taskFactory, args, butler=butler) 

334 self.assertEqual(AddTask.countExec, nQuanta) 

335 

336 def testShowPipeline(self): 

337 """Test for --show options for pipeline. 

338 """ 

339 fwk = CmdLineFwk() 

340 

341 actions = [ 

342 _ACTION_ADD_TASK("testUtil.AddTask:task"), 

343 _ACTION_CONFIG("task:addend=100") 

344 ] 

345 args = _makeArgs(pipeline_actions=actions) 

346 pipeline = fwk.makePipeline(args) 

347 

348 args.show = ["pipeline"] 

349 fwk.showInfo(args, pipeline) 

350 args.show = ["config"] 

351 fwk.showInfo(args, pipeline) 

352 args.show = ["history=task::addend"] 

353 fwk.showInfo(args, pipeline) 

354 args.show = ["tasks"] 

355 fwk.showInfo(args, pipeline) 

356 

357 def testShowGraph(self): 

358 """Test for --show options for quantum graph. 

359 """ 

360 fwk = CmdLineFwk() 

361 

362 nQuanta = 2 

363 butler, qgraph = makeSimpleQGraph(nQuanta) 

364 

365 args = _makeArgs() 

366 args.show = ["graph"] 

367 fwk.showInfo(args, pipeline=None, graph=qgraph) 

368 # TODO: cannot test "workflow" option presently, it instanciates 

369 # butler from command line options and there is no way to pass butler 

370 # mock to that code. 

371 

372 

373class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase): 

374 pass 

375 

376 

377def setup_module(module): 

378 lsst.utils.tests.init() 

379 

380 

381if __name__ == "__main__": 381 ↛ 382line 381 didn't jump to line 382, because the condition on line 381 was never true

382 lsst.utils.tests.init() 

383 unittest.main()