Coverage for python/lsst/ctrl/bps/generic_workflow.py: 34%

368 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-04-11 02:16 -0700

1# This file is part of ctrl_bps. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

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

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

7# for details of code ownership. 

8# 

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

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

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

12# (at your option) any later version. 

13# 

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

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

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

17# GNU General Public License for more details. 

18# 

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

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

21 

22"""Class definitions for a Generic Workflow Graph. 

23""" 

24 

25__all__ = ["GenericWorkflow", "GenericWorkflowFile", "GenericWorkflowJob", "GenericWorkflowExec"] 

26 

27 

28import dataclasses 

29import itertools 

30import logging 

31import pickle 

32from collections import Counter, defaultdict 

33from typing import Optional 

34 

35from lsst.utils.iteration import ensure_iterable 

36from networkx import DiGraph, topological_sort 

37from networkx.algorithms.dag import is_directed_acyclic_graph 

38 

39from .bps_draw import draw_networkx_dot 

40 

41_LOG = logging.getLogger(__name__) 

42 

43 

44@dataclasses.dataclass 

45class GenericWorkflowFile: 

46 """Information about a file that may be needed by various workflow 

47 management services. 

48 """ 

49 

50 name: str 

51 """Lookup key (logical file name) of file/directory. Must be unique 

52 within run. 

53 """ 

54 

55 src_uri: str or None # don't know that need ResourcePath 

56 """Original location of file/directory. 

57 """ 

58 

59 wms_transfer: bool 

60 """Whether the WMS should ignore file or not. Default is False. 

61 """ 

62 

63 job_access_remote: bool 

64 """Whether the job can remotely access file (using separately specified 

65 file access protocols). Default is False. 

66 """ 

67 

68 job_shared: bool 

69 """Whether job requires its own copy of this file. Default is False. 

70 """ 

71 

72 # As of python 3.7.8, can't use __slots__ + dataclass if give default 

73 # values, so writing own __init__. 

74 def __init__( 

75 self, 

76 name: str, 

77 src_uri: str = None, 

78 wms_transfer: bool = False, 

79 job_access_remote: bool = False, 

80 job_shared: bool = False, 

81 ): 

82 self.name = name 

83 self.src_uri = src_uri 

84 self.wms_transfer = wms_transfer 

85 self.job_access_remote = job_access_remote 

86 self.job_shared = job_shared 

87 

88 __slots__ = ("name", "src_uri", "wms_transfer", "job_access_remote", "job_shared") 

89 

90 def __hash__(self): 

91 return hash(self.name) 

92 

93 

94@dataclasses.dataclass 

95class GenericWorkflowExec: 

96 """Information about an executable that may be needed by various workflow 

97 management services. 

98 """ 

99 

100 name: str 

101 """Lookup key (logical file name) of executable. Must be unique 

102 within run. 

103 """ 

104 

105 src_uri: str or None # don't know that need ResourcePath 

106 """Original location of executable. 

107 """ 

108 

109 transfer_executable: bool 

110 """Whether the WMS/plugin is responsible for staging executable to 

111 location usable by job. 

112 """ 

113 

114 # As of python 3.7.8, can't use __slots__ + dataclass if give default 

115 # values, so writing own __init__. 

116 def __init__(self, name: str, src_uri: str = None, transfer_executable: bool = False): 

117 self.name = name 

118 self.src_uri = src_uri 

119 self.transfer_executable = transfer_executable 

120 

121 __slots__ = ("name", "src_uri", "transfer_executable") 

122 

123 def __hash__(self): 

124 return hash(self.name) 

125 

126 

127@dataclasses.dataclass 

128class GenericWorkflowJob: 

129 """Information about a job that may be needed by various workflow 

130 management services. 

131 """ 

132 

133 name: str 

134 """Name of job. Must be unique within workflow. 

135 """ 

136 

137 label: Optional[str] 

138 """Primary user-facing label for job. Does not need to be unique 

139 and may be used for summary reports. 

140 """ 

141 

142 quanta_counts: Optional[Counter] 

143 """Counts of quanta per task label in job. 

144 """ 

145 

146 tags: Optional[dict] 

147 """Other key/value pairs for job that user may want to use as a filter. 

148 """ 

