Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# This file is part of ctrl_mpexec. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22"""Simple unit test for cmdLineFwk module. 

23""" 

24 

25import click 

26from types import SimpleNamespace 

27import contextlib 

28import copy 

29from dataclasses import dataclass 

30import logging 

31import os 

32import pickle 

33import re 

34import shutil 

35import tempfile 

36from typing import NamedTuple 

37import unittest 

38 

39from lsst.ctrl.mpexec.cmdLineFwk import CmdLineFwk 

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

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

42 _ACTION_ADD_TASK, 

43 _ACTION_CONFIG, 

44 _ACTION_CONFIG_FILE, 

45 _ACTION_ADD_INSTRUMENT, 

46 PipetaskCommand, 

47) 

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

49from lsst.daf.butler.registry import RegistryConfig 

50from lsst.obs.base import Instrument 

51import lsst.pex.config as pexConfig 

52from lsst.pipe.base import (Pipeline, PipelineTaskConfig, QuantumGraph, TaskDef, PipelineTaskConnections) 

53import lsst.pipe.base.connectionTypes as cT 

54import lsst.utils.tests 

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

56 AddTaskFactoryMock, 

57 makeSimpleButler, 

58 makeSimplePipeline, 

59 makeSimpleQGraph, 

60 populateButler) 

61from lsst.utils.tests import temporaryDirectory 

62 

63 

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

65 

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

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

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

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

70 

71 

72@contextlib.contextmanager 

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

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

75 

76 Temporary file is deleted on exiting context. 

77 

78 Parameters 

79 ---------- 

80 contents : `bytes` 

81 Data to write into a file. 

82 """ 

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

84 if contents: 

85 os.write(fd, contents) 

86 os.close(fd) 

87 yield tmpname 

88 with contextlib.suppress(OSError): 

89 os.remove(tmpname) 

90 

91 

92@contextlib.contextmanager 

93def makeSQLiteRegistry(create=True): 

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

95 

96 Yields 

97 ------ 

98 config : `RegistryConfig` 

99 Registry configuration for initialized registry database. 

100 """ 

101 with temporaryDirectory() as tmpdir: 

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

103 config = RegistryConfig() 

104 config["db"] = uri 

105 if create: 

106 Registry.createFromConfig(config) 

107 yield config 

108 

109 

110class SimpleConnections(PipelineTaskConnections, dimensions=(), 

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

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

113 name="{template}schema", 

114 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 args = SimpleNamespace(**args) 

165 

166 # override butler_config with our defaults 

167 if "butler_config" not in kwargs: 

168 args.butler_config = Config() 

169 if registryConfig: 

170 args.butler_config["registry"] = registryConfig 

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

172 # some root here for it to use 

173 args.butler_config.configFile = "." 

174 

175 # override arguments from keyword parameters 

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

177 setattr(args, key, value) 

178 return args 

179 

180 

181class FakeDSType(NamedTuple): 

182 name: str 

183 

184 

185@dataclass(frozen=True) 

186class FakeDSRef: 

187 datasetType: str 

188 dataId: tuple 

189 

190 

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

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

193 

194 

195def _makeQGraph(): 

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

197 

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

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

200 

201 Returns 

202 ------- 

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

204 """ 

205 taskDef = TaskDef(taskName=_TASK_CLASS, config=SimpleConfig()) 

206 quanta = [Quantum(taskName=_TASK_CLASS, 

207 inputs={FakeDSType("A"): [FakeDSRef("A", (1, 2))]})] # type: ignore 

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

209 return qgraph 

210 

211 

212class CmdLineFwkTestCase(unittest.TestCase): 

213 """A test case for CmdLineFwk 

214 """ 

215 

216 def testMakePipeline(self): 

217 """Tests for CmdLineFwk.makePipeline method 

218 """ 

219 fwk = CmdLineFwk() 

220 

221 # make empty pipeline 

222 args = _makeArgs() 

223 pipeline = fwk.makePipeline(args) 

224 self.assertIsInstance(pipeline, Pipeline) 

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

226 

227 # few tests with serialization 

228 with makeTmpFile() as tmpname: 

229 # make empty pipeline and store it in a file 

230 args = _makeArgs(save_pipeline=tmpname) 

231 pipeline = fwk.makePipeline(args) 

232 self.assertIsInstance(pipeline, Pipeline) 

233 

234 # read pipeline from a file 

235 args = _makeArgs(pipeline=tmpname) 

236 pipeline = fwk.makePipeline(args) 

237 self.assertIsInstance(pipeline, Pipeline) 

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

239 

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

241 actions = [ 

242 _ACTION_ADD_TASK("TaskOne:task1") 

243 ] 

244 args = _makeArgs(pipeline_actions=actions) 

245 pipeline = fwk.makePipeline(args) 

246 self.assertIsInstance(pipeline, Pipeline) 

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

248 

249 # many task pipeline 

250 actions = [ 

251 _ACTION_ADD_TASK("TaskOne:task1a"), 

252 _ACTION_ADD_TASK("TaskTwo:task2"), 

253 _ACTION_ADD_TASK("TaskOne:task1b") 

254 ] 

255 args = _makeArgs(pipeline_actions=actions) 

256 pipeline = fwk.makePipeline(args) 

257 self.assertIsInstance(pipeline, Pipeline) 

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

259 

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

261 actions = [ 

262 _ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), 

263 _ACTION_CONFIG("task:addend=100") 

264 ] 

265 args = _makeArgs(pipeline_actions=actions) 

266 pipeline = fwk.makePipeline(args) 

267 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

270 

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

272 with makeTmpFile(overrides) as tmpname: 

273 actions = [ 

274 _ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), 

275 _ACTION_CONFIG_FILE("task:" + tmpname) 

276 ] 

277 args = _makeArgs(pipeline_actions=actions) 

278 pipeline = fwk.makePipeline(args) 

279 taskDefs = list(pipeline.toExpandedPipeline()) 

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

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

282 

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

284 # crash. 

285 actions = [ 

286 _ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), 

287 _ACTION_ADD_INSTRUMENT("Instrument") 

288 ] 

289 args = _makeArgs(pipeline_actions=actions) 

290 pipeline = fwk.makePipeline(args) 

291 

292 def testMakeGraphFromSave(self): 

293 """Tests for CmdLineFwk.makeGraph method. 

294 

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

296 building. 

297 """ 

298 fwk = CmdLineFwk() 

299 

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

301 

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

303 qgraph = _makeQGraph() 

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

305 qgraph.save(saveFile) 

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

307 qgraph = fwk.makeGraph(None, args) 

308 self.assertIsInstance(qgraph, QuantumGraph) 

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

310 

311 # will fail if graph id does not match 

312 args = _makeArgs( 

313 qgraph=tmpname, 

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

315 registryConfig=registryConfig, 

316 execution_butler_location=None 

317 ) 

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

319 fwk.makeGraph(None, args) 

320 

321 # save with wrong object type 

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

323 pickle.dump({}, saveFile) 

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

325 with self.assertRaises(ValueError): 

326 fwk.makeGraph(None, args) 

327 

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

329 # will return None and make a warning 

330 qgraph = QuantumGraph(dict()) 

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

332 qgraph.save(saveFile) 

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

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

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

336 qgraph = fwk.makeGraph(None, args) 

337 self.assertIs(qgraph, None) 

338 

339 def testShowPipeline(self): 

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

341 """ 

342 fwk = CmdLineFwk() 

343 

344 actions = [ 

345 _ACTION_ADD_TASK(f"{_TASK_CLASS}:task"), 

346 _ACTION_CONFIG("task:addend=100") 

347 ] 

348 args = _makeArgs(pipeline_actions=actions) 

349 pipeline = fwk.makePipeline(args) 

350 

351 args.show = ["pipeline"] 

352 fwk.showInfo(args, pipeline) 

353 args.show = ["config"] 

354 fwk.showInfo(args, pipeline) 

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

356 fwk.showInfo(args, pipeline) 

357 args.show = ["tasks"] 

358 fwk.showInfo(args, pipeline) 

359 

360 

361class CmdLineFwkTestCaseWithButler(unittest.TestCase): 

362 """A test case for CmdLineFwk 

363 """ 

364 

365 def setUp(self): 

366 super().setUpClass() 

367 self.root = tempfile.mkdtemp() 

368 self.nQuanta = 5 

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

370 

371 def tearDown(self): 

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

373 super().tearDownClass() 

374 

375 def testSimpleQGraph(self): 

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

377 """ 

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

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

380 populateButler(self.pipeline, butler) 

381 

382 fwk = CmdLineFwk() 

383 taskFactory = AddTaskFactoryMock() 

384 

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

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

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

388 

389 # run whole thing 

390 fwk.runPipeline(qgraph, taskFactory, args) 

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

392 

393 def testSimpleQGraphNoSkipExisting_inputs(self): 

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

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

