Add ability to capture stdout

Author: TheGreatCabbage

scheduler/ProcessTask.py

@@ -27,9 +27,14 @@
 
 class ProcessTask(Task):
     def __init__(
-        self, process: Process, queue: Queue, exc_queue=None, subtasks: int = 0
+        self,
+        process: Process,
+        queue: Queue,
+        exc_queue=None,
+        stdout_queue=None,
+        subtasks: int = 0,
     ):
-        super(ProcessTask, self).__init__(queue, exc_queue, subtasks)
+        super(ProcessTask, self).__init__(queue, exc_queue, stdout_queue, subtasks)
 
         self.process = process
 

scheduler/Scheduler.py

@@ -36,7 +36,7 @@
 from scheduler.ProcessTask import ProcessTask
 from scheduler.Task import Task
 from scheduler.ThreadTask import ThreadTask
-from scheduler.utils import SchedulerException, TaskFailedException
+from scheduler.utils import SchedulerException, TaskFailedException, StdOut
 
 
 class Scheduler:
@@ -62,6 +62,7 @@ def __init__(
         shared_memory_threshold: int = 1e7,
         run_in_thread: bool = False,
         raise_exceptions: bool = False,
+        capture_stdout: bool = False,
     ):
         """
         :param progress_callback: a function taking the number of finished tasks and the total number of tasks, which is
@@ -72,13 +73,17 @@ def __init__(
         to be below the threshold, the number of simultaneous tasks will be increased
         :param cpu_update_interval: the time, in seconds, between consecutive CPU usage checks when `dynamic` is enabled
         :param shared_memory: whether to use shared memory if possible
-        :param shared_memory_threshold: the minimum size of a Numpy array which will cause it to be transferred using shared memory if possible
-        :param run_in_thread: if True, a single task will be run in a thread instead of a process.
+        :param shared_memory_threshold: the minimum size of a Numpy array which will cause it to be transferred
+        using shared memory if possible
+        :param run_in_thread: if True, a single task will be run in a thread instead of a process. This reduces
+        the overhead (caused by spawning processes instead of forking) on Windows/macOS systems
         :param raise_exceptions: if True, Exceptions raised in processes will also be raised in
         the process which the Scheduler was started in.
-        This reduces the overhead (caused by spawning processes instead of forking) on Windows/macOS systems
+        :param capture_stdout: if True, `stdout` from processes will be captured and written to the main process'
+        `stdout`.
         """
         self.raise_exceptions = raise_exceptions
+        self.capture_stdout = capture_stdout
 
         self.run_in_thread = run_in_thread
         if self.run_in_thread and shared_memory:
@@ -178,7 +183,13 @@ def add(
             queue = MTQueue()
             exc_queue = MTQueue()
 
-            _args = (queue, self.mgr, self.shared_memory_threshold, exc_queue) + args
+            _args = (
+                queue,
+                self.mgr,
+                self.shared_memory_threshold,
+                exc_queue,
+                None,
+            ) + args
             _wrapper = functools.partial(wrapper, target)
 
             task = ThreadTask(
@@ -190,12 +201,25 @@ def add(
         else:
             queue = queue_type()
             exc_queue = queue_type()
+            stdout_queue = queue_type()
 
-            _args = (queue, self.mgr, self.shared_memory_threshold, exc_queue) + args
+            _args = (
+                queue,
+                self.mgr,
+                self.shared_memory_threshold,
+                exc_queue,
+                stdout_queue,
+            ) + args
             _wrapper = functools.partial(wrapper, target)
 
             process = process_type(target=_wrapper, args=_args)
-            task = ProcessTask(process, queue, exc_queue=exc_queue, subtasks=subtasks)
+            task = ProcessTask(
+                process,
+                queue,
+                exc_queue=exc_queue,
+                stdout_queue=stdout_queue,
+                subtasks=subtasks,
+            )
 
         self.tasks.append(task)
 
@@ -337,6 +361,7 @@ def terminate(self) -> None:
         """Terminates all running tasks by killing their processes."""
         if not self.terminated:
             [t.terminate() for t in self.tasks]
+            [self.stdout(t) for t in self.tasks]
             self.terminated = True
 
         self._shutdown()
@@ -412,8 +437,11 @@ def _update(self) -> None:
 
         for t in self.running_tasks:
             t.update()
+            self.stdout(t)
+
             if self.raise_exceptions and t.failed:
                 self.failed = True
+
                 if t.exception_tb:
                     raise TaskFailedException(t.exception_tb)
                 else:
@@ -431,6 +459,13 @@ def _update(self) -> None:
         if schedule_new_tasks:
             self._schedule_tasks()
 
+    def stdout(self, task: Task) -> None:
+        if self.capture_stdout:
+            text = task.get_stdout()
+
+            if text:
+                sys.stdout.write(text)
+
     def _start(self) -> None:
         """
         Starts the scheduler running its tasks.
@@ -514,16 +549,29 @@ def wrapper(
     manager: Optional["SharedMemoryManager"],
     threshold: int,
     exc_queue: Queue = None,
+    stdout_queue: Queue = None,
     *args: Any,
 ) -> None:
     """
     Wrapper which calls a function with its specified arguments and puts the output in a queue.
 
+    This function will be the Callable executed in a process/thread.
+
     :param function: the function which will be executed
     :param queue: a Queue object which may be used to transfer data between processes
     :param manager: a SharedMemoryManager or None; used to handle shared memory between processes
+    :param threshold:
+    :param exc_queue:
+    :param stdout_queue:
     """
+    stdout = None
+
     try:
+        if stdout_queue:
+            stdout = StdOut(stdout_queue)
+            sys.stdout = stdout
+            sys.stderr = stdout
+
         result = function(*args)
         out = []
 
@@ -542,6 +590,9 @@ def wrapper(
         else:
             out = tuple(out)
 
+        if stdout:
+            stdout.update(force=True)
+
         queue.put(out)
 
     except Exception as e:

scheduler/Task.py

@@ -20,6 +20,7 @@
 #  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 #  SOFTWARE.
 from abc import ABC, abstractmethod
+from typing import Optional
 
 
 class Task(ABC):
@@ -28,16 +29,18 @@ class Task(ABC):
     passed to the process, so that the process can put its output in the queue.
     """
 
-    def __init__(self, queue, exc_queue=None, subtasks: int = 0):
+    def __init__(self, queue, exc_queue=None, stdout_queue=None, subtasks: int = 0):
         """
-
         :param queue: Queue for data to be passed.
         :param exc_queue: Queue for exceptions to be passed.
         :param subtasks: Number of subtasks.
         """
         self.queue = queue
         self.exc_queue = exc_queue
 
+        self.stdout_queue = stdout_queue
+        self.stdout_text = None
+
         self.running = False
         self.finished = False
 
@@ -89,3 +92,29 @@ def has_exception(self) -> bool:
         Returns whether the task has raised an exception.
         """
         return self.exc_queue and not self.exc_queue.empty()
+
+    def has_stdout(self) -> bool:
+        """
+        Returns whether the task has provided any `stdout`.
+        """
+        return self.stdout_queue and not self.stdout_queue.empty()
+
+    def get_stdout(self) -> Optional[str]:
+        """
+        Returns the `stdout` from the task.
+        """
+        text = self.stdout_text or ""
+        if self.has_stdout():
+            lines = self.stdout_queue.get()
+
+            if lines:
+                if isinstance(lines, tuple):
+                    lines = "\n".join(lines)
+
+                text = f"{text}{lines}"
+
+        if text:
+            self.stdout_text = None
+            return text
+
+        return None

scheduler/ThreadTask.py

@@ -28,7 +28,8 @@
 
 class ThreadTask(Task):
     def __init__(self, thread: Thread, queue: Queue, exc_queue=None, subtasks: int = 0):
-        super(ThreadTask, self).__init__(queue, exc_queue, subtasks)
+        # stdout queue is not necessary for thread tasks.
+        super(ThreadTask, self).__init__(queue, exc_queue, None, subtasks)
 
         self.thread = thread
 

scheduler/utils.py

@@ -19,11 +19,40 @@
 #  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 #  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 #  SOFTWARE.
+import sys
+import time
 from multiprocessing import Process
 
 import psutil
 
 
+class StdOut:
+    def __init__(self, queue):
+        self.period = 1
+
+        self.last_update = time.time()
+        self.text = []
+
+        self.queue = queue
+
+    def write(self, text: str) -> None:
+        self.text.append(text)
+        self.update()
+
+    def update(self, force: bool = False) -> None:
+        if self.text and (force or time.time() - self.last_update > self.period):
+            out = "".join(self.text)
+
+            if out.strip():
+                self.queue.put(out)
+
+                self.last_update = time.time()
+                self.text = []
+
+    def flush(self) -> None:
+        return
+
+
 def terminate_tree(process: Process):
     """
     Terminates a process along with all of its child processes.

test/impl.py

@@ -21,13 +21,15 @@
 #  SOFTWARE.
 
 import asyncio
+import sys
+import time
 from multiprocessing import Process, Queue
 from typing import List, Tuple
 
 import multiprocess
 
 from scheduler.ProcessTask import ProcessTask
-from scheduler.utils import SchedulerException, TaskFailedException
+from scheduler.utils import TaskFailedException
 from test.utils import (
     _get_input_output_numpy,
     _func_numpy,
@@ -44,6 +46,7 @@
     _func_no_return,
     _long_task,
     _func_raise_exception,
+    _func_print,
 )
 
 
@@ -342,3 +345,27 @@ def test_raise_exception(scheduler):
         assert isinstance(e, TaskFailedException)
 
     assert scheduler.failed
+
+
+text = ""
+
+
+def test_stdout(scheduler):
+    global text
+
+    expected = "\n".join([f"{i}" for i in range(1000)]) + "\n"
+
+    def temp(_stdout):
+        global text
+        text += _stdout
+
+    write = sys.stdout.write
+    sys.stdout.write = temp
+
+    scheduler.map_blocking(target=_func_print, args=[(expected,),])
+
+    while not text:
+        time.sleep(0.01)
+
+    sys.stdout.write = write
+    assert text == expected, "Text from stdout is incorrect."

test/test_stdout.py

@@ -0,0 +1,35 @@
+#  MIT License
+#
+#  Copyright (c) 2020 Sam McCormack
+#
+#  Permission is hereby granted, free of charge, to any person obtaining a copy
+#  of this software and associated documentation files (the "Software"), to deal
+#  in the Software without restriction, including without limitation the rights
+#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+#  copies of the Software, and to permit persons to whom the Software is
+#  furnished to do so, subject to the following conditions:
+#
+#  The above copyright notice and this permission notice shall be included in all
+#  copies or substantial portions of the Software.
+#
+#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+#  SOFTWARE.
+from scheduler.Scheduler import Scheduler
+from test import impl
+
+
+def scheduler():
+    return Scheduler(capture_stdout=True)
+
+
+def test_stdout():
+    impl.test_stdout(scheduler())
+
+
+if __name__ == "__main__":
+    test_stdout()

test/utils.py

@@ -118,6 +118,12 @@ def _func_raise_exception(x, y, z) -> None:
     raise TestException("Test exception.")
 
 
+def _func_print(text: str) -> None:
+    for i in range(1000):
+        print(i)
+        time.sleep(0.001)
+
+
 class TestException(Exception):
     """
     Exception raised for testing.