Wrap metadata read file operations into a retry.

PiperOrigin-RevId: 767092333

Author: pierrot0

tensorflow_datasets/core/constants.py

@@ -49,3 +49,28 @@
 # Filepath for mapping between TFDS datasets and PapersWithCode entries.
 PWC_FILENAME = 'tfds_to_pwc_links.json'
 PWC_LINKS_PATH = f'scripts/documentation/{PWC_FILENAME}'
+
+# Retry parameters. Delays are in seconds.
+TFDS_RETRY_TRIES = int(os.environ.get('TFDS_RETRY_TRIES', 3))
+TFDS_RETRY_INITIAL_DELAY = int(os.environ.get('TFDS_RETRY_INITIAL_DELAY', 1))
+# How much to multiply the delay by for each subsequent try
+TFDS_RETRY_DELAY_MULTIPLIER = int(
+    os.environ.get('TFDS_RETRY_DELAY_MULTIPLIER', 2)
+)
+# Random noise to add to the delay (random pick between 0 and noise).
+TFDS_RETRY_NOISE = float(os.environ.get('TFDS_RETRY_NOISE', 0.5))
+# If the error message contains any of these substrings, retry.
+TFDS_RETRY_MSG_SUBSTRINGS = os.environ.get(
+    'TFDS_RETRY_MSG_SUBSTRINGS',
+    (
+        'deadline_exceeded,'
+        '408 Request Timeout,'
+        '429 Too Many Requests,'
+        '500 Internal Server Error,'
+        '502 Bad Gateway,'
+        '503 Service Unavailable,'
+        '504 Gateway Timeout,'
+        '509 Bandwidth Limit Exceeded,'
+        '599 Gateway Error'
+    ),
+).split(',')

tensorflow_datasets/core/dataset_builder.py

@@ -64,6 +64,7 @@
   from tensorflow_datasets.core.utils import file_utils
   from tensorflow_datasets.core.utils import gcs_utils
   from tensorflow_datasets.core.utils import read_config as read_config_lib
+  from tensorflow_datasets.core.utils import retry
   from tensorflow_datasets.core.utils import type_utils
   # pylint: enable=g-import-not-at-top
 
@@ -290,7 +291,8 @@ def __init__(
     # Compute the base directory (for download) and dataset/version directory.
     self._data_dir_root, self._data_dir = self._build_data_dir(data_dir)
     # If the dataset info is available, use it.
-    if dataset_info.dataset_info_path(self.data_path).exists():
+    dataset_info_path = dataset_info.dataset_info_path(self.data_path)
+    if retry.retry(dataset_info_path.exists):
       self.info.read_from_directory(self._data_dir)
     else:  # Use the code version (do not restore data)
       self.info.initialize_from_bucket()
@@ -466,8 +468,8 @@ def _checksums_path(cls) -> epath.Path | None:
         # zipfile.Path does not have `.parts`. Additionally, `os.fspath`
         # will extract the file, so use `str`.
         "tensorflow_datasets" in str(new_path)
-        and legacy_path.exists()
-        and not new_path.exists()
+        and retry.retry(legacy_path.exists)
+        and not retry.retry(new_path.exists)
     ):
       return legacy_path
     else:
@@ -484,7 +486,7 @@ def url_infos(cls) -> dict[str, download.checksums.UrlInfo] | None:
     # Search for the url_info file.
     checksums_path = cls._checksums_path
     # If url_info file is found, load the urls
-    if checksums_path and checksums_path.exists():
+    if checksums_path and retry.retry(checksums_path.exists):
       return download.checksums.load_url_infos(checksums_path)
     else:
       return None
@@ -624,7 +626,7 @@ def download_and_prepare(
 
     download_config = download_config or download.DownloadConfig()
     data_path = self.data_path
-    data_exists = data_path.exists()
+    data_exists = retry.retry(data_path.exists)
 
     # Saving nondeterministic_order in the DatasetInfo for documentation.
     if download_config.nondeterministic_order:
@@ -640,7 +642,7 @@ def download_and_prepare(
             "Deleting pre-existing dataset %s (%s)", self.name, self.data_dir
         )
         data_path.rmtree()  # Delete pre-existing data.
-        data_exists = data_path.exists()
+        data_exists = retry.retry(data_path.exists)
       else:
         logging.info("Reusing dataset %s (%s)", self.name, self.data_dir)
         return
@@ -805,7 +807,7 @@ def _post_download_and_prepare_hook(self) -> None:
   def _update_dataset_info(self) -> None:
     """Updates the `dataset_info.json` file in the dataset dir."""
     info_file = self.data_path / constants.DATASET_INFO_FILENAME
-    if not info_file.exists():
+    if not retry.retry(info_file.exists):
       raise AssertionError(f"To update {info_file}, it must already exist.")
     new_info = self.info
     new_info.read_from_directory(self.data_path)
@@ -1020,7 +1022,7 @@ def as_dataset(
     self.assert_is_not_blocked()
 
     # pylint: enable=line-too-long
-    if not self.data_path.exists():
+    if not retry.retry(self.data_path.exists):
       raise AssertionError(
           "Dataset %s: could not find data in %s. Please make sure to call "
           "dataset_builder.download_and_prepare(), or pass download=True to "
@@ -1817,7 +1819,7 @@ def read_text_file(
     """Returns the text in the given file and records the lineage."""
     filename = epath.Path(filename)
     self.info.add_file_data_source_access(filename)
-    return filename.read_text(encoding=encoding)
+    return retry.retry(filename.read_text, encoding=encoding)
 
   def read_tfrecord_as_dataset(
       self,
@@ -2057,9 +2059,9 @@ def _save_default_config_name(
 def load_default_config_name(builder_dir: epath.Path) -> str | None:
   """Load `builder_cls` metadata (common to all builder configs)."""
   config_path = builder_dir / ".config" / constants.METADATA_FILENAME
-  if not config_path.exists():
+  if not retry.retry(config_path.exists):
     return None
-  data = json.loads(config_path.read_text())
+  data = json.loads(retry.retry(config_path.read_text))
   return data.get("default_config_name")
 
 

tensorflow_datasets/core/dataset_info.py

@@ -61,6 +61,7 @@
   # pylint: disable=g-import-not-at-top
   from tensorflow_datasets.core.utils import file_utils
   from tensorflow_datasets.core.utils import gcs_utils
+  from tensorflow_datasets.core.utils import retry
 
   from google.protobuf import json_format
   # pylint: enable=g-import-not-at-top
@@ -1123,7 +1124,7 @@ def read_from_json(path: epath.PathLike) -> dataset_info_pb2.DatasetInfo:
     DatasetInfoFileError: If the dataset info file cannot be read.
   """
   try:
-    json_str = epath.Path(path).read_text()
+    json_str = retry.retry(epath.Path(path).read_text)
   except OSError as e:
     raise DatasetInfoFileError(
         f"Could not read dataset info from {path}"

tensorflow_datasets/core/features/feature.py

@@ -37,6 +37,7 @@
 from tensorflow_datasets.core.utils import dtype_utils
 from tensorflow_datasets.core.utils import np_utils
 from tensorflow_datasets.core.utils import py_utils
+from tensorflow_datasets.core.utils import retry
 from tensorflow_datasets.core.utils import tf_utils
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
@@ -658,7 +659,7 @@ def from_config(cls, root_dir: str) -> FeatureConnector:
     Returns:
       The reconstructed feature instance.
     """
-    content = json.loads(make_config_path(root_dir).read_text())
+    content = json.loads(retry.retry(make_config_path(root_dir).read_text))
     feature = FeatureConnector.from_json(content)
     feature.load_metadata(root_dir, feature_name=None)
     return feature

tensorflow_datasets/core/splits.py

@@ -36,6 +36,7 @@
 from tensorflow_datasets.core import proto as proto_lib
 from tensorflow_datasets.core import units
 from tensorflow_datasets.core import utils
+from tensorflow_datasets.core.utils import retry
 from tensorflow_datasets.core.utils import shard_utils
 
 from tensorflow_metadata.proto.v0 import statistics_pb2
@@ -149,7 +150,7 @@ def get_available_shards(
         pattern = filename_template.glob_pattern(num_shards=self.num_shards)
       else:
         pattern = filename_template.sharded_filepaths_pattern(num_shards=None)
-      return list(data_dir.glob(pattern))
+      return list(retry.retry(data_dir.glob, pattern))
     else:
       raise ValueError(f'Filename template for split {self.name} is empty.')
 

tensorflow_datasets/core/utils/retry.py

@@ -0,0 +1,54 @@
+# coding=utf-8
+# Copyright 2025 The TensorFlow Datasets Authors.
+#
+# 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.
+
+"""To add retry logic to operations suceptible to transient failures."""
+
+import random
+import time
+from typing import Callable, ParamSpec, TypeVar
+
+from absl import logging
+from tensorflow_datasets.core import constants
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def retry(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
+  """Returns a decorator that retries the function."""
+  # We purposely don't use flags, as this code might be run before flags are
+  # parsed.
+  tries = constants.TFDS_RETRY_TRIES
+  delay = constants.TFDS_RETRY_INITIAL_DELAY
+  multiplier = constants.TFDS_RETRY_DELAY_MULTIPLIER
+  noise = constants.TFDS_RETRY_NOISE
+  msg_substrings = constants.TFDS_RETRY_MSG_SUBSTRINGS
+  for trial in range(1, tries + 1):
+    try:
+      return func(*args, **kwargs)
+    except BaseException as err:  # pylint: disable=broad-except
+      if trial >= tries:
+        raise err
+      msg = str(err)
+      for msg_substring in msg_substrings:
+        if msg_substring in msg:
+          break
+      else:
+        raise err
+      delay = delay + random.uniform(0, noise)
+      logging.warning("%s, retrying in %s seconds...", msg, delay)
+      time.sleep(delay)
+      delay *= multiplier