Coverage for python/lsst/pipe/base/mp_graph_executor.py: 84%

411 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-07-23 08:59 +0000

1# This file is part of pipe_base. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://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 <http://www.gnu.org/licenses/>. 

27 

28from __future__ import annotations 

29 

30__all__ = ["MPGraphExecutor", "MPGraphExecutorError", "MPTimeoutError"] 

31 

32import enum 

33import importlib 

34import logging 

35import multiprocessing 

36import pickle 

37import signal 

38import sys 

39import threading 

40import time 

41import uuid 

42from contextlib import ExitStack 

43from typing import Literal, cast 

44 

45import networkx 

46 

47from lsst.daf.butler import DataCoordinate, Quantum 

48from lsst.daf.butler.cli.cliLog import CliLog 

49from lsst.daf.butler.logging import ButlerLogRecords 

50from lsst.utils.threads import disable_implicit_threading 

51 

52from ._status import InvalidQuantumError, RepeatableQuantumError 

53from ._task_metadata import TaskMetadata 

54from .execution_graph_fixup import ExecutionGraphFixup 

55from .graph import QuantumGraph 

56from .graph_walker import GraphWalker 

57from .log_on_close import LogOnClose 

58from .pipeline_graph import TaskNode 

59from .quantum_graph import PredictedQuantumGraph, PredictedQuantumInfo, ProvenanceQuantumGraphWriter 

60from .quantum_graph_executor import QuantumExecutor, QuantumGraphExecutor 

61from .quantum_reports import ExecutionStatus, QuantumReport, Report 

62 

63_LOG = logging.getLogger(__name__) 

64 

65 

66class JobState(enum.Enum): 

67 """Possible state for an executing task.""" 

68 

69 PENDING = enum.auto() 

70 """The job has not started yet.""" 

71 

72 RUNNING = enum.auto() 

73 """The job is currently executing.""" 

74 

75 FINISHED = enum.auto() 

76 """The job finished successfully.""" 

77 

78 FAILED = enum.auto() 

79 """The job execution failed (process returned non-zero status).""" 

80 

81 TIMED_OUT = enum.auto() 

82 """The job was killed due to too long execution time.""" 

83 

84 FAILED_DEP = enum.auto() 

85 """One of the dependencies of this job failed or timed out.""" 

86 

87 

88class _Job: 

89 """Class representing a job running single task. 

90 

91 Parameters 

92 ---------- 

93 quantum_id : `uuid.UUID` 

94 ID of the quantum this job executes. 

95 quantum : `lsst.daf.butler.Quantum` 

96 Description of the inputs and outputs. 

97 task_node : `.pipeline_graph.TaskNode` 

98 Description of the task and configuration. 

99 """ 

100 

101 def __init__(self, quantum_id: uuid.UUID, quantum: Quantum, task_node: TaskNode): 

102 self.quantum_id = quantum_id 

103 self.quantum = quantum 

104 self.task_node = task_node 

105 self.process: multiprocessing.process.BaseProcess | None = None 

106 self._state = JobState.PENDING 

107 self.started: float = 0.0 

108 self._rcv_conn: multiprocessing.connection.Connection | None = None 

109 self._terminated = False 

110 

111 @property 

112 def state(self) -> JobState: 

113 """Job processing state (JobState).""" 

114 return self._state 

115 

116 @property 

117 def terminated(self) -> bool: 

118 """Return `True` if job was killed by stop() method and negative exit 

119 code is returned from child process (`bool`). 

120 """ 

121 if self._terminated: 

122 assert self.process is not None, "Process must be started" 

123 if self.process.exitcode is not None: 123 ↛ 125line 123 didn't jump to line 125 because the condition on line 123 was always true

124 return self.process.exitcode < 0 

125 return False 

126 

127 def start( 

128 self, 

129 quantumExecutor: QuantumExecutor, 

130 startMethod: Literal["spawn"] | Literal["forkserver"], 

131 fail_fast: bool, 

132 ) -> None: 

133 """Start process which runs the task. 

134 

135 Parameters 

136 ---------- 

137 quantumExecutor : `QuantumExecutor` 

138 Executor for single quantum. 

139 startMethod : `str`, optional 

140 Start method from `multiprocessing` module. 

141 fail_fast : `bool`, optional 

142 If `True` then kill subprocess on RepeatableQuantumError. 

143 """ 

144 # Unpickling of quantum has to happen after butler/executor, also we 

145 # want to setup logging before unpickling anything that can generate 

146 # messages, this is why things are pickled manually here. 

147 qe_pickle = pickle.dumps(quantumExecutor) 

148 task_node_pickle = pickle.dumps(self.task_node) 

149 quantum_pickle = pickle.dumps(self.quantum) 

