From 3d2aff962a8db5466af7215ac029dbac30ef65af Mon Sep 17 00:00:00 2001
From: Yating Jing <ytjing@google.com>
Date: Fri, 7 Jan 2022 20:59:06 +0000
Subject: [PATCH 1/9] conditionally import tensorflow_io

---
 .../backend/event_processing/data_ingester.py | 34 +++++++++++++++++++
 .../event_processing/data_ingester_test.py    | 20 +++++++++++
 2 files changed, 54 insertions(+)

diff --git a/tensorboard/backend/event_processing/data_ingester.py b/tensorboard/backend/event_processing/data_ingester.py
index 3c74c770e48..803372440ea 100644
--- a/tensorboard/backend/event_processing/data_ingester.py
+++ b/tensorboard/backend/event_processing/data_ingester.py
@@ -30,6 +30,7 @@
 from tensorboard.plugins.pr_curve import metadata as pr_curve_metadata
 from tensorboard.plugins.scalar import metadata as scalar_metadata
 from tensorboard.util import tb_logging
+from tensorboard.compat import tf
 
 
 DEFAULT_SIZE_GUIDANCE = {
@@ -45,6 +46,10 @@
     pr_curve_metadata.PLUGIN_NAME: 100,
 }
 
+# TensorFlow I/O file systems of interest (not available in TensorFlow's
+# built-in support).
+_CLOUD_FILESYSTEMS = ["s3"]
+
 logger = tb_logging.get_logger()
 
 
@@ -79,6 +84,19 @@ def __init__(self, flags):
         else:
             self._path_to_run = _parse_event_files_spec(flags.logdir_spec)
 
+        # Conditionally import tensorflow_io.
+        if not getattr(tf, "__version__", "stub") == "stub":
+            if not _cloud_fs_supported():
+                try:
+                    import tensorflow_io
+                except:
+                    warning_msg = (
+                        "`tensorflow_io` is not installed, for additional file "
+                        + "system support (https://www.tensorflow.org/io), "
+                        + "run `pip install tensorflow_io`."
+                    )
+                    print(warning_msg)
+
     @property
     def data_provider(self):
         return self._data_provider
@@ -196,3 +214,19 @@ def _parse_event_files_spec(logdir_spec):
             path = os.path.realpath(os.path.expanduser(path))
         files[path] = run_name
     return files
+
+
+def _cloud_fs_supported() -> bool:
+    """Checks if the cloud filesystems are supported."""
+    for fs in _CLOUD_FILESYSTEMS:
+        try:
+            if fs not in tf.io.gfile.get_registered_schemes():
+                return False
+        except:
+            # `tf.io.gfile.get_registered_schemes` API is not available,
+            # fall back to `tf.io.gfile.exists`.
+            try:
+                tf.io.gfile.exists(fs + "://tmp")
+            except (tf.errors.UnimplementedError, NotImplementedError):
+                return False
+    return True
diff --git a/tensorboard/backend/event_processing/data_ingester_test.py b/tensorboard/backend/event_processing/data_ingester_test.py
index 5413bd12c48..97f3c1bd1d4 100644
--- a/tensorboard/backend/event_processing/data_ingester_test.py
+++ b/tensorboard/backend/event_processing/data_ingester_test.py
@@ -22,6 +22,7 @@
 
 from tensorboard import test as tb_test
 from tensorboard.backend.event_processing import data_ingester
+from tensorboard.compat import tf
 
 
 class FakeFlags(object):
@@ -251,5 +252,24 @@ def testSingleLetterGroup(self):
             )
 
 
+class CloudFileSystemsSupport(tb_test.TestCase):
+    def testCloudFsSupported(self):
+        with mock.patch.object(
+            tf.io.gfile, "get_registered_schemes", return_value=["gs", "s3"]
+        ):
+            self.assertTrue(data_ingester._cloud_fs_supported())
+
+    def testCloudFsNotSupported(self):
+        with mock.patch.object(
+            tf.io.gfile, "get_registered_schemes", side_effect=Exception("oops")
+        ):
+            with mock.patch.object(
+                tf.io.gfile,
+                "exists",
+                side_effect=tf.errors.UnimplementedError(None, None, "noop"),
+            ):
+                self.assertFalse(data_ingester._cloud_fs_supported())
+
+
 if __name__ == "__main__":
     tb_test.main()

From e807a00b709b540fb72584565c1a2fa827a48a90 Mon Sep 17 00:00:00 2001
From: Yating Jing <ytjing@google.com>
Date: Fri, 7 Jan 2022 21:13:00 +0000
Subject: [PATCH 2/9] lint

---
 tensorboard/backend/event_processing/data_ingester.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorboard/backend/event_processing/data_ingester.py b/tensorboard/backend/event_processing/data_ingester.py
index 803372440ea..c3631e30bb0 100644
--- a/tensorboard/backend/event_processing/data_ingester.py
+++ b/tensorboard/backend/event_processing/data_ingester.py
@@ -88,7 +88,7 @@ def __init__(self, flags):
         if not getattr(tf, "__version__", "stub") == "stub":
             if not _cloud_fs_supported():
                 try:
