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

418 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.obs.base import Instrument 

55from lsst.pipe.base import Pipeline, PipelineTaskConfig, PipelineTaskConnections, QuantumGraph, TaskDef 

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

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

58 AddTask, 

59 AddTaskFactoryMock, 

60 makeSimpleButler, 

61 makeSimplePipeline, 

62 makeSimpleQGraph, 

63 populateButler, 

64) 

65from lsst.utils.tests import temporaryDirectory 

66 

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

68 

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

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

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

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

73 

74 

75@contextlib.contextmanager 

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

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

78 

79 Temporary file is deleted on exiting context. 

80 

81 Parameters 

82 ---------- 

83 contents : `bytes` 

84 Data to write into a file. 

85 """ 

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

87 if contents: 

88 os.write(fd, contents) 

89 os.close(fd) 

90 yield tmpname 

91 with contextlib.suppress(OSError): 

92 os.remove(tmpname) 

93 

94 

95@contextlib.contextmanager 

96def makeSQLiteRegistry(create=True): 

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

98 

99 Yields 

100 ------ 

101 config : `RegistryConfig` 

102 Registry configuration for initialized registry database. 

103 """ 

104 with temporaryDirectory() as tmpdir: 

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

106 config = RegistryConfig() 

107 config["db"] = uri 

108 if create: 

109 Registry.createFromConfig(config) 

110 yield config 

111 

112 

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

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

115 

116 

117class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections): 

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

119 

120 def setDefaults(self): 

121 PipelineTaskConfig.setDefaults(self) 

122 

123 

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

125 """Return parsed command line arguments. 

126 

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

128 it can be overridden completely by keyword argument. 

129 

130 Parameters 

131 ---------- 

132 cmd : `str`, optional 

133 Produce arguments for this pipetask command. 

134 registryConfig : `RegistryConfig`, optional 

135 Override for registry configuration. 

136 **kwargs 

137 Overrides for other arguments. 

138 """ 

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

140 

141 mock = unittest.mock.Mock() 

142 

143 @click.command(cls=PipetaskCommand) 

144 @run_options() 

145 def fake_run(ctx, **kwargs): 

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

147 

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

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

150 """ 

151 mock(**kwargs) 

152 

153 runner = click.testing.CliRunner() 

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

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

156 if result.exit_code != 0: 

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

158 mock.assert_called_once() 

159 args = mock.call_args[1] 

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

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

162 if "pipeline_actions" not in args: 

163 args["pipeline_actions"] = [] 

164 if "mock_configs" not in args: 

165 args["mock_configs"] = [] 

166 args = SimpleNamespace(**args) 

167 

168 # override butler_config with our defaults 

169 if "butler_config" not in kwargs: 

170 args.butler_config = Config() 

171 if registryConfig: 

172 args.butler_config["registry"] = registryConfig 

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

174 # some root here for it to use 

175 args.butler_config.configFile = "." 

176 

177 # override arguments from keyword parameters 

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

179 setattr(args, key, value) 

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

181 return args 

182 

183 

184class FakeDSType(NamedTuple): 

185 name: str 

186 

187 

188@dataclass(frozen=True) 

189class FakeDSRef: 

190 datasetType: str 

191 dataId: tuple 

192 

193 def isComponent(self): 

194 return False 

195 

196 

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

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

199 

200 

201def _makeQGraph(): 

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

203 

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

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

206 

207 Returns 

208 ------- 

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

210 """ 

211 config = Config( 

212 { 

213 "version": 1, 

214 "skypix": { 

215 "common": "htm7", 

216 "htm": { 

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

218 "max_level": 24, 

219 }, 

220 }, 

221 "elements": { 

222 "A": { 

223 "keys": [ 

224 { 

225 "name": "id", 

226 "type": "int", 

227 } 

228 ], 

229 "storage": { 

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

231 }, 

232 }, 

233 "B": { 

234 "keys": [ 

235 { 

236 "name": "id", 

237 "type": "int", 

238 } 

239 ], 

240 "storage": { 

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

242 }, 

243 }, 

244 }, 

245 "packers": {}, 

246 } 

247 ) 

248 universe = DimensionUniverse(config=config) 

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

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

251 quanta = [ 

252 Quantum( 

253 taskName=_TASK_CLASS, 

254 inputs={ 

255 fakeDSType: [ 

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

257 ] 

258 }, 

259 ) 

260 ] # type: ignore 

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

262 return qgraph 

263 

264 

265class CmdLineFwkTestCase(unittest.TestCase): 

266 """A test case for CmdLineFwk""" 

267 

268 def testMakePipeline(self): 

269 """Tests for CmdLineFwk.makePipeline method""" 

270 fwk = CmdLineFwk() 

271 

272 # make empty pipeline 