150 self._rcv_conn, snd_conn = multiprocessing.Pipe(False) 

151 logConfigState = CliLog.configState 

152 

153 mp_ctx = multiprocessing.get_context(startMethod) 

154 self.process = mp_ctx.Process( # type: ignore[attr-defined] 

155 target=_Job._executeJob, 

156 args=( 

157 qe_pickle, 

158 task_node_pickle, 

159 quantum_pickle, 

160 self.quantum_id, 

161 logConfigState, 

162 snd_conn, 

163 fail_fast, 

164 ), 

165 name=f"task-{self.quantum.dataId}", 

166 ) 

167 # mypy is getting confused by multiprocessing. 

168 assert self.process is not None 

169 self.process.start() 

170 self.started = time.time() 

171 self._state = JobState.RUNNING 

172 

173 @staticmethod 

174 def _executeJob( 

175 quantumExecutor_pickle: bytes, 

176 task_node_pickle: bytes, 

177 quantum_pickle: bytes, 

178 quantum_id: uuid.UUID, 

179 logConfigState: list, 

180 snd_conn: multiprocessing.connection.Connection, 

181 fail_fast: bool, 

182 ) -> None: 

183 """Execute a job with arguments. 

184 

185 Parameters 

186 ---------- 

187 quantumExecutor_pickle : `bytes` 

188 Executor for single quantum, pickled. 

189 task_node_pickle : `bytes` 

190 Task definition structure, pickled. 

191 quantum_pickle : `bytes` 

192 Quantum for this task execution in pickled form. 

193 quantum_id : `uuid.UUID` 

194 Unique ID for the quantum. 

195 logConfigState : `list` 

196 Logging state from parent process. 

197 snd_conn : `multiprocessing.Connection` 

198 Connection to send job report to parent process. 

199 fail_fast : `bool` 

200 If `True` then kill subprocess on RepeatableQuantumError. 

201 """ 

202 # This process was started with the spawn method, so runtime thread 

203 # limits applied in the parent process do not carry over and implicit 

204 # threading has to be disabled again. Inherited environment variables 

205 # only cover libraries that have not yet created their thread pools. 

206 disable_implicit_threading() 

207 

208 # This terrible hack is a workaround for Python threading bug: 

209 # https://github.com/python/cpython/issues/102512. Should be removed 

210 # when fix for that bug is deployed. Inspired by 

211 # https://github.com/QubesOS/qubes-core-admin-client/pull/236/files. 

212 thread = threading.current_thread() 

213 if isinstance(thread, threading._DummyThread): 213 ↛ 214line 213 didn't jump to line 214 because the condition on line 213 was never true

214 if getattr(thread, "_tstate_lock", "") is None: 

215 thread._set_tstate_lock() # type: ignore[attr-defined] 

216 

217 if logConfigState and not CliLog.configState: 217 ↛ 220line 217 didn't jump to line 220 because the condition on line 217 was never true

218 # means that we are in a new spawned Python process and we have to 

219 # re-initialize logging 

220 CliLog.replayConfigState(logConfigState) 

221 

222 quantumExecutor: QuantumExecutor = pickle.loads(quantumExecutor_pickle) 

223 task_node: TaskNode = pickle.loads(task_node_pickle) 

224 quantum = pickle.loads(quantum_pickle) 

225 report: QuantumReport | None = None 

226 # Catch a few known failure modes and stop the process immediately, 

227 # with exception-specific exit code. 

228 try: 

229 _, report = quantumExecutor.execute(task_node, quantum, quantum_id=quantum_id) 

230 except RepeatableQuantumError as exc: 

231 report = QuantumReport.from_exception( 

232 quantumId=quantum_id, 

233 exception=exc, 

234 dataId=quantum.dataId, 

235 taskLabel=task_node.label, 

236 exitCode=exc.EXIT_CODE if fail_fast else None, 

237 ) 

238 if fail_fast: 

239 _LOG.warning("Caught repeatable quantum error for %s (%s):", task_node.label, quantum.dataId) 

240 _LOG.warning(exc, exc_info=True) 

241 sys.exit(exc.EXIT_CODE) 

242 else: 

243 raise 

244 except InvalidQuantumError as exc: 

245 _LOG.fatal("Invalid quantum error for %s (%s): %s", task_node.label, quantum.dataId) 

246 _LOG.fatal(exc, exc_info=True) 

247 report = QuantumReport.from_exception( 

248 quantumId=quantum_id, 

249 exception=exc, 

250 dataId=quantum.dataId, 

251 taskLabel=task_node.label, 

252 exitCode=exc.EXIT_CODE, 

253 ) 

254 sys.exit(exc.EXIT_CODE) 

255 except Exception as exc: 

256 _LOG.debug("exception from task %s dataId %s: %s", task_node.label, quantum.dataId, exc) 

257 report = QuantumReport.from_exception( 

258 quantumId=quantum_id, 

259 exception=exc, 

260 dataId=quantum.dataId, 

261 taskLabel=task_node.label, 

262 ) 

263 raise 

264 finally: 

265 if report is not None: 265 ↛ exitline 265 didn't return from function '_executeJob' because the condition on line 265 was always true

266 # If sending fails we do not want this new exception to be 

267 # exposed. 

268 try: 

269 _LOG.debug("sending report for task %s dataId %s", task_node.label, quantum.dataId) 

270 snd_conn.send(report) 

271 except Exception: 

272 pass 

273 

274 def stop(self) -> None: 

275 """Stop the process.""" 

276 assert self.process is not None, "Process must be started" 

277 self.process.terminate() 

278 # give it 1 second to finish or KILL 

279 for _ in range(10): 279 ↛ 284line 279 didn't jump to line 284 because the loop on line 279 didn't complete

280 time.sleep(0.1) 

281 if not self.process.is_alive(): 281 ↛ 279line 281 didn't jump to line 279 because the condition on line 281 was always true

282 break 

283 else: 

284 _LOG.debug("Killing process %s", self.process.name) 

285 self.process.kill() 

286 self._terminated = True 

287 

288 def cleanup(self) -> None: 

289 """Release processes resources, has to be called for each finished 

290 process. 

291 """ 

292 if self.process and not self.process.is_alive(): 292 ↛ exitline 292 didn't return from function 'cleanup' because the condition on line 292 was always true

293 self.process.close() 

294 self.process = None 

295 self._rcv_conn = None 

296 

297 def report(self) -> QuantumReport: 

298 """Return task report, should be called after process finishes and 

299 before cleanup(). 

300 """ 

301 assert self.process is not None, "Process must be started" 

302 assert self._rcv_conn is not None, "Process must be started" 

303 try: 

304 report = self._rcv_conn.recv() 

305 report.exitCode = self.process.exitcode 

306 except Exception: 

307 # Likely due to the process killed, but there may be other reasons. 

308 # Exit code should not be None, this is to keep mypy happy. 

309 exitcode = self.process.exitcode if self.process.exitcode is not None else -1 

310 assert self.quantum.dataId is not None, "Quantum DataId cannot be None" 

311 report = QuantumReport.from_exit_code( 

312 quantumId=self.quantum_id, 

313 exitCode=exitcode, 

314 dataId=self.quantum.dataId, 

315 taskLabel=self.task_node.label, 

316 ) 

317 if self.terminated: 

318 # Means it was killed, assume it's due to timeout 

319 report.status = ExecutionStatus.TIMEOUT 

320 return report 

321 

322 def failMessage(self) -> str: 

323 """Return a message describing task failure.""" 

324 assert self.process is not None, "Process must be started" 

325 assert self.process.exitcode is not None, "Process has to finish" 

326 exitcode = self.process.exitcode 

327 if exitcode < 0: 

328 # Negative exit code means it is killed by signal 

329 signum = -exitcode 

330 msg = f"Task {self} failed, killed by signal {signum}" 

331 # Just in case this is some very odd signal, expect ValueError 

332 try: 

333 strsignal = signal.strsignal(signum) 

334 msg = f"{msg} ({strsignal})" 

335 except ValueError: 

336 pass 

337 elif exitcode > 0: 337 ↛ 340line 337 didn't jump to line 340 because the condition on line 337 was always true

338 msg = f"Task {self} failed, exit code={exitcode}" 

339 else: 

340 msg = "" 

341 return msg 

342 

343 def __str__(self) -> str: 

344 return f"<{self.task_node.label} dataId={self.quantum.dataId}>" 

345 

346 

347class _JobList: 

348 """Simple list of _Job instances with few convenience methods. 

349 

350 Parameters 

351 ---------- 

352 xgraph : `networkx.DiGraph` 

353 Directed acyclic graph of quantum IDs. 

354 """ 

355 

356 def __init__(self, xgraph: networkx.DiGraph): 

357 self.jobs = { 

358 quantum_id: _Job( 

359 quantum_id=quantum_id, 

360 quantum=xgraph.nodes[quantum_id]["quantum"], 

361 task_node=xgraph.nodes[quantum_id]["pipeline_node"], 

362 ) 

363 for quantum_id in xgraph 

364 } 

365 self.walker: GraphWalker[uuid.UUID] = GraphWalker(xgraph.copy()) 

366 self.pending = set(next(self.walker, ())) 

367 self.running: set[uuid.UUID] = set() 

368 self.finished: set[uuid.UUID] = set() 

369 self.failed: set[uuid.UUID] = set() 

