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 contextlib 

26import copy 

27import logging 

28import os 

29import pickle 

30import shutil 

31import tempfile 

32import unittest 

33 

34from lsst.ctrl.mpexec.cmdLineFwk import CmdLineFwk 

35import lsst.ctrl.mpexec.cmdLineParser as parser_mod 

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

37 _ACTION_CONFIG_FILE, _ACTION_ADD_INSTRUMENT) 

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

39from lsst.daf.butler.registry import RegistryConfig 

40from lsst.obs.base import Instrument 

41import lsst.pex.config as pexConfig 

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

43 QuantumGraph, QuantumGraphTaskNodes, 

44 TaskDef, TaskFactory, PipelineTaskConnections) 

45import lsst.pipe.base.connectionTypes as cT 

46import lsst.utils.tests 

47from lsst.pipe.base.tests.simpleQGraph import (AddTaskFactoryMock, makeSimpleQGraph) 

48from lsst.utils.tests import temporaryDirectory 

49 

50 

51logging.basicConfig(level=logging.INFO) 

52 

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

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

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

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

57 

58 

59@contextlib.contextmanager 

60def makeTmpFile(contents=None): 

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

62 

63 Temporary file is deleted on exiting context. 

64 

65 Parameters 

66 ---------- 

67 contents : `bytes` 

68 Data to write into a file. 

69 """ 

70 fd, tmpname = tempfile.mkstemp() 

71 if contents: 

72 os.write(fd, contents) 

73 os.close(fd) 

74 yield tmpname 

75 with contextlib.suppress(OSError): 

76 os.remove(tmpname) 

77 

78 

79@contextlib.contextmanager 

80def makeSQLiteRegistry(create=True): 

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

82 

83 Yields 

84 ------ 

85 config : `RegistryConfig` 

86 Registry configuration for initialized registry database. 

87 """ 

88 with temporaryDirectory() as tmpdir: 

89 uri = f"sqlite:///{tmpdir}/gen3.sqlite" 

90 config = RegistryConfig() 

91 config["db"] = uri 

92 if create: 

93 Registry.fromConfig(config, create=True) 

94 yield config 

95 

96 

97class SimpleConnections(PipelineTaskConnections, dimensions=(), 

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

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

100 name="{template}schema", 

101 storageClass="SourceCatalog") 

102 

103 

104class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections): 

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

106 

107 def setDefaults(self): 

108 PipelineTaskConfig.setDefaults(self) 

109 

110 

111class TaskOne(PipelineTask): 

112 ConfigClass = SimpleConfig 

113 _DefaultName = "taskOne" 

114 

115 

116class TaskTwo(PipelineTask): 

117 ConfigClass = SimpleConfig 

118 _DefaultName = "taskTwo" 

119 

120 

121class TaskFactoryMock(TaskFactory): 

122 def loadTaskClass(self, taskName): 

123 if taskName == "TaskOne": 

124 return TaskOne, "TaskOne" 

125 elif taskName == "TaskTwo": 

126 return TaskTwo, "TaskTwo" 

127 

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

129 if config is None: 

130 config = taskClass.ConfigClass() 

131 if overrides: 

132 overrides.applyTo(config) 

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

134 

135 

136def _makeArgs(cmd="run", registryConfig=None, **kwargs): 

137 """Return parsed command line arguments. 

138 

139 By default butler_config is set to `Config` populated with some defaults, 

140 it can be overriden completely by keyword argument. 

141 

142 Parameters 

143 ---------- 

144 cmd : `str`, optional 

145 Produce arguments for this pipetask command. 

146 registryConfig : `RegistryConfig`, optional 

147 Override for registry configuration. 

148 **kwargs 

149 Overrides for other arguments. 

150 """ 

151 # call parser for "run" command to set defaults for all arguments 

152 parser = parser_mod.makeParser() 

153 args = parser.parse_args([cmd]) 

154 # override butler_config with our defaults 

155 args.butler_config = Config() 

156 if registryConfig: 

157 args.butler_config["registry"] = registryConfig 

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

