RSDK-10717 - support modules on windows (#924)

Co-authored-by: Alan Davidson <[email protected]>

Author: stuqdog

src/viam/module/module.py

@@ -32,7 +32,7 @@
 from viam.resource.registry import Registry
 from viam.resource.types import API, RESOURCE_TYPE_COMPONENT, RESOURCE_TYPE_SERVICE, Model, ResourceName, resource_name_from_string
 from viam.robot.client import RobotClient
-from viam.rpc.dial import DialOptions
+from viam.rpc.dial import DialOptions, _host_port_from_url
 from viam.rpc.server import Server
 
 # These imports are required to register build-in resources with the registry
@@ -73,7 +73,8 @@ def _parse_module_args() -> argparse.Namespace:
     p = argparse.ArgumentParser(description="Start this viam python module")
     p.add_argument("socket_path", help="path where this module will serve a unix socket")
     p.add_argument("--log-level", type=lambda name: pylogging._nameToLevel[name.upper()], default=logging.INFO)
-    return p.parse_args()
+    p.add_argument("--tcp-mode", action='store_true')
+    return p.parse_known_args()[0]
 
 
 class Module:
@@ -82,6 +83,7 @@ class Module:
     _ready: bool
     _log_level: int
     _lock: Lock
+    _tcp_mode: bool
     parent: Optional[RobotClient] = None
     server: Server
     logger: pylogging.Logger
@@ -100,7 +102,7 @@ def from_args(cls) -> Self:
             Module: a new Module instance
         """
         args = _parse_module_args()
-        return cls(args.socket_path, log_level=args.log_level)
+        return cls(args.socket_path, log_level=args.log_level, tcp_mode=args.tcp_mode)
 
     @classmethod
     async def run_with_models(cls, *models: ResourceBase):
@@ -132,7 +134,7 @@ async def run_from_registry(cls):
             module.add_model_from_registry(*key.split("/"))  # pyright: ignore [reportArgumentType]
         await module.start()
 
-    def __init__(self, address: str, *, log_level: int = logging.INFO) -> None:
+    def __init__(self, address: str, *, log_level: int = logging.INFO, tcp_mode: bool = False) -> None:
         # When a module is launched by viam-server, its stdout is not connected to a tty.  In
         # response, python disables line buffering, which prevents `print` statements from being
         # immediately flushed to viam-server. This behavior can be confusing, interfere with
@@ -143,6 +145,7 @@ def __init__(self, address: str, *, log_level: int = logging.INFO) -> None:
         if isinstance(sys.stderr, io.TextIOWrapper):
             sys.stderr.reconfigure(line_buffering=True)
         self._address = address
+        self._tcp_mode = tcp_mode
         self.server = Server(resources=[], module_service=ModuleRPCService(self))
         self._log_level = log_level
 
@@ -188,7 +191,11 @@ async def _get_dependencies(self, dependencies: Sequence[str]) -> Mapping[Resour
     async def start(self):
         """Start the module service and gRPC server"""
         try:
-            await self.server.serve(log_level=self._log_level, path=self._address)
+            if self._tcp_mode:
+                host, port = _host_port_from_url(self._address)
+                await self.server.serve(log_level=self._log_level, host=host, port=port)
+            else:
+                await self.server.serve(log_level=self._log_level, path=self._address)
         finally:
             await self.stop()
 

src/viam/rpc/server.py

@@ -1,5 +1,7 @@
 from typing import TYPE_CHECKING, Callable, List, Optional
 
+import sys
+
 from grpclib import GRPCError, Status
 from grpclib._typing import IServable
 from grpclib.const import Handler
@@ -82,6 +84,25 @@ async def rcv_and_log_msg():
 
         event.method_func = log_resource_name
 
+    async def _serve(
+        self,
+        host: Optional[str] = "localhost",
+        port: Optional[int] = 9090,
+        log_level: Optional[int] = logging.INFO,
+        *,
+        path: Optional[str] = None,
+    ):
+
+        if path:
+            await self._server.start(path=path)
+            LOGGER.info(f"Serving on {path}")
+        else:
+            await self._server.start(host, port)
+            LOGGER.info(f"Serving on {host}:{port}")
+        await self._server.wait_closed()
+        await self.close()
+        LOGGER.debug("gRPC server closed")
+
     async def serve(
         self,
         host: Optional[str] = "localhost",
@@ -105,16 +126,11 @@ async def serve(
             logging.setLevel(log_level)
         listen(self._server, RecvRequest, self._grpc_recvrequest_handler)
 
-        with graceful_exit([self._server]):
-            if path:
-                await self._server.start(path=path)
-                LOGGER.info(f"Serving on {path}")
-            else:
-                await self._server.start(host, port)
-                LOGGER.info(f"Serving on {host}:{port}")
-            await self._server.wait_closed()
-            await self.close()
-            LOGGER.debug("gRPC server closed")
+        if sys.platform != 'win32':
+            with graceful_exit([self._server]):
+                await self._serve(host=host, port=port, log_level=log_level, path=path)
+        else:
+            await self._serve(host=host, port=port, log_level=log_level, path=path)
 
     @classmethod
     async def create_and_serve(