Coverage for tests/test_cmdLineFwk.py: 17%

437 statements  

« prev     ^ index     » next       coverage.py v6.4.2, created at 2022-07-13 02:54 -0700

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 re 

31import shutil 

32import tempfile 

33import unittest 

34from dataclasses import dataclass 

35from types import SimpleNamespace 

36from typing import NamedTuple 

37 

38import click 

39import lsst.pex.config as pexConfig 

40import lsst.pipe.base.connectionTypes as cT 

41import lsst.utils.tests 

42from lsst.ctrl.mpexec import CmdLineFwk, MPGraphExecutorError 

43from lsst.ctrl.mpexec.cli.opt import run_options 

44from lsst.ctrl.mpexec.cli.utils import ( 

45 _ACTION_ADD_INSTRUMENT, 

46 _ACTION_ADD_TASK, 

47 _ACTION_CONFIG, 

48 _ACTION_CONFIG_FILE, 

49 PipetaskCommand, 

50) 

51from lsst.daf.butler import Config, DataCoordinate, DatasetRef, DimensionUniverse, Quantum, Registry 

52from lsst.daf.butler.core.datasets.type import DatasetType 

53from lsst.daf.butler.registry import RegistryConfig 

54from lsst.pipe.base import ( 

55 Instrument, 

56 Pipeline, 

57 PipelineTaskConfig, 

58 PipelineTaskConnections, 

59 QuantumGraph, 

60 TaskDef, 

61) 

62from lsst.pipe.base.graphBuilder import DatasetQueryConstraintVariant as DQCVariant 

63from lsst.pipe.base.tests.simpleQGraph import ( 

64 AddTask, 

65 AddTaskFactoryMock, 

66 makeSimpleButler, 

67 makeSimplePipeline, 

68 makeSimpleQGraph, 

69 populateButler, 

70) 

71from lsst.utils.tests import temporaryDirectory 

72 

73logging.basicConfig(level=getattr(logging, os.environ.get("UNIT_TEST_LOGGING_LEVEL", "INFO"), logging.INFO)) 

74 

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

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

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

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

79 

80 

81@contextlib.contextmanager 

82def makeTmpFile(contents=None, suffix=None): 

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

84 

85 Temporary file is deleted on exiting context. 

86 

87 Parameters 

88 ---------- 

89 contents : `bytes` 

90 Data to write into a file. 

91 """ 

92 fd, tmpname = tempfile.mkstemp(suffix=suffix) 

93 if contents: 

94 os.write(fd, contents) 

95 os.close(fd) 

96 yield tmpname 

97 with contextlib.suppress(OSError): 

98 os.remove(tmpname) 

99 

100 

101@contextlib.contextmanager 

102def makeSQLiteRegistry(create=True): 

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

104 

105 Yields 

106 ------ 

107 config : `RegistryConfig` 

108 Registry configuration for initialized registry database. 

109 """ 

110 with temporaryDirectory() as tmpdir: 

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

112 config = RegistryConfig() 

113 config["db"] = uri 

114 if create: 

115 Registry.createFromConfig(config) 

116 yield config 

117 

118 

119class SimpleConnections(PipelineTaskConnections, dimensions=(), defaultTemplates={"template": "simple"}): 

120 schema = cT.InitInput(doc="Schema", name="{template}schema", storageClass="SourceCatalog") 

121 

122 

123class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections): 

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

125 

126 def setDefaults(self): 

127 PipelineTaskConfig.setDefaults(self) 

128 

129 

130def _makeArgs(registryConfig=None, **kwargs): 

131 """Return parsed command line arguments. 

132 

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

134 it can be overridden completely by keyword argument. 

135 

136 Parameters 

137 ---------- 

138 cmd : `str`, optional 

139 Produce arguments for this pipetask command. 

140 registryConfig : `RegistryConfig`, optional 

141 Override for registry configuration. 

142 **kwargs 

143 Overrides for other arguments. 

144 """ 

145 # Use a mock to get the default value of arguments to 'run'. 

146 

147 mock = unittest.mock.Mock() 

148 

149 @click.command(cls=PipetaskCommand) 

150 @run_options() 

151 def fake_run(ctx, **kwargs): 

152 """Fake "pipetask run" command for gathering input arguments. 

153 

154 The arguments & options should always match the arguments & options in 

155 the "real" command function `lsst.ctrl.mpexec.cli.cmd.run`. 

156 """ 