-                    import tensorflow_io
+                    import tensorflow_io as tfio
                 except:
                     warning_msg = (
                         "`tensorflow_io` is not installed, for additional file "

From 7aef46f547a59074fd02c826130716c43efb2dc2 Mon Sep 17 00:00:00 2001
From: Yating Jing <ytjing@google.com>
Date: Fri, 7 Jan 2022 21:15:26 +0000
Subject: [PATCH 3/9] fix unused import error

---
 tensorboard/backend/event_processing/data_ingester.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorboard/backend/event_processing/data_ingester.py b/tensorboard/backend/event_processing/data_ingester.py
index c3631e30bb0..fd7c821b05a 100644
--- a/tensorboard/backend/event_processing/data_ingester.py
+++ b/tensorboard/backend/event_processing/data_ingester.py
@@ -88,7 +88,7 @@ def __init__(self, flags):
         if not getattr(tf, "__version__", "stub") == "stub":
             if not _cloud_fs_supported():
                 try:
-                    import tensorflow_io as tfio
+                    import tensorflow_io  # pylint: disable=unused-import
                 except:
                     warning_msg = (
                         "`tensorflow_io` is not installed, for additional file "

From 3ae6cf64f24aed5d735e2c5fd852b063ea82c646 Mon Sep 17 00:00:00 2001
From: Yating Jing <ytjing@google.com>
Date: Fri, 7 Jan 2022 21:20:04 +0000
Subject: [PATCH 4/9] fix unused import error

---
 tensorboard/backend/event_processing/data_ingester.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorboard/backend/event_processing/data_ingester.py b/tensorboard/backend/event_processing/data_ingester.py
index fd7c821b05a..bc4bfeaf721 100644
--- a/tensorboard/backend/event_processing/data_ingester.py
+++ b/tensorboard/backend/event_processing/data_ingester.py
@@ -88,7 +88,7 @@ def __init__(self, flags):
         if not getattr(tf, "__version__", "stub") == "stub":
             if not _cloud_fs_supported():
                 try:
-                    import tensorflow_io  # pylint: disable=unused-import
+                    import tensorflow_io  # noqa: F401
                 except:
                     warning_msg = (
                         "`tensorflow_io` is not installed, for additional file "

From a8f4c480a0a7229f1a64fd35efa50afe0e379a6c Mon Sep 17 00:00:00 2001
From: Yating Jing <ytjing@google.com>
Date: Tue, 11 Jan 2022 23:11:52 +0000
Subject: [PATCH 5/9] check whether user requested TF I/O schemes before
 importing tensorflow_io

---
 .../backend/event_processing/data_ingester.py | 106 +++++++++++----
 .../event_processing/data_ingester_test.py    | 126 +++++++++++++++---
 2 files changed, 193 insertions(+), 39 deletions(-)

diff --git a/tensorboard/backend/event_processing/data_ingester.py b/tensorboard/backend/event_processing/data_ingester.py
index bc4bfeaf721..ecbf80cdd9f 100644
--- a/tensorboard/backend/event_processing/data_ingester.py
+++ b/tensorboard/backend/event_processing/data_ingester.py
@@ -48,7 +48,7 @@
 
 # TensorFlow I/O file systems of interest (not available in TensorFlow's
 # built-in support).
-_CLOUD_FILESYSTEMS = ["s3"]
+_CLOUD_FILESYSTEMS = {"gs", "s3"}
 
 logger = tb_logging.get_logger()
 
@@ -86,16 +86,10 @@ def __init__(self, flags):
 
         # Conditionally import tensorflow_io.
         if not getattr(tf, "__version__", "stub") == "stub":
-            if not _cloud_fs_supported():
-                try:
-                    import tensorflow_io  # noqa: F401
-                except:
-                    warning_msg = (
-                        "`tensorflow_io` is not installed, for additional file "
-                        + "system support (https://www.tensorflow.org/io), "
-                        + "run `pip install tensorflow_io`."
-                    )
-                    print(warning_msg)
+            tfio_filesystems = _check_filesystem_support(
+                self._path_to_run.keys()
+            )
+            _try_to_support_tfio(self._path_to_run.keys())
 
     @property
     def data_provider(self):
@@ -216,17 +210,81 @@ def _parse_event_files_spec(logdir_spec):
     return files
 
 
-def _cloud_fs_supported() -> bool:
-    """Checks if the cloud filesystems are supported."""
-    for fs in _CLOUD_FILESYSTEMS:
-        try:
-            if fs not in tf.io.gfile.get_registered_schemes():
-                return False
-        except:
-            # `tf.io.gfile.get_registered_schemes` API is not available,
-            # fall back to `tf.io.gfile.exists`.
+def _get_filesystem_scheme(path):
+    """Extracts filesystem from a given path.
+
+    Args:
+        path: A strings representing an input log directory.
+    Returns:
+        Filesystem scheme (followed by `://`), None if the path doesn't contain
+        a scheme.
+    """
+    if "://" not in path:
+        return None
+    return path.split("://")[0]
+
+
+def _check_filesystem_support(paths):
+    """Examines the list of filesystems user requested.
+
+    Args:
+        paths: A list of strings representing input log directories.
+    Returns:
+        A list of TF I/O filesystems of interest.
+    """
+    tfio_filesystems = []
+    for path in paths:
+        fs = _get_filesystem_scheme(path)
+        if fs is not None and fs in _CLOUD_FILESYSTEMS:
+            tfio_filesystems.append(fs)
+    return tfio_filesystems
+
+
+def _try_to_support_tfio(paths):
+    """Try to support TF I/O.
+
+    Args:
+        paths: A list of strings representing input log directories.
+    """
+    get_registered_schemes = getattr(
+        tf.io.gfile, "get_registered_schemes", None
+    )
+    registered_schemes = (
+        None if get_registered_schemes is None else get_registered_schemes()
+    )
+    missing_fs = None
+    for path in paths:
+        fs = _get_filesystem_scheme(path)
+        if fs is None:
+            continue
+        # Use `tf.io.gfile.exists.get_registered_schemes` if possible.
+        if registered_schemes is not None:
+            if fs not in registered_schemes:
+                missing_fs = fs
+                break
+        else:
+            # Fall back to `tf.io.gfile.exists`.
             try:
-                tf.io.gfile.exists(fs + "://tmp")
-            except (tf.errors.UnimplementedError, NotImplementedError):
-                return False
-    return True
+                tf.io.gfile.exists(path)
+            except tf.errors.UnimplementedError:
+                missing_fs = fs
+                break
+
+    if missing_fs:
+        try:
+            import tensorflow_io  # noqa: F401
+        except ImportError:
+            supported_schemes_msg = (
+                " (supported schemes: {})".format(registered_schemes)
+                if registered_schemes
+                else ""
+            )
+            raise tf.errors.UnimplementedError(
+                None,
+                None,
+                (
+                    "Error: Unsupported filename scheme '{}'{}. For additional"
+                    + " filesystem support, consider installing TensorFlow I/O"
+                    + " (https://www.tensorflow.org/io) via `pip install tensorflow-io`."
+                ).format(missing_fs, supported_schemes_msg),
+            )
diff --git a/tensorboard/backend/event_processing/data_ingester_test.py b/tensorboard/backend/event_processing/data_ingester_test.py
index 97f3c1bd1d4..37a6fc433b9 100644
--- a/tensorboard/backend/event_processing/data_ingester_test.py
+++ b/tensorboard/backend/event_processing/data_ingester_test.py
@@ -252,23 +252,119 @@ def testSingleLetterGroup(self):
             )
 
 
-class CloudFileSystemsSupport(tb_test.TestCase):
-    def testCloudFsSupported(self):
-        with mock.patch.object(
-            tf.io.gfile, "get_registered_schemes", return_value=["gs", "s3"]
-        ):
-            self.assertTrue(data_ingester._cloud_fs_supported())
+class TfioSupport(tb_test.TestCase):
+    def testCheckFilesystemSupport(self):
+        self.assertEqual(
+            ["s3"],
+            data_ingester._check_filesystem_support(
+                ["tmp/demo", "s3://bucket/123", "file://tmp"]
+            ),
+        )
 
-    def testCloudFsNotSupported(self):
+    def testTryToSupportTfio(self):
+        with mock.patch.object(tf.io, "gfile") as gfile_mock:
+            gfile_mock.return_value.get_registered_schemes = mock.MagicMock()
+            with mock.patch(
+                "builtins.__import__", side_effect=mock.MagicMock()
+            ) as mock_import:
+                with mock.patch.object(
+                    gfile_mock,
+                    "get_registered_schemes",
+                    return_value=["gs", "s3"],
+                ) as mock_get_registered_schemes:
+                    data_ingester._try_to_support_tfio(
+                        ["tmp/demo", "s3://bucket/123"]
+                    )
+        mock_import.assert_not_called()
+        mock_get_registered_schemes.assert_called_once()
+
+    def testTryToSupportTfio_import(self):
+        with mock.patch.object(tf.io, "gfile") as gfile_mock:
+            gfile_mock.return_value.get_registered_schemes = mock.MagicMock()
+            with mock.patch(
+                "builtins.__import__", side_effect=mock.MagicMock()
+            ) as mock_import:
+                with mock.patch.object(
+                    gfile_mock,
+                    "get_registered_schemes",
+                    return_value=["file", ""],
+                ) as mock_get_registered_schemes:
+                    data_ingester._try_to_support_tfio(
+                        ["tmp/demo", "s3://bucket/123"]
+                    )
+        self.assertEqual("tensorflow_io", mock_import.call_args[0][0])
+        mock_get_registered_schemes.assert_called_once()
+
+    def testTryToSupportTfio_raiseError(self):
+        with mock.patch.object(tf.io, "gfile") as gfile_mock:
+            gfile_mock.return_value.get_registered_schemes = mock.MagicMock()
+            with mock.patch(
+                "builtins.__import__",
+                side_effect=ImportError,
+            ) as mock_import:
+                with mock.patch.object(
+                    gfile_mock,
+                    "get_registered_schemes",
+                    return_value=["file", "ram"],
+                ) as mock_get_registered_schemes:
+                    err_msg = (
+                        "Error: Unsupported filename scheme 's3' (supported schemes: ['file', 'ram'])."
+                        + " For additional filesystem support, consider installing TensorFlow I/O"
+                        + " (https://www.tensorflow.org/io) via `pip install tensorflow-io`."
+                    )
+                    with self.assertRaisesWithLiteralMatch(
+                        tf.errors.UnimplementedError, err_msg
+                    ):
+                        data_ingester._try_to_support_tfio(
+                            ["tmp/demo", "s3://bucket/123"]
+                        )
+        self.assertEqual("tensorflow_io", mock_import.call_args[0][0])
+        mock_get_registered_schemes.assert_called_once()
+
+    def testTryToSupportTfio_fallback(self):
+        with mock.patch.object(
+            tf.io.gfile, "get_registered_schemes"
+        ) as mock_get_registered_schemes:
+            mock_get_registered_schemes.return_value = None
+            with mock.patch(
+                "builtins.__import__", side_effect=mock.MagicMock()
+            ) as mock_import:
+                with mock.patch.object(
+                    tf.io.gfile,
+                    "exists",
+                    return_value=True,
+                ) as mock_gfile_exists:
+                    data_ingester._try_to_support_tfio(["gs://bucket/abc"])
+        mock_import.assert_not_called()
+        mock_gfile_exists.assert_called_once_with("gs://bucket/abc")
+
+    def testTryToSupportTfio_fallback_raiseError(self):
         with mock.patch.object(
-            tf.io.gfile, "get_registered_schemes", side_effect=Exception("oops")
-        ):
-            with mock.patch.object(
-                tf.io.gfile,
-                "exists",
-                side_effect=tf.errors.UnimplementedError(None, None, "noop"),
-            ):
-                self.assertFalse(data_ingester._cloud_fs_supported())
+            tf.io.gfile, "get_registered_schemes"
+        ) as mock_get_registered_schemes:
+            mock_get_registered_schemes.return_value = None
+            with mock.patch(
+                "builtins.__import__",
+                side_effect=ImportError,
+            ) as mock_import:
+                with mock.patch.object(
+                    tf.io.gfile,
+                    "exists",
+                    side_effect=tf.errors.UnimplementedError(
+                        None, None, "oops"
+                    ),
+                ) as mock_gfile_exists:
+                    err_msg = (
+                        "Error: Unsupported filename scheme 'gs'."
+                        + " For additional filesystem support, consider installing TensorFlow I/O"
+                        + " (https://www.tensorflow.org/io) via `pip install tensorflow-io`."
+                    )
+                    with self.assertRaisesWithLiteralMatch(
+                        tf.errors.UnimplementedError, err_msg
+                    ):
+                        data_ingester._try_to_support_tfio(["gs://bucket/abc"])
+        self.assertEqual("tensorflow_io", mock_import.call_args[0][0])
+        mock_gfile_exists.assert_called_once_with("gs://bucket/abc")
 
 
 if __name__ == "__main__":

From f62e4955a7eb2e4fbdace3d8db4421eb0ffa4c11 Mon Sep 17 00:00:00 2001
From: Yating Jing <ytjing@google.com>
Date: Wed, 12 Jan 2022 17:44:04 +0000
Subject: [PATCH 6/9] refactor

---
 .../backend/event_processing/data_ingester.py |  48 +++---
 .../event_processing/data_ingester_test.py    | 150 +++++++-----------
 2 files changed, 77 insertions(+), 121 deletions(-)

diff --git a/tensorboard/backend/event_processing/data_ingester.py b/tensorboard/backend/event_processing/data_ingester.py
index ecbf80cdd9f..dac9fdd4511 100644
--- a/tensorboard/backend/event_processing/data_ingester.py
+++ b/tensorboard/backend/event_processing/data_ingester.py
@@ -46,10 +46,6 @@
     pr_curve_metadata.PLUGIN_NAME: 100,
 }
 
