Relax protobuf<5.0.0 restriction

Instead of restricting `protobuf<5.0.0` due to the removal to
`including_default_value_fields` kwarg in `json_format.MessageToJson`
(replaced by `always_print_fields_with_nopresence`), we can instead
detect the currently installed protobuf version and use the appropriate
kwarg.

Contributes to tensorflow#6808 and should resolve tensorflow#6887.

tensorboard/pip_package/requirements.txt

@@ -20,14 +20,18 @@ absl-py >= 0.4
 grpcio >= 1.48.2
 markdown >= 2.6.8
 numpy >= 1.12.0
+# NOTE: The packaging dependency was introduced in order to compare installed
+# package versions and conditionally use different supported kwargs
+# (specifically the protobuf dependency). If we restrict protobuf >= 5.0.0 we
+# can get rid of the packaging dependency.
+packaging
 # NOTE: this version must be >= the protoc version in our WORKSPACE file.
 # At the same time, any constraints we specify here must allow at least some
 # version to be installed that is also compatible with TensorFlow's constraints:
 # https://github.com/tensorflow/tensorflow/blob/25adc4fccb4b0bb5a933eba1d246380e7b87d7f7/tensorflow/tools/pip_package/setup.py#L101
 # 4.24.0 had an issue that broke our tests, so we should avoid that release:
 # https://github.com/protocolbuffers/protobuf/issues/13485
-# 5.26.0 introduced a breaking change, so we restricted it for now. See issue #6808 for details.
-protobuf >= 3.19.6, != 4.24.0, < 5.0.0
+protobuf >= 3.19.6, != 4.24.0
 setuptools >= 41.0.0  # Note: provides pkg_resources as well as setuptools
 # A dependency of our vendored packages. This lower bound has not been correctly
 # vetted, but I wanted to be the least restrictive we can, since it's not a new

tensorboard/plugin_util.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Provides utilities that may be especially useful to plugins."""
 
-
+from importlib import metadata
+from packaging import version
 import threading
 
 from bleach.sanitizer import Cleaner
@@ -193,6 +194,17 @@ def experiment_id(environ):
     return environ.get(_experiment_id.WSGI_ENVIRON_KEY, "")
 
 
+def package_version_greater_than_or_equal(package_name, version_string):
+  """Compare the package name is greater than or equal to the provided version.
+
+  Args:
+    package_name: The package name to be checked.
+    version: The version string to compare the package against.
+  """
+  current_version = metadata.version(package_name)
+  return version.parse(current_version) >= version.parse(version_string)
+
+
 class _MetadataVersionChecker:
     """TensorBoard-internal utility for warning when data is too new.
 

tensorboard/plugin_util_test.py

@@ -14,6 +14,7 @@
 
 
 import textwrap
+from unittest import mock
 
 
 from tensorboard import context
@@ -229,5 +230,24 @@ def test_present(self):
         self.assertEqual(plugin_util.experiment_id(environ), "123")
 
 
+class PackageVersionTest(tb_test.TestCase):
+  """Tests for `plugin_util.package_version_greater_than_or_equal`."""
+
+  def test_package_version_greater_than(self):
+    with mock.patch("importlib.metadata.version") as mock_version:
+      mock_version.return_value = "5.0.1"
+      self.assertTrue(plugin_util.package_version_greater_than_or_equal("test", "5.0.0"))
+
+  def test_package_version_equal_to(self):
+    with mock.patch("importlib.metadata.version") as mock_version:
+      mock_version.return_value = "5.0.0"
+      self.assertTrue(plugin_util.package_version_greater_than_or_equal("test", "5.0.0"))
+
+  def test_package_version_less_than(self):
+    with mock.patch("importlib.metadata.version") as mock_version:
+      mock_version.return_value = "4.9.9"
+      self.assertFalse(plugin_util.package_version_greater_than_or_equal("test", "5.0.0"))
+
+
 if __name__ == "__main__":
     tb_test.main()

tensorboard/plugins/custom_scalar/custom_scalars_plugin.py

@@ -21,7 +21,6 @@
 this plugin.
 """
 
-
 import re
 
 from google.protobuf import json_format
@@ -318,9 +317,14 @@ def layout_impl(self, ctx, experiment):
                     title_to_category[category.title] = category
 
         if merged_layout:
-            return json_format.MessageToJson(
-                merged_layout, including_default_value_fields=True
-            )
+            if plugin_util.package_version_greater_than_or_equal("protobuf", "5.0"):
+                response = json_format.MessageToJson(
+                    merged_layout, always_print_fields_with_no_presence=True,
+                )
+            else:
+                response = json_format.MessageToJson(
+                    merged_layout, including_default_value_fields=True,
+                )
         else:
             # No layout was found.
             return {}

tensorboard/plugins/hparams/hparams_plugin.py

@@ -18,10 +18,8 @@
 this plugin.
 """
 
-
 import json
 
-
 import werkzeug
 from werkzeug import wrappers
 
@@ -116,14 +114,23 @@ def get_experiment_route(self, request):
             request_proto = _parse_request_argument(
                 request, api_pb2.GetExperimentRequest
             )
+            response_proto = get_experiment.Handler(
+                ctx,
+                self._context,
+                experiment_id,
+                request_proto,
+            ).run()
+            if plugin_util.package_version_greater_than_or_equal("protobuf", "5.0.0"):
+                response = json_format.MessageToJson(
+                    response_proto, always_print_fields_with_no_presence=True,
+                )
+            else:
+                response = json_format.MessageToJson(
+                    response_proto, including_default_value_fields=True,
+                )
             return http_util.Respond(
                 request,
-                json_format.MessageToJson(
-                    get_experiment.Handler(
-                        ctx, self._context, experiment_id, request_proto
-                    ).run(),
-                    including_default_value_fields=True,
-                ),
+                response,
                 "application/json",
             )
         except error.HParamsError as e:
@@ -139,14 +146,23 @@ def list_session_groups_route(self, request):
             request_proto = _parse_request_argument(
                 request, api_pb2.ListSessionGroupsRequest
             )
+            response_proto = list_session_groups.Handler(
+                ctx,
+                self._context,
+                experiment_id,
+                request_proto,
+            ).run()
+            if plugin_util.package_version_greater_than_or_equal("protobuf", "5.0.0"):
+                response = json_format.MessageToJson(
+                    response_proto, always_print_fields_with_no_presence=True,
+                )
+            else:
+                response = json_format.MessageToJson(
+                    response_proto, including_default_value_fields=True,
+                )
             return http_util.Respond(
                 request,
-                json_format.MessageToJson(
-                    list_session_groups.Handler(
-                        ctx, self._context, experiment_id, request_proto
-                    ).run(),
-                    including_default_value_fields=True,
-                ),
+                response,
                 "application/json",
             )
         except error.HParamsError as e: