Coverage for tests/test_run.py: 99%

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 contextlib 

29import logging 

30import os 

31import tempfile 

32import time 

33import unittest 

34import unittest.mock 

35 

36import click.testing 

37 

38import lsst.utils.tests 

39from lsst.ctrl.mpexec import PipelineGraphFactory 

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

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

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

43from lsst.ctrl.mpexec.showInfo import ShowInfo 

44from lsst.daf.butler import CollectionType, MissingCollectionError 

45from lsst.daf.butler.cli.utils import LogCliRunner 

46from lsst.pipe.base.mp_graph_executor import MPGraphExecutorError 

47from lsst.pipe.base.script import transfer_from_graph 

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

49 

50 

51# Copied from test_build.py 

52@contextlib.contextmanager 

53def make_tmp_file(contents=None, suffix=None): 

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

55 

56 Temporary file is deleted on exiting context. 

57 

58 Parameters 

59 ---------- 

60 contents : `bytes` or `None`, optional 

61 Data to write into a file. 

62 suffix : `str` or `None`, optional 

63 Suffix to use for temporary file. 

64 

65 Yields 

66 ------ 

67 `str` 

68 Name of the temporary file. 

69 """ 

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

71 if contents: 71 ↛ 73line 71 didn't jump to line 73 because the condition on line 71 was always true

72 os.write(fd, contents) 

73 os.close(fd) 

74 yield tmpname 

75 with contextlib.suppress(OSError): 

76 os.remove(tmpname) 

77 

78 

79class RunTestCase(unittest.TestCase): 

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

81 

82 @staticmethod 

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

84 mock = unittest.mock.Mock() 

85 

86 @click.command(cls=PipetaskCommand) 

87 @opt.run_options() 

88 @opt.config_search_path_option() 

89 @opt.no_existing_outputs_option() 

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

91 kwargs = collect_pipeline_actions(ctx, **kwargs) 

92 mock(**kwargs) 

93 

94 # At least one tests requires that we enable INFO logging so use 

95 # the specialist runner. 

96 runner = LogCliRunner() 

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

98 if result.exit_code != 0: 98 ↛ 99line 98 didn't jump to line 99 because the condition on line 98 was never true

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

100 mock.assert_called_once() 

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

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

103 result.update(kwargs) 

104 return result 

105 

106 def test_missing_options(self): 

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

108 fails. 

109 """ 

110 

111 @click.command() 

112 @opt.pipeline_build_options() 

113 def cli(**kwargs): 

114 script.run(**kwargs) 

115 

116 runner = click.testing.CliRunner() 

117 result = runner.invoke(cli) 

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

119 # than are defined by pipeline_build_options. 

120 self.assertNotEqual(result.exit_code, 0) 

121 

122 def test_simple_qg(self): 

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

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

125 helper.add_task() 

126 helper.add_task() 

127 helper.insert_datasets("dataset_auto0") 

128 kwargs = self._make_run_args( 

129 "-b", 

130 root, 

131 "-i", 

132 helper.input_chain, 

133 "-o", 

134 "output", 

135 "--register-dataset-types", 

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

137 ) 

138 qg = script.qgraph(**kwargs) 

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

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

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

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

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

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

145 # Execute the graph and check for output existence. 

146 script.run(qg, **kwargs) 

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

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

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

150 # Test that we've disabled implicit threading 

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

152 

153 def test_simple_qg_rebase(self): 

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

155 force redefinition of the output collection. 

156 """ 

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

158 helper.add_task() 

159 helper.add_task() 

160 helper.insert_datasets("dataset_auto0") 

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

162 # override it later. 

163 kwargs = self._make_run_args( 

164 "-b", 

165 root, 

166 "-i", 

167 helper.input_chain, 

168 "-o", 

169 "output", 

170 "--register-dataset-types", 

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

172 ) 

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

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

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

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

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

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

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

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

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

182 helper.butler.collections.redefine_chain( 

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

184 ) 

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

186 # are an error. 

187 with self.assertRaises(ValueError): 

188 script.qgraph(**kwargs) 

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

190 kwargs["rebase"] = True 

191 qg = script.qgraph(**kwargs) 

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

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

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

195 # pipeline execution. 

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

197 # Execute the graph and check for output existence. 

198 script.run(qg, **kwargs) 

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

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

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

202 

203 def test_simple_qgraph_qbb(self): 

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

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

206 helper.add_task() 

207 helper.add_task() 

208 helper.insert_datasets("dataset_auto0") 

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

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

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

212 kwargs = self._make_run_args( 

213 "-b", 

214 root, 

215 "-i", 

216 helper.input_chain, 

217 "-o", 

218 "output", 

219 "--register-dataset-types", 

220 "--qgraph-datastore-records", 

221 "--save-qgraph", 

222 qg_file_1, 

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

224 ) 

225 qg = script.qgraph(**kwargs) 

226 output_run = qg.header.output_run 

227 output = qg.header.output 

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

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

230 # Execute with QBB. 

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

232 script.pre_exec_init_qbb(**kwargs) 

233 script.run_qbb(**kwargs) 

234 # Transfer the datasets to the butler. 

235 n1 = transfer_from_graph( 

236 qg_file_1, 

237 root, 

238 register_dataset_types=True, 

239 transfer_dimensions=False, 

240 update_output_chain=True, 

241 dry_run=False, 

242 dataset_type=(), 

243 ) 

244 self.assertEqual(n1, 9) 

245 # Check that the expected outputs exist. 

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

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

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

249 # Check that some metadata keys were written. 

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

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

252 f"{some_task_label}_metadata", 

253 limit=1, 

254 collections=output, 

255 ) 

256 some_metadata = helper.butler.get(some_metadata_ref) 

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

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

259 

260 # Update the output run and try again. 

261 new_output_run = output_run + "_new" 

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

263 script.update_graph_run(qg_file_1, new_output_run, qg_file_2) 

264 kwargs.update(qgraph=qg_file_2) 

265 # Execute with QBB again. 

266 script.pre_exec_init_qbb(**kwargs) 

267 script.run_qbb(**kwargs) 

268 # Transfer the datasets to the butler. 

269 n2 = transfer_from_graph( 

270 qg_file_2, 

271 root, 

272 register_dataset_types=True, 

273 transfer_dimensions=False, 

274 update_output_chain=False, 

275 dry_run=False, 

276 dataset_type=(), 

277 ) 

278 self.assertEqual(n2, 9) 

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

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

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

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

283 

284 def test_empty_qg(self): 

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

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

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

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

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

290 helper.butler.registry.registerDatasetType( 

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

292 ) 

293 kwargs = self._make_run_args( 

294 "-b", 

295 root, 

296 "-i", 

297 helper.input_chain, 

298 "-o", 

299 "output", 

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

301 ) 

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

303 # how we'll force an empty graph. 

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

305 qg = script.qgraph(**kwargs) 

306 self.assertRegex( 

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

308 ) 

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

310 self.assertIsNone(qg) 

311 

312 def test_simple_qg_no_skip_existing_inputs(self): 

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

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

315 option is present. 

316 """ 

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

318 helper.add_task() 

319 helper.add_task() 

320 helper.insert_datasets("dataset_auto0") 

321 kwargs = self._make_run_args( 

322 "-b", 

323 root, 

324 "-i", 

325 helper.input_chain, 

326 "-o", 

327 "output", 

328 "--register-dataset-types", 

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

330 ) 

331 qg1 = script.qgraph(**kwargs) 

332 run1 = qg1.header.output_run 

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

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

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

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

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

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

339 kwargs["output_run"] = run1 

340 # Execute the graph and check for output existence. 

341 script.run(qg1, **kwargs) 

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

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

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

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

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

347 kwargs["output_run"] = None 

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

349 qg2 = script.qgraph(**kwargs) 

350 run2 = qg2.header.output_run 

351 self.assertNotEqual(run1, run2) 

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

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

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

355 kwargs["output_run"] = run2 

356 script.run(qg2, **kwargs) 

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

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

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

360 

361 def test_simple_qg_skip_existing_inputs(self): 

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

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

364 option is present. 

