From 7a07109d5c372ed57699434b1dc8f8b59a652bd2 Mon Sep 17 00:00:00 2001
From: Peyton Murray <peynmurray@gmail.com>
Date: Thu, 2 Mar 2023 14:16:08 -0800
Subject: [PATCH] Guard against using ipywidgets in google colab (#32841)

This PR guards against using ipywidgets in google colab due to an incompatibility. It does this by adding a decorator to all `_ipython_display_` definitions which simply prints the object `repr` instead of displaying a widget.
---
 python/ray/data/dataset.py                |  3 ++-
 python/ray/train/data_parallel_trainer.py |  3 ++-
 python/ray/widgets/util.py                | 19 +++++++++++++++++++
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py
index 83271d0b278a7..4ac3085b5d69c 100644
--- a/python/ray/data/dataset.py
+++ b/python/ray/data/dataset.py
@@ -128,7 +128,7 @@
 from ray.util.annotations import DeveloperAPI, PublicAPI
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 from ray.widgets import Template
-from ray.widgets.util import ensure_notebook_deps
+from ray.widgets.util import ensure_notebook_deps, fallback_if_colab
 
 if sys.version_info >= (3, 8):
     from typing import Literal
@@ -4321,6 +4321,7 @@ def _aggregate_result(self, result: Union[Tuple, TableRow]) -> U:
     @ensure_notebook_deps(
         ["ipywidgets", "8"],
     )
+    @fallback_if_colab
     def _ipython_display_(self):
         from ipywidgets import HTML, VBox, Layout
         from IPython.display import display
diff --git a/python/ray/train/data_parallel_trainer.py b/python/ray/train/data_parallel_trainer.py
index 4169f52aacf1d..4f0b95f17ae58 100644
--- a/python/ray/train/data_parallel_trainer.py
+++ b/python/ray/train/data_parallel_trainer.py
@@ -21,7 +21,7 @@
 from ray.train.trainer import BaseTrainer, GenDataset
 from ray.util.annotations import DeveloperAPI
 from ray.widgets import Template
-from ray.widgets.util import ensure_notebook_deps
+from ray.widgets.util import ensure_notebook_deps, fallback_if_colab
 
 if TYPE_CHECKING:
     from ray.data.preprocessor import Preprocessor
@@ -447,6 +447,7 @@ def get_dataset_config(self) -> Dict[str, DatasetConfig]:
         ["tabulate", None],
         ["ipywidgets", "8"],
     )
+    @fallback_if_colab
     def _ipython_display_(self):
         from ipywidgets import HTML, VBox, Tab, Layout
         from IPython.display import display
diff --git a/python/ray/widgets/util.py b/python/ray/widgets/util.py
index 212fb1c39c969..0c6e406a189e1 100644
--- a/python/ray/widgets/util.py
+++ b/python/ray/widgets/util.py
@@ -167,3 +167,22 @@ def _has_outdated(
         logger.warning(f"Outdated packages:\n{outdated_str}\n{message}", stacklevel=3)
 
     return outdated
+
+
+@DeveloperAPI
+def fallback_if_colab(func: F) -> Callable[[F], F]:
+    try:
+        ipython = get_ipython()
+    except NameError:
+        ipython = None
+
+    @wraps(func)
+    def wrapped(self, *args, **kwargs):
+        if ipython and "google.colab" not in str(ipython):
+            return func(self, *args, **kwargs)
+        elif hasattr(self, "__repr__"):
+            return print(self.__repr__(*args, **kwargs))
+        else:
+            return None
+
+    return wrapped