Lazily import legacy dataset builders

PiperOrigin-RevId: 638251592

Author: jblespiau

tensorflow_datasets/__init__.py

@@ -30,17 +30,62 @@
 
 * 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:
-  from tensorflow_datasets import rlds  # pylint: disable=g-bad-import-order
+  # 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
+  )
 
   # Public API to create and generate a dataset
   from tensorflow_datasets.public_api import *  # pylint: disable=wildcard-import
@@ -49,4 +94,12 @@
   __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,6 +19,7 @@
 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
@@ -27,7 +28,7 @@
 class LazyBuilderImport:
   """Lazy load DatasetBuilder from given name from legacy locations."""
 
-  def __init__(self, dataset_name: str):
+  def __init__(self, dataset_name: Text):
     object.__setattr__(self, "_dataset_name", dataset_name)
     object.__setattr__(self, "_dataset_cls", None)
 

tensorflow_datasets/core/registered.py

@@ -17,22 +17,17 @@
 
 import abc
 import collections
-from collections.abc import Iterator
 import contextlib
 import functools
 import importlib
 import inspect
 import os.path
-import time
-from typing import ClassVar, Type
+from typing import ClassVar, Dict, Iterator, List, Type, Text, Tuple
 
-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
 
@@ -43,7 +38,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)
@@ -56,7 +51,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".
@@ -85,70 +80,6 @@ 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."""
@@ -198,24 +129,23 @@ 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."""
-  _import_dataset_collections()
-  all_dataset_collections = list(_DATASET_COLLECTION_REGISTRY.keys())
+  all_dataset_collections = [
+      dataset_collection_name
+      for dataset_collection_name, dataset_collection_cls in _DATASET_COLLECTION_REGISTRY.items()
+  ]
   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.')
 
@@ -294,9 +224,8 @@ 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()
@@ -307,8 +236,8 @@ def list_imported_builders() -> list[str]:
 
 @functools.lru_cache(maxsize=None)
 def _get_existing_dataset_packages(
-    datasets_dir: str,
-) -> dict[str, tuple[epath.Path, str]]:
+    datasets_dir: Text,
+) -> Dict[Text, Tuple[epath.Path, Text]]:
   """Returns existing datasets.
 
   Args:
@@ -364,12 +293,7 @@ def imported_builder_cls(name: str) -> Type[RegisteredDataset]:
     raise AssertionError(f'Dataset {name} is an abstract class.')
 
   if name not in _DATASET_REGISTRY:
-    # 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.')
+    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,10 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""RLDS utils.
-
-Note that the datasets are imported lazily so they don't need to be in here.
-"""
+"""Datasets generated with RLDS."""
 
+from tensorflow_datasets.rlds import datasets
 from tensorflow_datasets.rlds import envlogger_reader
 from tensorflow_datasets.rlds import rlds_base