365 """ 

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

367 helper.add_task() 

368 helper.insert_datasets("dataset_auto0") 

369 kwargs = self._make_run_args( 

370 "-b", 

371 root, 

372 "-i", 

373 helper.input_chain, 

374 "-o", 

375 "output", 

376 "--register-dataset-types", 

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

378 ) 

379 qg1 = script.qgraph(**kwargs) 

380 run1 = qg1.header.output_run 

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

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

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

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

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

386 kwargs["output_run"] = run1 

387 # Execute the graph and check for output existence. 

388 script.run(qg1, **kwargs) 

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

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

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

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

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

394 helper.add_task() 

395 kwargs = self._make_run_args( 

396 "-b", 

397 root, 

398 "-i", 

399 helper.input_chain, 

400 "-o", 

401 "output", 

402 "--register-dataset-types", 

403 "--skip-existing-in", 

404 "output", 

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

406 ) 

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

408 qg2 = script.qgraph(**kwargs) 

409 run2 = qg2.header.output_run 

410 self.assertNotEqual(run1, run2) 

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

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

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

414 kwargs["output_run"] = run2 

415 script.run(qg2, **kwargs) 

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

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

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

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

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

421 

422 def test_simple_qg_extend_run(self): 

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

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

425 """ 

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

427 helper.add_task() 

428 helper.insert_datasets("dataset_auto0") 

429 kwargs = self._make_run_args( 

430 "-b", 

431 root, 

432 "-i", 

433 helper.input_chain, 

434 "-o", 

435 "output", 

436 "--register-dataset-types", 

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

438 ) 

439 qg1 = script.qgraph(**kwargs) 

440 run1 = qg1.header.output_run 

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

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

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

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

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

446 kwargs["output_run"] = run1 

447 # Execute the graph and check for output existence. 

448 script.run(qg1, **kwargs) 

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

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

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

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

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

454 helper.add_task() 

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 "--extend-run", 

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

465 ) 

466 qg2 = script.qgraph(**kwargs) 

467 run2 = qg2.header.output_run 

468 self.assertEqual(run1, run2) 

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

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

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

472 kwargs["output_run"] = run2 

473 script.run(qg2, **kwargs) 

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

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

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

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

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

479 

480 def test_simple_qg_clobber(self): 

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

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

483 to skip it. 

484 """ 

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

486 helper.add_task() 

487 helper.insert_datasets("dataset_auto0") 

488 kwargs = self._make_run_args( 

489 "-b", 

490 root, 

491 "-i", 

492 helper.input_chain, 

493 "-o", 

494 "output", 

495 "--register-dataset-types", 

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

497 ) 

498 qg1 = script.qgraph(**kwargs) 

499 run1 = qg1.header.output_run 

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

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

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

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

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

505 kwargs["output_run"] = run1 

506 # Execute the graph and check for output existence. 

507 script.run(qg1, **kwargs) 

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

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

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

511 # logic path instead of the clobbering one. 

512 helper.butler.pruneDatasets( 

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

514 purge=True, 

515 unstore=True, 

516 disassociate=True, 

517 ) 

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

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

520 # tasks should be run. 

521 helper.add_task() 

522 kwargs = self._make_run_args( 

523 "-b", 

524 root, 

525 "-i", 

526 helper.input_chain, 

527 "-o", 

528 "output", 

529 "--register-dataset-types", 

530 "--extend-run", 

531 "--clobber-outputs", 

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

533 ) 

534 qg2 = script.qgraph(**kwargs) 

535 run2 = qg2.header.output_run 

536 self.assertEqual(run1, run2) 

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

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

539 kwargs["output_run"] = run2 

540 script.run(qg2, **kwargs) 

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

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

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

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

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

546 

547 def test_simple_qg_replace_run(self): 

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

549 --replace-run. 

550 """ 

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

552 helper.add_task() 

553 helper.insert_datasets("dataset_auto0") 

554 kwargs = self._make_run_args( 

555 "-b", 

556 root, 

557 "-i", 

558 helper.input_chain, 

559 "-o", 

560 "output", 

561 "--register-dataset-types", 

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

563 ) 

564 qg1 = script.qgraph(**kwargs) 

565 run1 = qg1.header.output_run 

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

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

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

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

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

571 kwargs["output_run"] = run1 

572 # Execute the graph and check for output existence. 

573 script.run(qg1, **kwargs) 

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

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

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

577 # logic path instead of the clobbering one. 

578 helper.butler.pruneDatasets( 

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

580 purge=True, 

581 unstore=True, 

582 disassociate=True, 

583 ) 

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

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

586 # tasks should be run. 

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

588 kwargs = self._make_run_args( 

589 "-b", 

590 root, 

591 "-i", 

592 helper.input_chain, 

593 "-o", 

594 "output", 

595 "--replace-run", 

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

597 ) 

598 qg2 = script.qgraph(**kwargs) 

599 run2 = qg2.header.output_run 

600 self.assertNotEqual(run1, run2) 

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

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

603 kwargs["output_run"] = run2 

604 script.run(qg2, **kwargs) 

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

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

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

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

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

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

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

612 kwargs = self._make_run_args( 

613 "-b", 

614 root, 

615 "-i", 

616 helper.input_chain, 

617 "-o", 

618 "output", 

619 "--replace-run", 

620 "--prune-replaced", 

621 "purge", 

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

623 ) 

624 qg3 = script.qgraph(**kwargs) 

625 run3 = qg3.header.output_run 

626 self.assertNotEqual(run2, run3) 

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

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

629 kwargs["output_run"] = run3 

630 script.run(qg3, **kwargs) 

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

632 with self.assertRaises(MissingCollectionError): 

633 helper.butler.collections.get_info(run2) 

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

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

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

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

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

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

640 # data repo. 

641 kwargs = self._make_run_args( 

642 "-b", 

643 root, 

644 "-i", 

645 run1, 

646 "-o", 

647 "output", 

648 "--replace-run", 

649 "--prune-replaced", 

650 "purge", 

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

652 ) 

653 with self.assertRaises(ValueError): 

654 script.qgraph(**kwargs) 

655 

656 def test_qg_partial_failure(self): 

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

658 should continue. 

659 """ 

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

661 helper.add_task( 

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

663 ) 

664 helper.insert_datasets("dataset_auto0") 

665 kwargs = self._make_run_args( 

666 "-b", 

667 root, 

668 "-i", 

669 helper.input_chain, 

670 "-o", 

671 "output", 

672 "--register-dataset-types", 

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

674 ) 

675 qg = script.qgraph(**kwargs) 

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

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

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

679 # Execute the graph and check for output existence. 

680 with self.assertRaises(MPGraphExecutorError): 

681 script.run(qg, **kwargs) 

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

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

684 

685 def test_retained_dataset_types_option(self): 

686 """--retained-dataset-types accepts a file path.""" 

687 with make_tmp_file(b"- '*_metadata'\n- '*_log'\n", suffix=".yaml") as retained_path: 

688 kwargs = self._make_run_args( 

689 "-b", 

690 "fake_repo", 

691 "-i", 

692 "fake_input", 

693 "-o", 

694 "fake_output", 

695 "--retained-dataset-types", 

696 retained_path, 

697 ) 

698 self.assertEqual(kwargs["retained_dataset_types"], retained_path) 

699 

700 def test_simple_qg_retained_forces_rerun(self): 

701 """With --retained-dataset-types listing only metadata types, when 

702 task_auto2 has no metadata and must run, task_auto1 is forced to rerun 

703 because dataset_auto1 is not retained. 

704 """ 

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

706 helper.add_task() 

707 helper.add_task() 

708 helper.insert_datasets("dataset_auto0") 

709 kwargs = self._make_run_args( 

710 "-b", 

711 root, 

712 "-i", 

713 helper.input_chain, 

714 "-o", 

715 "output", 

716 "--register-dataset-types", 

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

718 ) 

719 qg1 = script.qgraph(**kwargs) 

720 run1 = qg1.header.output_run 

721 kwargs["output_run"] = run1 

722 script.run(qg1, **kwargs) 

723 # Simulate: task_auto1 ran (metadata kept) but intermediate output 

724 # not retained; task_auto2 failed (no metadata, no output). 

725 helper.butler.pruneDatasets( 

726 helper.butler.query_datasets("dataset_auto1", collections=run1), 

727 purge=True, 

728 unstore=True, 

729 disassociate=True, 

730 ) 

731 helper.butler.pruneDatasets( 

732 helper.butler.query_datasets("task_auto2_metadata", collections=run1), 

733 purge=True, 

734 unstore=True, 

735 disassociate=True, 

736 ) 

737 helper.butler.pruneDatasets( 

738 helper.butler.query_datasets("dataset_auto2", collections=run1), 

739 purge=True, 

740 unstore=True, 

741 disassociate=True,