Coverage for tests/test_run.py: 11%

314 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-10-14 01:44 -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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

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

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

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

18# (at your option) any later version. 

19# 

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

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

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

23# GNU General Public License for more details. 

24# 

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

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

27 

28import logging 

29import os 

30import time 

31import unittest 

32import unittest.mock 

33 

34import click.testing 

35 

36import lsst.utils.tests 

37from lsst.ctrl.mpexec import PipelineGraphFactory 

38from lsst.ctrl.mpexec.cli import opt, script 

39from lsst.ctrl.mpexec.cli.cmd.commands import PipetaskCommand, coverage_context 

40from lsst.ctrl.mpexec.cli.utils import collect_pipeline_actions 

41from lsst.ctrl.mpexec.showInfo import ShowInfo 

42from lsst.daf.butler import CollectionType, MissingCollectionError 

43from lsst.pipe.base.mp_graph_executor import MPGraphExecutorError 

44from lsst.pipe.base.script import transfer_from_graph 

45from lsst.pipe.base.tests.mocks import DirectButlerRepo, DynamicTestPipelineTaskConfig 

46 

47 

48class RunTestCase(unittest.TestCase): 

49 """Test pipetask run command-line.""" 

50 

51 @staticmethod 

52 def _make_run_args(*args: str, **kwargs: object) -> dict[str, object]: 

53 mock = unittest.mock.Mock() 

54 

55 @click.command(cls=PipetaskCommand) 

56 @opt.run_options() 

57 @opt.config_search_path_option() 

58 @opt.no_existing_outputs_option() 

59 def fake_run(ctx: click.Context, **kwargs: object): 

60 kwargs = collect_pipeline_actions(ctx, **kwargs) 

61 mock(**kwargs) 

62 

63 runner = click.testing.CliRunner() 

64 result = runner.invoke(fake_run, args, catch_exceptions=False) 

65 if result.exit_code != 0: 

66 raise RuntimeError(f"Failure getting default args for 'run': {result}") 

67 mock.assert_called_once() 

68 result: dict[str, object] = mock.call_args[1] 

69 result["show"] = ShowInfo([]) 

70 result.update(kwargs) 

71 return result 

72 

73 def test_missing_options(self): 

74 """Test that if options for the run script are missing that it 

75 fails. 

76 """ 

77 

78 @click.command() 

79 @opt.pipeline_build_options() 

80 def cli(**kwargs): 

81 script.run(**kwargs) 

82 

83 runner = click.testing.CliRunner() 

84 result = runner.invoke(cli) 

85 # The cli call should fail, because qgraph.run takes more options 

86 # than are defined by pipeline_build_options. 

87 self.assertNotEqual(result.exit_code, 0) 

88 

89 def test_simple_qg(self): 

90 """Test execution of a trivial quantum graph.""" 

91 with DirectButlerRepo.make_temporary() as (helper, root): 

92 helper.add_task() 

93 helper.add_task() 

94 helper.insert_datasets("dataset_auto0") 

95 kwargs = self._make_run_args( 

96 "-b", 

97 root, 

98 "-i", 

99 helper.input_chain, 

100 "-o", 

101 "output", 

102 "--register-dataset-types", 

103 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

104 ) 

105 qg = script.qgraph(**kwargs) 

106 self.assertEqual(len(qg.quanta_by_task), 2) 

107 self.assertEqual(len(qg), 2) 

108 # Ensure that the output run used in the graph is also used in the 

109 # pipeline execution. It is possible for 'qgraph' and 'run' to 

110 # calculate time-stamped runs across a second boundary. 

111 kwargs["output_run"] = qg.header.output_run 

112 # Execute the graph and check for output existence. 

113 script.run(qg, **kwargs) 

114 with helper.butler.query() as query: 

115 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

116 self.assertEqual(query.datasets("dataset_auto2", collections=["output"]).count(), 1) 

117 # Test that we've disabled implicit threading 

118 self.assertEqual(os.environ["OMP_NUM_THREADS"], "1") 

119 

120 def test_simple_qg_rebase(self): 

121 """Test execution of a trivial quantum graph, with --rebase used to 

122 force redefinition of the output collection. 

123 """ 

124 with DirectButlerRepo.make_temporary(input_chain="test1") as (helper, root): 

125 helper.add_task() 

126 helper.add_task() 

127 helper.insert_datasets("dataset_auto0") 

128 # Pass one input collection here for the usual test setup; we'll 

129 # override it later. 

130 kwargs = self._make_run_args( 

131 "-b", 

132 root, 

133 "-i", 

134 helper.input_chain, 

135 "-o", 

136 "output", 

137 "--register-dataset-types", 

138 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

139 ) 

140 # We'll actually pass two input collections in. One is empty, but 

141 # the stuff we're testing here doesn't care. 

142 kwargs["input"] = ["test2", "test1"] 

143 helper.butler.collections.register("test2", CollectionType.RUN) 

144 # Set up the output collection with a sequence that doesn't end the 

145 # same way as the input collection. This is normally an error. 

146 helper.butler.collections.register("output", CollectionType.CHAINED) 

147 helper.butler.collections.register("unexpected_input", CollectionType.RUN) 

148 helper.butler.collections.register("output/run0", CollectionType.RUN) 

149 helper.butler.collections.redefine_chain( 

150 "output", ["test2", "unexpected_input", "test1", "output/run0"] 

151 ) 

152 # Without --rebase, the inconsistent input and output collections 

153 # are an error. 

154 with self.assertRaises(ValueError): 

155 script.qgraph(**kwargs) 

156 # With --rebase, the output collection gets redefined. 

157 kwargs["rebase"] = True 

158 qg = script.qgraph(**kwargs) 

159 self.assertEqual(len(qg.quanta_by_task), 2) 

160 self.assertEqual(len(qg), 2) 

161 # Ensure that the output run used in the graph is also used in the 

162 # pipeline execution. 

163 kwargs["output_run"] = qg.header.output_run 

164 # Execute the graph and check for output existence. 

165 script.run(qg, **kwargs) 

166 with helper.butler.query() as query: 

167 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

168 self.assertEqual(query.datasets("dataset_auto2", collections=["output"]).count(), 1) 

169 

170 def test_simple_qgraph_qbb(self): 

171 """Test execution of a trivial quantum graph in QBB mode.""" 

172 with DirectButlerRepo.make_temporary() as (helper, root): 

173 helper.add_task() 

174 helper.add_task() 

175 helper.insert_datasets("dataset_auto0") 

176 # It's unusual to put a QG in a butler root, but since we've 

177 # already got a temp dir, we might as well use it. 

178 qg_file_1 = os.path.join(root, "test1.qg") 

179 kwargs = self._make_run_args( 

180 "-b", 

181 root, 

182 "-i", 

183 helper.input_chain, 

184 "-o", 

185 "output", 

186 "--register-dataset-types", 

187 "--qgraph-datastore-records", 

188 "--save-qgraph", 

189 qg_file_1, 

190 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

191 ) 

192 qg = script.qgraph(**kwargs) 

193 output_run = qg.header.output_run 

194 output = qg.header.output 

195 self.assertEqual(len(qg.quanta_by_task), 2) 

196 self.assertEqual(len(qg), 2) 

197 # Execute with QBB. 

198 kwargs.update(output_run=output_run, qgraph=qg_file_1) 

199 script.pre_exec_init_qbb(**kwargs) 

200 script.run_qbb(**kwargs) 

201 # Transfer the datasets to the butler. 

202 n1 = transfer_from_graph( 

203 qg_file_1, 

204 root, 

205 register_dataset_types=True, 

206 transfer_dimensions=False, 

207 update_output_chain=True, 

208 dry_run=False, 

209 dataset_type=(), 

210 ) 

