RSDK-9623: Add Discover Service and GetModelsFromModules to Python (#827

)

Co-authored-by: Naveed Jooma <[email protected]>

src/viam/robot/client.py

@@ -29,11 +29,14 @@
     GetCloudMetadataResponse,
     GetMachineStatusRequest,
     GetMachineStatusResponse,
+    GetModelsFromModulesRequest,
+    GetModelsFromModulesResponse,
     GetOperationsRequest,
     GetOperationsResponse,
     GetVersionRequest,
     GetVersionResponse,
     LogRequest,
+    ModuleModel,
     Operation,
     ResourceNamesRequest,
     ResourceNamesResponse,
@@ -771,6 +774,31 @@ async def discover_components(
         )
         return list(response.discovery)
 
+    #################
+    # MODULE MODELS #
+    #################
+
+    async def get_models_from_modules(
+        self,
+    ) -> List[ModuleModel]:
+        """
+
+        Get a list of module models.
+
+        ::
+            # Get module models
+            module_modles = await machine.get_models_from_modules(qs)
+
+        Args:
+
+        Returns:
+            List[ModuleModel]: A list of discovered models.
+
+        """
+        request = GetModelsFromModulesRequest()
+        response: GetModelsFromModulesResponse = await self._client.GetModelsFromModules(request)
+        return list(response.models)
+
     ############
     # STOP ALL #
     ############

src/viam/services/discovery/__init__.py

@@ -0,0 +1,12 @@
+from viam.resource.registry import Registry, ResourceRegistration
+from viam.services.discovery.service import DiscoveryRPCService
+
+from .client import DiscoveryClient
+from .discovery import Discovery
+
+__all__ = [
+    "DiscoveryClient",
+    "Discovery",
+]
+
+Registry.register_api(ResourceRegistration(Discovery, DiscoveryRPCService, lambda name, channel: DiscoveryClient(name, channel)))

src/viam/services/discovery/client.py

@@ -0,0 +1,55 @@
+from typing import Any, List, Mapping, Optional
+
+from grpclib.client import Channel
+
+from viam.proto.app.robot import ComponentConfig
+from viam.proto.common import DoCommandRequest, DoCommandResponse
+from viam.proto.service.discovery import (
+    DiscoverResourcesRequest,
+    DiscoverResourcesResponse,
+    DiscoveryServiceStub,
+)
+from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
+from viam.utils import ValueTypes, dict_to_struct, struct_to_dict
+
+from .discovery import Discovery
+
+
+class DiscoveryClient(Discovery, ReconfigurableResourceRPCClientBase):
+    """
+    Connect to the Discovery service, which allows you to discover resources on a machine.
+    """
+
+    client: DiscoveryServiceStub
+
+    def __init__(self, name: str, channel: Channel):
+        super().__init__(name)
+        self.channel = channel
+        self.client = DiscoveryServiceStub(channel)
+
+    async def discover_resources(
+        self,
+        *,
+        extra: Optional[Mapping[str, Any]] = None,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ) -> List[ComponentConfig]:
+        md = kwargs.get("metadata", self.Metadata()).proto
+        request = DiscoverResourcesRequest(
+            name=self.name,
+            extra=dict_to_struct(extra),
+        )
+        response: DiscoverResourcesResponse = await self.client.DiscoverResources(request, timeout=timeout, metadata=md)
+        return list(response.discoveries)
+
+    async def do_command(
+        self,
+        command: Mapping[str, ValueTypes],
+        *,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ) -> Mapping[str, ValueTypes]:
+        md = kwargs.get("metadata", self.Metadata()).proto
+        request = DoCommandRequest(name=self.name, command=dict_to_struct(command))
+        response: DoCommandResponse = await self.client.DoCommand(request, timeout=timeout, metadata=md)
+        return struct_to_dict(response.result)

src/viam/services/discovery/discovery.py

