Merge branch 'master' into nbrunner-cleanup

Netflix · May 27, 2024 · 34db634 · 34db634
2 parents f20381a + 56d984c
commit 34db634
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 19 deletions.
diff --git a/metaflow/runner/metaflow_runner.py b/metaflow/runner/metaflow_runner.py
@@ -188,7 +188,7 @@ class Runner(object):
     def __init__(
         self,
         flow_file: str,
-        show_output: bool = False,
+        show_output: bool = True,
         profile: Optional[str] = None,
         env: Optional[Dict] = None,
         cwd: Optional[str] = None,
@@ -214,8 +214,8 @@ def __init__(
         ----------
         flow_file : str
             Path to the flow file to run
-        show_output : bool, default False
-            Suppress the 'stdout' and 'stderr' to the console by default,
+        show_output : bool, default True
+            Show the 'stdout' and 'stderr' to the console by default,
             Only applicable for synchronous 'run' and 'resume' functions.
         profile : Optional[str], default None
             Metaflow profile to use to run this run. If not specified, the default

diff --git a/metaflow/runner/nbrun.py b/metaflow/runner/nbrun.py
@@ -1,22 +1,20 @@
 import ast
 import os
-import shutil
 import tempfile
 from typing import Dict, Optional
 
 from metaflow import Runner
 
-try:
-    from IPython import get_ipython
+DEFAULT_DIR = tempfile.gettempdir()
 
-    ipython = get_ipython()
-except ModuleNotFoundError:
-    print("'nbrun' requires an interactive python environment (such as Jupyter)")
 
-DEFAULT_DIR = tempfile.gettempdir()
+class NBRunnerInitializationError(Exception):
+    """Custom exception for errors during NBRunner initialization."""
+
+    pass
 
 
-def get_current_cell():
+def get_current_cell(ipython):
     if ipython:
         return ipython.history_manager.input_hist_raw[-1]
     return None
@@ -50,13 +48,22 @@ class NBRunner(object):
     def __init__(
         self,
         flow,
-        show_output: bool = False,
+        show_output: bool = True,
         profile: Optional[str] = None,
         env: Optional[Dict] = None,
         base_dir: str = DEFAULT_DIR,
         **kwargs,
     ):
-        self.cell = get_current_cell()
+        try:
+            from IPython import get_ipython
+
+            ipython = get_ipython()
+        except ModuleNotFoundError:
+            raise NBRunnerInitializationError(
+                "'NBRunner' requires an interactive Python environment (such as Jupyter)"
+            )
+
+        self.cell = get_current_cell(ipython)
         self.flow = flow
         self.show_output = show_output
 
@@ -93,18 +100,12 @@ def __init__(
         )
 
     def nbrun(self, **kwargs):
-        self.old_val_show_output = self.show_output
-        self.runner.show_output = True
         result = self.runner.run(**kwargs)
-        self.runner.show_output = self.old_val_show_output
         self.cleanup()
         return result.run
 
     def nbresume(self, **kwargs):
-        self.old_val_show_output = self.show_output
-        self.runner.show_output = True
         result = self.runner.resume(**kwargs)
-        self.runner.show_output = self.old_val_show_output
         self.cleanup()
         return result.run