370 self.timed_out: set[uuid.UUID] = set() 

371 

372 def submit( 

373 self, 

374 quantumExecutor: QuantumExecutor, 

375 startMethod: Literal["spawn"] | Literal["forkserver"], 

376 fail_fast: bool = False, 

377 ) -> _Job: 

378 """Submit a pending job for execution. 

379 

380 Parameters 

381 ---------- 

382 quantumExecutor : `QuantumExecutor` 

383 Executor for single quantum. 

384 startMethod : `str`, optional 

385 Start method from `multiprocessing` module. 

386 fail_fast : `bool`, optional 

387 If `True` then kill subprocess on RepeatableQuantumError. 

388 

389 Returns 

390 ------- 

391 job : `_Job` 

392 The job that was submitted. 

393 """ 

394 quantum_id = self.pending.pop() 

395 job = self.jobs[quantum_id] 

396 job.start(quantumExecutor, startMethod, fail_fast=fail_fast) 

397 self.running.add(job.quantum_id) 

398 return job 

399 

400 def setJobState(self, job: _Job, state: JobState) -> list[_Job]: 

401 """Update job state. 

402 

403 Parameters 

404 ---------- 

405 job : `_Job` 

406 Job to submit. 

407 state : `JobState` 

408 New job state; note that only the FINISHED, FAILED, and TIMED_OUT 

409 states are acceptable. 

410 

411 Returns 

412 ------- 

413 blocked : `list` [ `_Job` ] 

414 Additional jobs that have been marked as failed because this job 

415 was upstream of them and failed or timed out. 

416 """ 

417 allowedStates = (JobState.FINISHED, JobState.FAILED, JobState.TIMED_OUT) 

418 assert state in allowedStates, f"State {state} not allowed here" 

419 

420 # remove job from pending/running lists 

421 if job.state == JobState.PENDING: 421 ↛ 422line 421 didn't jump to line 422 because the condition on line 421 was never true

422 self.pending.remove(job.quantum_id) 

423 elif job.state == JobState.RUNNING: 423 ↛ 426line 423 didn't jump to line 426 because the condition on line 423 was always true

424 self.running.remove(job.quantum_id) 

425 

426 quantum_id = job.quantum_id 

427 # it should not be in any of these, but just in case 

428 self.finished.discard(quantum_id) 

429 self.failed.discard(quantum_id) 

430 self.timed_out.discard(quantum_id) 

431 job._state = state 

432 match job.state: 

433 case JobState.FINISHED: 

434 self.finished.add(quantum_id) 

435 self.walker.finish(quantum_id) 

436 self.pending.update(next(self.walker, ())) 

437 return [] 

438 case JobState.FAILED: 

439 self.failed.add(quantum_id) 

440 case JobState.TIMED_OUT: 440 ↛ 443line 440 didn't jump to line 443 because the pattern on line 440 always matched

441 self.failed.add(quantum_id) 

442 self.timed_out.add(quantum_id) 

443 case _: 

444 raise ValueError(f"Unexpected state value: {state}") 

445 blocked: list[_Job] = [] 

446 for downstream_quantum_id in self.walker.fail(quantum_id): 

447 self.failed.add(downstream_quantum_id) 

448 blocked.append(self.jobs[downstream_quantum_id]) 

449 self.jobs[downstream_quantum_id]._state = JobState.FAILED_DEP 

450 return blocked 

451 

452 def cleanup(self) -> None: 

453 """Do periodic cleanup for jobs that did not finish correctly. 

454 

455 If timed out jobs are killed but take too long to stop then regular 

456 cleanup will not work for them. Here we check all timed out jobs 

457 periodically and do cleanup if they managed to die by this time. 

458 """ 

459 for quantum_id in self.timed_out: 

460 job = self.jobs[quantum_id] 

461 assert job.state == JobState.TIMED_OUT, "Job state should be consistent with the set it's in." 

462 if job.process is not None: 462 ↛ 463line 462 didn't jump to line 463 because the condition on line 462 was never true

463 job.cleanup() 

464 

465 

466class MPGraphExecutorError(Exception): 

467 """Exception class for errors raised by MPGraphExecutor.""" 

468 

469 pass 

470 

471 

472class MPTimeoutError(MPGraphExecutorError): 

473 """Exception raised when task execution times out.""" 

474 

475 pass 

476 

477 

478class MPGraphExecutor(QuantumGraphExecutor): 

479 """Implementation of QuantumGraphExecutor using same-host multiprocess 

480 execution of Quanta. 

481 

482 Parameters 

483 ---------- 

484 num_proc : `int` 

485 Number of processes to use for executing tasks. 

486 timeout : `float` 

487 Time in seconds to wait for tasks to finish. 

488 quantum_executor : `.quantum_graph_executor.QuantumExecutor` 

489 Executor for single quantum. For multiprocess-style execution when 

490 ``num_proc`` is greater than one this instance must support pickle. 

491 start_method : `str`, optional 

492 Start method from `multiprocessing` module, `None` selects the best 

493 one for current platform. 

494 fail_fast : `bool`, optional 

495 If set to ``True`` then stop processing on first error from any task. 

496 pdb : `str`, optional 

497 Debugger to import and use (via the ``post_mortem`` function) in the 

498 event of an exception. 

499 execution_graph_fixup : `.execution_graph_fixup.ExecutionGraphFixup`, \ 

500 optional 

501 Instance used for modification of execution graph. 

502 """ 

503 

504 def __init__( 

505 self, 

506 *, 

507 num_proc: int, 

508 timeout: float, 

509 quantum_executor: QuantumExecutor, 

510 start_method: Literal["spawn"] | Literal["forkserver"] | None = None, 

511 fail_fast: bool = False, 

512 pdb: str | None = None, 

513 execution_graph_fixup: ExecutionGraphFixup | None = None, 

514 ): 

515 self._num_proc = num_proc 

516 self._timeout = timeout 

517 self._quantum_executor = quantum_executor 

518 self._fail_fast = fail_fast 

519 self._pdb = pdb 

520 self._execution_graph_fixup = execution_graph_fixup 

521 self._report: Report | None = None 

522 

523 # We set default start method as spawn for all platforms. 

524 if start_method is None: 

525 start_method = "spawn" 

526 self._start_method = start_method 

527 

528 def execute( 

529 self, graph: QuantumGraph | PredictedQuantumGraph, *, provenance_graph_file: str | None = None 

530 ) -> None: 

531 # Docstring inherited from QuantumGraphExecutor.execute 

532 old_graph: QuantumGraph | None = None 

533 if isinstance(graph, QuantumGraph): 

534 old_graph = graph 

535 new_graph = PredictedQuantumGraph.from_old_quantum_graph(old_graph) 

536 else: 

537 new_graph = graph 

538 xgraph = self._make_xgraph(new_graph, old_graph) 

539 self._report = Report(qgraphSummary=new_graph._make_summary()) 

540 err: MPGraphExecutorError | None = None 

541 with ExitStack() as exit_stack: 

542 provenance_writer: ProvenanceQuantumGraphWriter | None = None 

543 if provenance_graph_file is not None: 

544 if provenance_graph_file is not None and self._num_proc > 1: 

545 raise NotImplementedError( 

546 "Provenance writing is not implemented for multiprocess execution." 

547 ) 

548 provenance_writer = ProvenanceQuantumGraphWriter( 

549 provenance_graph_file, 

550 exit_stack=exit_stack, 

551 log_on_close=LogOnClose(_LOG.log), 

552 predicted=new_graph, 

553 ) 

554 try: 

555 if self._num_proc > 1: 

556 self._execute_quanta_mp(xgraph, self._report) 

557 else: 

558 self._execute_quanta_in_process(xgraph, self._report, provenance_writer) 

559 except MPGraphExecutorError as exc: 

560 self._report.set_exception(exc) 

561 err = exc 

562 # Defer re-raising this exception only to let provenance writes 

563 # finish as the ExitStack closes. The original traceback for 

564 # this exception isn't useful anyway. 

565 except Exception as exc: 

566 self._report.set_exception(exc) 

567 raise 

568 if provenance_writer is not None: 

569 provenance_writer.write_overall_inputs() 

570 provenance_writer.write_packages() 

571 provenance_writer.write_init_outputs(assume_existence=True) 

572 if err is not None: 

573 raise err 

574 

575 def _make_xgraph( 

576 self, new_graph: PredictedQuantumGraph, old_graph: QuantumGraph | None 

577 ) -> networkx.DiGraph: 

578 """Obtain a networkx DAG from a quantum graph, applying any fixup and 

579 adding `lsst.daf.butler.Quantum` and `~.pipeline_graph.TaskNode` 

580 attributes. 

581 

582 Parameters 

583 ---------- 

584 new_graph : `.quantum_graph.PredictedQuantumGraph` 

585 New quantum graph object. 

586 old_graph : `.QuantumGraph` or `None` 

587 Equivalent old quantum graph object. 

588 

589 Returns 

590 ------- 

591 xgraph : `networkx.DiGraph` 

592 NetworkX DAG with quantum IDs as node keys. 

593 

594 Raises 

595 ------ 

596 MPGraphExecutorError 

597 Raised if execution graph cannot be ordered after modification, 

598 i.e. it has dependency cycles. 

599 """ 

600 new_graph.build_execution_quanta() 

601 xgraph = new_graph.quantum_only_xgraph.copy() 

602 if self._execution_graph_fixup: 

603 try: 

604 self._execution_graph_fixup.fixup_graph(xgraph, new_graph.quanta_by_task) 

605 except NotImplementedError: 

606 # Backwards compatibility. 

607 if old_graph is None: 

608 old_graph = new_graph.to_old_quantum_graph() 

609 old_graph = self._execution_graph_fixup.fixupQuanta(old_graph) 

610 # Adding all of the edges from old_graph is overkill, but the 

611 # only option we really have to make sure we add any new ones. 

612 xgraph.update([(a.nodeId, b.nodeId) for a, b in old_graph.graph.edges]) 

613 if networkx.dag.has_cycle(xgraph): 613 ↛ 614line 613 didn't jump to line 614 because the condition on line 613 was never true

614 raise MPGraphExecutorError("Updated execution graph has dependency cycle.") 

615 return xgraph 

616 

617 def _execute_quanta_in_process( 

618 self, xgraph: networkx.DiGraph, report: Report, provenance_writer: ProvenanceQuantumGraphWriter | None 

619 ) -> None: 

620 """Execute all Quanta in current process. 

621 

622 Parameters 

623 ---------- 

624 xgraph : `networkx.DiGraph` 

625 DAG to execute. Should have quantum IDs for nodes and ``quantum`` 

626 (`lsst.daf.butler.Quantum`) and ``pipeline_node`` 

627 (`lsst.pipe.base.pipeline_graph.TaskNode`) attributes in addition 

628 to those provided by 

629 `.quantum_graph.PredictedQuantumGraph.quantum_only_xgraph`. 

630 report : `Report` 

631 Object for reporting execution status. 

632 provenance_writer : `.quantum_graph.ProvenanceQuantumGraphWriter` or \ 

633 `None` 

634 Object for recording provenance. 

635 """ 

636 

637 def tiebreaker_sort_key(quantum_id: uuid.UUID) -> tuple: 

638 node_state = xgraph.nodes[quantum_id] 

639 return (node_state["task_label"],) + node_state["data_id"].required_values 

640 

641 success_count, failed_count, total_count = 0, 0, len(xgraph.nodes) 

642 walker = GraphWalker[uuid.UUID](xgraph.copy()) 

643 for unblocked_quanta in walker: 

644 for quantum_id in sorted(unblocked_quanta, key=tiebreaker_sort_key): 

645 node_state: PredictedQuantumInfo = xgraph.nodes[quantum_id] 

646 data_id = node_state["data_id"] 

647 task_node = node_state["pipeline_node"] 

648 quantum = node_state["quantum"] 

649 

650 _LOG.debug("Executing %s (%s@%s)", quantum_id, task_node.label, data_id) 

651 fail_exit_code: int | None = None 

652 task_metadata: TaskMetadata | None = None 

653 task_logs = ButlerLogRecords([]) 

654 try: 

655 # For some exception types we want to exit immediately with 

656 # exception-specific exit code, but we still want to start 

657 # debugger before exiting if debugging is enabled. 

658 try: 

659 execution_result = self._quantum_executor.execute( 

660 task_node, quantum, quantum_id=quantum_id, log_records=task_logs 

661 ) 

662 if execution_result.report: 662 ↛ 664line 662 didn't jump to line 664 because the condition on line 662 was always true

663 report.quantaReports.append(execution_result.report) 

664 task_metadata = execution_result.task_metadata 

665 success_count += 1 

666 walker.finish(quantum_id) 

667 except RepeatableQuantumError as exc: 

668 if self._fail_fast: 

669 _LOG.warning( 

670 "Caught repeatable quantum error for %s (%s@%s):", 

671 quantum_id, 

672 task_node.label, 

673 data_id, 

674 ) 

675 _LOG.warning(exc, exc_info=True) 

676 fail_exit_code = exc.EXIT_CODE 

677 raise 

678 except InvalidQuantumError as exc: 

679 _LOG.fatal( 

680 "Invalid quantum error for %s (%s@%s):", quantum_id, task_node.label, data_id 

681 ) 

682 _LOG.fatal(exc, exc_info=True) 

683 fail_exit_code = exc.EXIT_CODE 

684 raise 

685 except Exception as exc: 

686 quantum_report = QuantumReport.from_exception( 

687 exception=exc, 

688 dataId=data_id, 

689 taskLabel=task_node.label, 

690 ) 

691 report.quantaReports.append(quantum_report) 

692 

693 if self._pdb and sys.stdin.isatty() and sys.stdout.isatty(): 693 ↛ 694line 693 didn't jump to line 694 because the condition on line 693 was never true

