Improve error message on non-0 rank when index file download failed (#…

…723) * a * lint
mosaicml · Jul 12, 2024 · 5f939c9 · 5f939c9
1 parent 55e83ec
commit 5f939c9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/streaming/base/stream.py b/streaming/base/stream.py
@@ -454,8 +454,8 @@ def get_shards(self, world: World, allow_unsafe_types: bool) -> List[Reader]:
                 wait_for_file_to_exist(
                     filename, TICK, self.download_timeout,
                     f'Index file {os.path.join(self.remote or "", self.split or "", basename)} ' +
-                    f'-> {filename} took too long to download. Either increase the ' +
-                    f'`download_timeout` value or check the other traceback.')
+                    f'-> {filename} took too long to download or failed to download. Either increase the '
+                    + f'`download_timeout` value or check the local rank 0 traceback.')
 
         # Load the index.
         try: