Makes file_system.glob support multiple patterns.

axlearn/common/file_system.py

@@ -12,7 +12,7 @@
 import contextlib
 import functools
 import os
-from typing import IO, TypeVar, Union
+from typing import IO, Sequence, TypeVar, Union
 from urllib.parse import urlparse
 
 import tensorflow as tf
@@ -99,8 +99,14 @@ def listdir(path: str) -> list[str]:
 
 
 @_wrap_tf_errors
-def glob(pattern: str) -> list[str]:
+def glob(pattern: Union[str, Sequence[str]]) -> list[str]:
     """Analogous to tf.io.gfile.glob."""
+    if isinstance(pattern, (list, tuple)):
+        results = set()
+        for p in pattern:
+            results.update(glob(p))
+        return list(results)
+
     parsed = urlparse(pattern)
 
     # tf.io.gfile.glob is prohibitively slow for gs paths containing many prefixes.

axlearn/common/file_system_test.py

@@ -86,6 +86,20 @@ def test_glob(self):
             ],
             fs.glob(os.path.join(self._temp_root.name, "test_dir/*_file")),
         )
+        # Globbing multiple patterns.
+        self.assertCountEqual(
+            [
+                os.path.join(self._temp_root.name, "test_file"),
+                os.path.join(self._temp_root.name, "test_dir"),
+                os.path.join(self._temp_root.name, "other_file"),
+            ],
+            fs.glob(
+                [
+                    os.path.join(self._temp_root.name, "test_*"),
+                    os.path.join(self._temp_root.name, "*_file"),
+                ]
+            ),
+        )
         # Globbing a non-existent path is different from listdir.
         self.assertEqual([], fs.glob(os.path.join(self._temp_root.name, "fake_dir")))
         self.assertEqual([], fs.glob(os.path.join(self._temp_root.name, "fake_dir/*_file")))