149 

150 executable: Optional[GenericWorkflowExec] 

151 """Executable for job. 

152 """ 

153 

154 arguments: Optional[str] 

155 """Command line arguments for job. 

156 """ 

157 

158 cmdvals: Optional[dict] 

159 """Values for variables in cmdline when using lazy command line creation. 

160 """ 

161 

162 memory_multiplier: Optional[float] 

163 """Memory growth rate between retries. 

164 """ 

165 

166 request_memory: Optional[int] # MB 

167 """Max memory (in MB) that the job is expected to need. 

168 """ 

169 

170 request_memory_max: Optional[int] # MB 

171 """Max memory (in MB) that the job should ever use. 

172 """ 

173 

174 request_cpus: Optional[int] # cores 

175 """Max number of cpus that the job is expected to need. 

176 """ 

177 

178 request_disk: Optional[int] # MB 

179 """Max amount of job scratch disk (in MB) that the job is expected to need. 

180 """ 

181 

182 request_walltime: Optional[str] # minutes 

183 """Max amount of time (in seconds) that the job is expected to need. 

184 """ 

185 

186 compute_site: Optional[str] 

187 """Key to look up site-specific information for running the job. 

188 """ 

189 

190 accounting_group: Optional[str] 

191 """Name of the accounting group to use. 

192 """ 

193 

194 accounting_user: Optional[str] 

195 """Name of the user to use for accounting purposes. 

196 """ 

197 

198 mail_to: Optional[str] 

199 """Comma separated list of email addresses for emailing job status. 

200 """ 

201 

202 when_to_mail: Optional[str] 

203 """WMS-specific terminology for when to email job status. 

204 """ 

205 

206 number_of_retries: Optional[int] 

207 """Number of times to automatically retry a failed job. 

208 """ 

209 

210 retry_unless_exit: Optional[int] 

211 """Exit code for job that means to not automatically retry. 

212 """ 

213 

214 abort_on_value: Optional[int] 

215 """Job exit value for signals to abort the entire workflow. 

216 """ 

217 

218 abort_return_value: Optional[int] 

219 """Exit value to use when aborting the entire workflow. 

220 """ 

221 

222 priority: Optional[str] 

223 """Initial priority of job in WMS-format. 

224 """ 

225 

226 category: Optional[str] 

227 """WMS-facing label of job within single workflow (e.g., can be used for 

228 throttling jobs within a single workflow). 

229 """ 

230 

231 concurrency_limit: Optional[str] 

232 """Names of concurrency limits that the WMS plugin can appropriately 

233 translate to limit the number of this job across all running workflows. 

234 """ 

235 

236 queue: Optional[str] 

237 """Name of queue to use. Different WMS can translate this concept 

238 differently. 

239 """ 

240 

241 pre_cmdline: Optional[str] 

242 """Command line to be executed prior to executing job. 

243 """ 

244 

245 post_cmdline: Optional[str] 

246 """Command line to be executed after job executes. 

247 

248 Should be executed regardless of exit status. 

249 """ 

250 

251 preemptible: Optional[bool] 

252 """The flag indicating whether the job can be preempted. 

253 """ 

254 

255 profile: Optional[dict] 

256 """Nested dictionary of WMS-specific key/value pairs with primary key being 

257 WMS key (e.g., pegasus, condor, panda). 

258 """ 

259 

260 attrs: Optional[dict] 

261 """Key/value pairs of job attributes (for WMS that have attributes in 

262 addition to commands). 

263 """ 

264 

265 environment: Optional[dict] 

266 """Environment variable names and values to be explicitly set inside job. 

267 """ 

268 

269 compute_cloud: Optional[str] 

270 """Key to look up cloud-specific information for running the job. 

271 """ 

272 

273 # As of python 3.7.8, can't use __slots__ if give default values, so 

274 # writing own __init__. 

275 def __init__(self, name, label="UNK"): 

276 self.name = name 

277 self.label = label 

278 self.quanta_counts = Counter() 

279 self.tags = {} 

280 self.executable = None 

281 self.arguments = None 

282 self.cmdvals = {} 

283 self.memory_multiplier = None 

284 self.request_memory = None 

285 self.request_memory_max = None 

286 self.request_cpus = None 

287 self.request_disk = None 

288 self.request_walltime = None 

289 self.compute_site = None 

290 self.accounting_group = None 

291 self.accounting_user = None 

292 self.mail_to = None 

293 self.when_to_mail = None 

294 self.number_of_retries = None 

295 self.retry_unless_exit = None 

296 self.abort_on_value = None 

297 self.abort_return_value = None 

298 self.priority = None 

299 self.category = None 

300 self.concurrency_limit = None 

301 self.queue = None 

302 self.pre_cmdline = None 

303 self.post_cmdline = None 

304 self.preemptible = None 

305 self.profile = {} 

306 self.attrs = {} 

307 self.environment = {} 

308 self.compute_cloud = None 

309 

310 __slots__ = ( 

311 "name", 

312 "label", 

313 "quanta_counts", 

314 "tags", 

315 "mail_to", 

316 "when_to_mail", 

317 "executable", 

318 "arguments", 

319 "cmdvals", 

320 "memory_multiplier", 

321 "request_memory", 

322 "request_memory_max", 

323 "request_cpus", 

324 "request_disk", 

325 "request_walltime", 

326 "number_of_retries", 

327 "retry_unless_exit", 

328 "abort_on_value", 

329 "abort_return_value", 

330 "compute_site", 

331 "accounting_group", 

332 "accounting_user", 

333 "environment", 

334 "priority", 

335 "category", 

336 "concurrency_limit", 

337 "queue", 

338 "pre_cmdline", 

339 "post_cmdline", 

340 "preemptible", 

341 "profile", 

342 "attrs", 

343 "compute_cloud", 

344 ) 

345 

346 def __hash__(self): 

347 return hash(self.name) 

348 

349 

350class GenericWorkflow(DiGraph): 

351 """A generic representation of a workflow used to submit to specific 

352 workflow management systems. 

353 

354 Parameters 

355 ---------- 

356 name : `str` 

357 Name of generic workflow. 

358 incoming_graph_data : `Any`, optional 

359 Data used to initialized graph that is passed through to DiGraph 

360 constructor. Can be any type supported by networkx.DiGraph. 

361 attr : `dict` 

362 Keyword arguments passed through to DiGraph constructor. 

363 """ 

364 

365 def __init__(self, name, incoming_graph_data=None, **attr): 

366 super().__init__(incoming_graph_data, **attr) 

367 self._name = name 

368 self.run_attrs = {} 

369 self._job_labels = GenericWorkflowLabels() 

370 self._files = {} 

371 self._executables = {} 

372 self._inputs = {} # mapping job.names to list of GenericWorkflowFile 

373 self._outputs = {} # mapping job.names to list of GenericWorkflowFile 

374 self.run_id = None 

375 self._final = None 

376 

377 @property 

378 def name(self): 

379 """Retrieve name of generic workflow. 

380 

381 Returns 

382 ------- 

383 name : `str` 

384 Name of generic workflow. 

385 """ 

386 return self._name 

387 

388 @property 

389 def quanta_counts(self): 

390 """Count of quanta per task label (`collections.Counter`).""" 

391 qcounts = Counter() 

392 for job_name in self: 

393 gwjob = self.get_job(job_name) 

394 if gwjob.quanta_counts is not None: 

395 qcounts += gwjob.quanta_counts 

396 return qcounts 

397 

398 @property 

399 def labels(self): 

400 """Job labels (`list` [`str`], read-only)""" 

401 return self._job_labels.labels 

402 

403 def regenerate_labels(self): 

404 """Regenerate the list of job labels.""" 

405 self._job_labels = GenericWorkflowLabels() 

406 for job_name in self: 

407 job = self.get_job(job_name) 

408 self._job_labels.add_job( 

409 job, 

410 [self.get_job(p).label for p in self.predecessors(job.name)], 

411 [self.get_job(p).label for p in self.successors(job.name)], 

412 ) 

413 

414 @property 

415 def job_counts(self): 

416 """Count of jobs per job label (`collections.Counter`).""" 

417 jcounts = self._job_labels.job_counts 

418 

419 # Final is separate 

420 final = self.get_final() 

421 if final: 

422 if isinstance(final, GenericWorkflow): 

423 jcounts.update(final.job_counts()) 

424 else: 

425 jcounts[final.label] += 1 

426 

427 return jcounts 

428 

429 def __iter__(self): 

430 """Return iterator of job names in topologically sorted order.""" 

431 return topological_sort(self) 

432 

433 def get_files(self, data=False, transfer_only=True): 

434 """Retrieve files from generic workflow. 

435 

436 Need API in case change way files are stored (e.g., make 

437 workflow a bipartite graph with jobs and files nodes). 

438 

439 Parameters 

440 ---------- 

441 data : `bool`, optional 

442 Whether to return the file data as well as the file object name. 

443 (The defaults is False.) 

444 transfer_only : `bool`, optional 

445 Whether to only return files for which a workflow management system 

446 would be responsible for transferring. 

447 

448 Returns 

449 ------- 

450 files : `list` [`lsst.ctrl.bps.GenericWorkflowFile`] or `list` [`str`] 

451 File names or objects from generic workflow meeting specifications. 

452 """ 

453 files = [] 

454 for filename, file in self._files.items(): 

455 if not transfer_only or file.wms_transfer: 

456 if not data: 

457 files.append(filename) 

458 else: 

459 files.append(file) 

460 return files 

461 

462 def add_job(self, job, parent_names=None, child_names=None): 

463 """Add job to generic workflow. 

464 

465 Parameters 

466 ---------- 

467 job : `lsst.ctrl.bps.GenericWorkflowJob` 

468 Job to add to the generic workflow. 

469 parent_names : `list` [`str`], optional 

470 Names of jobs that are parents of given job 

471 child_names : `list` [`str`], optional 

472 Names of jobs that are children of given job 

473 """ 

474 _LOG.debug("job: %s (%s)", job.name, job.label) 

475 _LOG.debug("parent_names: %s", parent_names) 

476 _LOG.debug("child_names: %s", child_names) 

477 if not isinstance(job, GenericWorkflowJob): 

478 raise RuntimeError(f"Invalid type for job to be added to GenericWorkflowGraph ({type(job)}).") 

479 if self.has_node(job.name): 

480 raise RuntimeError(f"Job {job.name} already exists in GenericWorkflowGraph.") 

481 super().add_node(job.name, job=job) 

482 self.add_job_relationships(parent_names, job.name) 

483 self.add_job_relationships(job.name, child_names) 

484 self.add_executable(job.executable) 

485 self._job_labels.add_job( 

486 job, 

487 [self.get_job(p).label for p in self.predecessors(job.name)], 

488 [self.get_job(p).label for p in self.successors(job.name)], 

489 ) 

490 

491 def add_node(self, node_for_adding, **attr): 

492 """Override networkx function to call more specific add_job function. 

493 

494 Parameters 

495 ---------- 

496 node_for_adding : `lsst.ctrl.bps.GenericWorkflowJob` 

497 Job to be added to generic workflow. 

498 attr : 

499 Needed to match original networkx function, but not used. 

500 """ 

501 self.add_job(node_for_adding) 

502 

503 def add_job_relationships(self, parents, children): 

504 """Add dependencies between parent and child jobs. All parents will 

505 be connected to all children. 

506 

507 Parameters 

508 ---------- 

509 parents : `list` [`str`] 

510 Parent job names. 

511 children : `list` [`str`] 

512 Children job names. 

513 """ 

514 if parents is not None and children is not None: 

515 self.add_edges_from(itertools.product(ensure_iterable(parents), ensure_iterable(children))) 

516 self._job_labels.add_job_relationships( 

517 [self.get_job(n).label for n in ensure_iterable(parents)], 

518 [self.get_job(n).label for n in ensure_iterable(children)], 

519 ) 

520 

521 def add_edges_from(self, ebunch_to_add, **attr): 

522 """Add several edges between jobs in the generic workflow. 

523 

524 Parameters 

525 ---------- 

526 ebunch_to_add : Iterable [`tuple`] 

527 Iterable of job name pairs between which a dependency should be 

528 saved. 

529 attr : keyword arguments, optional 

530 Data can be assigned using keyword arguments (not currently used). 

531 """ 

532 for edge_to_add in ebunch_to_add: 

533 self.add_edge(edge_to_add[0], edge_to_add[1], **attr) 

534 

535 def add_edge(self, u_of_edge: str, v_of_edge: str, **attr): 

536 """Add edge connecting jobs in workflow. 

537 

538 Parameters 

539 ---------- 

540 u_of_edge : `str` 

541 Name of parent job. 

542 v_of_edge : `str` 

543 Name of child job. 

544 attr : keyword arguments, optional 

545 Attributes to save with edge. 

546 """ 

547 if u_of_edge not in self: 

548 raise RuntimeError(f"{u_of_edge} not in GenericWorkflow") 

549 if v_of_edge not in self: 

550 raise RuntimeError(f"{v_of_edge} not in GenericWorkflow") 

551 super().add_edge(u_of_edge, v_of_edge, **attr) 

552 

553 def get_job(self, job_name: str): 

554 """Retrieve job by name from workflow. 

555 

556 Parameters 

557 ---------- 

558 job_name : `str` 

559 Name of job to retrieve. 

560 

561 Returns 

562 ------- 

563 job : `lsst.ctrl.bps.GenericWorkflowJob` 

564 Job matching given job_name. 

565 """ 

566 return self.nodes[job_name]["job"] 

567 

568 def del_job(self, job_name: str): 

569 """Delete job from generic workflow leaving connected graph. 

570 

571 Parameters 

572 ---------- 

573 job_name : `str` 

574 Name of job to delete from workflow. 

575 """ 

576 job = self.get_job(job_name) 

577 

578 # Remove from job labels 

579 self._job_labels.del_job(job) 

580 

581 # Connect all parent jobs to all children jobs. 

582 parents = self.predecessors(job_name) 

583 children = self.successors(job_name) 

584 self.add_job_relationships(parents, children) 

585 

586 # Delete job node (which deletes edges). 

587 self.remove_node(job_name) 

588 

589 def add_job_inputs(self, job_name, files): 

590 """Add files as inputs to specified job. 

591 

592 Parameters 

593 ---------- 

594 job_name : `str` 

595 Name of job to which inputs should be added 

596 files : `lsst.ctrl.bps.GenericWorkflowFile` or \ 

597 `list` [`lsst.ctrl.bps.GenericWorkflowFile`] 

598 File object(s) to be added as inputs to the specified job. 

599 """ 

600 self._inputs.setdefault(job_name, []) 

601 for file in ensure_iterable(files): 

602 # Save the central copy 

603 if file.name not in self._files: 

604 self._files[file.name] = file 

605 

606 # Save the job reference to the file 

607 self._inputs[job_name].append(file) 

608 

609 def get_file(self, name): 

610 """Retrieve a file object by name. 

611 

612 Parameters 

613 ---------- 

614 name : `str` 

615 Name of file object 

616 

617 Returns 

618 ------- 

619 gwfile : `lsst.ctrl.bps.GenericWorkflowFile` 

620 File matching given name. 

621 """ 

622 return self._files[name] 

623 

624 def add_file(self, gwfile): 

625 """Add file object. 

626 

627 Parameters 

628 ---------- 

629 gwfile : `lsst.ctrl.bps.GenericWorkflowFile` 

630 File object to add to workflow 

631 """ 

632 if gwfile.name not in self._files: 

633 self._files[gwfile.name] = gwfile 

634 else: 

635 _LOG.debug("Skipped add_file for existing file %s", gwfile.name) 

636 

637 def get_job_inputs(self, job_name, data=True, transfer_only=False): 

638 """Return the input files for the given job. 

639 

640 Parameters 

641 ---------- 

642 job_name : `str` 

643 Name of the job. 

644 data : `bool`, optional 

645 Whether to return the file data as well as the file object name. 

646 transfer_only : `bool`, optional 

647 Whether to only return files for which a workflow management system 

648 would be responsible for transferring. 

649 

650 Returns 

651 ------- 

652 inputs : `list` [`lsst.ctrl.bps.GenericWorkflowFile`] 

653 Input files for the given job. If no input files for the job, 

654 returns an empty list. 

655 """ 

656 inputs = [] 

657 if job_name in self._inputs: 

658 for gwfile in self._inputs[job_name]: 

659 if not transfer_only or gwfile.wms_transfer: 

660 if not data: 

661 inputs.append(gwfile.name) 

662 else: 

663 inputs.append(gwfile) 

664 return inputs 

665 

666 def add_job_outputs(self, job_name, files): 

667 """Add output files to a job. 

668 

669 Parameters 

670 ---------- 

671 job_name : `str` 

672 Name of job to which the files should be added as outputs. 

673 files : `list` [`lsst.ctrl.bps.GenericWorkflowFile`] 

674 File objects to be added as outputs for specified job. 

675 """ 

676 self._outputs.setdefault(job_name, []) 

677 

678 for file_ in ensure_iterable(files): 

679 # Save the central copy 

680 if file_.name not in self._files: 

681 self._files[file_.name] = file_ 

682 

683 # Save the job reference to the file 

684 self._outputs[job_name].append(file_) 

685 

686 def get_job_outputs(self, job_name, data=True, transfer_only=False): 

687 """Return the output files for the given job. 

688 

689 Parameters 

690 ---------- 

691 job_name : `str` 

692 Name of the job. 

693 data : `bool` 

694 Whether to return the file data as well as the file object name. 

695 It defaults to `True` thus returning file data as well. 

696 transfer_only : `bool` 

697 Whether to only return files for which a workflow management system 

698 would be responsible for transferring. It defaults to `False` thus 

699 returning all output files. 

700 

701 Returns 

702 ------- 

703 outputs : `list` [`lsst.ctrl.bps.GenericWorkflowFile`] 

704 Output files for the given job. If no output files for the job, 

705 returns an empty list. 

706 """ 

707 outputs = [] 

708 

709 if job_name in self._outputs: 

710 for file_name in self._outputs[job_name]: 

711 file = self._files[file_name] 

712 if not transfer_only or file.wms_transfer: 

713 if not data: 

714 outputs.append(file_name) 

715 else: 

716 outputs.append(self._files[file_name]) 

717 return outputs 

718 

719 def draw(self, stream, format_="dot"): 

720 """Output generic workflow in a visualization format. 

721 

722 Parameters 

723 ---------- 

724 stream : `str` or `io.BufferedIOBase` 

725 Stream to which the visualization should be written. 

726 format_ : `str`, optional 

727 Which visualization format to use. It defaults to the format for 

728 the dot program. 

729 """ 

730 draw_funcs = {"dot": draw_networkx_dot} 

731 if format_ in draw_funcs: 

732 draw_funcs[format_](self, stream) 

733 else: 

734 raise RuntimeError(f"Unknown draw format ({format_}") 

735 

736 def save(self, stream, format_="pickle"): 

737 """Save the generic workflow in a format that is loadable. 

738 

739 Parameters 

740 ---------- 

741 stream : `str` or `io.BufferedIOBase` 

742 Stream to pass to the format-specific writer. Accepts anything 

743 that the writer accepts. 

744 

745 format_ : `str`, optional 

746 Format in which to write the data. It defaults to pickle format. 

747 """ 

748 if format_ == "pickle": 

749 pickle.dump(self, stream) 

750 else: 

751 raise RuntimeError(f"Unknown format ({format_})") 

752 

753 @classmethod 

754 def load(cls, stream, format_="pickle"): 

755 """Load a GenericWorkflow from the given stream 

756 

757 Parameters 

758 ---------- 

759 stream : `str` or `io.BufferedIOBase` 

760 Stream to pass to the format-specific loader. Accepts anything that 

761 the loader accepts. 

762 format_ : `str`, optional 

763 Format of data to expect when loading from stream. It defaults 

764 to pickle format. 

765 

766 Returns 

767 ------- 

768 generic_workflow : `lsst.ctrl.bps.GenericWorkflow` 

769 Generic workflow loaded from the given stream 

770 """ 

771 if format_ == "pickle": 

772 return pickle.load(stream) 

773 

774 raise RuntimeError(f"Unknown format ({format_})") 

775 

776 def validate(self): 

777 """Run checks to ensure that the generic workflow graph is valid.""" 

778 # Make sure a directed acyclic graph 

