Add tests and support S3 with tensorboard-notf

tensorboard/backend/event_processing/directory_watcher.py

@@ -179,7 +179,7 @@ def _SetPath(self, path):
       path: The full path of the file to watch.
     """
     old_path = self._path
-    if old_path and not io_wrapper.IsGCSPath(old_path):
+    if old_path and not io_wrapper.IsCloudPath(old_path):
       try:
         # We're done with the path, so store its size.
         size = tf.io.gfile.stat(old_path).length
@@ -211,7 +211,7 @@ def _GetNextPath(self):
 
     # Don't bother checking if the paths are GCS (which we can't check) or if
     # we've already detected an OOO write.
-    if not io_wrapper.IsGCSPath(paths[0]) and not self._ooo_writes_detected:
+    if not io_wrapper.IsCloudPath(paths[0]) and not self._ooo_writes_detected:
       # Check the previous _OOO_WRITE_CHECK_COUNT paths for out of order writes.
       current_path_index = bisect.bisect_left(paths, self._path)
       ooo_check_start = max(0, current_path_index - self._OOO_WRITE_CHECK_COUNT)

tensorboard/backend/event_processing/event_accumulator_test.py

@@ -767,8 +767,8 @@ def FakeScalarSummary(tag, value):
 
     directory = os.path.join(self.get_temp_dir(), 'values_dir')
     if tf.io.gfile.isdir(directory):
-      tf.io.gfile.rmtree(directory)
-    tf.io.gfile.mkdir(directory)
+      os.removedirs(directory)
+    os.mkdir(directory)
 
     writer = test_util.FileWriter(directory, max_queue=100)
 
@@ -845,8 +845,8 @@ def testGraphFromMetaGraphBecomesAvailable(self):
 
     directory = os.path.join(self.get_temp_dir(), 'metagraph_test_values_dir')
     if tf.io.gfile.isdir(directory):
-      tf.io.gfile.rmtree(directory)
-    tf.io.gfile.mkdir(directory)
+      os.removedirs(directory)
+    os.mkdir(directory)
 
     writer = test_util.FileWriter(directory, max_queue=100)
 

tensorboard/backend/event_processing/event_multiplexer_test.py

@@ -29,7 +29,7 @@
 
 def _AddEvents(path):
   if not tf.io.gfile.isdir(path):
-    tf.io.gfile.makedirs(path)
+    os.makedirs(path)
   fpath = os.path.join(path, 'hypothetical.tfevents.out')
   with tf.io.gfile.GFile(fpath, 'w') as f:
     f.write('')
@@ -38,8 +38,8 @@ def _AddEvents(path):
 
 def _CreateCleanDirectory(path):
   if tf.io.gfile.isdir(path):
-    tf.io.gfile.rmtree(path)
-  tf.io.gfile.mkdir(path)
+    os.removedirs(path)
+  os.mkdir(path)
 
 
 class _FakeAccumulator(object):
@@ -206,7 +206,7 @@ def testAddRunsFromDirectory(self):
     self.assertEqual(x.Runs(), {}, 'loading empty directory had no effect')
 
     path1 = join(realdir, 'path1')
-    tf.io.gfile.mkdir(path1)
+    os.mkdir(path1)
     x.AddRunsFromDirectory(realdir)
     self.assertEqual(x.Runs(), {}, 'creating empty subdirectory had no effect')
 

tensorboard/backend/event_processing/io_wrapper.py

@@ -32,16 +32,15 @@
 _ESCAPE_GLOB_CHARACTERS_REGEX = re.compile('([*?[])')
 
 
-# TODO(chihuahua): Rename this method to use camel-case for GCS (Gcs).
-def IsGCSPath(path):
-  return path.startswith("gs://")
-
-
-def IsCnsPath(path):
-  return path.startswith("/cns/")
+def IsCloudPath(path):
+  return (
+    path.startswith("gs://") or
+    path.startswith("s3://") or
+    path.startswith("/cns/")
+  )
 
 def PathSeparator(path):
-  return '/' if IsGCSPath(path) or IsCnsPath(path) else os.sep
+  return '/' if IsCloudPath(path) else os.sep
 
 def IsTensorFlowEventsFile(path):
   """Check the path name to see if it is probably a TF Events file.
@@ -184,7 +183,7 @@ def GetLogdirSubdirectories(path):
     raise ValueError('GetLogdirSubdirectories: path exists and is not a '
                      'directory, %s' % path)
 
-  if IsGCSPath(path) or IsCnsPath(path):
+  if IsCloudPath(path):
     # Glob-ing for files can be significantly faster than recursively
     # walking through directories for some file systems.
     logger.info(

tensorboard/backend/event_processing/io_wrapper_test.py

@@ -33,17 +33,20 @@ def setUp(self):
   def tearDown(self):
     self.stubs.CleanUp()
 
-  def testIsGcsPathIsTrue(self):
-    self.assertTrue(io_wrapper.IsGCSPath('gs://bucket/foo'))
+  def testIsCloudPathGcsIsTrue(self):
+    self.assertTrue(io_wrapper.IsCloudPath('gs://bucket/foo'))
 
-  def testIsGcsPathIsFalse(self):
-    self.assertFalse(io_wrapper.IsGCSPath('/tmp/foo'))
+  def testIsCloudPathS3IsTrue(self):
+    self.assertTrue(io_wrapper.IsCloudPath('s3://bucket/foo'))
 
-  def testIsCnsPathTrue(self):
-    self.assertTrue(io_wrapper.IsCnsPath('/cns/foo/bar'))
+  def testIsCloudPathCnsIsTrue(self):
+    self.assertTrue(io_wrapper.IsCloudPath('/cns/foo/bar'))
 
-  def testIsCnsPathFalse(self):
-    self.assertFalse(io_wrapper.IsCnsPath('/tmp/foo'))
+  def testIsCloudPathFileIsFalse(self):
+    self.assertFalse(io_wrapper.IsCloudPath('file:///tmp/foo'))
+
+  def testIsCloudPathLocalIsFalse(self):
+    self.assertFalse(io_wrapper.IsCloudPath('/tmp/foo'))
 
   def testPathSeparator(self):
     # In nix systems, path separator would be the same as that of CNS/GCS

tensorboard/backend/event_processing/plugin_event_accumulator_test.py

@@ -558,8 +558,8 @@ def FakeScalarSummary(tag, value):
 
     directory = os.path.join(self.get_temp_dir(), 'values_dir')
     if tf.io.gfile.isdir(directory):
-      tf.io.gfile.rmtree(directory)
-    tf.io.gfile.mkdir(directory)
+      os.removedirs(directory)
+    os.mkdir(directory)
 
     writer = test_util.FileWriter(directory, max_queue=100)
 
@@ -636,8 +636,8 @@ def testGraphFromMetaGraphBecomesAvailable(self):
 
     directory = os.path.join(self.get_temp_dir(), 'metagraph_test_values_dir')
     if tf.io.gfile.isdir(directory):
-      tf.io.gfile.rmtree(directory)
-    tf.io.gfile.mkdir(directory)
+      os.removedirs(directory)
+    os.mkdir(directory)
 
     writer = test_util.FileWriter(directory, max_queue=100)
 

tensorboard/backend/event_processing/plugin_event_multiplexer_test.py

@@ -29,7 +29,7 @@
 
 def _AddEvents(path):
   if not tf.io.gfile.isdir(path):
-    tf.io.gfile.makedirs(path)
+    os.makedirs(path)
   fpath = os.path.join(path, 'hypothetical.tfevents.out')
   with tf.io.gfile.GFile(fpath, 'w') as f:
     f.write('')
@@ -38,8 +38,8 @@ def _AddEvents(path):
 
 def _CreateCleanDirectory(path):
   if tf.io.gfile.isdir(path):
-    tf.io.gfile.rmtree(path)
-  tf.io.gfile.mkdir(path)
+    os.removedirs(path)
+  os.mkdir(path)
 
 
 class _FakeAccumulator(object):
@@ -178,7 +178,7 @@ def testAddRunsFromDirectory(self):
     self.assertEqual(x.Runs(), {}, 'loading empty directory had no effect')
 
     path1 = join(realdir, 'path1')
-    tf.io.gfile.mkdir(path1)
+    os.mkdir(path1)
     x.AddRunsFromDirectory(realdir)
     self.assertEqual(x.Runs(), {}, 'creating empty subdirectory had no effect')
 

tensorboard/compat/tensorflow_stub/BUILD

@@ -13,6 +13,8 @@ py_library(
         "*.py",
         "compat/__init__.py",
         "compat/v1/__init__.py",
+        "io/__init__.py",
+        "io/gfile.py",
     ]),
     srcs_version = "PY2AND3",
     visibility = ["//visibility:public"],
@@ -23,3 +25,27 @@ py_library(
         "@org_pythonhosted_six",
     ],
 )
+
+py_test(
+    name = "gfile_test",
+    size = "small",
+    srcs = ["io/gfile_test.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":tensorflow_stub",
+    ],
+)
+
+py_test(
+    name = "gfile_s3_test",
+    size = "small",
+    srcs = ["io/gfile_s3_test.py"],
+    srcs_version = "PY2AND3",
+    tags = [
+        "manual",
+        "notap",
+    ],
+    deps = [
+        ":tensorflow_stub",
+    ],
+)

tensorboard/compat/tensorflow_stub/__init__.py

@@ -31,9 +31,12 @@
 from . import error_codes  # noqa
 from . import errors  # noqa
 from . import flags  # noqa
-from . import gfile  # noqa
+from . import io  # noqa
 from . import pywrap_tensorflow  # noqa
 from . import tensor_shape  # noqa
 
+# Set pywrap_tensorflow on v1 and avoid cycles on some imports
+compat.v1.pywrap_tensorflow = pywrap_tensorflow
+
 # Set a fake __version__ to help distinguish this as our own stub API.
 __version__ = 'stub'

tensorboard/compat/tensorflow_stub/compat/v1/__init__.py

@@ -17,4 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensorboard.compat.tensorflow_stub import pywrap_tensorflow  # noqa
+# Set this in tensorboard/compat/tensorflow_stub/__init__.py to eliminate
+# any cycles on import
+#
+# from tensorboard.compat.tensorflow_stub import pywrap_tensorflow  # noqa