manta-node 0.5b0__py3-none-any.whl

manta_node/infrastructure/grpc/client.py

@@ -0,0 +1,373 @@
+from __future__ import annotations
+
+from logging import getLogger
+from pathlib import Path
+from typing import (
+    AsyncIterable,
+    AsyncIterator,
+    Iterable,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
+
+from manta_common.base_client import GrpcClientBase, MetadataDict
+from manta_common.build.common.informations import (
+    DatasetOverviews,
+    MetricsSnapshot,
+    NodeOverview,
+    NodeStatus,
+    Registration,
+)
+from manta_common.build.common.results import (
+    ChunkResultValues,
+    GlobalTag,
+    GlobalUpdate,
+    GlobalValue,
+    Result,
+    ResultQuery,
+)
+from manta_common.build.common.system import Empty, Logs, Ok
+from manta_common.build.common.tasks import TaskScheduling
+from manta_common.build.node.node_service import NodeServiceStub, StopSwarmProcess
+from manta_common.errors import MantaError
+from manta_common.retry import RetryPolicy, StreamingRetryPolicy
+from manta_common.traces import Tracer
+
+__all__ = ["NodeClient"]
+
+T = TypeVar("T")
+
+
+class NodeClient(GrpcClientBase):
+    """
+    Node client facilitates management of requests to Manager with retry
+    and connection management capabilities.
+    """
+
+    #: Default retry policy for regular methods
+    DEFAULT_RETRY_POLICY = RetryPolicy(
+        max_retries=3,
+        initial_delay=0.2,
+        max_delay=10.0,
+        backoff_factor=2.0,
+        jitter_factor=0.2,
+    )
+
+    #: Default retry policy for streaming methods
+    DEFAULT_STREAMING_RETRY_POLICY = StreamingRetryPolicy(
+        max_retries=3,
+        initial_delay=0.2,
+        max_delay=10.0,
+        backoff_factor=2.0,
+        jitter_factor=0.2,
+        retry_if_no_items_processed=True,
+    )
+
+    def _get_stub_class(self) -> Type[NodeServiceStub]:
+        """
+        Get the stub class for this client.
+        """
+        return NodeServiceStub
+
+    @property
+    def metadata(self) -> MetadataDict:
+        """
+        Get metadata for gRPC requests.
+
+        Subclasses can override this property to provide custom metadata.
+        Similar to ssl_context property pattern.
+
+        Returns
+        -------
+        MetadataDict
+            Metadata dictionary for gRPC requests
+        """
+        return {"authorization": self.secured_token}
+
+    def __init__(
+        self,
+        manager_host: str,
+        manager_port: int,
+        secured_token: str,
+        ssl_cert_folder: Optional[Path] = None,
+        retry_policy: Optional[RetryPolicy] = None,
+        streaming_retry_policy: Optional[StreamingRetryPolicy] = None,
+    ) -> None:
+        """
+        Initialize Node client
+
+        Parameters
+        ----------
+
+            Node identifier
+        manager_host : str
+            Manager hostname
+        manager_port : int
+            Manager port
+        secured_token : str
+            Secured token for the Manager first connection
+        ssl_cert_folder : Optional[Path]
+            SSL certificate folder
+        retry_policy : Optional[RetryPolicy]
+            Custom retry policy for operations
+        streaming_retry_policy : Optional[StreamingRetryPolicy]
+            Custom retry policy for streaming operations
+        """
+        self.ssl_cert_folder = ssl_cert_folder
+        self.secured_token = secured_token
+        self.tracer = Tracer(getLogger(__name__))
+
+        # Log connection configuration for troubleshooting
+        security_mode = (
+            "TLS"
+            if (ssl_cert_folder is not None or manager_port == 443)
+            else "insecure"
+        )
+        self.tracer.info(
+            f"Initializing NodeClient for manager at {security_mode}://{manager_host}:{manager_port}"
+        )
+
+        # Log JWT token status without exposing the actual token
+        if secured_token:
+            token_length = len(secured_token)
+            # Basic JWT format validation (should have 3 parts separated by .)
+            token_parts = secured_token.count(".") + 1
+            if token_parts == 3:
+                self.tracer.debug(
+                    f"JWT token present: {token_length} chars, format appears valid (3 parts)"
+                )
+            else:
+                self.tracer.warning(
+                    f"JWT token present: {token_length} chars, but format may be invalid "
+                    f"(expected 3 parts, got {token_parts})"
+                )
+        else:
+            self.tracer.error("JWT token is empty or None - authentication will fail")
+
+        self.tracer.info(str(self.metadata))
+
+        super().__init__(
+            host=manager_host,
+            port=manager_port,
+            secure=ssl_cert_folder is not None or manager_port == 443,
+            tracer=self.tracer,
+            retry_policy=retry_policy,
+            streaming_retry_policy=streaming_retry_policy,
+        )
+
+    async def is_available(self, request: Empty) -> Ok:
+        """
+        Check if the Manager is available
+
+        Parameters
+        ----------
+        request : Empty
+            Empty request
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        return await self.call_service_method("is_available", request)
+
+    async def register_node(self, request: NodeOverview) -> Registration:
+        """
+        Register the node with the Manager via gRPC
+        and get back MQTT information
+
+        Parameters
+        ----------
+        request : NodeOverview
+            Node ID and its data names
+
+        Returns
+        -------
+        Registration
+            Node ID and MQTT information
+        """
+        return await self.call_service_method("register_node", request)
+
+    async def get_task_result(
+        self, request: ResultQuery
+    ) -> AsyncIterator[ChunkResultValues]:
+        """
+        Get the swarm results
+
+        Parameters
+        ----------
+        request : ResultQuery
+            Result query
+
+        Returns
+        -------
+        AsyncIterator[ChunkResultValues]
+            Result values
+        """
+        async for chunk in self.stream_service_method("get_task_result", request):
+            yield chunk
+
+    async def add_task_result(self, request: AsyncIterable[Result]) -> Ok:
+        """
+        Set a swarm result
+
+        Parameters
+        ----------
+        request : AsyncIterable[Result]
+            Task ID, global tag and value
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        return await self.call_service_method("add_task_result", request)
+
+    async def get_global(self, request: GlobalTag) -> AsyncIterator[GlobalValue]:
+        """
+        Get global values
+
+        Parameters
+        ----------
+        request : GlobalTag
+            Swarm ID and global tag
+
+        Returns
+        -------
+        AsyncIterator[GlobalValue]
+            Global values
+        """
+        async for value in self.stream_service_method("get_global", request):
+            yield value
+
+    async def set_global(
+        self,
+        request: Union[AsyncIterable[GlobalUpdate], Iterable[GlobalUpdate]],
+    ) -> Ok:
+        """
+        Set global values
+
+        Parameters
+        ----------
+        request : Union[AsyncIterable[GlobalUpdate], Iterable[GlobalUpdate]]
+            Task ID, global tag and value
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        return await self.call_service_method("set_global", request)
+
+    async def set_node_status(self, request: NodeStatus) -> Ok:
+        """
+        Set node status
+
+        Parameters
+        ----------
+        request : NodeStatus
+            Node ID, its status
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        return await self.call_service_method("set_node_status", request)
+
+    async def set_system_summary(self, request: MetricsSnapshot) -> Ok:
+        """
+        Set system summary
+
+        Parameters
+        ----------
+        request : MetricsSnapshot
+            Node ID and its system information
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        return await self.call_service_method("set_system_summary", request)
+
+    async def set_overviews(self, request: DatasetOverviews) -> Ok:
+        """
+        Set dataset overviews
+
+        Parameters
+        ----------
+        request : DatasetOverviews
+            Node ID and its dataset overviews
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        return await self.call_service_method("set_overviews", request)
+
+    async def stop_swarm(self, request: StopSwarmProcess) -> Ok:
+        """
+        Stop the swarm
+
+        Parameters
+        ----------
+        request : StopSwarmProcess
+            Swarm ID
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        return await self.call_service_method("stop_swarm", request)
+
+    async def add_task_logs(self, request: AsyncIterable[Logs]) -> Ok:
+        """
+        Add task logs to the database
+
+        Parameters
+        ----------
+        request : AsyncIterable[Logs]
+            Task logs
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        return await self.call_service_method("add_task_logs", request)
+
+    async def schedule_task(self, request: TaskScheduling) -> Ok:
+        """
+        Schedule next task
+
+        Parameters
+        ----------
+        request : TaskScheduling
+            Task scheduling containing node IDs
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        return await self.call_service_method("schedule_task", request)
+
+    async def send_error(self, manta_error: MantaError):
+        """
+        Send the error to the Manager
+
+        Parameters
+        ----------
+        error : MantaError
+            Error object
+        """
+        try:
+            self.tracer.debug("Sending MantaError to the Manager")
+            await self.call_service_method("send_error", manta_error.to_err())
+        except Exception as exc:
+            self.tracer.exception(f"Cannot send error to the Manager: {repr(exc)}")

manta_node/infrastructure/grpc/local_servicer.py

@@ -0,0 +1,151 @@
+from logging import getLogger
+from typing import AsyncIterator
+
+from manta_common.build.common.system import Binary
+from manta_common.build.node.light_service import (
+    DataRequest,
+    DirRequest,
+    ExistsResponse,
+    ListDirResponse,
+    LocalBase,
+    ProtoPath,
+    ReadText,
+)
+from manta_common.conversions import ID
+from manta_common.decorators import catch_batch_errors, catch_streaming_errors
+from manta_common.traces import Tracer
+from ..filesystem.dataset_manager import DatasetManager
+from ..grpc.client import NodeClient
+
+__all__ = ["LocalServicer"]
+
+
+class LocalServicer(LocalBase):
+    """
+    Local Servicer that provides access to the local data
+
+    Parameters
+    ----------
+
+        Node ID
+    node_client : NodeClient
+        Node client instance
+    dataset_manager : DatasetManager
+        Dataset manager for handling data operations
+    """
+
+    DTYPES = {
+        ".txt": "text",
+        ".csv": "csv",
+        ".npx": "numpy",
+        ".npy": "numpy",
+        ".npz": "numpy",
+        ".np": "numpy",
+    }
+
+    def __init__(
+        self,
+        node_client: NodeClient,
+        dataset_manager: DatasetManager,
+    ):
+        self.tracer = Tracer(getLogger(__name__))
+        self.node_client = node_client
+        self.dataset_manager = dataset_manager
+        super().__init__()
+
+    def _set_task_and_swarm_ids(self, data_request: DataRequest):
+        self.tracer.set_metadata(
+            task_id=ID(data_request.task_id),
+            swarm_id=ID(data_request.swarm_id),
+        )
+
+    @catch_streaming_errors
+    async def get_binary_data(self, data_request: DataRequest) -> AsyncIterator[Binary]:
+        """
+        Get the binary data
+
+        Parameters
+        ----------
+        data_request : DataRequest
+            Data Request containing the name of the file
+
+        Returns
+        -------
+        AsyncIterator[Binary]
+            Binary data
+
+        Raises
+        ------
+        MantaLocalError
+            If the file is not found
+        """
+        self._set_task_and_swarm_ids(data_request)
+        for chunk in self.dataset_manager.get_dataset_as_bytes(data_request.name):
+            yield Binary(content=chunk)
+
+    @catch_batch_errors
+    async def list_dir(self, request: DirRequest):
+        """
+        List the directories in the data folder
+
+        Parameters
+        ----------
+        request : DirRequest
+            Data Request
+
+        Returns
+        -------
+        ListString
+            List of directories
+        """
+        self._set_task_and_swarm_ids(request.data_request)
+        files = self.dataset_manager.list_files(request.data_request.name, request.path)
+        paths = [ProtoPath(value=str(f["name"]), is_file=f["is_file"]) for f in files]
+        self.tracer.info(f"Found {len(paths)} entries in directory")
+        return ListDirResponse(paths=paths)
+
+    @catch_streaming_errors
+    async def read_file_lines(self, request: ReadText) -> AsyncIterator[Binary]:
+        """
+        Read the lines of a file
+
+        Parameters
+        ----------
+        request : ReadText
+            Data Request containing the name of the file
+
+        Returns
+        -------
+        ListString
+            List of lines
+        """
+        self._set_task_and_swarm_ids(request.data_request)
+        for line in self.dataset_manager.read_file_lines(
+            request.data_request.name,
+            request.path.value,
+            encoding=request.encoding,
+            errors=request.errors,
+            newline=request.newline,
+        ):
+            yield Binary(content=line)
+
+    @catch_batch_errors
+    async def exists(self, request: DirRequest) -> ExistsResponse:
+        """
+        Check if a file exists
+
+        Parameters
+        ----------
+        request : DirRequest
+            Data Request with path
+
+        Returns
+        -------
+        ExistsResponse
+            Response with boolean indicating if the file exists
+        """
+        self._set_task_and_swarm_ids(request.data_request)
+        exists = self.dataset_manager.check_exists(
+            request.data_request.name, request.path
+        )
+        return ExistsResponse(value=exists)