396 option is present. 

397 """ 

398 args = _makeArgs( 

399 butler_config=self.root, 

400 input="test", 

401 output="output", 

402 ) 

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

404 populateButler( 

405 self.pipeline, butler, 

406 datasetTypes={args.input: [ 

407 "add_dataset0", 

408 "add_dataset1", "add2_dataset1", 

409 "add_init_output1", 

410 "task0_config", 

411 "task0_metadata", 

412 "task0_log", 

413 ]} 

414 ) 

415 

416 fwk = CmdLineFwk() 

417 taskFactory = AddTaskFactoryMock() 

418 

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

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

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

422 # executing one quantum is skipped. 

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

424 

425 # run whole thing 

426 fwk.runPipeline(qgraph, taskFactory, args) 

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

428 

429 def testSimpleQGraphSkipExisting_inputs(self): 

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

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

432 for this case. 

433 """ 

434 args = _makeArgs( 

435 butler_config=self.root, 

436 input="test", 

437 output="output", 

438 skip_existing_in=("test", ), 

439 ) 

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

441 populateButler( 

442 self.pipeline, butler, 

443 datasetTypes={args.input: [ 

444 "add_dataset0", 

445 "add_dataset1", "add2_dataset1", 

446 "add_init_output1", 

447 "task0_config", 

448 "task0_metadata", 

449 "task0_log", 

450 ]} 

451 ) 

452 

453 fwk = CmdLineFwk() 

454 taskFactory = AddTaskFactoryMock() 

455 

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

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

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

459 

460 # run whole thing 

461 fwk.runPipeline(qgraph, taskFactory, args) 

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

463 

464 def testSimpleQGraphSkipExisting_outputs(self): 

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

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

467 for this case. 

468 """ 

469 args = _makeArgs( 

470 butler_config=self.root, 

471 input="test", 

472 output_run="output/run", 

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

474 ) 

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

476 populateButler( 

477 self.pipeline, butler, datasetTypes={ 

478 args.input: ["add_dataset0"], 

479 args.output_run: [ 

480 "add_dataset1", "add2_dataset1", 

481 "add_init_output1", 

482 "task0_metadata", 

483 "task0_log", 

484 ] 

485 } 

486 ) 

487 

488 fwk = CmdLineFwk() 

489 taskFactory = AddTaskFactoryMock() 

490 

491 # fails without --extend-run 

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

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

494 

495 # retry with --extend-run 

496 args.extend_run = True 

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

498 

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

500 # Graph does not include quantum for first task 

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

502 

503 # run whole thing 

504 fwk.runPipeline(qgraph, taskFactory, args) 

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

506 

507 def testSimpleQGraphOutputsFail(self): 

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

509 outputs. 

510 """ 

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

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

513 populateButler(self.pipeline, butler) 

514 

515 fwk = CmdLineFwk() 

516 taskFactory = AddTaskFactoryMock(stopAt=3) 

517 

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

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

520 

521 # run first three quanta 

522 with self.assertRaises(RuntimeError): 

523 fwk.runPipeline(qgraph, taskFactory, args) 

524 self.assertEqual(taskFactory.countExec, 3) 

525 

526 butler.registry.refresh() 

527 

528 # drop one of the two outputs from one task 

529 ref1 = butler.registry.findDataset("add2_dataset2", collections=args.output, 

530 instrument="INSTR", detector=0) 

531 self.assertIsNotNone(ref1) 

532 # also drop the metadata output 

533 ref2 = butler.registry.findDataset("task1_metadata", collections=args.output, 

534 instrument="INSTR", detector=0) 

535 self.assertIsNotNone(ref2) 

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

537 

538 taskFactory.stopAt = -1 

539 args.skip_existing_in = (args.output, ) 

540 args.extend_run = True 

541 args.no_versions = True 

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

543 with self.assertRaisesRegex(RuntimeError, excRe): 

544 fwk.runPipeline(qgraph, taskFactory, args) 

545 

546 def testSimpleQGraphClobberOutputs(self): 

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

548 --clobber-outputs. 

