Coverage for python / lsst / ctrl / bps / generic_workflow.py: 28%
553 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-19 01:16 -0800
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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <https://www.gnu.org/licenses/>.
28"""Class definitions for a Generic Workflow Graph."""
30__all__ = [
31 "GenericWorkflow",
32 "GenericWorkflowExec",
33 "GenericWorkflowFile",
34 "GenericWorkflowGroup",
35 "GenericWorkflowJob",
36 "GenericWorkflowNode",
37 "GenericWorkflowNodeType",
38 "GenericWorkflowNoopJob",
39]
42import dataclasses
43import itertools
44import logging
45import pickle
46from collections import Counter, defaultdict
47from collections.abc import Iterable, Iterator
48from enum import IntEnum, auto
49from typing import IO, Any, BinaryIO, Literal, cast, overload
51from networkx import DiGraph, topological_sort
52from networkx.algorithms.dag import is_directed_acyclic_graph
54from lsst.utils.iteration import ensure_iterable
56from .bps_draw import draw_networkx_dot
57from .bps_utils import subset_dimension_values
59_LOG = logging.getLogger(__name__)
62@dataclasses.dataclass(slots=True)
63class GenericWorkflowFile:
64 """Information about a file that may be needed by various workflow
65 management services.
66 """
68 name: str
69 """Lookup key (logical file name) of file/directory. Must be unique
70 within run.
71 """
73 src_uri: str | None = None # don't know that need ResourcePath
74 """Original location of file/directory.
75 """
77 wms_transfer: bool = False
78 """Whether the WMS should ignore file or not. Default is False.
79 """
81 job_access_remote: bool = False
82 """Whether the job can remotely access file (using separately specified
83 file access protocols). Default is False.
84 """
86 job_shared: bool = False
87 """Whether job requires its own copy of this file. Default is False.
88 """
90 def __hash__(self) -> int:
91 return hash(self.name)
94@dataclasses.dataclass(slots=True)
95class GenericWorkflowExec:
96 """Information about an executable that may be needed by various workflow
97 management services.
98 """
100 name: str
101 """Lookup key (logical file name) of executable. Must be unique
102 within run.
103 """
105 src_uri: str | None = None # don't know that need ResourcePath
106 """Original location of executable.
107 """
109 transfer_executable: bool = False
110 """Whether the WMS/plugin is responsible for staging executable to
111 location usable by job.
112 """
114 def __hash__(self) -> int:
115 return hash(self.name)
118class GenericWorkflowNodeType(IntEnum):
119 """Type of valid types for nodes in the GenericWorkflow."""
121 NOOP = auto()
122 """Does nothing, but enforces special dependencies."""
124 PAYLOAD = auto()
125 """Typical workflow job."""
127 GROUP = auto()
128 """A special group (subdag) of jobs."""
131@dataclasses.dataclass(slots=True)
132class GenericWorkflowNode:
133 """Base class for nodes in the GenericWorkflow."""
135 name: str
136 """Name of node. Must be unique within workflow."""
138 label: str
139 """"Primary user-facing label for job. Does not need to be unique and
140 may be used for summary reports or to group nodes."""
142 def __hash__(self) -> int:
143 return hash(self.name)
145 @property
146 def node_type(self) -> GenericWorkflowNodeType:
147 """Type of node."""
148 raise NotImplementedError(f"{type(self).__name__} needs to override node_type.")
151@dataclasses.dataclass(slots=True)
152class GenericWorkflowNoopJob(GenericWorkflowNode):
153 """Job that does no work. Used for special dependencies."""
155 @property
156 def node_type(self) -> GenericWorkflowNodeType:
157 """Indicate this is a noop job."""
158 return GenericWorkflowNodeType.NOOP
161@dataclasses.dataclass(slots=True)
162class GenericWorkflowJob(GenericWorkflowNode):
163 """Information about a job that may be needed by various workflow
164 management services.
165 """
167 quanta_counts: Counter[str] = dataclasses.field(default_factory=Counter)
168 """Counts of quanta per task label in job.
169 """
171 tags: dict[str, Any] = dataclasses.field(default_factory=dict)
172 """Other key/value pairs for job that user may want to use as a filter.
173 """
175 executable: GenericWorkflowExec | None = None
176 """Executable for job.
177 """
179 arguments: str | None = None
180 """Command line arguments for job.
181 """
183 cmdvals: dict[str, Any] = dataclasses.field(default_factory=dict)
184 """Values for variables in cmdline when using lazy command line creation.
185 """
187 memory_multiplier: float | None = None
188 """Memory growth rate between retries.
189 """
191 request_memory: int | None = None # MB
192 """Max memory (in MB) that the job is expected to need.
193 """
195 request_memory_max: int | None = None # MB
196 """Max memory (in MB) that the job should ever use.
197 """
199 request_cpus: int | None = None # cores
200 """Max number of cpus that the job is expected to need.
201 """
203 request_disk: int | None = None # MB
204 """Max amount of job scratch disk (in MB) that the job is expected to need.
205 """
207 request_walltime: str | None = None # minutes
208 """Max amount of time (in seconds) that the job is expected to need.
209 """
211 compute_site: str | None = None
212 """Key to look up site-specific information for running the job.
213 """
215 accounting_group: str | None = None
216 """Name of the accounting group to use.
217 """
219 accounting_user: str | None = None
220 """Name of the user to use for accounting purposes.
221 """
223 mail_to: str | None = None
224 """Comma separated list of email addresses for emailing job status.
225 """
227 when_to_mail: str | None = None
228 """WMS-specific terminology for when to email job status.
229 """
231 number_of_retries: int | None = None
232 """Number of times to automatically retry a failed job.
233 """
235 retry_unless_exit: int | list[int] | None = None
236 """Exit code(s) for job that means to not automatically retry.
237 """
239 abort_on_value: int | None = None
240 """Job exit value for signals to abort the entire workflow.
241 """
243 abort_return_value: int | None = None
244 """Exit value to use when aborting the entire workflow.
245 """
247 priority: str | None = None
248 """Initial priority of job in WMS-format.
249 """
251 category: str | None = None
252 """WMS-facing label of job within single workflow (e.g., can be used for
253 throttling jobs within a single workflow).
254 """
256 concurrency_limit: str | None = None
257 """Names of concurrency limits that the WMS plugin can appropriately
258 translate to limit the number of this job across all running workflows.
259 """
261 queue: str | None = None
262 """Name of queue to use. Different WMS can translate this concept
263 differently.
264 """
266 pre_cmdline: str | None = None
267 """Command line to be executed prior to executing job.
268 """
270 post_cmdline: str | None = None
271 """Command line to be executed after job executes.
273 Should be executed regardless of exit status.
274 """
276 preemptible: bool | None = None
277 """The flag indicating whether the job can be preempted.
278 """
280 profile: dict[str, Any] = dataclasses.field(default_factory=dict)
281 """Nested dictionary of WMS-specific key/value pairs with primary key being
282 WMS key (e.g., pegasus, condor, panda).
283 """
285 attrs: dict[str, Any] = dataclasses.field(default_factory=dict)
286 """Key/value pairs of job attributes (for WMS that have attributes in
287 addition to commands).
288 """
290 environment: dict[str, Any] = dataclasses.field(default_factory=dict)
291 """Environment variable names and values to be explicitly set inside job.
292 """
294 compute_cloud: str | None = None
295 """Key to look up cloud-specific information for running the job.
296 """
298 @property
299 def node_type(self) -> GenericWorkflowNodeType:
300 """Indicate this is a payload job."""
301 return GenericWorkflowNodeType.PAYLOAD
304class GenericWorkflow(DiGraph):
305 """A generic representation of a workflow used to submit to specific
306 workflow management systems.
308 Parameters
309 ----------
310 name : `str`
311 Name of generic workflow.
312 incoming_graph_data : `~typing.Any`, optional
313 Data used to initialized graph that is passed through to DiGraph
314 constructor. Can be any type supported by networkx.DiGraph.
315 **attr : `dict`
316 Keyword arguments passed through to DiGraph constructor.
317 """
319 def __init__(self, name: str, incoming_graph_data: Any | None = None, **attr: Any) -> None:
320 super().__init__(incoming_graph_data, **attr)
321 self._name = name
322 self.run_attrs: dict[str, str] = {}
323 self._job_labels = GenericWorkflowLabels()
324 self._files: dict[str, GenericWorkflowFile] = {}
325 self._executables: dict[str, GenericWorkflowExec] = {}
326 self._inputs: dict[
327 str, list[GenericWorkflowFile]
328 ] = {} # mapping job.names to list of GenericWorkflowFile
329 self._outputs: dict[
330 str, list[GenericWorkflowFile]
331 ] = {} # mapping job.names to list of GenericWorkflowFile
332 self.run_id = None
333 self._final: GenericWorkflowJob | GenericWorkflow | None = None
335 @property
336 def name(self) -> str:
337 """Retrieve name of generic workflow.
339 Returns
340 -------
341 name : `str`
342 Name of generic workflow.
343 """
344 return self._name
346 @property
347 def quanta_counts(self) -> Counter[str]:
348 """Count of quanta per task label (`collections.Counter`)."""
349 qcounts: Counter[str] = Counter()
350 for job_name in self:
351 gwjob = self.get_job(job_name)
352 if hasattr(gwjob, "quanta_counts"):
353 qcounts += gwjob.quanta_counts
354 return qcounts
356 @property
357 def labels(self) -> list[str]:
358 """Job labels (`list` [`str`], read-only)."""
359 return self._job_labels.labels
361 def regenerate_labels(self) -> None:
362 """Regenerate the list of job labels."""
363 self._job_labels = GenericWorkflowLabels()
364 for job_name in self:
365 job = self.get_job(job_name)
366 if job.node_type == GenericWorkflowNodeType.PAYLOAD:
367 job = cast(GenericWorkflowJob, job)
368 parents_labels: list[str] = []
369 for parent_name in self.predecessors(job.name):
370 parent_job = self.get_job(parent_name)
371 if parent_job.node_type == GenericWorkflowNodeType.PAYLOAD:
372 # parent_job = cast(GenericWorkflowJob, parent_job)
373 parents_labels.append(parent_job.label)
374 children_labels: list[str] = []
375 for child_name in self.successors(job.name):
376 child_job = self.get_job(child_name)
377 if child_job.node_type == GenericWorkflowNodeType.PAYLOAD:
378 # child_job = cast(GenericWorkflowJob, child_job)
379 children_labels.append(child_job.label)
380 self._job_labels.add_job(job, parents_labels, children_labels)
382 @property
383 def job_counts(self) -> Counter[str]:
384 """Count of jobs per job label (`collections.Counter`)."""
385 jcounts = self._job_labels.job_counts
387 # Final is separate
388 final = self.get_final()
389 if final:
390 if isinstance(final, GenericWorkflow):
391 jcounts.update(final.job_counts)
392 else:
393 jcounts[final.label] += 1
395 return jcounts
397 def __iter__(self) -> Iterator[str]:
398 """Return iterator of job names in topologically sorted order."""
399 return topological_sort(self)
401 @overload
402 def get_files(self, data: Literal[False], transfer_only: bool = True) -> list[str]: ... 402 ↛ exitline 402 didn't return from function 'get_files' because
404 @overload
405 def get_files(self, data: Literal[True], transfer_only: bool = True) -> list[GenericWorkflowFile]: ... 405 ↛ exitline 405 didn't return from function 'get_files' because
407 def get_files(
408 self, data: bool = False, transfer_only: bool = True
409 ) -> list[GenericWorkflowFile] | list[str]:
410 """Retrieve files from generic workflow.
412 Need API in case change way files are stored (e.g., make
413 workflow a bipartite graph with jobs and files nodes).
415 Parameters
416 ----------
417 data : `bool`, optional
418 Whether to return the file data as well as the file object name
419 (The default is False).
420 transfer_only : `bool`, optional
421 Whether to only return files for which a workflow management system
422 would be responsible for transferring.
424 Returns
425 -------
426 files : `list` [`lsst.ctrl.bps.GenericWorkflowFile`] or `list` [`str`]
427 File names or objects from generic workflow meeting specifications.
428 """
429 files: list[Any] = [] # Any for mypy to allow the different append lines.
430 for filename, file in self._files.items():
431 if not transfer_only or file.wms_transfer:
432 if not data:
433 files.append(filename)
434 else:
435 files.append(file)
436 return files
438 def add_job(
439 self,
440 job: GenericWorkflowNode,
441 parent_names: str | list[str] | None = None,
442 child_names: str | list[str] | None = None,
443 ) -> None:
444 """Add job to generic workflow.
446 Parameters
447 ----------
448 job : `lsst.ctrl.bps.GenericWorkflowNode`
449 Job to add to the generic workflow.
450 parent_names : `str` | `list` [`str`], optional
451 Names of jobs that are parents of given job.
452 child_names : `str` | `list` [`str`], optional
453 Names of jobs that are children of given job.
454 """
455 _LOG.debug("job: %s (%s)", job.name, job.label)
456 _LOG.debug("parent_names: %s", parent_names)
457 _LOG.debug("child_names: %s", child_names)
458 if not isinstance(job, GenericWorkflowNode):
459 raise RuntimeError(f"Invalid type for job to be added to GenericWorkflowGraph ({type(job)}).")
460 if self.has_node(job.name):
461 raise RuntimeError(f"Job {job.name} already exists in GenericWorkflowGraph.")
462 super().add_node(job.name, job=job)
463 self.add_job_relationships(parent_names, job.name)
464 self.add_job_relationships(job.name, child_names)
465 if job.node_type == GenericWorkflowNodeType.PAYLOAD:
466 job = cast(GenericWorkflowJob, job)
467 self.add_executable(job.executable)
468 self._job_labels.add_job(
469 job,
470 [self.get_job(p).label for p in self.predecessors(job.name)],
471 [self.get_job(p).label for p in self.successors(job.name)],
472 )
474 def add_node(self, node_for_adding: GenericWorkflowNode, **attr: Any) -> None:
475 """Override networkx function to call more specific add_job function.
477 Parameters
478 ----------
479 node_for_adding : `lsst.ctrl.bps.GenericWorkflowJob`
480 Job to be added to generic workflow.
481 **attr
482 Needed to match original networkx function, but not used.
483 """
484 self.add_job(node_for_adding)
486 def add_job_relationships(
487 self, parents: str | list[str] | None, children: str | list[str] | None
488 ) -> None:
489 """Add dependencies between parent and child jobs. All parents will
490 be connected to all children.
492 Parameters
493 ----------
494 parents : `str` or `list` [`str`], optional
495 Parent job names.
496 children : `str` or `list` [`str`], optional
497 Children job names.
498 """
499 # Allow this to be a noop if no parents or no children
500 if parents is not None and children is not None:
501 self.add_edges_from(itertools.product(ensure_iterable(parents), ensure_iterable(children)))
502 self._job_labels.add_job_relationships(
503 [self.get_job(n).label for n in ensure_iterable(parents)],
504 [self.get_job(n).label for n in ensure_iterable(children)],
505 )
507 def add_edges_from(self, ebunch_to_add: Iterable[tuple[str, str]], **attr: Any) -> None:
508 """Add several edges between jobs in the generic workflow.
510 Parameters
511 ----------
512 ebunch_to_add : Iterable [`tuple` [`str`, `str`]]
513 Iterable of job name pairs between which a dependency should be
514 saved.
515 **attr : keyword arguments, optional
516 Data can be assigned using keyword arguments (not currently used).
517 """
518 for edge_to_add in ebunch_to_add:
519 self.add_edge(edge_to_add[0], edge_to_add[1], **attr)
521 def add_edge(self, u_of_edge: str, v_of_edge: str, **attr: Any) -> None:
522 """Add edge connecting jobs in workflow.
524 Parameters
525 ----------
526 u_of_edge : `str`
527 Name of parent job.
528 v_of_edge : `str`
529 Name of child job.
530 **attr : keyword arguments, optional
531 Attributes to save with edge.
532 """
533 if u_of_edge not in self:
534 raise RuntimeError(f"{u_of_edge} not in GenericWorkflow")
535 if v_of_edge not in self:
536 raise RuntimeError(f"{v_of_edge} not in GenericWorkflow")
537 super().add_edge(u_of_edge, v_of_edge, **attr)
539 def get_job(self, job_name: str) -> GenericWorkflowNode:
540 """Retrieve job by name from workflow.
542 Parameters
543 ----------
544 job_name : `str`
545 Name of job to retrieve.
547 Returns
548 -------
549 job : `lsst.ctrl.bps.GenericWorkflowNode`
550 Job matching given job_name.
551 """
552 return self.nodes[job_name]["job"]
554 def del_job(self, job_name: str) -> None:
555 """Delete job from generic workflow leaving connected graph.
557 Parameters
558 ----------
559 job_name : `str`
560 Name of job to delete from workflow.
561 """
562 job = self.get_job(job_name)
564 # Remove from job labels
565 if isinstance(job, GenericWorkflowJob):
566 self._job_labels.del_job(job)
568 # Connect all parent jobs to all children jobs.
569 parents = list(self.predecessors(job_name))
570 children = list(self.successors(job_name))
571 self.add_job_relationships(parents, children)
573 # Delete job node (which deletes edges).
574 self.remove_node(job_name)
576 def add_job_inputs(self, job_name: str, files: GenericWorkflowFile | list[GenericWorkflowFile]) -> None:
577 """Add files as inputs to specified job.
579 Parameters
580 ----------
581 job_name : `str`
582 Name of job to which inputs should be added.
583 files : `lsst.ctrl.bps.GenericWorkflowFile` or \
584 `list` [`lsst.ctrl.bps.GenericWorkflowFile`]
585 File object(s) to be added as inputs to the specified job.
586 """
587 self._inputs.setdefault(job_name, [])
588 for file in ensure_iterable(files):
589 # Save the central copy
590 if file.name not in self._files:
591 self._files[file.name] = file
593 # Save the job reference to the file
594 self._inputs[job_name].append(file)
596 def get_file(self, name: str) -> GenericWorkflowFile:
597 """Retrieve a file object by name.
599 Parameters
600 ----------
601 name : `str`
602 Name of file object.
604 Returns
605 -------
606 gwfile : `lsst.ctrl.bps.GenericWorkflowFile`
607 File matching given name.
608 """
609 return self._files[name]
611 def add_file(self, gwfile: GenericWorkflowFile) -> None:
612 """Add file object.
614 Parameters
615 ----------
616 gwfile : `lsst.ctrl.bps.GenericWorkflowFile`
617 File object to add to workflow.
618 """
619 if gwfile.name not in self._files:
620 self._files[gwfile.name] = gwfile
621 else:
622 _LOG.debug("Skipped add_file for existing file %s", gwfile.name)
624 @overload
625 def get_job_inputs( 625 ↛ exitline 625 didn't return from function 'get_job_inputs' because
626 self, job_name: str, data: Literal[False], transfer_only: bool = False
627 ) -> list[str]: ...
629 @overload
630 def get_job_inputs( 630 ↛ exitline 630 didn't return from function 'get_job_inputs' because
631 self, job_name: str, data: Literal[True], transfer_only: bool = False
632 ) -> list[GenericWorkflowFile]: ...
634 def get_job_inputs(
635 self, job_name: str, data: bool = True, transfer_only: bool = False
636 ) -> list[GenericWorkflowFile] | list[str]:
637 """Return the input files for the given job.
639 Parameters
640 ----------
641 job_name : `str`
642 Name of the job.
643 data : `bool`, optional
644 Whether to return the file data as well as the file object name.
645 transfer_only : `bool`, optional
646 Whether to only return files for which a workflow management system
647 would be responsible for transferring.
649 Returns
650 -------
651 inputs : `list` [`lsst.ctrl.bps.GenericWorkflowFile`] or `list` [`str`]
652 Input files for the given job. If no input files for the job,
653 returns an empty list.
654 """
655 inputs: list[Any] = [] # Any for mypy to allow the different append lines.
656 if job_name in self._inputs:
657 for gwfile in self._inputs[job_name]:
658 if not transfer_only or gwfile.wms_transfer:
659 if not data:
660 inputs.append(gwfile.name)
661 else:
662 inputs.append(gwfile)
663 return inputs
665 def add_job_outputs(self, job_name: str, files: list[GenericWorkflowFile]) -> None:
666 """Add output files to a job.
668 Parameters
669 ----------
670 job_name : `str`
671 Name of job to which the files should be added as outputs.
672 files : `list` [`lsst.ctrl.bps.GenericWorkflowFile`]
673 File objects to be added as outputs for specified job.
674 """
675 self._outputs.setdefault(job_name, [])
677 for file_ in ensure_iterable(files):
678 # Save the central copy
679 if file_.name not in self._files:
680 self._files[file_.name] = file_
682 # Save the job reference to the file
683 self._outputs[job_name].append(file_)
685 @overload
686 def get_job_outputs( 686 ↛ exitline 686 didn't return from function 'get_job_outputs' because
687 self, job_name: str, data: Literal[False], transfer_only: bool = False
688 ) -> list[str]: ...
690 @overload
691 def get_job_outputs( 691 ↛ exitline 691 didn't return from function 'get_job_outputs' because
692 self, job_name: str, data: Literal[True], transfer_only: bool = False
693 ) -> list[GenericWorkflowFile]: ...
695 def get_job_outputs(
696 self, job_name: str, data: bool = True, transfer_only: bool = False
697 ) -> list[GenericWorkflowFile] | list[str]:
698 """Return the output files for the given job.
700 Parameters
701 ----------
702 job_name : `str`
703 Name of the job.
704 data : `bool`
705 Whether to return the file data as well as the file object name.
706 It defaults to `True` thus returning file data as well.
707 transfer_only : `bool`
708 Whether to only return files for which a workflow management system
709 would be responsible for transferring. It defaults to `False` thus
710 returning all output files.
712 Returns
713 -------
714 outputs : `list` [`lsst.ctrl.bps.GenericWorkflowFile`] or \
715 `list` [`str`]
716 Output files for the given job. If no output files for the job,
717 returns an empty list.
718 """
719 outputs: list[Any] = [] # Any for mypy to allow the different append lines.
720 if not data:
721 outputs = cast(list[str], outputs)
722 else:
723 outputs = cast(list[GenericWorkflowFile], outputs)
725 if job_name in self._outputs:
726 for gwfile in self._outputs[job_name]:
727 if not transfer_only or gwfile.wms_transfer:
728 if not data:
729 outputs.append(gwfile.name)
730 else:
731 outputs.append(gwfile)
732 return outputs
734 def draw(self, stream: str | IO[str], format_: str = "dot") -> None:
735 """Output generic workflow in a visualization format.
737 Parameters
738 ----------
739 stream : `str` or `io.BufferedIOBase`
740 Stream to which the visualization should be written.
741 format_ : `str`, optional
742 Which visualization format to use. It defaults to the format for
743 the dot program.
744 """
745 draw_funcs = {"dot": draw_networkx_dot}
746 if format_ in draw_funcs:
747 draw_funcs[format_](self, stream)
748 else:
749 raise RuntimeError(f"Unknown draw format ({format_})")
751 def save(self, stream: str | IO[bytes], format_: str = "pickle") -> None:
752 """Save the generic workflow in a format that is loadable.
754 Parameters
755 ----------
756 stream : `str` or `io.BufferedIOBase`
757 Stream to pass to the format-specific writer. Accepts anything
758 that the writer accepts.
759 format_ : `str`, optional
760 Format in which to write the data. It defaults to pickle format.
761 """
762 if format_ == "pickle":
763 stream = cast(BinaryIO, stream)
764 pickle.dump(self, stream)
765 else:
766 raise RuntimeError(f"Unknown format ({format_})")
768 @classmethod
769 def load(cls, stream: str | IO[bytes], format_: str = "pickle") -> "GenericWorkflow":
770 """Load a GenericWorkflow from the given stream.
772 Parameters
773 ----------
774 stream : `str` or `io.BufferedIOBase`
775 Stream to pass to the format-specific loader. Accepts anything that
776 the loader accepts.
777 format_ : `str`, optional
778 Format of data to expect when loading from stream. It defaults
779 to pickle format.
781 Returns
782 -------
783 generic_workflow : `lsst.ctrl.bps.GenericWorkflow`
784 Generic workflow loaded from the given stream.
785 """
786 if format_ == "pickle":
787 stream = cast(BinaryIO, stream)
788 object_ = pickle.load(stream)
789 assert isinstance(object_, GenericWorkflow) # for mypy
790 return object_
792 raise RuntimeError(f"Unknown format ({format_})")
794 def validate(self) -> None:
795 """Run checks to ensure that the generic workflow graph is valid."""
796 # Make sure a directed acyclic graph
797 assert is_directed_acyclic_graph(self)
799 def add_workflow_source(self, workflow: "GenericWorkflow") -> None:
800 """Add given workflow as new source to this workflow.
802 Parameters
803 ----------
804 workflow : `lsst.ctrl.bps.GenericWorkflow`
805 The given workflow.
806 """
807 # Find source nodes in self.
808 self_sources = [n for n in self if self.in_degree(n) == 0]
809 _LOG.debug("self_sources = %s", self_sources)
811 # Find sink nodes of workflow.
812 new_sinks = [n for n in workflow if workflow.out_degree(n) == 0]
813 _LOG.debug("new sinks = %s", new_sinks)
815 # Add new workflow nodes to self graph and make new edges.
816 self.add_nodes_from(workflow.nodes(data=True))
817 self.add_edges_from(workflow.edges())
818 for source in self_sources:
819 for sink in new_sinks:
820 self.add_edge(sink, source)
822 # Add separately stored info
823 for job_name in workflow:
824 job = self.get_job(job_name)
825 # Add job labels
826 if isinstance(job, GenericWorkflowJob):
827 self._job_labels.add_job(
828 job,
829 [self.get_job(p).label for p in self.predecessors(job.name)],
830 [self.get_job(p).label for p in self.successors(job.name)],
831 )
832 # Executables are stored separately so copy them.
833 self.add_executable(job.executable)
835 # Files are stored separately so copy them.
836 self.add_job_inputs(job_name, workflow.get_job_inputs(job_name, data=True))
837 self.add_job_outputs(job_name, workflow.get_job_outputs(job_name, data=True))
839 def add_final(self, final: "GenericWorkflowJob | GenericWorkflow") -> None:
840 """Add special final job/workflow to the generic workflow.
842 Parameters
843 ----------
844 final : `lsst.ctrl.bps.GenericWorkflowJob` or \
845 `lsst.ctrl.bps.GenericWorkflow`
846 Information needed to execute the special final job(s), the
847 job(s) to be executed after all jobs that can be executed
848 have been executed regardless of exit status of any of the
849 jobs.
850 """
851 if not isinstance(final, GenericWorkflowJob) and not isinstance(final, GenericWorkflow):
852 raise TypeError("Invalid type for GenericWorkflow final ({type(final)})")
854 self._final = final
855 if isinstance(final, GenericWorkflowJob):
856 self.add_executable(final.executable)
858 def get_final(self) -> "GenericWorkflowJob | GenericWorkflow | None":
859 """Return job/workflow to be executed after all jobs that can be
860 executed have been executed regardless of exit status of any of
861 the jobs.
863 Returns
864 -------
865 final : `lsst.ctrl.bps.GenericWorkflowJob` or \
866 `lsst.ctrl.bps.GenericWorkflow`
867 Information needed to execute final job(s).
868 """
869 return self._final
871 def add_executable(self, executable: GenericWorkflowExec | None) -> None:
872 """Add executable to workflow's list of executables.
874 Parameters
875 ----------
876 executable : `lsst.ctrl.bps.GenericWorkflowExec`
877 Executable object to be added to workflow.
878 """
879 if executable is not None:
880 self._executables[executable.name] = executable
881 else:
882 _LOG.warning("executable not specified (None); cannot add to the workflow's list of executables")
884 @overload
885 def get_executables(self, data: Literal[False], transfer_only: bool = False) -> list[str]: ... 885 ↛ exitline 885 didn't return from function 'get_executables' because
887 @overload
888 def get_executables( 888 ↛ exitline 888 didn't return from function 'get_executables' because
889 self, data: Literal[True], transfer_only: bool = False
890 ) -> list[GenericWorkflowExec]: ...
892 def get_executables(
893 self, data: bool = False, transfer_only: bool = True
894 ) -> list[GenericWorkflowExec] | list[str]:
895 """Retrieve executables from generic workflow.
897 Parameters
898 ----------
899 data : `bool`, optional
900 Whether to return the executable data as well as the exec object
901 name (The defaults is False).
902 transfer_only : `bool`, optional
903 Whether to only return executables for which transfer_executable
904 is True.
906 Returns
907 -------
908 execs : `list` [`lsst.ctrl.bps.GenericWorkflowExec`] or `list` [`str`]
909 Filtered executable names or objects from generic workflow.
910 """
911 execs: list[Any] = [] # This and cast lines for mypy
912 if not data:
913 execs = cast(list[str], execs)
914 else:
915 execs = cast(list[GenericWorkflowExec], execs)
917 for name, executable in self._executables.items():
918 if not transfer_only or executable.transfer_executable:
919 if not data:
920 execs.append(name)
921 else:
922 execs.append(executable)
923 return execs
925 def get_jobs_by_label(self, label: str) -> list[GenericWorkflowJob]:
926 """Retrieve jobs by label from workflow.
928 Parameters
929 ----------
930 label : `str`
931 Label of jobs to retrieve.
933 Returns
934 -------
935 jobs : list[`lsst.ctrl.bps.GenericWorkflowNode`]
936 Jobs having given label.
937 """
938 return self._job_labels.get_jobs_by_label(label)
940 def _check_job_ordering_config(self, ordering_config: dict[str, Any]) -> dict[str, DiGraph]:
941 """Check configuration related to job ordering.
943 Parameters
944 ----------
945 ordering_config : `dict` [`str`, `~typing.Any`]
946 Job ordering configuration to check.
948 Returns
949 -------
950 group_to_label_subgraph : `dict` [`str`, `network.DiGraph`]
951 Mapping of group name to a graph of the job labels in the group.
952 """
953 group_to_label_subgraph = {}
954 job_label_to_group: dict[str, str] = {} # Checking label appears only in one group
955 for group, group_vals in ordering_config.items():
956 implementation = group_vals.get("implementation", "group")
957 if implementation not in ["noop", "group"]:
958 raise RuntimeError(f"Invalid implementation for {group}: {implementation}")
959 ordering_type = group_vals.get("ordering_type", "sort")
960 if ordering_type != "sort":
961 raise RuntimeError(f"Invalid ordering_type for {group}: {ordering_type}")
962 if "dimensions" not in group_vals:
963 raise KeyError(f"Missing dimensions entry in ordering group {group}")
965 job_labels = [x.strip() for x in group_vals["labels"].split(",")]
966 _LOG.debug("group %s: job_labels=%s", group, job_labels)
967 unused_labels = []
968 for job_label in job_labels:
969 if job_label not in self._job_labels.labels:
970 unused_labels.append(job_label)
971 elif job_label in job_label_to_group:
972 raise RuntimeError(
973 f"Job label {job_label} appears in more than one job ordering group "
974 f"({group} {job_label_to_group[job_label]})"
975 )
976 else:
977 job_label_to_group[job_label] = group
979 if unused_labels:
980 _LOG.info("Workflow job labels = %s", ",".join(self._job_labels.labels))
981 raise RuntimeError(
982 f"Job label(s) ({','.join(unused_labels)}) from job ordering group "
983 f"{group} does not exist in workflow. Aborting."
984 )
986 label_subgraph = self._job_labels.subgraph(job_labels)
987 group_to_label_subgraph[group] = label_subgraph
989 return group_to_label_subgraph
991 def _group_jobs_by_values(
992 self, group: str, group_config: dict[str, Any], label_subgraph: DiGraph
993 ) -> dict[tuple[Any, ...], list[str]]:
994 """Create job mapping of special sortable dimension key
995 to job name by comparing dimension values.
997 Parameters
998 ----------
999 group : `str`
1000 Name of group for which creating mapping.
1001 group_config : `dict` [`str`, `~typing.Any`]
1002 Config for group for which creating mapping.
1003 label_subgraph : `networkx.DiGraph`
1004 The graph of job labels to be used in mapping.
1006 Returns
1007 -------
1008 dims_to_jobs : `dict` [`tuple` [`~typing.Any`, ...], `list` [`str`]]
1009 Mapping of dimensions to job names.
1010 """
1011 dims_to_jobs: dict[tuple[Any, ...], list[str]] = {}
1012 for job_label in label_subgraph:
1013 jobs = self.get_jobs_by_label(job_label)
1014 for job in jobs:
1015 job_dim_values = subset_dimension_values(
1016 f"Job {job.name}",
1017 f"order group {group}",
1018 group_config["dimensions"],
1019 job.tags,
1020 group_config.get("equalDimensions", None),
1021 )
1022 job_dims = []
1023 for dim in [d.strip() for d in group_config["dimensions"].split(",")]:
1024 job_dims.append(job_dim_values[dim])
1025 dims_job_list = dims_to_jobs.setdefault(tuple(job_dims), [])
1026 dims_job_list.append(job.name)
1027 return dims_to_jobs
1029 def _group_jobs_by_dependencies(
1030 self, group: str, group_config: dict[str, Any], label_subgraph: DiGraph
1031 ) -> dict[tuple[Any, ...], list[str]]:
1032 """Create job mapping of special sortable dimension key
1033 to job name by following dependencies.
1035 Parameters
1036 ----------
1037 group : `str`
1038 Name of group for which creating mapping.
1039 group_config : `dict` [`str`, `~typing.Any`]
1040 Config for group for which creating mapping.
1041 label_subgraph : `networkx.DiGraph`
1042 The graph of job labels to be used in mapping.
1044 Returns
1045 -------
1046 dims_to_jobs : `dict` [`tuple` [`~typing.Any`, ...], `list` [`str`]]
1047 Mapping of dimensions to job names.
1048 """
1049 method = group_config["findDependencyMethod"]
1050 dim_labels = list(topological_sort(label_subgraph))
1051 match method:
1052 case "source":
1053 find_potential_jobs = self.successors
1054 case "sink":
1055 find_potential_jobs = self.predecessors
1056 dim_labels.reverse()
1057 case _:
1058 raise RuntimeError(f"Invalid findDependencyMethod ({method})")
1060 jobs_seen: set[str] = set()
1061 dims_to_jobs: dict[tuple[Any, ...], list[str]] = {}
1062 for label in dim_labels:
1063 jobs = self.get_jobs_by_label(label)
1064 for job in jobs:
1065 if job.name not in jobs_seen:
1066 jobs_seen.add(job.name)
1067 job_dim_values = subset_dimension_values(
1068 f"Job {job.name}",
1069 f"order group {group}",
1070 group_config["dimensions"],
1071 job.tags,
1072 group_config.get("equalDimensions", None),
1073 )
1074 job_dims = []
1075 for dim in [d.strip() for d in group_config["dimensions"].split(",")]:
1076 job_dims.append(job_dim_values[dim])
1077 dims_job_list = dims_to_jobs.setdefault(tuple(job_dims), [])
1078 dims_job_list.append(job.name)
1079 # Use dependencies to find other quantum to add
1080 # Note: in testing, using the following code was faster
1081 # than using networkx descendants and ancestors functions
1082 # While traversing the subgraph, nodes may appear
1083 # repeatedly in potential_jobs.
1084 jobs_to_use = [job]
1085 while jobs_to_use:
1086 job_to_use = jobs_to_use.pop()
1088 potential_job_names = find_potential_jobs(job_to_use.name)
1089 for potential_job_name in potential_job_names:
1090 potential_job = cast(GenericWorkflowJob, self.get_job(potential_job_name))
1091 if potential_job.label in label_subgraph:
1092 if potential_job.name not in dims_job_list:
1093 _LOG.debug(
1094 "Adding potential job %s (%s) to group %s",
1095 potential_job.name,
1096 potential_job.label,
1097 group,
1098 )
1099 dims_job_list.append(potential_job.name)
1100 jobs_to_use.append(potential_job)
1101 jobs_seen.add(potential_job.name)
1102 else:
1103 _LOG.debug(
1104 "label (%s) not in ordered_tasks. Not adding potential quantum %s",
1105 potential_job.label,
1106 potential_job.name,
1107 )
1108 return dims_to_jobs
1110 def _update_by_group_sort(
1111 self,
1112 group: str,
1113 dims_to_jobs: dict[tuple[Any, ...], list[str]],
1114 blocking: bool = False,
1115 ) -> None:
1116 """Update portion of workflow for special job ordering using sort.
1118 Parameters
1119 ----------
1120 group : `str`
1121 Ordering group label used for job name and messages.
1122 dims_to_jobs: `dict` [`tuple` [`~typing.Any`, ...], `list` [`str`]]
1123 Mapping of special dimension keys to workflow job names.
1124 The sort for the special ordering is over the keys.
1125 blocking: `bool`
1126 Whether a failure in a group blocks execution of remaining groups.
1127 """
1128 group_job_names: list[str] = []
1129 for dim_key, job_list in sorted(dims_to_jobs.items()):
1130 _LOG.debug("group %s: dim_key=%s", group, dim_key)
1131 group_job_name = f"group_{group}_{'=='.join([str(dk) for dk in dim_key])}"
1132 prev_name = group_job_names[-1] if group_job_names else None
1133 group_job_names.append(group_job_name)
1135 self._replace_subgraph_with_job_group(group_job_name, group, job_list, blocking)
1137 # ordering between groups
1138 if prev_name:
1139 self.add_edge(prev_name, group_job_name)
1141 def _replace_subgraph_with_job_group(
1142 self, group_name: str, group_label: str, job_list: list[str], blocking: bool
1143 ) -> None:
1144 """Update portion of workflow for special job ordering using groups.
1146 Parameters
1147 ----------
1148 group_name : `str`
1149 Ordering group name.
1150 group_label : `str`
1151 Ordering group label.
1152 job_list: `list` [`str`]
1153 List of job names to put in the group
1154 blocking: `bool`
1155 Whether a failure in a group blocks execution of remaining groups.
1156 """
1157 job_group = GenericWorkflowGroup(group_name, group_label, blocking=blocking)
1158 self.add_node(job_group)
1160 # Add jobs, files, and executables first
1161 # then add edges later to avoid order issues
1162 for job_name in job_list:
1163 job = cast(GenericWorkflowJob, self.get_job(job_name))
1164 job_group.add_job(job)
1165 files = self.get_job_inputs(job_name, data=True)
1166 job_group.add_job_inputs(job_name, files)
1167 files = self.get_job_outputs(job_name, data=True)
1168 job_group.add_job_outputs(job_name, files)
1169 job_group.add_executable(job.executable)
1171 # Can't remove edges while looping through edge view,
1172 # so save to remove after loops
1173 edges_to_remove: list[tuple[str, str]] = []
1174 for job_name in job_list:
1175 in_edges = self.in_edges(job_name)
1176 for u, v in in_edges:
1177 if u in job_list:
1178 job_group.add_edge(u, v)
1179 else:
1180 self.add_edge(u, group_name)
1181 edges_to_remove.append((u, v))
1183 out_edges = self.out_edges(job_name)
1184 for u, v in out_edges:
1185 if v in job_list:
1186 job_group.add_edge(u, v)
1187 else:
1188 self.add_edge(group_name, v)
1189 edges_to_remove.append((u, v))
1191 # Remove edges collected earlier
1192 self.remove_edges_from(edges_to_remove)
1194 # Remove nodes from main GenericWorkflow
1195 self.remove_nodes_from(job_list)
1197 def add_special_job_ordering(self, ordering: dict[str, Any]) -> None:
1198 """Add special nodes and dependencies to enforce given ordering.
1200 Parameters
1201 ----------
1202 ordering : `dict` [`str`, `~typing.Any`]
1203 Description of the job ordering to enforce.
1204 """
1205 group_to_label_subgraph = self._check_job_ordering_config(ordering)
1207 for group, group_vals in ordering.items():
1208 if "findDependencyMethod" in group_vals:
1209 job_grouping_func = self._group_jobs_by_dependencies
1210 else:
1211 job_grouping_func = self._group_jobs_by_values
1213 dims = [x.strip() for x in group_vals["dimensions"].split(",")]
1214 _LOG.debug("group %s: dims=%s", group, dims)
1215 job_groups = job_grouping_func(group, group_vals, group_to_label_subgraph[group])
1217 # Update the workflow
1218 implementation = group_vals.get("implementation", "group")
1219 ordering_type = group_vals.get("ordering_type", "sort")
1220 match (implementation, ordering_type):
1221 case ("noop", "sort"):
1222 self._update_by_noop_sort(group, job_groups)
1223 case ("group", "sort"):
1224 blocking = group_vals.get("blocking", False)
1225 self._update_by_group_sort(group, job_groups, blocking)
1226 case _:
1227 raise RuntimeError(
1228 f"Invalid implementation, ordering_type pair for group ({implementation},"
1229 f" {ordering_type})"
1230 )
1232 def _update_by_noop_sort(self, group: str, dims_to_jobs: dict[tuple[Any, ...], list[str]]) -> None:
1233 """Update portion of workflow for special ordering of jobs using sort.
1235 Parameters
1236 ----------
1237 group : `str`
1238 Ordering group label used for job name and messages.
1239 dims_to_jobs: `dict` [`tuple`[`str`,...], `list` [`str`]]
1240 Mapping of special dimension keys to workflow job names.
1241 The sort for the special ordering is over the keys.
1242 """
1243 noop_names: list[str] = []
1244 prev_name: str | None = None
1245 for dim_key, job_list in sorted(dims_to_jobs.items()):
1246 _LOG.debug("group %s: dim_key=%s", group, dim_key)
1247 noop_name = f"noop_{group}_{'=='.join([str(dk) for dk in dim_key])}"
1248 if noop_names:
1249 prev_name = noop_names[-1]
1250 noop_names.append(noop_name)
1252 self._update_single_noop(job_list, noop_name, group, prev_name)
1254 # As implemented, loop adds one last NOOP job with
1255 # 0 children. Remove it as not useful.
1256 assert self.out_degree(noop_names[-1]) == 0
1257 self.remove_node(noop_names[-1])
1259 def _update_single_noop(
1260 self, job_list: list[str], order_node_name: str, order_label: str, prev_order_name: str | None
1261 ) -> None:
1262 """Update the workflow around the single job making special NOOP
1263 jobs when necessary.
1265 Parameters
1266 ----------
1267 job_list : `list` [`str`]
1268 Current jobs involved in special ordering
1269 order_node_name : `str`
1270 Name for the order NOOP job. If it does not exist in workflow, a
1271 new NOOP job with this name will be created and added.
1272 order_label : `str`
1273 Label for the order NOOP job.
1274 prev_order_name: `str` or None
1275 Name of the previous order NOOP job used to add edge as predecessor
1276 of current job.
1277 """
1278 if order_node_name not in self:
1279 _LOG.debug("Adding new ordering node %s", order_node_name)
1280 order_node = GenericWorkflowNoopJob(order_node_name, order_label)
1281 self.add_job(order_node)
1283 subgraph = DiGraph(self).subgraph(job_list)
1284 sinks = [n for n in job_list if subgraph.out_degree(n) == 0]
1285 for job_name in sinks:
1286 self.add_edge(job_name, order_node_name)
1288 if prev_order_name:
1289 sources = [n for n in job_list if subgraph.in_degree(n) == 0]
1290 for job_name in sources:
1291 _LOG.debug("Adding edge %s to %s", prev_order_name, job_name)
1292 self.add_edge(prev_order_name, job_name)
1295@dataclasses.dataclass(slots=True)
1296class GenericWorkflowGroup(GenericWorkflowNode, GenericWorkflow):
1297 """Node representing a group of jobs. Used for special dependencies.
1299 Parameters
1300 ----------
1301 name : `str`
1302 Name of node. Must be unique within workflow.
1303 label : `str`
1304 Primary user-facing label for job. Does not need to be unique and
1305 may be used for summary reports or to group nodes.
1306 blocking : `bool`
1307 Whether a failure inside group prunes executions of remaining groups.
1308 """
1310 blocking: bool = False
1311 """Whether a failure inside group prunes executions of remaining groups."""
1313 @property
1314 def node_type(self) -> GenericWorkflowNodeType:
1315 """Indicate this is a group job."""
1316 return GenericWorkflowNodeType.GROUP
1318 def __init__(self, name: str, label: str, blocking: bool = False) -> None:
1319 """Initialize each parent class."""
1320 _LOG.debug("%s %s %s", name, label, blocking)
1321 GenericWorkflowNode.__init__(self, name, label)
1322 GenericWorkflow.__init__(self, name)
1323 self.blocking = blocking
1326class GenericWorkflowLabels:
1327 """Label-oriented representation of the GenericWorkflowJobs."""
1329 def __init__(self) -> None:
1330 self._label_graph = DiGraph() # Dependency graph of job labels
1331 self._label_to_jobs: defaultdict[str, list[GenericWorkflowJob]] = defaultdict(
1332 list
1333 ) # mapping job label to list of GenericWorkflowJob
1335 @property
1336 def labels(self) -> list[str]:
1337 """List of job labels (`list` [`str`], read-only)."""
1338 return list(topological_sort(self._label_graph))
1340 @property
1341 def job_counts(self) -> Counter[str]:
1342 """Count of jobs per job label (`collections.Counter`)."""
1343 return Counter({label: len(self._label_to_jobs[label]) for label in self.labels})
1345 def get_jobs_by_label(self, label: str) -> list[GenericWorkflowJob]:
1346 """Retrieve jobs by label from workflow.
1348 Parameters
1349 ----------
1350 label : `str`
1351 Label of jobs to retrieve.
1353 Returns
1354 -------
1355 jobs : list[`lsst.ctrl.bps.GenericWorkflowJob`]
1356 Jobs having given label.
1357 """
1358 return self._label_to_jobs[label]
1360 def add_job(self, job: GenericWorkflowJob, parent_labels: list[str], child_labels: list[str]) -> None:
1361 """Add job's label to labels.
1363 Parameters
1364 ----------
1365 job : `lsst.ctrl.bps.GenericWorkflowJob`
1366 The job to add to the job labels.
1367 parent_labels : `list` [`str`]
1368 Parent job labels.
1369 child_labels : `list` [`str`]
1370 Children job labels.
1371 """
1372 _LOG.debug("job: %s (%s)", job.name, job.label)
1373 _LOG.debug("parent_labels: %s", parent_labels)
1374 _LOG.debug("child_labels: %s", child_labels)
1375 self._label_to_jobs[job.label].append(job)
1376 self._label_graph.add_node(job.label)
1377 for parent in parent_labels:
1378 self._label_graph.add_edge(parent, job.label)
1379 for child in child_labels:
1380 self._label_graph.add_edge(job.label, child)
1382 def add_job_relationships(self, parent_labels: list[str], children_labels: list[str]) -> None:
1383 """Add dependencies between parent and child job labels.
1384 All parents will be connected to all children.
1386 Parameters
1387 ----------
1388 parent_labels : `list` [`str`]
1389 Parent job labels.
1390 children_labels : `list` [`str`]
1391 Children job labels.
1392 """
1393 # Since labels, must ensure not adding edge from label to itself.
1394 edges = [
1395 e
1396 for e in itertools.product(ensure_iterable(parent_labels), ensure_iterable(children_labels))
1397 if e[0] != e[1]
1398 ]
1400 self._label_graph.add_edges_from(edges)
1402 def del_job(self, job: GenericWorkflowJob) -> None:
1403 """Delete job and its label from job labels.
1405 Parameters
1406 ----------
1407 job : `lsst.ctrl.bps.GenericWorkflowJob`
1408 The job to delete from the job labels.
1409 """
1410 self._label_to_jobs[job.label].remove(job)
1411 # Don't leave keys around if removed last job
1412 if not self._label_to_jobs[job.label]:
1413 del self._label_to_jobs[job.label]
1415 parents = self._label_graph.predecessors(job.label)
1416 children = self._label_graph.successors(job.label)
1417 self._label_graph.remove_node(job.label)
1418 self._label_graph.add_edges_from(
1419 itertools.product(ensure_iterable(parents), ensure_iterable(children))
1420 )
1422 def subgraph(self, labels: Iterable[str]) -> DiGraph:
1423 """Create subgraph of workflow label graph with given labels.
1425 Parameters
1426 ----------
1427 labels : Iterable [`str`]
1428 Labels to appear in subgraph.
1430 Returns
1431 -------
1432 subgraph : `networkx.DiGraph`
1433 Subgraph of workflow label graph with given labels.
1434 """
1435 return self._label_graph.subgraph(labels)