157 mock(**kwargs) 

158 

159 runner = click.testing.CliRunner() 

160 # --butler-config is the only required option 

161 result = runner.invoke(fake_run, "--butler-config /") 

162 if result.exit_code != 0: 

163 raise RuntimeError(f"Failure getting default args from 'fake_run': {result}") 

164 mock.assert_called_once() 

165 args = mock.call_args[1] 

166 args["enableLsstDebug"] = args.pop("debug") 

167 args["execution_butler_location"] = args.pop("save_execution_butler") 

168 if "pipeline_actions" not in args: 

169 args["pipeline_actions"] = [] 

170 if "mock_configs" not in args: 

171 args["mock_configs"] = [] 

172 args = SimpleNamespace(**args) 

173 

174 # override butler_config with our defaults 

175 if "butler_config" not in kwargs: 

176 args.butler_config = Config() 

177 if registryConfig: 

178 args.butler_config["registry"] = registryConfig 

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

180 # some root here for it to use 

181 args.butler_config.configFile = "." 

182 

183 # override arguments from keyword parameters 

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

185 setattr(args, key, value) 

186 args.dataset_query_constraint = DQCVariant.fromExpression(args.dataset_query_constraint) 

187 return args 

188 

189 

190class FakeDSType(NamedTuple): 

191 name: str 

192 

193 

194@dataclass(frozen=True) 

195class FakeDSRef: 

196 datasetType: str 

197 dataId: tuple 

198 

199 def isComponent(self): 

200 return False 

201 

202 

203# Task class name used by tests, needs to be importable 

204_TASK_CLASS = "lsst.pipe.base.tests.simpleQGraph.AddTask" 

205 

206 

207def _makeQGraph(): 

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

209 

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

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

212 

213 Returns 

214 ------- 

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

216 """ 

217 config = Config( 

218 { 

219 "version": 1, 

220 "namespace": "ctrl_mpexec_test", 

221 "skypix": { 

222 "common": "htm7", 

223 "htm": { 

224 "class": "lsst.sphgeom.HtmPixelization", 

225 "max_level": 24, 

226 }, 

227 }, 

228 "elements": { 

229 "A": { 

230 "keys": [ 

231 { 

232 "name": "id", 

233 "type": "int", 

234 } 

235 ], 

236 "storage": { 

237 "cls": "lsst.daf.butler.registry.dimensions.table.TableDimensionRecordStorage", 

238 }, 

239 }, 

240 "B": { 

241 "keys": [ 

242 { 

243 "name": "id", 

244 "type": "int", 

245 } 

246 ], 

247 "storage": { 

248 "cls": "lsst.daf.butler.registry.dimensions.table.TableDimensionRecordStorage", 

249 }, 

250 }, 

251 }, 

252 "packers": {}, 

253 } 

254 ) 

255 universe = DimensionUniverse(config=config) 

256 fakeDSType = DatasetType("A", tuple(), storageClass="ExposureF", universe=universe) 

257 taskDef = TaskDef(taskName=_TASK_CLASS, config=AddTask.ConfigClass(), taskClass=AddTask) 

258 quanta = [ 

259 Quantum( 

260 taskName=_TASK_CLASS, 

261 inputs={ 

262 fakeDSType: [ 

263 DatasetRef(fakeDSType, DataCoordinate.standardize({"A": 1, "B": 2}, universe=universe)) 

264 ] 

265 }, 

266 ) 

267 ] # type: ignore 

268 qgraph = QuantumGraph({taskDef: set(quanta)}) 

269 return qgraph 

270 

271 

272class CmdLineFwkTestCase(unittest.TestCase): 

273 """A test case for CmdLineFwk""" 

274 

275 def testMakePipeline(self): 

276 """Tests for CmdLineFwk.makePipeline method""" 

277 fwk = CmdLineFwk() 

278 

279 # make empty pipeline 

280 args = _makeArgs() 

281 pipeline = fwk.makePipeline(args) 

282 self.assertIsInstance(pipeline, Pipeline) 

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

284 

285 # few tests with serialization 

286 with makeTmpFile() as tmpname: 

287 # make empty pipeline and store it in a file 

288 args = _makeArgs(save_pipeline=tmpname) 

289 pipeline = fwk.makePipeline(args) 

290 self.assertIsInstance(pipeline, Pipeline) 

291 

292 # read pipeline from a file 

293 args = _makeArgs(pipeline=tmpname) 

294 pipeline = fwk.makePipeline(args) 

295 self.assertIsInstance(pipeline, Pipeline) 

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

297 

298 # single task pipeline, task name can be anything here 

299 actions = [_ACTION_ADD_TASK("TaskOne:task1")] 

300 args = _makeArgs(pipeline_actions=actions) 

301 pipeline = fwk.makePipeline(args) 

302 self.assertIsInstance(pipeline, Pipeline) 

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

304 

305 # many task pipeline 

306 actions = [ 

307 _ACTION_ADD_TASK("TaskOne:task1a"), 

308 _ACTION_ADD_TASK("TaskTwo:task2"), 

309 _ACTION_ADD_TASK("TaskOne:task1b"), 

310 ] 

311 args = _makeArgs(pipeline_actions=actions) 

312 pipeline = fwk.makePipeline(args) 

313 self.assertIsInstance(pipeline, Pipeline) 

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

315 

316 # single task pipeline with config overrides, need real task class 

317 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG("task:addend=100")] 

318 args = _makeArgs(pipeline_actions=actions) 

319 pipeline = fwk.makePipeline(args) 

320 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

323 

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

325 with makeTmpFile(overrides) as tmpname: 

326 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG_FILE("task:" + tmpname)] 

327 args = _makeArgs(pipeline_actions=actions) 

328 pipeline = fwk.makePipeline(args) 

329 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

332 

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

334 # crash. 

335 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_ADD_INSTRUMENT("Instrument")] 

336 args = _makeArgs(pipeline_actions=actions) 

337 pipeline = fwk.makePipeline(args) 

338 

339 def testMakeGraphFromSave(self): 

340 """Tests for CmdLineFwk.makeGraph method. 

