Repository: NVIDIA-Merlin/NVTabular

Issue:
[BUG] excessive memory usage when reading parquet file from GCS bucket
**Describe the bug**
Cannot read 11GB parquet file from GCS with 4x 16GB T4 GPU, while same file reads just fine from local storage (persistent disk).

We suspect the parquet file got read into GPU as a whole, the expected behavior should be reading row groups and not the entire file.

**Steps/Code to reproduce bug**
Wget 2 python script here: https://gist.github.com/mengdong/8526f2765707be492b1b697c3fb8c687 and run python3 analyze_dataset_op.py

root@dongm-debug-fsspec:/scripts# python3 analyze_dataset_op.py
INFO:root:Creating Dask cluster.
INFO:numba.cuda.cudadrv.driver:init
distributed.preloading - INFO - Import preload module: dask_cuda.initialize
distributed.preloading - INFO - Import preload module: dask_cuda.initialize
distributed.preloading - INFO - Import preload module: dask_cuda.initialize
distributed.preloading - INFO - Import preload module: dask_cuda.initialize
INFO:root:Creating transformation workflow.
INFO:root:Creating dataset to be analysed.
distributed.worker - WARNING - Worker is at 82% memory usage. Pausing worker.  Process memory: 24.44 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker.html#memtrim for more information. -- Unmanaged memory: 24.44 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker.html#memtrim for more information. -- Unmanaged memory: 24.51 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker.html#memtrim for more information. -- Unmanaged memory: 24.65 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker.html#memtrim for more information. -- Unmanaged memory: 24.68 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker.html#memtrim for more information. -- Unmanaged memory: 24.85 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker.html#memtrim for more information. -- Unmanaged memory: 25.03 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Worker is at 4% memory usage. Resuming worker. Process memory: 1.33 GiB -- Worker memory limit: 29.51 GiB
INFO:root:Starting workflow fitting for fsspec_parquet split.
distributed.worker - WARNING - Worker is at 85% memory usage. Pausing worker.  Process memory: 25.30 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Worker is at 6% memory usage. Resuming worker. Process memory: 2.02 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - ERROR - [Errno 2] No such file or directory: '/scripts/dask-worker-space/storage/%28%27split-c6e5ab25962e13a1755fd762b62bbfe9%27%2C%205%2C%2011%2C%200%29'
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 2623, in release_key
    if key in self.data:
  File "/usr/lib/python3.8/_collections_abc.py", line 666, in __contains__
    self[key]
  File "/usr/local/lib/python3.8/dist-packages/dask_cuda/device_host_file.py", line 246, in __getitem__
    return self.host_buffer[key]
  File "/usr/local/lib/python3.8/dist-packages/zict-2.0.0-py3.8.egg/zict/buffer.py", line 78, in __getitem__
    return self.slow_to_fast(key)
  File "/usr/local/lib/python3.8/dist-packages/zict-2.0.0-py3.8.egg/zict/buffer.py", line 65, in slow_to_fast
    value = self.slow[key]
  File "/usr/local/lib/python3.8/dist-packages/zict-2.0.0-py3.8.egg/zict/func.py", line 38, in __getitem__
    return self.load(self.d[key])
  File "/usr/local/lib/python3.8/dist-packages/zict-2.0.0-py3.8.egg/zict/file.py", line 75, in __getitem__
    with open(os.path.join(self.directory, _safe_key(key)), "rb") as f:
FileNotFoundError: [Errno 2] No such file or directory: '/scripts/dask-worker-space/storage/%28%27split-c6e5ab25962e13a1755fd762b62bbfe9%27%2C%205%2C%2011%2C%200%29'
distributed.worker - ERROR -
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 1857, in transition_flight_memory
    self.batched_stream.send({"op": "add-keys", "keys": [ts.key]})
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/batched.py", line 136, in send
    raise CommClosedError()
