[BUZZOK-25745] Add support for streaming drum server run outputs (#1419)

* Add support for streaming drum server outputs

* Lint and update logic

* Fix test

* Add tests

* Bump version

* Keep version

* Rename

* Change print to logger

* Update tests to logger

* Tweak changes

* Lint

Author: mjnitz02

custom_model_runner/datarobot_drum/drum/root_predictors/drum_server_utils.py

@@ -39,13 +39,14 @@ def wait_for_server(url, timeout):
             raise TimeoutError("Server failed to start: url: {}".format(url))
 
 
-def _run_server_thread(cmd, process_obj_holder, verbose=True):
+def _run_server_thread(cmd, process_obj_holder, verbose=True, stream_output=False):
     _exec_shell_cmd(
         cmd,
         "Failed in {} command line! {}".format(ArgumentsOptions.MAIN_COMMAND, cmd),
         assert_if_fail=False,
         process_obj_holder=process_obj_holder,
         verbose=verbose,
+        stream_output=stream_output,
     )
 
 
@@ -89,6 +90,7 @@ def __init__(
         max_workers=None,
         cmd_override=None,
         port: Optional[int] = None,
+        stream_output: Optional[bool] = False,
     ):
         self.port = port or DrumUtils.find_free_port()
         self.server_address = "localhost:{}".format(self.port)
@@ -123,12 +125,18 @@ def __init__(
         self._wait_for_server_timeout = wait_for_server_timeout
         self._max_workers = max_workers
         self._cmd_override = cmd_override
+        self._stream_output = stream_output
 
     def __enter__(self):
         self._server_thread = self._thread_class(
             name="DRUM Server",
             target=_run_server_thread,
-            args=(self.get_command(), self._process_object_holder, self._verbose),
+            args=(
+                self.get_command(),
+                self._process_object_holder,
+                self._verbose,
+                self._stream_output,
+            ),
         )
         self._server_thread.start()
         time.sleep(0.5)

custom_model_runner/datarobot_drum/drum/root_predictors/utils.py

@@ -10,6 +10,8 @@
 import shlex
 import subprocess
 import time
+from queue import Queue, Empty
+from threading import Thread
 
 from datarobot_drum.drum.common import get_drum_logger
 from datarobot_drum.drum.enum import (
@@ -102,6 +104,47 @@ def _create_custom_model_dir(
     return custom_model_dir
 
 
+def _queue_output(stdout, stderr, queue):
+    """Helper function to stream output from subprocess to a queue."""
+    for line in iter(stdout.readline, b""):
+        queue.put(line)
+    for line in iter(stderr.readline, b""):
+        queue.put(line)
+    stdout.close()
+    stderr.close()
+
+
+def _stream_p_open(subprocess_popen: subprocess.Popen):
+    """Wraps the Popen object to stream output in a separate thread.
+    This realtime output of the stdout and stderr of the process
+    is streamed to the terminal.
+
+    This logic was added because there is not a direct python execution flow for DRUM.
+    SEE: https://datarobot.atlassian.net/browse/RAPTOR-12510
+    """
+    logger_queue = Queue()
+    logger_thread = Thread(
+        target=_queue_output,
+        daemon=True,
+        args=(subprocess_popen.stdout, subprocess_popen.stderr, logger_queue),
+    )
+    logger_thread.start()
+    while True:
+        try:
+            # Stream output if available
+            line = logger_queue.get_nowait()
+            logger.info(line.strip()) if len(line.strip()) > 0 else None
+        except Empty:
+            # Check if the process has terminated
+            if subprocess_popen.poll() is not None:
+                break
+            time.sleep(1)
+        except Exception:
+            break
+    # Output has already been displayed
+    return "", ""
+
+
 def _exec_shell_cmd(
     cmd,
     err_msg,
@@ -110,6 +153,7 @@ def _exec_shell_cmd(
     env=os.environ,
     verbose=True,
     capture_output=True,
+    stream_output=False,
 ):
     """
     Wrapper used by tests and validation to run shell command.
@@ -132,7 +176,10 @@ def _exec_shell_cmd(
         process_obj_holder.process = p
 
     if capture_output:
-        (stdout, stderr) = p.communicate()
+        if stream_output:
+            (stdout, stderr) = _stream_p_open(p)
+        else:
+            (stdout, stderr) = p.communicate()
     else:
         stdout, stderr = None, None
 

tests/unit/datarobot_drum/drum/root_predictors/test_drum_server_utils.py

@@ -5,6 +5,12 @@
 #  This is proprietary source code of DataRobot, Inc. and its affiliates.
 #  Released under the terms of DataRobot Tool and Utility Agreement.
 #
+import io
+import os
+import subprocess
+import threading
+import time
+from queue import Queue
 from typing import Tuple, Any
 from unittest.mock import patch, Mock
 
@@ -13,6 +19,11 @@
 from datarobot_drum.drum.enum import TargetType, ArgumentsOptions
 from datarobot_drum.drum.root_predictors.drum_server_utils import DrumServerRun, DrumServerProcess
 
+from custom_model_runner.datarobot_drum.drum.root_predictors.utils import (
+    _queue_output,
+    _stream_p_open,
+)
+
 
 @pytest.fixture
 def module_under_test():
@@ -110,7 +121,7 @@ class TestingThread:
     def __init__(self, name, target, args: Tuple[str, DrumServerProcess, Any]):
         self.name = name
         self.target = target
-        self.command, self.process_object_holder, self.verbose = args
+        self.command, self.process_object_holder, self.verbose, self.stream_output = args
 
     def start(self):
         self.process_object_holder.process = Mock(pid=123)
@@ -141,3 +152,161 @@ def test_calls_thread_correctly(self, mock_get_command, runner):
         with runner:
             pass
         assert runner.server_thread.command == mock_get_command.return_value
+
+
+class TestStreamingOutput:
+    def test_queue_output_handles_stdout_and_stderr(self):
+        # Create mock objects for stdout, stderr and queue
+        mock_stdout = Mock()
+        mock_stderr = Mock()
+        test_queue = Queue()
+
+        # Configure stdout mock to return a sequence of lines and then stop
+        stdout_lines = [b"stdout line 1\n", b"stdout line 2\n", b""]
+        mock_stdout.readline.side_effect = stdout_lines
+
+        # Configure stderr mock to return a sequence of lines and then stop
+        stderr_lines = [b"stderr line 1\n", b"stderr line 2\n", b""]
+        mock_stderr.readline.side_effect = stderr_lines
+
+        # Call the function
+        _queue_output(mock_stdout, mock_stderr, test_queue)
+
+        # Verify the queue contains all the expected output lines
+        expected_lines = [
+            b"stdout line 1\n",
+            b"stdout line 2\n",
+            b"stderr line 1\n",
+            b"stderr line 2\n",
+        ]
+
+        # Get all items from the queue and compare with expected
+        actual_lines = []
+        while not test_queue.empty():
+            actual_lines.append(test_queue.get())
+
+        assert actual_lines == expected_lines
+
+        # Verify close was called on both streams
+        mock_stdout.close.assert_called_once()
+        mock_stderr.close.assert_called_once()
+
+    @patch("custom_model_runner.datarobot_drum.drum.root_predictors.utils.logger.info")
+    def test_stream_p_open_handles_process_output(self, mock_logger):
+        # Create a mock Popen object
+        mock_process = Mock(spec=subprocess.Popen)
+
+        # Create real BytesIO objects that will be read by the real _queue_output function
+        mock_process.stdout = io.BytesIO(b"stdout line 1\nstdout line 2\n")
+        mock_process.stderr = io.BytesIO(b"stderr line 1\nstderr line 2\n")
+
+        # Configure process to terminate after first poll
+        mock_process.poll.side_effect = [None, 0]
+
+        # Call the function with real stdout/stderr streams
+        stdout, stderr = _stream_p_open(mock_process)
+
+        # Verify the results
+        assert stdout == ""
+        assert stderr == ""
+
+        # Verify logger.info was called with the expected content
+        expected_outputs = [
+            b"stdout line 1",
+            b"stdout line 2",
+            b"stderr line 1",
+            b"stderr line 2",
+        ]
+
+        # Get the actual arguments passed to logger.info
+        actual_outputs = [call_args[0][0] for call_args in mock_logger.call_args_list]
+
+        # Verify all expected outputs were logged
+        assert len(actual_outputs) == len(expected_outputs)
+        for expected in expected_outputs:
+            assert expected in actual_outputs
+
+        # Verify poll was called to check for process termination
+        assert mock_process.poll.call_count == 2
+
+    @patch("custom_model_runner.datarobot_drum.drum.root_predictors.utils.logger.info")
+    def test_stream_p_open_handles_empty_streams(self, mock_logger):
+        # Create a mock Popen object with empty streams
+        mock_process = Mock(spec=subprocess.Popen)
+        mock_process.stdout = io.BytesIO(b"")
+        mock_process.stderr = io.BytesIO(b"")
+
+        # Configure process to terminate immediately
+        mock_process.poll.return_value = 0
+
+        # Call the function
+        stdout, stderr = _stream_p_open(mock_process)
+
+        # Verify results
+        assert stdout == ""
+        assert stderr == ""
+        mock_logger.assert_not_called()  # No output to log
+        mock_process.poll.assert_called_once()  # Process termination was checked
+
+    @patch("custom_model_runner.datarobot_drum.drum.root_predictors.utils.logger.info")
+    def test_stream_p_open_handles_delayed_output(self, mock_logger):
+        # Create a mock process
+        mock_process = Mock(spec=subprocess.Popen)
+
+        # Create pipes that we can write to after the function starts
+        r_stdout, w_stdout = os.pipe()
+        r_stderr, w_stderr = os.pipe()
+
+        # Convert file descriptors to file-like objects
+        mock_process.stdout = os.fdopen(r_stdout, "rb")
+        mock_process.stderr = os.fdopen(r_stderr, "rb")
+
+        # Set up process to run and then terminate
+        mock_process.poll.side_effect = [None, None, 0]
+
+        # Create a thread to run the function we're testing
+        def run_stream_p_open():
+            return _stream_p_open(mock_process)
+
+        thread = threading.Thread(target=run_stream_p_open)
+        thread.daemon = True
+        thread.start()
+
+        # Write data to the pipes
+        time.sleep(0.1)  # Give the function time to start
+        os.write(w_stdout, b"delayed stdout\n")
+        os.write(w_stderr, b"delayed stderr\n")
+        os.close(w_stdout)
+        os.close(w_stderr)
+
+        # Wait for the thread to finish
+        thread.join(timeout=2)
+        assert not thread.is_alive()
+
+        # Verify logger.info was called with the expected content
+        mock_logger.assert_any_call(b"delayed stdout")
+        mock_logger.assert_any_call(b"delayed stderr")
+
+    @patch("custom_model_runner.datarobot_drum.drum.root_predictors.utils.logger.info")
+    def test_stream_p_open_ignores_empty_lines(self, mock_logger):
+        # Create a mock Popen object
+        mock_process = Mock(spec=subprocess.Popen)
+
+        # Create streams with empty lines
+        mock_process.stdout = io.BytesIO(b"regular line\n   \n\n")
+        mock_process.stderr = io.BytesIO(b"error line\n  \n")
+
+        # Configure process to terminate after output is processed
+        mock_process.poll.side_effect = [None, 0]
+
+        # Call the function
+        stdout, stderr = _stream_p_open(mock_process)
+
+        # Verify results
+        assert stdout == ""
+        assert stderr == ""
+
+        # Verify only non-empty lines were logged
+        mock_logger.assert_any_call(b"regular line")
+        mock_logger.assert_any_call(b"error line")
+        assert mock_logger.call_count == 2  # Only two lines should be logged