159 # some root here for it to use 

160 args.butler_config.configFile = "." 

161 # override arguments from keyword parameters 

162 for key, value in kwargs.items(): 

163 setattr(args, key, value) 

164 return args 

165 

166 

167def _makeQGraph(): 

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

169 

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

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

172 

173 Returns 

174 ------- 

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

176 """ 

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

178 quanta = [Quantum()] 

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

180 qgraph = QuantumGraph([taskNodes]) 

181 return qgraph 

182 

183 

184class CmdLineFwkTestCase(unittest.TestCase): 

185 """A test case for CmdLineFwk 

186 """ 

187 

188 def testMakePipeline(self): 

189 """Tests for CmdLineFwk.makePipeline method 

190 """ 

191 fwk = CmdLineFwk() 

192 

193 # make empty pipeline 

194 args = _makeArgs() 

195 pipeline = fwk.makePipeline(args) 

196 self.assertIsInstance(pipeline, Pipeline) 

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

198 

199 # few tests with serialization 

200 with makeTmpFile() as tmpname: 

201 # make empty pipeline and store it in a file 

202 args = _makeArgs(save_pipeline=tmpname) 

203 pipeline = fwk.makePipeline(args) 

204 self.assertIsInstance(pipeline, Pipeline) 

205 

206 # read pipeline from a file 

207 args = _makeArgs(pipeline=tmpname) 

208 pipeline = fwk.makePipeline(args) 

209 self.assertIsInstance(pipeline, Pipeline) 

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

211 

212 # single task pipeline 

213 actions = [ 

214 _ACTION_ADD_TASK("TaskOne:task1") 

215 ] 

216 args = _makeArgs(pipeline_actions=actions) 

217 pipeline = fwk.makePipeline(args) 

218 self.assertIsInstance(pipeline, Pipeline) 

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

220 

221 # many task pipeline 

222 actions = [ 

223 _ACTION_ADD_TASK("TaskOne:task1a"), 

224 _ACTION_ADD_TASK("TaskTwo:task2"), 

225 _ACTION_ADD_TASK("TaskOne:task1b") 

226 ] 

227 args = _makeArgs(pipeline_actions=actions) 

228 pipeline = fwk.makePipeline(args) 

229 self.assertIsInstance(pipeline, Pipeline) 

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

231 

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

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

234 actions = [ 

235 _ACTION_ADD_TASK("lsst.pipe.base.tests.simpleQGraph.AddTask:task"), 

236 _ACTION_CONFIG("task:addend=100") 

237 ] 

238 args = _makeArgs(pipeline_actions=actions) 

239 pipeline = fwk.makePipeline(args) 

240 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

243 

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

245 with makeTmpFile(overrides) as tmpname: 

246 actions = [ 

247 _ACTION_ADD_TASK("lsst.pipe.base.tests.simpleQGraph.AddTask:task"), 

248 _ACTION_CONFIG_FILE("task:" + tmpname) 

249 ] 

250 args = _makeArgs(pipeline_actions=actions) 

251 pipeline = fwk.makePipeline(args) 

252 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

255 

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

257 actions = [ 

258 _ACTION_ADD_TASK("lsst.pipe.base.tests.simpleQGraph.AddTask:task"), 

259 _ACTION_ADD_INSTRUMENT("Instrument") 

260 ] 

261 args = _makeArgs(pipeline_actions=actions) 

262 pipeline = fwk.makePipeline(args) 

263 

264 def testMakeGraphFromPickle(self): 

265 """Tests for CmdLineFwk.makeGraph method. 

266 

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

268 building. 

269 """ 

270 fwk = CmdLineFwk() 

271 

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

273 

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

275 qgraph = _makeQGraph() 

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

277 qgraph.save(pickleFile) 

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

279 qgraph = fwk.makeGraph(None, args) 

280 self.assertIsInstance(qgraph, QuantumGraph) 

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

282 

283 # pickle with wrong object type 

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

285 pickle.dump({}, pickleFile) 

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

287 with self.assertRaises(TypeError): 

288 fwk.makeGraph(None, args) 

289 

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

291 # will return None and make a warning 

292 qgraph = QuantumGraph() 

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

294 qgraph.save(pickleFile) 

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

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

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

298 qgraph = fwk.makeGraph(None, args) 

299 self.assertIs(qgraph, None) 

300 

301 def testShowPipeline(self): 

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

303 """ 

