Refac: Add generic process status check for ProcessRollingLogger (#1443)

## Description

This repository needs to make the `status` method that is declared on
ProcessRollingLogger to be placed 1 structure above, so that any
ProcessActor might benefit from it.

- [x] I have reviewed the [Guidelines for Contributing](CONTRIBUTING.md)
and the [Code of Conduct](CODE_OF_CONDUCT.md).

---------

Co-authored-by: Anthony Naddeo <[email protected]>

Author: murilommen

python/Makefile

@@ -28,7 +28,7 @@ pre-commit: proto
 	poetry run pyright
 
 .PHONY: dist clean clean-test help format lint test install coverage docs livedocs default proto proto-v0 github release
-.PHONY: test-system-python format-fix blackd jupyter-kernel docs livedocs
+.PHONY: test-system-python format-fix blackd jupyter-kernel docs livedocs fix
 
 ifeq ($(shell which poetry), )
 	$(error "Can't find poetry on the path. Install it at https://python-poetry.org/docs.")
@@ -133,6 +133,8 @@ lint-fix: ## Automatically fix linting issues.
 	@$(call i, Running the linter)
 	poetry run autoflake --in-place --remove-all-unused-imports --remove-unused-variables $(src.python) $(tst.python)
 
+fix: lint-fix format-fix ## Automatically fix linting and formatting issues.
+
 format: ## Check style formatting.
 	@$(call w, format is deprecated running pre-commit instead)
 	poetry run pre-commit run --all-files

python/tests/api/logger/experimental/logger/actor/test_actor_loggers.py

@@ -402,7 +402,7 @@ def test_process_throws_after_killed(actor: Tuple[DataLogger, FakeWriter]) -> No
     won't have time to start returning false. This tests the case where something kills the process
     while something else is trying to use it.
     """
-    logger, writer = actor
+    logger, _ = actor
     if isinstance(logger, ProcessRollingLogger):
         logger = cast(ProcessRollingLogger, logger)  # type: ignore
         ms = 1689881671000
@@ -423,7 +423,7 @@ def test_process_throws_after_killed_delay(actor: Tuple[DataLogger, FakeWriter])
     Very similar to test_process_throws_after_killed but there is a delay after the process is killed
     before logging so the log() call will throw before doing any actual work with a clear error message.
     """
-    logger, writer = actor
+    logger, _ = actor
     if isinstance(logger, ProcessRollingLogger):
         logger = cast(ProcessRollingLogger, logger)  # type: ignore
         ms = 1689881671000

python/tests/api/logger/experimental/logger/actor/test_actors.py

@@ -1,7 +1,7 @@
 import multiprocessing as mp
 import os
 from dataclasses import dataclass
-from typing import List, Type, Union
+from typing import List, Type, TypeVar, Union
 
 import pytest
 
@@ -23,6 +23,9 @@ class Message2:
     pass
 
 
+StatusType = TypeVar("StatusType")
+
+
 Messages = Union[Message1, Message2]
 
 
@@ -52,7 +55,7 @@ def process_batch(
             raise Exception(f"Unknown batch type: {batch_type}")
 
 
-class CountingMPProcessActor(ProcessActor[Messages]):
+class CountingMPProcessActor(ProcessActor[Messages, StatusType]):
     def __init__(self, queue_config: QueueConfig = QueueConfig()) -> None:
         super().__init__(queue_config, queue_type=QueueType.MP)
         self.counter = Counter()
@@ -63,7 +66,7 @@ def process_batch(
         self.counter.process_batch(batch, batch_type)
 
 
-class CountingFasterFifoProcessActor(ProcessActor[Messages]):
+class CountingFasterFifoProcessActor(ProcessActor[Messages, StatusType]):
     def __init__(self, queue_config: QueueConfig = QueueConfig()) -> None:
         super().__init__(queue_config, queue_type=QueueType.FASTER_FIFO)
         self.counter = Counter()

python/whylogs/api/fugue/__init__.py

@@ -1,6 +1,7 @@
 # flake8: noqa
 from whylogs.api.usage_stats import emit_usage
 
-from .profiler import fugue_profile
+# This import has a side effect
+from .profiler import fugue_profile  # type: ignore
 
 emit_usage("fugue")

python/whylogs/api/logger/experimental/logger/actor/actor.py

@@ -184,8 +184,11 @@ def _load_messages(self) -> Optional[List[Union[MessageType, CloseMessage]]]:
 
         while self._polling_condition(len(batch), max, last_message_time, self._queue.size()):
             try:
-                batch += self._queue.get_many(timeout=self._queue_config.message_poll_wait, max=max)
-                self._logger.info(f"Got {len(batch)} messages. {self._queue.size()} remaining")
+                next_batch = self._queue.get_many(timeout=self._queue_config.message_poll_wait, max=max)
+                batch += next_batch
+                self._logger.debug(
+                    f"Adding {len(next_batch)} to poll batch of length {len(batch)}. {self._queue.size()} remaining"
+                )
             except queue.Empty:
                 if self.is_closed() and self.close_message_handled():
                     self._logger.info("Queue closed and no more messages to process.")

python/whylogs/api/logger/experimental/logger/actor/pipe_signaler.py

@@ -0,0 +1,113 @@
+import logging
+import threading as th
+from concurrent.futures import Future
+from queue import Empty
+from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar
+
+from whylogs.api.logger.experimental.logger.actor.actor import QueueConfig
+from whylogs.api.logger.experimental.logger.actor.faster_fifo_queue_wrapper import (
+    FasterQueueWrapper,
+)
+
+T = TypeVar("T")
+
+
+class PipeSignaler(th.Thread, Generic[T]):
+    """
+    A thread that listens on a pipe for messages and signals the corresponding futures.
+    This class is used in the process logger to enable synchronous logging requests across processes.
+    It's essentially a dictionary of futures that are registered by the main process and signaled by the
+    child process. A lot of the behavior is implicit because it involves properties of processes, so it's
+    worth documenting here.
+    - This thread has to be started from the main process, which means it has to be started right before the
+        process logger is started (before the os.fork under the hood). It has to be started from the main process
+        because the main process will be registering futures on it, and those can't cross the process boundary.
+    - The parent and child process each have references to the pipes and they each need to close their references,
+        which means close_child has to be called from the child process and close has to be called from the parent.
+        Calling close_child in the main processing code will have right effect.
+    - The process actor does message batching so multiple ids may be signaled even though a single batch was processed
+        because that batch could have contained multiple messages.
+    - The signaler uses Events under the hood to know when to stop working. They can be th.Events even though this
+        is being used in a multiprocessing environment because nothing the child does can affect them. Keep in mind
+        that introducing any behavior on the child side that depends on knowing whether those events are set won't work
+        though, they would have to be switched to mp.Events for that.
+    This class should really never be used by anyone in most cases. It will just slow down the main process by making
+    it wait for logging to complete, but it enables a lot of testing and debugging.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.daemon = True
+        self._logger = logging.getLogger(__name__)
+        self._queue_config = QueueConfig()
+        self.queue: FasterQueueWrapper[Tuple[str, Optional[Exception], Optional[T]]] = FasterQueueWrapper(
+            self._queue_config
+        )
+        self.futures: Dict[str, "Future[Any]"] = {}
+        self._end_polling = th.Event()
+        self._done = th.Event()
+
+    def signal(self, result: Tuple[str, Optional[Exception], Optional[T]]) -> None:
+        """
+        Signal that a message was handled by sending a tuple of (message id, exception, data).
+        data and exception can be None.
+        This should be called from the child process.
+        """
+        self.queue.send(result)
+
+    def signal_many(self, results: List[Tuple[str, Optional[Exception], Optional[T]]]) -> None:
+        self.queue.send_many(results)
+
+    def register(self, future: "Future[T]", message_id: str) -> None:
+        """
+        Register a future to be signaled when the message id is received.
+        This should be called from the parent process.
+        """
+        self._logger.debug(f"Received register request for id {message_id}")
+        self.futures[message_id] = future
+
+    def _start_poll_conn(self) -> None:
+        while not self._end_polling.is_set():
+            try:
+                messages = self.queue.get_many(
+                    timeout=self._queue_config.message_poll_wait,
+                    max=self._queue_config.max_batch_size,
+                )
+
+                for message_id, exception, data in messages:
+                    self._logger.debug(f"Received message id {message_id}")
+                    future: Optional["Future[Any]"] = self.futures.pop(message_id, None)
+                    if future is not None:
+                        self._logger.debug(f"Setting result for message id {message_id} {exception}")
+                        if exception is not None:
+                            future.set_exception(exception)
+                        else:
+                            # Doing this in the else avoid testing `data is None`, which would be
+                            # wrong because data can potentially be None if T is Optional
+                            future.set_result(data)
+
+            except Empty:
+                continue
+            except Exception as e:
+                self._logger.exception(f"Error in queue {e}")
+                break
+
+        self._done.set()
+
+    def run(self) -> None:
+        self._start_poll_conn()
+
+    def close_child(self) -> None:
+        """
+        This method is no longer needed as queues do not require manual closing of file descriptors.
+        """
+
+    def close(self) -> None:
+        """
+        Closes the thread and all resources. This should be
+        called from the parent side.
+        """
+        self._end_polling.set()
+        self._done.wait()
+        self.queue.close()  # type: ignore[reportUnknownMemberType]
+        self.join()

python/whylogs/api/logger/experimental/logger/actor/process_actor.py

@@ -1,34 +1,46 @@
 import multiprocessing as mp
 import signal
 import sys
+from concurrent.futures import Future
 from enum import Enum
-from typing import Generic, TypeVar
+from typing import Generic, Optional, Sequence, TypeVar, Union
 
 from whylogs.api.logger.experimental.logger.actor.actor import (
     Actor,
     QueueConfig,
     QueueWrapper,
 )
+from whylogs.api.logger.experimental.logger.actor.future_util import wait_result_while
+from whylogs.api.logger.experimental.logger.actor.pipe_signaler import PipeSignaler
+from whylogs.api.logger.experimental.logger.actor.process_rolling_logger_messages import (
+    ProcessStatusMessage,
+    SyncMessage,
+)
 from whylogs.api.logger.experimental.logger.actor.signal_util import suspended_signals
 
+StatusType = TypeVar("StatusType")
+ProcessMessageType = TypeVar("ProcessMessageType")
+
 
 class QueueType(Enum):
     MP = "MP"
     FASTER_FIFO = "FASTER_FIFO"
 
 
-ProcessMessageType = TypeVar("ProcessMessageType")
-
-
-class ProcessActor(Actor[ProcessMessageType], mp.Process, Generic[ProcessMessageType]):
+class ProcessActor(
+    Actor[Union[ProcessMessageType, ProcessStatusMessage]], mp.Process, Generic[ProcessMessageType, StatusType]
+):
     """
     Subclass of Actor that uses a process to process messages.
     """
 
     _wrapper: QueueWrapper[ProcessMessageType]
 
     def __init__(
-        self, queue_config: QueueConfig = QueueConfig(), queue_type: QueueType = QueueType.FASTER_FIFO
+        self,
+        queue_config: QueueConfig = QueueConfig(),
+        queue_type: QueueType = QueueType.FASTER_FIFO,
+        sync_enabled: bool = False,
     ) -> None:
         if queue_type == QueueType.MP:
             from whylogs.api.logger.experimental.logger.actor.mp_queue_wrapper import (
@@ -45,6 +57,8 @@ def __init__(
         else:
             raise ValueError(f"Unknown queue type: {queue_type}")
 
+        self._sync_enabled = sync_enabled
+        self._pipe_signaler: Optional[PipeSignaler[StatusType]] = PipeSignaler() if self._sync_enabled is True else None
         self._event = mp.Event()
         self._is_closed = mp.Event()
         self._close_handled = mp.Event()
@@ -87,6 +101,31 @@ def close(self) -> None:
         super().close()
         self._wrapper.close()
 
+    def _signal(self, messages: Sequence[SyncMessage] = [], error: Optional[Exception] = None) -> None:
+        if self._pipe_signaler is None:
+            return
+
+        for message in messages:
+            if message.sync:
+                self._pipe_signaler.signal((message.id, error, None))
+
+    def status(self, timeout: Optional[float] = 1.0) -> StatusType:
+        """
+        Get the internal status of the Process Actor. Used for diagnostics and debugging.
+        This is always synchronous and requires the ProcessActor to be created with sync_enabled=True.
+        """
+        if self._pipe_signaler is None:
+            raise Exception(
+                "Can't log synchronously without a pipe signaler. Initialize the process logger with sync_enabled=True."
+            )
+
+        # add a sync flag to Message
+        message = ProcessStatusMessage()
+        future: "Future[StatusType]" = Future()
+        self._pipe_signaler.register(future, message.id)
+        self.send(message)
+        return wait_result_while(future, self.is_alive)
+
     def run(self) -> None:
         try:
             with suspended_signals(signal.SIGINT, signal.SIGTERM):