341 

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

343 building. 

344 """ 

345 fwk = CmdLineFwk() 

346 

347 with makeTmpFile(suffix=".qgraph") as tmpname, makeSQLiteRegistry() as registryConfig: 

348 

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

350 qgraph = _makeQGraph() 

351 with open(tmpname, "wb") as saveFile: 

352 qgraph.save(saveFile) 

353 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None) 

354 qgraph = fwk.makeGraph(None, args) 

355 self.assertIsInstance(qgraph, QuantumGraph) 

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

357 

358 # will fail if graph id does not match 

359 args = _makeArgs( 

360 qgraph=tmpname, 

361 qgraph_id="R2-D2 is that you?", 

362 registryConfig=registryConfig, 

363 execution_butler_location=None, 

364 ) 

365 with self.assertRaisesRegex(ValueError, "graphID does not match"): 

366 fwk.makeGraph(None, args) 

367 

368 # save with wrong object type 

369 with open(tmpname, "wb") as saveFile: 

370 pickle.dump({}, saveFile) 

371 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None) 

372 with self.assertRaises(ValueError): 

373 fwk.makeGraph(None, args) 

374 

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

376 # will return None and make a warning 

377 qgraph = QuantumGraph(dict()) 

378 with open(tmpname, "wb") as saveFile: 

379 qgraph.save(saveFile) 

380 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig, execution_butler_location=None) 

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

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

383 qgraph = fwk.makeGraph(None, args) 

384 self.assertIs(qgraph, None) 

385 

386 def testShowPipeline(self): 

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

388 fwk = CmdLineFwk() 

389 

390 actions = [_ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), _ACTION_CONFIG("task:addend=100")] 

391 args = _makeArgs(pipeline_actions=actions) 

392 pipeline = fwk.makePipeline(args) 

393 

394 args.show = ["pipeline"] 

395 fwk.showInfo(args, pipeline) 

396 args.show = ["config"] 

397 fwk.showInfo(args, pipeline) 

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

399 fwk.showInfo(args, pipeline) 

400 args.show = ["tasks"] 

401 fwk.showInfo(args, pipeline) 

402 

403 

404class CmdLineFwkTestCaseWithButler(unittest.TestCase): 

405 """A test case for CmdLineFwk""" 

406 

407 def setUp(self): 

408 super().setUpClass() 

409 self.root = tempfile.mkdtemp() 

410 self.nQuanta = 5 

411 self.pipeline = makeSimplePipeline(nQuanta=self.nQuanta) 

412 

413 def tearDown(self): 

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

415 super().tearDownClass() 

416 

417 def testSimpleQGraph(self): 

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

419 args = _makeArgs(butler_config=self.root, input="test", output="output") 

420 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

421 populateButler(self.pipeline, butler) 

422 

423 fwk = CmdLineFwk() 

424 taskFactory = AddTaskFactoryMock() 

425 

426 qgraph = fwk.makeGraph(self.pipeline, args) 

427 self.assertEqual(len(qgraph.taskGraph), self.nQuanta) 

428 self.assertEqual(len(qgraph), self.nQuanta) 

429 

430 # run whole thing 

431 fwk.runPipeline(qgraph, taskFactory, args) 

432 self.assertEqual(taskFactory.countExec, self.nQuanta) 

433 

434 def testSimpleQGraphNoSkipExisting_inputs(self): 

435 """Test for case when output data for one task already appears in 

