ntaive redirect for notebooks (#191)

Co-authored-by: Jeff Raubitschek <jeff@wandb.com>

wandb/lib/redirect.py

@@ -30,6 +30,33 @@ def __getattr__(self, attr):
         return getattr(self.stream, attr)
 
 
+class StreamFork(object):
+
+    def __init__(self, output_streams, unbuffered=False):
+        self.output_streams = output_streams
+        self.unbuffered = unbuffered
+
+    def write(self, data):
+        output_streams = object.__getattribute__(self, 'output_streams')
+        unbuffered = object.__getattribute__(self, 'unbuffered')
+        for stream in output_streams:
+            stream.write(data)
+            if unbuffered:
+                stream.flush()
+
+    def writelines(self, datas):
+        output_streams = object.__getattribute__(self, 'output_streams')
+        unbuffered = object.__getattribute__(self, 'unbuffered')
+        for stream in output_streams:
+            stream.writelines(datas)
+            if unbuffered:
+                stream.flush()
+
+    def __getattr__(self, attr):
+        output_streams = object.__getattribute__(self, 'output_streams')
+        return getattr(output_streams[0], attr)
+
+
 class StreamWrapper(object):
     def __init__(self, name, cb, output_writer=None):
         self.name = name
@@ -119,9 +146,17 @@ def _redirect(self, to_fd, unbuffered=False, close=False):
             # Do not close old filedescriptor as others might be using it
             fp.close()
         os.dup2(to_fd, self._old_fd)
-        setattr(sys, self._stream, os.fdopen(self._old_fd, "w"))
-        if unbuffered:
-            setattr(sys, self._stream, Unbuffered(getattr(sys, self._stream)))
+        if getattr(sys, self._stream) == getattr(sys, "__%s__" % self._stream):
+            setattr(sys, self._stream, os.fdopen(self._old_fd, "w"))
+            if unbuffered:
+                setattr(sys, self._stream, Unbuffered(getattr(sys, self._stream)))
+        else:
+            if close:
+                setattr(sys, self._stream, getattr(sys, self._stream).output_streams[0])
+            else:
+                setattr(sys, self._stream, StreamFork([getattr(sys, self._stream),
+                                                      os.fdopen(self._old_fd, "w")],
+                                                      unbuffered=unbuffered))
 
     def install(self):
         if self._installed:
@@ -137,7 +172,7 @@ def install(self):
 
         logger.info("install start")
 
-        fp = getattr(sys, self._stream)
+        fp = getattr(sys, "__%s__" % self._stream)
         fd = fp.fileno()
         old_fp = os.fdopen(os.dup(fd), "w")
 

wandb/sdk/wandb_run.py

@@ -990,8 +990,8 @@ def _redirect(self, stdout_slave_fd, stderr_slave_fd):
             err_redir = redirect.Redirect(
                 src="stderr", dest=err_cap, unbuffered=True, tee=True
             )
-        elif console == self._settings.Console.NOTEBOOK:
-            logger.info("Redirecting notebook output.")
+        elif console == self._settings.Console.WRAP:
+            logger.info("Wrapping output streams.")
             out_redir = redirect.StreamWrapper(
                 name="stdout", cb=self._redirect_cb, output_writer=self._output_writer
             )

wandb/sdk/wandb_settings.py

@@ -172,7 +172,7 @@ def get_wandb_dir(root_dir: str):
 @enum.unique
 class SettingsConsole(enum.Enum):
     OFF = 0
-    NOTEBOOK = 1
+    WRAP = 1
     REDIRECT = 2
 
 
@@ -374,14 +374,12 @@ def _windows(self) -> bool:
     def _console(self) -> SettingsConsole:
         convert_dict: Dict[str, SettingsConsole] = dict(
             off=SettingsConsole.OFF,
-            notebook=SettingsConsole.NOTEBOOK,
+            wrap=SettingsConsole.WRAP,
             redirect=SettingsConsole.REDIRECT,
         )
         console: str = self.console
         if console == "auto":
-            if self._jupyter:
-                console = "notebook"
-            elif self._windows:
+            if self._windows:
                 legacy_env_var = "PYTHONLEGACYWINDOWSSTDIO"
                 if sys.version_info >= (3, 6) and legacy_env_var not in os.environ:
                     msg = (
@@ -390,7 +388,7 @@ def _console(self) -> SettingsConsole:
                         "back to monkey patching stdout/err." % legacy_env_var
                     )
                     logger.info(msg)
-                    console = "notebook"
+                    console = "wrap"
                 else:
                     console = "redirect"
             else:
@@ -464,7 +462,7 @@ def _validate_console(self, value):
             "auto",
             "redirect",
             "off",
-            "notebook",
+            "wrap",
         }
         if value in choices:
             return

wandb/sdk_py27/wandb_run.py

@@ -990,8 +990,8 @@ def _redirect(self, stdout_slave_fd, stderr_slave_fd):
             err_redir = redirect.Redirect(
                 src="stderr", dest=err_cap, unbuffered=True, tee=True
             )
-        elif console == self._settings.Console.NOTEBOOK:
-            logger.info("Redirecting notebook output.")
+        elif console == self._settings.Console.WRAP:
+            logger.info("Wrapping output streams.")
             out_redir = redirect.StreamWrapper(
                 name="stdout", cb=self._redirect_cb, output_writer=self._output_writer
             )

wandb/sdk_py27/wandb_settings.py

@@ -172,7 +172,7 @@ def get_wandb_dir(root_dir):
 @enum.unique
 class SettingsConsole(enum.Enum):
     OFF = 0
-    NOTEBOOK = 1
+    WRAP = 1
     REDIRECT = 2
 
 
@@ -374,14 +374,12 @@ def _windows(self):
     def _console(self):
         convert_dict = dict(
             off=SettingsConsole.OFF,
-            notebook=SettingsConsole.NOTEBOOK,
+            wrap=SettingsConsole.WRAP,
             redirect=SettingsConsole.REDIRECT,
         )
         console = self.console
         if console == "auto":
-            if self._jupyter:
-                console = "notebook"
-            elif self._windows:
+            if self._windows:
                 legacy_env_var = "PYTHONLEGACYWINDOWSSTDIO"
                 if sys.version_info >= (3, 6) and legacy_env_var not in os.environ:
                     msg = (
@@ -390,7 +388,7 @@ def _console(self):
                         "back to monkey patching stdout/err." % legacy_env_var
                     )
                     logger.info(msg)
-                    console = "notebook"
+                    console = "wrap"
                 else:
                     console = "redirect"
             else:
@@ -464,7 +462,7 @@ def _validate_console(self, value):
             "auto",
             "redirect",
             "off",
-            "notebook",
+            "wrap",
         }
         if value in choices:
             return