Cache GetAllDatasetParams.

Every GetDatasetParams(self, dataset) call calls GetAllDatasetParams(), which
instantiates all dataset and tasks. In MMASR+intrainer case, there are 80
dataset and corresponding Task_XXX, which are 2B params model. It takes
more than 10 mins.

The initial GetAllDatasetParams() API is unfortunate. It supposed to be
a generator or keep only factory methods, instead of actual instantiation.

PiperOrigin-RevId: 492103447

Author: ds-hwang

lingvo/core/base_model_params.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 """BaseModelParams class definition."""
-
+import functools
 
 from lingvo import datasets
 from lingvo.core import base_input_generator
@@ -35,6 +35,13 @@ class DatasetError(Exception):
   pass
 
 
+# Cache the GetAllDatasetParams() results, which may include many Task_XXX.
+# Note: datasets inspect _BaseModelParams, so it can not have a cached method.
+@functools.lru_cache(maxsize=8)
+def GetCachedAllDatasetParams(obj):
+  return obj.GetAllDatasetParams()
+
+
 class _BaseModelParams:
   """Base class for storing model Params for a single experiment."""
 
@@ -45,12 +52,10 @@ def GetAllDatasetParams(self):
     be treated as dataset specifications.
 
     Returns:
-      A dict of {dataset_name: dataset_params}.
-
-    Raises:
-      GetAllDatasetParamsNotImplementedError: by default.
+      A dict of {dataset_name: dataset_params}. If not implemented, returns
+      None.
     """
-    raise datasets.GetAllDatasetParamsNotImplementedError(type(self))
+    return None
 
   def GetDatasetParams(self, dataset):
     """Convenience function that returns the param for the given dataset name.
@@ -65,16 +70,13 @@ def GetDatasetParams(self, dataset):
     Raises:
       DatasetError: if there is not a `${dataset}` method defined under `cls`.
     """
-    try:
-      all_datasets = self.GetAllDatasetParams()
-      if dataset not in all_datasets:
+    all_datasets = GetCachedAllDatasetParams(self)
+    if all_datasets is not None:
+      if dataset not in all_datasets.keys():
         # When `GetAllDatasetParams` is defined, all public methods are ignored.
         raise DatasetError(f'Dataset {dataset} not found; '
                            f'available datasets are: {all_datasets.keys()}')
       return all_datasets.get(dataset)
-    except datasets.GetAllDatasetParamsNotImplementedError:
-      # Fall through the legacy path.
-      pass
 
     try:
       f = getattr(self, dataset)

lingvo/datasets.py

@@ -31,10 +31,6 @@ class DatasetFunctionError(TypeError):
   pass
 
 
-class GetAllDatasetParamsNotImplementedError(NotImplementedError):
-  pass
-
-
 def GetDatasets(cls: Any, warn_on_error: bool = True) -> List[str]:
   """Returns the list of dataset functions (e.g., Train, Dev, ...).
 
@@ -55,7 +51,6 @@ def GetDatasets(cls: Any, warn_on_error: bool = True) -> List[str]:
     DatasetFunctionError: if the cls contains public methods that cannot be used
       as datasets, and warn_on_error is False.
   """
-
   mdl_params = None
   if inspect.isclass(cls):
     try:
@@ -65,13 +60,11 @@ def GetDatasets(cls: Any, warn_on_error: bool = True) -> List[str]:
   else:
     mdl_params = cls
 
-  if mdl_params:
-    try:
-      all_datasets = mdl_params.GetAllDatasetParams()
+  if mdl_params and hasattr(mdl_params, 'GetAllDatasetParams'):
+    all_datasets = mdl_params.GetAllDatasetParams()
+    if all_datasets is not None:
       # When `GetAllDatasetParams` is defined, all public methods are ignored.
       return sorted(list(all_datasets.keys()))
-    except GetAllDatasetParamsNotImplementedError:
-      pass
 
   datasets = []
   for name, _ in inspect.getmembers(cls, inspect.isroutine):

lingvo/model_registry.py

@@ -243,6 +243,22 @@ def GetClass(cls, class_key):
                         class_key)
     return all_params[class_key]
 
+  @classmethod
+  @functools.lru_cache(maxsize=8)
+  def GetCachedObject(cls, class_key):
+    """Returns a ModelParams class object with the given `class_key`.
+
+    To reuse the same object for given `class_key`, it uses
+    functools.lru_cache annotation.
+
+    Args:
+      class_key: string key of the ModelParams subclass to return.
+
+    Returns:
+      A instance of `~.base_model_params._BaseModelParams`.
+    """
+    return cls.GetClass(class_key)()
+
   @classmethod
   def GetParamsFromModelParamsObject(cls, model_params_obj, dataset_name):
     """Returns a `Params` object for given _BaseModelParams instance.
@@ -281,7 +297,7 @@ def GetParamsFromModelParamsObject(cls, model_params_obj, dataset_name):
     return cfg
 
   @classmethod
-  def GetParams(cls, class_key, dataset_name):
+  def GetParams(cls, class_key, dataset_name, cache=True):
     """Constructs a `Params` object for given model and dataset, obeying flags.
 
     In case of default model, params may be updated based on the flags
@@ -290,12 +306,17 @@ def GetParams(cls, class_key, dataset_name):
     Args:
       class_key: String class key (i.e. `image.mnist.LeNet5`).
       dataset_name: Method to generate dataset params (i.e. 'Test').
+      cache: Whether to cache the params given class_key and dataset_name.
 
     Returns:
       Full `~.hyperparams.Params` for the model class.
     """
-    model_params_cls = cls.GetClass(class_key)
-    return cls.GetParamsFromModelParamsObject(model_params_cls(), dataset_name)
+    if cache:
+      # Reuse the cached object to reused the cached obj.GetAllDatasetParams().
+      model_params_obj = cls.GetCachedObject(class_key)
+    else:
+      model_params_obj = cls.GetClass(class_key)()
+    return cls.GetParamsFromModelParamsObject(model_params_obj, dataset_name)
 
   @classmethod
   def GetProgramScheduleFromModelParamsObject(cls, model_params_obj):

lingvo/models_test_helper.py

@@ -129,9 +129,8 @@ def _ValidateEMA(self, name, mdl):
   def _testOneModelParams(self, registry, name):
     with tf.Graph().as_default():
       model_params = registry.GetClass(name)()
-      try:
-        all_datasets = model_params.GetAllDatasetParams()
-      except datasets.GetAllDatasetParamsNotImplementedError:
+      all_datasets = model_params.GetAllDatasetParams()
+      if all_datasets is None:
         all_datasets = {}
         for dataset_name in datasets.GetDatasets(model_params):
           try: