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

188 args.fail_fast = False 

189 return args 

190 

191 

192def _makeQGraph(): 

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

194 

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

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

197 

198 Returns 

199 ------- 

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

201 """ 

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

203 quanta = [Quantum()] 

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

205 qgraph = QuantumGraph([taskNodes]) 

206 return qgraph 

207 

208 

209class CmdLineFwkTestCase(unittest.TestCase): 

210 """A test case for CmdLineFwk 

211 """ 

212 

213 def testMakePipeline(self): 

214 """Tests for CmdLineFwk.makePipeline method 

215 """ 

216 fwk = CmdLineFwk() 

217 

218 # make empty pipeline 

219 args = _makeArgs() 

220 pipeline = fwk.makePipeline(args) 

221 self.assertIsInstance(pipeline, Pipeline) 

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

223 

224 # few tests with serialization 

225 with makeTmpFile() as tmpname: 

226 # make empty pipeline and store it in a file 

227 args = _makeArgs(save_pipeline=tmpname) 

228 pipeline = fwk.makePipeline(args) 

229 self.assertIsInstance(pipeline, Pipeline) 

230 

231 # read pipeline from a file 

232 args = _makeArgs(pipeline=tmpname) 

233 pipeline = fwk.makePipeline(args) 

234 self.assertIsInstance(pipeline, Pipeline) 

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

236 

237 # single task pipeline 

238 actions = [ 

239 _ACTION_ADD_TASK("TaskOne:task1") 

240 ] 

241 args = _makeArgs(pipeline_actions=actions) 

242 pipeline = fwk.makePipeline(args) 

243 self.assertIsInstance(pipeline, Pipeline) 

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

245 

246 # many task pipeline 

247 actions = [ 

248 _ACTION_ADD_TASK("TaskOne:task1a"), 

249 _ACTION_ADD_TASK("TaskTwo:task2"), 

250 _ACTION_ADD_TASK("TaskOne:task1b") 

251 ] 

252 args = _makeArgs(pipeline_actions=actions) 

253 pipeline = fwk.makePipeline(args) 

254 self.assertIsInstance(pipeline, Pipeline) 

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

256 

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

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

259 actions = [ 

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

261 _ACTION_CONFIG("task:addend=100") 

262 ] 

263 args = _makeArgs(pipeline_actions=actions) 

264 pipeline = fwk.makePipeline(args) 

265 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

268 

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

270 with makeTmpFile(overrides) as tmpname: 

271 actions = [ 

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

273 _ACTION_CONFIG_FILE("task:" + tmpname) 

274 ] 

275 args = _makeArgs(pipeline_actions=actions) 

276 pipeline = fwk.makePipeline(args) 

277 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

280 

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

282 actions = [ 

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

284 _ACTION_ADD_INSTRUMENT("Instrument") 

285 ] 

286 args = _makeArgs(pipeline_actions=actions) 

287 pipeline = fwk.makePipeline(args) 

288 

289 def testMakeGraphFromPickle(self): 

290 """Tests for CmdLineFwk.makeGraph method. 

291 

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

293 building. 

294 """ 

295 fwk = CmdLineFwk() 

296 

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

298 

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

300 qgraph = _makeQGraph() 

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

302 qgraph.save(pickleFile) 

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

304 qgraph = fwk.makeGraph(None, args) 

305 self.assertIsInstance(qgraph, QuantumGraph) 

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

307 

308 # pickle with wrong object type 

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

310 pickle.dump({}, pickleFile) 

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

312 with self.assertRaises(TypeError): 

313 fwk.makeGraph(None, args) 

314 

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

316 # will return None and make a warning 

317 qgraph = QuantumGraph() 

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

319 qgraph.save(pickleFile) 

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

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

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

323 qgraph = fwk.makeGraph(None, args) 

324 self.assertIs(qgraph, None) 

325 

326 def testSimpleQGraph(self): 

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

328 """ 

329 

330 nQuanta = 5 

331 butler, qgraph = makeSimpleQGraph(nQuanta) 

332 

333 # should have one task and number of quanta 

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

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

336 

337 args = _makeArgs() 

338 fwk = CmdLineFwk() 

339 taskFactory = AddTaskFactoryMock() 

340 

341 # run whole thing 

342 AddTask.countExec = 0 

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

344 self.assertEqual(AddTask.countExec, nQuanta) 

345 

346 def testSimpleQGraphSkipExisting(self): 

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

348 """ 

349 

350 nQuanta = 5 

351 butler, qgraph = makeSimpleQGraph(nQuanta) 

352 

353 # should have one task and number of quanta 

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

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

356 

357 args = _makeArgs() 

358 fwk = CmdLineFwk() 

359 taskFactory = AddTaskFactoryMock() 

360 

361 # run whole thing 

362 AddTask.countExec = 0 

363 AddTask.stopAt = 3 

364 with self.assertRaises(RuntimeError): 

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

366 self.assertEqual(AddTask.countExec, 3) 

367 

368 AddTask.stopAt = -1 

369 args.skip_existing = True 

370 args.no_versions = True 

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

372 self.assertEqual(AddTask.countExec, nQuanta) 

373 

374 def testShowPipeline(self): 

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

376 """ 

377 fwk = CmdLineFwk() 

378 

379 actions = [ 

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

381 _ACTION_CONFIG("task:addend=100") 

382 ] 

383 args = _makeArgs(pipeline_actions=actions) 

384 pipeline = fwk.makePipeline(args) 

385 

386 args.show = ["pipeline"] 

387 fwk.showInfo(args, pipeline) 

388 args.show = ["config"] 

389 fwk.showInfo(args, pipeline) 

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

391 fwk.showInfo(args, pipeline) 

392 args.show = ["tasks"] 

393 fwk.showInfo(args, pipeline) 

394 

395 def testShowGraph(self): 

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

397 """ 

398 fwk = CmdLineFwk() 

399 

400 nQuanta = 2 

401 butler, qgraph = makeSimpleQGraph(nQuanta) 

402 

403 args = _makeArgs() 

404 args.show = ["graph"] 

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

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

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

408 # mock to that code. 

409 

410 

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

412 pass 

413 

414 

415def setup_module(module): 

416 lsst.utils.tests.init() 

417 

418 

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

420 lsst.utils.tests.init() 

421 unittest.main()