fix(wit): Redirect WIT plugin to LIT (#6343)

groszewn · web-flow · commit 33abcb54d79e · 2023-04-24T15:31:11.000-04:00
The WIT is no longer actively maintained, so redirect users to install and use the newer [Learning Interoperability Tool (LIT)](https://pair-code.github.io/lit/) instead via a dynamic plugin. I attempted to get the WIT package building again to push out an updated release that does this redirection outside of the TB code, but that proved more difficult than expected. The resulting dashboard can be seen below, which links to the [LIT homepage](https://pair-code.github.io/lit/) and [setup guide](https://pair-code.github.io/lit/setup/), respectively. ![image](https://user-images.githubusercontent.com/13921855/234000839-7affc80a-dc89-478f-8a9e-98e5522cc456.png) Contributes to #6195
diff --git a/docs/what_if_tool.md b/docs/what_if_tool.md
@@ -1,5 +1,11 @@
 # Model Understanding with the What-If Tool Dashboard
 
+> **Warning**
+> This documentation only applies to TensorBoard 2.11 and earlier, as the
+> What-If Tool is no longer actively maintained. Please check out the actively
+> maintained [Learning Interpretability Tool
+> (LIT)](https://pair-code.github.io/lit/) instead.
+
 ![What-If Tool](./images/what_if_tool.png)
 
 The What-If Tool (WIT) provides an easy-to-use interface for expanding
diff --git a/tensorboard/BUILD b/tensorboard/BUILD
@@ -320,6 +320,7 @@ py_library(
         "//tensorboard/plugins/profile_redirect:profile_redirect_plugin",
         "//tensorboard/plugins/scalar:scalars_plugin",
         "//tensorboard/plugins/text:text_plugin",
+        "//tensorboard/plugins/wit_redirect:wit_redirect_plugin",
     ],
 )
 
diff --git a/tensorboard/components/BUILD b/tensorboard/components/BUILD
@@ -21,6 +21,7 @@ tf_ts_library(
         "//tensorboard/plugins/profile_redirect/tf_profile_redirect_dashboard",
         "//tensorboard/plugins/scalar/tf_scalar_dashboard",
         "//tensorboard/plugins/text/tf_text_dashboard",
+        "//tensorboard/plugins/wit_redirect/tf_wit_redirect_dashboard",
     ],
 )
 
diff --git a/tensorboard/components/polymer3_lib.ts b/tensorboard/components/polymer3_lib.ts
@@ -26,4 +26,5 @@ import '../plugins/profile_redirect/tf_profile_redirect_dashboard/tf-profile-red
 import '../plugins/pr_curve/tf_pr_curve_dashboard/tf-pr-curve-dashboard';
 import '../plugins/scalar/tf_scalar_dashboard/tf-scalar-dashboard';
 import '../plugins/text/tf_text_dashboard/tf-text-dashboard';
+import '../plugins/wit_redirect/tf_wit_redirect_dashboard/tf-wit-redirect-dashboard';
 import './polymer3_interop_helper';
diff --git a/tensorboard/default.py b/tensorboard/default.py
@@ -44,6 +44,7 @@
 from tensorboard.plugins.scalar import scalars_plugin
 from tensorboard.plugins.text import text_plugin
 from tensorboard.plugins.mesh import mesh_plugin
+from tensorboard.plugins.wit_redirect import wit_redirect_plugin
 
 
 logger = logging.getLogger(__name__)
@@ -67,6 +68,7 @@
     profile_redirect_plugin.ProfileRedirectPluginLoader,
     hparams_plugin.HParamsPlugin,
     mesh_plugin.MeshPlugin,
+    wit_redirect_plugin.WITRedirectPluginLoader,
 ]
 
 
diff --git a/tensorboard/pip_package/requirements.txt b/tensorboard/pip_package/requirements.txt
@@ -30,6 +30,5 @@ protobuf >= 3.19.6
 requests >= 2.21.0, < 3
 setuptools >= 41.0.0  # Note: provides pkg_resources as well as setuptools
 tensorboard-data-server >= 0.7.0, < 0.8.0
-tensorboard-plugin-wit >= 1.6.0
 werkzeug >= 1.0.1
 wheel >= 0.26
diff --git a/tensorboard/plugins/interactive_inference/README.md b/tensorboard/plugins/interactive_inference/README.md
@@ -1,5 +1,10 @@
 # What-If Tool
 
