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

manta/light/globals.py

@@ -0,0 +1,245 @@
+import io
+import logging
+import os
+from typing import AsyncIterable, Optional
+
+import grpclib.exceptions
+
+from ..clients.world_client import WorldClient
+from manta_common.build.node.light_service import GetGlobalRequest, SetGlobalRequest
+from manta_common.const import CHUNK_SIZE
+from manta_common.conversions import ID
+from manta_common.event_loop import EventLoopManager
+from .utils import bytes_to_dict, dict_to_bytes
+
+__all__ = ["Globals"]
+
+
+class Globals:
+    """
+    Class which deals with global manipulation (getter and setter)
+
+    Parameters
+    ----------
+    host : str
+        Manager host
+    port : int
+        Manager port
+    task_id : ID
+        Task ID
+    swarm_id : ID
+        Swarm ID
+    chunk_size : int
+        Chunk size
+    """
+
+    __slots__ = [
+        "world_client",
+        "swarm_id",
+        "task_id",
+        "logger",
+        "chunk_size",
+        "_loop",
+        "_loop_lock",
+        "loop_manager",
+    ]
+
+    def __init__(
+        self,
+        world_client: Optional[WorldClient] = None,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+        swarm_id: Optional[ID] = None,
+        task_id: Optional[ID] = None,
+        chunk_size: int = CHUNK_SIZE,
+    ):
+        # Retrieve env variables for RPC host and port
+        if world_client is None:
+            self.world_client = WorldClient(
+                host=host or os.getenv("RPC_HOST", "host.docker.internal"),
+                port=int(port or os.getenv("RPC_PORT", 50051)),
+            )
+        else:
+            self.world_client = world_client
+
+        self.task_id: ID = task_id or ID(os.getenv("TASK_ID"))
+        self.swarm_id: ID = swarm_id or ID(os.getenv("SWARM_ID"))
+
+        self.chunk_size = chunk_size
+
+        self.logger = logging.getLogger(__name__)
+
+        self.loop_manager = EventLoopManager.get_instance()
+
+    def __getitem__(self, tag: str):
+        """
+        Get the results of a Task
+
+        Parameters
+        ----------
+        tag : str
+            The tag of the result to get
+
+        Returns
+        -------
+        dict
+            The response from the world service
+
+        Examples
+        --------
+
+        Inside a :class:`Task <manta_light.task.Task>` class, you can
+        access globals defined by a :class:`Swarm <manta.swarm.Swarm>`
+        by using the attribute :code:`self.world` automatically defined
+        by :class:`Task <manta_light.task.Task>`:
+
+        >>> weights = self.world.globals["global_model_params"]
+        """
+        return self.loop_manager.run_coroutine(self.async_get(tag))
+
+    def __setitem__(self, tag: str, result: dict):
+        """
+        Set a result of a task
+
+        Parameters
+        ----------
+        tag : str
+            The tag of the result to set
+        result : dict
+            The result to set
+
+        Examples
+        --------
+
+        Inside a :class:`Task <manta_light.task.Task>` class, you can
+        change globals defined by a :class:`Swarm <manta.swarm.Swarm>`
+        by using the attribute :code:`self.world` automatically defined
+        by :class:`Task <manta_light.task.Task>`:
+
+        >>> self.world.globals["global_model_params"] = new_weights
+        """
+        self.loop_manager.run_coroutine(self.async_set(tag, result))
+
+    def __str__(self):  # pragma: no cover
+        return f"Globals(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)
+
+    # Asynchonous methods
+
+    async def async_get(self, tag: str):
+        """
+        Get the results of tasks
+
+        Parameters
+        ----------
+        tag : str
+            The tag of the result to get
+
+        Returns
+        -------
+        dict
+            The response from the world service
+
+        Examples
+        --------
+
+        Same as :meth:`__getitem__ <manta_light.globals.Globals.__getitem__>`
+        but asynchronous.
+
+        >>> weights = await self.world.globals.async_get(
+        ...     "global_model_params"
+        ... )
+        """
+        # Initialize a buffer to accumulate chunks
+        buffer = io.BytesIO()
+
+        self.logger.info(
+            f"Getting global '{tag}' for task {self.task_id} and swarm {self.swarm_id}"
+        )
+        try:
+            # Iterate over the chunks
+            async for chunk in self.world_client.get_global(
+                GetGlobalRequest(
+                    task_id=self.task_id.oid, swarm_id=self.swarm_id.oid, tag=tag
+                )
+            ):
+                self.logger.debug(
+                    f"Received chunk for global '{tag}' of size {len(chunk.data)}"
+                )
+                # Append the chunk to the buffer
+                buffer.write(chunk.data)
+
+            self.logger.info(f"Successfully received all chunks for global '{tag}'")
+            return bytes_to_dict(buffer.getvalue())
+        except grpclib.exceptions.GRPCError as e:
+            self.logger.error(
+                f"GRPC error while getting global '{tag}': {e!r}", exc_info=True
+            )
+            raise  # Re-raise the original gRPC error
+        except Exception as e:
+            # This catches other potential errors, e.g., from bytes_to_dict if the stream was successful
+            # but data was malformed, or other unexpected errors within this block.
+            self.logger.error(
+                f"Failed to process global '{tag}' after fetching: {e!r}", exc_info=True
+            )
+            raise
+
+    async def chunked_global_update(
+        self, request: SetGlobalRequest
+    ) -> AsyncIterable[SetGlobalRequest]:
+        """
+        This function chunks the data into smaller pieces and yields GlobalUpdate messages.
+
+        Parameters
+        ----------
+        request : SetGlobalRequest
+            Set global request
+
+        Returns
+        -------
+        AsyncIterable[SetGlobalRequest]
+            Async iterable of SetGlobalRequest messages
+        """
+        data_stream = io.BytesIO(request.data)
+        while chunk := data_stream.read(self.chunk_size):
+            yield SetGlobalRequest(
+                task_id=request.task_id,
+                swarm_id=request.swarm_id,
+                tag=request.tag,
+                data=chunk,
+            )
+
+    async def async_set(self, tag: str, result: dict):
+        """
+        Set a result of a task
+
+        Parameters
+        ----------
+        tag : str
+            The tag of the result to set
+        result : dict
+            The result to set
+
+        Examples
+        --------
+
+        Same as :meth:`__setitem__ <manta_light.globals.Globals.__setitem__>`
+        but asynchronous.
+
+        >>> await self.world.globals.async_set(
+        ...     "global_model_params",
+        ...     new_weights
+        ... )
+        """
+        await self.world_client.set_global(
+            self.chunked_global_update(
+                SetGlobalRequest(
+                    task_id=self.task_id.oid,
+                    swarm_id=self.swarm_id.oid,
+                    tag=tag,
+                    data=dict_to_bytes(result),
+                )
+            )
+        )

