Repository: NVIDIA-Merlin/NVTabular

Issue:
[BUG] OOM error while trying to train on Outbrain dataset with Tensorflow
**Describe the bug**
A V100 16GB GPU now runs out memory when trying to run the [Outbrain TF training notebook](https://github.com/NVIDIA/NVTabular/blob/main/examples/advanced-ops-outbrain/03-Training-with-TF.ipynb). This is a new problem that has cropped up in the past few days, so may be related to recent TF dataloader changes.

**Steps/Code to reproduce bug**
Run the Outbrain example notebooks with a single V100 16GB GPU.

**Expected behavior**
Example should complete successfully without running out of memory

**Environment details (please complete the following information):**
 TBD

**Additional context**

Stack trace:
```
--------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
<ipython-input-18-b1725d7df4b2> in <module>
      5     experimental_run_tf_function=False
      6 )
----> 7 history = wide_and_deep_model.fit(train_dataset_tf, epochs=1)
/opt/conda/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)
   1048          training_utils.RespectCompiledTrainableState(self):
   1049       # Creates a `tf.data.Dataset` and handles batch and epoch iteration.
-> 1050       data_handler = data_adapter.DataHandler(
   1051           x=x,
   1052           y=y,
/opt/conda/lib/python3.8/site-packages/tensorflow/python/keras/engine/data_adapter.py in __init__(self, x, y, sample_weight, batch_size, steps_per_epoch, initial_epoch, epochs, shuffle, class_weight, max_queue_size, workers, use_multiprocessing, model, steps_per_execution)
   1098 
   1099     adapter_cls = select_data_adapter(x, y)
-> 1100     self._adapter = adapter_cls(
   1101         x,
   1102         y,
/opt/conda/lib/python3.8/site-packages/tensorflow/python/keras/engine/data_adapter.py in __init__(self, x, y, sample_weights, shuffle, workers, use_multiprocessing, max_queue_size, model, **kwargs)
    896                        "`keras.utils.Sequence` as input.")
    897 
--> 898     self._size = len(x)
    899     self._shuffle_sequence = shuffle
    900     self._keras_sequence = x
/nvtabular/nvtabular/loader/tensorflow.py in __len__(self)
    237         # TODO: what's a better way to do this inheritance
    238         # of the appropriate methods? A Metaclass?
--> 239         return DataLoader.__len__(self)
    240 
    241     def __getitem__(self, idx):
/nvtabular/nvtabular/loader/backend.py in __len__(self)
    203 
    204     def __len__(self):
--> 205         return _num_steps(len(self._buff), self.batch_size)
    206 
    207     @property
/nvtabular/nvtabular/loader/backend.py in __len__(self)
     61 
     62     def __len__(self):
---> 63         return len(self.itr)
     64 
     65     @property
/nvtabular/nvtabular/io/dataset.py in __len__(self)
    766 
    767     def __len__(self):
--> 768         return len(self._ddf.partitions[self.indices])
    769 
    770     def __iter__(self):
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/dataframe/core.py in __len__(self)
   3654             return super().__len__()
   3655         else:
-> 3656             return len(s)
   3657 
   3658     def __contains__(self, key):
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/dataframe/core.py in __len__(self)
    555 
    556     def __len__(self):
--> 557         return self.reduction(
    558             len, np.sum, token="len", meta=int, split_every=False
    559         ).compute()
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/base.py in compute(self, **kwargs)
    281         dask.base.compute
    282         """
--> 283         (result,) = compute(self, traverse=False, **kwargs)
    284         return result
    285 
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/base.py in compute(*args, **kwargs)
    563         postcomputes.append(x.__dask_postcompute__())
    564 
--> 565     results = schedule(dsk, keys, **kwargs)
    566     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    567 
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/threaded.py in get(dsk, result, cache, num_workers, pool, **kwargs)
     74                 pools[thread][num_workers] = pool
     75 
---> 76     results = get_async(
     77         pool.apply_async,
     78         len(pool._pool),
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/local.py in get_async(apply_async, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, **kwargs)
    485                         _execute_task(task, data)  # Re-execute locally
    486                     else:
--> 487                         raise_exception(exc, tb)
    488                 res, worker_id = loads(res_info)
    489                 state["cache"][key] = res
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/local.py in reraise(exc, tb)
    315     if exc.__traceback__ is not tb:
    316         raise exc.with_traceback(tb)
--> 317     raise exc
    318 
    319 
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/local.py in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
    220     try:
    221         task, data = loads(task_info)
--> 222         result = _execute_task(task, data)
    223         id = get_id()
    224         result = dumps((result, id))
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    119         # temporaries by their reference count and can execute certain
    120         # operations in-place.
--> 121         return func(*(_execute_task(a, cache) for a in args))
    122     elif not ishashable(arg):
    123         return arg
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/core.py in <genexpr>(.0)
    119         # temporaries by their reference count and can execute certain
    120         # operations in-place.
--> 121         return func(*(_execute_task(a, cache) for a in args))
    122     elif not ishashable(arg):
    123         return arg
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    113     """
    114     if isinstance(arg, list):
--> 115         return [_execute_task(a, cache) for a in arg]
    116     elif istask(arg):
    117         func, args = arg[0], arg[1:]
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/core.py in <listcomp>(.0)
    113     """
    114     if isinstance(arg, list):
--> 115         return [_execute_task(a, cache) for a in arg]
    116     elif istask(arg):
    117         func, args = arg[0], arg[1:]
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    119         # temporaries by their reference count and can execute certain
    120         # operations in-place.
--> 121         return func(*(_execute_task(a, cache) for a in args))
    122     elif not ishashable(arg):
    123         return arg
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/core.py in <genexpr>(.0)
    119         # temporaries by their reference count and can execute certain
    120         # operations in-place.
--> 121         return func(*(_execute_task(a, cache) for a in args))
    122     elif not ishashable(arg):
    123         return arg
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    119         # temporaries by their reference count and can execute certain
    120         # operations in-place.
--> 121         return func(*(_execute_task(a, cache) for a in args))
    122     elif not ishashable(arg):
    123         return arg
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/dataframe/io/parquet/core.py in read_parquet_part(fs, func, meta, part, columns, index, kwargs)
    381 
    382     if isinstance(part, list):
--> 383         dfs = [
    384             func(fs, rg, columns.copy(), index, **toolz.merge(kwargs, kw))
    385             for (rg, kw) in part
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask/dataframe/io/parquet/core.py in <listcomp>(.0)
    382     if isinstance(part, list):
    383         dfs = [
--> 384             func(fs, rg, columns.copy(), index, **toolz.merge(kwargs, kw))
    385             for (rg, kw) in part
    386         ]
/opt/conda/envs/rapids/lib/python3.8/site-packages/dask_cudf/io/parquet.py in read_partition(fs, piece, columns, index, categories, partitions, **kwargs)
     59         strings_to_cats = kwargs.get("strings_to_categorical", False)
     60         if cudf.utils.ioutils._is_local_filesystem(fs):
---> 61             df = cudf.read_parquet(
     62                 path,
     63                 engine="cudf",
/opt/conda/envs/rapids/lib/python3.8/site-packages/cudf/io/parquet.py in read_parquet(filepath_or_buffer, engine, columns, filters, row_groups, skiprows, num_rows, strings_to_categorical, use_pandas_metadata, *args, **kwargs)
    249 
    250     if engine == "cudf":
--> 251         return libparquet.read_parquet(
    252             filepaths_or_buffers,
    253             columns=columns,
cudf/_lib/parquet.pyx in cudf._lib.parquet.read_parquet()
cudf/_lib/parquet.pyx in cudf._lib.parquet.read_parquet()
MemoryError: std::bad_alloc: CUDA error at: /opt/conda/envs/rapids/include/rmm/mr/device/cuda_memory_resource.hpp:69: cudaErrorMemoryAllocation out of memory
```



Discussion:
We think this is related to changes to the way that the `__len__` of a Dataset is computed in order to account for the possibility that the dataset may be sharded across multiple processes. A first pass partial fix might be to check if the Dataset holds less partition indices than are available in the Dataset, and if not, then use the earlier length computation code.
@rjzamora says:
> self._ddf.partitions[self.indices] will produce a dask DataFrame (a subset of the original ddf), and the len will probably alias to len(<ddf>.index).  So, it will probably try to read from the files to produce a index for each partition.  This can probably be optimized to use the metadata for a parquet file, but I don't think this is being done right now.

> we can probably provide our own optimization until it gets into Dask

> We want to use the metadata in the same way we do in [num_rows](https://github.com/NVIDIA/NVTabular/blob/d7a9373504f3e680d81bf899fb97bd1427c3a99f/nvtabular/io/parquet.py#L149), but we need to map the row-group metadata onto the expected ddf partitions.  So, we will need to use the kind of tricky logic used in [_file_partition_map](https://github.com/NVIDIA/NVTabular/blob/d7a9373504f3e680d81bf899fb97bd1427c3a99f/nvtabular/io/parquet.py#L115)

> Note that the file_partition_map code is generating a map between files and partitions, but something similar could also map row-groups to partition lengths
@karlhigley and @jperez999 I could train the nb after pulling the recent changes you folks did, and reducing the buffer size on a 16 GB V100.
Okay, cool! So the partial fix to avoid the issue in single-GPU mode seems to have worked, now we just need to make the changes @rjzamora suggested to handle it for large datasets in multi-GPU mode.
I have been scoping this out today, and there doesn't seem to be a "trivial" solution here.  I have worked out a way to do this but there are two primary problems:

(1) It requires an API change to add `Dataset.shuffle` and to remove the `shuffle=` option from `Dataset.to_ddf`.  The purpose of this is to make it easier to track partition reordering after a shuffle is performed.  If the shuffle returns a Dataset object (rather than a ddf), it is easier to record the mapping between parquet row-group metadata and output ddf partitions.

(2) There is no guarantee that the metadata-informed partition sizes will be the same after pre-processing (i.e. general NVTabular Ops).  So, we would still need to pull data into memory to calculate the length for an end-to-end pre-processing + training workflow.

Im curious, why do we need to know the length of the dataset anyway?

Resolving patch:
diff --git a/nvtabular/io/dataset.py b/nvtabular/io/dataset.py
--- a/nvtabular/io/dataset.py
+++ b/nvtabular/io/dataset.py
@@ -362,6 +362,10 @@ def to_ddf(self, columns=None, shuffle=False, seed=None):
     def file_partition_map(self):
         return self.engine._file_partition_map
 
+    @property
+    def partition_lens(self):
+        return self.engine._partition_lens
+
     def to_cpu(self):
         warnings.warn(
             "Changing an NVTabular Dataset to CPU mode."
@@ -477,7 +481,7 @@ def shuffle_by_keys(self, keys, hive_data=None, npartitions=None):
         # Fall back to dask.dataframe algorithm
         return Dataset(ddf.shuffle(keys, npartitions=npartitions))
 
-    def to_iter(self, columns=None, indices=None, shuffle=False, seed=None):
+    def to_iter(self, columns=None, indices=None, shuffle=False, seed=None, use_file_metadata=None):
         """Convert `Dataset` object to a `cudf.DataFrame` iterator.
 
         Note that this method will use `to_ddf` to produce a
@@ -502,12 +506,39 @@ def to_iter(self, columns=None, indices=None, shuffle=False, seed=None):
             The random seed to use if `shuffle=True`.  If nothing
             is specified, the current system time will be used by the
             `random` std library.
+        use_file_metadata : bool; Optional
+            Whether to allow the returned ``DataFrameIter`` object to
+            use file metadata from the ``base_dataset`` to estimate
+            the row-count. By default, the file-metadata
+            optimization will only be used if the current Dataset is
+            backed by a file-based engine. Otherwise, it is possible
+            that an intermediate transform has modified the row-count.
         """
         if isinstance(columns, str):
             columns = [columns]
 
+        # Try to extract the row-size metadata
+        # if we are not shuffling
+        partition_lens_meta = None
+        if not shuffle and use_file_metadata is not False:
+            # We are allowed to use file metadata to calculate
+            # partition sizes.  If `use_file_metadata` is None,
+            # we only use metadata if `self` is backed by a
+            # file-based engine (like "parquet").  Otherwise,
+            # we cannot be "sure" that the metadata row-count
+            # is correct.
+            try:
+                if use_file_metadata:
+                    partition_lens_meta = self.base_dataset.partition_lens
+                else:
+                    partition_lens_meta = self.partition_lens
+            except AttributeError:
+                pass
+
         return DataFrameIter(
-            self.to_ddf(columns=columns, shuffle=shuffle, seed=seed), indices=indices
+            self.to_ddf(columns=columns, shuffle=shuffle, seed=seed),
+            indices=indices,
+            partition_lens=partition_lens_meta,
         )
 
     def to_parquet(
@@ -869,12 +900,19 @@ def _set_dtypes(chunk, dtypes):
 
 
 class DataFrameIter:
-    def __init__(self, ddf, columns=None, indices=None):
+    def __init__(self, ddf, columns=None, indices=None, partition_lens=None):
         self.indices = indices if isinstance(indices, list) else range(ddf.npartitions)
         self._ddf = ddf
         self.columns = columns
+        self.partition_lens = partition_lens
 
     def __len__(self):
+        if self.partition_lens:
+            # Use metadata-based partition-size information
+            # if/when it is available.  Note that this metadata
+            # will not be correct if rows where added or dropped
+            # after IO (within Ops).
+            return sum(self.partition_lens[i] for i in self.indices)
         if len(self.indices) < self._ddf.npartitions:
             return len(self._ddf.partitions[self.indices])
         return len(self._ddf)
diff --git a/nvtabular/io/parquet.py b/nvtabular/io/parquet.py
--- a/nvtabular/io/parquet.py
+++ b/nvtabular/io/parquet.py
@@ -67,6 +67,8 @@ def __init__(
         cpu=False,
     ):
         super().__init__(paths, part_size, cpu=cpu, storage_options=storage_options)
+        self._pp_map = None
+        self._pp_nrows = None
         if row_groups_per_part is None:
             path0 = self._dataset.pieces[0].path
             with self.fs.open(path0, "rb") as f0:
@@ -111,8 +113,40 @@ def _dataset(self):
         return dataset
 
     @property
-    @functools.lru_cache(1)
     def _file_partition_map(self):
+        if self._pp_map is None:
+            self._process_parquet_metadata()
+        return self._pp_map
+
+    @property
+    def _partition_lens(self):
+        if self._pp_nrows is None:
+            self._process_parquet_metadata()
+        return self._pp_nrows
+
+    @property
+    def num_rows(self):
+        # TODO: Avoid parsing metadata once upstream dask
+        # can get the length efficiently (in all practical cases)
+        return sum(self._partition_lens)
+
+    def _process_parquet_metadata(self):
+        # Utility shared by `_file_partition_map` and `_partition_lens`
+        # to collect useful information from the parquet metadata
+
+        _pp_nrows = []
+
+        def _update_partition_lens(part_count, md, num_row_groups, rg_offset=None):
+            # Helper function to calculate the row count for each
+            # output partition (and add it to `_pp_nrows`)
+            rg_offset = rg_offset or 0
+            for rg_i in range(0, part_count, self.row_groups_per_part):
+                rg_f = min(rg_i + self.row_groups_per_part, num_row_groups)
+                _pp_nrows.append(
+                    sum([md.row_group(rg + rg_offset).num_rows for rg in range(rg_i, rg_f)])
+                )
+            return
+
         dataset = self._dataset
         if dataset.metadata:
             # We have a metadata file.
@@ -124,41 +158,29 @@ def _file_partition_map(self):
 
             # Convert the per-file row-group count to the
             # file-to-partition mapping
-            ind = 0
+            ind, rg = 0, 0
             _pp_map = defaultdict(list)
             for fn, num_row_groups in _path_row_groups.items():
                 part_count = math.ceil(num_row_groups / self.row_groups_per_part)
                 _pp_map[fn] = np.arange(ind, ind + part_count)
+                _update_partition_lens(part_count, dataset.metadata, num_row_groups, rg_offset=rg)
                 ind += part_count
-
+                rg += num_row_groups
         else:
             # No metadata file. Construct file-to-partition map manually
             ind = 0
             _pp_map = {}
             for piece in dataset.pieces:
-                num_row_groups = piece.get_metadata().num_row_groups
+                md = piece.get_metadata()
+                num_row_groups = md.num_row_groups
                 part_count = math.ceil(num_row_groups / self.row_groups_per_part)
                 fn = piece.path.split(self.fs.sep)[-1]
                 _pp_map[fn] = np.arange(ind, ind + part_count)
+                _update_partition_lens(part_count, md, num_row_groups)
                 ind += part_count
 
-        return _pp_map
-
-    @property
-    @functools.lru_cache(1)
-    def num_rows(self):
-        # TODO: Avoid parsing metadata here if we can confirm upstream dask
-        # can get the length efficiently (in all practical cases)
-        dataset = self._dataset
-        if dataset.metadata:
-            # We have a metadata file
-            return dataset.metadata.num_rows
-        else:
-            # Sum up row-group sizes manually
-            num_rows = 0
-            for piece in dataset.pieces:
-                num_rows += piece.get_metadata().num_rows
-            return num_rows
+        self._pp_map = _pp_map
+        self._pp_nrows = _pp_nrows
 
     def to_ddf(self, columns=None, cpu=None):
 


Test patch:
diff --git a/tests/unit/test_io.py b/tests/unit/test_io.py
--- a/tests/unit/test_io.py
+++ b/tests/unit/test_io.py
@@ -73,15 +73,17 @@ def test_dask_dataset_itr(tmpdir, datasets, engine, gpu_memory_frac):
     else:
         columns = mycols_csv
 
-    dd = nvtabular.io.Dataset(
+    size = 0
+    ds = nvtabular.io.Dataset(
         paths[0], engine=engine, part_mem_fraction=gpu_memory_frac, dtypes=dtypes
     )
-    size = 0
-    for chunk in dd.to_iter(columns=columns):
+    my_iter = ds.to_iter(columns=columns)
+    for chunk in my_iter:
         size += chunk.shape[0]
         assert chunk["id"].dtype == np.int32
 
     assert size == df1.shape[0]
+    assert len(my_iter) == size
 
 
 @pytest.mark.parametrize("engine", ["csv", "parquet", "csv-no-header"])
@@ -603,6 +605,40 @@ def test_dataset_conversion(tmpdir, cpu, preserve_files):
     assert not glob.glob(os.path.join(pq_path, "*.parquet"))
 
 
+@pytest.mark.parametrize("use_file_metadata", [True, None])
+@pytest.mark.parametrize("shuffle", [True, False])
+def test_parquet_iterator_len(tmpdir, shuffle, use_file_metadata):
+
+    ddf1 = dask.datasets.timeseries(
+        start="2000-01-01",
+        end="2000-01-6",
+        freq="600s",
+        partition_freq="1d",
+        id_lam=10,
+        seed=42,
+    ).shuffle("id")
+
+    # Write to parquet dataset
+    ddf1.to_parquet(str(tmpdir))
+
+    # Initialize Dataset
+    ds = nvt.Dataset(str(tmpdir), engine="parquet")
+
+    # Convert ds -> ds2
+    ds2 = nvt.Dataset(ds.to_ddf())
+
+    # Check that iterator lengths match the partition lengths
+    ddf2 = ds2.to_ddf(shuffle=shuffle, seed=42)
+    for i in range(ddf2.npartitions):
+        _iter = ds2.to_iter(
+            shuffle=shuffle,
+            seed=42,
+            indices=[i],
+            use_file_metadata=use_file_metadata,
+        )
+        assert len(ddf2.partitions[i]) == len(_iter)
+
+
 @pytest.mark.parametrize("cpu", [True, False])
 def test_hive_partitioned_data(tmpdir, cpu):