Experimental tool to disable SeqIO registries.

seqio/experimental.py

@@ -32,6 +32,69 @@
 TaskRegistry = dataset_providers.TaskRegistry
 
 
+def _enfore_empty_registries():
+  """Enforces that the seqio TaskRegistry and MixtureRegistry are empty."""
+  non_empty_registries = []
+  if TaskRegistry.names():
+    logging.error('TaskRegistry has %s tasks', len(TaskRegistry.names()))
+    for task_name in TaskRegistry.names():
+      logging.error('Registered Task: %s', task_name)
+    non_empty_registries.append('TaskRegistry')
+  if MixtureRegistry.names():
+    logging.error(
+        'MixtureRegistry has %s mixtures', len(MixtureRegistry.names())
+    )
+    for mixture_name in MixtureRegistry.names():
+      logging.error('Registered Mixture: %s', mixture_name)
+    non_empty_registries.append('MixtureRegistry')
+  if non_empty_registries:
+    non_empty_registries = ', '.join(non_empty_registries)
+    raise ValueError(
+        f'The {non_empty_registries} is non-empty. Please invoke'
+        ' `disable_registry()` before any Task/Mixtures are registered.'
+    )
+
+
+def _get_argument(name, position, args, kwargs, default_val=''):
+  if name in kwargs:
+    return kwargs.get(name)
+  if len(args) > position:
+    return args[position]
+  return default_val
+
+
+def _no_op_task_registry_add(*args, **kwargs):
+  logging.info(
+      'No-op Task register call: %s', _get_argument('name', 0, args, kwargs)
+  )
+
+
+def _no_op_mixture_registry_add(*args, **kwargs):
+  logging.info(
+      'No-op Mixture register call: %s', _get_argument('name', 0, args, kwargs)
+  )
+
+
+def _no_op_task_registry_get(*args, **kwargs):
+  name = _get_argument('name', 0, args, kwargs)
+  raise ValueError(f'Disabled TaskRegistry.get call: {name}')
+
+
+def _no_op_mixture_registry_get(*args, **kwargs):
+  name = _get_argument('name', 0, args, kwargs)
+  raise ValueError(f'Disabled MixtureRegistry.get call: {name}')
+
+
+def disable_registry():
+  _enfore_empty_registries()
+  dataset_providers.TaskRegistry.add = _no_op_task_registry_add
+  dataset_providers.TaskRegistry.add_provider = _no_op_task_registry_add
+  dataset_providers.TaskRegistry.get = _no_op_task_registry_get
+  dataset_providers.MixtureRegistry.add = _no_op_mixture_registry_add
+  dataset_providers.MixtureRegistry.add_provider = _no_op_mixture_registry_add
+  dataset_providers.MixtureRegistry.get = _no_op_mixture_registry_get
+
+
 def _get_fully_cached_name(
     original_name: str, sequence_length: Mapping[str, int]
 ) -> str:

seqio/experimental_test.py

@@ -14,6 +14,7 @@
 
 """Tests for seqio.preprocessors."""
 
+import contextlib
 from absl.testing import absltest
 from seqio import dataset_providers
 from seqio import experimental
@@ -868,5 +869,119 @@ def test_add_sentinel_task(self):
       )
 
 
