diff --git a/python/spider-py/src/spider_py/core/__init__.py b/python/spider-py/src/spider_py/core/__init__.py index d417d95e2..26e822ca3 100644 --- a/python/spider-py/src/spider_py/core/__init__.py +++ b/python/spider-py/src/spider_py/core/__init__.py @@ -1 +1,32 @@ """Spider core package.""" + +from .data import Data, DataId +from .task import ( + Task, + TaskId, + TaskInput, + TaskInputData, + TaskInputOutput, + TaskInputValue, + TaskOutput, + TaskOutputData, + TaskOutputValue, + TaskState, +) +from .task_graph import TaskGraph + +__all__ = [ + "Data", + "DataId", + "Task", + "TaskGraph", + "TaskId", + "TaskInput", + "TaskInputData", + "TaskInputOutput", + "TaskInputValue", + "TaskOutput", + "TaskOutputData", + "TaskOutputValue", + "TaskState", +] diff --git a/python/spider-py/src/spider_py/core/data.py b/python/spider-py/src/spider_py/core/data.py new file mode 100644 index 000000000..7a0586712 --- /dev/null +++ b/python/spider-py/src/spider_py/core/data.py @@ -0,0 +1,14 @@ +"""Data module for Spider.""" + +from dataclasses import dataclass +from uuid import UUID + +DataId = UUID + + +@dataclass +class Data: + """Represents a data object.""" + + id: DataId + value: bytes diff --git a/python/spider-py/src/spider_py/core/task.py b/python/spider-py/src/spider_py/core/task.py index 855f2a72b..d9c81e68d 100644 --- a/python/spider-py/src/spider_py/core/task.py +++ b/python/spider-py/src/spider_py/core/task.py @@ -1,5 +1,50 @@ """Task module for Spider.""" +from dataclasses import dataclass, field +from enum import IntEnum +from uuid import UUID, uuid4 +from spider_py.core.data import DataId + +TaskId = UUID + + +@dataclass +class TaskInputOutput: + """Represents a task input that references the output of another task by its ID and position.""" + + task_id: TaskId + position: int + + +TaskInputValue = bytes +TaskInputData = DataId +TaskInput = TaskInputOutput | TaskInputValue | TaskInputData + +TaskOutputValue = bytes +TaskOutputData = DataId +TaskOutput = TaskOutputValue | TaskOutputData + + +class TaskState(IntEnum): + """Represents the state of a task""" + + Pending = 0 + Ready = 1 + Running = 2 + Succeeded = 3 + Failed = 4 + Cancelled = 5 + + +@dataclass class Task: """Represents a task in Spider.""" + + task_id: TaskId = field(default_factory=uuid4) + function_name: str = "" + state: TaskState = TaskState.Pending + timeout: float = 0 + max_retries: int = 0 + task_inputs: list[TaskInput] = field(default_factory=list) + task_outputs: list[TaskOutput] = field(default_factory=list) diff --git a/python/spider-py/src/spider_py/core/task_graph.py b/python/spider-py/src/spider_py/core/task_graph.py new file mode 100644 index 000000000..845bbd7c0 --- /dev/null +++ b/python/spider-py/src/spider_py/core/task_graph.py @@ -0,0 +1,62 @@ +"""TaskGraph module for Spider.""" + +from spider_py.core.task import Task, TaskId + + +class TaskGraph: + """ + Represents a task graph in Spider. + TaskGraph represents a directed acyclic graph (DAG) of tasks. + It stores: + - tasks: A dictionary mapping task ids to Task objects. + - dependencies: A list of tuples representing the dependencies between tasks. Each tuple + contains: + - parent task id + - child task id + - input_tasks: A set of task ids that have no parents (input tasks). + - output_tasks: A set of task ids that have no children (output tasks). + """ + + def __init__(self) -> None: + """Initializes an empty task graph.""" + self.tasks: dict[TaskId, Task] = {} + self.dependencies: list[tuple[TaskId, TaskId]] = [] + self.input_tasks: set[TaskId] = set() + self.output_tasks: set[TaskId] = set() + + def add_task( + self, task: Task, parents: list[TaskId] | None = None, children: list[TaskId] | None = None + ) -> None: + """ + Adds a task to the graph. + :param task: The task to add. + :param parents: The parent ids of the task. Must be already in the task graph. + :param children: The children ids of the task. Must be already in the task graph. + """ + self.tasks[task.task_id] = task + if parents is not None and len(parents) > 0: + for parent in parents: + self.dependencies.append((parent, task.task_id)) + self.output_tasks.discard(parent) + else: + self.input_tasks.add(task.task_id) + if children is not None and len(children) > 0: + for child in children: + self.dependencies.append((task.task_id, child)) + self.input_tasks.discard(child) + else: + self.output_tasks.add(task.task_id) + + def get_parents(self, task_id: TaskId) -> list[Task]: + """ + :param task_id: + :return: Parent tasks of the task identified by `task_id`. + """ + return [self.tasks[parent] for (parent, child) in self.dependencies if child == task_id] + + def get_children(self, task_id: TaskId) -> list[Task]: + """ + :param task_id: + :return: Child tasks of the task identified by `task_id`. + """ + return [self.tasks[child] for (parent, child) in self.dependencies if parent == task_id]