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 testUtil import (AddTaskFactoryMock, makeSimpleQGraph) 

48 

49 

50logging.basicConfig(level=logging.INFO) 

51 

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

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

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

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

56 

57 

58@contextlib.contextmanager 

59def makeTmpFile(contents=None): 

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

61 

62 Temporary file is deleted on exiting context. 

63 

64 Parameters 

65 ---------- 

66 contents : `bytes` 

67 Data to write into a file. 

68 """ 

69 fd, tmpname = tempfile.mkstemp() 

70 if contents: 

71 os.write(fd, contents) 

72 os.close(fd) 

73 yield tmpname 

74 with contextlib.suppress(OSError): 

75 os.remove(tmpname) 

76 

77 

78@contextlib.contextmanager 

79def temporaryDirectory(): 

80 """Context manager that creates and destroys temporary directory. 

81 

82 Difference from `tempfile.TemporaryDirectory` is that it ignores errors 

83 when deleting a directory, which may happen with some filesystems. 

84 """ 

85 tmpdir = tempfile.mkdtemp() 

86 yield tmpdir 

87 shutil.rmtree(tmpdir, ignore_errors=True) 

88 

89 

90@contextlib.contextmanager 

91def makeSQLiteRegistry(create=True): 

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

93 

94 Yields 

95 ------ 

96 config : `RegistryConfig` 

97 Registry configuration for initialized registry database. 

98 """ 

99 with temporaryDirectory() as tmpdir: 

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

101 config = RegistryConfig() 

102 config["db"] = uri 

103 if create: 

104 Registry.fromConfig(config, create=True) 

105 yield config 

106 

107 

108class SimpleConnections(PipelineTaskConnections, dimensions=(), 

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

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

111 name="{template}schema", 

112 storageClass="SourceCatalog") 

113 

114 

115class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections): 

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

117 

118 def setDefaults(self): 

119 PipelineTaskConfig.setDefaults(self) 

120 

121 

122class TaskOne(PipelineTask): 

123 ConfigClass = SimpleConfig 

124 _DefaultName = "taskOne" 

125 

126 

127class TaskTwo(PipelineTask): 

128 ConfigClass = SimpleConfig 

129 _DefaultName = "taskTwo" 

130 

131 

132class TaskFactoryMock(TaskFactory): 

133 def loadTaskClass(self, taskName): 

134 if taskName == "TaskOne": 

135 return TaskOne, "TaskOne" 

136 elif taskName == "TaskTwo": 

137 return TaskTwo, "TaskTwo" 

138 

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

140 if config is None: 

141 config = taskClass.ConfigClass() 

142 if overrides: 

143 overrides.applyTo(config) 

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

145 

146 

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

148 """Return parsed command line arguments. 

149 

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

151 it can be overriden completely by keyword argument. 

152 

153 Parameters 

154 ---------- 

155 cmd : `str`, optional 

156 Produce arguments for this pipetask command. 

157 registryConfig : `RegistryConfig`, optional 

158 Override for registry configuration. 

159 **kwargs 

160 Overrides for other arguments. 

161 """ 

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

163 parser = parser_mod.makeParser() 

164 args = parser.parse_args([cmd]) 

165 # override butler_config with our defaults 

166 args.butler_config = Config() 

167 if registryConfig: 

168 args.butler_config["registry"] = registryConfig 

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

170 # some root here for it to use 

171 args.butler_config.configFile = "." 

172 # override arguments from keyword parameters 

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

174 setattr(args, key, value) 

175 return args 

176 

177 

178def _makeQGraph(): 

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

180 

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

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

183 

184 Returns 

185 ------- 

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

187 """ 

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

189 quanta = [Quantum()] 

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

191 qgraph = QuantumGraph([taskNodes]) 

192 return qgraph 

193 

194 

195class CmdLineFwkTestCase(unittest.TestCase): 

196 """A test case for CmdLineFwk 

197 """ 

198 

199 def testMakePipeline(self): 

200 """Tests for CmdLineFwk.makePipeline method 

201 """ 

202 fwk = CmdLineFwk() 

203 

204 # make empty pipeline 

205 args = _makeArgs() 

206 pipeline = fwk.makePipeline(args) 

207 self.assertIsInstance(pipeline, Pipeline) 

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

209 

210 # few tests with serialization 

211 with makeTmpFile() as tmpname: 

212 # make empty pipeline and store it in a file 

213 args = _makeArgs(save_pipeline=tmpname) 

214 pipeline = fwk.makePipeline(args) 

215 self.assertIsInstance(pipeline, Pipeline) 

216 

217 # read pipeline from a file 

218 args = _makeArgs(pipeline=tmpname) 

219 pipeline = fwk.makePipeline(args) 

220 self.assertIsInstance(pipeline, Pipeline) 

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

222 

223 # single task pipeline 

224 actions = [ 

225 _ACTION_ADD_TASK("TaskOne:task1") 

226 ] 

227 args = _makeArgs(pipeline_actions=actions) 

228 pipeline = fwk.makePipeline(args) 

229 self.assertIsInstance(pipeline, Pipeline) 

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

231 

232 # many task pipeline 

233 actions = [ 

234 _ACTION_ADD_TASK("TaskOne:task1a"), 

235 _ACTION_ADD_TASK("TaskTwo:task2"), 

236 _ACTION_ADD_TASK("TaskOne:task1b") 

237 ] 

238 args = _makeArgs(pipeline_actions=actions) 

239 pipeline = fwk.makePipeline(args) 

240 self.assertIsInstance(pipeline, Pipeline) 

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

242 

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

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

245 actions = [ 

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

247 _ACTION_CONFIG("task:addend=100") 

248 ] 

249 args = _makeArgs(pipeline_actions=actions) 

250 pipeline = fwk.makePipeline(args) 

251 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

254 

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

256 with makeTmpFile(overrides) as tmpname: 

257 actions = [ 

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

259 _ACTION_CONFIG_FILE("task:" + tmpname) 

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

266 

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

268 actions = [ 

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

270 _ACTION_ADD_INSTRUMENT("Instrument") 

271 ] 

272 args = _makeArgs(pipeline_actions=actions) 

273 pipeline = fwk.makePipeline(args) 

274 

275 def testMakeGraphFromPickle(self): 

276 """Tests for CmdLineFwk.makeGraph method. 

277 

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

279 building. 

280 """ 

281 fwk = CmdLineFwk() 

282 

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

284 

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

286 qgraph = _makeQGraph() 

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

288 qgraph.save(pickleFile) 

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

290 qgraph = fwk.makeGraph(None, args) 

291 self.assertIsInstance(qgraph, QuantumGraph) 

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

293 

294 # pickle with wrong object type 

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

296 pickle.dump({}, pickleFile) 

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

298 with self.assertRaises(TypeError): 

299 fwk.makeGraph(None, args) 

300 

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

302 # will return None and make a warning 

303 qgraph = QuantumGraph() 

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

305 qgraph.save(pickleFile) 

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

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

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

309 qgraph = fwk.makeGraph(None, args) 

310 self.assertIs(qgraph, None) 

311 

312 def testShowPipeline(self): 

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

314 """ 

315 fwk = CmdLineFwk() 

316 

317 actions = [ 

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

319 _ACTION_CONFIG("task:addend=100") 

320 ] 

321 args = _makeArgs(pipeline_actions=actions) 

322 pipeline = fwk.makePipeline(args) 

323 

324 args.show = ["pipeline"] 

325 fwk.showInfo(args, pipeline) 

326 args.show = ["config"] 

327 fwk.showInfo(args, pipeline) 

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

329 fwk.showInfo(args, pipeline) 

330 args.show = ["tasks"] 

331 fwk.showInfo(args, pipeline) 

332 

333 

334class CmdLineFwkTestCaseWithButler(unittest.TestCase): 

335 """A test case for CmdLineFwk 

336 """ 

337 

338 def setUp(self): 

339 super().setUpClass() 

340 self.root = tempfile.mkdtemp() 

341 

342 def tearDown(self): 

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

344 super().tearDownClass() 

345 

346 def testSimpleQGraph(self): 

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

348 """ 

349 

350 nQuanta = 5 

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

352 

353 # should have one task and number of quanta 

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

355 

356 args = _makeArgs() 

357 fwk = CmdLineFwk() 

358 taskFactory = AddTaskFactoryMock() 

359 

360 # run whole thing 

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

362 self.assertEqual(taskFactory.countExec, nQuanta) 

363 

364 def testSimpleQGraphSkipExisting(self): 

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

366 """ 

367 

368 nQuanta = 5 

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

370 

371 # should have one task and number of quanta 

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

373 

374 args = _makeArgs() 

375 fwk = CmdLineFwk() 

376 taskFactory = AddTaskFactoryMock(stopAt=3) 

377 

378 # run first three quanta 

379 with self.assertRaises(RuntimeError): 

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

381 self.assertEqual(taskFactory.countExec, 3) 

382 

383 # run remaining ones 

384 taskFactory.stopAt = -1 

385 args.skip_existing = True 

386 args.no_versions = True 

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

388 self.assertEqual(taskFactory.countExec, nQuanta) 

389 

390 def testSimpleQGraphPartialOutputsFail(self): 

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

392 outputs. 

393 """ 

394 

395 nQuanta = 5 

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

397 

398 # should have one task and number of quanta 

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

400 

401 args = _makeArgs() 

402 fwk = CmdLineFwk() 

403 taskFactory = AddTaskFactoryMock(stopAt=3) 

404 

405 # run first three quanta 

406 with self.assertRaises(RuntimeError): 

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

408 self.assertEqual(taskFactory.countExec, 3) 

409 

410 # drop one of the two outputs from one task 

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

412 self.assertIsNotNone(ref) 

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

414 

415 taskFactory.stopAt = -1 

416 args.skip_existing = True 

417 args.no_versions = True 

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

419 with self.assertRaisesRegex(RuntimeError, excRe): 

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

421 

422 def testSimpleQGraphClobberPartialOutputs(self): 

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

424 --clobber-partial-outputs. 

425 """ 

426 

427 nQuanta = 5 

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

429 

430 # should have one task and number of quanta 

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

432 

433 args = _makeArgs() 

434 fwk = CmdLineFwk() 

435 taskFactory = AddTaskFactoryMock(stopAt=3) 

436 

437 # run first three quanta 

438 with self.assertRaises(RuntimeError): 

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

440 self.assertEqual(taskFactory.countExec, 3) 

441 

442 # drop one of the two outputs from one task 

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

444 self.assertIsNotNone(ref) 

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

446 

447 taskFactory.stopAt = -1 

448 args.skip_existing = True 

449 args.clobber_partial_outputs = True 

450 args.no_versions = True 

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

452 # number of executed quanta is incremented 

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

454 

455 def testSimpleQGraphReplaceRun(self): 

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

457 --replace-run. 

458 """ 

459 

460 # need non-memory registry in this case 

461 nQuanta = 5 

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

463 

464 # should have one task and number of quanta 

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

466 

467 fwk = CmdLineFwk() 

468 taskFactory = AddTaskFactoryMock() 

469 

470 # run whole thing 

471 args = _makeArgs( 

472 butler_config=self.root, 

473 input="test", 

474 output="output", 

475 output_run="output/run1") 

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

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

478 self.assertEqual(taskFactory.countExec, nQuanta) 

479 

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

481 butler.registry._collections.refresh() 

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

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

484 

485 # number of datasets written by pipeline: 

486 # - nQuanta of init_outputs 

487 # - nQuanta of configs 

488 # - packages (single dataset) 

489 # - nQuanta * two output datasets 

490 # - nQuanta of metadata 

491 n_outputs = nQuanta * 5 + 1 

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

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

494 

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

496 args.input = None 

497 args.replace_run = True 

498 args.output_run = "output/run2" 

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

500 

501 butler.registry._collections.refresh() 

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

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

504 

505 # new output collection 

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

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

508 

509 # old output collection is still there 

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

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

512 

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

514 args.input = None 

515 args.replace_run = True 

516 args.prune_replaced = "unstore" 

517 args.output_run = "output/run3" 

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

519 

520 butler.registry._collections.refresh() 

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

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

523 

524 # new output collection 

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

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

527 

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

529 # they are not in datastore 

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

531 refs = list(refs) 

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

533 with self.assertRaises(FileNotFoundError): 

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

535 

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

537 args.input = None 

538 args.replace_run = True 

539 args.prune_replaced = "purge" 

540 args.output_run = "output/run4" 

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

542 

543 butler.registry._collections.refresh() 

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

545 # output/run3 should disappear now 

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

547 

548 # new output collection 

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

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

551 

552 def testShowGraph(self): 

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

554 """ 

555 fwk = CmdLineFwk() 

556 

557 nQuanta = 2 

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

559 

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

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

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

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

564 # mock to that code. 

565 

566 

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

568 pass 

569 

570 

571def setup_module(module): 

572 lsst.utils.tests.init() 

573 

574 

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

576 lsst.utils.tests.init() 

577 unittest.main()