273 args = _makeArgs() 

274 pipeline = fwk.makePipeline(args) 

275 self.assertIsInstance(pipeline, Pipeline) 

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

277 

278 # few tests with serialization 

279 with makeTmpFile() as tmpname: 

280 # make empty pipeline and store it in a file 

281 args = _makeArgs(save_pipeline=tmpname) 

282 pipeline = fwk.makePipeline(args) 

283 self.assertIsInstance(pipeline, Pipeline) 

284 

285 # read pipeline from a file 

286 args = _makeArgs(pipeline=tmpname) 

287 pipeline = fwk.makePipeline(args) 

288 self.assertIsInstance(pipeline, Pipeline) 

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

290 

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

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

293 args = _makeArgs(pipeline_actions=actions) 

294 pipeline = fwk.makePipeline(args) 

295 self.assertIsInstance(pipeline, Pipeline) 

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

297 

298 # many task pipeline 

299 actions = [ 

300 _ACTION_ADD_TASK("TaskOne:task1a"), 

301 _ACTION_ADD_TASK("TaskTwo:task2"), 

302 _ACTION_ADD_TASK("TaskOne:task1b"), 

303 ] 

304 args = _makeArgs(pipeline_actions=actions) 

305 pipeline = fwk.makePipeline(args) 

306 self.assertIsInstance(pipeline, Pipeline) 

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

308 

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

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

311 args = _makeArgs(pipeline_actions=actions) 

312 pipeline = fwk.makePipeline(args) 

313 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

316 

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

318 with makeTmpFile(overrides) as tmpname: 

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

320 args = _makeArgs(pipeline_actions=actions) 

321 pipeline = fwk.makePipeline(args) 

322 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

325 

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

327 # crash. 

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

329 args = _makeArgs(pipeline_actions=actions) 

330 pipeline = fwk.makePipeline(args) 

331 

332 def testMakeGraphFromSave(self): 

333 """Tests for CmdLineFwk.makeGraph method. 

334 

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

336 building. 

337 """ 

338 fwk = CmdLineFwk() 

339 

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

341 

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

343 qgraph = _makeQGraph() 

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

345 qgraph.save(saveFile) 

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

347 qgraph = fwk.makeGraph(None, args) 

348 self.assertIsInstance(qgraph, QuantumGraph) 

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

350 

351 # will fail if graph id does not match 

352 args = _makeArgs( 

353 qgraph=tmpname, 

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

355 registryConfig=registryConfig, 

356 execution_butler_location=None, 

357 ) 

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

359 fwk.makeGraph(None, args) 

360 

361 # save with wrong object type 

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

363 pickle.dump({}, saveFile) 

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

365 with self.assertRaises(ValueError): 

366 fwk.makeGraph(None, args) 

367 

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

369 # will return None and make a warning 

370 qgraph = QuantumGraph(dict()) 

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

372 qgraph.save(saveFile) 

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

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

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

376 qgraph = fwk.makeGraph(None, args) 

377 self.assertIs(qgraph, None) 

378 

379 def testShowPipeline(self): 

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

381 fwk = CmdLineFwk() 

382 

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

384 args = _makeArgs(pipeline_actions=actions) 

385 pipeline = fwk.makePipeline(args) 

386 

387 args.show = ["pipeline"] 

388 fwk.showInfo(args, pipeline) 

389 args.show = ["config"] 

390 fwk.showInfo(args, pipeline) 

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

392 fwk.showInfo(args, pipeline) 

393 args.show = ["tasks"] 

394 fwk.showInfo(args, pipeline) 

395 

396 

397class CmdLineFwkTestCaseWithButler(unittest.TestCase): 

398 """A test case for CmdLineFwk""" 

399 

400 def setUp(self): 

401 super().setUpClass() 

402 self.root = tempfile.mkdtemp() 

403 self.nQuanta = 5 

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

405 

406 def tearDown(self): 

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

408 super().tearDownClass() 

409 

410 def testSimpleQGraph(self): 

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

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

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

414 populateButler(self.pipeline, butler) 

415 

416 fwk = CmdLineFwk() 

417 taskFactory = AddTaskFactoryMock() 

418 

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

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

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

422 

423 # run whole thing 

424 fwk.runPipeline(qgraph, taskFactory, args) 

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

426 

427 def testSimpleQGraphNoSkipExisting_inputs(self): 

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

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

430 option is present. 