@@ -0,0 +1,52 @@
+import abc
+from typing import Final, List, Mapping, Optional
+
+from viam.proto.app.robot import ComponentConfig
+from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_SERVICE, API
+from viam.utils import ValueTypes
+
+from ..service_base import ServiceBase
+
+
+class Discovery(ServiceBase):
+    """
+    Discovery represents a Discovery service.
+
+    This acts as an abstract base class for any drivers representing specific
+    discovery implementations. This cannot be used on its own. If the ``__init__()`` function is
+    overridden, it must call the ``super().__init__()`` function.
+
+    """
+
+    API: Final = API(  # pyright: ignore [reportIncompatibleVariableOverride]
+        RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_SERVICE, "discovery"
+    )
+
+    @abc.abstractmethod
+    async def discover_resources(
+        self,
+        *,
+        extra: Optional[Mapping[str, ValueTypes]] = None,
+        timeout: Optional[float] = None,
+    ) -> List[ComponentConfig]:
+        """Get all component configs of discovered resources on a machine
+
+        ::
+
+            my_discovery = DiscoveryClient.from_robot(machine, "my_discovery")
+
+            # Get the discovered resources
+            result = await my_discovery.discover_resources(
+                "my_discovery",
+            )
+            discoveries = result.discoveries
+
+        Args:
+            name (str): The name of the discover service
+
+        Returns:
+            List[ComponentConfig]: A list of ComponentConfigs that describe
+            the components found by a discover service
+
+        """
+        ...

src/viam/services/discovery/service.py

@@ -0,0 +1,43 @@
+from grpclib.server import Stream
+
+from viam.proto.common import DoCommandRequest, DoCommandResponse
+from viam.proto.service.discovery import (
+    DiscoverResourcesRequest,
+    DiscoverResourcesResponse,
+    UnimplementedDiscoveryServiceBase,
+)
+from viam.resource.rpc_service_base import ResourceRPCServiceBase
+from viam.utils import dict_to_struct, struct_to_dict
+
+from .discovery import Discovery
+
+
+class DiscoveryRPCService(UnimplementedDiscoveryServiceBase, ResourceRPCServiceBase):
+    """
+    gRPC service for a Discovery service
+    """
+
+    RESOURCE_TYPE = Discovery
+
+    async def DiscoverResources(self, stream: Stream[DiscoverResourcesRequest, DiscoverResourcesResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        discovery = self.get_resource(request.name)
+        extra = struct_to_dict(request.extra)
+        timeout = stream.deadline.time_remaining() if stream.deadline else None
+        result = await discovery.discover_resources(
+            extra=extra,
+            timeout=timeout,
+        )
+        response = DiscoverResourcesResponse(
+            discoveries=result,
+        )
+        await stream.send_message(response)
+
+    async def DoCommand(self, stream: Stream[DoCommandRequest, DoCommandResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        discovery = self.get_resource(request.name)
+        timeout = stream.deadline.time_remaining() if stream.deadline else None
+        result = await discovery.do_command(struct_to_dict(request.command), timeout=timeout)
+        await stream.send_message(DoCommandResponse(result=dict_to_struct(result)))

tests/mocks/services.py

@@ -266,6 +266,7 @@
     UnimplementedMLTrainingServiceBase,
 )
 from viam.proto.app.packages import PackageType
+from viam.proto.app.robot import ComponentConfig
 from viam.proto.common import (
     DoCommandRequest,
     DoCommandResponse,
@@ -324,6 +325,7 @@
 from viam.proto.service.navigation import MapType, Mode, Path, Waypoint
 from viam.proto.service.slam import MappingMode, SensorInfo, SensorType
 from viam.proto.service.vision import Classification, Detection
+from viam.services.discovery import Discovery
 from viam.services.generic import Generic as GenericService
 from viam.services.mlmodel import File, LabelType, Metadata, MLModel, TensorInfo
 from viam.services.mlmodel.utils import flat_tensors_to_ndarrays, ndarrays_to_flat_tensors
@@ -432,6 +434,31 @@ async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Option
         return {"cmd": command}
 
 
+class MockDiscovery(Discovery):
+    def __init__(
+        self,
+        name: str,
+    ):
+        self.extra: Optional[Mapping[str, Any]] = None
+        self.timeout: Optional[float] = None
+        super().__init__(name)
+
+    async def discover_resources(
+        self,
+        *,
+        extra: Optional[Mapping[str, ValueTypes]] = None,
+        timeout: Optional[float] = None,
+    ) -> List[ComponentConfig]:
+        self.extra = extra
+        self.timeout = timeout
+        result = ComponentConfig()
+        return [result]
+
+    async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Optional[float] = None) -> Mapping[str, ValueTypes]:
+        self.timeout = timeout
+        return {"cmd": command}
+
+
 class MockMLModel(MLModel):
     INT8_NDARRAY = np.array([[0, -1], [8, 8]], dtype=np.int8)
     INT16_NDARRAY = np.array([1, 0, 0, 69, -1, 16], dtype=np.int16)

tests/test_discovery_service.py

@@ -0,0 +1,86 @@
+import pytest
+from grpclib.testing import ChannelFor
+
+from viam.proto.app.robot import ComponentConfig
+from viam.proto.common import (
+    DoCommandRequest,
+    DoCommandResponse,
+)
+from viam.proto.service.discovery import (
+    DiscoverResourcesRequest,
+    DiscoverResourcesResponse,
+    DiscoveryServiceStub,
+)
+from viam.resource.manager import ResourceManager
+from viam.services.discovery import DiscoveryClient
+from viam.services.discovery.service import DiscoveryRPCService
+from viam.utils import dict_to_struct, struct_to_dict
+
+from .mocks.services import MockDiscovery
+
+DISCOVERIES = [ComponentConfig()]
+
+
+DISCOVERY_SERVICE_NAME = "discovery1"
+
+
+@pytest.fixture(scope="function")
+def discovery() -> MockDiscovery:
+    return MockDiscovery(
+        DISCOVERY_SERVICE_NAME,
+    )
+
+
+@pytest.fixture(scope="function")
+def service(discovery: MockDiscovery) -> DiscoveryRPCService:
+    rm = ResourceManager([discovery])
+    return DiscoveryRPCService(rm)
+
+
+class TestDiscovery:
+    async def test_discover_resources(self, discovery: MockDiscovery):
+        extra = {"foo": "discovery"}
+        response = await discovery.discover_resources(extra=extra)
+        assert discovery.extra == extra
+        assert response == DISCOVERIES
+
+    async def test_do(self, discovery: MockDiscovery):
+        command = {"command": "args"}
+        response = await discovery.do_command(command)
+        assert response["cmd"] == command
+
+
+class TestService:
+    async def test_discover_resources(self, discovery: MockDiscovery, service: DiscoveryRPCService):
+        async with ChannelFor([service]) as channel:
+            client = DiscoveryServiceStub(channel)
+            extra = {"cmd": "discovery"}
+            request = DiscoverResourcesRequest(name=discovery.name, extra=dict_to_struct(extra))
+            response: DiscoverResourcesResponse = await client.DiscoverResources(request)
+            assert discovery.extra == extra
+            assert response.discoveries == DISCOVERIES
+
+    async def test_do(self, discovery: MockDiscovery, service: DiscoveryRPCService):
+        async with ChannelFor([service]) as channel:
+            client = DiscoveryServiceStub(channel)
+            command = {"command": "args"}
+            request = DoCommandRequest(name=discovery.name, command=dict_to_struct(command))
+            response: DoCommandResponse = await client.DoCommand(request)
+            assert struct_to_dict(response.result)["cmd"] == command
+
+
+class TestClient:
+    async def test_discover_resources(self, discovery: MockDiscovery, service: DiscoveryRPCService):
+        async with ChannelFor([service]) as channel:
+            client = DiscoveryClient(DISCOVERY_SERVICE_NAME, channel)
+            extra = {"foo": "discovery"}
+            response = await client.discover_resources(name=DISCOVERY_SERVICE_NAME, extra=extra)
+            assert response == DISCOVERIES
+            assert discovery.extra == extra
+
+    async def test_do(self, service: DiscoveryRPCService):
+        async with ChannelFor([service]) as channel:
+            client = DiscoveryClient(DISCOVERY_SERVICE_NAME, channel)
+            command = {"command": "args"}
+            response = await client.do_command(command)
+            assert response["cmd"] == command