mantatech-sdk 0.5b0.dev65__py3-none-any.whl

manta/apis/graph.py

@@ -0,0 +1,498 @@
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, Iterator, List, Optional, Set, Tuple, Union
+from uuid import UUID
+
+from manta_common.conversions import ID, u64x, x2u64
+from .module import Module, Scheduling
+
+__all__ = ["Graph", "Edge", "Condition", "merge", "Driver", "Task"]
+
+
+class Condition(Enum):
+    FINISHED = "FINISHED"
+    RESULT = "RESULT"
+
+
+class Driver(Enum):
+    HOST = "host"
+    BRIDGE = "bridge"
+    OVERLAY = "overlay"
+
+
+@dataclass(unsafe_hash=True)
+class Edge:
+    """
+    Edge between two module IDs
+
+    Attributes
+    ----------
+    start : UUID
+        Head Module ID
+    end : UUID
+        Tail Module ID
+    condition : Condition
+        Condition for going from head module to tail module
+    """
+
+    start: UUID
+    end: UUID
+    condition: Condition
+
+
+@dataclass
+class Connection:
+    """
+    Class which holds connection before and after a node in the `Graph`
+
+    Attributes
+    ----------
+    previous : Set[Edge]
+        Downstream edges
+    next : Set[Edge]
+        Upstream edges
+    """
+
+    previous: Set[Edge] = field(default_factory=set)
+    next: Set[Edge] = field(default_factory=set)
+
+
+class Graph:
+    """
+    The :code:`Graph` is a task graph where a task is described by a
+    :class:`Task <manta.graph.Task>` in order to be sent to
+    the Manager
+
+    Parameters
+    ----------
+    task_table : Optional[Dict[UUID, Task]]
+        Table for descriptions of :class:`Task <manta.graph.Task>`
+    connectivity_table : Optional[Dict[UUID, Connection]]
+        Table for connectivities between :class:`Task <manta.graph.Task>`
+    """
+
+    __slots__ = ["task_table", "connectivity_table"]
+
+    def __init__(
+        self,
+        task_table: Optional[Dict[UUID, "Task"]] = None,
+        connectivity_table: Optional[Dict[UUID, Connection]] = None,
+    ):
+        self.connectivity_table = connectivity_table or {}
+        self.task_table = task_table or {}
+
+    def add_to_task_table(self, uuid: UUID, task: "Task"):
+        """
+        Add a task description associated to its ID
+
+        Parameters
+        ----------
+        uuid : UUID
+            ID of the task
+        task : Task
+            Description of the task
+
+        Raises
+        ------
+        KeyError
+            If the ID already exists, it cannot be added twice.
+        """
+        if uuid in self.task_table:
+            raise KeyError(f"{uuid} already in the task table.")
+        self.task_table[uuid] = task
+
+    def add_to_connectivity_table(self, edge: Edge):
+        """
+        Add an edge
+
+        Parameters
+        ----------
+        edge : Edge
+            Edge between two IDs
+
+        Raises
+        ------
+        KeyError
+            If one of the IDs into the edge is not in the task table
+        """
+        # Check if IDs are associated to a module
+        if edge.start not in self.task_table:
+            raise KeyError(f"{edge.start} missing in the task table.")
+        if edge.end not in self.task_table:
+            raise KeyError(f"{edge.end} missing in the task table.")
+
+        # Set connection if they don't exist in the connectivity table
+        connection_start = self.connectivity_table.setdefault(edge.start, Connection())
+        connection_end = self.connectivity_table.setdefault(edge.end, Connection())
+
+        # Add edges to them
+        connection_start.next.add(edge)
+        connection_end.previous.add(edge)
+
+    def check_networks(self, networks: Dict[str, str]):
+        """
+        Check if all network names are found in :code:`networks`
+
+        Parameters
+        ----------
+        networks : Dict[str, str]
+            Networks from the :code:`Swarm`
+
+        Raises
+        ------
+        RuntimeError
+            Network name in a :code:`Task` was not found in :code:`networks`
+        """
+        for task in self.task_table.values():
+            if (network := task.network) and network not in networks:
+                raise RuntimeError(
+                    f"Network {network} not found in declared swarm networks"
+                )
+
+    def starting_uuids(self) -> Iterator[UUID]:
+        """
+        Yield the starting UUIDs of the graph
+
+        Returns
+        -------
+        Iterator[UUID]
+           IDs of starting nodes in the graph
+        """
+        for task_uuid, connection in self.connectivity_table.items():
+            if len(connection.previous) == 0:
+                yield task_uuid
+
+    def propagation(
+        self, acc: list, uuid: UUID, visited_uuids: set, order: int
+    ) -> Tuple[List[Dict[str, Any]], set]:
+        """
+        Recursive function which explore a UUID in the graph and returns a succession
+        of tasks ordered by the definition of the graph
+
+        Parameters
+        ----------
+        acc : list
+            Output
+        uuid : UUID
+            UUID to explore
+        visited_uuids : set
+            Already visited uuids
+        order : int
+            Depth variable
+
+        Returns
+        -------
+        Tuple[list, set]
+            :code:`(acc, visited_uuids)` updated
+        """
+        connection = self.connectivity_table[uuid]
+        task = self.task_table[uuid]
+
+        task_dict = {
+            "task_id": ID(uuid.hex).oid,
+            "module": task.module,
+            "scheduling": task.scheduling.to_proto(),
+            "previous_tasks": [ID(edge.start.hex).oid for edge in connection.previous],
+            "next_tasks": [ID(edge.end.hex).oid for edge in connection.next],
+            "conditions": [str(edge.condition) for edge in connection.next],
+            "order": order,
+            "network": task.network,
+            "alias": task.alias,
+        }
+
+        acc.append(task_dict)
+
+        next_uuids = [edge.end for edge in connection.next]
+        visited_uuids.add(uuid)
+        if len(next_uuids) != 0:
+            for uuid in next_uuids:
+                if uuid not in visited_uuids:
+                    acc, visited_uuids = self.propagation(
+                        acc, uuid, visited_uuids, order + 1
+                    )
+        return acc, visited_uuids
+
+    def to_proto(self) -> List[Dict[str, Any]]:
+        """
+        Make the graph as a protobuf structure
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            List of tasks
+        """
+        acc: List[Dict[str, Any]] = []
+        visited_uuids = set()
+
+        if len(self.connectivity_table) == 0:
+            assert len(self.task_table) == 1, (
+                "Swarm without connectivity must have only one task."
+            )
+            main_uuid = next(iter(self.task_table))
+            self.connectivity_table[main_uuid] = Connection()
+
+        for starting_uuid in self.starting_uuids():
+            acc, visited_uuids = self.propagation(acc, starting_uuid, visited_uuids, 0)
+
+        return acc
+
+    def __str__(self) -> str:
+        """
+        Organize the graph structure into a string by showing the UUIDs of the
+        tasks and their connections.
+
+        It will print the UUID of each task and how they are connected through the
+        `connectivity_table`, showing both the previous and next connections. Additionally,
+        it will display the associated network for each task (if applicable).
+        """
+
+        def alias_or_uuid(uuid):
+            task = self.task_table[uuid]
+            return str(task.alias or uuid)
+
+        if len(self.connectivity_table) == 0:
+            return "Graph()"
+
+        string = "Graph(\n"
+        for task_uuid, connection in self.connectivity_table.items():
+            string += (
+                f"    {alias_or_uuid(task_uuid)!r}: (\n"
+                f"        previous : {[alias_or_uuid(edge.start) for edge in connection.previous]},\n"
+                f"        next     : {[alias_or_uuid(edge.end) for edge in connection.next]},\n"
+                "    ),\n"
+            )
+        string += ")"
+        return string
+
+
+def merge(graph_list: List[Graph]) -> Graph:
+    """
+    Merge a list of graphs into one graph
+
+    Parameters
+    ----------
+    graph_list : List[Graph]
+        List of graphs to be merged
+
+    Returns
+    -------
+    Graph
+        Merged graph
+
+    Raises
+    ------
+    ValueError
+        If two graphs contain the same :code`TaskID`
+    """
+    task_table = {}
+    connectivity_table = {}
+    for graph in graph_list:
+        # if task_table.keys() & graph.task_table.keys():
+        #     print(task_table.keys())
+        #     print(graph.task_table.keys())
+        #     raise ValueError("Not allow to merge two graphs with same `IDs`")
+        task_table.update(graph.task_table)
+        connectivity_table.update(graph.connectivity_table)
+
+    # print(task_table)
+    # print(connectivity_table)
+
+    return Graph(task_table, connectivity_table)
+
+
+def fold(graph_list: List[Graph]) -> Graph:
+    """
+    Like :code:`merge` function, it generates a graph from a list
+    of graphs. However similarities, such as same `TaskIDs` are
+    put together
+
+    Parameters
+    ----------
+    graph_list : List[Graph]
+        List of graphs to be merged
+
+    Returns
+    -------
+    Graph
+        Merged graph based on similarities
+    """
+    task_table = {}
+    connectivity_table = {}
+    for graph in graph_list:
+        task_table.update(graph.task_table)
+        for edge, connection in graph.connectivity_table.items():
+            local_connection = connectivity_table.setdefault(edge, Connection())
+            local_connection.previous |= connection.previous
+            local_connection.next |= connection.next
+
+    return Graph(task_table, connectivity_table)
+
+
+class Task:
+    """
+    Class which holds the description of a task
+
+    Parameters
+    ----------
+    python_program : Union[str, Path]
+        Python program code
+    image : str
+        Image name of the container in which the program will be executed
+    method : str
+        Currently, only :code:`"any"` is available
+    fixed : bool
+        If :code:`True`, the chosen nodes must remain the same over the execution of the :code:`Swarm`
+    excepted_ids : Optional[list]
+        List of ids that must not be chosen for this module
+    specified_ids : Optional[list]
+        List of ids that must be chosen for this module
+    maximum : Optional[Union[int, float]]
+        Maximum of ids that can be chosen for this module.
+        If it is a float, it is a percentage between :code:`0.0` and :code:`1.0`
+    network : Optional[str]
+        It indicates in which network the task will be executed.
+        Network name must exist in :code:`Swarm.networks`.
+    gpu : bool
+        If :code:`True`, the module will be executed on GPU
+
+    Attributes
+    ----------
+    module : Module
+        Description of a task
+    graph : Graph
+        Hold the connectivity with previous tasks
+    uuid : UUID
+        ID of the task
+
+    Examples
+    --------
+
+    >>> aggregator = Task(
+    ...     module=Module("modules/aggregator.py", "manta-demo:latest"),
+    ...     method="any",
+    ...     fixed=False,
+    ...     maximum=1,
+    ...     alias="aggregator",
+    ... )
+    >>> worker = Task(module="worker_module_id")
+    """
+
+    def __init__(
+        self,
+        module: Module,
+        method: str = "all",
+        fixed: bool = False,
+        excepted_ids: Optional[list] = None,
+        specified_ids: Optional[list] = None,
+        maximum: Optional[Union[int, float]] = -1,
+        network: Optional[str] = None,
+        gpu: bool = False,
+        graph: Optional[Graph] = None,
+        alias: Optional[str] = None,
+    ):
+        self.module = module
+        self.scheduling = Scheduling(
+            method=method,
+            fixed=fixed,
+            excepted_ids=excepted_ids,
+            specified_ids=specified_ids,
+            maximum=maximum,
+            gpu=gpu,
+        )
+        self.network = network
+        self.alias = alias or self.module.name
+
+        self.graph = Graph() if graph is None else graph
+        self.uuid = uuid.uuid4()
+
+        # If this task has been initialized with module data,
+        # immediately register it in the graph
+        self.take_in(self.module)
+
+    def take_in(self, module: Module):
+        """
+        Add the module definition of the task to its graph
+
+        Parameters
+        ----------
+        module : Module
+            Module of a task
+        """
+        if module.name is not None:
+            self.uuid = uuid.UUID(hex=u64x(x2u64(module.name)))
+        self.graph.add_to_task_table(self.uuid, self)
+
+    def add_edge(self, uuid: uuid.UUID, condition: Condition):
+        """
+        Add an edge to its graph where it is executed by the Manager
+        given the :code:`Condition`
+
+        Parameters
+        ----------
+        uuid : uuid.UUID
+            ID of the task
+        condition : Condition
+            Condition which tells to the Manager when next task should be executed
+        """
+        self.graph.add_to_connectivity_table(Edge(uuid, self.uuid, condition))
+
+    def __call__(
+        self,
+        task: Optional[Union[List["Task"], "Task"]] = None,
+        condition: Condition = Condition.FINISHED,
+    ) -> "Task":
+        """
+        Update its own graph with connections to other modules
+
+        Parameters
+        ----------
+        task : Optional[Union[List[Task], Task]]
+            One or multiple tasks to connect with itself
+        condition : Condition
+            Condition which tells to the Manager when next task should be executed
+
+        Returns
+        -------
+        Task
+            Return itself with updated graph
+
+        Examples
+        --------
+
+        >>> aggregator = Task(
+        ...     Module(
+        ...         "modules/aggregator.py",
+        ...         "manta-demo:latest",
+        ...     ),
+        ...     method="any",
+        ...     fixed=False,
+        ...     maximum=1,
+        ...     alias="aggregator",
+        ... )
+        >>> worker = Task(
+        ...     Module(
+        ...         "modules/worker",
+        ...         "manta-demo:latest",
+        ...     ),
+        ...     method="all",
+        ...     fixed=False,
+        ...     maximum=-1,
+        ...     alias="worker",
+        ... )
+        >>> m = aggregator(worker) # start worker and follow by aggregator
+        """
+        if task is None:
+            return self
+        elif isinstance(task, Task):
+            self.graph = merge([self.graph, task.graph])
+            self.add_edge(task.uuid, condition)
+        else:
+            # Handle list of tasks
+            self.graph = merge([self.graph] + [t.graph for t in task])
+            for t in task:
+                self.add_edge(t.uuid, condition)
+        return self
+
+    def __str__(self):  # pragma: no cover
+        return f"Task({self.uuid})"