`test_download_and_prepare_as_dataset` fails with `apache-beam==2.65.0`

[GaetanLepage]

Short description

Context: Bumping apache-beam in nixpkgs: NixOS/nixpkgs#408537

Since updating apache-beam to 2.65.0, the following tests have started to fail:

=========================== short test summary info ============================
FAILED tensorflow_datasets/datasets/robonet/robonet_dataset_builder_test.py::Robonet64Test::test_download_and_prepare_as_dataset - TypeError: unsupported operand type(s) for +: 'Mock' and 'str'
FAILED tensorflow_datasets/structured/web_graph/web_graph_test.py::WebGraphTest::test_download_and_prepare_as_dataset - RuntimeError: Unable to pickle fn CallableWrapperDoFn(<function Map.<locals...
FAILED tensorflow_datasets/rl_unplugged/rlu_rwrl/rlu_rwrl_test.py::RluRwrlTest::test_download_and_prepare_as_dataset - TypeError: unsupported operand type(s) for +: 'Mock' and 'str'
FAILED tensorflow_datasets/robotics/dataset_importer_builder_test.py::MockRDSDatasetTestRlds::test_download_and_prepare_as_dataset - RuntimeError: Unable to pickle fn CallableWrapperDoFn(<function DatasetImpo...
FAILED tensorflow_datasets/structured/covid19/covid19_test.py::Covid19Test::test_download_and_prepare_as_dataset - RuntimeError: Unable to pickle fn CallableWrapperDoFn(<function Covid19._ge...
FAILED tensorflow_datasets/video/youtube_vis/youtube_vis_test.py::YoutubeVisTest::test_download_and_prepare_as_dataset - TypeError: unsupported operand type(s) for +: 'Mock' and 'str'
FAILED tensorflow_datasets/vision_language/grounded_scan/grounded_scan_test.py::GroundedScanTest::test_download_and_prepare_as_dataset - TypeError: unsupported operand type(s) for +: 'Mock' and 'str'
FAILED tensorflow_datasets/vision_language/wit_kaggle/wit_kaggle_test.py::WitKaggleTestTrain::test_download_and_prepare_as_dataset - RuntimeError: Unable to pickle fn CallableWrapperDoFn(<function WitKaggle._...
FAILED tensorflow_datasets/vision_language/wit_kaggle/wit_kaggle_test.py::WitKaggleTestTest::test_download_and_prepare_as_dataset - RuntimeError: Unable to pickle fn CallableWrapperDoFn(<function WitKaggle._...
ERROR tensorflow_datasets/core/data_sources/base_test.py
= 9 failed, 3828 passed, 588 skipped, 229 warnings, 1 error in 161.74s (0:02:41) =

Environment information

• Operating System: NixOS

• Python version: 3.12.10

• tensorflow-datasets/tfds-nightly version: 4.9.8

• tensorflow/tf-nightly version: 2.19.0

• Does the issue still exists with the last tfds-nightly package (pip install --upgrade tfds-nightly) ?

Reproduction instructions

Run pytest.

Link to logs

____________ WitKaggleTestTest.test_download_and_prepare_as_dataset ____________
[gw6] linux -- Python 3.12.10 /nix/store/8w718rm43x7z73xhw9d6vh8s4snrq67h-python3-3.12.10/bin/python3.12

self = <ParDo(PTransform) label=[ParDo(CallableWrapperDoFn)] at 0x7ffe4c178b30>
fn = CallableWrapperDoFn(<function WitKaggle._generate_examples.<locals>._read_resnet_rows at 0x7ffe54357ec0>)
args = (), kwargs = {}

    def __init__(self, fn, *args, **kwargs):
      # type: (WithTypeHints, *Any, **Any) -> None
      if isinstance(fn, type) and issubclass(fn, WithTypeHints):
        # Don't treat Fn class objects as callables.
        raise ValueError('Use %s() not %s.' % (fn.__name__, fn.__name__))
      self.fn = self.make_fn(fn, bool(args or kwargs))
      # Now that we figure out the label, initialize the super-class.
      super().__init__()

      if (any(isinstance(v, pvalue.PCollection) for v in args) or
          any(isinstance(v, pvalue.PCollection) for v in kwargs.values())):
        raise error.SideInputError(
            'PCollection used directly as side input argument. Specify '
            'AsIter(pcollection) or AsSingleton(pcollection) to indicate how the '
            'PCollection is to be used.')
      self.args, self.kwargs, self.side_inputs = util.remove_objects_from_args(
          args, kwargs, pvalue.AsSideInput)
      self.raw_side_inputs = args, kwargs

      # Prevent name collisions with fns of the form '<function <lambda> at ...>'
      self._cached_fn = self.fn

      # Ensure fn and side inputs are picklable for remote execution.
      try:
>       self.fn = pickler.loads(pickler.dumps(self.fn))

/nix/store/bpfl79aspdjrzvmyhipzavdfafbvan9x-python3.12-apache-beam-2.65.0/lib/python3.12/site-packages/apache_beam/transforms/ptransform.py:872:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/nix/store/bpfl79aspdjrzvmyhipzavdfafbvan9x-python3.12-apache-beam-2.65.0/lib/python3.12/site-packages/apache_beam/internal/pickler.py:57: in loads
    return desired_pickle_lib.loads(
/nix/store/bpfl79aspdjrzvmyhipzavdfafbvan9x-python3.12-apache-beam-2.65.0/lib/python3.12/site-packages/apache_beam/internal/cloudpickle_pickler.py:168: in loads
    unpickled = cloudpickle.loads(s)
/nix/store/0dfvvls5mhlgi0cqqmjc4y8gr8qj0lkr-python3.12-etils-1.12.2/lib/python3.12/site-packages/etils/epy/lazy_imports_utils.py:114: in __getattr__
    if name in self._submodules:
/nix/store/0dfvvls5mhlgi0cqqmjc4y8gr8qj0lkr-python3.12-etils-1.12.2/lib/python3.12/site-packages/etils/epy/lazy_imports_utils.py:114: in __getattr__
    if name in self._submodules:
/nix/store/0dfvvls5mhlgi0cqqmjc4y8gr8qj0lkr-python3.12-etils-1.12.2/lib/python3.12/site-packages/etils/epy/lazy_imports_utils.py:114: in __getattr__
    if name in self._submodules:
E   RecursionError: maximum recursion depth exceeded
!!! Recursion detected (same locals & position)

Expected behavior
Tests succeed.

[fineguy]

Thanks for reporting this issue! Please downgrade apache-beam, we have done the same in TFDS: 2034928

[GaetanLepage]

Thanks for reporting this issue! Please downgrade apache-beam, we have done the same in TFDS: 2034928

Thank you for the tip. Actually, we cannot have duplicated versions of each package in nixpkgs. Hence, we ended up bumping apache-beam nonetheless and skipping the failing test on TFDS.

Not a big deal. We will wait for the next TFDS release to un-skip the test.