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

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: 

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 terrible hack is a workaround for Python threading bug: 

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

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

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

206 thread = threading.current_thread() 

207 if isinstance(thread, threading._DummyThread): 

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

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

210 

211 if logConfigState and not CliLog.configState: 

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

213 # re-initialize logging 

214 CliLog.replayConfigState(logConfigState) 

215 

216 quantumExecutor: QuantumExecutor = pickle.loads(quantumExecutor_pickle) 

217 task_node: TaskNode = pickle.loads(task_node_pickle) 

218 quantum = pickle.loads(quantum_pickle) 

219 report: QuantumReport | None = None 

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

221 # with exception-specific exit code. 

222 try: 

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

224 except RepeatableQuantumError as exc: 

225 report = QuantumReport.from_exception( 

226 quantumId=quantum_id, 

227 exception=exc, 

228 dataId=quantum.dataId, 

229 taskLabel=task_node.label, 

230 exitCode=exc.EXIT_CODE if fail_fast else None, 

231 ) 

232 if fail_fast: 

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

234 _LOG.warning(exc, exc_info=True) 

235 sys.exit(exc.EXIT_CODE) 

236 else: 

237 raise 

238 except InvalidQuantumError as exc: 

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

240 _LOG.fatal(exc, exc_info=True) 

241 report = QuantumReport.from_exception( 

242 quantumId=quantum_id, 

243 exception=exc, 

244 dataId=quantum.dataId, 

245 taskLabel=task_node.label, 

246 exitCode=exc.EXIT_CODE, 

247 ) 

248 sys.exit(exc.EXIT_CODE) 

249 except Exception as exc: 

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

251 report = QuantumReport.from_exception( 

252 quantumId=quantum_id, 

253 exception=exc, 

254 dataId=quantum.dataId, 

255 taskLabel=task_node.label, 

256 ) 

257 raise 

258 finally: 

259 if report is not None: 

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

261 # exposed. 

262 try: 

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

264 snd_conn.send(report) 

265 except Exception: 

266 pass 

267 

268 def stop(self) -> None: 

269 """Stop the process.""" 

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

271 self.process.terminate() 

272 # give it 1 second to finish or KILL 

273 for _ in range(10): 

274 time.sleep(0.1) 

275 if not self.process.is_alive(): 

276 break 

277 else: 

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

279 self.process.kill() 

280 self._terminated = True 

281 

282 def cleanup(self) -> None: 

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

284 process. 

285 """ 

286 if self.process and not self.process.is_alive(): 

287 self.process.close() 

288 self.process = None 

289 self._rcv_conn = None 

290 

291 def report(self) -> QuantumReport: 

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

293 before cleanup(). 

294 """ 

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

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

297 try: 

298 report = self._rcv_conn.recv() 

299 report.exitCode = self.process.exitcode 

300 except Exception: 

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

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

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

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

305 report = QuantumReport.from_exit_code( 

306 quantumId=self.quantum_id, 

307 exitCode=exitcode, 

308 dataId=self.quantum.dataId, 

309 taskLabel=self.task_node.label, 

310 ) 

311 if self.terminated: 

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

313 report.status = ExecutionStatus.TIMEOUT 

314 return report 

315 

316 def failMessage(self) -> str: 

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

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

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

320 exitcode = self.process.exitcode 

321 if exitcode < 0: 

322 # Negative exit code means it is killed by signal 

323 signum = -exitcode 

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

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

326 try: 

327 strsignal = signal.strsignal(signum) 

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

329 except ValueError: 

330 pass 

331 elif exitcode > 0: 

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

333 else: 

334 msg = "" 

335 return msg 

336 

337 def __str__(self) -> str: 

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

339 

340 

341class _JobList: 

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

343 

344 Parameters 

345 ---------- 

346 xgraph : `networkx.DiGraph` 

347 Directed acyclic graph of quantum IDs. 

348 """ 

349 

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

351 self.jobs = { 

352 quantum_id: _Job( 

353 quantum_id=quantum_id, 

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

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

356 ) 

357 for quantum_id in xgraph 

358 } 

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

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

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

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

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

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

365 

366 def submit( 

367 self, 

368 quantumExecutor: QuantumExecutor, 

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

370 fail_fast: bool = False, 

371 ) -> _Job: 

372 """Submit a pending job for execution. 

373 

374 Parameters 

375 ---------- 

376 quantumExecutor : `QuantumExecutor` 

377 Executor for single quantum. 

378 startMethod : `str`, optional 

379 Start method from `multiprocessing` module. 