-# TensorFlow I/O file systems of interest (not available in TensorFlow's
-# built-in support).
-_CLOUD_FILESYSTEMS = {"gs", "s3"}
-
 logger = tb_logging.get_logger()
 
 
@@ -216,8 +212,9 @@ def _get_filesystem_scheme(path):
     Args:
         path: A strings representing an input log directory.
     Returns:
-        Filesystem scheme (followed by `://`), None if the path doesn't contain
-        a scheme.
+        Filesystem scheme, None if the path doesn't contain one. The filesystem
+        scheme is usually separated by `://` from the local filesystem path if
+        given. For example, the scheme of `file://tmp/tf` is `file`.
     """
     if "://" not in path:
         return None
@@ -227,21 +224,7 @@ def _get_filesystem_scheme(path):
 def _check_filesystem_support(paths):
     """Examines the list of filesystems user requested.
 
-    Args:
-        paths: A list of strings representing input log directories.
-    Returns:
-        A list of TF I/O filesystems of interest.
-    """
-    tfio_filesystems = []
-    for path in paths:
-        fs = _get_filesystem_scheme(path)
-        if fs is not None and fs in _CLOUD_FILESYSTEMS:
-            tfio_filesystems.append(fs)
-    return tfio_filesystems
-
-
-def _try_to_support_tfio(paths):
-    """Try to support TF I/O.
+    If TF I/O schemes are requested, try to import tensorflow_io module.
 
     Args:
         paths: A list of strings representing input log directories.
