# # This file is distributed under the MIT License. See LICENSE.md for details. # from __future__ import annotations from dataclasses import dataclass from itertools import chain from typing import Dict, Generator, Generic, List, Mapping, Optional, Set, TypeVar import yaml from .analysis import AnalysisBinding from .container import Container, ContainerDeclaration from .graph import Graph from .model import Model, ReadOnlyModel from .object import ObjectID, ObjectSet from .pipeline_node import PipelineConfiguration, PipelineNode from .schedule.schedule import Schedule from .schedule.scheduled_task import ScheduledTask from .storage.storage_provider import InvalidatedObjects, SavePointsRange, StorageProvider from .task.pipe import Pipe from .task.requests import Requests from .task.savepoint import SavePoint from .task.task import TaskArgumentAccess from .utils.default_dict_from_key import DefaultDictFromKey from .utils.logger import pypeline_logger from .utils.registry import get_singleton @dataclass(frozen=True, slots=True) class Artifact: """ An artifact is a container in a certain point of the pipeline with some extra metadata, such as a name. It's designed to mark interesting results in the pipeline, so that users an easily obtain them. """ name: str node: PipelineNode container: ContainerDeclaration description: Optional[str] = None def is_cacheable(self) -> bool: """An artifact is cacheable if it's backed by a savepoint.""" return isinstance(self.node.task, SavePoint) C = TypeVar("C", bound=Model) class Pipeline(Generic[C]): """ A pipeline is a tree of tasks. Given a set of requests, a model and a configuration of the pipes, it can produce a schedule that fulfills the requests. """ __slots__ = ("declarations", "root", "artifacts", "analyses", "savepoint_id_to_artifact") def __init__( self, declarations: set[ContainerDeclaration], root: PipelineNode, artifacts: Optional[set[Artifact]] = None, analyses: Optional[set[AnalysisBinding]] = None, ): self.root = root self.declarations = set(declarations) self.savepoint_id_to_artifact: dict[int, Artifact] = {} """ """ self.artifacts: Mapping[str, Artifact] = {} """ The artifacts, indexed by their name for easy access. """ for artifact in artifacts or set(): if artifact.name in self.artifacts: raise ValueError( f"Artifact {artifact.name} is defined multiple times in the pipeline" ) self.artifacts[artifact.name] = artifact self.analyses: Mapping[str, AnalysisBinding] = {} """ The analyses, indexed by their name for easy access. """ for analysis in analyses or set(): if analysis.analysis.name in self.analyses: raise ValueError( f"Analysis {analysis.analysis.name} is defined multiple times in the pipeline" ) self.analyses[analysis.analysis.name] = analysis Pipeline.assign_savepoint_ranges(root) id_index = 0 # Set the dependencies field for node in self.walk_pipeline(stable=True): node.id = id_index id_index += 1 if node is self.root: node.pipe_dependencies = set() continue node.pipe_dependencies = set( chain.from_iterable(n.pipe_dependencies for n in node.predecessors) ) if isinstance(node.task, Pipe): node.pipe_dependencies.add(node.task) for name, artifact in self.artifacts.items(): self.savepoint_id_to_artifact[artifact.node.id] = artifact if name != artifact.name: raise ValueError( f"Artifact name {artifact.name} does not match the key " f"{name} in the artifacts map." ) @staticmethod def assign_savepoint_ranges(node: PipelineNode, current_id: int = 0) -> int: """ Assigns savepoint ranges to the nodes in the pipeline tree. A savepoint range is a pair of integers that represent the start and end of the savepoint in a subtree. The idea is to deduplicate things common to a savepoint and all its children, so that we can efficiently represent a subtree of savepoints as a continuous range of integers. """ # We ID only the savepoints, so we increment the ID only if the current node is a savepoint. # This implies that, while the ids start at 0, the first savepoint will have id 1. # This is needed in the case the first node is a pipe, which is not a savepoint if isinstance(node.task, SavePoint): current_id += 1 # Save the id on preorder visit start = current_id # Recurse on the children, but do it in a deterministic way end = current_id for child in node.sorted_successors(): current_id = Pipeline.assign_savepoint_ranges(child, current_id) assert child.savepoint_range is not None, ( f"Child {child.task.name} does not have a savepoint range assigned:" f" {child.savepoint_range}" ) end = max(end, child.savepoint_range.end) # The end of the savepoint range is assigned on postorder visit. # Therefore the end is inclusive assert node.savepoint_range is None, ( f"SavePoint {node.task.name} already has a savepoint range assigned" ": {node.savepoint_range}" ) node.savepoint_range = SavePointsRange(start, end) # Return the new id return current_id def walk_pipeline( self, start: Optional[PipelineNode] = None, forward: bool = True, stable: bool = False ) -> Generator[PipelineNode, None, None]: """BFS walk of pipeline nodes""" assert int(not forward) + int(stable) < 2, "forward=False,stable=True is unsupported" to_visit: List[PipelineNode] = [start or self.root] visited: Set[PipelineNode] = set() if forward: def successors(node): return node.sorted_successors() if stable else node.successors else: def successors(node): return node.predecessors while len(to_visit) > 0: node = to_visit.pop() yield node visited.add(node) for child_node in successors(node): if child_node not in visited: to_visit.append(child_node) def graph(self, container_edges=False) -> Graph: """A graph for debugging purposes.""" graph = Graph() nodes_map: Dict[PipelineNode, Graph.Node] = {} access_to_str = { TaskArgumentAccess.READ: "R", TaskArgumentAccess.WRITE: "W", TaskArgumentAccess.READ_WRITE: "RW", } for node in self.walk_pipeline(): if isinstance(node.task, Pipe): new_node = Graph.Node(node.task.name) for argument in node.arguments_with_access: new_node.entries.append(f"{argument.name} [{access_to_str[argument.access]}]") else: new_node = Graph.Node(node.task.name, color="#2A52BE", bgcolor="#6C83BE") for argument in node.arguments_with_access: new_node.entries.append(argument.name) nodes_map[node] = new_node graph.nodes.add(new_node) for predecessor in node.predecessors: graph.edges.add( Graph.Edge( nodes_map[predecessor], new_node, source_port=-1, destination_port=-1, style="bold", ) ) if not container_edges: continue node_inputs: list[ContainerDeclaration] = list(node.arguments) taken_inputs: set[int] = set() for parent_node in self.walk_pipeline(node, forward=False): if parent_node is node: continue for source_index, parent_argument in enumerate(parent_node.arguments_with_access): parent_container_decl = parent_argument.to_container_decl() if ( parent_argument.access == TaskArgumentAccess.READ or parent_container_decl not in node_inputs ): continue destination_index = node_inputs.index(parent_container_decl) if destination_index in taken_inputs: continue taken_inputs.add(destination_index) if parent_node in node.predecessors: continue graph.edges.add( Graph.Edge( nodes_map[parent_node], new_node, source_port=source_index, destination_port=destination_index, style="dashed", color="#FD6D53", ) ) return graph def schedule( self, model: ReadOnlyModel, target_node: PipelineNode, requests: Requests, pipeline_configuration: PipelineConfiguration, storage_provider: StorageProvider, ) -> Schedule: tasks: DefaultDictFromKey[PipelineNode, ScheduledTask] = DefaultDictFromKey(ScheduledTask) # The pipeline is a tree, so we can just unroll the predecessors, # When we parallelize, we will make a subclass that overrides this method, # and probably it will first call it to produce the initial schedule and # then add the parallelization logic node: PipelineNode = target_node node_outgoing_requests = requests while not node_outgoing_requests.empty(): pypeline_logger.debug_log(f"Scheduling node {node}") pypeline_logger.debug_log(f"Outgoing requests: {node_outgoing_requests}") orig = repr(node_outgoing_requests) # Each node should remove the requests it can handle # and add the requests it needs to satisfy the task node_ingoing_requests = node.prerequisites_for( model=model, requests=node_outgoing_requests, pipeline_configuration=pipeline_configuration, storage_provider=storage_provider, ) assert orig == repr(node_outgoing_requests), ( f"Node {node} modified the outgoing requests, which is not allowed. " f"Original: {orig}, modified: {node_outgoing_requests}" ) pypeline_logger.debug_log(f"Computed Ingoing requests: {node_ingoing_requests}") # Store the computed requests so that we can use them in the run method tasks[node].request(node_ingoing_requests, node_outgoing_requests) # If the node has no predecessors, we are done, but # we need to check that it has no requests left if not node.predecessors: assert node_ingoing_requests.empty(), ( f"Node {node} has no predecessors, but it still has " f"requests: {node_ingoing_requests.minimize()}" ) break assert len(node.predecessors) == 1, ( f"Node {node} has multiple predecessors, but we assume a tree structure. " f"Predecessors: {node.predecessors}" ) predecessor = node.predecessors[0] tasks[node].depends_on.append(tasks[predecessor]) # Recurse on THE predecessor, its outgoing requests will be the # ingoing requests of the current node node = predecessor node_outgoing_requests = node_ingoing_requests # Here the task list has been decided, these are now iterated to apply # the following optimizations: # * Reduce the set of container declarations to those that are actually # necessary to run the schedule # * Compact incoming and outgoing of each task to their reduced version # * Skip tasks where nothing in outgoing is actually written by the task # # These could be derived by using a dataflow-based approach to the # pipeline, where each pipe binds its inputs to the predecessor's pipe # output, but in the general case (with RW pipes) this turns into # having a mini-programming language and applying SSA analysis and the # like to deduplicate the containers (that, in the general case, need # to be duplicated at each node of the dataflow). scheduled_task: ScheduledTask | None = tasks[target_node] parent_scheduled_task: ScheduledTask | None = None used_declatations: set[str] = set() while scheduled_task is not None: # Assume that the schedule is a straight line, so a task has at # most one dependency assert len(scheduled_task.depends_on) in (0, 1) # Reduce incoming and outgoing scheduled_task.incoming = scheduled_task.incoming.minimize() scheduled_task.outgoing = scheduled_task.outgoing.minimize() # Figure out if a task will actually produce outputs written_out = Requests() for argument in scheduled_task.node.arguments_with_access: if argument.access != TaskArgumentAccess.READ: container_decl = argument.to_container_decl() if container_decl in scheduled_task.outgoing: written_out[container_decl] = scheduled_task.outgoing[container_decl] if written_out.empty() and parent_scheduled_task is not None: # If here the task does not actually need to produce anything # and can be skipped, by "glueing" its dependencies onto the parent assert [scheduled_task] == parent_scheduled_task.depends_on parent_scheduled_task.depends_on = scheduled_task.depends_on else: used_declatations.update(x.name for x in scheduled_task.node.arguments) parent_scheduled_task = scheduled_task if len(scheduled_task.depends_on) == 1: scheduled_task = scheduled_task.depends_on[0] else: scheduled_task = None return Schedule( {v for v in self.declarations if v.name in used_declatations}, tasks[target_node], pipeline_configuration, ) def get_artifact( self, model: ReadOnlyModel, artifact: Artifact, requests: ObjectSet, pipeline_configuration: PipelineConfiguration, storage_provider: StorageProvider, ) -> Container: schedule = self.schedule( model=model, target_node=artifact.node, requests=Requests({artifact.container: requests}), pipeline_configuration=pipeline_configuration, storage_provider=storage_provider, ) return schedule.run( model=model, storage_provider=storage_provider, )[artifact.container] def run_analysis( self, model: ReadOnlyModel, analysis_name: str, requests: Requests, analysis_configuration: str, pipeline_configuration: PipelineConfiguration, storage_provider: StorageProvider, ) -> tuple[Model, InvalidatedObjects]: """ Run an analysis on the pipeline, given a model and a set of requests. The analysis will return the new potentially modified model, and set it in the storage provider. """ if analysis_name not in self.analyses: raise ValueError(f"Analysis {analysis_name} not found in the pipeline") analysis_info = self.analyses[analysis_name] if len(requests) != len(analysis_info.bindings): raise ValueError( f"Expected {len(analysis_info.bindings)} requests for analysis " f"{analysis_name}, but got {len(requests)}: {requests}" ) for req in requests: if req not in analysis_info.bindings: raise ValueError( f"Request {req} but it's not compatible with in the " f"analysis bindings: {analysis_info.bindings}" ) schedule = self.schedule( model=model, target_node=analysis_info.node, requests=requests, pipeline_configuration=pipeline_configuration, storage_provider=storage_provider, ) all_containers = schedule.run(model=model, storage_provider=storage_provider) # The analysis modifies the model, but we need to invalidate # the changes, so we make a clone of the model. # This also allows to keep the original model intact # in case the analysis fails new_model = model.clone() analysis_info.analysis.run( model=new_model, containers=[all_containers[decl] for decl in analysis_info.bindings], incoming=[requests.get(decl) for decl in analysis_info.bindings], configuration=analysis_configuration, ) invalidated = storage_provider.invalidate(ReadOnlyModel(new_model).diff(model)) storage_provider.set_model(new_model.serialize()) return new_model, invalidated def deserialize_schedule(self, schedule: str) -> Schedule: schedule_dict = yaml.safe_load(schedule) declarations = set() for container in schedule_dict["containers"]: for declaration in self.declarations: if container["name"] == declaration.name: declarations.add(declaration) break else: raise ValueError() pipeline_nodes = list(self.walk_pipeline(stable=True)) container_map = {x.name: x for x in self.declarations} obj_id_type = get_singleton(ObjectID) # type: ignore[type-abstract] configuration = {} scheduled_tasks: list[ScheduledTask] = [] for task in schedule_dict["tasks"]: pipeline_node: PipelineNode = pipeline_nodes[task["node_id"]] outgoing = Requests() incoming = Requests() if task["type"] == "Pipe": assert isinstance(pipeline_node.task, Pipe) assert task["name"] == pipeline_node.task.name configuration[pipeline_node.task] = task["dynamic_config"] for index, arg in enumerate(task["args"]): container_declaration = pipeline_node.bindings[index] assert container_declaration.name == arg["name"] container_kind = container_declaration.container_type.kind incoming[container_declaration] = ObjectSet( container_kind, {obj_id_type.deserialize(x) for x in arg["incoming"]} ) outgoing[container_declaration] = ObjectSet( container_kind, {obj_id_type.deserialize(x) for x in arg["outgoing"]} ) elif task["type"] == "SavePoint": assert isinstance(pipeline_node.task, SavePoint) assert task["name"] == pipeline_node.task.name for container in task["containers"]: container_declaration = container_map[container["name"]] container_kind = container_declaration.container_type.kind incoming[container_declaration] = ObjectSet( container_kind, {obj_id_type.deserialize(x) for x in container["incoming"]} ) outgoing[container_declaration] = ObjectSet( container_kind, {obj_id_type.deserialize(x) for x in container["outgoing"]} ) else: raise ValueError(f"Unknown task type: \"{task['type']}\"") depends_on = [scheduled_tasks[i] for i in task["depends_on"]] scheduled_task = ScheduledTask(pipeline_node, False, outgoing, incoming, depends_on) scheduled_tasks.append(scheduled_task) return Schedule(declarations, scheduled_tasks[-1], configuration)