431 """ 

432 args = _makeArgs( 

433 butler_config=self.root, 

434 input="test", 

435 output="output", 

436 ) 

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

438 populateButler( 

439 self.pipeline, 

440 butler, 

441 datasetTypes={ 

442 args.input: [ 

443 "add_dataset0", 

444 "add_dataset1", 

445 "add2_dataset1", 

446 "add_init_output1", 

447 "task0_config", 

448 "task0_metadata", 

449 "task0_log", 

450 ] 

451 }, 

452 ) 

453 

454 fwk = CmdLineFwk() 

455 taskFactory = AddTaskFactoryMock() 

456 

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

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

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

460 # executing one quantum is skipped. 

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

462 

463 # run whole thing 

464 fwk.runPipeline(qgraph, taskFactory, args) 

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

466 

467 def testSimpleQGraphSkipExisting_inputs(self): 

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

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

470 for this case. 

471 """ 

472 args = _makeArgs( 

473 butler_config=self.root, 

474 input="test", 

475 output="output", 

476 skip_existing_in=("test",), 

477 ) 

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

479 populateButler( 

480 self.pipeline, 

481 butler, 

482 datasetTypes={ 

483 args.input: [ 

484 "add_dataset0", 

485 "add_dataset1", 

486 "add2_dataset1", 

487 "add_init_output1", 

488 "task0_config", 

489 "task0_metadata", 

490 "task0_log", 

491 ] 

492 }, 

493 ) 

494 

495 fwk = CmdLineFwk() 

496 taskFactory = AddTaskFactoryMock() 

497 

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

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

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

501 

502 # run whole thing 

503 fwk.runPipeline(qgraph, taskFactory, args) 

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

505 

506 def testSimpleQGraphSkipExisting_outputs(self): 

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

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

509 for this case. 

510 """ 

511 args = _makeArgs( 

512 butler_config=self.root, 

513 input="test", 

514 output_run="output/run", 

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

516 ) 

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

518 populateButler( 

519 self.pipeline, 

520 butler, 

521 datasetTypes={ 

522 args.input: ["add_dataset0"], 

523 args.output_run: [ 

524 "add_dataset1", 

525 "add2_dataset1", 

526 "add_init_output1", 

527 "task0_metadata", 

528 "task0_log", 

529 ], 

530 }, 

531 ) 

532 

533 fwk = CmdLineFwk() 

534 taskFactory = AddTaskFactoryMock() 

535 

536 # fails without --extend-run 

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

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

539 

540 # retry with --extend-run 

541 args.extend_run = True 

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

543 

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

545 # Graph does not include quantum for first task 

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

547 

548 # run whole thing 

549 fwk.runPipeline(qgraph, taskFactory, args) 

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

551 

552 def testSimpleQGraphOutputsFail(self): 

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

554 outputs. 

555 """ 

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

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

558 populateButler(self.pipeline, butler) 

559 

560 fwk = CmdLineFwk() 

561 taskFactory = AddTaskFactoryMock(stopAt=3) 

562 

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

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

565 

566 # run first three quanta 

567 with self.assertRaises(MPGraphExecutorError): 

568 fwk.runPipeline(qgraph, taskFactory, args) 

569 self.assertEqual(taskFactory.countExec, 3) 

570 

571 butler.registry.refresh() 

572 

573 # drop one of the two outputs from one task 

574 ref1 = butler.registry.findDataset( 

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

576 ) 

577 self.assertIsNotNone(ref1) 

578 # also drop the metadata output 

579 ref2 = butler.registry.findDataset( 

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

581 ) 

582 self.assertIsNotNone(ref2) 

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

584 

585 taskFactory.stopAt = -1 

586 args.skip_existing_in = (args.output,) 

587 args.extend_run = True 

588 args.no_versions = True 

589 with self.assertRaises(MPGraphExecutorError): 

590 fwk.runPipeline(qgraph, taskFactory, args) 

591 

592 def testSimpleQGraphClobberOutputs(self): 

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

594 --clobber-outputs. 

595 """ 

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

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

598 populateButler(self.pipeline, butler) 

599 

600 fwk = CmdLineFwk() 

601 taskFactory = AddTaskFactoryMock(stopAt=3) 

602 

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

604 

605 # should have one task and number of quanta 

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

607 

608 # run first three quanta 

609 with self.assertRaises(MPGraphExecutorError): 

610 fwk.runPipeline(qgraph, taskFactory, args) 

611 self.assertEqual(taskFactory.countExec, 3) 

612 

613 butler.registry.refresh() 

614 

615 # drop one of the two outputs from one task 

616 ref1 = butler.registry.findDataset( 

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

618 ) 

619 self.assertIsNotNone(ref1) 

620 # also drop the metadata output 

621 ref2 = butler.registry.findDataset( 

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

623 ) 

624 self.assertIsNotNone(ref2) 

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

626 

627 taskFactory.stopAt = -1 

628 args.skip_existing = True 

629 args.extend_run = True 

630 args.clobber_outputs = True 

631 args.no_versions = True 

632 fwk.runPipeline(qgraph, taskFactory, args) 

633 # number of executed quanta is incremented 

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

635 

636 def testSimpleQGraphReplaceRun(self): 

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

638 --replace-run. 

639 """ 

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

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

642 populateButler(self.pipeline, butler) 

643 

644 fwk = CmdLineFwk() 

645 taskFactory = AddTaskFactoryMock() 

646 

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

648 

649 # should have one task and number of quanta 

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

651 

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

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

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

655 

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

657 butler.registry.refresh() 

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

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

660 

661 # number of datasets written by pipeline: 

662 # - nQuanta of init_outputs 

663 # - nQuanta of configs 

664 # - packages (single dataset) 

665 # - nQuanta * two output datasets 

666 # - nQuanta of metadata 

667 # - nQuanta of log output 

668 n_outputs = self.nQuanta * 6 + 1 

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

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

671 

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

673 # changed) 

674 args.replace_run = True 

675 args.output_run = "output/run2" 

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

677 

678 butler.registry.refresh() 

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

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

681 

682 # new output collection 

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

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

685 

686 # old output collection is still there 

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

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

689 

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

691 args.replace_run = True 

692 args.prune_replaced = "unstore" 

693 args.output_run = "output/run3" 

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

695 

696 butler.registry.refresh() 

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

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

699 

700 # new output collection 

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

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

703 

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

705 # non-InitOutputs are not in datastore 

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

707 refs = list(refs) 

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

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

710 for ref in refs: 

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

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

713 else: 

714 with self.assertRaises(FileNotFoundError): 

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

716 

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

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

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

720 args.input = None 

721 args.replace_run = True 

722 args.prune_replaced = "purge" 

723 args.output_run = "output/run4" 

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

725 

726 butler.registry.refresh() 

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

728 # output/run3 should disappear now 

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

730 

731 # new output collection 

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

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

734 

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

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

737 with self.assertRaises(ValueError): 

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

739 args.prune_replaced = None 

740 args.replace_run = True 

741 args.output_run = "output/run5" 

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

743 butler.registry.refresh() 

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

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

746 with self.assertRaises(ValueError): 

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

748 args.prune_replaced = None 

749 args.replace_run = True 

750 args.output_run = "output/run6" 

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

752 butler.registry.refresh() 

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

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

755 

756 def testMockTask(self): 

757 """Test --mock option.""" 

758 args = _makeArgs( 

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

760 ) 

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

762 populateButler(self.pipeline, butler) 

763 

764 fwk = CmdLineFwk() 

765 taskFactory = AddTaskFactoryMock() 

766 

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

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

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

770 

771 # run whole thing 

772 fwk.runPipeline(qgraph, taskFactory, args) 

773 # None of the actual tasks is executed 

774 self.assertEqual(taskFactory.countExec, 0) 

775 

776 # check dataset types 

777 butler.registry.refresh() 

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

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

780 

781 def testMockTaskFailure(self): 

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

783 args = _makeArgs( 

784 butler_config=self.root, 

785 input="test", 

786 output="output", 

787 mock=True, 

788 register_dataset_types=True, 

789 mock_configs=[ 

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

791 ], 

792 fail_fast=True, 

793 ) 

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

795 populateButler(self.pipeline, butler) 

796 

797 fwk = CmdLineFwk() 

798 taskFactory = AddTaskFactoryMock() 

799 

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

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

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

803 

804 with self.assertRaises(MPGraphExecutorError) as cm: 

805 fwk.runPipeline(qgraph, taskFactory, args) 

806 

807 self.assertIsNotNone(cm.exception.__cause__) 

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

809 

810 def testSubgraph(self): 

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

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

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

814 populateButler(self.pipeline, butler) 

815 

816 fwk = CmdLineFwk() 

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

818 

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

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

821 nNodes = 2 

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

823 nodeIds = nodeIds[:nNodes] 

824 

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

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

827 

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

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

830 qgraph.save(saveFile) 

831 

832 args = _makeArgs( 

833 qgraph=tmpname, 

834 qgraph_node_id=nodeIds, 

835 registryConfig=registryConfig, 

836 execution_butler_location=None, 

837 ) 

838 fwk = CmdLineFwk() 

839 

840 # load graph, should only read a subset 

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

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

843 

844 def testShowGraph(self): 

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

846 fwk = CmdLineFwk() 

847 

848 nQuanta = 2 

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

850 

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

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

853 

854 def testShowGraphWorkflow(self): 

855 fwk = CmdLineFwk() 

856 

857 nQuanta = 2 

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

859 

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

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

862 

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

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

865 # mock to that code. 

866 

867 

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

869 pass 

870 

871 

872def setup_module(module): 

873 lsst.utils.tests.init() 

874 

875 

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

877 lsst.utils.tests.init() 

878 unittest.main()