380 fail_fast : `bool`, optional 

381 If `True` then kill subprocess on RepeatableQuantumError. 

382 

383 Returns 

384 ------- 

385 job : `_Job` 

386 The job that was submitted. 

387 """ 

388 quantum_id = self.pending.pop() 

389 job = self.jobs[quantum_id] 

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

391 self.running.add(job.quantum_id) 

392 return job 

393 

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

395 """Update job state. 

396 

397 Parameters 

398 ---------- 

399 job : `_Job` 

400 Job to submit. 

401 state : `JobState` 

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

403 states are acceptable. 

404 

405 Returns 

406 ------- 

407 blocked : `list` [ `_Job` ] 

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

409 was upstream of them and failed or timed out. 

410 """ 

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

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

413 

414 # remove job from pending/running lists 

415 if job.state == JobState.PENDING: 

416 self.pending.remove(job.quantum_id) 

417 elif job.state == JobState.RUNNING: 

418 self.running.remove(job.quantum_id) 

419 

420 quantum_id = job.quantum_id 

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

422 self.finished.discard(quantum_id) 

423 self.failed.discard(quantum_id) 

424 self.timed_out.discard(quantum_id) 

425 job._state = state 

426 match job.state: 

427 case JobState.FINISHED: 

428 self.finished.add(quantum_id) 

429 self.walker.finish(quantum_id) 

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

431 return [] 

432 case JobState.FAILED: 

433 self.failed.add(quantum_id) 

434 case JobState.TIMED_OUT: 

435 self.failed.add(quantum_id) 

436 self.timed_out.add(quantum_id) 

437 case _: 

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

439 blocked: list[_Job] = [] 

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

441 self.failed.add(downstream_quantum_id) 

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

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

444 return blocked 

445 

446 def cleanup(self) -> None: 

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

448 

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

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

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

452 """ 

453 for quantum_id in self.timed_out: 

454 job = self.jobs[quantum_id] 

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

456 if job.process is not None: 

457 job.cleanup() 

458 

459 

460class MPGraphExecutorError(Exception): 

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

462 

463 pass 

464 

465 

466class MPTimeoutError(MPGraphExecutorError): 

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

468 

469 pass 

470 

471 

472class MPGraphExecutor(QuantumGraphExecutor): 

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

474 execution of Quanta. 

475 

476 Parameters 

477 ---------- 

478 num_proc : `int` 

479 Number of processes to use for executing tasks. 

480 timeout : `float` 

481 Time in seconds to wait for tasks to finish. 

482 quantum_executor : `.quantum_graph_executor.QuantumExecutor` 

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

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

485 start_method : `str`, optional 

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

487 one for current platform. 

488 fail_fast : `bool`, optional 

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

490 pdb : `str`, optional 

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

492 event of an exception. 

493 execution_graph_fixup : `.execution_graph_fixup.ExecutionGraphFixup`, \ 

494 optional 

495 Instance used for modification of execution graph. 

496 """ 

497 

498 def __init__( 

499 self, 

500 *, 

501 num_proc: int, 

502 timeout: float, 

503 quantum_executor: QuantumExecutor, 

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

505 fail_fast: bool = False, 

506 pdb: str | None = None, 

507 execution_graph_fixup: ExecutionGraphFixup | None = None, 

508 ): 

509 self._num_proc = num_proc 

510 self._timeout = timeout 

511 self._quantum_executor = quantum_executor 

512 self._fail_fast = fail_fast 

513 self._pdb = pdb 

514 self._execution_graph_fixup = execution_graph_fixup 

515 self._report: Report | None = None 

516 

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

518 if start_method is None: 

519 start_method = "spawn" 

520 self._start_method = start_method 

521 

522 def execute( 

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

524 ) -> None: 

525 # Docstring inherited from QuantumGraphExecutor.execute 

526 old_graph: QuantumGraph | None = None 

527 if isinstance(graph, QuantumGraph): 

528 old_graph = graph 

529 new_graph = PredictedQuantumGraph.from_old_quantum_graph(old_graph) 

530 else: 

531 new_graph = graph 

532 xgraph = self._make_xgraph(new_graph, old_graph) 

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

534 with ExitStack() as exit_stack: 

535 provenance_writer: ProvenanceQuantumGraphWriter | None = None 

536 if provenance_graph_file is not None: 

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

538 raise NotImplementedError( 

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

540 ) 

541 provenance_writer = ProvenanceQuantumGraphWriter( 

542 provenance_graph_file, 

543 exit_stack=exit_stack, 

544 log_on_close=LogOnClose(_LOG.log), 

545 predicted=new_graph, 

546 ) 

547 try: 

548 if self._num_proc > 1: 

549 self._execute_quanta_mp(xgraph, self._report) 

550 else: 

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

552 except Exception as exc: 

553 self._report.set_exception(exc) 

554 raise 

555 if provenance_writer is not None: 

556 provenance_writer.write_overall_inputs() 

557 provenance_writer.write_packages() 

558 provenance_writer.write_init_outputs(assume_existence=True) 

559 

560 def _make_xgraph( 

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

562 ) -> networkx.DiGraph: 

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

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

565 attributes. 

566 

567 Parameters 

568 ---------- 

569 new_graph : `.quantum_graph.PredictedQuantumGraph` 

570 New quantum graph object. 

571 old_graph : `.QuantumGraph` or `None` 

572 Equivalent old quantum graph object. 

573 

574 Returns 

575 ------- 

576 xgraph : `networkx.DiGraph` 

577 NetworkX DAG with quantum IDs as node keys. 

578 

579 Raises 

580 ------ 

581 MPGraphExecutorError 

582 Raised if execution graph cannot be ordered after modification, 

583 i.e. it has dependency cycles. 

584 """ 