549 """ 

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

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

552 populateButler(self.pipeline, butler) 

553 

554 fwk = CmdLineFwk() 

555 taskFactory = AddTaskFactoryMock(stopAt=3) 

556 

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

558 

559 # should have one task and number of quanta 

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

561 

562 # run first three quanta 

563 with self.assertRaises(RuntimeError): 

564 fwk.runPipeline(qgraph, taskFactory, args) 

565 self.assertEqual(taskFactory.countExec, 3) 

566 

567 butler.registry.refresh() 

568 

569 # drop one of the two outputs from one task 

570 ref1 = butler.registry.findDataset("add2_dataset2", collections=args.output, 

571 dataId=dict(instrument="INSTR", detector=0)) 

572 self.assertIsNotNone(ref1) 

573 # also drop the metadata output 

574 ref2 = butler.registry.findDataset("task1_metadata", collections=args.output, 

575 dataId=dict(instrument="INSTR", detector=0)) 

576 self.assertIsNotNone(ref2) 

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

578 

579 taskFactory.stopAt = -1 

580 args.skip_existing = True 

581 args.extend_run = True 

582 args.clobber_outputs = True 

583 args.no_versions = True 

584 fwk.runPipeline(qgraph, taskFactory, args) 

585 # number of executed quanta is incremented 

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

587 

588 def testSimpleQGraphReplaceRun(self): 

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

590 --replace-run. 

591 """ 

592 args = _makeArgs( 

593 butler_config=self.root, 

594 input="test", 

595 output="output", 

596 output_run="output/run1") 

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

598 populateButler(self.pipeline, butler) 

599 

600 fwk = CmdLineFwk() 

601 taskFactory = AddTaskFactoryMock() 

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 # deep copy is needed because quanta are updated in place 

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

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

611 

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

613 butler.registry.refresh() 

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

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

616 

617 # number of datasets written by pipeline: 

618 # - nQuanta of init_outputs 

619 # - nQuanta of configs 

620 # - packages (single dataset) 

621 # - nQuanta * two output datasets 

622 # - nQuanta of metadata 

623 # - nQuanta of log output 

624 n_outputs = self.nQuanta * 6 + 1 

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

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

627 

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

629 # changed) 

630 args.replace_run = True 

631 args.output_run = "output/run2" 

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

633 

634 butler.registry.refresh() 

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

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

637 

638 # new output collection 

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

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

641 

642 # old output collection is still there 

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

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

645 

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

647 args.replace_run = True 

648 args.prune_replaced = "unstore" 

649 args.output_run = "output/run3" 

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

651 

652 butler.registry.refresh() 

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

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

655 

656 # new output collection 

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

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

659 

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

661 # non-InitOutputs are not in datastore 

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

663 refs = list(refs) 

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

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

666 for ref in refs: 

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

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

669 else: 

670 with self.assertRaises(FileNotFoundError): 

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

672 

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

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

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

676 args.input = None 

677 args.replace_run = True 

678 args.prune_replaced = "purge" 

679 args.output_run = "output/run4" 

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

681 

682 butler.registry.refresh() 

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

684 # output/run3 should disappear now 

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

686 

687 # new output collection 

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

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

690 

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

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

693 with self.assertRaises(ValueError): 

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

695 args.prune_replaced = None 

696 args.replace_run = True 

697 args.output_run = "output/run5" 

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

699 butler.registry.refresh() 

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

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

702 with self.assertRaises(ValueError): 

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

704 args.prune_replaced = None 

705 args.replace_run = True 

706 args.output_run = "output/run6" 

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

708 butler.registry.refresh() 

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

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

711 

712 def testSubgraph(self): 

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

714 """ 

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

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

717 populateButler(self.pipeline, butler) 

718 

719 fwk = CmdLineFwk() 

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

721 

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

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

724 nNodes = 2 

725 nodeIds = [node.nodeId.number for node in qgraph] 

726 nodeIds = nodeIds[:nNodes] 

727 

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

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

730 

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

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

733 qgraph.save(saveFile) 

734 

735 args = _makeArgs(qgraph=tmpname, qgraph_node_id=nodeIds, registryConfig=registryConfig, 

736 execution_butler_location=None) 

737 fwk = CmdLineFwk() 

738 

739 # load graph, should only read a subset 

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

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

742 

743 def testShowGraph(self): 

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

745 """ 

746 fwk = CmdLineFwk() 

747 

748 nQuanta = 2 

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

750 

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

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

753 

754 def testShowGraphWorkflow(self): 

755 fwk = CmdLineFwk() 

756 

757 nQuanta = 2 

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

759 

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

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

762 

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

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

765 # mock to that code. 

766 

767 

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

769 pass 

770 

771 

772def setup_module(module): 

773 lsst.utils.tests.init() 

774 

775 

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

777 lsst.utils.tests.init() 

778 unittest.main()