436 _input_ collection, but no ``--extend-run`` or ``-skip-existing`` 

437 option is present. 

438 """ 

439 args = _makeArgs( 

440 butler_config=self.root, 

441 input="test", 

442 output="output", 

443 ) 

444 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

445 populateButler( 

446 self.pipeline, 

447 butler, 

448 datasetTypes={ 

449 args.input: [ 

450 "add_dataset0", 

451 "add_dataset1", 

452 "add2_dataset1", 

453 "add_init_output1", 

454 "task0_config", 

455 "task0_metadata", 

456 "task0_log", 

457 ] 

458 }, 

459 ) 

460 

461 fwk = CmdLineFwk() 

462 taskFactory = AddTaskFactoryMock() 

463 

464 qgraph = fwk.makeGraph(self.pipeline, args) 

465 self.assertEqual(len(qgraph.taskGraph), self.nQuanta) 

466 # With current implementation graph has all nQuanta quanta, but when 

467 # executing one quantum is skipped. 

468 self.assertEqual(len(qgraph), self.nQuanta) 

469 

470 # run whole thing 

471 fwk.runPipeline(qgraph, taskFactory, args) 

472 self.assertEqual(taskFactory.countExec, self.nQuanta) 

473 

474 def testSimpleQGraphSkipExisting_inputs(self): 

475 """Test for ``--skip-existing`` with output data for one task already 

476 appears in _input_ collection. No ``--extend-run`` option is needed 

477 for this case. 

478 """ 

479 args = _makeArgs( 

480 butler_config=self.root, 

481 input="test", 

482 output="output", 

483 skip_existing_in=("test",), 

484 ) 

485 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

486 populateButler( 

487 self.pipeline, 

488 butler, 

489 datasetTypes={ 

490 args.input: [ 

491 "add_dataset0", 

492 "add_dataset1", 

493 "add2_dataset1", 

494 "add_init_output1", 

495 "task0_config", 

496 "task0_metadata", 

497 "task0_log", 

498 ] 

499 }, 

500 ) 

501 

502 fwk = CmdLineFwk() 

503 taskFactory = AddTaskFactoryMock() 

504 

505 qgraph = fwk.makeGraph(self.pipeline, args) 

506 self.assertEqual(len(qgraph.taskGraph), self.nQuanta) 

507 self.assertEqual(len(qgraph), self.nQuanta - 1) 

508 

509 # run whole thing 

510 fwk.runPipeline(qgraph, taskFactory, args) 

511 self.assertEqual(taskFactory.countExec, self.nQuanta - 1) 

512 

513 def testSimpleQGraphSkipExisting_outputs(self): 

514 """Test for ``--skip-existing`` with output data for one task already 

515 appears in _output_ collection. The ``--extend-run`` option is needed 

516 for this case. 

517 """ 

518 args = _makeArgs( 

519 butler_config=self.root, 

520 input="test", 

521 output_run="output/run", 

522 skip_existing_in=("output/run",), 

523 ) 

524 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

525 populateButler( 

526 self.pipeline, 

527 butler, 

528 datasetTypes={ 

529 args.input: ["add_dataset0"], 

530 args.output_run: [ 

531 "add_dataset1", 

532 "add2_dataset1", 

533 "add_init_output1", 

534 "task0_metadata", 

535 "task0_log", 

536 ], 

537 }, 

538 ) 

539 

540 fwk = CmdLineFwk() 

541 taskFactory = AddTaskFactoryMock() 

542 

543 # fails without --extend-run 

544 with self.assertRaisesRegex(ValueError, "--extend-run was not given"): 

545 qgraph = fwk.makeGraph(self.pipeline, args) 

546 

547 # retry with --extend-run 

548 args.extend_run = True 

549 qgraph = fwk.makeGraph(self.pipeline, args) 

550 

551 self.assertEqual(len(qgraph.taskGraph), self.nQuanta) 

552 # Graph does not include quantum for first task 

553 self.assertEqual(len(qgraph), self.nQuanta - 1) 

554 

555 # run whole thing 

556 fwk.runPipeline(qgraph, taskFactory, args) 

557 self.assertEqual(taskFactory.countExec, self.nQuanta - 1) 

558 

559 def testSimpleQGraphOutputsFail(self): 

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

561 outputs. 

562 """ 

563 args = _makeArgs(butler_config=self.root, input="test", output="output") 

564 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

565 populateButler(self.pipeline, butler) 

566 

567 fwk = CmdLineFwk() 

568 taskFactory = AddTaskFactoryMock(stopAt=3) 

569 

570 qgraph = fwk.makeGraph(self.pipeline, args) 

571 self.assertEqual(len(qgraph), self.nQuanta) 

572 

573 # run first three quanta 

574 with self.assertRaises(MPGraphExecutorError): 

575 fwk.runPipeline(qgraph, taskFactory, args) 

576 self.assertEqual(taskFactory.countExec, 3) 

577 

578 butler.registry.refresh() 

579 

580 # drop one of the two outputs from one task 

581 ref1 = butler.registry.findDataset( 

582 "add2_dataset2", collections=args.output, instrument="INSTR", detector=0 

583 ) 

584 self.assertIsNotNone(ref1) 

585 # also drop the metadata output 

586 ref2 = butler.registry.findDataset( 

587 "task1_metadata", collections=args.output, instrument="INSTR", detector=0 

588 ) 

589 self.assertIsNotNone(ref2) 

590 butler.pruneDatasets([ref1, ref2], disassociate=True, unstore=True, purge=True) 

591 

592 taskFactory.stopAt = -1 

593 args.skip_existing_in = (args.output,) 

594 args.extend_run = True 

595 args.no_versions = True 

596 with self.assertRaises(MPGraphExecutorError): 

597 fwk.runPipeline(qgraph, taskFactory, args) 

598 

599 def testSimpleQGraphClobberOutputs(self): 

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

601 --clobber-outputs. 

602 """ 

603 args = _makeArgs(butler_config=self.root, input="test", output="output") 

604 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

605 populateButler(self.pipeline, butler) 

606 

607 fwk = CmdLineFwk() 

608 taskFactory = AddTaskFactoryMock(stopAt=3) 

609 

610 qgraph = fwk.makeGraph(self.pipeline, args) 

611 

612 # should have one task and number of quanta 

613 self.assertEqual(len(qgraph), self.nQuanta) 

614 

615 # run first three quanta 

616 with self.assertRaises(MPGraphExecutorError): 

617 fwk.runPipeline(qgraph, taskFactory, args) 

618 self.assertEqual(taskFactory.countExec, 3) 

619 

620 butler.registry.refresh() 

621 

622 # drop one of the two outputs from one task 

623 ref1 = butler.registry.findDataset( 

624 "add2_dataset2", collections=args.output, dataId=dict(instrument="INSTR", detector=0) 

625 ) 

626 self.assertIsNotNone(ref1) 

627 # also drop the metadata output 

628 ref2 = butler.registry.findDataset( 

629 "task1_metadata", collections=args.output, dataId=dict(instrument="INSTR", detector=0) 

630 ) 

631 self.assertIsNotNone(ref2) 

632 butler.pruneDatasets([ref1, ref2], disassociate=True, unstore=True, purge=True) 

633 

634 taskFactory.stopAt = -1 

635 args.skip_existing = True 

636 args.extend_run = True 

637 args.clobber_outputs = True 

638 args.no_versions = True 

639 fwk.runPipeline(qgraph, taskFactory, args) 

640 # number of executed quanta is incremented 

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

642 

643 def testSimpleQGraphReplaceRun(self): 

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

645 --replace-run. 

646 """ 

647 args = _makeArgs(butler_config=self.root, input="test", output="output", output_run="output/run1") 

648 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

649 populateButler(self.pipeline, butler) 

650 

651 fwk = CmdLineFwk() 

652 taskFactory = AddTaskFactoryMock() 

653 

654 qgraph = fwk.makeGraph(self.pipeline, args) 

655 

656 # should have one task and number of quanta 

657 self.assertEqual(len(qgraph), self.nQuanta) 

658 

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

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

661 self.assertEqual(taskFactory.countExec, self.nQuanta) 

662 

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

664 butler.registry.refresh() 

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

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

667 

668 # number of datasets written by pipeline: 

669 # - nQuanta of init_outputs 

670 # - nQuanta of configs 

671 # - packages (single dataset) 

672 # - nQuanta * two output datasets 

673 # - nQuanta of metadata 

674 # - nQuanta of log output 

675 n_outputs = self.nQuanta * 6 + 1 

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

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

678 

679 # re-run with --replace-run (--inputs is ignored, as long as it hasn't 

680 # changed) 

681 args.replace_run = True 

682 args.output_run = "output/run2" 

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

684 

685 butler.registry.refresh() 

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

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

688 

689 # new output collection 

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

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

692 

693 # old output collection is still there 

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

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

696 

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

698 args.replace_run = True 

699 args.prune_replaced = "unstore" 

700 args.output_run = "output/run3" 

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

702 

703 butler.registry.refresh() 

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

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

706 

707 # new output collection 

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

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

710 

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

712 # non-InitOutputs are not in datastore 

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

714 refs = list(refs) 

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

716 initOutNameRe = re.compile("packages|task.*_config|add_init_output.*") 

717 for ref in refs: 

718 if initOutNameRe.fullmatch(ref.datasetType.name): 

719 butler.get(ref, collections="output/run2") 

720 else: 

721 with self.assertRaises(FileNotFoundError): 

722 butler.get(ref, collections="output/run2") 

723 

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

725 # This time also remove --input; passing the same inputs that we 

726 # started with and not passing inputs at all should be equivalent. 

727 args.input = None 

728 args.replace_run = True 

729 args.prune_replaced = "purge" 

730 args.output_run = "output/run4" 

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

732 

733 butler.registry.refresh() 

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

735 # output/run3 should disappear now 

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

737 

738 # new output collection 

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

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

741 

742 # Trying to run again with inputs that aren't exactly what we started 

743 # with is an error, and the kind that should not modify the data repo. 

744 with self.assertRaises(ValueError): 

745 args.input = ["test", "output/run2"] 

746 args.prune_replaced = None 

747 args.replace_run = True 

748 args.output_run = "output/run5" 

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

750 butler.registry.refresh() 

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

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

753 with self.assertRaises(ValueError): 

754 args.input = ["output/run2", "test"] 

755 args.prune_replaced = None 

756 args.replace_run = True 

757 args.output_run = "output/run6" 

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

759 butler.registry.refresh() 

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

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

762 

763 def testMockTask(self): 

764 """Test --mock option.""" 

765 args = _makeArgs( 

766 butler_config=self.root, input="test", output="output", mock=True, register_dataset_types=True 

767 ) 

768 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

769 populateButler(self.pipeline, butler) 

770 

771 fwk = CmdLineFwk() 

772 taskFactory = AddTaskFactoryMock() 

773 

774 qgraph = fwk.makeGraph(self.pipeline, args) 

775 self.assertEqual(len(qgraph.taskGraph), self.nQuanta) 

776 self.assertEqual(len(qgraph), self.nQuanta) 

777 

778 # run whole thing 

779 fwk.runPipeline(qgraph, taskFactory, args) 

780 # None of the actual tasks is executed 

781 self.assertEqual(taskFactory.countExec, 0) 

782 

783 # check dataset types 

784 butler.registry.refresh() 

785 datasetTypes = list(butler.registry.queryDatasetTypes(re.compile("^_mock_.*"))) 

786 self.assertEqual(len(datasetTypes), self.nQuanta * 2) 

787 

788 def testMockTaskFailure(self): 

789 """Test --mock option and configure one of the tasks to fail.""" 

790 args = _makeArgs( 

791 butler_config=self.root, 

792 input="test", 

793 output="output", 

794 mock=True, 

795 register_dataset_types=True, 

796 mock_configs=[ 

797 _ACTION_CONFIG("task3-mock:failCondition='detector = 0'"), 

798 ], 

799 fail_fast=True, 

800 ) 

801 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

802 populateButler(self.pipeline, butler) 

803 

804 fwk = CmdLineFwk() 

805 taskFactory = AddTaskFactoryMock() 

806 

807 qgraph = fwk.makeGraph(self.pipeline, args) 

808 self.assertEqual(len(qgraph.taskGraph), self.nQuanta) 

809 self.assertEqual(len(qgraph), self.nQuanta) 

810 

811 with self.assertRaises(MPGraphExecutorError) as cm: 

812 fwk.runPipeline(qgraph, taskFactory, args) 

813 

814 self.assertIsNotNone(cm.exception.__cause__) 

815 self.assertRegex(str(cm.exception.__cause__), "Simulated failure: task=task3") 

816 

817 def testSubgraph(self): 

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

819 args = _makeArgs(butler_config=self.root, input="test", output="output") 

820 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

821 populateButler(self.pipeline, butler) 

822 

823 fwk = CmdLineFwk() 

824 qgraph = fwk.makeGraph(self.pipeline, args) 

825 

826 # Select first two nodes for execution. This depends on node ordering 

827 # which I assume is the same as execution order. 

828 nNodes = 2 

829 nodeIds = [node.nodeId for node in qgraph] 

830 nodeIds = nodeIds[:nNodes] 

831 

832 self.assertEqual(len(qgraph.taskGraph), self.nQuanta) 

833 self.assertEqual(len(qgraph), self.nQuanta) 

834 

835 with makeTmpFile(suffix=".qgraph") as tmpname, makeSQLiteRegistry() as registryConfig: 

836 with open(tmpname, "wb") as saveFile: 

837 qgraph.save(saveFile) 

838 

839 args = _makeArgs( 

840 qgraph=tmpname, 

841 qgraph_node_id=nodeIds, 

842 registryConfig=registryConfig, 

843 execution_butler_location=None, 

844 ) 

845 fwk = CmdLineFwk() 

846 

847 # load graph, should only read a subset 

848 qgraph = fwk.makeGraph(pipeline=None, args=args) 

849 self.assertEqual(len(qgraph), nNodes) 

850 

851 def testShowGraph(self): 

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

853 fwk = CmdLineFwk() 

854 

855 nQuanta = 2 

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

857 

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

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

860 

861 def testShowGraphWorkflow(self): 

862 fwk = CmdLineFwk() 

863 

864 nQuanta = 2 

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

866 

867 args = _makeArgs(show=["workflow"]) 

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

869 

870 # TODO: cannot test "uri" option presently, it instanciates 

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

872 # mock to that code. 

873 

874 def testSimpleQGraphDatastoreRecords(self): 

875 """Test quantum graph generation with --qgraph-datastore-records.""" 

876 args = _makeArgs( 

877 butler_config=self.root, input="test", output="output", qgraph_datastore_records=True 

878 ) 

879 butler = makeSimpleButler(self.root, run=args.input, inMemory=False) 

880 populateButler(self.pipeline, butler) 

881 

882 fwk = CmdLineFwk() 

883 qgraph = fwk.makeGraph(self.pipeline, args) 

884 self.assertEqual(len(qgraph), self.nQuanta) 

885 for i, qnode in enumerate(qgraph): 

886 quantum = qnode.quantum 

887 self.assertIsNotNone(quantum.datastore_records) 

888 # only the first quantum has a pre-existing input 

889 if i == 0: 

890 datastore_name = "FileDatastore@<butlerRoot>" 

891 self.assertEqual(set(quantum.datastore_records.keys()), {datastore_name}) 

892 records_data = quantum.datastore_records[datastore_name] 

893 records = dict(records_data.records) 

894 self.assertEqual(len(records), 1) 

895 _, records = records.popitem() 

896 records = records["file_datastore_records"] 

897 self.assertEqual( 

898 [record.path for record in records], 

899 ["test/add_dataset0/add_dataset0_INSTR_det0_test.pickle"], 

900 ) 

901 else: 

902 self.assertEqual(quantum.datastore_records, {}) 

903 

904 

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

906 pass 

907 

908 

909def setup_module(module): 

910 lsst.utils.tests.init() 

911 

912 

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

914 lsst.utils.tests.init() 

915 unittest.main()