tensorbored 2.21.0rc1769983804__py3-none-any.whl

tensorbored/backend/event_processing/directory_watcher.py

@@ -0,0 +1,272 @@
+# Copyright 2015 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.
+# ==============================================================================
+
+"""Contains the implementation for the DirectoryWatcher class."""
+
+import bisect
+
+from tensorbored.backend.event_processing import io_wrapper
+from tensorbored.compat import tf
+from tensorbored.util import io_util
+from tensorbored.util import tb_logging
+
+logger = tb_logging.get_logger()
+
+
+class DirectoryWatcher:
+    """A DirectoryWatcher wraps a loader to load from a sequence of paths.
+
+    A loader reads a path and produces some kind of values as an iterator. A
+    DirectoryWatcher takes a directory, a factory for loaders, and optionally a
+    path filter and watches all the paths inside that directory.
+
+    This class is only valid under the assumption that only one path will be
+    written to by the data source at a time and that once the source stops writing
+    to a path, it will start writing to a new path that's lexicographically
+    greater and never come back. It uses some heuristics to check whether this is
+    true based on tracking changes to the files' sizes, but the check can have
+    false negatives. However, it should have no false positives.
+    """
+
+    def __init__(self, directory, loader_factory, path_filter=lambda x: True):
+        """Constructs a new DirectoryWatcher.
+
+        Args:
+          directory: The directory to load files from.
+          loader_factory: A factory for creating loaders. The factory should take a
+            path and return an object that has a Load method returning an
+            iterator that will yield all events that have not been yielded yet.
+          path_filter: If specified, only paths matching this filter are loaded.
+
+        Raises:
+          ValueError: If path_provider or loader_factory are None.
+        """
+        if directory is None:
+            raise ValueError("A directory is required")
+        if loader_factory is None:
+            raise ValueError("A loader factory is required")
+        self._directory = directory
+        self._path = None
+        self._loader_factory = loader_factory
+        self._loader = None
+        self._path_filter = path_filter
+        self._ooo_writes_detected = False
+        # The file size for each file at the time it was finalized.
+        self._finalized_sizes = {}
+
+    def Load(self):
+        """Loads new values.
+
+        The watcher will load from one path at a time; as soon as that path stops
+        yielding events, it will move on to the next path. We assume that old paths
+        are never modified after a newer path has been written. As a result, Load()
+        can be called multiple times in a row without losing events that have not
+        been yielded yet. In other words, we guarantee that every event will be
+        yielded exactly once.
+
+        Yields:
+          All values that have not been yielded yet.
+
+        Raises:
+          DirectoryDeletedError: If the directory has been permanently deleted
+            (as opposed to being temporarily unavailable).
+        """
+        try:
+            for event in self._LoadInternal():
+                yield event
+        except tf.errors.OpError:
+            if not tf.io.gfile.exists(self._directory):
+                raise DirectoryDeletedError(
+                    "Directory %s has been permanently deleted"
+                    % self._directory
+                )
+
+    def _LoadInternal(self):
+        """Internal implementation of Load().
+
+        The only difference between this and Load() is that the latter will throw
+        DirectoryDeletedError on I/O errors if it thinks that the directory has been
+        permanently deleted.
+
+        Yields:
+          All values that have not been yielded yet.
+        """
+
+        # If the loader exists, check it for a value.
+        if not self._loader:
+            self._InitializeLoader()
+
+        # If it still doesn't exist, there is no data
+        if not self._loader:
+            return
+
+        while True:
+            # Yield all the new events in the path we're currently loading from.
+            for event in self._loader.Load():
+                yield event
+
+            next_path = self._GetNextPath()
+            if not next_path:
+                logger.info("No path found after %s", self._path)
+                # Current path is empty and there are no new paths, so we're done.
+                return
+
+            # There's a new path, so check to make sure there weren't any events
+            # written between when we finished reading the current path and when we
+            # checked for the new one. The sequence of events might look something
+            # like this:
+            #
+            # 1. Event #1 written to path #1.
+            # 2. We check for events and yield event #1 from path #1
+            # 3. We check for events and see that there are no more events in path #1.
+            # 4. Event #2 is written to path #1.
+            # 5. Event #3 is written to path #2.
+            # 6. We check for a new path and see that path #2 exists.
+            #
+            # Without this loop, we would miss event #2. We're also guaranteed by the
+            # loader contract that no more events will be written to path #1 after
+            # events start being written to path #2, so we don't have to worry about
+            # that.
+            for event in self._loader.Load():
+                yield event
+
+            logger.info(
+                "Directory watcher advancing from %s to %s",
+                self._path,
+                next_path,
+            )
+
+            # Advance to the next path and start over.
+            self._SetPath(next_path)
+
+    # The number of paths before the current one to check for out of order writes.
+    _OOO_WRITE_CHECK_COUNT = 20
+
+    def OutOfOrderWritesDetected(self):
+        """Returns whether any out-of-order writes have been detected.
+
+        Out-of-order writes are only checked as part of the Load() iterator. Once an
+        out-of-order write is detected, this function will always return true.
+
+        Note that out-of-order write detection is not performed on GCS paths, so
+        this function will always return false.
+
+        Returns:
+          Whether any out-of-order write has ever been detected by this watcher.
+        """
+        return self._ooo_writes_detected
+
+    def _InitializeLoader(self):
+        path = self._GetNextPath()
+        if path:
+            self._SetPath(path)
+
+    def _SetPath(self, path):
+        """Sets the current path to watch for new events.
+
+        This also records the size of the old path, if any. If the size can't be
+        found, an error is logged.
+
+        Args:
+          path: The full path of the file to watch.
+        """
+        old_path = self._path
+        if old_path and not io_util.IsCloudPath(old_path):
+            try:
+                # We're done with the path, so store its size.
+                size = tf.io.gfile.stat(old_path).length
+                logger.debug("Setting latest size of %s to %d", old_path, size)
+                self._finalized_sizes[old_path] = size
+            except tf.errors.OpError as e:
+                logger.error("Unable to get size of %s: %s", old_path, e)
+
+        self._path = path
+        self._loader = self._loader_factory(path)
+
+    def _GetNextPath(self):
+        """Gets the next path to load from.
+
+        This function also does the checking for out-of-order writes as it iterates
+        through the paths.
+
+        Returns:
+          The next path to load events from, or None if there are no more paths.
+        """
+        paths = sorted(
+            path
+            for path in io_wrapper.ListDirectoryAbsolute(self._directory)
+            if self._path_filter(path)
+        )
+        if not paths:
+            return None
+
+        if self._path is None:
+            return paths[0]
+
+        # 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_util.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
+            )
+            for path in paths[ooo_check_start:current_path_index]:
+                if self._HasOOOWrite(path):
+                    self._ooo_writes_detected = True
+                    break
+
+        next_paths = list(
+            path for path in paths if self._path is None or path > self._path
+        )
+        if next_paths:
+            return min(next_paths)
+        else:
+            return None
+
+    def _HasOOOWrite(self, path):
+        """Returns whether the path has had an out-of-order write."""
+        # Check the sizes of each path before the current one.
+        size = tf.io.gfile.stat(path).length
+        old_size = self._finalized_sizes.get(path, None)
+        if size != old_size:
+            if old_size is None:
+                logger.error(
+                    "File %s created after file %s even though it's "
+                    "lexicographically earlier",
+                    path,
+                    self._path,
+                )
+            else:
+                logger.error(
+                    "File %s updated even though the current file is %s",
+                    path,
+                    self._path,
+                )
+            return True
+        else:
+            return False
+
+
+class DirectoryDeletedError(Exception):
+    """Thrown by Load() when the directory is *permanently* gone.
+
+    We distinguish this from temporary errors so that other code can
+    decide to drop all of our data only when a directory has been
+    intentionally deleted, as opposed to due to transient filesystem
+    errors.
+    """
+
+    pass