Coverage for tests/test_cmdLineFwk.py: 19%

Shortcuts 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

417 statements  

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 "skypix": { 

221 "common": "htm7", 

222 "htm": { 

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

224 "max_level": 24, 

225 }, 

226 }, 

227 "elements": { 

228 "A": { 

229 "keys": [ 

230 { 

231 "name": "id", 

232 "type": "int", 

233 } 

234 ], 

235 "storage": { 

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

237 }, 

238 }, 

239 "B": { 

240 "keys": [ 

241 { 

242 "name": "id", 

243 "type": "int", 

244 } 

245 ], 

246 "storage": { 

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

248 }, 

249 }, 

250 }, 

251 "packers": {}, 

252 } 

253 ) 

254 universe = DimensionUniverse(config=config) 

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

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

257 quanta = [ 

258 Quantum( 

259 taskName=_TASK_CLASS, 

260 inputs={ 

261 fakeDSType: [ 

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

263 ] 

264 }, 

265 ) 

266 ] # type: ignore 

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

268 return qgraph 

269 

270 

271class CmdLineFwkTestCase(unittest.TestCase): 

272 """A test case for CmdLineFwk""" 

273 

274 def testMakePipeline(self): 

275 """Tests for CmdLineFwk.makePipeline method""" 

276 fwk = CmdLineFwk() 

277 

278 # make empty pipeline 

279 args = _makeArgs() 

280 pipeline = fwk.makePipeline(args) 

281 self.assertIsInstance(pipeline, Pipeline) 

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

283 

284 # few tests with serialization 

285 with makeTmpFile() as tmpname: 

286 # make empty pipeline and store it in a file 

287 args = _makeArgs(save_pipeline=tmpname) 

288 pipeline = fwk.makePipeline(args) 

289 self.assertIsInstance(pipeline, Pipeline) 

290 

291 # read pipeline from a file 

292 args = _makeArgs(pipeline=tmpname) 

293 pipeline = fwk.makePipeline(args) 

294 self.assertIsInstance(pipeline, Pipeline) 

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

296 

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

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

299 args = _makeArgs(pipeline_actions=actions) 

300 pipeline = fwk.makePipeline(args) 

301 self.assertIsInstance(pipeline, Pipeline) 

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

303 

304 # many task pipeline 

305 actions = [ 

306 _ACTION_ADD_TASK("TaskOne:task1a"), 

307 _ACTION_ADD_TASK("TaskTwo:task2"), 

308 _ACTION_ADD_TASK("TaskOne:task1b"), 

309 ] 

310 args = _makeArgs(pipeline_actions=actions) 

311 pipeline = fwk.makePipeline(args) 

312 self.assertIsInstance(pipeline, Pipeline) 

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

314 

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

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

317 args = _makeArgs(pipeline_actions=actions) 

318 pipeline = fwk.makePipeline(args) 

319 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

322 

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

324 with makeTmpFile(overrides) as tmpname: 

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

326 args = _makeArgs(pipeline_actions=actions) 

327 pipeline = fwk.makePipeline(args) 

328 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

331 

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

333 # crash. 

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

335 args = _makeArgs(pipeline_actions=actions) 

336 pipeline = fwk.makePipeline(args) 

337 

338 def testMakeGraphFromSave(self): 

339 """Tests for CmdLineFwk.makeGraph method. 

340 

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

342 building. 

343 """ 

344 fwk = CmdLineFwk() 

345 

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

347 

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

349 qgraph = _makeQGraph() 

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

351 qgraph.save(saveFile) 

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

353 qgraph = fwk.makeGraph(None, args) 

354 self.assertIsInstance(qgraph, QuantumGraph) 

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

356 

357 # will fail if graph id does not match 

358 args = _makeArgs( 

359 qgraph=tmpname, 

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

361 registryConfig=registryConfig, 

362 execution_butler_location=None, 

363 ) 

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

365 fwk.makeGraph(None, args) 

