Make sure tf.io.gfile.glob is only executed once in SeqIO's list_shards

Multiple threads could trigger executing the same glob multiple times, which is slow.

PiperOrigin-RevId: 739922342

Author: tomvdw

seqio/dataset_providers.py

@@ -23,7 +23,6 @@
 import collections
 import dataclasses
 import functools
-import glob
 import inspect
 import json
 import numbers
@@ -61,8 +60,6 @@
 Feature = utils.Feature
 
 
-
-
 @dataclasses.dataclass(frozen=True)
 class ContinuousFeature(Feature):
   """A container for multi-modal output features of data providers."""
@@ -592,12 +589,6 @@ def _get_filename(info):
 
 
 
-def _list_files(pattern: str) -> Sequence[str]:
-  # Ensure that all machines observe the list of files in the same order and
-  # unique.
-  return sorted(set(tf.io.gfile.glob(pattern)))
-
-
 class FileDataSource(DataSource):
   """A `DataSource` that reads a file to provide the input dataset."""
 
@@ -714,22 +705,9 @@ def get_dataset(
         num_parallel_calls=tf.data.experimental.AUTOTUNE,
     )
 
-  @functools.lru_cache(maxsize=1024)
   def list_shards(self, split: str) -> Sequence[str]:
     filepattern = self._split_to_filepattern[split]
-    if isinstance(filepattern, str):
-      return _list_files(pattern=filepattern)
-
-    filepattern = list(filepattern)
-
-    if not any(glob.has_magic(f) for f in filepattern):
-      return filepattern
-    else:
-      assert isinstance(filepattern, Iterable)
-      ret = []
-      for f in filepattern:
-        ret.extend(_list_files(pattern=f))
-      return ret
+    return utils.list_files(filepattern)
 
 
 

seqio/dataset_providers_test.py

@@ -2121,12 +2121,12 @@ def test_repr(self):
         " block_length=16)",
     )
 
-  @mock.patch.object(dataset_providers, "_list_files")
+  @mock.patch.object(utils, "_list_files_for_glob")
   def test_file_data_source_shuffle_buffer_low(self, mock_list_files):
     mock_list_files.return_value = [f"{i}" for i in range(20)]
     fds = dataset_providers.FileDataSource(
         read_file_fn=lambda x: tf.data.Dataset.from_tensor_slices([x]),
-        split_to_filepattern={"train": "filepattern"},
+        split_to_filepattern={"train": "filepattern*"},
         file_shuffle_buffer_size=2,
     )
     for _ in range(10):
@@ -2162,12 +2162,12 @@ def test_file_data_source_shuffle_buffer_low(self, mock_list_files):
           ],
       )
 
-  @mock.patch.object(dataset_providers, "_list_files")
+  @mock.patch.object(utils, "_list_files_for_glob")
   def test_file_data_source_shuffle_buffer_full(self, mock_list_files):
     mock_list_files.return_value = [f"{i}" for i in range(20)]
     fds = dataset_providers.FileDataSource(
         read_file_fn=lambda x: tf.data.Dataset.from_tensor_slices([x]),
-        split_to_filepattern={"train": "filepattern"},
+        split_to_filepattern={"train": "filepattern*"},
         file_shuffle_buffer_size=None,
     )
     for _ in range(10):

seqio/utils.py

@@ -18,9 +18,11 @@
 import contextlib
 import dataclasses
 import functools
+import glob
 import inspect
 import os
 import re
+import threading
 import types
 from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Sequence, Tuple, Type, Union
 
@@ -1597,3 +1599,45 @@ def function_name(function) -> str:
   return ""
 
 
+# Lock to prevent multiple threads from calling tf.io.gfile.glob on the same
+# file pattern at the same time.
+_LIST_SHARD_LOCKS: dict[str, threading.Lock] = collections.defaultdict(
+    threading.Lock
+)
+
+
+def list_files(file_patterns: Union[str, Iterable[str]]) -> list[str]:
+  """Returns a sorted list of file for the given file pattern.
+
+  Note that shard patterns like `foo@2` are not expanded. Only glob patterns
+  are expanded, e.g., `foo*`.
+
+  Args:
+    file_patterns: A string or an iterable of strings, each of which is a file
+      pattern to expand.
+  """
+  if isinstance(file_patterns, str):
+    file_patterns = [file_patterns]
+
+  result = []
+  for file_pattern in file_patterns:
+    if not glob.has_magic(file_pattern):
+      result.append(file_pattern)
+      continue
+    # Make sure only one thread is calling tf.io.gfile.glob on the same file
+    # pattern at the same time. The other threads will wait for the lock to be
+    # released, and then use the cached result.
+    with _LIST_SHARD_LOCKS[file_pattern]:
+      result.extend(_list_files_for_glob(file_pattern))
+  return result
+
+
+
+
+@functools.lru_cache(maxsize=1024)
+def _list_files_for_glob(file_pattern: str) -> list[str]:
+  """Returns a sorted list of files for the given glob file pattern."""
+  file_names = set(tf.io.gfile.glob(file_pattern))
+  return sorted(file_names)
+
+

seqio/utils_test.py

@@ -1160,6 +1160,16 @@ def test_shift_right_by_one_without_default_bos(self):
     self.assertAllEqual(actual, expected)
     self.assertEqual(actual.dtype, np.int32)
 
+  @mock.patch.object(tf.io.gfile, "glob", autospec=True)
+  def test_list_files(self, mock_tf_glob):
+    mock_tf_glob.return_value = ["/bar1", "/bar2"]
+    result = utils.list_files(["/foo@2", "/bar*"])
+    self.assertEqual(result, ["/foo@2", "/bar1", "/bar2"])
+    _ = utils.list_files(["/bar*"])
+    _ = utils.list_files(["/foo@2"])
+    mock_tf_glob.assert_called_once_with("/bar*")
+
+
 
 class MixtureRateTest(test_utils.FakeTaskTest):