@@ -252,25 +235,30 @@ def _try_to_support_tfio(paths):
     registered_schemes = (
         None if get_registered_schemes is None else get_registered_schemes()
     )
-    missing_fs = None
-    for path in paths:
-        fs = _get_filesystem_scheme(path)
-        if fs is None:
+
+    # Only need to check one path for each scheme.
+    scheme_to_path = {_get_filesystem_scheme(path): path for path in paths}
+    missing_scheme = None
+    for scheme, path in scheme_to_path.items():
+        if scheme is None:
             continue
         # Use `tf.io.gfile.exists.get_registered_schemes` if possible.
         if registered_schemes is not None:
-            if fs not in registered_schemes:
-                missing_fs = fs
+            if scheme not in registered_schemes:
+                missing_scheme = scheme
                 break
         else:
             # Fall back to `tf.io.gfile.exists`.
             try:
                 tf.io.gfile.exists(path)
             except tf.errors.UnimplementedError:
-                missing_fs = fs
+                missing_scheme = scheme
                 break
+            except tf.errors.OpError:
+                # Swallow other errors; we aren't concerned about them at this point.
+                pass
 
-    if missing_fs:
+    if missing_scheme:
         try:
             import tensorflow_io  # noqa: F401
         except ImportError:
@@ -286,5 +274,5 @@ def _try_to_support_tfio(paths):
                     "Error: Unsupported filename scheme '{}'{}. For additional"
                     + " filesystem support, consider installing TensorFlow I/O"
                     + " (https://www.tensorflow.org/io) via `pip install tensorflow-io`."
-                ).format(missing_fs, supported_schemes_msg),
+                ).format(missing_scheme, supported_schemes_msg),
             )
diff --git a/tensorboard/backend/event_processing/data_ingester_test.py b/tensorboard/backend/event_processing/data_ingester_test.py
index 37a6fc433b9..22c7cc6895d 100644
--- a/tensorboard/backend/event_processing/data_ingester_test.py
+++ b/tensorboard/backend/event_processing/data_ingester_test.py
@@ -252,119 +252,87 @@ def testSingleLetterGroup(self):
             )
 
 
-class TfioSupport(tb_test.TestCase):
+class FileSystemSupport(tb_test.TestCase):
     def testCheckFilesystemSupport(self):
-        self.assertEqual(
-            ["s3"],
-            data_ingester._check_filesystem_support(
-                ["tmp/demo", "s3://bucket/123", "file://tmp"]
-            ),
-        )
-
-    def testTryToSupportTfio(self):
-        with mock.patch.object(tf.io, "gfile") as gfile_mock:
-            gfile_mock.return_value.get_registered_schemes = mock.MagicMock()
-            with mock.patch(
-                "builtins.__import__", side_effect=mock.MagicMock()
-            ) as mock_import:
-                with mock.patch.object(
-                    gfile_mock,
-                    "get_registered_schemes",
-                    return_value=["gs", "s3"],
-                ) as mock_get_registered_schemes:
-                    data_ingester._try_to_support_tfio(
-                        ["tmp/demo", "s3://bucket/123"]
-                    )
+        with mock.patch.object(
+            tf.io.gfile,
+            "get_registered_schemes",
+            return_value=["gs", "s3"],
+        ) as mock_get_registered_schemes:
+            with mock.patch("builtins.__import__") as mock_import:
+                data_ingester._check_filesystem_support(
+                    ["tmp/demo", "s3://bucket/123"]
+                )
         mock_import.assert_not_called()
         mock_get_registered_schemes.assert_called_once()
 
-    def testTryToSupportTfio_import(self):
-        with mock.patch.object(tf.io, "gfile") as gfile_mock:
-            gfile_mock.return_value.get_registered_schemes = mock.MagicMock()
-            with mock.patch(
-                "builtins.__import__", side_effect=mock.MagicMock()
-            ) as mock_import:
-                with mock.patch.object(
-                    gfile_mock,
-                    "get_registered_schemes",
-                    return_value=["file", ""],
-                ) as mock_get_registered_schemes:
-                    data_ingester._try_to_support_tfio(
-                        ["tmp/demo", "s3://bucket/123"]
-                    )
+    def testCheckFilesystemSupport_importTFIO(self):
+        with mock.patch.object(
+            tf.io.gfile,
+            "get_registered_schemes",
+            return_value=["file", ""],
+        ) as mock_get_registered_schemes:
+            with mock.patch("builtins.__import__") as mock_import:
+                data_ingester._check_filesystem_support(
+                    ["tmp/demo", "s3://bucket/123"]
+                )
         self.assertEqual("tensorflow_io", mock_import.call_args[0][0])
         mock_get_registered_schemes.assert_called_once()
 
