# # This file is distributed under the MIT License. See LICENSE.md for details. # from __future__ import annotations import hashlib from hashlib import sha256 from typing import Annotated, List, Mapping, Optional, Sequence, Set, Union, cast, overload from revng.pypeline.analysis import Analysis from .container import Configuration, ConfigurationId, Container, ContainerDeclaration from .model import ReadOnlyModel from .object import ObjectSet from .storage.file_provider import FileProvider from .storage.storage_provider import SavePointsRange, StorageProvider from .task.pipe import Pipe, PipeDependencies from .task.requests import Requests from .task.savepoint import SavePoint from .task.task import TaskArgument PipelineConfiguration = Annotated[ Mapping[Pipe | Analysis, str], "Configuration for an execution of a schedule or analysis, this describes " "the dynamic configuration of each component of the pipeline", ] class PipelineNode: """ A node of the tree of elements composing a pipeline. Mostly a wrapper around a Task. """ __slots__ = ( "successors", "predecessors", "task", "bindings", "reverse_bindings", "pipe_dependencies", "savepoint_range", "id", "_hash", ) @overload def __init__(self, task: Pipe, bindings: Sequence[ContainerDeclaration]): ... @overload def __init__(self, task: SavePoint): ... def __init__( self, task: Union[Pipe, SavePoint], bindings: Sequence[ContainerDeclaration] | None = None, ): """ The bindings are used to map between the pipeline declarations and the task arguments and are needed only for Pipe tasks as SavePoints only live in the pipeline namespace. """ self.successors: List[PipelineNode] = [] self.predecessors: List[PipelineNode] = [] self.task = task self.pipe_dependencies: Set[Pipe] = set() self.bindings: list[ContainerDeclaration] = [] self.savepoint_range: Optional[SavePointsRange] = None self.id = -1 # This is the hash of the node, it's used to guarantee that # `sorted_successors` is stable. It will be computed at a later time # once the successors of the node are determined. self._hash: str | None = None if bindings is None: assert isinstance( task, SavePoint ), f"A pipeline node with no bindings can only be a Savepoint, got {task}" self.bindings = [] else: assert isinstance(task, Pipe), f"Bindings are only allowed for Pipe tasks, got {task}" if len(bindings) != len(task.arguments): raise ValueError( f"Expected {len(task.arguments)} bindings but got " f"{len(bindings)} for task {task}. arguments: " f"{task.arguments}, bindings: {bindings}" ) for bind, arg in zip(bindings, task.arguments): if bind.container_type != arg.container_type: raise TypeError( f"Binding {bind} is not compatible with argument {arg} " f"for task {task}" ) # These maps the declarations of the pipeline to the task arguments declarations self.bindings = list(bindings) assert len(self.bindings) == len(task.arguments) @property def arguments(self) -> list[TaskArgument]: if isinstance(self.task, SavePoint): # SavePoints do not have bindings, so we return the task arguments directly return self.task.arguments elif isinstance(self.task, Pipe): # For Pipes, we return the bindings, which are the pipeline declarations # that map to the task arguments result = [] for index, argument in enumerate(self.bindings): pipe_arguments = self.task.arguments[index] result.append( TaskArgument(argument.name, argument.container_type, pipe_arguments.access) ) return result else: raise TypeError(f"Unsupported task type: {type(self.task)}") @property def argument_declarations(self) -> list[ContainerDeclaration]: return [x.declaration() for x in self.arguments] def add_successor(self, node: PipelineNode) -> PipelineNode: node.predecessors.append(self) self.successors.append(node) return node def _compute_node_hash(self): # Compute the hash of this node, this requires that the successors of # this node have had their hash pre-computed assert self._hash is None for node in self.successors: assert node._hash is not None hasher = hashlib.sha256() if isinstance(self.task, SavePoint): hasher.update(b"Savepoint\0") elif isinstance(self.task, Pipe): hasher.update(b"Pipe\0") hasher.update(self.task.__class__.name.encode() + b"\0") hasher.update(self.task.static_configuration.encode() + b"\0") hasher.update(b"Arguments\0") for argument in self.argument_declarations: hasher.update(argument.name.encode()) hasher.update(b"\0") hasher.update(argument.container_type.name.encode()) hasher.update(b"\0") for node in self.sorted_successors(): hasher.update(cast(str, node._hash).encode() + b"\0") self._hash = hasher.hexdigest() def sorted_successors(self) -> list[PipelineNode]: """ Return the successors of this node sorted by their deterministic sort key. This is needed to ensure that the savepoint ranges are assigned in a deterministic order, otherwise the caches could cause problems. """ return sorted(self.successors, key=lambda x: cast(str, x._hash)) def prerequisites_for( self, model: ReadOnlyModel, requests: Requests, configuration: PipelineConfiguration, storage_provider: StorageProvider, ) -> Requests: # In order to execute the task we might need to remap the requests # from the pipeline declarations to the task arguments declarations. # So we apply a remap before and after if isinstance(self.task, SavePoint): # SavePoints have a different semantics for prerequisites, so we call the method # from the SavePoint class assert ( self.savepoint_range is not None ), "SavePoint range must be set before calling prerequisites_for on a SavePoint" return self.task.prerequisites_for( requests=requests, configuration_id=self.configuration_id(configuration), storage_provider=storage_provider, savepoint_range=self.savepoint_range, ) elif isinstance(self.task, Pipe): # Use the bindings to map the requests to the pipe arguments pipe_requests: list[ObjectSet] = [requests.get(decl) for decl in self.bindings] # Ask the pipe for its prerequisites outgoing = self.task.prerequisites_for( model=model, requests=pipe_requests, ) # Map back the requests to the pipeline declarations result = requests.clone() for decl, request in zip(self.bindings, outgoing): result[decl] = request return result else: raise TypeError(f"Unsupported task type: {type(self.task)}") def configuration_id(self, configuration: PipelineConfiguration) -> ConfigurationId: hasher = sha256() if isinstance(self.task, Pipe) and self.task in configuration: hasher.update(configuration[self.task].encode()) elif isinstance(self.task, SavePoint): for pipe in sorted(self.pipe_dependencies, key=lambda x: (x.name, x.signature())): hasher.update(configuration.get(pipe, "").encode()) return hasher.hexdigest() def __repr__(self): return f"" def to_dict(self) -> dict: """Convert the data into a dictionary representation.""" return { "id": self.id, "type": "pipe" if isinstance(self.task, Pipe) else "savepoint", "class": self.task.__class__.__name__, "arguments": [binding.name for binding in self.bindings], "successors": sorted([succ.id for succ in self.successors]), "predecessors": sorted([pred.id for pred in self.predecessors]), "savepoint_range": ( { "start": self.savepoint_range.start, "end": self.savepoint_range.end, } if self.savepoint_range else None ), } class DummyPipelineNode(PipelineNode): """Dummy pipeline node, to be used in cases where multiple PipelineNode branches need to be merged into one""" class DummyPipe(Pipe): name = "dummy-pipe" def __init__(self, static_configuration: str): self.static_configuration = "" @classmethod def signature(cls) -> tuple[TaskArgument, ...]: return () def run( self, file_provider: FileProvider, model: ReadOnlyModel, containers: list[Container], incoming: list[ObjectSet], outgoing: list[ObjectSet], configuration: Configuration, ) -> PipeDependencies: return PipeDependencies([]) def __init__(self): # Create a new class object of DummyPipe, with the `name` class # attribute changed with the provided one super().__init__(self.__class__.DummyPipe(""), [])