779 assert is_directed_acyclic_graph(self) 

780 

781 def add_workflow_source(self, workflow): 

782 """Add given workflow as new source to this workflow. 

783 

784 Parameters 

785 ---------- 

786 workflow : `lsst.ctrl.bps.GenericWorkflow` 

787 """ 

788 # Find source nodes in self. 

789 self_sources = [n for n in self if self.in_degree(n) == 0] 

790 _LOG.debug("self_sources = %s", self_sources) 

791 

792 # Find sink nodes of workflow. 

793 new_sinks = [n for n in workflow if workflow.out_degree(n) == 0] 

794 _LOG.debug("new sinks = %s", new_sinks) 

795 

796 # Add new workflow nodes to self graph and make new edges. 

797 self.add_nodes_from(workflow.nodes(data=True)) 

798 self.add_edges_from(workflow.edges()) 

799 for source in self_sources: 

800 for sink in new_sinks: 

801 self.add_edge(sink, source) 

802 

803 # Add separately stored info 

804 for job_name in workflow: 

805 job = self.get_job(job_name) 

806 # Add job labels 

807 self._job_labels.add_job( 

808 job, 

809 [self.get_job(p).label for p in self.predecessors(job.name)], 

810 [self.get_job(p).label for p in self.successors(job.name)], 

811 ) 

812 # Files are stored separately so copy them. 

813 self.add_job_inputs(job_name, workflow.get_job_inputs(job_name, data=True)) 

814 self.add_job_outputs(job_name, workflow.get_job_outputs(job_name, data=True)) 

815 # Executables are stored separately so copy them. 

816 self.add_job_inputs(job_name, workflow.get_job_inputs(job_name, data=True)) 

817 self.add_executable(workflow.get_job(job_name).executable) 

818 

819 def add_final(self, final): 

820 """Add special final job/workflow to the generic workflow. 

821 

822 Parameters 

823 ---------- 

824 final : `lsst.ctrl.bps.GenericWorkflowJob` or \ 

825 `lsst.ctrl.bps.GenericWorkflow` 

826 Information needed to execute the special final job(s), the 

827 job(s) to be executed after all jobs that can be executed 

828 have been executed regardless of exit status of any of the 

829 jobs. 

830 """ 

831 if not isinstance(final, GenericWorkflowJob) and not isinstance(final, GenericWorkflow): 

832 raise TypeError("Invalid type for GenericWorkflow final ({type(final)})") 

833 

834 self._final = final 

835 if isinstance(final, GenericWorkflowJob): 

836 self.add_executable(final.executable) 

837 

838 def get_final(self): 

839 """Return job/workflow to be executed after all jobs that can be 

840 executed have been executed regardless of exit status of any of 

841 the jobs. 

842 

843 Returns 

844 ------- 

845 final : `lsst.ctrl.bps.GenericWorkflowJob` or \ 

846 `lsst.ctrl.bps.GenericWorkflow` 

847 Information needed to execute final job(s). 

848 """ 

849 return self._final 

850 

851 def add_executable(self, executable): 

852 """Add executable to workflow's list of executables. 

853 

854 Parameters 

855 ---------- 

856 executable : `lsst.ctrl.bps.GenericWorkflowExec` 

857 Executable object to be added to workflow. 

858 """ 

859 if executable is not None: 

860 self._executables[executable.name] = executable 

861 else: 

862 _LOG.warning("executable not specified (None); cannot add to the workflow's list of executables") 

863 

864 def get_executables(self, data=False, transfer_only=True): 

865 """Retrieve executables from generic workflow. 

866 

867 Parameters 

868 ---------- 

869 data : `bool`, optional 

870 Whether to return the executable data as well as the exec object 

871 name. (The defaults is False.) 

872 transfer_only : `bool`, optional 

873 Whether to only return executables for which transfer_executable 

874 is True. 

875 

876 Returns 

877 ------- 

878 execs : `list` [`lsst.ctrl.bps.GenericWorkflowExec`] or `list` [`str`] 

879 Filtered executable names or objects from generic workflow. 

880 """ 

881 execs = [] 

882 for name, executable in self._executables.items(): 

883 if not transfer_only or executable.transfer_executable: 

884 if not data: 

885 execs.append(name) 

886 else: 

887 execs.append(executable) 

888 return execs 

889 

890 def get_jobs_by_label(self, label: str): 

891 """Retrieve jobs by label from workflow. 

892 

893 Parameters 

894 ---------- 

895 label : `str` 

896 Label of jobs to retrieve. 

897 

898 Returns 

899 ------- 

900 jobs : list[`lsst.ctrl.bps.GenericWorkflowJob`] 

901 Jobs having given label. 

902 """ 

903 return self._job_labels.get_jobs_by_label(label) 

904 

905 

906class GenericWorkflowLabels: 

907 """Label-oriented representation of the GenericWorkflow.""" 

908 

909 def __init__(self): 

910 self._label_graph = DiGraph() # Dependency graph of job labels 

911 self._label_to_jobs = defaultdict(list) # mapping job label to list of GenericWorkflowJob 

912 

913 @property 

914 def labels(self): 

915 """List of job labels (`list` [`str`], read-only)""" 

916 return list(topological_sort(self._label_graph)) 

917 

918 @property 

919 def job_counts(self): 

920 """Count of jobs per job label (`collections.Counter`).""" 

921 jcounts = Counter({label: len(jobs) for label, jobs in self._label_to_jobs.items()}) 

922 return jcounts 

923 

924 def get_jobs_by_label(self, label: str): 

925 """Retrieve jobs by label from workflow. 

926 

927 Parameters 

928 ---------- 

929 label : `str` 

930 Label of jobs to retrieve. 

931 

932 Returns 

933 ------- 

934 jobs : list[`lsst.ctrl.bps.GenericWorkflowJob`] 

935 Jobs having given label. 

936 """ 

937 return self._label_to_jobs[label] 

938 

939 def add_job(self, job, parent_labels, child_labels): 

940 """Add job's label to labels. 

941 

942 Parameters 

943 ---------- 

944 job : `lsst.ctrl.bps.GenericWorkflowJob` 

945 The job to delete from the job labels. 

946 parent_labels : `list` [`str`] 

947 Parent job labels. 

948 children_labels : `list` [`str`] 

949 Children job labels. 

950 """ 

951 _LOG.debug("job: %s (%s)", job.name, job.label) 

952 _LOG.debug("parent_labels: %s", parent_labels) 

953 _LOG.debug("child_labels: %s", child_labels) 

954 self._label_to_jobs[job.label].append(job) 

955 self._label_graph.add_node(job.label) 

956 for parent in parent_labels: 

957 self._label_graph.add_edge(parent, job.label) 

958 for child in child_labels: 

959 self._label_graph.add_edge(job.label, child) 

960 

961 def add_job_relationships(self, parent_labels, children_labels): 

962 """Add dependencies between parent and child job labels. 

963 All parents will be connected to all children. 

964 

965 Parameters 

966 ---------- 

967 parent_labels : `list` [`str`] 

968 Parent job labels. 

969 children_labels : `list` [`str`] 

970 Children job labels. 

971 """ 

972 if parent_labels is not None and children_labels is not None: 

973 # Since labels, must ensure not adding edge from label to itself. 

974 edges = [ 

975 e 

976 for e in itertools.product(ensure_iterable(parent_labels), ensure_iterable(children_labels)) 

977 if e[0] != e[1] 

978 ] 

979 

980 self._label_graph.add_edges_from(edges) 

981 

982 def del_job(self, job): 

983 """Delete job and its label from job labels. 

984 

985 Parameters 

986 ---------- 

987 job : `lsst.ctrl.bps.GenericWorkflowJob` 

988 The job to delete from the job labels. 

989 """ 

990 self._label_to_jobs[job.label].remove(job) 

991 # Don't leave keys around if removed last job 

992 if not self._label_to_jobs[job.label]: 

993 del self._label_to_jobs[job.label] 

994 

995 parents = self._label_graph.predecessors(job.label) 

996 children = self._label_graph.successors(job.label) 

997 self._label_graph.remove_node(job.label) 

998 self._label_graph.add_edges_from( 

999 itertools.product(ensure_iterable(parents), ensure_iterable(children)) 

1000 )