+> **Warning**
+> The What-If Tool is no longer actively maintained. Please use the actively
+> maintained [Learning Interpretability Tool (LIT)](https://pair-code.github.io/lit/)
+> instead.
+
 The What-If Tool code and documentation has moved to https://github.com/pair-code/what-if-tool.
 
 The What-If Tool TensorBoard plugin has been converted to a dynamic plugin, through the tensorboard-plugin-wit pip package.
diff --git a/tensorboard/plugins/wit_redirect/BUILD b/tensorboard/plugins/wit_redirect/BUILD
@@ -0,0 +1,26 @@
+# Description:
+# Plugin with installation instructions for dynamic interpretability plugin
+
+package(default_visibility = ["//tensorboard:internal"])
+
+licenses(["notice"])
+
+py_library(
+    name = "wit_redirect_plugin",
+    srcs = ["wit_redirect_plugin.py"],
+    srcs_version = "PY3",
+    deps = [
+        "//tensorboard/plugins:base_plugin",
+    ],
+)
+
+py_test(
+    name = "wit_redirect_plugin_test",
+    srcs = ["wit_redirect_plugin_test.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":wit_redirect_plugin",
+        "//tensorboard:test",
+        "//tensorboard/plugins:base_plugin",
+    ],
+)
diff --git a/tensorboard/plugins/wit_redirect/__init__.py b/tensorboard/plugins/wit_redirect/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
diff --git a/tensorboard/plugins/wit_redirect/tf_wit_redirect_dashboard/BUILD b/tensorboard/plugins/wit_redirect/tf_wit_redirect_dashboard/BUILD
@@ -0,0 +1,17 @@
+load("//tensorboard/defs:defs.bzl", "tf_ts_library")
+
+package(default_visibility = ["//tensorboard:internal"])
+
+licenses(["notice"])
+
+tf_ts_library(
+    name = "tf_wit_redirect_dashboard",
+    srcs = ["tf-wit-redirect-dashboard.ts"],
+    strict_checks = False,
+    deps = [
+        "//tensorboard/components/polymer:irons_and_papers",
+        "//tensorboard/components/polymer:legacy_element_mixin",
+        "@npm//@polymer/decorators",
+        "@npm//@polymer/polymer",
+    ],
+)
diff --git a/tensorboard/plugins/wit_redirect/tf_wit_redirect_dashboard/tf-wit-redirect-dashboard.ts b/tensorboard/plugins/wit_redirect/tf_wit_redirect_dashboard/tf-wit-redirect-dashboard.ts
@@ -0,0 +1,67 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+import {customElement, property} from '@polymer/decorators';
+import {html, PolymerElement} from '@polymer/polymer';
+import '../../../components/polymer/irons_and_papers';
+import {LegacyElementMixin} from '../../../components/polymer/legacy_element_mixin';
+
+/**
+ * A frontend that directs users to install the Learning Interoperability Plugin
+ * (LIT) instead of the What-If Tools, since the latter is no longer maintained.
+ */
+@customElement('tf-wit-redirect-dashboard')
+class TfWITRedirectDashboard extends LegacyElementMixin(PolymerElement) {
+  static readonly template = html`
+    <div class="message">
+      <h3>The What-If Tool is no longer supported.</h3>
+      <p>
+        The
+        <a href="https://pair-code.github.io/lit/"
+          >Learning Interpretability Tool (LIT)</a
+        >
+        is an actively maintained alternative. Please follow the instructions
+        <a href="https://pair-code.github.io/lit/setup/">here</a> to install and
+        use this tool.
+      </p>
+      <style>
+        :host {
+          display: flex;
+        }
+
+        .message {
+          margin: 80px auto 0 auto;
+          max-width: 540px;
+        }
+        #commandTextarea {
+          margin-top: 1ex;
+          padding: 1ex 1em;
+          resize: vertical;
+          width: 100%;
+        }
+        #copyContainer {
+          display: flex;
+        }
+        #copiedMessage {
+          align-self: center;
+          flex-grow: 1;
+          font-style: italic;
+          padding-right: 1em;
+          text-align: right;
+        }
+      </style>
+    </div>
+  `;
+}
diff --git a/tensorboard/plugins/wit_redirect/wit_redirect_plugin.py b/tensorboard/plugins/wit_redirect/wit_redirect_plugin.py
@@ -0,0 +1,50 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Plugin that only displays a message with installation instructions."""
+
+
+from tensorboard.plugins import base_plugin
+
+
+class WITRedirectPluginLoader(base_plugin.TBLoader):
+    """Load the redirect notice iff the dynamic plugin is unavailable."""
+
+    def load(self, context):
+        try:
+            import tensorboard_plugin_wit  # noqa: F401
+
+            # If we successfully load the dynamic plugin, don't show
+            # this redirect plugin at all.
+            return None
+        except ImportError:
+            return _WITRedirectPlugin(context)
+
+
+class _WITRedirectPlugin(base_plugin.TBPlugin):
+    """Redirect notice pointing users to the new dynamic LIT plugin."""
+
+    plugin_name = "wit_redirect"
+
+    def get_plugin_apps(self):
+        return {}
+
+    def is_active(self):
+        return False
+
+    def frontend_metadata(self):
+        return base_plugin.FrontendMetadata(
+            element_name="tf-wit-redirect-dashboard",
+            tab_name="What-If Tool",
+        )
diff --git a/tensorboard/plugins/wit_redirect/wit_redirect_plugin_test.py b/tensorboard/plugins/wit_redirect/wit_redirect_plugin_test.py
@@ -0,0 +1,91 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `wit_redirect_plugin`."""
+
+
+import contextlib
+import sys
+from unittest import mock
+
+from tensorboard.plugins import base_plugin
+from tensorboard.plugins.wit_redirect import wit_redirect_plugin
+from tensorboard import test as tb_test
+
+
+_DYNAMIC_PLUGIN_MODULE = "tensorboard_plugin_wit"
+
+
+class WITRedirectPluginLoaderTest(tb_test.TestCase):
+    """Tests for `WITRedirectPluginLoader`."""
+
+    def test_loads_when_no_dynamic_plugin(self):
+        with contextlib.ExitStack() as stack:
+            stack.enter_context(mock.patch.dict(sys.modules))
+            sys.modules.pop(_DYNAMIC_PLUGIN_MODULE, None)
+
+            real_import = __import__
+
+            def fake_import(name, *args, **kwargs):
+                if name == _DYNAMIC_PLUGIN_MODULE:
+                    raise ImportError("Pretend I'm not here")
+                else:
+                    return real_import(name, *args, **kwargs)
+
+            stack.enter_context(mock.patch("builtins.__import__", fake_import))
+
+            plugin_class = wit_redirect_plugin._WITRedirectPlugin
+            plugin_init = stack.enter_context(
+                mock.patch.object(plugin_class, "__init__", return_value=None)
+            )
+
+            loader = wit_redirect_plugin.WITRedirectPluginLoader()
+            context = base_plugin.TBContext()
+            result = loader.load(context)
+            self.assertIsInstance(result, plugin_class)
+            plugin_init.assert_called_once_with(context)
+
+    def test_does_not_load_when_dynamic_plugin_present(self):
+        with contextlib.ExitStack() as stack:
+            stack.enter_context(mock.patch.dict(sys.modules))
+            sys.modules.pop(_DYNAMIC_PLUGIN_MODULE, None)
+
+            real_import = __import__
+
+            def fake_import(name, *args, **kwargs):
+                if name == _DYNAMIC_PLUGIN_MODULE:
+                    arbitrary_module = sys
+                    sys.modules.setdefault(
+                        _DYNAMIC_PLUGIN_MODULE, arbitrary_module
+                    )
+                    return arbitrary_module
+                else:
+                    return real_import(name, *args, **kwargs)
+
+            stack.enter_context(mock.patch("builtins.__import__", fake_import))
+
+            plugin_class = wit_redirect_plugin._WITRedirectPlugin
+            plugin_init = stack.enter_context(
+                mock.patch.object(plugin_class, "__init__", return_value=None)
+            )
+
+            loader = wit_redirect_plugin.WITRedirectPluginLoader()
+            context = base_plugin.TBContext()
+            result = loader.load(context)
+            self.assertIsNone(result)
+            plugin_init.assert_not_called()
+
+
+if __name__ == "__main__":
+    tb_test.main()
diff --git a/tensorboard/tools/diagnose_tensorboard.py b/tensorboard/tools/diagnose_tensorboard.py
@@ -254,23 +254,6 @@ def installed_packages():
         )
         yield Suggestion("Fix conflicting installations", message)
 
-    wit_version = packages.get("tensorboard-plugin-wit")
-    if wit_version == "tensorboard-plugin-wit==1.6.0.post2":
-        # This is only incompatible with TensorBoard prior to 2.2.0, but
-        # we just issue a blanket warning so that we don't have to pull
-        # in a `pkg_resources` dep to parse the version.
-        preamble = reflow(
-            """
-            Versions of the What-If Tool (`tensorboard-plugin-wit`)
-            prior to 1.6.0.post3 are incompatible with some versions of
-            TensorBoard. Please upgrade this package to the latest
-            version to resolve any startup errors:
-            """
-        )
-        command = "pip install -U tensorboard-plugin-wit"
-        message = "%s\n\n\t%s" % (preamble, command)
-        yield Suggestion("Upgrade `tensorboard-plugin-wit`", message)
-
 
 @check
 def tensorboard_python_version():

Original file line number	Diff line number	Diff line change
`@@ -320,6 +320,7 @@ py_library(`
`320`	`320`	`"//tensorboard/plugins/profile_redirect:profile_redirect_plugin",`
`321`	`321`	`"//tensorboard/plugins/scalar:scalars_plugin",`
`322`	`322`	`"//tensorboard/plugins/text:text_plugin",`
	`323`	`+ "//tensorboard/plugins/wit_redirect:wit_redirect_plugin",`
`323`	`324`	`],`
`324`	`325`	`)`
`325`	`326`
Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@ tf_ts_library(`
`21`	`21`	`"//tensorboard/plugins/profile_redirect/tf_profile_redirect_dashboard",`
`22`	`22`	`"//tensorboard/plugins/scalar/tf_scalar_dashboard",`
`23`	`23`	`"//tensorboard/plugins/text/tf_text_dashboard",`
	`24`	`+ "//tensorboard/plugins/wit_redirect/tf_wit_redirect_dashboard",`
`24`	`25`	`],`
`25`	`26`	`)`
`26`	`27`