fix(wit): Redirect WIT plugin to LIT

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

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",
     ],
 )
 

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",
     ],
 )
 

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';

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,
 ]
 
 

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

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.

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",
+    ],
+)

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.
+# ==============================================================================

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",
+    ],
+)

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>
+  `;
+}

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",
+        )

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()

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():