+class DisableRegistryTest(absltest.TestCase):
+
+  @contextlib.contextmanager
+  def _restore_registry_after_test(self):
+    og_tr_add = TaskRegistry.add
+    og_tr_add_p = TaskRegistry.add_provider
+    og_tr_get = TaskRegistry.get
+    og_mr_add = MixtureRegistry.add
+    og_mr_add_p = MixtureRegistry.add_provider
+    og_mr_get = MixtureRegistry.get
+    try:
+      yield
+    finally:
+      TaskRegistry.add = og_tr_add
+      TaskRegistry.add_provider = og_tr_add_p
+      TaskRegistry.get = og_tr_get
+      MixtureRegistry.add = og_mr_add
+      MixtureRegistry.add_provider = og_mr_add_p
+      MixtureRegistry.get = og_mr_get
+
+  def test_enforce_task_registry_empty(self):
+    TaskRegistry.reset()
+    MixtureRegistry.reset()
+    function_source = dataset_providers.FunctionDataSource(
+        dataset_fn=test_utils.get_fake_dataset, splits=['train', 'validation']
+    )
+    TaskRegistry.add(
+        'dummy_task',
+        source=function_source,
+        output_features=test_utils.FakeTaskTest.DEFAULT_OUTPUT_FEATURES,
+    )
+    with self.assertRaisesRegex(ValueError, 'The TaskRegistry is non-empty.*'):
+      with self._restore_registry_after_test():
+        experimental.disable_registry()
+
+  def test_enforce_both_registries_empty(self):
+    TaskRegistry.reset()
+    MixtureRegistry.reset()
+    function_source = dataset_providers.FunctionDataSource(
+        dataset_fn=test_utils.get_fake_dataset, splits=['train', 'validation']
+    )
+    TaskRegistry.add(
+        'dummy_task',
+        source=function_source,
+        output_features=test_utils.FakeTaskTest.DEFAULT_OUTPUT_FEATURES,
+    )
+    MixtureRegistry.add('dummy_mixture', ['dummy_task'], default_rate=1.0)
+    with self.assertRaisesRegex(
+        ValueError, 'The TaskRegistry, MixtureRegistry is non-empty.*'
+    ):
+      with self._restore_registry_after_test():
+        experimental.disable_registry()
+
+  def test_task_registry_add_no_op(self):
+    TaskRegistry.reset()
+    MixtureRegistry.reset()
+    with self._restore_registry_after_test():
+      experimental.disable_registry()
+      function_source = dataset_providers.FunctionDataSource(
+          dataset_fn=test_utils.get_fake_dataset, splits=['train', 'validation']
+      )
+      TaskRegistry.add(
+          'dummy_task',
+          source=function_source,
+          output_features=test_utils.FakeTaskTest.DEFAULT_OUTPUT_FEATURES,
+      )
+      task = dataset_providers.Task(
+          'dummy_task2',
+          source=function_source,
+          output_features=test_utils.FakeTaskTest.DEFAULT_OUTPUT_FEATURES,
+      )
+      TaskRegistry.add_provider('dummy_task2', task)
+      self.assertEmpty(TaskRegistry.names())
+
+  def test_task_registry_get_error(self):
+    TaskRegistry.reset()
+    MixtureRegistry.reset()
+    with self._restore_registry_after_test():
+      experimental.disable_registry()
+      with self.assertRaisesRegex(
+          ValueError, 'Disabled TaskRegistry.get call.*'
+      ):
+        TaskRegistry.get('dummy_task')
+
+  def test_mixture_registry_add_no_op(self):
+    TaskRegistry.reset()
+    MixtureRegistry.reset()
+    with self._restore_registry_after_test():
+      experimental.disable_registry()
+      function_source = dataset_providers.FunctionDataSource(
+          dataset_fn=test_utils.get_fake_dataset, splits=['train', 'validation']
+      )
+      task = dataset_providers.Task(
+          'dummy_task',
+          source=function_source,
+          output_features=test_utils.FakeTaskTest.DEFAULT_OUTPUT_FEATURES,
+      )
+      TaskRegistry.add_provider('dummy_task2', task)
+      MixtureRegistry.add('dummy_mixture', [task], default_rate=1.0)
+      mixture = dataset_providers.Mixture('dummy_mixture2', [[task, 1.0]])
+      MixtureRegistry.add_provider('dummy_mixture2', mixture)
+      self.assertEmpty(TaskRegistry.names())
+      self.assertEmpty(MixtureRegistry.names())
+
+  def test_mixture_registry_get_error(self):
+    TaskRegistry.reset()
+    MixtureRegistry.reset()
+    with self.assertRaisesRegex(
+        ValueError, 'Disabled MixtureRegistry.get call.*'
+    ):
+      with self._restore_registry_after_test():
+        experimental.disable_registry()
+        MixtureRegistry.get('dummy_mixture')
+
 if __name__ == '__main__':
   absltest.main()