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 shutil 

31import tempfile 

32import unittest 

33 

34from lsst.ctrl.mpexec.cmdLineFwk import CmdLineFwk 

35from lsst.ctrl.mpexec.cmdLineParser import (_ACTION_ADD_TASK, _ACTION_CONFIG, 

36 _ACTION_CONFIG_FILE, _ACTION_ADD_INSTRUMENT) 

37from lsst.daf.butler import Config, Quantum, Registry 

38from lsst.daf.butler.registry import RegistryConfig 

39from lsst.obs.base import Instrument 

40import lsst.pex.config as pexConfig 

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

42 QuantumGraph, QuantumGraphTaskNodes, 

43 TaskDef, TaskFactory, PipelineTaskConnections) 

44import lsst.pipe.base.connectionTypes as cT 

45import lsst.utils.tests 

46from testUtil import (AddTaskFactoryMock, makeSimpleQGraph) 

47 

48 

49logging.basicConfig(level=logging.INFO) 

50 

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

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

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

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

55 

56 

57@contextlib.contextmanager 

58def makeTmpFile(contents=None): 

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

60 

61 Temporary file is deleted on exiting context. 

62 

63 Parameters 

64 ---------- 

65 contents : `bytes` 

66 Data to write into a file. 

67 """ 

68 fd, tmpname = tempfile.mkstemp() 

69 if contents: 

70 os.write(fd, contents) 

71 os.close(fd) 

72 yield tmpname 

73 with contextlib.suppress(OSError): 

74 os.remove(tmpname) 

75 

76 

77@contextlib.contextmanager 

78def makeSQLiteRegistry(): 

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

80 

81 Yields 

82 ------ 

83 config : `RegistryConfig` 

84 Registry configuration for initialized registry database. 

85 """ 

86 with makeTmpFile() as filename: 

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

88 config = RegistryConfig() 

89 config["db"] = uri 

90 Registry.fromConfig(config, create=True) 

91 yield config 

92 

93 

94class SimpleConnections(PipelineTaskConnections, dimensions=(), 

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

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

97 name="{template}schema", 

98 storageClass="SourceCatalog") 

99 

100 

101class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections): 

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

103 

104 def setDefaults(self): 

105 PipelineTaskConfig.setDefaults(self) 

106 

107 

108class TaskOne(PipelineTask): 

109 ConfigClass = SimpleConfig 

110 _DefaultName = "taskOne" 

111 

112 

113class TaskTwo(PipelineTask): 

114 ConfigClass = SimpleConfig 

115 _DefaultName = "taskTwo" 

116 

117 

118class TaskFactoryMock(TaskFactory): 

119 def loadTaskClass(self, taskName): 

120 if taskName == "TaskOne": 

121 return TaskOne, "TaskOne" 

122 elif taskName == "TaskTwo": 

123 return TaskTwo, "TaskTwo" 

124 

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

126 if config is None: 

127 config = taskClass.ConfigClass() 

128 if overrides: 

129 overrides.applyTo(config) 

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

131 

132 

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

134 save_pipeline="", save_qgraph="", save_single_quanta="", 

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

136 """Return parsed command line arguments. 

137 

138 Parameters 

139 ---------- 

140 pipeline : `str`, optional 

141 Name of the YAML file with pipeline. 

142 qgraph : `str`, optional 

143 Name of the pickle file with QGraph. 

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

145 order_pipeline : `bool` 

146 save_pipeline : `str` 

147 Name of the YAML file to store pipeline. 

148 save_qgraph : `str` 

149 Name of the pickle file to store QGraph. 

150 save_single_quanta : `str` 

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

152 pipeline_dot : `str` 

153 Name of the DOT file to write pipeline graph. 

154 qgraph_dot : `str` 

155 Name of the DOT file to write QGraph representation. 

156 """ 

157 args = argparse.Namespace() 

158 args.butler_config = Config() 

159 if registryConfig: 

160 args.butler_config["registry"] = registryConfig 

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

162 # some root here for it to use 

163 args.butler_config.configFile = "." 

164 args.pipeline = pipeline 

165 args.qgraph = qgraph 

166 args.pipeline_actions = pipeline_actions 

167 args.order_pipeline = order_pipeline 

168 args.save_pipeline = save_pipeline 

169 args.save_qgraph = save_qgraph 

170 args.save_single_quanta = save_single_quanta 

171 args.pipeline_dot = pipeline_dot 

172 args.qgraph_dot = qgraph_dot 

173 args.input = "" 

174 args.output = None 

175 args.output_run = None 

176 args.extend_run = False 

177 args.replace_run = False 

178 args.prune_replaced = False 

179 args.register_dataset_types = False 

180 args.skip_init_writes = False 

181 args.no_versions = False 

182 args.skip_existing = False 

183 args.clobber_partial_outputs = False 

184 args.init_only = False 

185 args.processes = 1 

186 args.profile = None 

187 args.enableLsstDebug = False 

188 args.graph_fixup = None 

189 args.timeout = None 

190 args.fail_fast = False 

191 return args 

192 

193 

194def _makeQGraph(): 

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

196 

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

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

199 

200 Returns 

201 ------- 

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

203 """ 

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

205 quanta = [Quantum()] 

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

207 qgraph = QuantumGraph([taskNodes]) 

208 return qgraph 

209 

210 

211class CmdLineFwkTestCase(unittest.TestCase): 

212 """A test case for CmdLineFwk 

213 """ 

214 

215 def testMakePipeline(self): 

216 """Tests for CmdLineFwk.makePipeline method 

217 """ 

218 fwk = CmdLineFwk() 

219 

220 # make empty pipeline 

221 args = _makeArgs() 

222 pipeline = fwk.makePipeline(args) 

223 self.assertIsInstance(pipeline, Pipeline) 

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

225 

226 # few tests with serialization 

227 with makeTmpFile() as tmpname: 

228 # make empty pipeline and store it in a file 

229 args = _makeArgs(save_pipeline=tmpname) 

230 pipeline = fwk.makePipeline(args) 

231 self.assertIsInstance(pipeline, Pipeline) 

232 

233 # read pipeline from a file 

234 args = _makeArgs(pipeline=tmpname) 

235 pipeline = fwk.makePipeline(args) 

236 self.assertIsInstance(pipeline, Pipeline) 

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

238 

239 # single task pipeline 

240 actions = [ 

241 _ACTION_ADD_TASK("TaskOne:task1") 

242 ] 

243 args = _makeArgs(pipeline_actions=actions) 

244 pipeline = fwk.makePipeline(args) 

245 self.assertIsInstance(pipeline, Pipeline) 

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

247 

248 # many task pipeline 

249 actions = [ 

250 _ACTION_ADD_TASK("TaskOne:task1a"), 

251 _ACTION_ADD_TASK("TaskTwo:task2"), 

252 _ACTION_ADD_TASK("TaskOne:task1b") 

253 ] 

254 args = _makeArgs(pipeline_actions=actions) 

255 pipeline = fwk.makePipeline(args) 

256 self.assertIsInstance(pipeline, Pipeline) 

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

258 

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

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

261 actions = [ 

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

263 _ACTION_CONFIG("task:addend=100") 

264 ] 

265 args = _makeArgs(pipeline_actions=actions) 

266 pipeline = fwk.makePipeline(args) 

267 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

270 

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

272 with makeTmpFile(overrides) as tmpname: 

273 actions = [ 

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

275 _ACTION_CONFIG_FILE("task:" + tmpname) 

276 ] 

277 args = _makeArgs(pipeline_actions=actions) 

278 pipeline = fwk.makePipeline(args) 

279 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

282 

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

284 actions = [ 

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

286 _ACTION_ADD_INSTRUMENT("Instrument") 

287 ] 

288 args = _makeArgs(pipeline_actions=actions) 

289 pipeline = fwk.makePipeline(args) 

290 

291 def testMakeGraphFromPickle(self): 

292 """Tests for CmdLineFwk.makeGraph method. 

293 

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

295 building. 

296 """ 

297 fwk = CmdLineFwk() 

298 

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

300 

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

302 qgraph = _makeQGraph() 

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

304 qgraph.save(pickleFile) 

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

306 qgraph = fwk.makeGraph(None, args) 

307 self.assertIsInstance(qgraph, QuantumGraph) 

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

309 

310 # pickle with wrong object type 

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

312 pickle.dump({}, pickleFile) 

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

314 with self.assertRaises(TypeError): 

315 fwk.makeGraph(None, args) 

316 

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

318 # will return None and make a warning 

319 qgraph = QuantumGraph() 

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

321 qgraph.save(pickleFile) 

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

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

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

325 qgraph = fwk.makeGraph(None, args) 

326 self.assertIs(qgraph, None) 

327 

328 def testShowPipeline(self): 

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

330 """ 

331 fwk = CmdLineFwk() 

332 

333 actions = [ 

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

335 _ACTION_CONFIG("task:addend=100") 

336 ] 

337 args = _makeArgs(pipeline_actions=actions) 