211 self.assertEqual(n1, 9) 

212 # Check that the expected outputs exist. 

213 with helper.butler.query() as query: 

214 self.assertEqual(query.datasets("dataset_auto1", collections=output).count(), 1) 

215 self.assertEqual(query.datasets("dataset_auto2", collections=output).count(), 1) 

216 # Check that some metadata keys were written. 

217 some_task_label = next(iter(qg.pipeline_graph.tasks)) 

218 (some_metadata_ref,) = helper.butler.query_datasets( 

219 f"{some_task_label}_metadata", 

220 limit=1, 

221 collections=output, 

222 ) 

223 some_metadata = helper.butler.get(some_metadata_ref) 

224 self.assertIn("qg_read_time", some_metadata["job"]) 

225 self.assertIn("qg_size", some_metadata["job"]) 

226 

227 # Update the output run and try again. 

228 new_output_run = output_run + "_new" 

229 qg_file_2 = os.path.join(root, "test2.qg") 

230 script.update_graph_run(qg_file_1, new_output_run, qg_file_2) 

231 kwargs.update(qgraph=qg_file_2) 

232 # Execute with QBB again. 

233 script.pre_exec_init_qbb(**kwargs) 

234 script.run_qbb(**kwargs) 

235 # Transfer the datasets to the butler. 

236 n2 = transfer_from_graph( 

237 qg_file_2, 

238 root, 

239 register_dataset_types=True, 

240 transfer_dimensions=False, 

241 update_output_chain=False, 

242 dry_run=False, 

243 dataset_type=(), 

244 ) 

245 self.assertEqual(n2, 9) 

246 # Check that the expected outputs exist in the new run. 

247 with helper.butler.query() as query: 

248 self.assertEqual(query.datasets("dataset_auto1", collections=new_output_run).count(), 1) 

249 self.assertEqual(query.datasets("dataset_auto2", collections=new_output_run).count(), 1) 

250 

251 def test_empty_qg(self): 

252 """Test that making an empty QG produces the right error messages.""" 

253 with DirectButlerRepo.make_temporary("base.yaml") as (helper, root): 

254 helper.add_task(dimensions=["instrument"]) 

255 helper.add_task(dimensions=["instrument"]) 

256 helper.pipeline_graph.resolve(registry=helper.butler.registry) 

257 helper.butler.registry.registerDatasetType( 

258 helper.pipeline_graph.dataset_types["dataset_auto0"].dataset_type 

259 ) 

260 kwargs = self._make_run_args( 

261 "-b", 

262 root, 

263 "-i", 

264 helper.input_chain, 

265 "-o", 

266 "output", 

267 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

268 ) 

269 # Note that we haven't inserted any datasets into this repo; that's 

270 # how we'll force an empty graph. 

271 with self.assertLogs(level=logging.ERROR) as cm: 

272 qg = script.qgraph(**kwargs) 

273 self.assertRegex( 

274 cm.output[0], ".*Initial data ID query returned no rows, so QuantumGraph will be empty.*" 

275 ) 

276 self.assertRegex(cm.output[0], ".*dataset_auto0.*input_run.*doomed to fail.") 

277 self.assertIsNone(qg) 

278 

279 def test_simple_qg_no_skip_existing_inputs(self): 

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

281 the *input* collection, but no ``--extend-run`` or ``-skip-existing`` 

282 option is present. 

