diff --git a/ceci/pipeline/__init__.py b/ceci/pipeline/__init__.py index b0920db..57cd327 100644 --- a/ceci/pipeline/__init__.py +++ b/ceci/pipeline/__init__.py @@ -11,5 +11,5 @@ from .dry_run import DryRunPipeline from .flow_chart import FlowChartPipeline from .mini import MiniPipeline -from .file_manager import FileManager from .sec import StageExecutionConfig +from .graph import StageStatus diff --git a/ceci/pipeline/dry_run.py b/ceci/pipeline/dry_run.py index 473bf73..375d40b 100644 --- a/ceci/pipeline/dry_run.py +++ b/ceci/pipeline/dry_run.py @@ -18,11 +18,11 @@ def initiate_run(self, overall_inputs): def should_skip_stage(self, stage): return False - def enqueue_job(self, stage, pipeline_files): + def enqueue_job(self, stage, input_paths): outputs = stage.find_outputs(self.run_config["output_dir"]) - sec = self.stage_execution_config[stage.instance_name] + sec = self.get_stage_execution_config(stage.instance_name) - cmd = sec.generate_full_command(pipeline_files, outputs, self.stages_config) + cmd = sec.generate_full_command(input_paths, outputs, self.stages_config) # Replace the first instance of the stage name with bold # text, but only if we are printing to screen. This helps the diff --git a/ceci/pipeline/file_manager.py b/ceci/pipeline/file_manager.py deleted file mode 100644 index 671bff1..0000000 --- a/ceci/pipeline/file_manager.py +++ /dev/null @@ -1,73 +0,0 @@ -class FileManager(dict): - """Small class to manage files within a particular Pipeline - - This is a dict which is used for tag to path mapping, - but has a couple of additional dicts to manage the reverse mapping and the - tag to type mapping. - - - The tag to path mapping is the thing that the Pipeline uses to set the - input paths for downstream stages that use the outputs of earlier stages, - i.e., everything in the pipeline can refer to a particular file by tag. - - The tag defaults to the input or output tag as define in the stage class attributes. - However, in the case that we want multiple stages of the same class in a pipeline we - have to alias the tags so that each stage can write to its own location (and so that - downstream stages can pick up that location correctly) - """ - - def __init__(self): - """Constructor, makes empty dictionaries""" - self._tag_to_type = {} - self._path_to_tag = {} - dict.__init__(self) - - def __setitem__(self, tag, path): - """Override dict.__setitem__() to also insert the reverse mapping""" - dict.__setitem__(self, tag, path) - self._path_to_tag[path] = tag - - def insert(self, tag, path=None, ftype=None): - """Insert a file, including the path and the file type - - Parameters - ---------- - tag : str - The tag by which this file will be identified - path : str - The path to this file - ftype : type - The file type for this file - """ - if path is not None: - self[tag] = path - self._path_to_tag[path] = tag - if tag not in self._tag_to_type: - self._tag_to_type[tag] = ftype - - def get_type(self, tag): - """Return the file type associated to a given tag""" - return self._tag_to_type[tag] - - def get_path(self, tag): - """Return the path associated to a give tag""" - return self[tag] - - def get_tag(self, path): - """Return the tag associated to a given path""" - return self._path_to_tag[path] - - def insert_paths(self, path_dict): - """Insert a set of paths from a dict that has tag, path pairs""" - for key, val in path_dict.items(): - self.insert(key, path=val) - - def insert_outputs(self, stage, outdir): - """Insert a set of tags and associated paths and file types that are output by a stage""" - stage_outputs = stage.find_outputs(outdir) - for tag, ftype in stage.outputs: - aliased_tag = stage.get_aliased_tag(tag) - path = stage_outputs[aliased_tag] - self.insert(aliased_tag, path=path, ftype=ftype) - return stage_outputs - diff --git a/ceci/pipeline/graph.py b/ceci/pipeline/graph.py index 0b1fb2b..6744c22 100644 --- a/ceci/pipeline/graph.py +++ b/ceci/pipeline/graph.py @@ -1,71 +1,156 @@ +import enum import networkx import collections -def build_graph(stage_names, stage_classes, stage_aliases, overall_inputs): + +class StageStatus(enum.Enum): + """The run status of a pipeline stage, stored on its graph node""" + + PENDING = "pending" + QUEUED = "queued" + RUNNING = "running" + SKIPPED = "skipped" + COMPLETE = "complete" + FAILED = "failed" + + +# Key in the graph attributes dictionary where the static stage ordering +# is cached. Functions in this module that change the stages or connections +# in a graph invalidate it. +STATIC_ORDERING = "static_ordering" + + +def invalidate_static_ordering(graph): + """Clear the cached static ordering of a graph, if there is one. + + This must be called when stages or edges are added to or removed from + the graph other than through the functions in this module. + + Parameters + ---------- + graph : `networkx.DiGraph` + The graph to modify """ - Build a directed graph of the pipeline stages and their inputs and outputs. + graph.graph.pop(STATIC_ORDERING, None) + + +def add_stage_node(graph, stage_name, stage_class, aliases, **attrs): + """ + Add a stage, and the output files it generates, to a pipeline graph. + + The inputs to the stage are not connected here, since they may not + exist in the graph yet; use `connect_stage_inputs` for that. - This also checks for various pipeline problems like missing inputs, - multiple stages generating the same output, and cyclic dependencies. - Parameters ---------- - stage_names : `list` of `str` - The names of the stages in the pipeline - stage_classes : `list` of `PipelineStage` subclasses - The classes of the stages in the pipeline - stage_aliases : `list` of `dict` - A list of dictionaries mapping tags to aliases for each stage - overall_inputs : `dict` of `str: str` - A list of overall input files that are not generated by any stage but are needed by the pipeline + graph : `networkx.DiGraph` + The graph to add to + stage_name : `str` + The name of the stage + stage_class : `PipelineStage` subclass + The class of the stage + aliases : `dict` + Dictionary mapping tags to aliases for this stage + **attrs + Any additional attributes to set on the stage node + """ + if stage_name in graph.nodes: + raise ValueError(f"Stage {stage_name} is used multiple times in the pipeline") + + # Check all the outputs before changing anything, so that + # a failure leaves the graph unchanged + outputs = [] + for tag1, ftype in stage_class.outputs_(): + tag = aliases.get(tag1, tag1) + node = graph.nodes.get(tag) + if node is None: + pass + elif node["type"] == "output": + raise ValueError(f"Output {tag} is generated by multiple stages: {[node['creator'], stage_name]}") + else: + raise ValueError(f"Tag {tag} is generated by pipeline stage {stage_name} but is also listed as an overall input to the pipeline") + outputs.append((tag, tag1, ftype)) + + invalidate_static_ordering(graph) + graph.add_node( + stage_name, + type="stage", + stage_class=stage_class, + aliases=aliases, + status=StageStatus.PENDING, + **attrs, + ) + for tag, tag1, ftype in outputs: + graph.add_node(tag, type="output", creator=stage_name, original_tag=tag1, ftype=ftype) + graph.add_edge(stage_name, tag) + + +def add_input_node(graph, tag, **attrs): + """ + Add (or update) an overall input file in a pipeline graph. - Returns - ------- + Parameters + ---------- graph : `networkx.DiGraph` - A directed graph representing the pipeline stages, their inputs, and outputs + The graph to add to + tag : `str` + The tag for the file + **attrs + Any additional attributes, e.g. path, to set on the node """ + node = graph.nodes.get(tag) + if node is not None and node["type"] != "input": + creator = node.get("creator", tag) + raise ValueError(f"Tag {tag} is generated by pipeline stage {creator} but is also listed as an overall input to the pipeline") + graph.add_node(tag, type="input", **attrs) - # Build a graph with all the stages and files in it - graph = networkx.DiGraph() - # First add all the overall inputs - for inp in overall_inputs: - graph.add_node(inp, type="input") - - # Now go through all the stages. We do this in two passes, - # one to add the stages themselves and their outputs, and one - # to connect in all their inputs - for stage_name, stage_class, aliases in zip(stage_names, stage_classes, stage_aliases): - if stage_name in graph.nodes: - raise ValueError(f"Stage {stage_name} is used multiple times in the pipeline") - graph.add_node(stage_name, type="stage", stage_class=stage_class) - - # Now set up all the outputs of the stage, first checking to see if it has been aliased - for tag1 in stage_class.output_tags(): - tag = aliases.get(tag1, tag1) - if tag in overall_inputs: - raise ValueError(f"Tag {tag} is generated by pipeline stage {stage_name} but is also listed as an overall input to the pipeline") - graph.add_node(tag, type="output", creator=stage_class, original_tag=tag1) - graph.add_edge(stage_name, tag) - - # Now check for outputs that are generated by multiple stages - for node in graph.nodes: - if graph.nodes[node]["type"] == "output": - parents = list(graph.predecessors(node)) - if len(parents) > 1: - raise ValueError(f"Output {node} is generated by multiple stages: {parents}") - - - # Now go through and add in all the inputs needed for each stage: +def remove_stage_node(graph, stage_name): + """ + Remove a stage and all the outputs it generates from a pipeline graph. + + Parameters + ---------- + graph : `networkx.DiGraph` + The graph to modify + stage_name : `str` + The name of the stage + """ + outputs = [n for n in graph.successors(stage_name) if graph.nodes[n]["type"] == "output"] + invalidate_static_ordering(graph) + graph.remove_nodes_from(outputs + [stage_name]) + + +def stage_nodes(graph): + """Return the names of the stage nodes in a graph, in insertion order""" + return [n for n, t in graph.nodes(data="type") if t == "stage"] + + +def connect_stage_inputs(graph): + """ + Connect the input files of each stage in the graph to that stage, + and check that the result is a valid pipeline, with no missing + inputs or cycles. + + This can safely be called more than once. + + Parameters + ---------- + graph : `networkx.DiGraph` + The graph to modify + """ + invalidate_static_ordering(graph) missing = collections.defaultdict(list) - for stage_name, stage_class, aliases in zip(stage_names, stage_classes, stage_aliases): - for tag in stage_class.input_tags(): + for stage_name in stage_nodes(graph): + node = graph.nodes[stage_name] + aliases = node["aliases"] + for tag in node["stage_class"].input_tags(): tag = aliases.get(tag, tag) if tag not in graph.nodes: missing[tag].append(stage_name) else: graph.add_edge(tag, stage_name) - + if missing: msg = [] for tag, stages in missing.items(): @@ -74,27 +159,79 @@ def build_graph(stage_names, stage_classes, stage_aliases, overall_inputs): msg = "\n".join(msg) raise ValueError(f"The following files are needed by the pipeline but not generated by any stage:\n{msg}") - # Now check for cycles in the graph - cycles = list(networkx.simple_cycles(graph)) - if cycles: + # Now check for cycles in the graph, only listing them if there are any + if not networkx.is_directed_acyclic_graph(graph): msg = [] - for cycle in cycles: + for cycle in networkx.simple_cycles(graph): cycle_text = " -> ".join(cycle) cycle_text += " -> " + cycle[0] msg.append(f"- {cycle_text}") + msg = "\n".join(msg) raise ValueError(f"Your pipeline is cyclic. The following cycles were found:\n{msg}") - + + +def build_graph(stage_names, stage_classes, stage_aliases, overall_inputs): + """ + Build a directed graph of the pipeline stages and their inputs and outputs. + + This also checks for various pipeline problems like missing inputs, + multiple stages generating the same output, and cyclic dependencies. + + Parameters + ---------- + stage_names : `list` of `str` + The names of the stages in the pipeline + stage_classes : `list` of `PipelineStage` subclasses + The classes of the stages in the pipeline + stage_aliases : `list` of `dict` + A list of dictionaries mapping tags to aliases for each stage + overall_inputs : `dict` of `str: str` + A list of overall input files that are not generated by any stage but are needed by the pipeline + + Returns + ------- + graph : `networkx.DiGraph` + A directed graph representing the pipeline stages, their inputs, and outputs + """ + graph = networkx.DiGraph() + for inp in overall_inputs: + add_input_node(graph, inp) + for stage_name, stage_class, aliases in zip(stage_names, stage_classes, stage_aliases): + add_stage_node(graph, stage_name, stage_class, aliases) + connect_stage_inputs(graph) return graph + def get_static_ordering(graph): - # get a topological ordering of the graph - # this will be a list of all the nodes in the graph, including the input and output files - ordering = list(networkx.topological_sort(graph)) + """ + Get an order in which the stages in a pipeline graph can be run. - # cut down to just the stages - ordering = [node for node in ordering if graph.nodes[node]["type"] == "stage"] - - return ordering + The result is cached in the graph attributes, and re-used until + it is invalidated by a change to the graph. + + Parameters + ---------- + graph : `networkx.DiGraph` + The pipeline graph + + Returns + ------- + ordering : `list` of `str` + The names of the stages in a valid execution order + """ + ordering = graph.graph.get(STATIC_ORDERING) + + if ordering is None: + # get a topological ordering of the graph + # this will be a list of all the nodes in the graph, including the input and output files + ordering = list(networkx.topological_sort(graph)) + + # cut down to just the stages + ordering = [node for node in ordering if graph.nodes[node]["type"] == "stage"] + graph.graph[STATIC_ORDERING] = ordering + + # return a copy so callers cannot modify the cached version + return list(ordering) def trim_pipeline_graph(graph, from_=None, to_=None): """ @@ -170,30 +307,23 @@ def trim_pipeline_graph(graph, from_=None, to_=None): # Create a subgraph with just these nodes trimmed_graph = graph.subgraph(nodes_to_keep).copy() + # The copy includes the graph attributes, so any cached + # ordering refers to the untrimmed graph + invalidate_static_ordering(trimmed_graph) + # Some stages that were previously outputs of existing stages # must now be overall inputs to the pipeline instead because # the stage that generated them has been trimmed. We need to keep # track of this converted_inputs = {} - for node_name in list(trimmed_graph.nodes): - node = trimmed_graph.nodes[node_name] + for node_name, node in trimmed_graph.nodes(data=True): if node["type"] == "output": # The node will be listed as an output but will have no # predecessor creating it now. - predecessors = list(trimmed_graph.predecessors(node_name)) - if not predecessors: - creator_stage_class = node["creator"] - original_tag = node["original_tag"] - # use the original tag to find the file type object - for t, typ in creator_stage_class.outputs: - if t == original_tag: - ftype = typ - break - else: - raise RuntimeError("This should never happen - please open an issue") + if trimmed_graph.in_degree(node_name) == 0: # the node_name is the aliased_tag, so this file name should be correct - file_name = ftype.make_name(node_name) + file_name = node["ftype"].make_name(node_name) converted_inputs[node_name] = file_name return trimmed_graph, converted_inputs diff --git a/ceci/pipeline/mini.py b/ceci/pipeline/mini.py index 5747370..a44ae9a 100644 --- a/ceci/pipeline/mini.py +++ b/ceci/pipeline/mini.py @@ -1,6 +1,7 @@ import sys from .. import minirunner from .pipeline import Pipeline, get_default_site +from .graph import StageStatus class MiniPipeline(Pipeline): """A pipeline subclass that uses Minirunner, a sub-module @@ -50,42 +51,49 @@ def build_dag(self, jobs): Parameters ---------- - stages: list[PipelineStage] - A list of stages to generate the DAG for + jobs: dict[str, minirunner.Job] + The jobs to generate the DAG for, keyed by stage name """ depend = {} # for each stage in our pipeline ... - for stage in self.stages[:]: - if stage.instance_name not in jobs: + for stage_name in self.stage_names: + if stage_name not in jobs: continue - job = jobs[stage.instance_name] + job = jobs[stage_name] depend[job] = [] - # check for each of the inputs for that stage ... - for tag in stage.input_tags(): - aliased_tag = stage.get_aliased_tag(tag) - for potential_parent in self.stages[:]: - # if that stage is supplied by another pipeline stage - if potential_parent.instance_name not in jobs: # pragma: no cover - continue - potential_parent_tags = [ - potential_parent.get_aliased_tag(tag_) - for tag_ in potential_parent.output_tags() - ] - if aliased_tag in potential_parent_tags: - depend[job].append(jobs[potential_parent.instance_name]) + # the stages it depends on are the ones that make + # its input files. Stages that are not being run + # are not included + for tag in self.graph.predecessors(stage_name): + for parent in self.graph.predecessors(tag): + if parent in jobs: + depend[job].append(jobs[parent]) return depend + def _status_callback(self, event_type, event_info): + """Record job events as stage status in the graph, and + then pass them on to any user callback.""" + status = { + minirunner.EVENT_LAUNCH: StageStatus.RUNNING, + minirunner.EVENT_COMPLETED: StageStatus.COMPLETE, + minirunner.EVENT_FAIL: StageStatus.FAILED, + }.get(event_type) + if status is not None: + self.set_stage_status(event_info["job"].name, status) + if self.callback is not None: + self.callback(event_type, event_info) + def initiate_run(self, overall_inputs): jobs = {} stages = [] return jobs, stages - def enqueue_job(self, stage, pipeline_files): - sec = self.stage_execution_config[stage.instance_name] + def enqueue_job(self, stage, input_paths): + sec = self.get_stage_execution_config(stage.instance_name) outputs = stage.find_outputs(self.run_config["output_dir"]) - cmd = sec.generate_full_command(pipeline_files, outputs, self.stages_config) + cmd = sec.generate_full_command(input_paths, outputs, self.stages_config) job = minirunner.Job( stage.instance_name, cmd, @@ -106,12 +114,12 @@ def run_jobs(self): # This pipeline, being mainly for testing, can only # run at a single site, so we can assume here that the # sites are all the same - sec = self.stage_execution_config[self.stage_names[0]] + sec = self.get_stage_execution_config(self.stage_names[0]) nodes = sec.site.info["nodes"] # Run under minirununer runner = minirunner.Runner( - nodes, graph, log_dir, callback=self.callback, sleep=self.sleep + nodes, graph, log_dir, callback=self._status_callback, sleep=self.sleep ) interval = self.launcher_config.get("interval", 3) try: diff --git a/ceci/pipeline/parsl.py b/ceci/pipeline/parsl.py index 70a67db..cee3d5e 100644 --- a/ceci/pipeline/parsl.py +++ b/ceci/pipeline/parsl.py @@ -2,6 +2,7 @@ import os from .pipeline import Pipeline +from .graph import StageStatus class ParslPipeline(Pipeline): @@ -10,7 +11,7 @@ class ParslPipeline(Pipeline): def initiate_run(self, overall_inputs): return [] # list of futures - def enqueue_job(self, stage, pipeline_files): + def enqueue_job(self, stage, input_paths): from parsl.data_provider.files import File @@ -22,7 +23,7 @@ def enqueue_job(self, stage, pipeline_files): # parsl wants. # The inputs that exist already need to be converted into Parsl File objects. # The ones that don't stay as data futures - inputs1 = stage.find_inputs(pipeline_files) + inputs1 = stage.find_inputs(input_paths) inputs = [ File(val) if isinstance(val, str) else val for val in inputs1.values() ] @@ -47,8 +48,10 @@ def run_jobs(self): try: # This waits for b/g pipeline completion. future.result() + self.set_stage_status(stage_name, StageStatus.COMPLETE) # Parsl emits this on any non-zero status code. except BashExitFailure: + self.set_stage_status(stage_name, StageStatus.FAILED) stdout_file = f"{log_dir}/{stage_name}.err" stderr_file = f"{log_dir}/{stage_name}.out" sys.stderr.write( @@ -138,7 +141,7 @@ def generate_app(self, stage, run_config): config = f"{{inputs[{config_index}]}}" # This includes all the "mpirun" stuff. - sec = self.stage_execution_config[stage.instance_name] + sec = self.get_stage_execution_config(stage.instance_name) executor = sec.site.info["executor"] # Construct the command line call diff --git a/ceci/pipeline/pipeline.py b/ceci/pipeline/pipeline.py index eb86cde..425446d 100644 --- a/ceci/pipeline/pipeline.py +++ b/ceci/pipeline/pipeline.py @@ -9,8 +9,17 @@ from ..stage import PipelineStage from ..sites import load, set_default_site, get_default_site from ..utils import extra_paths -from .graph import build_graph, get_static_ordering, trim_pipeline_graph, combine_output_nodes, combine_input_nodes -from .file_manager import FileManager +from .graph import ( + StageStatus, + add_stage_node, + add_input_node, + remove_stage_node, + connect_stage_inputs, + get_static_ordering, + trim_pipeline_graph, + combine_output_nodes, + combine_input_nodes, +) from .sec import StageExecutionConfig from .templates import read_and_apply_template @@ -131,15 +140,15 @@ def __init__(self, stages, launcher_config, **kwargs): self.overall_inputs = {} self.modules = kwargs.get("modules", "") - # These are populated as we add stages below - self.stage_execution_config = {} - self.stage_names = [] + # The graph is the main store of information about the pipeline. + # Stage nodes carry the execution config ("sec"), the stage + # configuration ("config") and run status ("status") of each stage, + # and file nodes carry the "path" and "ftype" of each file. + # It is populated as we add stages below. + self.graph = networkx.DiGraph() self.run_info = None self.run_config = None - self.stages = None - self.graph = None - self.pipeline_files = FileManager() self.pipeline_outputs = None self.stages_config = None self.stage_config_data = None @@ -262,15 +271,93 @@ def build_config( def __getitem__(self, name): """Get a particular stage by name""" - try: - return self.stage_execution_config[name].stage_obj - except Exception as msg: # pragma: no cover - raise AttributeError(f"Pipeline does not have stage {name}") from msg + # Use __dict__ directly to avoid recursion via __getattr__ if + # the graph does not exist yet + graph = self.__dict__.get("graph") + node = None if graph is None else graph.nodes.get(name) + if node is None or node["type"] != "stage": + raise AttributeError(f"Pipeline does not have stage {name}") + return node["sec"].stage_obj def __getattr__(self, name): """Get a particular stage by name""" return self.__getitem__(name) + @property + def stage_names(self): + """The names of the stages in the pipeline, in an order they can be run""" + return get_static_ordering(self.graph) + + @property + def stages(self): + """The PipelineStage objects in the pipeline, in an order they can be run""" + return [self.graph.nodes[name]["sec"].stage_obj for name in self.stage_names] + + def add_inputs(self, paths): + """Add overall input files to the pipeline, or update their paths + + Parameters + ---------- + paths : `Mapping` + Mapping of tag to path for the input files + """ + for tag, path in paths.items(): + add_input_node(self.graph, tag, path=path) + + def stage_input_paths(self, stage): + """Return the paths to the inputs of a stage, as known so far + + Parameters + ---------- + stage : `PipelineStage` + The stage in question + + Returns + ------- + paths : `dict` + Mapping of aliased tag to path, in the order of the stage's input tags. + Inputs whose paths are not yet known are omitted. + """ + paths = {} + for tag in stage.input_tags(): + aliased_tag = stage.get_aliased_tag(tag) + node = self.graph.nodes.get(aliased_tag, {}) + if "path" in node: + paths[aliased_tag] = node["path"] + return paths + + def _record_outputs(self, stage, paths): + """Record the paths of a stage's outputs on its output nodes + + Parameters + ---------- + stage : `PipelineStage` + The stage in question + paths : `Mapping` + Mapping of aliased tag to path for the outputs + """ + for tag in stage.output_tags(): + aliased_tag = stage.get_aliased_tag(tag) + self.graph.nodes[aliased_tag]["path"] = paths[aliased_tag] + + def _stage_node(self, stage_name): + node = self.graph.nodes.get(stage_name) + if node is None or node["type"] != "stage": + raise KeyError(f"Failed to find stage named {stage_name} in {self.stage_names}") + return node + + def get_stage_execution_config(self, stage_name): + """Return the StageExecutionConfig for the named stage""" + return self._stage_node(stage_name)["sec"] + + def get_stage_status(self, stage_name): + """Return the StageStatus for the named stage""" + return self._stage_node(stage_name)["status"] + + def set_stage_status(self, stage_name, status): + """Set the StageStatus for the named stage""" + self._stage_node(stage_name)["status"] = StageStatus(status) + def print_stages(self, stream=sys.stdout): """Print the list of stages in this pipeline to a stream""" for stage in self.stages: @@ -470,8 +557,7 @@ def add_stage(self, stage_info): else: sec = StageExecutionConfig(stage_info) sec.build_stage_class() - self.stage_execution_config[sec.name] = sec - self.stage_names.append(sec.name) + add_stage_node(self.graph, sec.name, sec.stage_class, sec.aliases, sec=sec, config={}) # If we are adding a pre-built stage then we insert the output files now. # Otherwise this happens later. Why does this happen now? Probably so that @@ -479,7 +565,9 @@ def add_stage(self, stage_info): # interactively. if sec.stage_obj is None: return {} - return self.pipeline_files.insert_outputs(sec.stage_obj, ".") + stage_outputs = sec.stage_obj.find_outputs(".") + self._record_outputs(sec.stage_obj, stage_outputs) + return stage_outputs def build_stage(self, stage_class, **kwargs): """Build a stage and add it to the pipeline @@ -496,15 +584,22 @@ def build_stage(self, stage_class, **kwargs): Notes ----- - The keyword arguments will be based to the `stage_class` constructor + The keyword arguments will be based to the `stage_class` constructor. + Paths to the stage's inputs are taken from the outputs of stages already + added and any inputs set with `add_inputs`, unless given explicitly. - The output files produced by this stage will be added to the - `Pipeline.pipeline_files` data member, so that they are available to later stages + The paths to the output files produced by this stage will be recorded + in the pipeline graph, so that they are available to later stages """ kwcopy = kwargs.copy() aliases = kwcopy.pop("aliases", {}) comm = kwcopy.pop("comm", None) - kwcopy.update(**self.pipeline_files) + kwcopy.setdefault("config", None) + for tag in stage_class.input_tags(): + aliased_tag = aliases.get(tag, tag) + node = self.graph.nodes.get(aliased_tag, {}) + if aliased_tag not in kwcopy and "path" in node: + kwcopy[aliased_tag] = node["path"] stage = stage_class(kwcopy, comm=comm, aliases=aliases) return self.add_stage(stage) @@ -518,8 +613,11 @@ def remove_stage(self, name): The name of the stage to remove. """ - self.stage_names.remove(name) - del self.stage_execution_config[name] + # Check that this is the name of a stage, raising a KeyError if not. + # Otherwise a file name would be accepted, and remove_stage_node + # would delete that file's node from the graph. + self._stage_node(name) + remove_stage_node(self.graph, name) def configure_stages(self, stages_config): @@ -534,35 +632,35 @@ def configure_stages(self, stages_config): self.read_stages_config(stages_config) stages = [] - all_inputs = self.overall_inputs.copy() for stage_name in self.stage_names: # Find the stage class and execution config - sec = self.stage_execution_config[stage_name] + node = self.graph.nodes[stage_name] + sec = node["sec"] stage_aliases = sec.aliases stage_class = sec.stage_class + # Record the configuration for this stage in the graph + stage_config = self.stage_config_data.get(stage_name, {}) + node["config"] = stage_config + # In interactive use the stage object might already be created # otherwise we need to create it if sec.stage_obj is not None: - orig_stage_config = self.stage_config_data.get(stage_name, {}).copy() stage = sec.stage_obj - stage.config.update(**orig_stage_config) + stage.config.update(**stage_config.copy()) else: # Find the inputs for this stage and set up the arguments # to the stage init method - args = self.stage_config_data.get(stage_name, {}) + args = stage_config.copy() for tag in stage_class.input_tags(): aliased_tag = stage_aliases.get(tag, tag) - args[aliased_tag] = all_inputs[aliased_tag] + args[aliased_tag] = self.graph.nodes[aliased_tag]["path"] args["config"] = self.stages_config - # Make the stage object and find the outputs, - # keeping track of them for the next stages + # Make the stage object and record the outputs + # in the graph for the next stages stage = sec.build_stage_object(args) - stage_outputs = stage.find_outputs(".") - for tag in stage_class.output_tags(): - aliased_tag = stage_aliases.get(tag, tag) - all_inputs[aliased_tag] = stage_outputs[aliased_tag] + self._record_outputs(stage, stage.find_outputs(".")) stages.append(stage) return stages @@ -596,15 +694,13 @@ def construct_pipeline_graph(self, overall_inputs, run_config): self.data_registry = self.setup_data_registry(registry_info) self.overall_inputs = self.process_overall_inputs(overall_inputs) - # Get the stages in the order we need. - # First build the graph that includes both the stages and the - # input and output files for the pipeline. - self.graph = build_graph( - self.stage_names, - [self.stage_execution_config[name].stage_class for name in self.stage_names], - [self.stage_execution_config[name].aliases for name in self.stage_names], - self.overall_inputs - ) + # The stages and their outputs are already in the graph. + # Add the overall inputs, recording their paths, and then + # connect everything together. This also checks the pipeline + # is valid. + for tag, path in self.overall_inputs.items(): + add_input_node(self.graph, tag, path=path) + connect_stage_inputs(self.graph) if "to" in self.run_config or "from" in self.run_config: to_ = self.run_config.get("to") @@ -628,15 +724,7 @@ def construct_pipeline_graph(self, overall_inputs, run_config): for t, path in converted_inputs.items(): print(f" - {t}: {path}") self.overall_inputs[t] = path - - - # Re-order the pipeline stages in a static order - self.stage_names = get_static_ordering(self.graph) - - # This is also a convenient place to record the location - # of the overall inputs in the file manager - self.pipeline_files.insert_paths(self.overall_inputs) - + self.graph.nodes[t]["path"] = path return self.stage_names @@ -698,6 +786,7 @@ def enqueue_jobs(self): self.run_info = self.initiate_run(self.overall_inputs) # Now we can queue up all the jobs. The subclass decides what this actually means. + output_dir = self.run_config["output_dir"] for stage in self.stages: # If we are in "resume" mode and the pipeline has already been run # then we re-use any existing outputs. User is responsible for making @@ -705,13 +794,15 @@ def enqueue_jobs(self): if self.should_skip_stage(stage): stage.already_finished() - self.pipeline_files.insert_outputs(stage, self.run_config["output_dir"]) + self._record_outputs(stage, stage.find_outputs(output_dir)) + self.set_stage_status(stage.instance_name, StageStatus.SKIPPED) # Otherwise, run the pipeline and register any outputs from the # pipe element. else: - stage_outputs = self.enqueue_job(stage, self.pipeline_files) - self.pipeline_files.insert_paths(stage_outputs) + stage_outputs = self.enqueue_job(stage, self.stage_input_paths(stage)) + self._record_outputs(stage, stage_outputs) + self.set_stage_status(stage.instance_name, StageStatus.QUEUED) def initialize(self, overall_inputs, run_config, stages_config): @@ -737,7 +828,7 @@ def initialize(self, overall_inputs, run_config, stages_config): # Configure the individual stages. Reads the config information for them # and creates PipelineStage objects for each one. - self.stages = self.configure_stages(stages_config) + self.configure_stages(stages_config) # Tell the subclass to preapre to run the pipeline by queueing up # the jobs. This does a few other miscellaneous things too. @@ -772,8 +863,9 @@ def initiate_run(self, overall_inputs): # pragma: no cover raise NotImplementedError() @abstractmethod - def enqueue_job(self, stage, pipeline_files): # pragma: no cover - """Setup the job for a single stage, and return stage specific information""" + def enqueue_job(self, stage, input_paths): # pragma: no cover + """Setup the job for a single stage, given a mapping of its input tags to paths, + and return a mapping of its output tags to paths (or futures)""" raise NotImplementedError() @abstractmethod @@ -836,7 +928,8 @@ def save(self, pipefile, stagefile=None, reduce_config=False, **kwargs): site = None - for key, val in self.stage_execution_config.items(): + for key in self.stage_names: + val = self.get_stage_execution_config(key) if val.stage_obj is None: # pragma: no cover raise ValueError(f"Stage {key} has not been built, can not save") if site is None: @@ -902,8 +995,15 @@ def make_flow_chart(self, filename=None): graphviz.AGraph The graphviz object representing the pipeline stages """ + # The graph nodes carry python objects that we do not want + # to write to the chart, so make a copy with just the structure + # and the node types. + structure = networkx.DiGraph() + structure.add_nodes_from((n, {"type": t}) for n, t in self.graph.nodes(data="type")) + structure.add_edges_from(self.graph.edges) + # generate a graphviz object from the networkx graph - graph = networkx.nx_agraph.to_agraph(self.graph) + graph = networkx.nx_agraph.to_agraph(structure) # Some stages have long lists of output nodes like plots that # are final outputs not used in the rest of the pipeline. We @@ -941,23 +1041,12 @@ def generate_stage_command(self, stage_name, **kwargs): kwargs: dict Used to override pipeline inputs """ - try: - sec = self.stage_execution_config[stage_name] - except KeyError as msg: - raise KeyError(f'Failed to find stage named {stage_name} in {list(self.stage_execution_config.keys())}') from msg - - if sec.stage_obj is not None: - the_stage = sec.stage_obj - else: - for i, stage_name_ in self.stage_names: - if stage_name == stage_name_: - idx = i - break - if idx is None: - raise KeyError(f'Failed to find stage named {stage_name} in {self.stage_names}') - the_stage = self.stages[idx] - - all_inputs = self.pipeline_files.copy() + sec = self.get_stage_execution_config(stage_name) + the_stage = sec.stage_obj + if the_stage is None: + raise ValueError(f"Stage {stage_name} has not been configured yet") + + all_inputs = self.stage_input_paths(the_stage) all_inputs.update(**kwargs) outputs = the_stage.find_outputs(self.run_config["output_dir"]) diff --git a/ceci/stage.py b/ceci/stage.py index 3de3c1f..65d440e 100644 --- a/ceci/stage.py +++ b/ceci/stage.py @@ -1360,15 +1360,15 @@ def get_config_dict(self, ignore=None, reduce_config=False): out_dict[key] = cast_to_streamable(val) return out_dict - def find_inputs(self, pipeline_files): - """Find and retrun all the inputs associated to this stage in the FileManager + def find_inputs(self, input_paths): + """Find and retrun all the inputs associated to this stage in a mapping of tags to paths These are returned as a dictionary of tag : path pairs """ ret_dict = {} for tag, _ in self.inputs_(): aliased_tag = self.get_aliased_tag(tag) - ret_dict[aliased_tag] = pipeline_files[aliased_tag] + ret_dict[aliased_tag] = input_paths[aliased_tag] return ret_dict def find_outputs(self, outdir): diff --git a/nb/ceci_interactive_example.ipynb b/nb/ceci_interactive_example.ipynb index 4abc0f3..358f8fb 100644 --- a/nb/ceci_interactive_example.ipynb +++ b/nb/ceci_interactive_example.ipynb @@ -350,13 +350,10 @@ "pipe2 = Pipeline.interactive()\n", "overall_inputs = {'DM':'./tests/inputs/dm.txt',\n", " 'fiducial_cosmology':'./tests/inputs/fiducial_cosmology.txt'}\n", - "inputs = overall_inputs.copy()\n", - "inputs['metacalibration'] = True\n", - "inputs['config'] = None\n", "\n", - "pipe2.pipeline_files.update(**inputs)\n", + "pipe2.add_inputs(overall_inputs)\n", "pipe2.build_stage(PZEstimationPipe)\n", - "pipe2.build_stage(shearMeasurementPipe, apply_flag=False)\n", + "pipe2.build_stage(shearMeasurementPipe, metacalibration=True, apply_flag=False)\n", "pipe2.build_stage(WLGCSelector, zbin_edges=[0.2, 0.3, 0.5], ra_range=[-5, 5])\n", "pipe2.build_stage(SysMapMaker)\n", "pipe2.build_stage(SourceSummarizer)\n", diff --git a/tests/test_interactive.py b/tests/test_interactive.py index fa4cd65..e66edfd 100644 --- a/tests/test_interactive.py +++ b/tests/test_interactive.py @@ -79,13 +79,10 @@ def test_interactive_pipeline(): pipe2 = Pipeline.interactive() overall_inputs = {'DM':'./tests/inputs/dm.txt', 'fiducial_cosmology':'./tests/inputs/fiducial_cosmology.txt'} - inputs = overall_inputs.copy() - inputs['metacalibration'] = True - inputs['config'] = None - pipe2.pipeline_files.update(**inputs) + pipe2.add_inputs(overall_inputs) pipe2.build_stage(PZEstimationPipe) - pipe2.build_stage(shearMeasurementPipe, apply_flag=False) + pipe2.build_stage(shearMeasurementPipe, metacalibration=True, apply_flag=False) pipe2.build_stage(WLGCSelector, zbin_edges=[0.2, 0.3, 0.5], ra_range=[-5, 5]) pipe2.build_stage(SysMapMaker) pipe2.build_stage(SourceSummarizer) @@ -105,9 +102,10 @@ def test_interactive_pipeline(): rpr = repr(pipe2.WLGCCov.config) - path = pipe2.pipeline_files.get_path('covariance_copy') - assert pipe2.pipeline_files.get_tag(path) == 'covariance_copy' - assert pipe2.pipeline_files.get_type('covariance_copy') == pipe2.WLGCCov.get_output_type('covariance_1') + node = pipe2.graph.nodes['covariance_copy'] + assert node['path'] == pipe2.WLGCCov.find_outputs(pipeline.run_config['output_dir'])['covariance_copy'] + assert node['creator'] == 'WLGCCov' + assert node['ftype'] == pipe2.WLGCCov.get_output_type('covariance_1') pipe2.run() @@ -118,10 +116,8 @@ def test_inter_pipe(): pipe2 = Pipeline.interactive() overall_inputs = {'DM':'./tests/inputs/dm.txt', 'fiducial_cosmology':'./tests/inputs/fiducial_cosmology.txt'} - inputs = overall_inputs.copy() - inputs['config'] = None - pipe2.pipeline_files.update(**inputs) + pipe2.add_inputs(overall_inputs) pipe2.build_stage(PZEstimationPipe, name='bob') assert isinstance(pipe2.bob, PZEstimationPipe) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 782f214..26da403 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -3,6 +3,7 @@ from ceci_example.types import TextFile from ceci.sites import load, reset_default_site from ceci.utils import extra_paths +from ceci.pipeline.graph import StageStatus import pytest from parsl import clear import yaml @@ -132,6 +133,7 @@ def run(self): def _return_value_test_(resume): expected_status = 0 if ((resume == "resume") or (resume is True)) else 1 + expected_stage_status = StageStatus.SKIPPED if expected_status == 0 else StageStatus.FAILED # Mini pipeline should not run launcher_config = {"interval": 0.5, "name": "mini"} run_config = { @@ -144,6 +146,7 @@ def _return_value_test_(resume): pipeline.initialize({}, run_config, "tests/config.yml") status = pipeline.run() assert status == expected_status + assert pipeline.get_stage_status("FailingStage") == expected_stage_status # Parsl pipeline should not run stage either launcher_config = {"name": "parsl"} @@ -154,6 +157,7 @@ def _return_value_test_(resume): pipeline.initialize({}, run_config, "tests/config.yml") status = pipeline.run() assert status == expected_status + assert pipeline.get_stage_status("FailingStage") == expected_stage_status clear() # clear parsl settings reset_default_site() # reset so default is minirunner again diff --git a/tests/test_pipeline_graph.py b/tests/test_pipeline_graph.py new file mode 100644 index 0000000..b28ebfe --- /dev/null +++ b/tests/test_pipeline_graph.py @@ -0,0 +1,220 @@ +from ceci import PipelineStage, MiniPipeline, Pipeline, FlowChartPipeline +from ceci.pipeline import StageExecutionConfig +from ceci.pipeline.graph import StageStatus, STATIC_ORDERING +from ceci_example.types import TextFile +import networkx +import pytest + + +class GraphStageX(PipelineStage): + name = "GraphStageX" + inputs = [("gx_in", TextFile)] + outputs = [("gx_out", TextFile)] + config_options = {"a": 1} + + def run(self): # pragma: no cover + pass + + +class GraphStageY(PipelineStage): + name = "GraphStageY" + inputs = [("gx_out", TextFile)] + outputs = [("gy_out", TextFile), ("gy_extra", TextFile)] + config_options = {"b": 2} + + def run(self): # pragma: no cover + pass + + +class GraphStageZ(PipelineStage): + name = "GraphStageZ" + inputs = [("gy_out", TextFile), ("gx_in", TextFile)] + outputs = [("gz_out", TextFile)] + config_options = {} + + def run(self): # pragma: no cover + pass + + +STAGES = [{"name": "GraphStageZ"}, {"name": "GraphStageY"}, {"name": "GraphStageX"}] +INPUTS = {"gx_in": "./tests/inputs/gx_in.txt"} +CONFIG = {"GraphStageX": {"a": 10}, "global": {"c": 3}} + + +def make_pipeline(tmp_path, **run_config): + rc = { + "output_dir": str(tmp_path), + "log_dir": str(tmp_path), + "resume": "restart", + } + rc.update(run_config) + pipeline = FlowChartPipeline(STAGES, {"name": "mini"}) + pipeline.initialize(INPUTS, rc, CONFIG) + return pipeline + + +def test_stage_nodes(tmp_path): + pipeline = make_pipeline(tmp_path) + graph = pipeline.graph + + assert pipeline.stage_names == ["GraphStageX", "GraphStageY", "GraphStageZ"] + assert [s.instance_name for s in pipeline.stages] == pipeline.stage_names + + for name in pipeline.stage_names: + node = graph.nodes[name] + assert node["type"] == "stage" + assert isinstance(node["sec"], StageExecutionConfig) + assert node["sec"].stage_obj is pipeline[name] + assert node["status"] == StageStatus.QUEUED + assert pipeline.get_stage_status(name) == StageStatus.QUEUED + + # stage configuration, including global options, is on the node + assert graph.nodes["GraphStageX"]["config"] == {"a": 10, "c": 3} + assert graph.nodes["GraphStageY"]["config"] == {} + assert pipeline.GraphStageX.config["a"] == 10 + + +def test_file_nodes(tmp_path): + pipeline = make_pipeline(tmp_path) + graph = pipeline.graph + + assert graph.nodes["gx_in"]["type"] == "input" + assert graph.nodes["gx_in"]["path"] == INPUTS["gx_in"] + + node = graph.nodes["gy_out"] + assert node["type"] == "output" + assert node["creator"] == "GraphStageY" + assert node["ftype"] is TextFile + assert node["path"] == f"{tmp_path}/gy_out.txt" + + # input paths for a stage come from the graph, in input tag order + assert pipeline.stage_input_paths(pipeline.GraphStageZ) == { + "gy_out": node["path"], + "gx_in": INPUTS["gx_in"], + } + + +def test_trimmed_stages_removed(tmp_path): + pipeline = make_pipeline(tmp_path, **{"from": "GraphStageY"}) + assert pipeline.stage_names == ["GraphStageY", "GraphStageZ"] + assert "GraphStageX" not in pipeline.graph + with pytest.raises(AttributeError): + pipeline.GraphStageX + + # the output of the trimmed stage is now an input from the previous run + assert pipeline.overall_inputs["gx_out"] == f"{tmp_path}/gx_out.txt" + assert pipeline.graph.nodes["gx_out"]["path"] == f"{tmp_path}/gx_out.txt" + + # Only the remaining stages are saved + pipeline.save(str(tmp_path / "trimmed.yml")) + saved = Pipeline.build_config(str(tmp_path / "trimmed.yml")) + assert [s["name"] for s in saved["stages"]] == ["GraphStageY", "GraphStageZ"] + + +def test_remove_stage(): + pipeline = Pipeline(STAGES, {"name": "mini"}) + assert "gy_extra" in pipeline.graph + pipeline.remove_stage("GraphStageY") + assert "GraphStageY" not in pipeline.graph + assert "gy_out" not in pipeline.graph + assert "gy_extra" not in pipeline.graph + with pytest.raises(KeyError): + pipeline.remove_stage("GraphStageY") + + +def test_static_ordering_cache(tmp_path, mocker): + sort = mocker.spy(networkx, "topological_sort") + pipeline = Pipeline(STAGES, {"name": "mini"}) + + # Before inputs are connected the order is the insertion order + assert pipeline.stage_names == ["GraphStageZ", "GraphStageY", "GraphStageX"] + assert pipeline.graph.graph[STATIC_ORDERING] == pipeline.stage_names + assert sort.call_count == 1 + + # Connecting the inputs invalidates the cache + pipeline.construct_pipeline_graph(INPUTS, {}) + assert pipeline.stage_names == ["GraphStageX", "GraphStageY", "GraphStageZ"] + assert pipeline.stage_names == ["GraphStageX", "GraphStageY", "GraphStageZ"] + assert sort.call_count == 2 + + # Modifying the returned list does not change the cache + pipeline.stage_names.remove("GraphStageX") + assert pipeline.stage_names == ["GraphStageX", "GraphStageY", "GraphStageZ"] + + # Removing and adding stages invalidates it + pipeline.remove_stage("GraphStageZ") + assert pipeline.stage_names == ["GraphStageX", "GraphStageY"] + pipeline.add_stage({"name": "GraphStageZ"}) + # the new stage is not connected yet, so its position is not fixed + assert sorted(pipeline.stage_names) == ["GraphStageX", "GraphStageY", "GraphStageZ"] + assert sort.call_count == 4 + + # So does trimming + pipeline = Pipeline(STAGES, {"name": "mini"}) + pipeline.stage_names + rc = {"output_dir": str(tmp_path), "from": "GraphStageY"} + pipeline.construct_pipeline_graph(INPUTS, rc) + assert pipeline.stage_names == ["GraphStageY", "GraphStageZ"] + + +def test_flow_chart_structure_only(tmp_path): + pipeline = make_pipeline(tmp_path) + dot_file = str(tmp_path / "chart.dot") + pipeline.make_flow_chart(dot_file) + text = open(dot_file).read() + assert "GraphStageX" in text + assert "StageExecutionConfig" not in text + assert "status" not in text + + +def test_interactive_inputs(): + pipeline = Pipeline.interactive() + pipeline.add_inputs({"gx_in": "gx_in.txt"}) + outputs = pipeline.build_stage(GraphStageX, a=5) + assert outputs == {"gx_out": "./gx_out.txt"} + assert pipeline.GraphStageX.get_input("gx_in") == "gx_in.txt" + assert pipeline.graph.nodes["gx_out"]["path"] == "./gx_out.txt" + + # outputs of earlier stages are used as inputs to later ones, + # unless given explicitly + pipeline.build_stage(GraphStageY) + assert pipeline.GraphStageY.get_input("gx_out") == "./gx_out.txt" + pipeline.build_stage(GraphStageZ, gy_out="elsewhere.txt") + assert pipeline.GraphStageZ.get_input("gy_out") == "elsewhere.txt" + + # Inputs cannot be set for files that a stage generates + with pytest.raises(ValueError): + pipeline.add_inputs({"gx_out": "x.txt"}) + + +def test_mini_status_callback(): + events = [] + pipeline = MiniPipeline( + STAGES, {"name": "mini"}, callback=lambda *args: events.append(args) + ) + + class FakeJob: + name = "GraphStageY" + + pipeline._status_callback("launch", {"job": FakeJob}) + assert pipeline.get_stage_status("GraphStageY") == StageStatus.RUNNING + pipeline._status_callback("completed", {"job": FakeJob}) + assert pipeline.get_stage_status("GraphStageY") == StageStatus.COMPLETE + pipeline._status_callback("fail", {"job": FakeJob}) + assert pipeline.get_stage_status("GraphStageY") == StageStatus.FAILED + # user callback still gets everything + assert [e[0] for e in events] == ["launch", "completed", "fail"] + + +def test_mini_dag(tmp_path): + rc = {"output_dir": str(tmp_path), "log_dir": str(tmp_path), "resume": "restart"} + pipeline = MiniPipeline(STAGES, {"name": "mini"}) + pipeline.initialize(INPUTS, rc, CONFIG) + jobs, _ = pipeline.run_info + dag = pipeline.build_dag(jobs) + deps = {job.name: sorted(j.name for j in parents) for job, parents in dag.items()} + assert deps == { + "GraphStageX": [], + "GraphStageY": ["GraphStageX"], + "GraphStageZ": ["GraphStageY"], + }