366 

367 # save with wrong object type 

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

369 pickle.dump({}, saveFile) 

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

371 with self.assertRaises(ValueError): 

372 fwk.makeGraph(None, args) 

373 

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

375 # will return None and make a warning 

376 qgraph = QuantumGraph(dict()) 

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

378 qgraph.save(saveFile) 

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

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

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

382 qgraph = fwk.makeGraph(None, args) 

383 self.assertIs(qgraph, None) 

384 

385 def testShowPipeline(self): 

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

387 fwk = CmdLineFwk() 

388 

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

390 args = _makeArgs(pipeline_actions=actions) 

391 pipeline = fwk.makePipeline(args) 

392 

393 args.show = ["pipeline"] 

394 fwk.showInfo(args, pipeline) 

395 args.show = ["config"] 

396 fwk.showInfo(args, pipeline) 

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

398 fwk.showInfo(args, pipeline) 

399 args.show = ["tasks"] 

400 fwk.showInfo(args, pipeline) 

401 

402 

403class CmdLineFwkTestCaseWithButler(unittest.TestCase): 

404 """A test case for CmdLineFwk""" 

405 

406 def setUp(self): 

407 super().setUpClass() 

408 self.root = tempfile.mkdtemp() 

409 self.nQuanta = 5 

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

411 

412 def tearDown(self): 

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

414 super().tearDownClass() 

415 

416 def testSimpleQGraph(self): 

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

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

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

420 populateButler(self.pipeline, butler) 

421 

422 fwk = CmdLineFwk() 

423 taskFactory = AddTaskFactoryMock() 

424 

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

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

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

428 

429 # run whole thing 

430 fwk.runPipeline(qgraph, taskFactory, args) 

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

432 

433 def testSimpleQGraphNoSkipExisting_inputs(self): 

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

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

436 option is present. 

437 """ 

438 args = _makeArgs( 

439 butler_config=self.root, 

440 input="test", 

441 output="output", 

442 ) 

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

444 populateButler( 

445 self.pipeline, 

446 butler, 

447 datasetTypes={ 

448 args.input: [ 

449 "add_dataset0", 

450 "add_dataset1", 

451 "add2_dataset1", 

452 "add_init_output1", 

453 "task0_config", 

454 "task0_metadata", 

455 "task0_log", 

456 ] 

457 }, 

458 ) 

459 

460 fwk = CmdLineFwk() 

461 taskFactory = AddTaskFactoryMock() 

462 

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

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

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

466 # executing one quantum is skipped. 

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

468 

469 # run whole thing 

470 fwk.runPipeline(qgraph, taskFactory, args) 

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

472 

473 def testSimpleQGraphSkipExisting_inputs(self): 

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

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

476 for this case. 

477 """ 

478 args = _makeArgs( 

479 butler_config=self.root, 

480 input="test", 

481 output="output", 

482 skip_existing_in=("test",), 

483 ) 

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

485 populateButler( 

486 self.pipeline, 

487 butler, 

488 datasetTypes={ 

489 args.input: [ 

490 "add_dataset0", 

491 "add_dataset1", 

492 "add2_dataset1", 

493 "add_init_output1", 

494 "task0_config", 

495 "task0_metadata", 

496 "task0_log", 

497 ] 

498 }, 

499 ) 

500 

501 fwk = CmdLineFwk() 

502 taskFactory = AddTaskFactoryMock() 

503 

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

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

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

507 

508 # run whole thing 

509 fwk.runPipeline(qgraph, taskFactory, args) 

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

511 

512 def testSimpleQGraphSkipExisting_outputs(self): 

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

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

515 for this case. 

516 """ 

517 args = _makeArgs( 

518 butler_config=self.root, 

519 input="test", 

520 output_run="output/run", 

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

522 ) 

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

524 populateButler( 

525 self.pipeline, 

526 butler, 

527 datasetTypes={ 

528 args.input: ["add_dataset0"], 

529 args.output_run: [ 

530 "add_dataset1", 

531 "add2_dataset1", 

532 "add_init_output1", 

533 "task0_metadata", 

534 "task0_log", 

535 ], 

536 }, 

537 ) 

538 

539 fwk = CmdLineFwk() 

540 taskFactory = AddTaskFactoryMock() 

541 

542 # fails without --extend-run 

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

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

545 

546 # retry with --extend-run 

547 args.extend_run = True 

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

549 

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

551 # Graph does not include quantum for first task 

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

553 

554 # run whole thing 

555 fwk.runPipeline(qgraph, taskFactory, args) 

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

557 

558 def testSimpleQGraphOutputsFail(self): 

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

560 outputs. 

561 """ 

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

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

564 populateButler(self.pipeline, butler) 

565 

566 fwk = CmdLineFwk() 

567 taskFactory = AddTaskFactoryMock(stopAt=3) 

568 

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

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

571 

572 # run first three quanta 

573 with self.assertRaises(MPGraphExecutorError): 

574 fwk.runPipeline(qgraph, taskFactory, args) 

575 self.assertEqual(taskFactory.countExec, 3) 

576 

577 butler.registry.refresh() 

578 

579 # drop one of the two outputs from one task 

580 ref1 = butler.registry.findDataset( 

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

582 ) 

583 self.assertIsNotNone(ref1) 

584 # also drop the metadata output 

585 ref2 = butler.registry.findDataset( 

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

587 ) 

588 self.assertIsNotNone(ref2) 

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

590 

591 taskFactory.stopAt = -1 

592 args.skip_existing_in = (args.output,) 

593 args.extend_run = True 

594 args.no_versions = True 

595 with self.assertRaises(MPGraphExecutorError): 

596 fwk.runPipeline(qgraph, taskFactory, args) 

597 

598 def testSimpleQGraphClobberOutputs(self): 

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

600 --clobber-outputs. 

601 """ 

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

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

604 populateButler(self.pipeline, butler) 

605 

606 fwk = CmdLineFwk() 

607 taskFactory = AddTaskFactoryMock(stopAt=3) 

608 

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

610 

611 # should have one task and number of quanta 

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

613 

614 # run first three quanta 

615 with self.assertRaises(MPGraphExecutorError): 

616 fwk.runPipeline(qgraph, taskFactory, args) 

617 self.assertEqual(taskFactory.countExec, 3) 

618 

619 butler.registry.refresh() 

620 

621 # drop one of the two outputs from one task 

622 ref1 = butler.registry.findDataset( 

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

624 ) 

625 self.assertIsNotNone(ref1) 

626 # also drop the metadata output 

627 ref2 = butler.registry.findDataset( 

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

629 ) 

630 self.assertIsNotNone(ref2) 

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

632 

633 taskFactory.stopAt = -1 

634 args.skip_existing = True 

635 args.extend_run = True 

636 args.clobber_outputs = True 

637 args.no_versions = True 

638 fwk.runPipeline(qgraph, taskFactory, args) 

639 # number of executed quanta is incremented 

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

641 

642 def testSimpleQGraphReplaceRun(self): 

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

644 --replace-run. 

645 """ 

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

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

648 populateButler(self.pipeline, butler) 

649 

650 fwk = CmdLineFwk() 

651 taskFactory = AddTaskFactoryMock() 

652 

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

654 

655 # should have one task and number of quanta 

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

657 

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

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

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

661 

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

663 butler.registry.refresh() 

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

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

666 

667 # number of datasets written by pipeline: 

668 # - nQuanta of init_outputs 

669 # - nQuanta of configs 

670 # - packages (single dataset) 

671 # - nQuanta * two output datasets 

672 # - nQuanta of metadata 

673 # - nQuanta of log output 

674 n_outputs = self.nQuanta * 6 + 1 

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

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

677 

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

679 # changed) 

680 args.replace_run = True 

681 args.output_run = "output/run2" 

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

683 

684 butler.registry.refresh() 

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

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

687 

688 # new output collection 

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

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

691 

692 # old output collection is still there 

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

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

695 

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

697 args.replace_run = True 

698 args.prune_replaced = "unstore" 

699 args.output_run = "output/run3" 

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

701 

702 butler.registry.refresh() 

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

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

705 

706 # new output collection 

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

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

709 

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

711 # non-InitOutputs are not in datastore 

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

713 refs = list(refs) 

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

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

716 for ref in refs: 

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

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

719 else: 

720 with self.assertRaises(FileNotFoundError): 

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

722 

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

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

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

726 args.input = None 

727 args.replace_run = True 

728 args.prune_replaced = "purge" 

729 args.output_run = "output/run4" 

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

731 

732 butler.registry.refresh() 

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

734 # output/run3 should disappear now 

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

736 

737 # new output collection 

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

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

740 

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

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

743 with self.assertRaises(ValueError): 

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

745 args.prune_replaced = None 

746 args.replace_run = True 

747 args.output_run = "output/run5" 

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

749 butler.registry.refresh() 

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

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

752 with self.assertRaises(ValueError): 

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

754 args.prune_replaced = None 

755 args.replace_run = True 

756 args.output_run = "output/run6" 

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

758 butler.registry.refresh() 

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

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

761 

762 def testMockTask(self): 

763 """Test --mock option.""" 

764 args = _makeArgs( 

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

766 ) 

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

768 populateButler(self.pipeline, butler) 

769 

770 fwk = CmdLineFwk() 

771 taskFactory = AddTaskFactoryMock() 

772 

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

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

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

776 

777 # run whole thing 

778 fwk.runPipeline(qgraph, taskFactory, args) 

779 # None of the actual tasks is executed 

780 self.assertEqual(taskFactory.countExec, 0) 

781 

782 # check dataset types 

783 butler.registry.refresh() 

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

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

786 

787 def testMockTaskFailure(self): 

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

789 args = _makeArgs( 

790 butler_config=self.root, 

791 input="test", 

792 output="output", 

793 mock=True, 

794 register_dataset_types=True, 

795 mock_configs=[ 

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

797 ], 

798 fail_fast=True, 

799 ) 

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

801 populateButler(self.pipeline, butler) 

802 

803 fwk = CmdLineFwk() 

804 taskFactory = AddTaskFactoryMock() 

805 

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

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

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

809 

810 with self.assertRaises(MPGraphExecutorError) as cm: 

811 fwk.runPipeline(qgraph, taskFactory, args) 

812 

813 self.assertIsNotNone(cm.exception.__cause__) 

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

815 

816 def testSubgraph(self): 

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

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

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

820 populateButler(self.pipeline, butler) 

821 

822 fwk = CmdLineFwk() 

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

824 

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

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

827 nNodes = 2 

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

829 nodeIds = nodeIds[:nNodes] 

830 

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

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

833 

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

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

836 qgraph.save(saveFile) 

837 

838 args = _makeArgs( 

839 qgraph=tmpname, 

840 qgraph_node_id=nodeIds, 

841 registryConfig=registryConfig, 

842 execution_butler_location=None, 

843 ) 

844 fwk = CmdLineFwk() 

845 

846 # load graph, should only read a subset 

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

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

849 

850 def testShowGraph(self): 

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

852 fwk = CmdLineFwk() 

853 

854 nQuanta = 2 

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

856 

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

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

859 

860 def testShowGraphWorkflow(self): 

861 fwk = CmdLineFwk() 

862 

863 nQuanta = 2 

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

865 

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

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

868 

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

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

871 # mock to that code. 

872 

873 

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

875 pass 

876 

877 

878def setup_module(module): 

879 lsst.utils.tests.init() 

880 

881 

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

883 lsst.utils.tests.init() 

884 unittest.main()