Improve logging infrastructure

Author: kirkbrauer

packages/jumpstarter/jumpstarter/driver/base.py

@@ -27,8 +27,9 @@
     MARKER_STREAMCALL,
     MARKER_STREAMING_DRIVERCALL,
 )
-from jumpstarter.common import Metadata
+from jumpstarter.common import LogSource, Metadata
 from jumpstarter.common.resources import ClientStreamResource, PresignedRequestResource, Resource, ResourceMetadata
+from jumpstarter.exporter.logging import get_logger
 from jumpstarter.common.serde import decode_value, encode_value
 from jumpstarter.common.streams import (
     DriverStreamRequest,
@@ -86,7 +87,7 @@ def __post_init__(self):
         if hasattr(super(), "__post_init__"):
             super().__post_init__()
 
-        self.logger = logging.getLogger(self.__class__.__name__)
+        self.logger = get_logger(f"driver.{self.__class__.__name__}", LogSource.DRIVER)
         self.logger.setLevel(self.log_level)
 
     def close(self):

packages/jumpstarter/jumpstarter/exporter/exporter.py

@@ -240,6 +240,8 @@ async def status(retries=5, backoff=3):
                         # Shield the post-lease hook from cancellation and await it
                         with CancelScope(shield=True):
                             await self._update_status(ExporterStatus.AFTER_LEASE_HOOK, "Running afterLease hooks")
+                            # Pass the current session to hook executor for logging
+                            self.hook_executor.main_session = self._current_session
                             await self.hook_executor.execute_post_lease_hook(hook_context)
                             await self._update_status(ExporterStatus.AVAILABLE, "Available for new lease")
 
@@ -279,6 +281,8 @@ async def run_before_lease_hook(hook_ctx):
                                     await self._update_status(
                                         ExporterStatus.BEFORE_LEASE_HOOK, "Running beforeLease hooks"
                                     )
+                                    # Pass the current session to hook executor for logging
+                                    self.hook_executor.main_session = self._current_session
                                     await self.hook_executor.execute_pre_lease_hook(hook_ctx)
                                     await self._update_status(ExporterStatus.LEASE_READY, "Ready for commands")
                                     logger.info("beforeLease hook completed successfully")
@@ -311,6 +315,8 @@ async def run_before_lease_hook(hook_ctx):
                         # Shield the post-lease hook from cancellation and await it
                         with CancelScope(shield=True):
                             await self._update_status(ExporterStatus.AFTER_LEASE_HOOK, "Running afterLease hooks")
+                            # Pass the current session to hook executor for logging
+                            self.hook_executor.main_session = self._current_session
                             await self.hook_executor.execute_post_lease_hook(hook_context)
                             await self._update_status(ExporterStatus.AVAILABLE, "Available for new lease")
 

packages/jumpstarter/jumpstarter/exporter/hooks.py

@@ -7,9 +7,11 @@
 from dataclasses import dataclass, field
 from typing import Callable
 
+from jumpstarter.common import LogSource
 from jumpstarter.config.env import JMP_DRIVERS_ALLOW, JUMPSTARTER_HOST
 from jumpstarter.config.exporter import HookConfigV1Alpha1
 from jumpstarter.driver import Driver
+from jumpstarter.exporter.logging import get_logger
 from jumpstarter.exporter.session import Session
 
 logger = logging.getLogger(__name__)
@@ -32,6 +34,7 @@ class HookExecutor:
 
     config: HookConfigV1Alpha1
     device_factory: Callable[[], Driver]
+    main_session: Session | None = field(default=None)
     timeout: int = field(init=False)
 
     def __post_init__(self):
@@ -63,17 +66,17 @@ async def _create_hook_environment(self, context: HookContext):
                     }
                 )
 
-                yield hook_env
+                yield session, hook_env
 
-    async def _execute_hook(self, command: str, context: HookContext) -> bool:
+    async def _execute_hook(self, command: str, context: HookContext, log_source: LogSource) -> bool:
         """Execute a single hook command."""
         if not command or not command.strip():
             logger.debug("Hook command is empty, skipping")
             return True
 
         logger.info("Executing hook: %s", command.strip().split("\n")[0][:100])
 
-        async with self._create_hook_environment(context) as hook_env:
+        async with self._create_hook_environment(context) as (session, hook_env):
             try:
                 # Execute the hook command using shell
                 process = await asyncio.create_subprocess_shell(
@@ -84,6 +87,12 @@ async def _execute_hook(self, command: str, context: HookContext) -> bool:
                 )
 
                 try:
+                    # Determine which session to use for logging - prefer main session if available
+                    logging_session = self.main_session if self.main_session is not None else session
+
+                    # Create a logger with automatic source registration
+                    hook_logger = get_logger(f"hook.{context.lease_name}", log_source, logging_session)
+
                     # Stream output line-by-line for real-time logging
                     output_lines = []
 
@@ -94,7 +103,8 @@ async def read_output():
                                 break
                             line_decoded = line.decode().rstrip()
                             output_lines.append(line_decoded)
-                            logger.info("[Hook Output] %s", line_decoded)
+                            # Route hook output through the logging system
+                            hook_logger.info(line_decoded)
 
                     # Run output reading and process waiting concurrently with timeout
                     await asyncio.wait_for(asyncio.gather(read_output(), process.wait()), timeout=self.timeout)
@@ -104,8 +114,6 @@ async def read_output():
                         return True
                     else:
                         logger.error("Hook failed with return code %d", process.returncode)
-                        if output_lines:
-                            logger.error("Hook output: %s", "\n".join(output_lines))
                         return False
 
                 except asyncio.TimeoutError:
@@ -129,7 +137,7 @@ async def execute_pre_lease_hook(self, context: HookContext) -> bool:
             return True
 
         logger.info("Executing pre-lease hook for lease %s", context.lease_name)
-        return await self._execute_hook(self.config.pre_lease, context)
+        return await self._execute_hook(self.config.pre_lease, context, LogSource.BEFORE_LEASE_HOOK)
 
     async def execute_post_lease_hook(self, context: HookContext) -> bool:
         """Execute the post-lease hook."""
@@ -138,4 +146,4 @@ async def execute_post_lease_hook(self, context: HookContext) -> bool:
             return True
 
         logger.info("Executing post-lease hook for lease %s", context.lease_name)
-        return await self._execute_hook(self.config.post_lease, context)
+        return await self._execute_hook(self.config.post_lease, context, LogSource.AFTER_LEASE_HOOK)

packages/jumpstarter/jumpstarter/exporter/logging.py

@@ -1,31 +1,90 @@
 import logging
 from collections import deque
+from contextlib import contextmanager
+from threading import RLock
+from typing import TYPE_CHECKING
 
 from jumpstarter_protocol import jumpstarter_pb2
 
 from jumpstarter.common import LogSource
 
+if TYPE_CHECKING:
+    from .session import Session
+
 
 class LogHandler(logging.Handler):
     def __init__(self, queue: deque, source: LogSource = LogSource.UNSPECIFIED):
         logging.Handler.__init__(self)
         self.queue = queue
         self.listener = None
         self.source = source  # LogSource enum value
+        self._lock = RLock()
+        self._child_handlers = {}  # Dict of logger_name -> LogSource mappings
+
+    def add_child_handler(self, logger_name: str, source: LogSource):
+        """Add a child handler that will route logs from a specific logger with a different source."""
+        with self._lock:
+            self._child_handlers[logger_name] = source
+
+    def remove_child_handler(self, logger_name: str):
+        """Remove a child handler mapping."""
+        with self._lock:
+            self._child_handlers.pop(logger_name, None)
+
+    def get_source_for_record(self, record):
+        """Determine the appropriate log source for a record."""
+        with self._lock:
+            # Check if this record comes from a logger with a specific source mapping
+            logger_name = record.name
+            for mapped_logger, source in self._child_handlers.items():
+                if logger_name.startswith(mapped_logger):
+                    return source
+            return self.source
 
     def enqueue(self, record):
         self.queue.append(record)
 
     def prepare(self, record):
+        source = self.get_source_for_record(record)
         return jumpstarter_pb2.LogStreamResponse(
             uuid="",
             severity=record.levelname,
             message=self.format(record),
-            source=self.source.value,  # Convert to proto value
+            source=source.value,  # Convert to proto value
         )
 
     def emit(self, record):
         try:
             self.enqueue(self.prepare(record))
         except Exception:
             self.handleError(record)
+
+    @contextmanager
+    def context_log_source(self, logger_name: str, source: LogSource):
+        """Context manager to temporarily set a log source for a specific logger."""
+        self.add_child_handler(logger_name, source)
+        try:
+            yield
+        finally:
+            self.remove_child_handler(logger_name)
+
+
+def get_logger(name: str, source: LogSource = LogSource.SYSTEM, session: "Session" = None) -> logging.Logger:
+    """
+    Get a logger with automatic LogSource mapping.
+
+    Args:
+        name: Logger name (e.g., __name__ or custom name)
+        source: The LogSource to associate with this logger
+        session: Optional session to register with immediately
+
+    Returns:
+        A standard Python logger instance
+    """
+    logger = logging.getLogger(name)
+
+    # If session provided, register the source mapping
+    if session:
+        session.add_logger_source(name, source)
+
+    return logger

packages/jumpstarter/jumpstarter/exporter/session.py

@@ -17,8 +17,7 @@
 )
 
 from .logging import LogHandler
-from jumpstarter.common import ExporterStatus, Metadata, TemporarySocket
-from jumpstarter.common.enums import LogSource
+from jumpstarter.common import ExporterStatus, LogSource, Metadata, TemporarySocket
 from jumpstarter.common.streams import StreamRequestMetadata
 from jumpstarter.driver import Driver
 from jumpstarter.streams.common import forward_stream
@@ -74,6 +73,9 @@ def __init__(self, *args, root_device, **kwargs):
         self._logging_handler = LogHandler(self._logging_queue, LogSource.SYSTEM)
         self._status_update_event = Event()
 
+        # Map all driver logs to DRIVER source
+        self._logging_handler.add_child_handler("driver.", LogSource.DRIVER)
+
     @asynccontextmanager
     async def serve_port_async(self, port):
         server = grpc.aio.server()
@@ -153,6 +155,18 @@ def update_status(self, status: int | ExporterStatus, message: str = ""):
             self._current_status = status
         self._status_message = message
 
+    def add_logger_source(self, logger_name: str, source: LogSource):
+        """Add a log source mapping for a specific logger."""
+        self._logging_handler.add_child_handler(logger_name, source)
+
+    def remove_logger_source(self, logger_name: str):
+        """Remove a log source mapping for a specific logger."""
+        self._logging_handler.remove_child_handler(logger_name)
+
+    def context_log_source(self, logger_name: str, source: LogSource):
+        """Context manager to temporarily set a log source for a specific logger."""
+        return self._logging_handler.context_log_source(logger_name, source)
+
     async def GetStatus(self, request, context):
         """Get the current exporter status."""
         logger.debug("GetStatus() -> %s", self._current_status)