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

manta/light/world.py

@@ -0,0 +1,177 @@
+import logging
+import os
+from typing import List, Optional
+
+from ..clients.world_client import WorldClient
+from manta_common.build.common.system import EnvId
+from manta_common.build.common.tasks import TaskScheduling
+from manta_common.conversions import ID
+from manta_common.event_loop import EventLoopManager
+from .globals import Globals
+from .logging_config import configure_logging
+from .results import Results
+
+__all__ = ["World"]
+
+configure_logging()
+
+
+class World:
+    """
+    World service for accessing values with other nodes
+
+    The :code:`World` service allows the user to interact with collective values
+    during a :class:`Swarm <manta.swarm.Swarm>` iteration
+
+    Parameters should not be fulfill manually. Instead, it is managed during the
+    deployment of the task by the :code:`Node`.
+
+    Parameters
+    ----------
+    host : Optional[str]
+        Manager host
+    port : Optional[int]
+        Manager port
+    swarm_id : Optional[ID]
+        Swarm ID
+    task_id : Optional[ID]
+        Task ID
+    """
+
+    __slots__ = [
+        "swarm_id",
+        "task_id",
+        "logger",
+        "results",
+        "globals",
+        "world_client",
+        "loop_manager",
+    ]
+
+    def __init__(
+        self,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+        swarm_id: Optional[ID] = None,
+        task_id: Optional[ID] = None,
+    ):
+        # Retrieve environment variables for RPC host and port
+        host = host or os.getenv("RPC_HOST", "host.docker.internal")
+        port = int(port or os.getenv("RPC_PORT", 50051))
+        self.task_id = task_id or ID(os.getenv("TASK_ID"))
+        self.swarm_id = swarm_id or ID(os.getenv("SWARM_ID"))
+
+        self.logger = logging.getLogger(__name__)
+
+        self.logger.info(
+            f"World configured with Host: {host}, Port: {port},"
+            f" Task ID: {self.task_id},"
+            f" Swarm ID: {self.swarm_id}"
+        )
+
+        self.world_client = WorldClient(host=host, port=port)
+        self.loop_manager = EventLoopManager.get_instance()
+
+        self.results = Results(
+            world_client=self.world_client,
+            swarm_id=self.swarm_id,
+            task_id=self.task_id,
+        )
+        self.globals = Globals(
+            world_client=self.world_client, swarm_id=self.swarm_id, task_id=self.task_id
+        )
+
+    def schedule_next_iter(self, node_ids: List[str], task_to_schedule_alias: str):
+        """
+        Schedule the next iteration
+
+        Parameters
+        ----------
+        node_ids : List[str]
+            The list of node ids to schedule the next iteration for
+        task_to_schedule_alias : str
+            The task to schedule for the next iteration
+
+        Examples
+        --------
+
+        From :class:`Task <manta_light.task.Task>` attribute:
+
+        >>> self.world.schedule_next_iter(
+        ...     node_ids=selected_nodes, task_to_schedule_alias="worker"
+        ... )
+        """
+        self.loop_manager.run_coroutine(
+            self.async_schedule_next_iter(node_ids, task_to_schedule_alias)
+        )
+
+    def stop_swarm(self):
+        """
+        Stop the swarm
+
+        Examples
+        --------
+
+        From :class:`Task <manta_light.task.Task>` attribute:
+
+        >>> self.world.stop_swarm()
+        """
+        self.loop_manager.run_coroutine(self.async_stop_swarm())
+
+    def __str__(self):  # pragma: no cover
+        return f"World(host={self.world_client.host}, port={self.world_client.port}, swarm_id={self.swarm_id}, task_id={self.task_id})"
+
+    def __repr__(self):  # pragma: no cover
+        return str(self)
+
+    # Asynchronous methods
+
+    async def async_schedule_next_iter(
+        self, node_ids: List[str], task_to_schedule_alias: str
+    ):
+        """
+        Schedule the next iteration
+
+        Parameters
+        ----------
+        node_ids : List[str]
+            The list of node ids to schedule the next iteration for
+        task_to_schedule_alias : str
+            The task to schedule for the next iteration
+
+        Examples
+        --------
+
+        Same as :meth:`schedule_next_iter <manta_light.world.World.schedule_next_iter>`
+        but asynchronous:
+
+        >>> await self.world.async_schedule_next_iter(
+        ...     node_ids=selected_nodes, task_to_schedule_alias="worker"
+        ... )
+        """
+        await self.world_client.schedule_task(
+            TaskScheduling(
+                task_id=self.task_id.oid,
+                node_ids=[ID(node_id).oid for node_id in node_ids],
+                task_to_schedule_alias=task_to_schedule_alias,
+                swarm_id=self.swarm_id.oid,
+            )
+        )
+        self.logger.info("Schedule next iteration")
+
+    async def async_stop_swarm(self):
+        """
+        Stop the swarm
+
+        Examples
+        --------
+
+        Same as :meth:`stop_swarm <manta_light.world.World.stop_swarm>`
+        but asynchronous:
+
+        >>> await self.world.async_stop_swarm()
+        """
+        await self.world_client.stop_swarm(
+            EnvId(swarm_id=self.swarm_id.oid, task_id=self.task_id.oid)
+        )
+        self.logger.info("Stop swarm response")