manta/light/local.py

@@ -0,0 +1,407 @@
+import io
+import logging
+import os
+from typing import TYPE_CHECKING, Optional, Union
+
+if TYPE_CHECKING:
+    import numpy
+
+from ..clients.local_client import LocalClient
+from manta_common.build.node.light_service import (
+    DataRequest,
+    DirRequest,
+    ListDirResponse,
+    ProtoPath,
+    ReadText,
+)
+from manta_common.conversions import ID
+from manta_common.event_loop import EventLoopManager
+from .logging_config import configure_logging
+
+__all__ = ["Local"]
+
+configure_logging()
+
+
+class Local:
+    """
+    Local service for accessing local data securely
+
+    The :code:`Local` service allows the user to get local data.
+
+    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", "local_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 env variables for RPC host and port
+        self.task_id: ID = task_id or ID(os.getenv("TASK_ID"))
+        self.swarm_id: ID = swarm_id or ID(os.getenv("SWARM_ID"))
+
+        self.loop_manager = EventLoopManager.get_instance()
+        self.local_client = LocalClient(
+            host=host or os.getenv("RPC_HOST", "host.docker.internal"),
+            port=int(port or os.getenv("RPC_PORT", 50051)),
+        )
+
+        self.logger = logging.getLogger(__name__)
+
+        self.logger.info(
+            f"Local configured with Host: {self.local_client.host}, Port: {self.local_client.port},"
+            f" Task ID: {self.task_id.xid},"
+            f" Swarm ID: {self.swarm_id.xid}"
+        )
+
+    def get_numpy_data(self, dataset_name: Union[str, ProtoPath]) -> "numpy.ndarray":  # type: ignore
+        """
+        Get the numpy data
+
+        Parameters
+        ----------
+        dataset_name : Union[str, ProtoPath]
+            The name of the dataset to get
+
+        Returns
+        -------
+        numpy.ndarray
+            The numpy data
+
+        Examples
+        --------
+
+        Inside a :class:`Task <manta_light.task.Task>` class, you can
+        access to local data by using the attribute :code:`self.local`
+        automatically created by :class:`Task <manta_light.task.Task>`:
+
+        >>> mnist_data = self.local.get_numpy_data("mnist.npz")
+        """
+        return self.loop_manager.run_coroutine(self.async_get_numpy_data(dataset_name))
+
+    def get_binary_data(self, dataset_name: Union[str, ProtoPath]) -> io.BytesIO:
+        """
+        Get the binary data
+
+        Parameters
+        ----------
+        dataset_name : Union[str, ProtoPath]
+            The name of the dataset to get
+
+        Returns
+        -------
+        io.BytesIO
+            The binary data
+
+        Examples
+        --------
+
+        Inside a :class:`Task <manta_light.task.Task>` class, you can
+        access to local data by using the attribute :code:`self.local`
+        automatically created by :class:`Task <manta_light.task.Task>`:
+
+        >>> mnist_data = self.local.get_binary_data("mnist.npz")
+        >>> data = np.load(mnist_data)
+        """
+        return self.loop_manager.run_coroutine(self.async_get_binary_data(dataset_name))
+
+    def list_dir(self, directory: Union[str, ProtoPath]) -> ListDirResponse:
+        """
+        List directories
+
+        Parameters
+        ----------
+        directory : Union[str, ProtoPath]
+            Directory name
+
+        Returns
+        -------
+        ListDirResponse
+            List of directories
+
+        Examples
+        --------
+
+        Inside a :class:`Task <manta_light.task.Task>` class, you can
+        access to local data by using the attribute :code:`self.local`
+        automatically created by :class:`Task <manta_light.task.Task>`:
+
+        >>> images = self.local.list_dir("images")
+        >>> print(images)
+        >>> ListDirResponse(
+        >>>     paths=[ProtoPath("image1.jpg", is_file=True), ProtoPath("image2.jpg", is_file=True)]
+        >>> )
+        """
+        if isinstance(directory, str):
+            proto_path = ProtoPath(name=directory, value="")
+        else:
+            proto_path = directory
+        return self.loop_manager.run_coroutine(self.async_list_dir(proto_path))
+
+    def read_file_lines(
+        self,
+        file_name: Union[str, ProtoPath],
+        encoding: Optional[str] = None,
+        errors: Optional[str] = None,
+        newline: Optional[str] = None,
+    ) -> io.StringIO:
+        """
+        Read file lines
+
+        Parameters
+        ----------
+        file_name : Union[str, ProtoPath]
+            File name
+        encoding : str, optional
+            Encoding, by default None
+        errors : str, optional
+            Errors, by default None
+        newline : str, optional
+            Newline, by default None
+
+        Returns
+        -------
+        io.StringIO
+            String data
+
+        Examples
+        --------
+
+        Inside a :class:`Task <manta_light.task.Task>` class, you can
+        access to local data by using the attribute :code:`self.local`
+        automatically created by :class:`Task <manta_light.task.Task>`:
+
+        >>> text = self.local.read_file_lines("text.txt")
+        >>> print(text)
+        >>> io.StringIO("Hello world")
+        """
+        if isinstance(file_name, str):
+            proto_path = ProtoPath(name=file_name, value="")
+        else:
+            proto_path = file_name
+        return self.loop_manager.run_coroutine(
+            self.async_read_file_lines(proto_path, encoding, errors, newline)
+        )
+
+    def exists(self, file_name: Union[str, ProtoPath]) -> bool:
+        """
+        Check if a file exists
+
+        Parameters
+        ----------
+        file_name : Union[str, ProtoPath]
+            File name
+
+        Returns
+        -------
+        bool
+            True if the file exists, False otherwise
+
+        Examples
+        --------
+
+        Inside a :class:`Task <manta_light.task.Task>` class, you can
+        access to local data by using the attribute :code:`self.local`
+        automatically created by :class:`Task <manta_light.task.Task>`:
+
+        >>> if self.local.exists("text.txt"):
+        >>>     print("File exists")
+        >>> else:
+        >>>     print("File does not exist")
+        """
+        if isinstance(file_name, str):
+            proto_path = ProtoPath(name=file_name, value="")
+        else:
+            proto_path = file_name
+        return self.loop_manager.run_coroutine(self.async_exists(proto_path))
+
+    def __str__(self):  # pragma: no cover
+        return f"Local(host={self.local_client.host}, port={self.local_client.port}, task_id={self.task_id}, swarm_id={self.swarm_id})"
+
+    def __repr__(self):  # pragma: no cover
+        return str(self)
+
+    # Asynchronous methods
+
+    async def async_get_numpy_data(
+        self, dataset_name: Union[str, ProtoPath]
+    ) -> "numpy.ndarray":  # type: ignore
+        """
+        Get Numpy Data
+
+        Parameters
+        ----------
+        dataset_name : Union[str, ProtoPath]
+            Dataset name
+
+        Returns
+        -------
+        numpy.ndarray
+            Numpy data
+
+        Examples
+        --------
+
+        Same as :meth:`get_numpy_data <manta_light.local.Local.get_numpy_data>`
+        but asynchronous:
+
+        >>> mnist_data = await self.local.async_get_numpy_data("mnist.npz")
+        """
+        try:
+            import numpy as np
+        except ImportError:
+            raise ImportError("numpy is not installed")
+
+        buffer = await self.async_get_binary_data(dataset_name)
+        return np.load(buffer)
+
+    async def async_get_binary_data(
+        self, dataset_name: Union[str, ProtoPath]
+    ) -> io.BytesIO:
+        """
+        Asynchronously get binary data
+
+        Parameters
+        ----------
+        dataset_name : Union[str, ProtoPath]
+            Dataset name
+
+        Returns
+        -------
+        io.BytesIO
+            Binary data
+
+        Examples
+        --------
+
+        Same as :meth:`get_binary_data <manta_light.local.Local.get_binary_data>`
+        but asynchronous:
+
+        >>> mnist_data = await self.local.async_get_binary_data("mnist.npz")
+        >>> data = np.load(mnist_data)
+        """
+        if isinstance(dataset_name, str):
+            name = dataset_name
+        else:
+            name = dataset_name.name
+
+        # Get binary data and convert to BytesIO for numpy.load
+        buffer = io.BytesIO()
+        async for chunk in self.local_client.get_binary_data(
+            DataRequest(
+                task_id=self.task_id.oid, swarm_id=self.swarm_id.oid, name=name
+            ),
+        ):
+            buffer.write(chunk.content)
+        buffer.seek(0)
+        return buffer
+
+    async def async_list_dir(self, directory: ProtoPath) -> ListDirResponse:
+        """
+        List directories
+
+        Parameters
+        ----------
+        directory : ProtoPath
+            Directory name
+
+        Returns
+        -------
+        ListDirResponse
+            List of directories
+        """
+        return await self.local_client.list_dir(
+            DirRequest(
+                data_request=DataRequest(
+                    task_id=self.task_id.oid,
+                    swarm_id=self.swarm_id.oid,
+                    name=directory.name,
+                ),
+                path=directory.value,
+            ),
+        )
+
+    async def async_read_file_lines(
+        self,
+        file_name: ProtoPath,
+        encoding: Optional[str] = None,
+        errors: Optional[str] = None,
+        newline: Optional[str] = None,
+    ) -> io.StringIO:
+        """
+        Read file lines
+
+        Parameters
+        ----------
+        file_name : ProtoPath
+            File name
+        encoding : Optional[str], optional
+            Encoding, by default None
+        errors : Optional[str], optional
+            Errors, by default None
+        newline : Optional[str], optional
+            Newline, by default None
+
+        Returns
+        -------
+        io.StringIO
+            String data
+        """
+        buffer = io.StringIO()
+        async for chunk in self.local_client.read_file_lines(
+            ReadText(
+                data_request=DataRequest(
+                    task_id=self.task_id.oid,
+                    swarm_id=self.swarm_id.oid,
+                    name=file_name.name,
+                ),
+                path=file_name,
+                encoding="" if encoding is None else encoding,
+                errors="" if errors is None else errors,
+                newline="" if newline is None else newline,
+            ),
+        ):
+            buffer.write(chunk.content.decode())
+        buffer.seek(0)
+        return buffer
+
+    async def async_exists(self, path: ProtoPath) -> bool:
+        """
+        Check if a file exists
+
+        Parameters
+        ----------
+        path : ProtoPath
+            Path
+
+        Returns
+        -------
+        bool
+            True if exists
+        """
+        response = await self.local_client.exists(
+            DirRequest(
+                data_request=DataRequest(
+                    task_id=self.task_id.oid, swarm_id=self.swarm_id.oid, name=path.name
+                ),
+                path=path.value,
+            ),
+        )
+        return response.value

manta/light/logging_config.py

@@ -0,0 +1,39 @@
+import logging
+import logging.config
+
+FORMAT = (
+    "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
+)
+
+MANTA_LOGGING_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "formatters": {
+        "standard": {"format": FORMAT},
+    },
+    "filters": {
+        "warnings_and_below": {
+            "()": "manta_common.logging_config.filter_maker",
+            "level": "WARNING",
+        },
+        "manta_filter": {
+            "()": "manta_common.logging_config.filter_manta",
+            "level": "INFO",
+        },
+    },
+    "handlers": {
+        "default": {
+            "class": "logging.StreamHandler",
+            "level": "INFO",
+            "formatter": "standard",
+        },
+    },
+    "root": {
+        "level": "DEBUG",
+        "handlers": ["default"],
+    },
+}
+
+
+def configure_logging():
+    logging.config.dictConfig(MANTA_LOGGING_CONFIG)