diff --git a/src/viam/robot/client.py b/src/viam/robot/client.py index ef1130439..9fce18502 100644 --- a/src/viam/robot/client.py +++ b/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 # ############ diff --git a/src/viam/services/discovery/__init__.py b/src/viam/services/discovery/__init__.py new file mode 100644 index 000000000..ed1d49e2e --- /dev/null +++ b/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))) diff --git a/src/viam/services/discovery/client.py b/src/viam/services/discovery/client.py new file mode 100644 index 000000000..486b62add --- /dev/null +++ b/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) diff --git a/src/viam/services/discovery/discovery.py b/src/viam/services/discovery/discovery.py new file mode 100644 index 000000000..d87efe2b8 --- /dev/null +++ b/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 + + """ + ... diff --git a/src/viam/services/discovery/service.py b/src/viam/services/discovery/service.py new file mode 100644 index 000000000..35eed96c1 --- /dev/null +++ b/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))) diff --git a/tests/mocks/services.py b/tests/mocks/services.py index cb1b42a42..a29bbb90a 100644 --- a/tests/mocks/services.py +++ b/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) diff --git a/tests/test_discovery_service.py b/tests/test_discovery_service.py new file mode 100644 index 000000000..1622d2727 --- /dev/null +++ b/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