338 pipeline = fwk.makePipeline(args) 

339 

340 args.show = ["pipeline"] 

341 fwk.showInfo(args, pipeline) 

342 args.show = ["config"] 

343 fwk.showInfo(args, pipeline) 

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

345 fwk.showInfo(args, pipeline) 

346 args.show = ["tasks"] 

347 fwk.showInfo(args, pipeline) 

348 

349 

350class CmdLineFwkTestCaseWithButler(unittest.TestCase): 

351 """A test case for CmdLineFwk 

352 """ 

353 

354 def setUp(self): 

355 super().setUpClass() 

356 self.root = tempfile.mkdtemp() 

357 

358 def tearDown(self): 

359 shutil.rmtree(self.root, ignore_errors=True) 

360 super().tearDownClass() 

361 

362 def testSimpleQGraph(self): 

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

364 """ 

365 

366 nQuanta = 5 

367 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root) 

368 

369 # should have one task and number of quanta 

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

371 

372 args = _makeArgs() 

373 fwk = CmdLineFwk() 

374 taskFactory = AddTaskFactoryMock() 

375 

376 # run whole thing 

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

378 self.assertEqual(taskFactory.countExec, nQuanta) 

379 

380 def testSimpleQGraphSkipExisting(self): 

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

382 """ 

383 

384 nQuanta = 5 

385 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root) 

386 

387 # should have one task and number of quanta 

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

389 

390 args = _makeArgs() 

391 fwk = CmdLineFwk() 

392 taskFactory = AddTaskFactoryMock(stopAt=3) 

393 

394 # run first three quanta 

395 with self.assertRaises(RuntimeError): 

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

397 self.assertEqual(taskFactory.countExec, 3) 

398 

399 # run remaining ones 

400 taskFactory.stopAt = -1 

401 args.skip_existing = True 

402 args.no_versions = True 

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

404 self.assertEqual(taskFactory.countExec, nQuanta) 

405 

406 def testSimpleQGraphPartialOutputsFail(self): 

407 """Test continuing execution of trivial quantum graph with partial 

408 outputs. 

409 """ 

410 

411 nQuanta = 5 

412 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root) 

413 

414 # should have one task and number of quanta 

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

416 

417 args = _makeArgs() 

418 fwk = CmdLineFwk() 

419 taskFactory = AddTaskFactoryMock(stopAt=3) 

420 

421 # run first three quanta 

422 with self.assertRaises(RuntimeError): 

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

424 self.assertEqual(taskFactory.countExec, 3) 

425 

426 # drop one of the two outputs from one task 

427 ref = butler._findDatasetRef("add2_dataset2", instrument="INSTR", detector=0) 

428 self.assertIsNotNone(ref) 

429 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True) 

430 

431 taskFactory.stopAt = -1 

432 args.skip_existing = True 

433 args.no_versions = True 

434 excRe = "Registry inconsistency while checking for existing outputs.*" 

435 with self.assertRaisesRegex(RuntimeError, excRe): 

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

437 

438 def testSimpleQGraphClobberPartialOutputs(self): 

439 """Test continuing execution of trivial quantum graph with 

440 --clobber-partial-outputs. 

441 """ 

442 

443 nQuanta = 5 

444 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root) 

445 

446 # should have one task and number of quanta 

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

448 

449 args = _makeArgs() 

450 fwk = CmdLineFwk() 

451 taskFactory = AddTaskFactoryMock(stopAt=3) 

452 

453 # run first three quanta 

454 with self.assertRaises(RuntimeError): 

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

456 self.assertEqual(taskFactory.countExec, 3) 

457 

458 # drop one of the two outputs from one task 

459 ref = butler._findDatasetRef("add2_dataset2", instrument="INSTR", detector=0) 

460 self.assertIsNotNone(ref) 

461 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True) 

462 

463 taskFactory.stopAt = -1 

464 args.skip_existing = True 

465 args.clobber_partial_outputs = True 

466 args.no_versions = True 

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

468 # number of executed quanta is incremented 

469 self.assertEqual(taskFactory.countExec, nQuanta + 1) 

470 

471 def testShowGraph(self): 

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

473 """ 

474 fwk = CmdLineFwk() 

475 

476 nQuanta = 2 

477 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root) 

478 

479 args = _makeArgs() 

480 args.show = ["graph"] 

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

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

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

484 # mock to that code. 

485 

486 

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

488 pass 

489 

490 

491def setup_module(module): 

492 lsst.utils.tests.init() 

493 

494 

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

496 lsst.utils.tests.init() 

497 unittest.main()