694 _LOG.error( 

695 "%s (%s@%s) failed; dropping into pdb.", 

696 quantum_id, 

697 task_node.label, 

698 data_id, 

699 exc_info=exc, 

700 ) 

701 try: 

702 pdb = importlib.import_module(self._pdb) 

703 except ImportError as imp_exc: 

704 raise MPGraphExecutorError( 

705 f"Unable to import specified debugger module ({self._pdb}): {imp_exc}" 

706 ) from exc 

707 if not hasattr(pdb, "post_mortem"): 

708 raise MPGraphExecutorError( 

709 f"Specified debugger module ({self._pdb}) can't debug with post_mortem", 

710 ) from exc 

711 pdb.post_mortem(exc.__traceback__) 

712 

713 report.status = ExecutionStatus.FAILURE 

714 failed_count += 1 

715 

716 # If exception specified an exit code then just exit with 

717 # that code, otherwise crash if fail-fast option is 

718 # enabled. 

719 if fail_exit_code is not None: 719 ↛ 720line 719 didn't jump to line 720 because the condition on line 719 was never true

720 sys.exit(fail_exit_code) 

721 if self._fail_fast: 721 ↛ 722line 721 didn't jump to line 722 because the condition on line 721 was never true

722 raise MPGraphExecutorError( 

723 f"Quantum {quantum_id} ({task_node.label}@{data_id}) failed." 

724 ) from exc 

725 else: 

726 _LOG.error( 

727 "%s (%s@%s) failed; processing will continue for remaining tasks.", 

728 quantum_id, 

729 task_node.label, 

730 data_id, 

731 exc_info=exc, 

732 ) 

733 

734 for downstream_quantum_id in walker.fail(quantum_id): 

735 downstream_node_state = xgraph.nodes[downstream_quantum_id] 

736 failed_quantum_report = QuantumReport( 

737 status=ExecutionStatus.SKIPPED, 

738 dataId=downstream_node_state["data_id"], 

739 taskLabel=downstream_node_state["task_label"], 

740 ) 

741 report.quantaReports.append(failed_quantum_report) 

742 if provenance_writer is not None: 

743 provenance_writer.write_blocked_quantum_provenance(downstream_quantum_id) 

744 _LOG.error( 

745 "Upstream job failed for task %s (%s@%s), skipping this quantum.", 

746 downstream_quantum_id, 

747 downstream_node_state["task_label"], 

748 downstream_node_state["data_id"], 

749 ) 

750 failed_count += 1 

751 

752 if provenance_writer is not None: 

753 provenance_writer.write_quantum_provenance( 

754 quantum_id, metadata=task_metadata, logs=task_logs 

755 ) 

756 

757 _LOG.info( 

758 "Executed %d quanta successfully, %d failed and %d remain out of total %d quanta.", 

759 success_count, 

760 failed_count, 

761 total_count - success_count - failed_count, 

762 total_count, 

763 ) 

764 

765 # Raise an exception if there were any failures. 

766 if failed_count: 

767 raise MPGraphExecutorError("One or more tasks failed during execution.") 

768 

769 def _execute_quanta_mp(self, xgraph: networkx.DiGraph, report: Report) -> None: 

770 """Execute all Quanta in separate processes. 

771 

772 Parameters 

773 ---------- 

774 xgraph : `networkx.DiGraph` 

775 DAG to execute. Should have quantum IDs for nodes and ``quantum`` 

776 (`lsst.daf.butler.Quantum`) and ``task_node`` 

777 (`lsst.pipe.base.pipeline_graph.TaskNode`) attributes in addition 

778 to those provided by 

779 `.quantum_graph.PredictedQuantumGraph.quantum_only_xgraph`. 

780 report : `Report` 

781 Object for reporting execution status. 

782 """ 

783 disable_implicit_threading() # To prevent thread contention 

784 

785 _LOG.debug("Using %r for multiprocessing start method", self._start_method) 

786 

787 # re-pack input quantum data into jobs list 

788 jobs = _JobList(xgraph) 

789 

790 # check that all tasks can run in sub-process 

791 for job in jobs.jobs.values(): 

792 if not job.task_node.task_class.canMultiprocess: 

793 raise MPGraphExecutorError( 

794 f"Task {job.task_node.label!r} does not support multiprocessing; use single process" 

795 ) 

796 

797 finishedCount, failedCount = 0, 0 

798 while jobs.pending or jobs.running: 

799 _LOG.debug("#pendingJobs: %s", len(jobs.pending)) 

800 _LOG.debug("#runningJobs: %s", len(jobs.running)) 

801 

802 # See if any jobs have finished 

803 for quantum_id in list(jobs.running): # iterate over a copy so we can remove. 

804 job = jobs.jobs[quantum_id] 

805 assert job.process is not None, "Process cannot be None" 