distributed.comm.core.CommClosedError
tornado.application - ERROR - Exception in callback functools.partial(<bound method IOLoop._discard_future_result of <tornado.platform.asyncio.AsyncIOLoop object at 0x7ff429eb9b20>>, <Task finished name='Task-1992' coro=<Worker.gather_dep() done, defined at /usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py:2267> exception=CommClosedError()>)
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/tornado/ioloop.py", line 741, in _run_callback
    ret = callback()
  File "/usr/local/lib/python3.8/dist-packages/tornado/ioloop.py", line 765, in _discard_future_result
    future.result()
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 2411, in gather_dep
    self.transition(ts, "memory", value=data[d])
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 1694, in transition
    state = func(ts, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 1857, in transition_flight_memory
    self.batched_stream.send({"op": "add-keys", "keys": [ts.key]})
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/batched.py", line 136, in send
    raise CommClosedError()
distributed.comm.core.CommClosedError
tornado.application - ERROR - Exception in callback functools.partial(<bound method IOLoop._discard_future_result of <tornado.platform.asyncio.AsyncIOLoop object at 0x7ff429eb9b20>>, <Task finished name='Task-1995' coro=<Worker.heartbeat() done, defined at /usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py:959> exception=RuntimeError('Set changed size during iteration')>)
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/tornado/ioloop.py", line 741, in _run_callback
    ret = callback()
  File "/usr/local/lib/python3.8/dist-packages/tornado/ioloop.py", line 765, in _discard_future_result
    future.result()
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 991, in heartbeat
    await self._register_with_scheduler()
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 905, in _register_with_scheduler
    types={k: typename(v) for k, v in self.data.items()},
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 905, in <dictcomp>
    types={k: typename(v) for k, v in self.data.items()},
  File "/usr/lib/python3.8/_collections_abc.py", line 743, in __iter__
    for key in self._mapping:
RuntimeError: Set changed size during iteration
distributed.worker - WARNING - Worker is at 87% memory usage. Pausing worker.  Process memory: 25.84 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker.html#memtrim for more information. -- Unmanaged memory: 25.84 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker.html#memtrim for more information. -- Unmanaged memory: 25.90 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - WARNING - Worker is at 7% memory usage. Resuming worker. Process memory: 2.07 GiB -- Worker memory limit: 29.51 GiB
distributed.worker - ERROR -
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 1857, in transition_flight_memory
    self.batched_stream.send({"op": "add-keys", "keys": [ts.key]})
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/batched.py", line 136, in send
    raise CommClosedError()
distributed.comm.core.CommClosedError
tornado.application - ERROR - Exception in callback functools.partial(<bound method IOLoop._discard_future_result of <tornado.platform.asyncio.AsyncIOLoop object at 0x7ff429eb9b20>>, <Task finished name='Task-2007' coro=<Worker.gather_dep() done, defined at /usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py:2267> exception=CommClosedError()>)
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/tornado/ioloop.py", line 741, in _run_callback
    ret = callback()
  File "/usr/local/lib/python3.8/dist-packages/tornado/ioloop.py", line 765, in _discard_future_result
    future.result()
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 2411, in gather_dep
    self.transition(ts, "memory", value=data[d])
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 1694, in transition
    state = func(ts, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/worker.py", line 1857, in transition_flight_memory
    self.batched_stream.send({"op": "add-keys", "keys": [ts.key]})
  File "/usr/local/lib/python3.8/dist-packages/distributed-2021.7.1-py3.8.egg/distributed/batched.py", line 136, in send
    raise CommClosedError()
distributed.comm.core.CommClosedError


**Expected behavior**
should be able to process 11 GB parquet with 4 T4
GCP - n1-standard-32+4xT4 - nvcr.io/nvidia/merlin/merlin-training:vnightly
Driver: NVIDIA-SMI 460.73.01 Driver Version: 460.73.01 CUDA Version: 11.2



Discussion:

Resolving patch:
diff --git a/nvtabular/io/fsspec_utils.py b/nvtabular/io/fsspec_utils.py
--- a/nvtabular/io/fsspec_utils.py
+++ b/nvtabular/io/fsspec_utils.py
@@ -109,30 +109,30 @@ def _optimized_read_remote(path, row_groups, columns, fs, **kwargs):
         path, row_groups, columns, fs, **kwargs
     )
 
-    # Transfer the required byte-ranges with fsspec.
-    # Store these blocks in a local dummy buffer
-    dummy_buffer = _fsspec_data_transfer(
-        path,
-        fs,
-        byte_ranges=byte_ranges,
-        footer=footer,
-        file_size=file_size,
-        add_par1_magic=True,
-        **kwargs,
-    )
-
     # Call cudf.read_parquet on the dummy buffer
     strings_to_cats = kwargs.get("strings_to_categorical", False)
-    df = cudf.read_parquet(
-        io.BytesIO(dummy_buffer),
+    return cudf.read_parquet(
+        # Wrap in BytesIO since cudf will sometimes use
+        # pyarrow to parse the metadata (and pyarrow
+        # cannot read from a bytes object)
+        io.BytesIO(
+            # Transfer the required bytes with fsspec
+            _fsspec_data_transfer(
+                path,
+                fs,
+                byte_ranges=byte_ranges,
+                footer=footer,
+                file_size=file_size,
+                add_par1_magic=True,
+                **kwargs,
+            )
+        ),
         engine="cudf",
         columns=columns,
         row_groups=row_groups,
         strings_to_categorical=strings_to_cats,
         **kwargs.get("read", {}),
     )
-    del dummy_buffer
-    return df
 
 
 def _get_parquet_byte_ranges(


Test patch: