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

manta/apis/swarm.py

@@ -0,0 +1,206 @@
+import io
+from collections.abc import Sequence
+from typing import Any, Dict, Iterable, List, Union
+
+from manta_common.build.common.results import Global
+from manta_common.build.common.system import Network, Networks
+from manta_common.const import CHUNK_SIZE
+from manta_common.conversions import ID, dict_to_bytes
+from .graph import Driver, Task, fold
+
+
+class Swarm:
+    """
+    This class defines a decentralized algorithm which must be sent
+    to a cluster
+
+    Attributes
+    ----------
+    globals : Dict[str, Any]
+        Initial global values
+    networks : Dict[str, str]
+        Networks used for peer-to-peer communications
+    """
+
+    name: str = "Swarm"
+
+    def __init__(self):
+        self.globals = {}
+        self.networks = {}
+
+    def execute(self) -> Union[Task, Sequence[Task]]:
+        """
+        This method must be defined by the User
+
+        Returns
+        -------
+        Task
+            Task which is the last Task executed in the graph
+            before iterating
+
+        Examples
+        --------
+
+
+        >>> def execute(self):
+        ...     \"\"\"
+        ...     Generation of the task graph
+        ...
+        ...     +--------+     +------------+     +-----------+ if has_converged
+        ...     | Worker | --> | Aggregator | --> | Scheduler | ----------------> END PROGRAM
+        ...     +--------+     +------------+     +-----------+
+        ...         |                                   | else
+        ...         +--<<<----------<<<----------<<<----+
+        ...     \"\"\"
+        ...     m = self.worker()
+        ...     m = self.aggregator(m)
+        ...     return self.scheduler(m)
+        """
+        raise NotImplementedError(
+            f'Swarm [{type(self).__name__}] is missing the required "Swarm" function'
+        )
+
+    def set_global(self, tag: str, value: Union[dict, bytes]):
+        """
+        Set a global value in the swarm
+
+        Examples
+        --------
+
+        >>> self.set_global(
+        ...     "hyperparameters",
+        ...     {
+        ...         "epochs": 1,
+        ...         "batch_size": 32,
+        ...         "loss": "CrossEntropyLoss",
+        ...         "loss_params": {},
+        ...         "optimizer": "SGD",
+        ...         "optimizer_params": {"lr": 0.01, "momentum": 0.9},
+        ...     },
+        ... )
+        """
+        self.globals[tag] = value
+
+    def add_network(self, name: str, driver: Driver = Driver.OVERLAY):
+        """
+        Add network in the swarm
+
+        Parameters
+        ----------
+        name : str
+            Name of the network
+        driver : Driver
+            Driver of the network
+
+        Raises
+        ------
+        KeyError
+            Name of network must be unique
+
+        Examples
+        --------
+
+        >>> self.add_network(
+        ...     "docker_network",
+        ...     Driver.BRIDGE,
+        ... )
+        """
+        if name in self.networks:
+            driver = self.networks[name]
+            raise KeyError(f"Name of network must be unique (found {name} => {driver})")
+        self.networks[name] = driver.value
+
+    def prepare_graph(self) -> List[Dict[str, Any]]:
+        """
+        Prepares the graph of the Swarm to a protobuf object
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            List of tasks
+        """
+        output = self.execute()
+        if isinstance(output, Sequence) and all(
+            map(lambda obj: isinstance(obj, Task), output)
+        ):
+            graph = fold([module.graph for module in output])
+        elif isinstance(output, Task):
+            graph = output.graph
+        else:
+            raise ValueError(
+                "Invalid output type. It must be 'Union[Task, Sequence[Task]]'."
+            )
+        graph.check_networks(self.networks)
+        return graph.to_proto()
+
+    def prepare_globals(self, swarm_id: str) -> Iterable[Iterable[Global]]:
+        """
+        Generates an iterable of asynchronous generators where each generator
+        yields global chunks of a global value.
+
+        Parameters
+        ----------
+        swarm_id : str
+            Swarm ID
+        Iterable[Iterable[Global]]
+            Iterable of asynchronous generators of global chunks
+        """
+        for tag, values in self.globals.items():
+            buffer = io.BytesIO(dict_to_bytes(values))
+
+            def generator():
+                while chunk := buffer.read(CHUNK_SIZE):
+                    yield Global(ID(swarm_id).oid, tag=tag, data=chunk)
+
+            yield generator()
+
+    def prepare_networks(self, swarm_id: str) -> Networks:
+        """
+        Casts networks attribute as protobuf object
+
+        Parameters
+        ----------
+        swarm_id : str
+            Swarm ID
+
+        Returns
+        -------
+        Networks
+            Networks protobuf object
+        """
+        return Networks(
+            ID(swarm_id).oid,
+            [Network(name, driver) for name, driver in self.networks.items()],
+        )
+
+    def __str__(self) -> str:
+        """
+        Convert the Swarm to string
+
+        Returns
+        -------
+        str
+            Swarm formatted into string
+        """
+        # TODO: use pprint
+        try:
+            execution = self.execute()
+            if isinstance(execution, Task):
+                graph = execution.graph
+            else:
+                graph = fold([task.graph for task in execution])
+            string = (
+                "Swarm(\n    "
+                + "\n    ".join(str(graph).split("\n"))
+                + f",\n    Networks({self.networks}),\n"
+                + f"    Globals(keys: {list(self.globals.keys())}),\n"
+                + ")"
+            )
+        except NotImplementedError:
+            string = (
+                "Swarm(\n"
+                + f"    Networks({self.networks}),\n"
+                + f"    Globals(keys: {list(self.globals.keys())}),\n"
+                + ")"
+            )
+        return string