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 args.graph_fixup = None 

154 return args 

155 

156 

157def _makeQGraph(): 

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

159 

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

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

162 

163 Returns 

164 ------- 

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

166 """ 

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

168 quanta = [Quantum()] 

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

170 qgraph = QuantumGraph([taskNodes]) 

171 return qgraph 

172 

173 

174class CmdLineFwkTestCase(unittest.TestCase): 

175 """A test case for CmdLineFwk 

176 """ 

177 

178 def testMakePipeline(self): 

179 """Tests for CmdLineFwk.makePipeline method 

180 """ 

181 fwk = CmdLineFwk() 

182 

183 # make empty pipeline 

184 args = _makeArgs() 

185 pipeline = fwk.makePipeline(args) 

186 self.assertIsInstance(pipeline, Pipeline) 

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

188 

189 # few tests with serialization 

190 with makeTmpFile() as tmpname: 

191 # make empty pipeline and store it in a file 

192 args = _makeArgs(save_pipeline=tmpname) 

193 pipeline = fwk.makePipeline(args) 

194 self.assertIsInstance(pipeline, Pipeline) 

195 

196 # read pipeline from a file 

197 args = _makeArgs(pipeline=tmpname) 

198 pipeline = fwk.makePipeline(args) 

199 self.assertIsInstance(pipeline, Pipeline) 

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

201 

202 # single task pipeline 

203 actions = [ 

204 _ACTION_ADD_TASK("TaskOne:task1") 

205 ] 

206 args = _makeArgs(pipeline_actions=actions) 

207 pipeline = fwk.makePipeline(args) 

208 self.assertIsInstance(pipeline, Pipeline) 

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

210 

211 # many task pipeline 

212 actions = [ 

213 _ACTION_ADD_TASK("TaskOne:task1a"), 

214 _ACTION_ADD_TASK("TaskTwo:task2"), 

215 _ACTION_ADD_TASK("TaskOne:task1b") 

216 ] 

217 args = _makeArgs(pipeline_actions=actions) 

218 pipeline = fwk.makePipeline(args) 

219 self.assertIsInstance(pipeline, Pipeline) 

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

221 

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

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

224 actions = [ 

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

226 _ACTION_CONFIG("task:addend=100") 

227 ] 

228 args = _makeArgs(pipeline_actions=actions) 

229 pipeline = fwk.makePipeline(args) 

230 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

233 

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

235 with makeTmpFile(overrides) as tmpname: 

236 actions = [ 

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

238 _ACTION_CONFIG_FILE("task:" + tmpname) 

239 ] 

240 args = _makeArgs(pipeline_actions=actions) 

241 pipeline = fwk.makePipeline(args) 

242 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

245 

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

247 actions = [ 

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

249 _ACTION_ADD_INSTRUMENT("Instrument") 

250 ] 

251 args = _makeArgs(pipeline_actions=actions) 

252 pipeline = fwk.makePipeline(args) 

253 

254 def testMakeGraphFromPickle(self): 

255 """Tests for CmdLineFwk.makeGraph method. 

256 

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

258 building. 

259 """ 

260 fwk = CmdLineFwk() 

261 

262 with makeTmpFile() as tmpname: 

263 

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

265 qgraph = _makeQGraph() 

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

267 pickle.dump(qgraph, pickleFile) 

268 args = _makeArgs(qgraph=tmpname) 

269 qgraph = fwk.makeGraph(None, args) 

270 self.assertIsInstance(qgraph, QuantumGraph) 

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

272 

273 # pickle with wrong object type 

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

275 pickle.dump({}, pickleFile) 

276 args = _makeArgs(qgraph=tmpname) 

277 with self.assertRaises(TypeError): 

278 fwk.makeGraph(None, args) 

279 

280 # reading empty graph from pickle should return None 

281 qgraph = QuantumGraph() 

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

283 pickle.dump(qgraph, pickleFile) 

284 args = _makeArgs(qgraph=tmpname) 

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

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

287 qgraph = fwk.makeGraph(None, args) 

288 self.assertIs(qgraph, None) 

289 

290 def testSimpleQGraph(self): 

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

292 """ 

293 

294 nQuanta = 5 

295 butler, qgraph = makeSimpleQGraph(nQuanta) 

296 

297 # should have one task and number of quanta 

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

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

300 

301 args = _makeArgs() 

302 fwk = CmdLineFwk() 

303 taskFactory = AddTaskFactoryMock() 

304 

305 # run whole thing 

306 AddTask.countExec = 0 

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

308 self.assertEqual(AddTask.countExec, nQuanta) 

309 

310 def testSimpleQGraphSkipExisting(self): 

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

312 """ 

313 

314 nQuanta = 5 

315 butler, qgraph = makeSimpleQGraph(nQuanta) 

316 

317 # should have one task and number of quanta 

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

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

320 

321 args = _makeArgs() 

322 fwk = CmdLineFwk() 

323 taskFactory = AddTaskFactoryMock() 

324 

325 # run whole thing 

326 AddTask.countExec = 0 

327 AddTask.stopAt = 3 

328 with self.assertRaises(RuntimeError): 

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

330 self.assertEqual(AddTask.countExec, 3) 

331 

332 AddTask.stopAt = -1 

333 args.skip_existing = True 

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

335 self.assertEqual(AddTask.countExec, nQuanta) 

336 

337 def testShowPipeline(self): 

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

339 """ 

340 fwk = CmdLineFwk() 

341 

342 actions = [ 

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

344 _ACTION_CONFIG("task:addend=100") 

345 ] 

346 args = _makeArgs(pipeline_actions=actions) 

347 pipeline = fwk.makePipeline(args) 

348 

349 args.show = ["pipeline"] 

350 fwk.showInfo(args, pipeline) 

351 args.show = ["config"] 

352 fwk.showInfo(args, pipeline) 

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

354 fwk.showInfo(args, pipeline) 

355 args.show = ["tasks"] 

356 fwk.showInfo(args, pipeline) 

357 

358 def testShowGraph(self): 

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

360 """ 

361 fwk = CmdLineFwk() 

362 

363 nQuanta = 2 

364 butler, qgraph = makeSimpleQGraph(nQuanta) 

365 

366 args = _makeArgs() 

367 args.show = ["graph"] 

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

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

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

371 # mock to that code. 

372 

373 

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

375 pass 

376 

377 

378def setup_module(module): 

379 lsst.utils.tests.init() 

380 

381 

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

383 lsst.utils.tests.init() 

384 unittest.main()