Check that a dataset's version (and config) are not blocked before running download_and_prepare.

PiperOrigin-RevId: 652794653

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/dataset_builder.py

@@ -557,6 +557,16 @@ def is_prepared(self) -> bool:
     """Returns whether this dataset is already downloaded and prepared."""
     return self.data_path.exists()
 
+  def assert_is_not_blocked(self) -> None:
+    """Checks that the dataset is not blocked."""
+    config_name = self.builder_config.name if self.builder_config else None
+    if blocked_versions := self.blocked_versions:
+      is_blocked = blocked_versions.is_blocked(
+          version=self.version, config=config_name
+      )
+      if is_blocked.result:
+        raise utils.DatasetVariantBlockedError(is_blocked.blocked_msg)
+
   @tfds_logging.download_and_prepare()
   def download_and_prepare(
       self,
@@ -581,7 +591,11 @@ def download_and_prepare(
     Raises:
       IOError: if there is not enough disk space available.
       RuntimeError: when the config cannot be found.
+      DatasetBlockedError: if the given version, or combination of version and
+        config, has been marked as blocked in the builder's BLOCKED_VERSIONS.
     """
+    self.assert_is_not_blocked()
+
 
     download_config = download_config or download.DownloadConfig()
     data_path = self.data_path

tensorflow_datasets/core/dataset_builder_test.py

@@ -94,6 +94,14 @@ def _generate_examples(self, range_):
       yield i, {"x": x}
 
 
+class DummyDatasetWithBlockedVersions(DummyDatasetWithConfigs):
+
+  BLOCKED_VERSIONS = utils.BlockedVersions(
+      versions={"0.0.1": "Version 0.0.1 is blocked"},
+      configs={"0.0.2": {"plus2": "plus2 is blocked for version 0.0.2"}},
+  )
+
+
 class DummyDatasetWithDefaultConfig(DummyDatasetWithConfigs):
   DEFAULT_BUILDER_CONFIG_NAME = "plus2"
 
@@ -393,6 +401,42 @@ def test_builder_configs_configs_with_multiple_versions(self):
         set(DummyDatasetWithVersionedConfigs.builder_configs.keys()),
     )
 
+  def test_assert_is_not_blocked(self):
+    with testing.tmp_dir(self.get_temp_dir()) as tmp_dir:
+      tmp_dir = epath.Path(tmp_dir)
+      builder_1 = DummyDatasetWithBlockedVersions(
+          config="plus1", version="0.0.1", data_dir=tmp_dir
+      )
+      builder_2 = DummyDatasetWithBlockedVersions(
+          config="plus2", version="0.0.2", data_dir=tmp_dir
+      )
+      builder_3 = DummyDatasetWithBlockedVersions(
+          config="plus2", version="0.0.1", data_dir=tmp_dir
+      )
+      not_blocked_builder = DummyDatasetWithConfigs(
+          config="plus1", version="0.0.1", data_dir=tmp_dir
+      )
+
+      assert builder_1.blocked_versions is not None
+      assert builder_2.blocked_versions is not None
+      assert builder_3.blocked_versions is not None
+      assert not_blocked_builder.blocked_versions is None
+
+      with pytest.raises(
+          utils.DatasetVariantBlockedError, match="Version 0.0.1 is blocked"
+      ):
+        assert builder_1.assert_is_not_blocked()
+      with pytest.raises(
+          utils.DatasetVariantBlockedError,
+          match="plus2 is blocked for version 0.0.2",
+      ):
+        assert builder_2.assert_is_not_blocked()
+      with pytest.raises(
+          utils.DatasetVariantBlockedError, match="Version 0.0.1 is blocked"
+      ):
+        assert builder_3.assert_is_not_blocked()
+      assert not_blocked_builder.assert_is_not_blocked() is None
+
   def test_versioned_configs(self):
     with testing.tmp_dir(self.get_temp_dir()) as tmp_dir:
       tmp_dir = epath.Path(tmp_dir)

tensorflow_datasets/core/utils/__init__.py

@@ -84,5 +84,6 @@
 from tensorflow_datasets.core.utils.tqdm_utils import TqdmStream
 from tensorflow_datasets.core.utils.type_utils import *
 from tensorflow_datasets.core.utils.version import BlockedVersions
+from tensorflow_datasets.core.utils.version import DatasetVariantBlockedError
 from tensorflow_datasets.core.utils.version import Experiment
 from tensorflow_datasets.core.utils.version import Version

tensorflow_datasets/core/utils/version.py

@@ -32,12 +32,28 @@
 _VERSION_RESOLVED_REG = re.compile(_VERSION_TMPL.format(v=_NO_LEADING_ZEROS))
 
 
+class DatasetVariantBlockedError(ValueError):
+  """Exception raised when a blocked version and/or config is requested."""
+
+
 # A dictionary of blocked versions or configs.
 # The key is a version or config string, the value is a short sentence
 # explaining why that version or config should not be used (or None).
 BlockedWithMsg = dict[str, str | None]
 
 
+@dataclasses.dataclass(frozen=True)
+class IsBlocked:
+  """Class to store information about a version or config being blocked.
+
+  Also contains an optional message explaining why the version or config is
+  blocked.
+  """
+
+  result: bool
+  blocked_msg: str | None = None
+
+
 @dataclasses.dataclass(frozen=True)
 class BlockedVersions:
   """Holds information on versions and configs that should not be used.
@@ -56,15 +72,31 @@ class BlockedVersions:
 
   def is_blocked(
       self, version: str | Version, config: str | None = None
-  ) -> bool:
-    """Checks whether a version or config is blocked."""
+  ) -> IsBlocked:
+    """Checks whether a version or config is blocked.
+
+    Args:
+      version: The version to check.
+      config: The config to check. If None, the version is checked.
+
+    Returns:
+      An IsBlocked object. If IsBlocked.result is True, IsBlocked.blocked_msg
+      contains the message explaining why the version or config is blocked, if
+      it exists, or a default message otherwise.
+    """
     if isinstance(version, Version):
       version = str(version)
     if version in self.versions:
-      return True
+      blocked_msg = self.versions[version] or f"Version {version} is blocked."
+      return IsBlocked(True, blocked_msg)
     if config is not None and version in self.configs:
-      return config in self.configs[version]
-    return False
+      if config in self.configs[version]:
+        blocked_msg = (
+            self.configs[version][config]
+            or f"Config {config} for version {version} is blocked."
+        )
+        return IsBlocked(True, blocked_msg)
+    return IsBlocked(False)
 
 
 class Experiment(enum.Enum):

tensorflow_datasets/core/utils/version_test.py

@@ -151,28 +151,35 @@ def test_str_to_version():
 
 
 @pytest.mark.parametrize(
-    'blocked_version, blocked_config, expected',
+    'blocked_version, blocked_config, expected_res, expected_msg',
     [
-        ('1.2.3', None, True),
-        ('1.0.0', None, False),
-        ('1.2.3', 'non_existing_config', True),
-        ('1.0.0', 'config_1', True),
-        ('1.0.0', 'config_2', True),
-        ('1.0.0', 'config_1', True),
-        ('1.1.0', 'config_2', False),
+        ('1.2.3', None, True, 'Version 1.2.3 is blocked.'),
+        ('1.0.0', None, False, None),
+        ('1.2.3', 'non_existing_config', True, 'Version 1.2.3 is blocked.'),
+        ('1.0.0', 'config_1', True, 'blocked_config'),
+        (
+            '1.0.0',
+            'config_2',
+            True,
+            'Config config_2 for version 1.0.0 is blocked.',
+        ),
+        ('1.1.0', 'config_1', True, 'blocked config in version 1.1.0'),
+        ('1.1.0', 'config_2', False, None),
     ],
 )
-def test_is_blocked(blocked_version, blocked_config, expected):
+def test_is_blocked(
+    blocked_version, blocked_config, expected_res, expected_msg
+):
   blocked_versions = version.BlockedVersions(
       versions={'1.2.3': None},
       configs={
           '1.0.0': {'config_1': 'blocked_config', 'config_2': None},
           '1.1.0': {'config_1': 'blocked config in version 1.1.0'},
       },
   )
-  assert (
-      blocked_versions.is_blocked(blocked_version, blocked_config) == expected
-  )
+  is_blocked = blocked_versions.is_blocked(blocked_version, blocked_config)
+  assert is_blocked.result == expected_res
+  assert is_blocked.blocked_msg == expected_msg
 
 
 if __name__ == '__main__':