304 fwk = CmdLineFwk() 

305 

306 actions = [ 

307 _ACTION_ADD_TASK("lsst.pipe.base.tests.simpleQGraph.AddTask:task"), 

308 _ACTION_CONFIG("task:addend=100") 

309 ] 

310 args = _makeArgs(pipeline_actions=actions) 

311 pipeline = fwk.makePipeline(args) 

312 

313 args.show = ["pipeline"] 

314 fwk.showInfo(args, pipeline) 

315 args.show = ["config"] 

316 fwk.showInfo(args, pipeline) 

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

318 fwk.showInfo(args, pipeline) 

319 args.show = ["tasks"] 

320 fwk.showInfo(args, pipeline) 

321 

322 

323class CmdLineFwkTestCaseWithButler(unittest.TestCase): 

324 """A test case for CmdLineFwk 

325 """ 

326 

327 def setUp(self): 

328 super().setUpClass() 

329 self.root = tempfile.mkdtemp() 

330 

331 def tearDown(self): 

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

333 super().tearDownClass() 

334 

335 def testSimpleQGraph(self): 

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

337 """ 

338 

339 nQuanta = 5 

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

341 

342 # should have one task and number of quanta 

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

344 

345 args = _makeArgs() 

346 fwk = CmdLineFwk() 

347 taskFactory = AddTaskFactoryMock() 

348 

349 # run whole thing 

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

351 self.assertEqual(taskFactory.countExec, nQuanta) 

352 

353 def testSimpleQGraphSkipExisting(self): 

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

355 """ 

356 

357 nQuanta = 5 

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

359 

360 # should have one task and number of quanta 

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

362 

363 args = _makeArgs() 

364 fwk = CmdLineFwk() 

365 taskFactory = AddTaskFactoryMock(stopAt=3) 

366 

367 # run first three quanta 

368 with self.assertRaises(RuntimeError): 

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

370 self.assertEqual(taskFactory.countExec, 3) 

371 

372 # run remaining ones 

373 taskFactory.stopAt = -1 

374 args.skip_existing = True 

375 args.no_versions = True 

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

377 self.assertEqual(taskFactory.countExec, nQuanta) 

378 

379 def testSimpleQGraphPartialOutputsFail(self): 

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

381 outputs. 

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 # drop one of the two outputs from one task 

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

401 self.assertIsNotNone(ref) 

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

403 

404 taskFactory.stopAt = -1 

405 args.skip_existing = True 

406 args.no_versions = True 

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

408 with self.assertRaisesRegex(RuntimeError, excRe): 

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

410 

411 def testSimpleQGraphClobberPartialOutputs(self): 

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

413 --clobber-partial-outputs. 

414 """ 

415 

416 nQuanta = 5 

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

418 

419 # should have one task and number of quanta 

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

421 

422 args = _makeArgs() 

423 fwk = CmdLineFwk() 

424 taskFactory = AddTaskFactoryMock(stopAt=3) 

425 

426 # run first three quanta 

427 with self.assertRaises(RuntimeError): 

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

429 self.assertEqual(taskFactory.countExec, 3) 

430 

431 # drop one of the two outputs from one task 

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

433 self.assertIsNotNone(ref) 

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

435 

436 taskFactory.stopAt = -1 

437 args.skip_existing = True 

438 args.clobber_partial_outputs = True 

439 args.no_versions = True 

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

441 # number of executed quanta is incremented 

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

443 

444 def testSimpleQGraphReplaceRun(self): 

445 """Test repeated execution of trivial quantum graph with 

446 --replace-run. 

447 """ 

448 

449 # need non-memory registry in this case 

450 nQuanta = 5 

451 butler, qgraph = makeSimpleQGraph(nQuanta, root=self.root, inMemory=False) 

452 

453 # should have one task and number of quanta 

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

455 

456 fwk = CmdLineFwk() 

457 taskFactory = AddTaskFactoryMock() 

458 

459 # run whole thing 

460 args = _makeArgs( 

461 butler_config=self.root, 

462 input="test", 

463 output="output", 

464 output_run="output/run1") 

465 # deep copy is needed because quanta are updated in place 

466 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args) 

467 self.assertEqual(taskFactory.countExec, nQuanta) 

468 

469 # need to refresh collections explicitly (or make new butler/registry) 

470 butler.registry._collections.refresh() 

471 collections = set(butler.registry.queryCollections(...)) 

472 self.assertEqual(collections, {"test", "output", "output/run1"}) 

473 

474 # number of datasets written by pipeline: 

475 # - nQuanta of init_outputs 

476 # - nQuanta of configs 

477 # - packages (single dataset) 

478 # - nQuanta * two output datasets 

479 # - nQuanta of metadata 

480 n_outputs = nQuanta * 5 + 1 

481 refs = butler.registry.queryDatasets(..., collections="output/run1") 

482 self.assertEqual(len(list(refs)), n_outputs) 

483 

484 # re-run with --replace-run (--inputs is not compatible) 

485 args.input = None 

486 args.replace_run = True 

487 args.output_run = "output/run2" 

488 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args) 

489 

490 butler.registry._collections.refresh() 

491 collections = set(butler.registry.queryCollections(...)) 

492 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2"}) 

493 

494 # new output collection 

495 refs = butler.registry.queryDatasets(..., collections="output/run2") 

496 self.assertEqual(len(list(refs)), n_outputs) 

497 

498 # old output collection is still there 

499 refs = butler.registry.queryDatasets(..., collections="output/run1") 

500 self.assertEqual(len(list(refs)), n_outputs) 

501 

502 # re-run with --replace-run and --prune-replaced=unstore 

503 args.input = None 

504 args.replace_run = True 

505 args.prune_replaced = "unstore" 

506 args.output_run = "output/run3" 

507 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args) 

508 

509 butler.registry._collections.refresh() 

510 collections = set(butler.registry.queryCollections(...)) 

511 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run3"}) 

512 

513 # new output collection 

514 refs = butler.registry.queryDatasets(..., collections="output/run3") 

515 self.assertEqual(len(list(refs)), n_outputs) 

516 

517 # old output collection is still there, and it has all datasets but 

518 # they are not in datastore 

519 refs = butler.registry.queryDatasets(..., collections="output/run2") 

520 refs = list(refs) 

521 self.assertEqual(len(refs), n_outputs) 

522 with self.assertRaises(FileNotFoundError): 

523 butler.get(refs[0], collections="output/run2") 

524 

525 # re-run with --replace-run and --prune-replaced=purge 

526 args.input = None 

527 args.replace_run = True 

528 args.prune_replaced = "purge" 

529 args.output_run = "output/run4" 

530 fwk.runPipeline(copy.deepcopy(qgraph), taskFactory, args) 

531 

532 butler.registry._collections.refresh() 

533 collections = set(butler.registry.queryCollections(...)) 

534 # output/run3 should disappear now 

535 self.assertEqual(collections, {"test", "output", "output/run1", "output/run2", "output/run4"}) 

536 

537 # new output collection 

538 refs = butler.registry.queryDatasets(..., collections="output/run4") 

539 self.assertEqual(len(list(refs)), n_outputs) 

540 

541 def testShowGraph(self): 

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

543 """ 

544 fwk = CmdLineFwk() 

545 

546 nQuanta = 2 

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

548 

549 args = _makeArgs(show=["graph"]) 

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

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

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

553 # mock to that code. 

554 

555 

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

557 pass 

558 

559 

560def setup_module(module): 

561 lsst.utils.tests.init() 

562 

563 

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

565 lsst.utils.tests.init() 

566 unittest.main()