806 blocked: list[_Job] = [] 

807 if not job.process.is_alive(): 

808 _LOG.debug("finished: %s", job) 

809 # finished 

810 exitcode = job.process.exitcode 

811 quantum_report = job.report() 

812 report.quantaReports.append(quantum_report) 

813 if exitcode == 0: 

814 jobs.setJobState(job, JobState.FINISHED) 

815 job.cleanup() 

816 _LOG.debug("success: %s took %.3f seconds", job, time.time() - job.started) 

817 else: 

818 if job.terminated: 

819 # Was killed due to timeout. 

820 if report.status == ExecutionStatus.SUCCESS: 

821 # Do not override global FAILURE status 

822 report.status = ExecutionStatus.TIMEOUT 

823 message = f"Timeout ({self._timeout} sec) for task {job}, task is killed" 

824 blocked = jobs.setJobState(job, JobState.TIMED_OUT) 

825 else: 

826 report.status = ExecutionStatus.FAILURE 

827 # failMessage() has to be called before cleanup() 

828 message = job.failMessage() 

829 blocked = jobs.setJobState(job, JobState.FAILED) 

830 

831 job.cleanup() 

832 _LOG.debug("failed: %s", job) 

833 if self._fail_fast or exitcode == InvalidQuantumError.EXIT_CODE: 

834 # stop all running jobs 

835 for stop_quantum_id in jobs.running: 

836 stop_job = jobs.jobs[stop_quantum_id] 

837 if stop_job is not job: 837 ↛ 835line 837 didn't jump to line 835 because the condition on line 837 was always true

838 stop_job.stop() 

839 if job.state is JobState.TIMED_OUT: 

840 raise MPTimeoutError(f"Timeout ({self._timeout} sec) for task {job}.") 

841 else: 

842 raise MPGraphExecutorError(message) 

843 else: 

844 _LOG.error("%s; processing will continue for remaining tasks.", message) 

845 else: 

846 # check for timeout 

847 now = time.time() 

848 if now - job.started > self._timeout: 

849 # Try to kill it, and there is a chance that it 

850 # finishes successfully before it gets killed. Exit 

851 # status is handled by the code above on next 

852 # iteration. 

853 _LOG.debug("Terminating job %s due to timeout", job) 

854 job.stop() 

855 

856 for downstream_job in blocked: 

857 quantum_report = QuantumReport( 

858 quantumId=downstream_job.quantum_id, 

859 status=ExecutionStatus.SKIPPED, 

860 dataId=cast(DataCoordinate, downstream_job.quantum.dataId), 

861 taskLabel=downstream_job.task_node.label, 

862 ) 

863 report.quantaReports.append(quantum_report) 

864 _LOG.error("Upstream job failed for task %s, skipping this task.", downstream_job) 

865 

866 # see if we can start more jobs 

867 while len(jobs.running) < self._num_proc and jobs.pending: 

868 job = jobs.submit(self._quantum_executor, self._start_method) 

869 _LOG.debug("Submitted %s", job) 

870 

871 # Do cleanup for timed out jobs if necessary. 

872 jobs.cleanup() 

873 

874 # Print progress message if something changed. 

875 newFinished, newFailed = len(jobs.finished), len(jobs.failed) 

876 if (finishedCount, failedCount) != (newFinished, newFailed): 

877 finishedCount, failedCount = newFinished, newFailed 

878 totalCount = len(jobs.jobs) 

879 _LOG.info( 

880 "Executed %d quanta successfully, %d failed and %d remain out of total %d quanta.", 

881 finishedCount, 

882 failedCount, 

883 totalCount - finishedCount - failedCount, 

884 totalCount, 

885 ) 

886 

887 # Here we want to wait until one of the running jobs completes 

888 # but multiprocessing does not provide an API for that, for now 

889 # just sleep a little bit and go back to the loop. 

890 if jobs.running: 

891 time.sleep(0.1) 

892 

893 if jobs.failed: 

894 # print list of failed jobs 

895 _LOG.error("Failed jobs:") 

896 for quantum_id in jobs.failed: 

897 job = jobs.jobs[quantum_id] 

898 _LOG.error(" - %s: %s", job.state.name, job) 

899 

900 # if any job failed raise an exception 

901 if jobs.failed == jobs.timed_out: 

902 raise MPTimeoutError("One or more tasks timed out during execution.") 

903 else: 

904 raise MPGraphExecutorError("One or more tasks failed or timed out during execution.") 

905 

906 def getReport(self) -> Report: 

907 # Docstring inherited from base class 

908 if self._report is None: 908 ↛ 909line 908 didn't jump to line 909 because the condition on line 908 was never true

909 raise RuntimeError("getReport() called before execute()") 

910 return self._report