283 """ 

284 with DirectButlerRepo.make_temporary() as (helper, root): 

285 helper.add_task() 

286 helper.add_task() 

287 helper.insert_datasets("dataset_auto0") 

288 kwargs = self._make_run_args( 

289 "-b", 

290 root, 

291 "-i", 

292 helper.input_chain, 

293 "-o", 

294 "output", 

295 "--register-dataset-types", 

296 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

297 ) 

298 qg1 = script.qgraph(**kwargs) 

299 run1 = qg1.header.output_run 

300 self.assertEqual(len(qg1.quanta_by_task["task_auto1"]), 1) 

301 self.assertEqual(len(qg1.quanta_by_task["task_auto2"]), 1) 

302 self.assertEqual(len(qg1), 2) 

303 # Ensure that the output run used in the graph is also used in the 

304 # pipeline execution. It is possible for 'qgraph' and 'run' to 

305 # calculate time-stamped runs across a second boundary. 

306 kwargs["output_run"] = run1 

307 # Execute the graph and check for output existence. 

308 script.run(qg1, **kwargs) 

309 with helper.butler.query() as query: 

310 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

311 self.assertEqual(query.datasets("dataset_auto2", collections=["output"]).count(), 1) 

312 # Make a new QG with the same output collection, but a new RUN 

313 # collection, it should run again, shadowing the previous outputs. 

314 kwargs["output_run"] = None 

315 time.sleep(1) # Make sure we don't get the same RUN timestamp. 

316 qg2 = script.qgraph(**kwargs) 

317 run2 = qg2.header.output_run 

318 self.assertNotEqual(run1, run2) 

319 self.assertEqual(len(qg1.quanta_by_task["task_auto1"]), 1) 

320 self.assertEqual(len(qg1.quanta_by_task["task_auto2"]), 1) 

321 self.assertEqual(len(qg2), 2) 

322 kwargs["output_run"] = run2 

323 script.run(qg2, **kwargs) 

324 with helper.butler.query() as query: 

325 self.assertEqual(query.datasets("dataset_auto1", collections=[run2]).count(), 1) 

326 self.assertEqual(query.datasets("dataset_auto2", collections=[run2]).count(), 1) 

327 

328 def test_simple_qg_skip_existing_inputs(self): 

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

330 the *input* collection, but no ``--extend-run`` or ``-skip-existing`` 

331 option is present. 

332 """ 

333 with DirectButlerRepo.make_temporary() as (helper, root): 

334 helper.add_task() 

335 helper.insert_datasets("dataset_auto0") 

336 kwargs = self._make_run_args( 

337 "-b", 

338 root, 

339 "-i", 

340 helper.input_chain, 

341 "-o", 

342 "output", 

343 "--register-dataset-types", 

344 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

345 ) 

346 qg1 = script.qgraph(**kwargs) 

347 run1 = qg1.header.output_run 

348 self.assertEqual(len(qg1.quanta_by_task["task_auto1"]), 1) 

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

350 # Ensure that the output run used in the graph is also used in the 

351 # pipeline execution. It is possible for 'qgraph' and 'run' to 

352 # calculate time-stamped runs across a second boundary. 

353 kwargs["output_run"] = run1 

354 # Execute the graph and check for output existence. 

355 script.run(qg1, **kwargs) 

356 with helper.butler.query() as query: 

357 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

358 # Make a new QG with the same output collection, but a new RUN 

359 # collection, with --skip-existing-in, and one more task. The 

360 # first task should be skipped and the second should be run. 

361 helper.add_task() 

362 kwargs = self._make_run_args( 

363 "-b", 

364 root, 

365 "-i", 

366 helper.input_chain, 

367 "-o", 

368 "output", 

369 "--register-dataset-types", 

370 "--skip-existing-in", 

371 "output", 

372 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

373 ) 

374 time.sleep(1) # Make sure we don't get the same RUN timestamp. 

375 qg2 = script.qgraph(**kwargs) 

376 run2 = qg2.header.output_run 

377 self.assertNotEqual(run1, run2) 

378 self.assertEqual(len(qg2.quanta_by_task["task_auto1"]), 0) 

379 self.assertEqual(len(qg2.quanta_by_task["task_auto2"]), 1) 

380 self.assertEqual(len(qg2), 1) 

381 kwargs["output_run"] = run2 

382 script.run(qg2, **kwargs) 

383 with helper.butler.query() as query: 