-    def testTryToSupportTfio_raiseError(self):
-        with mock.patch.object(tf.io, "gfile") as gfile_mock:
-            gfile_mock.return_value.get_registered_schemes = mock.MagicMock()
+    def testCheckFilesystemSupport_raiseError(self):
+        with mock.patch.object(
+            tf.io.gfile,
+            "get_registered_schemes",
+            return_value=["file", "ram"],
+        ) as mock_get_registered_schemes:
             with mock.patch(
                 "builtins.__import__",
                 side_effect=ImportError,
             ) as mock_import:
-                with mock.patch.object(
-                    gfile_mock,
-                    "get_registered_schemes",
-                    return_value=["file", "ram"],
-                ) as mock_get_registered_schemes:
-                    err_msg = (
-                        "Error: Unsupported filename scheme 's3' (supported schemes: ['file', 'ram'])."
-                        + " For additional filesystem support, consider installing TensorFlow I/O"
-                        + " (https://www.tensorflow.org/io) via `pip install tensorflow-io`."
+                err_msg = (
+                    "Error: Unsupported filename scheme 's3' (supported schemes: ['file', 'ram'])."
+                    + " For additional filesystem support, consider installing TensorFlow I/O"
+                    + " (https://www.tensorflow.org/io) via `pip install tensorflow-io`."
+                )
+                with self.assertRaisesWithLiteralMatch(
+                    tf.errors.UnimplementedError, err_msg
+                ):
+                    data_ingester._check_filesystem_support(
+                        ["tmp/demo", "s3://bucket/123"]
                     )
-                    with self.assertRaisesWithLiteralMatch(
-                        tf.errors.UnimplementedError, err_msg
-                    ):
-                        data_ingester._try_to_support_tfio(
-                            ["tmp/demo", "s3://bucket/123"]
-                        )
         self.assertEqual("tensorflow_io", mock_import.call_args[0][0])
         mock_get_registered_schemes.assert_called_once()
 
-    def testTryToSupportTfio_fallback(self):
-        with mock.patch.object(
-            tf.io.gfile, "get_registered_schemes"
-        ) as mock_get_registered_schemes:
-            mock_get_registered_schemes.return_value = None
-            with mock.patch(
-                "builtins.__import__", side_effect=mock.MagicMock()
-            ) as mock_import:
-                with mock.patch.object(
-                    tf.io.gfile,
-                    "exists",
-                    return_value=True,
-                ) as mock_gfile_exists:
-                    data_ingester._try_to_support_tfio(["gs://bucket/abc"])
+    def testCheckFilesystemSupport_fallback(self):
+        with mock.patch.object(tf.io, "gfile", autospec=True) as mock_gfile:
+            del mock_gfile.get_registered_schemes
+            with mock.patch("builtins.__import__") as mock_import:
+                mock_gfile.exists.return_value = True
+                data_ingester._check_filesystem_support(["gs://bucket/abc"])
         mock_import.assert_not_called()
-        mock_gfile_exists.assert_called_once_with("gs://bucket/abc")
+        mock_gfile.exists.assert_called_once_with("gs://bucket/abc")
 
-    def testTryToSupportTfio_fallback_raiseError(self):
-        with mock.patch.object(
-            tf.io.gfile, "get_registered_schemes"
-        ) as mock_get_registered_schemes:
-            mock_get_registered_schemes.return_value = None
+    def testCheckFilesystemSupport_fallback_raiseError(self):
+        with mock.patch.object(tf.io, "gfile", autospec=True) as mock_gfile:
+            del mock_gfile.get_registered_schemes
             with mock.patch(
                 "builtins.__import__",
                 side_effect=ImportError,
             ) as mock_import:
-                with mock.patch.object(
-                    tf.io.gfile,
-                    "exists",
-                    side_effect=tf.errors.UnimplementedError(
-                        None, None, "oops"
-                    ),
-                ) as mock_gfile_exists:
-                    err_msg = (
-                        "Error: Unsupported filename scheme 'gs'."
-                        + " For additional filesystem support, consider installing TensorFlow I/O"
-                        + " (https://www.tensorflow.org/io) via `pip install tensorflow-io`."
-                    )
-                    with self.assertRaisesWithLiteralMatch(
-                        tf.errors.UnimplementedError, err_msg
-                    ):
-                        data_ingester._try_to_support_tfio(["gs://bucket/abc"])
+                mock_gfile.exists.side_effect = tf.errors.UnimplementedError(
+                    None, None, "oops"
+                )
+                err_msg = (
+                    "Error: Unsupported filename scheme 'gs'."
+                    + " For additional filesystem support, consider installing TensorFlow I/O"
+                    + " (https://www.tensorflow.org/io) via `pip install tensorflow-io`."
+                )
+                with self.assertRaisesWithLiteralMatch(
+                    tf.errors.UnimplementedError, err_msg
+                ):
+                    data_ingester._check_filesystem_support(["gs://bucket/abc"])
         self.assertEqual("tensorflow_io", mock_import.call_args[0][0])
-        mock_gfile_exists.assert_called_once_with("gs://bucket/abc")
+        mock_gfile.exists.assert_called_once_with("gs://bucket/abc")
 
 
 if __name__ == "__main__":

From cc282d0347e85ced54943be5bef616eec5dfef48 Mon Sep 17 00:00:00 2001
From: Yating Jing <ytjing@google.com>
Date: Wed, 12 Jan 2022 17:46:38 +0000
Subject: [PATCH 7/9] oops

---
 tensorboard/backend/event_processing/data_ingester.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/tensorboard/backend/event_processing/data_ingester.py b/tensorboard/backend/event_processing/data_ingester.py
index dac9fdd4511..1c06a55a4a6 100644
--- a/tensorboard/backend/event_processing/data_ingester.py
+++ b/tensorboard/backend/event_processing/data_ingester.py
@@ -82,10 +82,7 @@ def __init__(self, flags):
 
         # Conditionally import tensorflow_io.
         if not getattr(tf, "__version__", "stub") == "stub":
-            tfio_filesystems = _check_filesystem_support(
-                self._path_to_run.keys()
-            )
-            _try_to_support_tfio(self._path_to_run.keys())
+            _check_filesystem_support(self._path_to_run.keys())
 
     @property
     def data_provider(self):

From ff09e209859f9560504a84c4e5602449c92527c0 Mon Sep 17 00:00:00 2001
From: Yating Jing <ytjing@google.com>
Date: Wed, 12 Jan 2022 20:21:16 +0000
Subject: [PATCH 8/9] fix error where attribute `get_registered_schemes` not
 found in gfile mock

---
 .../event_processing/data_ingester_test.py    | 33 +++++++++----------
 1 file changed, 15 insertions(+), 18 deletions(-)

diff --git a/tensorboard/backend/event_processing/data_ingester_test.py b/tensorboard/backend/event_processing/data_ingester_test.py
index 22c7cc6895d..1f49e139992 100644
--- a/tensorboard/backend/event_processing/data_ingester_test.py
+++ b/tensorboard/backend/event_processing/data_ingester_test.py
@@ -254,37 +254,34 @@ def testSingleLetterGroup(self):
 
 class FileSystemSupport(tb_test.TestCase):
     def testCheckFilesystemSupport(self):
-        with mock.patch.object(
-            tf.io.gfile,
-            "get_registered_schemes",
-            return_value=["gs", "s3"],
-        ) as mock_get_registered_schemes:
+        with mock.patch.object(tf.io, "gfile", autospec=True) as mock_gfile:
+            mock_gfile.get_registered_schemes = mock.MagicMock(
+                return_value=["gs", "s3"]
+            )
             with mock.patch("builtins.__import__") as mock_import:
                 data_ingester._check_filesystem_support(
                     ["tmp/demo", "s3://bucket/123"]
                 )
         mock_import.assert_not_called()
-        mock_get_registered_schemes.assert_called_once()
+        mock_gfile.get_registered_schemes.assert_called_once()
 
     def testCheckFilesystemSupport_importTFIO(self):
-        with mock.patch.object(
-            tf.io.gfile,
-            "get_registered_schemes",
-            return_value=["file", ""],
-        ) as mock_get_registered_schemes:
+        with mock.patch.object(tf.io, "gfile", autospec=True) as mock_gfile:
+            mock_gfile.get_registered_schemes = mock.MagicMock(
+                return_value=["file", ""]
+            )
             with mock.patch("builtins.__import__") as mock_import:
                 data_ingester._check_filesystem_support(
                     ["tmp/demo", "s3://bucket/123"]
                 )
         self.assertEqual("tensorflow_io", mock_import.call_args[0][0])
-        mock_get_registered_schemes.assert_called_once()
+        mock_gfile.get_registered_schemes.assert_called_once()
 
     def testCheckFilesystemSupport_raiseError(self):
-        with mock.patch.object(
-            tf.io.gfile,
-            "get_registered_schemes",
-            return_value=["file", "ram"],
-        ) as mock_get_registered_schemes:
+        with mock.patch.object(tf.io, "gfile", autospec=True) as mock_gfile:
+            mock_gfile.get_registered_schemes = mock.MagicMock(
+                return_value=["file", "ram"]
+            )
             with mock.patch(
                 "builtins.__import__",
                 side_effect=ImportError,
@@ -301,7 +298,7 @@ def testCheckFilesystemSupport_raiseError(self):
                         ["tmp/demo", "s3://bucket/123"]
                     )
         self.assertEqual("tensorflow_io", mock_import.call_args[0][0])
-        mock_get_registered_schemes.assert_called_once()
+        mock_gfile.get_registered_schemes.assert_called_once()
 
     def testCheckFilesystemSupport_fallback(self):
         with mock.patch.object(tf.io, "gfile", autospec=True) as mock_gfile:

From 388f449221e8cfb2a400636670722bd613ab3c29 Mon Sep 17 00:00:00 2001
From: Yating Jing <ytjing@google.com>
Date: Thu, 13 Jan 2022 23:21:26 +0000
Subject: [PATCH 9/9] refactor

---
 tensorboard/backend/event_processing/BUILD    |  3 ++
 .../backend/event_processing/data_ingester.py |  9 ++---
 .../event_processing/data_ingester_test.py    | 36 +++++++++++--------
 3 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/tensorboard/backend/event_processing/BUILD b/tensorboard/backend/event_processing/BUILD
index bcce833f7e8..d1adfe1f7fe 100644
--- a/tensorboard/backend/event_processing/BUILD
+++ b/tensorboard/backend/event_processing/BUILD
@@ -36,6 +36,7 @@ py_library(
         ":data_provider",
         ":event_multiplexer",
         ":tag_types",
+        "//tensorboard/compat:tensorflow",
         "//tensorboard/data:ingester",
         "//tensorboard/plugins/audio:metadata",
         "//tensorboard/plugins/histogram:metadata",
@@ -53,7 +54,9 @@ py_test(
     srcs_version = "PY3",
     deps = [
         ":data_ingester",
+        "//tensorboard:expect_tensorflow_installed",
         "//tensorboard:test",
+        "//tensorboard/compat:tensorflow",
     ],
 )
 
diff --git a/tensorboard/backend/event_processing/data_ingester.py b/tensorboard/backend/event_processing/data_ingester.py
index 1c06a55a4a6..2ff648713ee 100644
--- a/tensorboard/backend/event_processing/data_ingester.py
+++ b/tensorboard/backend/event_processing/data_ingester.py
@@ -204,14 +204,15 @@ def _parse_event_files_spec(logdir_spec):
 
 
 def _get_filesystem_scheme(path):
-    """Extracts filesystem from a given path.
+    """Extracts filesystem scheme from a given path.
+
+    The filesystem scheme is usually separated by `://` from the local filesystem
+    path if given. For example, the scheme of `file://tmp/tf` is `file`.
 
     Args:
         path: A strings representing an input log directory.
     Returns:
-        Filesystem scheme, None if the path doesn't contain one. The filesystem
-        scheme is usually separated by `://` from the local filesystem path if
-        given. For example, the scheme of `file://tmp/tf` is `file`.
+        Filesystem scheme, None if the path doesn't contain one.
     """
     if "://" not in path:
         return None
diff --git a/tensorboard/backend/event_processing/data_ingester_test.py b/tensorboard/backend/event_processing/data_ingester_test.py
index 1f49e139992..77b336a550a 100644
--- a/tensorboard/backend/event_processing/data_ingester_test.py
+++ b/tensorboard/backend/event_processing/data_ingester_test.py
@@ -254,34 +254,40 @@ def testSingleLetterGroup(self):
 
 class FileSystemSupport(tb_test.TestCase):
     def testCheckFilesystemSupport(self):
-        with mock.patch.object(tf.io, "gfile", autospec=True) as mock_gfile:
-            mock_gfile.get_registered_schemes = mock.MagicMock(
-                return_value=["gs", "s3"]
-            )
+        with mock.patch.object(
+            tf.io.gfile,
+            "get_registered_schemes",
+            autospec=True,
+            return_value=["g3", "s3"],
+        ) as mock_get_registered_schemes:
             with mock.patch("builtins.__import__") as mock_import:
                 data_ingester._check_filesystem_support(
                     ["tmp/demo", "s3://bucket/123"]
                 )
         mock_import.assert_not_called()
-        mock_gfile.get_registered_schemes.assert_called_once()
+        mock_get_registered_schemes.assert_called_once()
 
     def testCheckFilesystemSupport_importTFIO(self):
-        with mock.patch.object(tf.io, "gfile", autospec=True) as mock_gfile:
-            mock_gfile.get_registered_schemes = mock.MagicMock(
-                return_value=["file", ""]
-            )
+        with mock.patch.object(
+            tf.io.gfile,
+            "get_registered_schemes",
+            autospec=True,
+            return_value=["file", ""],
+        ) as mock_get_registered_schemes:
             with mock.patch("builtins.__import__") as mock_import:
                 data_ingester._check_filesystem_support(
                     ["tmp/demo", "s3://bucket/123"]
                 )
         self.assertEqual("tensorflow_io", mock_import.call_args[0][0])
-        mock_gfile.get_registered_schemes.assert_called_once()
+        mock_get_registered_schemes.assert_called_once()
 
     def testCheckFilesystemSupport_raiseError(self):
-        with mock.patch.object(tf.io, "gfile", autospec=True) as mock_gfile:
-            mock_gfile.get_registered_schemes = mock.MagicMock(
-                return_value=["file", "ram"]
-            )
+        with mock.patch.object(
+            tf.io.gfile,
+            "get_registered_schemes",
+            autospec=True,
+            return_value=["file", "ram"],
+        ) as mock_get_registered_schemes:
             with mock.patch(
                 "builtins.__import__",
                 side_effect=ImportError,
@@ -298,7 +304,7 @@ def testCheckFilesystemSupport_raiseError(self):
                         ["tmp/demo", "s3://bucket/123"]
                     )
         self.assertEqual("tensorflow_io", mock_import.call_args[0][0])
-        mock_gfile.get_registered_schemes.assert_called_once()
+        mock_get_registered_schemes.assert_called_once()
 
     def testCheckFilesystemSupport_fallback(self):
         with mock.patch.object(tf.io, "gfile", autospec=True) as mock_gfile: