pypeline: split ScheduledTask

Due to the drifting differences, split `ScheduledTask` into
`PipeScheduledTask` and `SavepointScheduledTask`, which allows each to
specialize on the task it has to accomplish.

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)

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,
+        )

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,

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,