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

manta/light/path.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import io
+from pathlib import Path
+from typing import List, Optional, Union
+
+from manta_common.build.node.light_service import ProtoPath
+from manta_common.event_loop import EventLoopManager
+from .local import Local
+
+_shared_local = Local()  # shared Local instance between files
+
+
+class MantaPath:
+    def __init__(self, path: str, is_file: bool = True):
+        self._is_file = is_file
+        self._path = path
+        self._local = _shared_local
+        self.loop_manager = EventLoopManager.get_instance()
+
+    def _to_proto(self):
+        return ProtoPath(value=str(self._path), is_file=self._is_file)
+
+    @property
+    def name(self):
+        return Path(self._path).name
+
+    @property
+    def stem(self):
+        return Path(self._path).stem
+
+    @property
+    def suffix(self):
+        return Path(self._path).suffix
+
+    def __str__(self):
+        return self._path
+
+    def __truediv__(self, other: Union[str, MantaPath]) -> MantaPath:
+        if isinstance(other, str):
+            other_path = Path(other)
+        elif isinstance(other, MantaPath):
+            other_path = Path(other._path)
+        else:
+            raise TypeError(f"Unsupported type (found: {type(other)})")
+
+        return MantaPath(str(Path(self._path) / other_path))
+
+    async def async_exists(self) -> bool:
+        response = await self._local.async_exists(self._to_proto())
+        return response
+
+    def exists(self) -> bool:
+        return self._local.exists(self._to_proto())
+
+    async def async_read_bytes(self) -> io.BytesIO:
+        return await self._local.async_get_binary_data(self._to_proto())
+
+    def read_bytes(self):
+        return self._local.get_binary_data(self._to_proto())
+
+    async def async_read_text(
+        self,
+        encoding: Optional[str] = None,
+        errors: Optional[str] = None,
+        newline: Optional[str] = None,
+    ) -> io.StringIO:
+        return await self._local.async_read_file_lines(
+            self._to_proto(), encoding, errors, newline
+        )
+
+    def read_text(
+        self,
+        encoding: Optional[str] = None,
+        errors: Optional[str] = None,
+        newline: Optional[str] = None,
+    ) -> io.StringIO:
+        """
+        Read Text
+
+        Parameters
+        ----------
+        encoding : Optional[str], optional
+            Specifies the encoding to use for decoding the file contents, by default None
+        errors : Optional[str], optional
+            Specifies how encoding/decoding errors should be handled, by default None
+        newline : Optional[str], optional
+            Specifies how newlines (`\n`, `\r\n`, `\r`) should be handled, by default None
+
+        Returns
+        -------
+        io.StringIO
+            _description_
+        """
+        return self._local.read_file_lines(self._to_proto(), encoding, errors, newline)
+
+    async def async_iterdir(self) -> List[MantaPath]:
+        response = await self._local.async_list_dir(self._to_proto())
+        return [MantaPath(p.value, is_file=p.is_file) for p in response.paths]
+
+    def iterdir(self) -> List[MantaPath]:
+        return self.loop_manager.run_coroutine(self.async_iterdir())
+
+    # async def async_glob(self, pattern, *, case_sensitive=None, recurse_symlinks=False):
+    #     request = Glob(self._to_proto(), pattern, case_sensitive, recurse_symlinks)
+    #     response = await self._local.glob(request)
+    #     return [MantaPath(p.value, is_file=p.is_file) for p in response.paths]
+
+    # def glob(self, pattern, *, case_sensitive=None, recurse_symlinks=False):
+    #     return asyncio.run(
+    #         self.async_glob(
+    #             pattern,
+    #             case_sensitive=case_sensitive,
+    #             recurse_symlinks=recurse_symlinks,
+    #         )
+    #     )

manta/light/results.py

