RSDK-6184: better module logs over gRPC (#557)

Author: purplenicole730

src/viam/logging.py

@@ -1,14 +1,66 @@
+import asyncio
 import logging
 import sys
 from copy import copy
+from datetime import datetime
 from logging import DEBUG, ERROR, FATAL, INFO, WARN, WARNING  # noqa: F401
-from typing import Dict
+from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Union
 
-LOG_LEVEL = INFO
-LOGGERS: Dict[str, logging.Logger] = {}
+from grpclib.exceptions import StreamTerminatedError
+
+import viam
 
+if TYPE_CHECKING:
+    from .robot.client import RobotClient
 
-class ColorFormatter(logging.Formatter):
+
+LOG_LEVEL = INFO
+LOGGERS: Dict[str, logging.Logger] = {}
+_MODULE_PARENT: Optional["RobotClient"] = None
+
+
+class _ModuleHandler(logging.Handler):
+    _parent: "RobotClient"
+    _logger: logging.Logger
+
+    def __init__(self, parent: "RobotClient"):
+        self._parent = parent
+        self._logger = logging.getLogger("ModuleLogger")
+        addHandlers(self._logger, True)
+        super().__init__()
+        self._logger.setLevel(self.level)
+
+    def setLevel(self, level: Union[int, str]) -> None:
+        self._logger.setLevel(level)
+        return super().setLevel(level)
+
+    def handle_task_result(self, task: asyncio.Task):
+        try:
+            _ = task.result()
+        except (asyncio.CancelledError, asyncio.InvalidStateError, StreamTerminatedError):
+            pass
+        except Exception:
+            self._logger.exception("Exception raised by task = %r", task)
+
+    def emit(self, record: logging.LogRecord):
+        assert isinstance(record, logging.LogRecord)
+        name = record.name.split(".")[-1]
+        message = f"{record.filename}:{record.lineno}\t{record.msg}"
+        stack = f"exc_info: {record.exc_info}, exc_text: {record.exc_text}, stack_info: {record.stack_info}"
+        time = datetime.fromtimestamp(record.created)
+
+        try:
+            assert self._parent is not None
+            asyncio.create_task(
+                self._parent.log(name, record.levelname, time, message, stack), name=f"{viam._TASK_PREFIX}-LOG-{record.created}"
+            ).add_done_callback(self.handle_task_result)
+        except Exception as err:
+            # If the module log fails, log using stdout/stderr handlers
+            self._logger.error(f"ModuleLogger failed for {record.name} - {err}")
+            self._logger.log(record.levelno, message)
+
+
+class _ColorFormatter(logging.Formatter):
     MAPPING = {
         "DEBUG": 37,  # white
         "INFO": 36,  # cyan
@@ -43,31 +95,56 @@ def getLogger(name: str) -> logging.Logger:
     return logger
 
 
-def addHandlers(logger: logging.Logger):
-    logger.handlers.clear()
+def addHandlers(logger: logging.Logger, use_default_handlers=False):
+    _addHandlers([logger], use_default_handlers)
 
-    format = ColorFormatter("%(asctime)s\t\t" + "%(levelname)s\t" + "%(name)s (%(filename)s:%(lineno)d)\t" + "%(message)s\t")
 
-    handler = logging.StreamHandler(stream=sys.stdout)
-    handler.setFormatter(format)
+def _addHandlers(loggers: Iterable[logging.Logger], use_default_handlers=False):
+    format = _ColorFormatter("%(asctime)s\t\t" + "%(levelname)s\t" + "%(name)s (%(filename)s:%(lineno)d)\t" + "%(message)s\t")
+
+    handlers: List[logging.Handler] = []
+
+    std_handler = logging.StreamHandler(stream=sys.stdout)
+    std_handler.setFormatter(format)
     # filter out logs at error level or above
-    handler.setLevel(LOG_LEVEL)
-    handler.addFilter(filter=lambda record: (record.levelno < ERROR))
-    logger.addHandler(handler)
+    std_handler.setLevel(LOG_LEVEL)
+    std_handler.addFilter(filter=lambda record: (record.levelno < ERROR))
 
     err_handler = logging.StreamHandler(stream=sys.stderr)
     err_handler.setFormatter(format)
     # filter out logs below error level
     err_handler.setLevel(max(ERROR, LOG_LEVEL))
-    logger.addHandler(err_handler)
+
+    if _MODULE_PARENT is not None and not use_default_handlers:
+        mod_handler = _ModuleHandler(_MODULE_PARENT)
+        mod_handler.setFormatter(format)
+        mod_handler.setLevel(LOG_LEVEL)
+        handlers = [mod_handler]
+    else:
+        handlers = [std_handler, err_handler]
+
+    for logger in loggers:
+        logger.handlers.clear()
+        if "viam.sessions_client" in LOGGERS and LOGGERS["viam.sessions_client"] == logger:
+            logger.addHandler(std_handler)
+            logger.addHandler(err_handler)
+        else:
+            for h in handlers:
+                logger.addHandler(h)
+
+
+def setParent(parent: "RobotClient"):
+    global _MODULE_PARENT
+    _MODULE_PARENT = parent
+    _addHandlers(LOGGERS.values())
 
 
 def setLevel(level: int):
     global LOG_LEVEL
     LOG_LEVEL = level
     for logger in LOGGERS.values():
         logger.setLevel(LOG_LEVEL)
-        addHandlers(logger)
+    _addHandlers(LOGGERS.values())
 
 
 def silence():

src/viam/module/module.py

@@ -81,6 +81,8 @@ async def _connect_to_parent(self):
                     log_level=self._log_level,
                 ),
             )
+            LOGGER.debug("Starting module logging")
+            logging.setParent(self.parent)
 
     async def _get_resource(self, name: ResourceName) -> ResourceBase:
         await self._connect_to_parent()
@@ -165,6 +167,7 @@ async def remove_resource(self, request: RemoveResourceRequest):
 
     async def ready(self, request: ReadyRequest) -> ReadyResponse:
         self._parent_address = request.parent_address
+        await self._connect_to_parent()
 
         svcname_to_models: Mapping[Tuple[str, Subtype], List[Model]] = {}
         for subtype_model_str in Registry.REGISTERED_RESOURCE_CREATORS().keys():

src/viam/proto/robot/__init__.py

@@ -2,6 +2,7 @@
 @generated by Viam.
 Do not edit manually!
 """
+from ...gen.common.v1.common_pb2 import LogEntry
 from ...gen.robot.v1.robot_grpc import RobotServiceBase, RobotServiceStub
 from ...gen.robot.v1.robot_pb2 import (
     BlockForOperationRequest,
@@ -72,6 +73,7 @@
     "GetSessionsResponse",
     "GetStatusRequest",
     "GetStatusResponse",
+    "LogEntry",
     "LogRequest",
     "LogResponse",
     "Operation",

src/viam/robot/client.py

@@ -1,5 +1,6 @@
 import asyncio
 from dataclasses import dataclass
+from datetime import datetime
 from threading import RLock
 from typing import Any, Dict, List, Optional, Union
 
@@ -29,6 +30,8 @@
     GetOperationsResponse,
     GetStatusRequest,
     GetStatusResponse,
+    LogEntry,
+    LogRequest,
     Operation,
     ResourceNamesRequest,
     ResourceNamesResponse,
@@ -46,7 +49,7 @@
 from viam.rpc.dial import DialOptions, ViamChannel, dial
 from viam.services.service_base import ServiceBase
 from viam.sessions_client import SessionsClient
-from viam.utils import dict_to_struct
+from viam.utils import datetime_to_timestamp, dict_to_struct
 
 LOGGER = logging.getLogger(__name__)
 
@@ -753,6 +756,26 @@ async def stop_all(self, extra: Dict[ResourceName, Dict[str, Any]] = {}):
         request = StopAllRequest(extra=ep)
         await self._client.StopAll(request)
 
+    #######
+    # LOG #
+    #######
+
+    async def log(self, name: str, level: str, time: datetime, log: str, stack: str):
+        """Send log from Python module over gRPC.
+
+        Create a LogEntry object from the log to send to RDK.
+
+        Args:
+            name (str): The logger's name.
+            level (str): The level of the log.
+            time (str): The log creation time.
+            log (str): The log message.
+            stack (str): The stack information of the log.
+        """
+        entry = LogEntry(level=level, time=datetime_to_timestamp(time), logger_name=name, message=log, stack=stack)
+        request = LogRequest(logs=[entry])
+        await self._client.Log(request)
+
     ######################
     # Get Cloud Metadata #
     ######################

tests/test_module.py

@@ -156,26 +156,27 @@ async def test_remove_resource(self, module: Module):
 
     @pytest.mark.asyncio
     async def test_ready(self, module: Module):
-        p_addr = "SOME_FAKE_ADDRESS"
-        assert module._parent_address != p_addr
-        req = ReadyRequest(parent_address=p_addr)
-        resp = await module.ready(req)
-        assert module._parent_address == p_addr
-        assert len(resp.handlermap.handlers) == 2
-
-        handler = resp.handlermap.handlers[0]
-        rn = Gizmo.get_resource_name("")
-        assert handler.subtype == ResourceRPCSubtype(subtype=rn, proto_service="acme.component.gizmo.v1.GizmoService")
-        assert len(handler.models) == 1
-        model = handler.models[0]
-        assert model == "acme:demo:mygizmo"
-
-        handler = resp.handlermap.handlers[1]
-        rn = SummationService.get_resource_name("")
-        assert handler.subtype == ResourceRPCSubtype(subtype=rn, proto_service="acme.service.summation.v1.SummationService")
-        assert len(handler.models) == 1
-        model = handler.models[0]
-        assert model == "acme:demo:mysum"
+        with mock.patch("viam.module.Module._connect_to_parent"):
+            p_addr = "SOME_FAKE_ADDRESS"
+            assert module._parent_address != p_addr
+            req = ReadyRequest(parent_address=p_addr)
+            resp = await module.ready(req)
+            assert module._parent_address == p_addr
+            assert len(resp.handlermap.handlers) == 2
+
+            handler = resp.handlermap.handlers[0]
+            rn = Gizmo.get_resource_name("")
+            assert handler.subtype == ResourceRPCSubtype(subtype=rn, proto_service="acme.component.gizmo.v1.GizmoService")
+            assert len(handler.models) == 1
+            model = handler.models[0]
+            assert model == "acme:demo:mygizmo"
+
+            handler = resp.handlermap.handlers[1]
+            rn = SummationService.get_resource_name("")
+            assert handler.subtype == ResourceRPCSubtype(subtype=rn, proto_service="acme.service.summation.v1.SummationService")
+            assert len(handler.models) == 1
+            model = handler.models[0]
+            assert model == "acme:demo:mysum"
 
     def test_add_model_from_registry(self):
         mod = Module("fake")