Real-Time worker logs (#60)

Co-authored-by: Damian Janiga <[email protected]>

Author: Hevagog

src/logger/__init__.py

@@ -1,10 +1,12 @@
 from logger.numeric_logger import get_csv_logger
 from logger.orchestration_logger import log_docker_logs
+from logger.stream_reader import stream_reader
 from logger.u_logger import configure_logger, get_logger
 
 __all__ = [
     "get_csv_logger",
     "configure_logger",
     "get_logger",
     "log_docker_logs",
+    "stream_reader",
 ]

src/logger/orchestration_logger.py

@@ -1,41 +1,12 @@
 import subprocess
 from logging import Logger
 from pathlib import Path
-from typing import List
 
-from logger.utils import get_external_console_logging
+from logger.utils import get_external_console_logging, get_services, get_services_id
 
 TAIL_LINES = 100
-WORKER_PREFIX = "worker"
-SIM_SERVICE_NAME = "simulation_server"
-
-
-def get_services() -> List[str]:
-    """Return a list of all docker compose service names."""
-    try:
-        container_ids = subprocess.check_output(
-            ["docker", "ps", "-q"], text=True
-        ).splitlines()
-
-        services = set()
-        for cid in container_ids:
-            inspect_output = subprocess.check_output(
-                [
-                    "docker",
-                    "inspect",
-                    "--format",
-                    '{{ index .Config.Labels "com.docker.compose.service" }}',
-                    cid,
-                ],
-                text=True,
-            ).strip()
-            if inspect_output and inspect_output != "<no value>":
-                services.add(inspect_output)
-        return list(services)
-    except Exception as e:
-        raise RuntimeError(
-            "Failed to get Docker services. Ensure Docker is running and you have access to it."
-        ) from e
+WORKER_PREFIX = "simulation-worker"
+SIM_SERVICE_NAME = "simulation-server"
 
 
 def _start_external_xterm_log_terminal(title: str, command: str) -> None:
@@ -57,6 +28,8 @@ def _start_external_log_terminal(title: str, command: str) -> None:
             title,
             command,
         )
+    except Exception as e:
+        raise RuntimeError(f"Failed to open external terminal for logging: {e}") from e
 
 
 def _start_file_logging(command: str) -> None:
@@ -69,7 +42,7 @@ def _start_file_logging(command: str) -> None:
     )
 
 
-def log_docker_logs(logger: Logger) -> None:
+def log_docker_logs(logger: Logger, disperse_worker_logs: bool = True) -> None:
     """Fetch and log docker logs for workers and simulation service."""
     assert logger is not None, "Logger must be initialized before logging Docker logs."
 
@@ -78,7 +51,7 @@ def log_docker_logs(logger: Logger) -> None:
     log_dir.mkdir(exist_ok=True)
 
     try:
-        services = get_services()
+        services = get_services(logger)
     except RuntimeError as e:
         logger.error(f"Error fetching Docker services: {e}")
         return
@@ -89,35 +62,69 @@ def log_docker_logs(logger: Logger) -> None:
         return
 
     worker_services = sorted([s for s in services if s.startswith(WORKER_PREFIX)])
-
-    worker_log_path = (log_dir / "docker_workers.log").resolve()
     sim_log_path = (log_dir / "docker_simulation_server.log").resolve()
 
+    logger.info(f"Logging Docker worker services: {worker_services}")
+
     if worker_services:
-        worker_cmd = f"docker compose logs --tail {TAIL_LINES} --follow {' '.join(worker_services)}"
+        if disperse_worker_logs:
+            for worker_service in worker_services:
+                worker_log_path = (
+                    log_dir / f'docker_{worker_service.replace("-", "_")}.log'
+                ).resolve()
+                worker_cmd = (
+                    f"docker logs --tail {TAIL_LINES} --follow {worker_service}"
+                )
+
+                if get_external_console_logging():
+                    _start_external_log_terminal(
+                        f"Docker Logs: {worker_service}",
+                        f"{worker_cmd} | stdbuf -oL -eL tee '{worker_log_path}'",
+                    )
+                    logger.info(
+                        f"Opened external terminal for {worker_service} logs (also logging to {worker_log_path})."
+                    )
+
+                _start_file_logging(f"{worker_cmd} > '{worker_log_path}'")
+                logger.info(f"Logging {worker_service} logs to file: {worker_log_path}")
+        else:
+            worker_log_path = (log_dir / "docker_workers.log").resolve()
+            worker_services = sorted(
+                [s for s in get_services_id() if s.startswith(WORKER_PREFIX)]
+            )
+
+            logger.info(f"Aggregating logs for worker services: {worker_services}")
+            worker_cmd = f"docker compose logs --tail {TAIL_LINES} --follow {' '.join(worker_services)}"
+
+            if get_external_console_logging():
+                _start_external_log_terminal(
+                    "Docker Worker Logs",
+                    f"{worker_cmd} | stdbuf -oL -eL tee '{worker_log_path}'",
+                )
+                logger.info(
+                    f"Opened external terminal for aggregated worker logs (also logging to {worker_log_path})."
+                )
+
+            _start_file_logging(f"{worker_cmd} > '{worker_log_path}'")
+            logger.info(f"Logging aggregated worker logs to file: {worker_log_path}")
+
+    if SIM_SERVICE_NAME in services:
+        sim_cmd = f"docker logs --tail {TAIL_LINES} --follow {SIM_SERVICE_NAME}"
 
         if get_external_console_logging():
             _start_external_log_terminal(
-                "Docker Worker Logs",
-                f"{worker_cmd} | stdbuf -oL -eL tee '{worker_log_path}'",
+                "Docker Simulation Server Logs",
+                f"{sim_cmd} | stdbuf -oL -eL tee '{sim_log_path}'",
             )
             logger.info(
-                f"Opened external terminal for worker logs (also logging to {worker_log_path})."
+                f"Opened external terminal for simulation logs (also logging to {sim_log_path})."
             )
 
-        _start_file_logging(f"{worker_cmd} > '{worker_log_path}'")
-        logger.info(f"Logging worker logs to file: {worker_log_path}")
-
-    sim_cmd = f"docker logs --tail {TAIL_LINES} --follow {SIM_SERVICE_NAME}"
-
-    if get_external_console_logging():
-        _start_external_log_terminal(
-            "Docker Simulation Server Logs",
-            f"{sim_cmd} | stdbuf -oL -eL tee '{sim_log_path}'",
+        _start_file_logging(f"{sim_cmd} > '{sim_log_path}'")
+        logger.info(f"Logging simulation logs to file: {sim_log_path}")
+    elif SIM_SERVICE_NAME not in services and any(
+        s.startswith(WORKER_PREFIX) for s in services
+    ):
+        logger.warning(
+            f"Simulation service '{SIM_SERVICE_NAME}' not found. Skipping simulation server logs."
         )
-        logger.info(
-            f"Opened external terminal for simulation logs (also logging to {sim_log_path})."
-        )
-
-    _start_file_logging(f"{sim_cmd} > '{sim_log_path}'")
-    logger.info(f"Logging simulation logs to file: {sim_log_path}")

src/logger/stream_reader.py

@@ -0,0 +1,17 @@
+import io
+from typing import Callable
+
+
+def stream_reader(
+    stream: io.TextIOWrapper | None,
+    output_list: list[str],
+    log_with_level: Callable,
+) -> None:
+    if stream:
+        try:
+            for line in iter(stream.readline, ""):
+                line_content = line.strip()
+                log_with_level(line_content)
+                output_list.append(line_content)
+        finally:
+            stream.close()

src/logger/utils.py

@@ -1,5 +1,7 @@
 import os
-from typing import Any, Dict
+import subprocess
+from logging import Logger
+from typing import Any, Dict, List
 
 
 def get_logger_profile() -> str:
@@ -87,3 +89,43 @@ def get_external_console_logging() -> bool:
         The value of external_docker_log_console.
     """
     return bool(get_logging_output()["external_docker_log_console"])
+
+
+def get_services(logger: Logger) -> List[str]:
+    """Return a list of all docker compose service names, including numbered workers."""
+    try:
+        compose_services = subprocess.check_output(
+            ["docker", "ps", "--format", "'{{.Names}}'"], text=True
+        ).splitlines()
+        return [service.strip("'") for service in compose_services if service.strip()]
+    except Exception as e:
+        raise RuntimeError(
+            "Failed to get Docker services. Ensure Docker is running and you have access to it."
+        ) from e
+
+
+def get_services_id() -> List[str]:
+    try:
+        container_ids = subprocess.check_output(
+            ["docker", "ps", "-q"], text=True
+        ).splitlines()
+
+        services = set()
+        for cid in container_ids:
+            inspect_output = subprocess.check_output(
+                [
+                    "docker",
+                    "inspect",
+                    "--format",
+                    '{{ index .Config.Labels "com.docker.compose.service" }}',
+                    cid,
+                ],
+                text=True,
+            ).strip()
+            if inspect_output and inspect_output != "<no value>":
+                services.add(inspect_output)
+        return list(services)
+    except Exception as e:
+        raise RuntimeError(
+            "Failed to get Docker services. Ensure Docker is running and you have access to it."
+        ) from e

src/services/simulation_service/core/connectors/conn_utils/__init__.py

@@ -0,0 +1,89 @@
+import subprocess
+import threading
+from typing import Callable
+
+
+class ManagedSubprocess:
+    def __init__(
+        self,
+        command_args: list[str],
+        stream_reader_func: Callable,
+        logger_info_func: Callable,
+        logger_error_func: Callable,
+        text: bool = True,
+    ):
+        self.command_args = command_args
+        self.text = text
+        self.stream_reader_func = stream_reader_func
+        self.logger_info_func = logger_info_func
+        self.logger_error_func = logger_error_func
+
+        self.process: subprocess.Popen | None = None
+        self.stdout_thread: threading.Thread | None = None
+        self.stderr_thread: threading.Thread | None = None
+        self.stdout_lines: list[str] = []
+        self.stderr_lines: list[str] = []
+
+    def __enter__(self) -> subprocess.Popen:
+        self.logger_info_func(f"Starting subprocess: {' '.join(self.command_args[:3])}")
+        try:
+            self.process = subprocess.Popen(
+                self.command_args,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=self.text,
+            )
+        except Exception:
+            self.logger_error_func(
+                f"Failed to start subprocess with command: {' '.join(self.command_args[:3])}"
+            )
+            self.process = None
+            raise  # Re-raise the exception from Popen
+
+        if self.process.stdout is None or self.process.stderr is None:
+            if self.process:
+                self.process.kill()
+                self.process.wait()
+            raise RuntimeError("Subprocess stdout/stderr streams are not available.")
+
+        self.stdout_thread = threading.Thread(
+            target=self.stream_reader_func,
+            args=(self.process.stdout, self.stdout_lines, self.logger_info_func),
+            daemon=True,
+        )
+        self.stderr_thread = threading.Thread(
+            target=self.stream_reader_func,
+            args=(self.process.stderr, self.stderr_lines, self.logger_error_func),
+            daemon=True,
+        )
+
+        self.stdout_thread.start()
+        self.stderr_thread.start()
+
+        return self.process
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.process and self.process.poll() is None:
+            self.logger_warning_func = getattr(
+                self, "logger_warning_func", self.logger_error_func
+            )
+            self.logger_warning_func(
+                f"Subprocess (PID: {self.process.pid}) still running upon exiting context. Exception: {exc_type}. Terminating."
+            )
+            self.process.terminate()
+            try:
+                self.process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                self.logger_warning_func(
+                    f"Subprocess (PID: {self.process.pid}) did not terminate gracefully. Killing."
+                )
+                self.process.kill()
+                self.process.wait()
+
+        # Join threads to ensure all output is processed and resources are released
+        if self.stdout_thread and self.stdout_thread.is_alive():
+            self.stdout_thread.join(timeout=2)
+        if self.stderr_thread and self.stderr_thread.is_alive():
+            self.stderr_thread.join(timeout=2)
+
+        return False