@@ -0,0 +1,236 @@
+import io
+import logging
+import os
+from typing import AsyncIterable, Dict, Optional
+
+from ..clients.world_client import WorldClient
+from manta_common.build.common.results import ResultMethod
+from manta_common.build.node.light_service import LightResult, LightResultQuery
+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__ = ["Results"]
+
+
+class Results:
+    """
+    Class for accessing results from the shared database or
+    adding results into the shared database
+
+    Parameters
+    ----------
+    host : str
+        Manager host
+    port : int
+        Manager port
+    swarm_id : ID
+        Swarm ID
+    task_id : ID
+        Task ID
+    chunk_size : int
+        Chunk size
+    """
+
+    __slots__ = [
+        "world_client",
+        "swarm_id",
+        "task_id",
+        "logger",
+        "chunk_size",
+        "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 select(self, tag: str, size: int = -1, method: ResultMethod = ResultMethod.ALL):
+        """
+        Get the results of a Task
+
+        Parameters
+        ----------
+        tag : str
+            The tag of the result to get
+        size : int, optional
+            The number of results to get
+        method : ResultMethod, optional
+            Method to use to select the results
+
+        Returns
+        -------
+        dict
+            The response from the world service
+
+        Examples
+        --------
+
+        Inside a :class:`Task <manta_light.task.Task>` class, you can
+        select results stored in the Manager database from the attribute
+        :code:`self.world` automatically created by
+        :class:`Task <manta_light.task.Task>`:
+
+        >>> params = self.world.results.select("model_params")
+        """
+        return self.loop_manager.run_coroutine(self.async_select(tag, size, method))
+
+    def add(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
+        add results to be stored in the Manager database from the
+        attribute :code:`self.world` automatically created by
+        :class:`Task <manta_light.task.Task>`:
+
+        >>> self.world.results.add("metrics", metrics)
+        """
+        self.loop_manager.run_coroutine(self.async_add(tag, result))
+
+    def __str__(self):  # pragma: no cover
+        return f"Results(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_select(self, tag: str, size: int, method: ResultMethod) -> dict:
+        """
+        Get the results of tasks
+
+        Parameters
+        ----------
+        tag : str
+            Tag of the result to get
+        size : int
+            Number of results to get
+        method : str
+            The method to use to select the results
+
+        Returns
+        -------
+        dict
+            Response from the world service
+
+        Examples
+        --------
+
+        Same as :meth:`select <manta_light.results.Results.select>`
+        but asynchronous:
+
+        >>> params = await self.world.results.async_select("model_params")
+        """
+        # Initialize a buffer to accumulate chunks
+        buffer_dict: Dict[str, io.BytesIO] = {}
+
+        # Iterate over the chunks
+        async for chunk in self.world_client.get_task_result(
+            LightResultQuery(
+                task_id=self.task_id.oid,
+                swarm_id=self.swarm_id.oid,
+                tag=tag,
+                size=size,
+                method=method,
+            )
+        ):
+            node_id = ID(chunk.node_id).xid
+            if node_id not in buffer_dict:
+                buffer_dict[node_id] = io.BytesIO()
+            buffer_dict[node_id].write(chunk.data)
+
+        self.logger.info("Results received")
+        return {
+            node_id: bytes_to_dict(buffer.getvalue())
+            for node_id, buffer in buffer_dict.items()
+        }
+
+    async def chunked_light_result(
+        self, request: LightResult
+    ) -> AsyncIterable[LightResult]:
+        """
+        This function chunks the data into smaller pieces and yields LightResult messages.
+
+        Parameters
+        ----------
+        request : LightResult
+            Request
+
+        Returns
+        -------
+        AsyncIterable[LightResult]
+            AsyncIterable of LightResult messages
+        """
+        data_stream = io.BytesIO(request.data)
+        while chunk := data_stream.read(self.chunk_size):
+            yield LightResult(
+                task_id=request.task_id,
+                swarm_id=request.swarm_id,
+                tag=request.tag,
+                data=chunk,
+            )
+
+    async def async_add(self, tag: str, result: dict):
+        """
+        Set a result of a task
+
+        Parameters
+        ----------
+        tag : str
+            Tag of the result to set
+        result : dict
+            Result to add
+
+        Examples
+        --------
+
+        Same as :meth:`select <manta_light.results.Results.add>`
+        but asynchronous:
+
+        >>> await self.world.results.async_add("metrics", metrics)
+        """
+        self.logger.info(f"Setting result for tag: {tag}")
+        await self.world_client.add_task_result(
+            self.chunked_light_result(
+                LightResult(
+                    task_id=self.task_id.oid,
+                    swarm_id=self.swarm_id.oid,
+                    tag=tag,
+                    data=dict_to_bytes(result),
+                )
+            )
+        )
+        self.logger.info("Set result response")

manta/light/task.py

@@ -0,0 +1,100 @@
+import logging
+import os
+from abc import ABC
+
+from manta_common.conversions import ID
+from manta_common.event_loop import EventLoopManager
+from .local import Local
+from .world import World
+
+__all__ = ["Task"]
+
+
+class Task(ABC):
+    """
+    Task abstract module.
+
+    - Set the name of the task
+    - Set the host and port for the RPC connection
+    - Set the Task ID
+    - Initialize the Local and World services
+    - Initialize the logger
+
+    Attributes
+    ----------
+    world: World
+        For accessing and sending global data
+    local: Local
+        For accessing local data
+    logger: logging.Logger
+        Convenient logger which are collected and stored in Manager
+    """
+
+    __slots__ = [
+        "name",
+        "host",
+        "port",
+        "task_id",
+        "swarm_id",
+        "local",
+        "world",
+        "logger",
+        "loop_manager",
+    ]
+
+    def __init__(self):
+        """Initialize the task."""
+        self.logger = logging.getLogger(__name__)
+        self.logger.debug("Initializing Task")
+
+        self.name: str = str(__class__)
+
+        self.host = os.getenv("RPC_HOST", "host.docker.internal")
+        self.port = int(os.getenv("RPC_PORT", 50051))
+        self.task_id = ID(os.getenv("TASK_ID"))
+        self.swarm_id = ID(os.getenv("SWARM_ID"))
+
+        self.local = Local(
+            host=self.host, port=self.port, swarm_id=self.swarm_id, task_id=self.task_id
+        )
+        self.world = World(
+            host=self.host, port=self.port, swarm_id=self.swarm_id, task_id=self.task_id
+        )
+
+        self.loop_manager = EventLoopManager.get_instance()
+
+    def __str__(self):  # pragma: no cover
+        return f"Task(name={self.name}, host={self.host}, port={self.port}, swarm_id={self.swarm_id}, task_id={self.task_id})"
+
+    def __repr__(self):  # pragma: no cover
+        return str(self)
+
+    def cleanup(self) -> None:
+        """Clean up any resources used by the task.
+
+        This method should be called when the task is done to ensure all
+        resources are properly released. It's especially important to
+        close any open event loops to prevent resource leaks.
+        """
+        self.logger.info("Cleaning up task resources")
+        try:
+            self.loop_manager.run_coroutine(self.local.local_client.disconnect())
+            self.loop_manager.run_coroutine(self.world.world_client.disconnect())
+            self.loop_manager.close()
+            self.logger.info("Event loop closed successfully")
+        except Exception as e:
+            self.logger.error(f"Error closing event loop: {e}")
+
+    def __del__(self) -> None:
+        """Destructor to ensure cleanup is called."""
+        try:
+            self.cleanup()
+        except Exception as e:
+            # Log the exception before suppressing to avoid silent failures
+            try:
+                self.logger.error(f"Error during task cleanup in destructor: {e}")
+            except Exception:
+                # If logging fails, at least print to stderr
+                import sys
+
+                print(f"Task cleanup error: {e}", file=sys.stderr)

manta/light/utils.py

@@ -0,0 +1,217 @@
+import io
+from typing import Union
+
+import numpy as np
+
+from manta_common.conversions import bytes_to_dict, dict_to_bytes
+
+__all__ = [
+    "dict_to_bytes",
+    "bytes_to_dict",
+    "numpy_to_bytes",
+    "bytes_to_numpy",
+    "torchmodel_to_bytes",
+    "bytes_to_torchmodel",
+]
+
+
+def numpy_to_bytes(data: Union[np.ndarray, list, dict]) -> Union[bytes, dict]:
+    """
+    Recursive function which converts a numpy array or iterator with numpy arrays to bytes
+
+    Parameters
+    ----------
+    data : np.array
+        The numpy array to convert
+
+    Returns
+    -------
+    bytes
+        The bytes representation of the numpy array
+
+    Examples
+    --------
+
+    * From :code:`np.array`
+
+    >>> import numpy as np
+    >>> from manta_light.utils import bytes_to_numpy, numpy_to_bytes
+    >>> np_array = np.array([1, 2, 3])
+    >>> numpy_to_bytes(np_array)
+
+    * From :code:`Dict[str, np.array]`
+
+    >>> import numpy as np
+    >>> from manta_light.utils import bytes_to_numpy, numpy_to_bytes
+    >>> dict_np_array = {"key1": np.array([1, 2, 3]), "key2": np.array([4, 5, 6])}
+    >>> numpy_to_bytes(dict_np_array)
+
+    * From :code:`Dict[str, list]`
+
+    >>> import numpy as np
+    >>> from manta_light.utils import bytes_to_numpy, numpy_to_bytes
+    >>> dict_list = {"key1": [1, 2, 3], "key2": [4, 5, 6]}
+    >>> numpy_to_bytes(dict_list)
+
+    * From :code:`List[np.array]`
+
+    >>> import numpy as np
+    >>> from manta_light.utils import bytes_to_numpy, numpy_to_bytes
+    >>> list_np_array = [np.array([1, 2, 3]), np.array([4, 5, 6])]
+    >>> numpy_to_bytes(list_np_array)
+    """
+    import numpy as np
+
+    if isinstance(data, np.ndarray) or isinstance(data, list):
+        buffer = io.BytesIO()
+        np.save(buffer, data)
+        return buffer.getvalue()
+    elif isinstance(data, dict):
+        return {key: numpy_to_bytes(value) for key, value in data.items()}
+    else:
+        raise ValueError(f"Unsupported type: {type(data)}")
+
+
+def bytes_to_numpy(b: Union[bytes, dict]) -> Union[np.ndarray, dict]:
+    """
+    Convert bytes to a numpy array
+
+    Parameters
+    ----------
+    b : bytes
+        The bytes to convert
+
+    Returns
+    -------
+    np.array
+        The numpy array representation of the bytes
+
+    Examples
+    --------
+
+    * From :code:`np.array`
+
+    >>> import numpy as np
+    >>> from manta_light.utils import bytes_to_numpy, numpy_to_bytes
+    >>> np_array = np.array([1, 2, 3])
+    >>> np_bytes = numpy_to_bytes(np_array)
+    >>> bytes_to_numpy(np_bytes)
+
+    * From :code:`Dict[str, np.array]`
+
+    >>> import numpy as np
+    >>> from manta_light.utils import bytes_to_numpy, numpy_to_bytes
+    >>> dict_np_array = {"key1": np.array([1, 2, 3]), "key2": np.array([4, 5, 6])}
+    >>> np_bytes = numpy_to_bytes(dict_np_array)
+    >>> bytes_to_numpy(np_bytes)
+
+    * From :code:`Dict[str, list]`
+
+    >>> import numpy as np
+    >>> from manta_light.utils import bytes_to_numpy, numpy_to_bytes
+    >>> dict_list = {"key1": [1, 2, 3], "key2": [4, 5, 6]}
+    >>> np_bytes = numpy_to_bytes(dict_list)
+    >>> bytes_to_numpy(np_bytes)
+
+    * From :code:`List[np.array]`
+
+    >>> import numpy as np
+    >>> from manta_light.utils import bytes_to_numpy, numpy_to_bytes
+    >>> list_np_array = [np.array([1, 2, 3]), np.array([4, 5, 6])]
+    >>> np_bytes = numpy_to_bytes(list_np_array)
+    >>> bytes_to_numpy(np_bytes)
+    """
+    import numpy as np
+
+    if isinstance(b, bytes):
+        buffer = io.BytesIO(b)
+        return np.load(buffer)
+    elif isinstance(b, dict):
+        return {key: bytes_to_numpy(value) for key, value in b.items()}
+    else:
+        raise ValueError(f"Unsupported type: {type(b)}")
+
+
+def torchmodel_to_bytes(model: "torch.nn.Module") -> bytes:  # type: ignore # noqa: F821
+    """
+    Transform a torch model into bytes
+
+    Parameters
+    ----------
+    model : "torch.nn.Module"
+        Torch model
+
+    Returns
+    -------
+    bytes
+        Bytes from the torch model
+
+    Examples
+    --------
+
+    >>> from torch.nn import Linear, ReLU, Sequential
+    >>> from manta_light.utils import bytes_to_torchmodel, torchmodel_to_bytes
+    >>> torch_model = Sequential(
+    ...     Linear(3, 2), ReLU(), Linear(2, 1), ReLU(), Linear(1, 1), ReLU()
+    ... )
+    >>> torchmodel_to_bytes(torch_model)
+    """
+    import torch
+
+    buffer = io.BytesIO()
+    torch.save(model, buffer)
+    return buffer.getvalue()
+
+
+def bytes_to_torchmodel(b: bytes) -> "torch.nn.Module":  # type: ignore # noqa: F821
+    """
+    Transform bytes to torch model
+
+    Parameters
+    ----------
+    b : bytes
+        Bytes from a torch model
+
+    Returns
+    -------
+    "torch.nn.Module"
+        Torch model
+
+    Examples
+    --------
+
+    >>> from torch.nn import Linear, ReLU, Sequential
+    >>> from manta_light.utils import bytes_to_torchmodel, torchmodel_to_bytes
+    >>> torch_model = Sequential(
+    ...     Linear(3, 2), ReLU(), Linear(2, 1), ReLU(), Linear(1, 1), ReLU()
+    ... )
+    >>> model_bytes = torchmodel_to_bytes(torch_model)
+    >>> bytes_to_torchmodel(model_bytes)
+
+    Security Note
+    -------------
+    This function uses torch.load with weights_only=False which can execute
+    arbitrary code. Only use with trusted model sources. In the Manta platform,
+    models should only come from authenticated users and trusted containers.
+    """
+    import pickle
+    import warnings
+
+    import torch
+
+    buffer = io.BytesIO(b)
+
+    # Try to load with weights_only=True first (safer)
+    try:
+        return torch.load(buffer, weights_only=True)
+    except (RuntimeError, TypeError, pickle.UnpicklingError):
+        # If that fails, fallback to full model loading with a warning
+        # This is necessary for complex models with custom layers
+        buffer.seek(0)  # Reset buffer position
+        warnings.warn(
+            "Loading PyTorch model with weights_only=False. "
+            "This can execute arbitrary code. Ensure the model source is trusted.",
+            UserWarning,
+            stacklevel=2,
+        )
+        return torch.load(buffer, weights_only=False)