384 self.assertEqual(query.datasets("dataset_auto1", collections=[run2]).count(), 0) 

385 self.assertEqual(query.datasets("dataset_auto2", collections=[run2]).count(), 1) 

386 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

387 self.assertEqual(query.datasets("dataset_auto2", collections=["output"]).count(), 1) 

388 

389 def test_simple_qg_extend_run(self): 

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

391 the output RUN collection, and `--extend-run` is used to skip it. 

392 """ 

393 with DirectButlerRepo.make_temporary() as (helper, root): 

394 helper.add_task() 

395 helper.insert_datasets("dataset_auto0") 

396 kwargs = self._make_run_args( 

397 "-b", 

398 root, 

399 "-i", 

400 helper.input_chain, 

401 "-o", 

402 "output", 

403 "--register-dataset-types", 

404 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

405 ) 

406 qg1 = script.qgraph(**kwargs) 

407 run1 = qg1.header.output_run 

408 self.assertEqual(len(qg1.quanta_by_task["task_auto1"]), 1) 

409 self.assertEqual(len(qg1), 1) 

410 # Ensure that the output run used in the graph is also used in the 

411 # pipeline execution. It is possible for 'qgraph' and 'run' to 

412 # calculate time-stamped runs across a second boundary. 

413 kwargs["output_run"] = run1 

414 # Execute the graph and check for output existence. 

415 script.run(qg1, **kwargs) 

416 with helper.butler.query() as query: 

417 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

418 # Make a new QG with the same output collection, but a new RUN 

419 # collection, with --extend-run, and one more task. The first task 

420 # should be skipped and the second should be run. 

421 helper.add_task() 

422 kwargs = self._make_run_args( 

423 "-b", 

424 root, 

425 "-i", 

426 helper.input_chain, 

427 "-o", 

428 "output", 

429 "--register-dataset-types", 

430 "--extend-run", 

431 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

432 ) 

433 qg2 = script.qgraph(**kwargs) 

434 run2 = qg2.header.output_run 

435 self.assertEqual(run1, run2) 

436 self.assertEqual(len(qg2.quanta_by_task["task_auto1"]), 0) 

437 self.assertEqual(len(qg2.quanta_by_task["task_auto2"]), 1) 

438 self.assertEqual(len(qg2), 1) 

439 kwargs["output_run"] = run2 

440 script.run(qg2, **kwargs) 

441 with helper.butler.query() as query: 

442 self.assertEqual(query.datasets("dataset_auto1", collections=[run2]).count(), 1) 

443 self.assertEqual(query.datasets("dataset_auto2", collections=[run2]).count(), 1) 

444 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

445 self.assertEqual(query.datasets("dataset_auto2", collections=["output"]).count(), 1) 

446 

447 def test_simple_qg_clobber(self): 

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

449 the output RUN collection, and `--extend-run --clobber-outputs` is used 

450 to skip it. 

451 """ 

452 with DirectButlerRepo.make_temporary() as (helper, root): 

453 helper.add_task() 

454 helper.insert_datasets("dataset_auto0") 

455 kwargs = self._make_run_args( 

456 "-b", 

457 root, 

458 "-i", 

459 helper.input_chain, 

460 "-o", 

461 "output", 

462 "--register-dataset-types", 

463 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

464 ) 

465 qg1 = script.qgraph(**kwargs) 

466 run1 = qg1.header.output_run 

467 self.assertEqual(len(qg1.quanta_by_task), 1) 

468 self.assertEqual(len(qg1), 1) 

469 # Ensure that the output run used in the graph is also used in the 

470 # pipeline execution. It is possible for 'qgraph' and 'run' to 

471 # calculate time-stamped runs across a second boundary. 

472 kwargs["output_run"] = run1 

473 # Execute the graph and check for output existence. 

474 script.run(qg1, **kwargs) 

475 with helper.butler.query() as query: 

476 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

477 # Delete the metadata output so we don't take the skip-existing 

