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

manta/clients/swarm_management_client.py

@@ -0,0 +1,562 @@
+"""
+Swarm Management Client implementation.
+
+This module provides the SwarmManagementClient class for interacting with the SwarmManagement service.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from logging import Logger
+from pathlib import Path
+from typing import Any, AsyncIterator, Dict, List, Optional, Type, TypeVar
+
+from manta_common.base_client import GrpcClientBase, MetadataDict
+from manta_common.build.common.results import Global
+from manta_common.build.common.swarms import Swarm
+from manta_common.build.common.system import (
+    CollectionIds,
+    Empty,
+    Networks,
+    Response,
+    StopSwarmRequest,
+    StrId,
+)
+from manta_common.build.core.user_services import (  # Request/Response messages; Service stub
+    DeleteResultsRequest,
+    GlobalRequest,
+    LogQuery,
+    LogResponse,
+    ResultsRequest,
+    ResultsResponse,
+    SelectGlobalResponse,
+    SendSwarmRequest,
+    SwarmManagementStub,
+    Tags,
+    TaskResponse,
+)
+from manta_common.retry import RetryPolicy
+
+T = TypeVar("T")
+
+
+class SwarmManagementClient(GrpcClientBase):
+    """
+    Client for interacting with the SwarmManagement Service.
+
+    This client provides a high-level interface for making gRPC calls to the SwarmManagement service,
+    handling swarm operations, data management, and log collection.
+    """
+
+    __slots__ = ("_secure", "cafile", "certfile", "keyfile")
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        jwt_token: str,
+        secure: bool = False,
+        cafile: Optional[str] = None,
+        certfile: Optional[str] = None,
+        keyfile: Optional[str] = None,
+        channel_options: Optional[Dict[str, Any]] = None,
+        logger: Optional[Logger] = None,
+        retry_policy: Optional[RetryPolicy] = None,
+        streaming_retry_policy: Optional[RetryPolicy] = None,
+    ):
+        """
+        Initialize the SwarmManagementClient.
+
+        Parameters
+        ----------
+        host : str
+            The host of the SwarmManagement service
+        port : int
+            The port of the SwarmManagement service
+        jwt_token : str
+            JWT token for authentication
+        secure : bool
+            Whether to use a secure connection
+        cafile : Optional[str]
+            Path to the CA certificate file for verifying the server
+        certfile : Optional[str]
+            Path to the client certificate file
+        keyfile : Optional[str]
+            Path to the client private key file
+        channel_options : Optional[Dict[str, Any]]
+            The channel options
+        logger : Optional[Logger]
+            The logger
+        retry_policy : Optional[RetryPolicy]
+            The retry policy
+        streaming_retry_policy : Optional[RetryPolicy]
+            The streaming retry policy
+        """
+        super().__init__(
+            host=host,
+            port=port,
+            secure=secure,
+            channel_options=channel_options,
+            tracer=logger,
+            retry_policy=retry_policy,
+            streaming_retry_policy=streaming_retry_policy,
+        )
+        self.jwt_token = jwt_token
+        self._secure = secure
+        self.cafile = Path(cafile) if cafile is not None else None
+        self.certfile = Path(certfile) if certfile is not None else None
+        self.keyfile = Path(keyfile) if keyfile is not None else None
+
+    @property
+    def metadata(self) -> MetadataDict:
+        return {"authorization": self.jwt_token}
+
+    def _get_stub_class(self) -> Type[SwarmManagementStub]:
+        """Get the stub class for the SwarmManagement service."""
+        return SwarmManagementStub
+
+    async def is_available(self) -> Response:
+        """
+        Check if the SwarmManagement service is available.
+
+        Returns
+        -------
+        Response
+            Service availability status
+        """
+        return await self.call_service_method(
+            "is_available",
+            Empty(),
+        )
+
+    async def send_swarm(self, request: SendSwarmRequest) -> Swarm:
+        """
+        Send a swarm graph to the service.
+
+        Parameters
+        ----------
+        request : SendSwarmRequest
+            The swarm request containing graph, name, and cluster_id
+
+        Returns
+        -------
+        Swarm
+            Overview of the registered swarm
+        """
+        return await self.call_service_method(
+            "send_swarm",
+            request,
+        )
+
+    async def list_swarm_ids(self) -> CollectionIds:
+        """
+        List all swarm IDs for the authenticated user.
+
+        Returns
+        -------
+        CollectionIds
+            Collection of swarm IDs
+        """
+        return await self.call_service_method(
+            "list_swarm_ids",
+            Empty(),
+        )
+
+    async def stream_swarms(self) -> AsyncIterator[Swarm]:
+        """
+        Stream all swarms for the authenticated user.
+
+        Yields
+        ------
+        Swarm
+            Swarm details
+        """
+        async for swarm in self.stream_service_method(
+            "stream_swarms",
+            Empty(),
+        ):
+            yield swarm
+
+    async def get_swarm(self, swarm_id: str) -> Swarm:
+        """
+        Get a swarm by its ID.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm
+
+        Returns
+        -------
+        Swarm
+            Swarm details
+        """
+        return await self.call_service_method(
+            "get_swarm",
+            StrId(id=swarm_id),
+        )
+
+    async def stream_swarm_tasks(self, swarm_id: str) -> AsyncIterator[TaskResponse]:
+        """
+        Stream tasks for a specific swarm.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm
+
+        Yields
+        ------
+        TaskResponse
+            Stream of task information
+        """
+        async for task in self.stream_service_method(
+            "stream_swarm_tasks",
+            StrId(id=swarm_id),
+        ):
+            yield task
+
+    async def start_swarm(self, swarm_id: str) -> Response:
+        """
+        Start a swarm.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm to start
+
+        Returns
+        -------
+        Response
+            Start operation response
+        """
+        return await self.call_service_method(
+            "start_swarm",
+            StrId(id=swarm_id),
+        )
+
+    async def stop_swarm(
+        self, swarm_id: str, force: bool = False, remove: bool = False
+    ) -> Response:
+        """
+        Stop a running swarm.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm to stop
+        force : bool
+            Whether to force stop the swarm
+        remove : bool
+            Whether to remove the swarm after stopping
+
+        Returns
+        -------
+        Response
+            Stop operation response
+        """
+        return await self.call_service_method(
+            "stop_swarm",
+            StopSwarmRequest(
+                swarm_id=swarm_id,
+                force=force,
+                remove=remove,
+            ),
+        )
+
+    async def remove_swarm(self, swarm_id: str) -> Response:
+        """
+        Remove a swarm from the system.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm to remove
+
+        Returns
+        -------
+        Response
+            Removal confirmation response
+        """
+        return await self.call_service_method(
+            "remove_swarm",
+            StrId(id=swarm_id),
+        )
+
+    async def select_results(
+        self, swarm_id: str, tag: str
+    ) -> AsyncIterator[ResultsResponse]:
+        """
+        Select results based on query criteria.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm
+        tag : str
+            The result tag
+
+        Yields
+        ------
+        ResultsResponse
+            Stream of results
+        """
+        async for result in self.stream_service_method(
+            "select_results",
+            ResultsRequest(swarm_id=swarm_id, tag=tag),
+        ):
+            yield result
+
+    async def stream_results(
+        self, swarm_id: str, tag: str
+    ) -> AsyncIterator[ResultsResponse]:
+        """
+        Stream results for a specific swarm and tag.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm
+        tag : str
+            The result tag
+
+        Yields
+        ------
+        ResultsResponse
+            Stream of results
+        """
+        async for result in self.stream_service_method(
+            "stream_results",
+            ResultsRequest(swarm_id=swarm_id, tag=tag),
+        ):
+            yield result
+
+    async def delete_results(
+        self,
+        swarm_id: str,
+        tag: str,
+        node_id: Optional[str] = None,
+        task_id: Optional[str] = None,
+        iteration: Optional[int] = None,
+    ) -> Response:
+        """
+        Delete results for a specific swarm and tag.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm
+        tag : str
+            The tag of the results to delete
+        node_id : Optional[str]
+            Filter by node ID
+        task_id : Optional[str]
+            Filter by task ID
+        iteration : Optional[int]
+            Filter by iteration number
+
+        Returns
+        -------
+        Response
+            Deletion confirmation response
+        """
+        request = DeleteResultsRequest(swarm_id=swarm_id, tag=tag)
+        if node_id:
+            request.node_id = node_id
+        if task_id:
+            request.task_id = task_id
+        if iteration is not None:
+            request.iteration = iteration
+        return await self.call_service_method(
+            "delete_results",
+            request,
+        )
+
+    async def list_result_tags(self, swarm_id: str) -> Tags:
+        """
+        List result tags for a specific swarm.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm
+
+        Returns
+        -------
+        Tags
+            Tags with usage information
+        """
+        return await self.call_service_method(
+            "list_result_tags",
+            StrId(id=swarm_id),
+        )
+
+    async def collect_logs(
+        self,
+        swarm_id: str,
+        node_ids: Optional[List[str]] = None,
+        task_ids: Optional[List[str]] = None,
+        severity: Optional[List[str]] = None,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+        limit: Optional[int] = None,
+        sort_order: Optional[str] = None,
+        iteration: Optional[int] = None,
+        circular: Optional[int] = None,
+    ) -> AsyncIterator[LogResponse]:
+        """
+        Collect logs based on a query.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm
+        node_ids : Optional[List[str]]
+            Filter by node IDs
+        task_ids : Optional[List[str]]
+            Filter by task IDs
+        severity : Optional[List[str]]
+            Filter by severity levels
+        start_time : Optional[datetime]
+            Start time for filtering logs
+        end_time : Optional[datetime]
+            End time for filtering logs
+        limit : Optional[int]
+            Maximum number of logs to return
+        sort_order : Optional[str]
+            Sort order ("asc" or "desc")
+        iteration : Optional[int]
+            Filter by specific iteration
+        circular : Optional[int]
+            Filter by specific circular
+
+        Yields
+        ------
+        LogResponse
+            Stream of log entries
+        """
+        request = LogQuery(swarm_id=swarm_id)
+        if node_ids:
+            request.node_ids.extend(node_ids)
+        if task_ids:
+            request.task_ids.extend(task_ids)
+        if severity:
+            request.severity.extend(severity)
+        if start_time:
+            request.start_time = start_time
+        if end_time:
+            request.end_time = end_time
+        if limit:
+            request.limit = limit
+        if sort_order:
+            request.sort_order = sort_order
+        if iteration is not None:
+            request.iteration = iteration
+        if circular is not None:
+            request.circular = circular
+        async for log in self.stream_service_method(
+            "collect_logs",
+            request,
+        ):
+            yield log
+
+    async def stream_logs(self, swarm_id: str) -> AsyncIterator[LogResponse]:
+        """
+        Stream logs for a specific swarm.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm
+
+        Yields
+        ------
+        LogResponse
+            Stream of log entries
+        """
+        async for log in self.stream_service_method(
+            "stream_logs",
+            StrId(id=swarm_id),
+        ):
+            yield log
+
+    async def initialize_global(self, globals: AsyncIterator[Global]) -> Response:
+        """
+        Initialize global values for a swarm.
+
+        Parameters
+        ----------
+        globals : AsyncIterator[Global]
+            Stream of global values
+
+        Returns
+        -------
+        Response
+            Initialization response
+        """
+        return await self.call_service_method(
+            "initialize_global",
+            globals,
+        )
+
+    async def select_global(
+        self, swarm_id: str, tag: str
+    ) -> AsyncIterator[SelectGlobalResponse]:
+        """
+        Select global data for a specific swarm and tag.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm
+        tag : str
+            The tag
+
+        Yields
+        ------
+        SelectGlobalResponse
+            Stream of global data
+        """
+        async for result in self.stream_service_method(
+            "select_global",
+            GlobalRequest(swarm_id=swarm_id, tag=tag),
+        ):
+            yield result
+
+    async def list_global_tags(self, swarm_id: str) -> Tags:
+        """
+        List global tags for a specific swarm.
+
+        Parameters
+        ----------
+        swarm_id : str
+            ID of the swarm
+
+        Returns
+        -------
+        Tags
+            Tags with usage information
+        """
+        return await self.call_service_method(
+            "list_global_tags",
+            StrId(id=swarm_id),
+        )
+
+    async def initialize_networks(self, networks: Networks) -> Response:
+        """
+        Initialize networks for a swarm.
+
+        Parameters
+        ----------
+        networks : Networks
+            Network configuration
+
+        Returns
+        -------
+        Response
+            Initialization response
+        """
+        return await self.call_service_method(
+            "initialize_networks",
+            networks,
+        )