Lazily import legacy dataset builders

PiperOrigin-RevId: 638201042

Author: tomvdw

tensorflow_datasets/__init__.py

@@ -30,62 +30,17 @@
 
 * These API docs
 * [Available datasets](https://www.tensorflow.org/datasets/catalog/overview)
-* [Colab tutorial](https://colab.research.google.com/github/tensorflow/datasets/blob/master/docs/overview.ipynb)
+* [Colab
+tutorial](https://colab.research.google.com/github/tensorflow/datasets/blob/master/docs/overview.ipynb)
 * [Add a dataset](https://www.tensorflow.org/datasets/add_dataset)
 """
 # pylint: enable=line-too-long
 # pylint: disable=g-import-not-at-top,g-bad-import-order,wrong-import-position,unused-import
 
-import time
-
-_TIMESTAMP_IMPORT_STARTS = time.time()
 from absl import logging
-import tensorflow_datasets.core.logging as _tfds_logging
-from tensorflow_datasets.core.logging import call_metadata as _call_metadata
-
-_metadata = _call_metadata.CallMetadata()
-_metadata.start_time_micros = int(_TIMESTAMP_IMPORT_STARTS * 1e6)
-_import_time_ms_dataset_builders = 0
 
 try:
-  # Imports for registration
-  _before_dataset_imports = time.time()
-  from tensorflow_datasets import dataset_collections
-
-  # pytype: disable=import-error
-  # For builds that don't include all dataset builders, we don't want to fail on
-  # import errors of dataset builders.
-  try:
-    from tensorflow_datasets import audio
-    from tensorflow_datasets import graphs
-    from tensorflow_datasets import image
-    from tensorflow_datasets import image_classification
-    from tensorflow_datasets import object_detection
-    from tensorflow_datasets import nearest_neighbors
-    from tensorflow_datasets import question_answering
-    from tensorflow_datasets import d4rl
-    from tensorflow_datasets import ranking
-    from tensorflow_datasets import recommendation
-    from tensorflow_datasets import rl_unplugged
-    from tensorflow_datasets import rlds
-    from tensorflow_datasets import robotics
-    from tensorflow_datasets import robomimic
-    from tensorflow_datasets import structured
-    from tensorflow_datasets import summarization
-    from tensorflow_datasets import text
-    from tensorflow_datasets import text_simplification
-    from tensorflow_datasets import time_series
-    from tensorflow_datasets import translate
-    from tensorflow_datasets import video
-    from tensorflow_datasets import vision_language
-
-  except ImportError:
-    pass
-  # pytype: enable=import-error
-
-  _import_time_ms_dataset_builders = int(
-      (time.time() - _before_dataset_imports) * 1000
-  )
+  from tensorflow_datasets import rlds  # pylint: disable=g-bad-import-order
 
   # Public API to create and generate a dataset
   from tensorflow_datasets.public_api import *  # pylint: disable=wildcard-import
@@ -94,12 +49,4 @@
   __all__ = public_api.__all__
 
 except Exception as exception:  # pylint: disable=broad-except
-  _metadata.mark_error()
   logging.exception(exception)
-finally:
-  _metadata.mark_end()
-  _tfds_logging.tfds_import(
-      metadata=_metadata,
-      import_time_ms_tensorflow=0,
-      import_time_ms_dataset_builders=_import_time_ms_dataset_builders,
-  )

tensorflow_datasets/core/lazy_builder_import.py

@@ -19,7 +19,6 @@
 defined under the "datasets/" directory and the legacy locations have been
 deprecated (ie: a new major release).
 """
-from typing import Text
 
 from absl import logging
 from tensorflow_datasets.core import registered
@@ -28,11 +27,12 @@
 class LazyBuilderImport:
   """Lazy load DatasetBuilder from given name from legacy locations."""
 
-  def __init__(self, dataset_name: Text):
+  def __init__(self, dataset_name: str):
     object.__setattr__(self, "_dataset_name", dataset_name)
     object.__setattr__(self, "_dataset_cls", None)
 
   def _get_builder_cls(self):
+    """Returns the DatasetBuilder class."""
     cls = object.__getattribute__(self, "_dataset_cls")
     if not cls:
       builder_name = object.__getattribute__(self, "_dataset_name")

tensorflow_datasets/core/registered.py

@@ -17,17 +17,22 @@
 
 import abc
 import collections
+from collections.abc import Iterator
 import contextlib
 import functools
 import importlib
 import inspect
 import os.path
-from typing import ClassVar, Dict, Iterator, List, Type, Text, Tuple
+import time
+from typing import ClassVar, Type
 
+from absl import logging
 from etils import epath
 from tensorflow_datasets.core import constants
 from tensorflow_datasets.core import naming
 from tensorflow_datasets.core import visibility
+import tensorflow_datasets.core.logging as _tfds_logging
+from tensorflow_datasets.core.logging import call_metadata as _call_metadata
 from tensorflow_datasets.core.utils import py_utils
 from tensorflow_datasets.core.utils import resource_utils
 
@@ -38,7 +43,7 @@
 # <str snake_cased_name, abstract DatasetBuilder subclass>
 _ABSTRACT_DATASET_REGISTRY = {}
 
-# Keep track of Dict[str (module name), List[DatasetBuilder]]
+# Keep track of dict[str (module name), list[DatasetBuilder]]
 # This is directly accessed by `tfds.community.builder_cls_from_module` when
 # importing community packages.
 _MODULE_TO_DATASETS = collections.defaultdict(list)
@@ -51,7 +56,7 @@
 # <str snake_cased_name, abstract DatasetCollectionBuilder subclass>
 _ABSTRACT_DATASET_COLLECTION_REGISTRY = {}
 
-# Keep track of Dict[str (module name), List[DatasetCollectionBuilder]]
+# Keep track of dict[str (module name), list[DatasetCollectionBuilder]]
 _MODULE_TO_DATASET_COLLECTIONS = collections.defaultdict(list)
 
 # eg for dataset "foo": "tensorflow_datasets.datasets.foo.foo_dataset_builder".
@@ -80,6 +85,70 @@ def skip_registration() -> Iterator[None]:
     _skip_registration = False
 
 
+@functools.cache
+def _import_legacy_builders() -> None:
+  """Imports legacy builders."""
+  modules_to_import = [
+      'audio',
+      'graphs',
+      'image',
+      'image_classification',
+      'object_detection',
+      'nearest_neighbors',
+      'question_answering',
+      'd4rl',
+      'ranking',
+      'recommendation',
+      'rl_unplugged',
+      'rlds.datasets',
+      'robotics',
+      'robomimic',
+      'structured',
+      'summarization',
+      'text',
+      'text_simplification',
+      'time_series',
+      'translate',
+      'video',
+      'vision_language',
+  ]
+
+  before_dataset_imports = time.time()
+  metadata = _call_metadata.CallMetadata()
+  metadata.start_time_micros = int(before_dataset_imports * 1e6)
+  try:
+    # For builds that don't include all dataset builders, we don't want to fail
+    # on import errors of dataset builders.
+    try:
+      for module in modules_to_import:
+        importlib.import_module(f'tensorflow_datasets.{module}')
+    except (ImportError, ModuleNotFoundError):
+      pass
+
+  except Exception as exception:  # pylint: disable=broad-except
+    metadata.mark_error()
+    logging.exception(exception)
+  finally:
+    import_time_ms_dataset_builders = int(
+        (time.time() - before_dataset_imports) * 1000
+    )
+    metadata.mark_end()
+    _tfds_logging.tfds_import(
+        metadata=metadata,
+        import_time_ms_tensorflow=0,
+        import_time_ms_dataset_builders=import_time_ms_dataset_builders,
+    )
+
+
+@functools.cache
+def _import_dataset_collections() -> None:
+  """Imports dataset collections."""
+  try:
+    importlib.import_module('tensorflow_datasets.dataset_collections')
+  except (ImportError, ModuleNotFoundError):
+    pass
+
+
 # The implementation of this class follows closely RegisteredDataset.
 class RegisteredDatasetCollection(abc.ABC):
   """Subclasses will be registered and given a `name` property."""
@@ -129,23 +198,24 @@ def __init_subclass__(cls, skip_registration=False, **kwargs):  # pylint: disabl
       _DATASET_COLLECTION_REGISTRY[cls.name] = cls
 
 
-def list_imported_dataset_collections() -> List[str]:
+def list_imported_dataset_collections() -> list[str]:
   """Returns the string names of all `tfds.core.DatasetCollection`s."""
-  all_dataset_collections = [
-      dataset_collection_name
-      for dataset_collection_name, dataset_collection_cls in _DATASET_COLLECTION_REGISTRY.items()
-  ]
+  _import_dataset_collections()
+  all_dataset_collections = list(_DATASET_COLLECTION_REGISTRY.keys())
   return sorted(all_dataset_collections)
 
 
 def is_dataset_collection(name: str) -> bool:
+  _import_dataset_collections()
   return name in _DATASET_COLLECTION_REGISTRY
 
 
 def imported_dataset_collection_cls(
     name: str,
 ) -> Type[RegisteredDatasetCollection]:
   """Returns the Registered dataset class."""
+  _import_dataset_collections()
+
   if name in _ABSTRACT_DATASET_COLLECTION_REGISTRY:
     raise AssertionError(f'DatasetCollection {name} is an abstract class.')
 
@@ -224,8 +294,9 @@ def _is_builder_available(builder_cls: Type[RegisteredDataset]) -> bool:
   return visibility.DatasetType.TFDS_PUBLIC.is_available()
 
 
-def list_imported_builders() -> List[str]:
+def list_imported_builders() -> list[str]:
   """Returns the string names of all `tfds.core.DatasetBuilder`s."""
+  _import_legacy_builders()
   all_builders = [
       builder_name
       for builder_name, builder_cls in _DATASET_REGISTRY.items()
@@ -236,8 +307,8 @@ def list_imported_builders() -> List[str]:
 
 @functools.lru_cache(maxsize=None)
 def _get_existing_dataset_packages(
-    datasets_dir: Text,
-) -> Dict[Text, Tuple[epath.Path, Text]]:
+    datasets_dir: str,
+) -> dict[str, tuple[epath.Path, str]]:
   """Returns existing datasets.
 
   Args:
@@ -293,7 +364,12 @@ def imported_builder_cls(name: str) -> Type[RegisteredDataset]:
     raise AssertionError(f'Dataset {name} is an abstract class.')
 
   if name not in _DATASET_REGISTRY:
-    raise DatasetNotFoundError(f'Dataset {name} not found.')
+    # Dataset not found in the registry, try to import legacy builders.
+    # Dataset builders are imported lazily to avoid slowing down the startup
+    # of the binary.
+    _import_legacy_builders()
+    if name not in _DATASET_REGISTRY:
+      raise DatasetNotFoundError(f'Dataset {name} not found.')
 
   builder_cls = _DATASET_REGISTRY[name]
   if not _is_builder_available(builder_cls):

tensorflow_datasets/rlds/__init__.py

@@ -13,8 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Datasets generated with RLDS."""
+"""RLDS utils.
+
+Note that the datasets are imported lazily so they don't need to be in here.
+"""
 
-from tensorflow_datasets.rlds import datasets
 from tensorflow_datasets.rlds import envlogger_reader
 from tensorflow_datasets.rlds import rlds_base