478 # logic path instead of the clobbering one. 

479 helper.butler.pruneDatasets( 

480 helper.butler.query_datasets("task_auto1_metadata", collections=run1), 

481 purge=True, 

482 unstore=True, 

483 disassociate=True, 

484 ) 

485 # Make a new QG with the same output collection, but a new RUN 

486 # collection, with --clobber-outputs, and one more task. Both 

487 # tasks should be run. 

488 helper.add_task() 

489 kwargs = self._make_run_args( 

490 "-b", 

491 root, 

492 "-i", 

493 helper.input_chain, 

494 "-o", 

495 "output", 

496 "--register-dataset-types", 

497 "--extend-run", 

498 "--clobber-outputs", 

499 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

500 ) 

501 qg2 = script.qgraph(**kwargs) 

502 run2 = qg2.header.output_run 

503 self.assertEqual(run1, run2) 

504 self.assertEqual(len(qg2.quanta_by_task), 2) 

505 self.assertEqual(len(qg2), 2) 

506 kwargs["output_run"] = run2 

507 script.run(qg2, **kwargs) 

508 with helper.butler.query() as query: 

509 self.assertEqual(query.datasets("dataset_auto1", collections=[run2]).count(), 1) 

510 self.assertEqual(query.datasets("dataset_auto2", collections=[run2]).count(), 1) 

511 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

512 self.assertEqual(query.datasets("dataset_auto2", collections=["output"]).count(), 1) 

513 

514 def test_simple_qg_replace_run(self): 

515 """Test repeated execution of a trivial quantum graph with 

516 --replace-run. 

517 """ 

518 with DirectButlerRepo.make_temporary() as (helper, root): 

519 helper.add_task() 

520 helper.insert_datasets("dataset_auto0") 

521 kwargs = self._make_run_args( 

522 "-b", 

523 root, 

524 "-i", 

525 helper.input_chain, 

526 "-o", 

527 "output", 

528 "--register-dataset-types", 

529 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

530 ) 

531 qg1 = script.qgraph(**kwargs) 

532 run1 = qg1.header.output_run 

533 self.assertEqual(len(qg1.quanta_by_task["task_auto1"]), 1) 

534 self.assertEqual(len(qg1), 1) 

535 # Ensure that the output run used in the graph is also used in the 

536 # pipeline execution. It is possible for 'qgraph' and 'run' to 

537 # calculate time-stamped runs across a second boundary. 

538 kwargs["output_run"] = run1 

539 # Execute the graph and check for output existence. 

540 script.run(qg1, **kwargs) 

541 with helper.butler.query() as query: 

542 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

543 # Delete the metadata output so we don't take the skip-existing 

544 # logic path instead of the clobbering one. 

545 helper.butler.pruneDatasets( 

546 helper.butler.query_datasets("task_auto1_metadata", collections=run1), 

547 purge=True, 

548 unstore=True, 

549 disassociate=True, 

550 ) 

551 # Make a new QG with the same output collection, but a new RUN 

552 # collection, with --clobber-outputs, and one more task. Both 

553 # tasks should be run. 

554 time.sleep(1) # Make sure we don't get the same RUN timestamp. 

555 kwargs = self._make_run_args( 

556 "-b", 

557 root, 

558 "-i", 

559 helper.input_chain, 

560 "-o", 

561 "output", 

562 "--replace-run", 

563 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

564 ) 

565 qg2 = script.qgraph(**kwargs) 

566 run2 = qg2.header.output_run 

567 self.assertNotEqual(run1, run2) 

568 self.assertEqual(len(qg2.quanta_by_task["task_auto1"]), 1) 

569 self.assertEqual(len(qg2), 1) 

570 kwargs["output_run"] = run2 

571 script.run(qg2, **kwargs) 

572 self.assertNotIn(run1, helper.butler.collections.get_info("output").children) 

573 self.assertIn(run2, helper.butler.collections.get_info("output").children) 

574 with helper.butler.query() as query: 

575 self.assertEqual(query.datasets("dataset_auto1", collections=[run2]).count(), 1) 

576 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

577 # Repeat once again with --prune-replaced as well. 

578 time.sleep(1) # Make sure we don't get the same RUN timestamp. 

579 kwargs = self._make_run_args( 

580 "-b", 

581 root, 

582 "-i", 

583 helper.input_chain, 

584 "-o", 

585 "output", 

586 "--replace-run", 

587 "--prune-replaced", 

588 "purge", 

589 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

590 ) 

591 qg3 = script.qgraph(**kwargs) 

592 run3 = qg3.header.output_run 

593 self.assertNotEqual(run2, run3) 

594 self.assertEqual(len(qg3.quanta_by_task["task_auto1"]), 1) 

595 self.assertEqual(len(qg3), 1) 

596 kwargs["output_run"] = run3 

597 script.run(qg3, **kwargs) 

598 self.assertNotIn(run2, helper.butler.collections.get_info("output").children) 

599 with self.assertRaises(MissingCollectionError): 

600 helper.butler.collections.get_info(run2) 

601 self.assertIn(run3, helper.butler.collections.get_info("output").children) 

602 with helper.butler.query() as query: 

603 self.assertEqual(query.datasets("dataset_auto1", collections=[run3]).count(), 1) 

604 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 1) 

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

606 # started with is an error, and the kind that should not modify the 

607 # data repo. 

608 kwargs = self._make_run_args( 

609 "-b", 

610 root, 

611 "-i", 

612 run1, 

613 "-o", 

614 "output", 

615 "--replace-run", 

616 "--prune-replaced", 

617 "purge", 

618 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

619 ) 

620 with self.assertRaises(ValueError): 

621 script.qgraph(**kwargs) 

622 

623 def test_qg_partial_failure(self): 

624 """Test execution of a quantum graph where one quantum fails but others 

625 should continue. 

626 """ 

627 with DirectButlerRepo.make_temporary("base.yaml") as (helper, root): 

628 helper.add_task( 

629 dimensions=["detector"], config=DynamicTestPipelineTaskConfig(fail_condition="detector=3") 

630 ) 

631 helper.insert_datasets("dataset_auto0") 

632 kwargs = self._make_run_args( 

633 "-b", 

634 root, 

635 "-i", 

636 helper.input_chain, 

637 "-o", 

638 "output", 

639 "--register-dataset-types", 

640 pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), 

641 ) 

642 qg = script.qgraph(**kwargs) 

643 self.assertEqual(len(qg.quanta_by_task), 1) 

644 self.assertEqual(len(qg), 4) 

645 kwargs["output_run"] = qg.header.output_run 

646 # Execute the graph and check for output existence. 

647 with self.assertRaises(MPGraphExecutorError): 

648 script.run(qg, **kwargs) 

649 with helper.butler.query() as query: 

650 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 3) 

651 

652 

653class CoverageTestCase(unittest.TestCase): 

654 """Test the coverage context manager.""" 

655 

656 @unittest.mock.patch.dict("sys.modules", coverage=unittest.mock.MagicMock()) 

657 def testWithCoverage(self): 

658 """Test that the coverage context manager runs when invoked.""" 

659 with coverage_context({"coverage": True}): 

660 self.assertTrue(True) 

661 

662 @unittest.mock.patch("lsst.ctrl.mpexec.cli.cmd.commands.import_module", side_effect=ModuleNotFoundError()) 

663 def testWithMissingCoverage(self, mock_import): # numpydoc ignore=PR01 

664 """Test that the coverage context manager complains when coverage is 

665 not available. 

666 """ 

667 with self.assertRaises(click.exceptions.ClickException): 

668 with coverage_context({"coverage": True}): 

669 pass 

670 

671 

672if __name__ == "__main__": 

673 lsst.utils.tests.init() 

674 unittest.main()