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

37from lsst.daf.butler.registry import RegistryConfig 

38from lsst.obs.base import Instrument 

39import lsst.pex.config as pexConfig 

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

41 QuantumGraph, QuantumGraphTaskNodes, 

42 TaskDef, TaskFactory, PipelineTaskConnections) 

43import lsst.pipe.base.connectionTypes as cT 

44import lsst.utils.tests 

45from testUtil import (AddTask, AddTaskFactoryMock, makeSimpleQGraph) 

46 

47 

48logging.basicConfig(level=logging.INFO) 

49 

50# Have to monkey-patch Instrument.fromName() to not retrieve non-existing 

51# instrument from registry, these tests can run fine without actual instrument 

52# and implementing full mock for Instrument is too complicated. 

53Instrument.fromName = lambda name, reg: None 53 ↛ exitline 53 didn't run the lambda on line 53

54 

55 

56@contextlib.contextmanager 

57def makeTmpFile(contents=None): 

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

59 

60 Temporary file is deleted on exiting context. 

61 

62 Parameters 

63 ---------- 

64 contents : `bytes` 

65 Data to write into a file. 

66 """ 

67 fd, tmpname = tempfile.mkstemp() 

68 if contents: 

69 os.write(fd, contents) 

70 os.close(fd) 

71 yield tmpname 

72 with contextlib.suppress(OSError): 

73 os.remove(tmpname) 

74 

75 

76@contextlib.contextmanager 

77def makeSQLiteRegistry(): 

78 """Context manager to create new empty registry database. 

79 

80 Yields 

81 ------ 

82 config : `RegistryConfig` 

83 Registry configuration for initialized registry database. 

84 """ 

85 with makeTmpFile() as filename: 

86 uri = f"sqlite:///{filename}" 

87 config = RegistryConfig() 

88 config["db"] = uri 

89 Registry.fromConfig(config, create=True) 

90 yield config 

91 

92 

93class SimpleConnections(PipelineTaskConnections, dimensions=(), 

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

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

96 name="{template}schema", 

97 storageClass="SourceCatalog") 

98 

99 

100class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections): 

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

102 

103 def setDefaults(self): 

104 PipelineTaskConfig.setDefaults(self) 

105 

106 

107class TaskOne(PipelineTask): 

108 ConfigClass = SimpleConfig 

109 _DefaultName = "taskOne" 

110 

111 

112class TaskTwo(PipelineTask): 

113 ConfigClass = SimpleConfig 

114 _DefaultName = "taskTwo" 

115 

116 

117class TaskFactoryMock(TaskFactory): 

118 def loadTaskClass(self, taskName): 

119 if taskName == "TaskOne": 

120 return TaskOne, "TaskOne" 

121 elif taskName == "TaskTwo": 

122 return TaskTwo, "TaskTwo" 

123 

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

125 if config is None: 

126 config = taskClass.ConfigClass() 

127 if overrides: 

128 overrides.applyTo(config) 

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

130 

131 

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

133 save_pipeline="", save_qgraph="", save_single_quanta="", 

134 pipeline_dot="", qgraph_dot="", registryConfig=None): 

135 """Return parsed command line arguments. 

136 

137 Parameters 

138 ---------- 

139 pipeline : `str`, optional 

140 Name of the YAML file with pipeline. 

141 qgraph : `str`, optional 

142 Name of the pickle file with QGraph. 

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

144 order_pipeline : `bool` 

145 save_pipeline : `str` 

146 Name of the YAML file to store pipeline. 

147 save_qgraph : `str` 

148 Name of the pickle file to store QGraph. 

149 save_single_quanta : `str` 

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

151 pipeline_dot : `str` 

152 Name of the DOT file to write pipeline graph. 

153 qgraph_dot : `str` 

154 Name of the DOT file to write QGraph representation. 

155 """ 

156 args = argparse.Namespace() 

157 args.butler_config = Config() 

158 if registryConfig: 

159 args.butler_config["registry"] = registryConfig 

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

161 # some root here for it to use 

162 args.butler_config.configFile = "." 

163 args.pipeline = pipeline 

164 args.qgraph = qgraph 

165 args.pipeline_actions = pipeline_actions 

166 args.order_pipeline = order_pipeline 

167 args.save_pipeline = save_pipeline 

168 args.save_qgraph = save_qgraph 

169 args.save_single_quanta = save_single_quanta 

170 args.pipeline_dot = pipeline_dot 

171 args.qgraph_dot = qgraph_dot 

172 args.input = "" 

173 args.output = None 

174 args.output_run = None 

175 args.extend_run = False 

176 args.replace_run = False 

177 args.prune_replaced = False 

178 args.register_dataset_types = False 

179 args.skip_init_writes = False 

180 args.no_versions = False 

181 args.skip_existing = False 

182 args.init_only = False 

183 args.processes = 1 

184 args.profile = None 

185 args.enableLsstDebug = False 

186 args.graph_fixup = None 

187 return args 

188 

189 

190def _makeQGraph(): 

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

192 

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

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

195 

196 Returns 

197 ------- 

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

199 """ 

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

201 quanta = [Quantum()] 

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

203 qgraph = QuantumGraph([taskNodes]) 

204 return qgraph 

205 

206 

207class CmdLineFwkTestCase(unittest.TestCase): 

208 """A test case for CmdLineFwk 

209 """ 

210 

211 def testMakePipeline(self): 

212 """Tests for CmdLineFwk.makePipeline method 

213 """ 

214 fwk = CmdLineFwk() 

215 

216 # make empty pipeline 

217 args = _makeArgs() 

218 pipeline = fwk.makePipeline(args) 

219 self.assertIsInstance(pipeline, Pipeline) 

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

221 

222 # few tests with serialization 

223 with makeTmpFile() as tmpname: 

224 # make empty pipeline and store it in a file 

225 args = _makeArgs(save_pipeline=tmpname) 

226 pipeline = fwk.makePipeline(args) 

227 self.assertIsInstance(pipeline, Pipeline) 

228 

229 # read pipeline from a file 

230 args = _makeArgs(pipeline=tmpname) 

231 pipeline = fwk.makePipeline(args) 

232 self.assertIsInstance(pipeline, Pipeline) 

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

234 

235 # single task pipeline 

236 actions = [ 

237 _ACTION_ADD_TASK("TaskOne:task1") 

238 ] 

239 args = _makeArgs(pipeline_actions=actions) 

240 pipeline = fwk.makePipeline(args) 

241 self.assertIsInstance(pipeline, Pipeline) 

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

243 

244 # many task pipeline 

245 actions = [ 

246 _ACTION_ADD_TASK("TaskOne:task1a"), 

247 _ACTION_ADD_TASK("TaskTwo:task2"), 

248 _ACTION_ADD_TASK("TaskOne:task1b") 

249 ] 

250 args = _makeArgs(pipeline_actions=actions) 

251 pipeline = fwk.makePipeline(args) 

252 self.assertIsInstance(pipeline, Pipeline) 

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

254 

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

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

257 actions = [ 

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

259 _ACTION_CONFIG("task:addend=100") 

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, 100) 

266 

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

268 with makeTmpFile(overrides) as tmpname: 

269 actions = [ 

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

271 _ACTION_CONFIG_FILE("task:" + tmpname) 

272 ] 

273 args = _makeArgs(pipeline_actions=actions) 

274 pipeline = fwk.makePipeline(args) 

275 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

278 

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

280 actions = [ 

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

282 _ACTION_ADD_INSTRUMENT("Instrument") 

283 ] 

284 args = _makeArgs(pipeline_actions=actions) 

285 pipeline = fwk.makePipeline(args) 

286 

287 def testMakeGraphFromPickle(self): 

288 """Tests for CmdLineFwk.makeGraph method. 

289 

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

291 building. 

292 """ 

293 fwk = CmdLineFwk() 

294 

295 with makeTmpFile() as tmpname, makeSQLiteRegistry() as registryConfig: 

296 

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

298 qgraph = _makeQGraph() 

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

300 qgraph.save(pickleFile) 

301 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig) 

302 qgraph = fwk.makeGraph(None, args) 

303 self.assertIsInstance(qgraph, QuantumGraph) 

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

305 

306 # pickle with wrong object type 

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

308 pickle.dump({}, pickleFile) 

309 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig) 

310 with self.assertRaises(TypeError): 

311 fwk.makeGraph(None, args) 

312 

313 # reading empty graph from pickle should work but makeGraph() 

314 # will return None and make a warning 

315 qgraph = QuantumGraph() 

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

317 qgraph.save(pickleFile) 

318 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig) 

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

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

321 qgraph = fwk.makeGraph(None, args) 

322 self.assertIs(qgraph, None) 

323 

324 def testSimpleQGraph(self): 

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

326 """ 

327 

328 nQuanta = 5 

329 butler, qgraph = makeSimpleQGraph(nQuanta) 

330 

331 # should have one task and number of quanta 

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

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

334 

335 args = _makeArgs() 

336 fwk = CmdLineFwk() 

337 taskFactory = AddTaskFactoryMock() 

338 

339 # run whole thing 

340 AddTask.countExec = 0 

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

342 self.assertEqual(AddTask.countExec, nQuanta) 

343 

344 def testSimpleQGraphSkipExisting(self): 

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

346 """ 

347 

348 nQuanta = 5 

349 butler, qgraph = makeSimpleQGraph(nQuanta) 

350 

351 # should have one task and number of quanta 

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

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

354 

355 args = _makeArgs() 

356 fwk = CmdLineFwk() 

357 taskFactory = AddTaskFactoryMock() 

358 

359 # run whole thing 

360 AddTask.countExec = 0 

361 AddTask.stopAt = 3 

362 with self.assertRaises(RuntimeError): 

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

364 self.assertEqual(AddTask.countExec, 3) 

365 

366 AddTask.stopAt = -1 

367 args.skip_existing = True 

368 args.no_versions = True 

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

370 self.assertEqual(AddTask.countExec, nQuanta) 

371 

372 def testShowPipeline(self): 

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

374 """ 

375 fwk = CmdLineFwk() 

376 

377 actions = [ 

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

379 _ACTION_CONFIG("task:addend=100") 

380 ] 

381 args = _makeArgs(pipeline_actions=actions) 

382 pipeline = fwk.makePipeline(args) 

383 

384 args.show = ["pipeline"] 

385 fwk.showInfo(args, pipeline) 

386 args.show = ["config"] 

387 fwk.showInfo(args, pipeline) 

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

389 fwk.showInfo(args, pipeline) 

390 args.show = ["tasks"] 

391 fwk.showInfo(args, pipeline) 

392 

393 def testShowGraph(self): 

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

395 """ 

396 fwk = CmdLineFwk() 

397 

398 nQuanta = 2 

399 butler, qgraph = makeSimpleQGraph(nQuanta) 

400 

401 args = _makeArgs() 

402 args.show = ["graph"] 

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

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

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

406 # mock to that code. 

407 

408 

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

410 pass 

411 

412 

413def setup_module(module): 

414 lsst.utils.tests.init() 

415 

416 

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

418 lsst.utils.tests.init() 

419 unittest.main()