585 new_graph.build_execution_quanta() 

586 xgraph = new_graph.quantum_only_xgraph.copy() 

587 if self._execution_graph_fixup: 

588 try: 

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

590 except NotImplementedError: 

591 # Backwards compatibility. 

592 if old_graph is None: 

593 old_graph = new_graph.to_old_quantum_graph() 

594 old_graph = self._execution_graph_fixup.fixupQuanta(old_graph) 

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

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

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

598 if networkx.dag.has_cycle(xgraph): 

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

600 return xgraph 

601 

602 def _execute_quanta_in_process( 

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

604 ) -> None: 

605 """Execute all Quanta in current process. 

606 

607 Parameters 

608 ---------- 

609 xgraph : `networkx.DiGraph` 

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

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

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

613 to those provided by 

614 `.quantum_graph.PredictedQuantumGraph.quantum_only_xgraph`. 

615 report : `Report` 

616 Object for reporting execution status. 

617 provenance_writer : `.quantum_graph.ProvenanceQuantumGraphWriter` or \ 

618 `None` 

619 Object for recording provenance. 

620 """ 

621 

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

623 node_state = xgraph.nodes[quantum_id] 

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

625 

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

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

628 for unblocked_quanta in walker: 

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

630 node_state: PredictedQuantumInfo = xgraph.nodes[quantum_id] 

631 data_id = node_state["data_id"] 

632 task_node = node_state["pipeline_node"] 

633 quantum = node_state["quantum"] 

634 

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

636 fail_exit_code: int | None = None 

637 task_metadata: TaskMetadata | None = None 

638 task_logs = ButlerLogRecords([]) 

639 try: 

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

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

642 # debugger before exiting if debugging is enabled. 

643 try: 

644 execution_result = self._quantum_executor.execute( 

645 task_node, quantum, quantum_id=quantum_id, log_records=task_logs 

646 ) 

647 if execution_result.report: 

648 report.quantaReports.append(execution_result.report) 

649 task_metadata = execution_result.task_metadata 

650 success_count += 1 

651 walker.finish(quantum_id) 

652 except RepeatableQuantumError as exc: 

653 if self._fail_fast: 

654 _LOG.warning( 

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

656 quantum_id, 

657 task_node.label, 

658 data_id, 

659 ) 

660 _LOG.warning(exc, exc_info=True) 

661 fail_exit_code = exc.EXIT_CODE 

662 raise 

663 except InvalidQuantumError as exc: 

664 _LOG.fatal( 

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

666 ) 

667 _LOG.fatal(exc, exc_info=True) 

668 fail_exit_code = exc.EXIT_CODE 

669 raise 

670 except Exception as exc: 

671 quantum_report = QuantumReport.from_exception( 

672 exception=exc, 

673 dataId=data_id, 

674 taskLabel=task_node.label, 

675 ) 

676 report.quantaReports.append(quantum_report) 

677 

678 if self._pdb and sys.stdin.isatty() and sys.stdout.isatty(): 

679 _LOG.error( 

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

681 quantum_id, 

682 task_node.label, 

683 data_id, 

684 exc_info=exc, 

685 ) 

686 try: 

687 pdb = importlib.import_module(self._pdb) 

688 except ImportError as imp_exc: 

689 raise MPGraphExecutorError( 

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

691 ) from exc 

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

693 raise MPGraphExecutorError( 

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

695 ) from exc 

696 pdb.post_mortem(exc.__traceback__) 

697 

698 report.status = ExecutionStatus.FAILURE 

699 failed_count += 1 

700 

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

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

703 # enabled. 

704 if fail_exit_code is not None: 

705 sys.exit(fail_exit_code) 

706 if self._fail_fast: 

707 raise MPGraphExecutorError( 

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

709 ) from exc 

710 else: 

711 _LOG.error( 

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

713 quantum_id, 

714 task_node.label, 

715 data_id, 

716 exc_info=exc, 

717 ) 

718 

719 for downstream_quantum_id in walker.fail(quantum_id): 

720 downstream_node_state = xgraph.nodes[downstream_quantum_id] 

721 failed_quantum_report = QuantumReport( 

722 status=ExecutionStatus.SKIPPED, 

723 dataId=downstream_node_state["data_id"], 

724 taskLabel=downstream_node_state["task_label"], 

725 ) 

726 report.quantaReports.append(failed_quantum_report) 

727 _LOG.error( 

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

729 downstream_quantum_id, 

730 downstream_node_state["task_label"], 

731 downstream_node_state["data_id"], 

732 ) 

733 failed_count += 1 

734 

735 if provenance_writer is not None: 

736 provenance_writer.write_quantum_provenance( 

737 quantum_id, metadata=task_metadata, logs=task_logs 

738 ) 

739 

740 _LOG.info( 

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

742 success_count, 

743 failed_count, 

744 total_count - success_count - failed_count, 

745 total_count, 

746 ) 

747 

748 # Raise an exception if there were any failures. 

749 if failed_count: 

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

751 

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

753 """Execute all Quanta in separate processes. 

754 

755 Parameters 

756 ---------- 

757 xgraph : `networkx.DiGraph` 

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

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

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

761 to those provided by 

762 `.quantum_graph.PredictedQuantumGraph.quantum_only_xgraph`. 

763 report : `Report` 

764 Object for reporting execution status. 

765 """ 

766 disable_implicit_threading() # To prevent thread contention 

767 

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

769 

770 # re-pack input quantum data into jobs list 

771 jobs = _JobList(xgraph) 

772 

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

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

775 if not job.task_node.task_class.canMultiprocess: 

776 raise MPGraphExecutorError( 

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

778 ) 

779 

780 finishedCount, failedCount = 0, 0 

781 while jobs.pending or jobs.running: 

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

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

784 

785 # See if any jobs have finished 

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

787 job = jobs.jobs[quantum_id] 

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

789 blocked: list[_Job] = [] 

790 if not job.process.is_alive(): 

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

792 # finished 

793 exitcode = job.process.exitcode 

794 quantum_report = job.report() 

795 report.quantaReports.append(quantum_report) 

796 if exitcode == 0: 

797 jobs.setJobState(job, JobState.FINISHED) 

798 job.cleanup() 

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

800 else: 

801 if job.terminated: 

802 # Was killed due to timeout. 

803 if report.status == ExecutionStatus.SUCCESS: 

804 # Do not override global FAILURE status 

805 report.status = ExecutionStatus.TIMEOUT 

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

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

808 else: 

809 report.status = ExecutionStatus.FAILURE 

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

811 message = job.failMessage() 

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

813 

814 job.cleanup() 

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

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

817 # stop all running jobs 

818 for stop_quantum_id in jobs.running: 

819 stop_job = jobs.jobs[stop_quantum_id] 

820 if stop_job is not job: 

821 stop_job.stop() 

822 if job.state is JobState.TIMED_OUT: 

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

824 else: 

825 raise MPGraphExecutorError(message) 

826 else: 

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

828 else: 

829 # check for timeout 

830 now = time.time() 

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

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

833 # finishes successfully before it gets killed. Exit 

834 # status is handled by the code above on next 

835 # iteration. 

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

837 job.stop() 

838 

839 for downstream_job in blocked: 

840 quantum_report = QuantumReport( 

841 quantumId=downstream_job.quantum_id, 

842 status=ExecutionStatus.SKIPPED, 

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

844 taskLabel=downstream_job.task_node.label, 

845 ) 

846 report.quantaReports.append(quantum_report) 

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

848 

849 # see if we can start more jobs 

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

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

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

853 

854 # Do cleanup for timed out jobs if necessary. 

855 jobs.cleanup() 

856 

857 # Print progress message if something changed. 

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

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

860 finishedCount, failedCount = newFinished, newFailed