diff --git a/python/revng/pypeline/pipeline.py b/python/revng/pypeline/pipeline.py index b236b1dd1..82997413d 100644 --- a/python/revng/pypeline/pipeline.py +++ b/python/revng/pypeline/pipeline.py @@ -20,7 +20,7 @@ from .model import Model, ModelDiff, ReadOnlyModel from .object import ObjectID, ObjectSet from .pipeline_node import PipelineConfiguration, PipelineNode from .schedule.schedule import Schedule -from .schedule.scheduled_task import ScheduledTask +from .schedule.scheduled_task import PipeScheduledTask, SavepointScheduledTask, ScheduledTask from .storage.storage_provider import InvalidatedObjects, ObjectsToInvalidate, SavePointsRange from .storage.storage_provider import StorageProvider from .task.pipe import Pipe @@ -367,9 +367,13 @@ class Pipeline: configuration: PipelineConfiguration, storage_provider: StorageProvider, ) -> Schedule: - tasks: DefaultDictFromKey[PipelineNode, ScheduledTask] = DefaultDictFromKey( - lambda pn: ScheduledTask(pn, model, storage_provider, configuration) - ) + def task_generator(pipeline_node: PipelineNode): + if isinstance(pipeline_node.task, Pipe): + return PipeScheduledTask(pipeline_node, model, storage_provider, configuration) + else: + return SavepointScheduledTask(pipeline_node, model, storage_provider, configuration) + + tasks: DefaultDictFromKey[PipelineNode, ScheduledTask] = DefaultDictFromKey(task_generator) # 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 @@ -673,6 +677,7 @@ class Pipeline: scheduled_tasks: list[ScheduledTask] = [] for task in schedule_dict["tasks"]: pipeline_node: PipelineNode = pipeline_nodes[task["node_id"]] + dependencies = [scheduled_tasks[i] for i in task["dependencies"]] outgoing = Requests() incoming = Requests() @@ -692,6 +697,17 @@ class Pipeline: outgoing[container_declaration] = ObjectSet( container_kind, {obj_id_type.deserialize(x) for x in arg["outgoing"]} ) + + scheduled_tasks.append( + PipeScheduledTask( + pipeline_node, + model, + storage_provider, + configuration, + (incoming, outgoing), + dependencies, + ) + ) elif task["type"] == "SavePoint": assert isinstance(pipeline_node.task, SavePoint) assert task["name"] == pipeline_node.task.name @@ -706,18 +722,18 @@ class Pipeline: outgoing[container_declaration] = ObjectSet( container_kind, {obj_id_type.deserialize(x) for x in container["outgoing"]} ) + + scheduled_tasks.append( + SavepointScheduledTask( + pipeline_node, + model, + storage_provider, + configuration, + (incoming, outgoing), + dependencies, + ) + ) else: raise ValueError(f"Unknown task type: \"{task['type']}\"") - dependencies = [scheduled_tasks[i] for i in task["dependencies"]] - scheduled_task = ScheduledTask( - pipeline_node, - model, - storage_provider, - configuration, - (incoming, outgoing), - dependencies, - ) - scheduled_tasks.append(scheduled_task) - return Schedule(declarations, scheduled_tasks[-1], configuration, model, storage_provider) diff --git a/python/revng/pypeline/schedule/scheduled_task.py b/python/revng/pypeline/schedule/scheduled_task.py index 0720d0bfb..b9d240b50 100644 --- a/python/revng/pypeline/schedule/scheduled_task.py +++ b/python/revng/pypeline/schedule/scheduled_task.py @@ -4,7 +4,8 @@ from __future__ import annotations -from typing import Generator +from abc import ABC, abstractmethod +from typing import Generator, cast from revng.pypeline.container import ContainerDeclaration, ContainerSet from revng.pypeline.model import ReadOnlyModel @@ -17,12 +18,14 @@ from revng.pypeline.task.savepoint import SavePoint from revng.pypeline.task.task import TaskArgumentAccess -class ScheduledTask: +class ScheduledTask(ABC): """ The scheduling works by figuring out which tasks to run, and which dependencies each task can fulfill. - The schedule is stored as a Directed Acyclic Graph (DAG) of ScheduledTask objects. + The schedule is stored as a Directed Acyclic Graph (DAG) of + ScheduledTask objects (which are specialized for the task that will be + executed: PipeScheduledTask and SavepointScheduledTask). While the schedule is originally a path of the pipeline tree, we can apply transformations to it that can make it a DAG. An example of this is kind of @@ -92,6 +95,11 @@ class ScheduledTask: self.incoming.merge(incoming) self.outgoing.merge(outgoing) + @abstractmethod + def _run_impl( + self, containers: ContainerSet, runner_context: RunnerContext + ) -> ScheduledTaskDependencies | None: ... + def run( self, containers: ContainerSet, runner_context: RunnerContext ) -> ScheduledTaskDependencies | None: @@ -103,75 +111,10 @@ class ScheduledTask: self.incoming.check(containers) - task = self.node.task - result: ScheduledTaskDependencies | None = None - for container_declaration in self.disposable_containers: containers[container_declaration].set_is_disposable() - if isinstance(task, SavePoint): - assert ( - self.node.savepoint_range is not None - ), "SavePoint range must be set before calling run on a SavePoint" - task.run( - containers=containers, - incoming=self.incoming, - outgoing=self.outgoing, - configuration_id=self.node.configuration_id(self.configuration), - storage_provider=self.storage_provider, - savepoint_range=self.node.savepoint_range, - ) - # SavePoints do not return any dependencies - elif isinstance(task, Pipe): - bindings = self.node.bindings - pipe_containers = [containers[decl] for decl in bindings] - pipe_incoming = [self.incoming.get(decl) for decl in bindings] - pipe_outgoing = [self.outgoing.get(decl) for decl in bindings] - configuration = self.configuration.get(task, "") - - pipe_output = runner_context.run_pipe( - pipe=task, - file_provider=StorageProviderFileProvider(self.storage_provider), - model=self.model, - containers=pipe_containers, - incoming=pipe_incoming, - outgoing=pipe_outgoing, - configuration=configuration, - ) - - for index, decl in enumerate(bindings): - containers[decl] = pipe_containers[index] - - result_pipe: ObjectDependencies = [] - for index, index_deps in enumerate(pipe_output.dependencies): - container_type = task.signature()[index] - if container_type.access == TaskArgumentAccess.READ: - assert len(index_deps) == 0, ( - "An read only container cannot produce new objects so it can't add " - f"dependencies. For container {container_type.name} got dependencies " - f"{index_deps}" - ) - result_pipe.extend( - (self.node.bindings[index].name, obj, path) for obj, path in index_deps - ) - - # Check if the pipe output - if not all(len(x) == 0 for x in pipe_output.custom_invalidation): - assert task.has_custom_invalidation(), ( - f"Pipe {task.name} returned advanced invalidation data" - "but did not override the 'invalidate' method" - ) - for index, invalidation in enumerate(pipe_output.custom_invalidation): - argument = task.signature()[index] - if argument.access == TaskArgumentAccess.READ: - assert len(invalidation) == 0, ( - f"Pipe {task.name} returned advanced " - "invalidation data for a read-only container" - ) - - result = ScheduledTaskDependencies(result_pipe, pipe_output.custom_invalidation) - else: - raise TypeError(f"Unsupported task type: {type(task)}") + result = self._run_impl(containers, runner_context) self.outgoing.check(containers) self.completed = True @@ -197,3 +140,81 @@ class ScheduledTask: Two instances of the same task will have different ids, and thus different hashes. """ return hash(id(self)) + + +class PipeScheduledTask(ScheduledTask): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert isinstance(self.node.task, Pipe) + + def _run_impl( + self, containers: ContainerSet, runner_context: RunnerContext + ) -> ScheduledTaskDependencies: + task: Pipe = cast(Pipe, self.node.task) + bindings = self.node.bindings + pipe_containers = [containers[decl] for decl in bindings] + pipe_incoming = [self.incoming.get(decl) for decl in bindings] + pipe_outgoing = [self.outgoing.get(decl) for decl in bindings] + configuration = self.configuration.get(task, "") + + pipe_output = runner_context.run_pipe( + pipe=task, + file_provider=StorageProviderFileProvider(self.storage_provider), + model=self.model, + containers=pipe_containers, + incoming=pipe_incoming, + outgoing=pipe_outgoing, + configuration=configuration, + ) + + for index, decl in enumerate(bindings): + containers[decl] = pipe_containers[index] + + result_pipe: ObjectDependencies = [] + for index, index_deps in enumerate(pipe_output.dependencies): + container_type = task.signature()[index] + if container_type.access == TaskArgumentAccess.READ: + assert len(index_deps) == 0, ( + "An read only container cannot produce new objects so it can't add " + f"dependencies. For container {container_type.name} got dependencies " + f"{index_deps}" + ) + result_pipe.extend( + (self.node.bindings[index].name, obj, path) for obj, path in index_deps + ) + + # Check if the pipe output + if not all(len(x) == 0 for x in pipe_output.custom_invalidation): + assert task.has_custom_invalidation(), ( + f"Pipe {task.name} returned advanced invalidation data" + "but did not override the 'invalidate' method" + ) + for index, invalidation in enumerate(pipe_output.custom_invalidation): + argument = task.signature()[index] + if argument.access == TaskArgumentAccess.READ: + assert len(invalidation) == 0, ( + f"Pipe {task.name} returned advanced " + "invalidation data for a read-only container" + ) + + return ScheduledTaskDependencies(result_pipe, pipe_output.custom_invalidation) + + +class SavepointScheduledTask(ScheduledTask): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert isinstance(self.node.task, SavePoint) + + def _run_impl(self, containers: ContainerSet, runner_context: RunnerContext) -> None: + task: SavePoint = cast(SavePoint, self.node.task) + assert ( + self.node.savepoint_range is not None + ), "SavePoint range must be set before calling run on a SavePoint" + task.run( + containers=containers, + incoming=self.incoming, + outgoing=self.outgoing, + configuration_id=self.node.configuration_id(self.configuration), + storage_provider=self.storage_provider, + savepoint_range=self.node.savepoint_range, + ) diff --git a/tests/pypeline/cpp/python_cpp_test.py b/tests/pypeline/cpp/python_cpp_test.py index d5a663544..884907404 100755 --- a/tests/pypeline/cpp/python_cpp_test.py +++ b/tests/pypeline/cpp/python_cpp_test.py @@ -18,7 +18,7 @@ from revng.pypeline.object import Kind, ObjectID, ObjectSet from revng.pypeline.pipeline import AnalysisBinding, Pipeline from revng.pypeline.pipeline_node import PipelineConfiguration, PipelineNode from revng.pypeline.runner_context import RunnerContext -from revng.pypeline.schedule.scheduled_task import ScheduledTask +from revng.pypeline.schedule.scheduled_task import SavepointScheduledTask from revng.pypeline.storage.memory import InMemoryStorageProvider from revng.pypeline.storage.storage_provider import ContainerLocation from revng.pypeline.task.pipe import Pipe @@ -159,7 +159,7 @@ def check_simple_pipeline(): foo_bar_request = Requests({child_cont: ObjectSet.from_list([foo, bar])}) # Pre-populate the storage - task = ScheduledTask( + task = SavepointScheduledTask( begin_node, ReadOnlyModel(model), storage_provider, diff --git a/tests/pypeline/python/test_all.py b/tests/pypeline/python/test_all.py index 07647da0d..97dc44afd 100644 --- a/tests/pypeline/python/test_all.py +++ b/tests/pypeline/python/test_all.py @@ -24,7 +24,7 @@ from revng.pypeline.pipeline import AnalysisBinding, Artifact, ArtifactCategory, from revng.pypeline.pipeline_node import PipelineConfiguration, PipelineNode from revng.pypeline.pipeline_parser import load_pipeline_yaml_file from revng.pypeline.runner_context import RunnerContext -from revng.pypeline.schedule.scheduled_task import ScheduledTask +from revng.pypeline.schedule.scheduled_task import SavepointScheduledTask from revng.pypeline.storage.local_provider import LocalStorageProvider from revng.pypeline.storage.memory import InMemoryStorageProvider from revng.pypeline.storage.storage_provider import ContainerLocation, SavePointsRange @@ -220,7 +220,7 @@ def test_pipeline_inplace(model, storage_provider): container.add_object(MyObjectID(MyKind.CHILD, "one")) container.add_object(MyObjectID(MyKind.CHILD, "two")) begin_configuration_id = begin_node.configuration_id(configuration) - task = ScheduledTask( + task = SavepointScheduledTask( begin_node, ReadOnlyModel(model), storage_provider, @@ -294,7 +294,7 @@ def test_pipeline_up_down(model, storage_provider): container = RootDictContainer() container.add_object(MyObjectID.root()) requests = Requests({root1: root_obj}) - task = ScheduledTask( + task = SavepointScheduledTask( begin_node, ReadOnlyModel(model), storage_provider, @@ -346,7 +346,7 @@ def test_artifact(model, storage_provider): container = ChildDictContainer() container.add_object(MyObjectID(MyKind.CHILD, "one")) requests = Requests({child1: one}) - task = ScheduledTask( + task = SavepointScheduledTask( begin_node, ReadOnlyModel(model), storage_provider, @@ -612,7 +612,7 @@ def test_schedule_serdes(model): container = RootDictContainer() container.add_object(MyObjectID.root()) requests = Requests({root1: root_obj}) - task = ScheduledTask( + task = SavepointScheduledTask( begin_node, ReadOnlyModel(model), storage_provider,