Repository: NVIDIA-Merlin/NVTabular

Issue:
[BUG]Criteo Preprocessing example failed with increase of number of parquet files.
**Describe the bug**
observe Exception: KeyError('C1') when using 8 A100 each output 8 files, however, same job succeed when using 8A100 each output 4 files

**Steps/Code to reproduce bug**
Follow https://github.com/jarokaz/merlin-on-vertex/blob/main/src/training/dataprep/README.md

docker run -it --rm --gpus all \
        -v /home/dongm/merlin-on-vertex/src/training/dataprep:/src \
        -w /src \
        -v /mnt/disks/data:/criteo_data \
        nvcr.io/nvidia/merlin/merlin-training:21.09 \
        python preprocess.py \
        --train_folder /criteo_data/criteo_raw_parquet_train \
        --valid_folder /criteo_data/criteo_raw_parquet_valid \
        --output_folder /criteo_data/criteo_processed_parquet \
        --devices 0,1,2,3,4,5,6,7 \
        --protocol tcp \
        --device_limit_frac 0.8 \
        --device_pool_frac 0.9 \
        --num_io_threads 8 \
        --part_mem_frac 0.08 \
        --out_files_per_proc 8 \
        --freq_limit 6 \
        --shuffle PER_PARTITION

Will fail
```
distributed.worker - WARNING - Compute Failed
Function:  _write_subgraph
args:      (<nvtabular.io.dask.DaskSubgraph object at 0x7ef650f80340>, ('part_56.parquet', 'part_57.parquet', 'part_58.parquet', 'part_59.parquet', 'part_60.parquet', 'part_61.parquet', 'part_62.parquet', 'part_63.parquet'), '/criteo_data/criteo_processed_parquet/train', <Shuffle.PER_PARTITION: 0>, <fsspec.implementations.local.LocalFileSystem object at 0x7f015df3c790>, ['C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9', 'C10', 'C11', 'C12', 'C13', 'C14', 'C15', 'C16', 'C17', 'C18', 'C19', 'C20', 'C21', 'C22', 'C23', 'C24', 'C25', 'C26'], ['I1', 'I2', 'I3', 'I4', 'I5', 'I6', 'I7', 'I8', 'I9', 'I10', 'I11', 'I12', 'I13'], ['label'], 'parquet', 8, False, '')
kwargs:    {}
Exception: KeyError('C1')

Traceback (most recent call last):
  File "preprocess.py", line 284, in <module>
    preprocess(args)
  File "preprocess.py", line 135, in preprocess
    workflow.transform(train_dataset).to_parquet(
  File "/nvtabular/nvtabular/io/dataset.py", line 886, in to_parquet
    _ddf_to_dataset(
  File "/nvtabular/nvtabular/io/dask.py", line 369, in _ddf_to_dataset
    out = client.compute(out).result()
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/client.py", line 228, in result
    raise exc.with_traceback(tb)
  File "/usr/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/nvtabular/nvtabular/io/dask.py", line 210, in _write_subgraph
    return writer.close()
  File "/nvtabular/nvtabular/io/writer.py", line 314, in close
    _general_meta = self.package_general_metadata()
  File "/nvtabular/nvtabular/io/writer.py", line 253, in package_general_metadata
    data["cats"].append({"col_name": c, "index": self.col_idx[c]})
KeyError: 'C1'
distributed.nanny - WARNING - Worker process still alive after 3 seconds, killing
distributed.nanny - WARNING - Worker process still alive after 3 seconds, killing
distributed.nanny - WARNING - Worker process still alive after 3 seconds, killing
distributed.nanny - WARNING - Worker process still alive after 3 seconds, killing
distributed.nanny - WARNING - Worker process still alive after 3 seconds, killing
distributed.nanny - WARNING - Worker process still alive after 3 seconds, killing
distributed.nanny - WARNING - Worker process still alive after 3 seconds, killing
distributed.core - ERROR - Exception while handling op register-client
Traceback (most recent call last):
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/core.py", line 502, in handle_comm
    result = await result
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/scheduler.py", line 5194, in add_client
    self.remove_client(client=client)
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/scheduler.py", line 5221, in remove_client
    self.client_releases_keys(
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/scheduler.py", line 4961, in client_releases_keys
    self.transitions(recommendations)
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/scheduler.py", line 6975, in transitions
    self.send_all(client_msgs, worker_msgs)
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/scheduler.py", line 5478, in send_all
    w = stream_comms[worker]
KeyError: 'tcp://127.0.0.1:36065'
tornado.application - ERROR - Exception in callback functools.partial(<function TCPServer._handle_connection.<locals>.<lambda> at 0x7f8c74443820>, <Task finished name='Task-72' coro=<BaseTCPListener._handle_stream() done, defined at /root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/comm/tcp.py:478> exception=KeyError('tcp://127.0.0.1:36065')>)
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/tcpserver.py", line 331, in <lambda>
    gen.convert_yielded(future), lambda f: f.result()
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/comm/tcp.py", line 495, in _handle_stream
    await self.comm_handler(comm)
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/core.py", line 502, in handle_comm
    result = await result
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/scheduler.py", line 5194, in add_client
    self.remove_client(client=client)
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/scheduler.py", line 5221, in remove_client
    self.client_releases_keys(
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/scheduler.py", line 4961, in client_releases_keys
    self.transitions(recommendations)
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/scheduler.py", line 6975, in transitions
    self.send_all(client_msgs, worker_msgs)
  File "/root/.local/lib/python3.8/site-packages/distributed-2021.7.1-py3.8.egg/distributed/scheduler.py", line 5478, in send_all
    w = stream_comms[worker]
KeyError: 'tcp://127.0.0.1:36065'
Exception in thread AsyncProcess Dask Worker process (from Nanny) watch process join:
Traceback (most recent call last):
  File "/usr/lib/python3.8/threading.py", line 932, in _bootstrap_inner

```


docker run -it --rm --gpus all \
        -v /home/dongm/merlin-on-vertex/src/training/dataprep:/src \
        -w /src \
        -v /mnt/disks/data:/criteo_data \
        nvcr.io/nvidia/merlin/merlin-training:21.09 \
        python preprocess.py \
        --train_folder /criteo_data/criteo_raw_parquet_train \
        --valid_folder /criteo_data/criteo_raw_parquet_valid \
        --output_folder /criteo_data/criteo_processed_parquet \
        --devices 0,1,2,3,4,5,6,7 \
        --protocol tcp \
        --device_limit_frac 0.8 \
        --device_pool_frac 0.9 \
        --num_io_threads 8 \
        --part_mem_frac 0.08 \
        --out_files_per_proc 4 \
        --freq_limit 6 \
        --shuffle PER_PARTITION

Will work

**Expected behavior**
The workflow transform should work when there are missing keys in certain partition.

**Environment details (please complete the following information):**
GCP A100 a2-high-8g, | NVIDIA-SMI 470.57.02    Driver Version: 470.57.02    CUDA Version: 11.4 
 nvcr.io/nvidia/merlin/merlin-training:21.09

Aha! Link: https://nvaiinfra.aha.io/features/MERLIN-505


Discussion:
@mengdong how are you generating these parquet files? Can you generate a consistent set of columns in each parquet file?

@rjzamora any ideas on how complicated this would be to implement? 


If we look at the code indicate by error: https://github.com/jarokaz/merlin-on-vertex/blob/main/src/training/dataprep/preprocess.py#L135-L143

We will see that we use nvt.Workflow(features, client=client)  to create the workflow, in which we pass only 1 set of features, which when distributed to Dask workers, should ideally let each worker know the columns per:

```
    features, dict_dtypes = create_preprocessing_workflow(
        categorical_columns=CATEGORICAL_COLUMNS,
        continuous_columns=CONTINUOUS_COLUMNS,
        label_columns=LABEL_COLUMNS,
        freq_limit=args.freq_limit, 
        stats_path=stats_output_folder)
```

No sure why this error would occur at all.
@benfred did you have a chance to review @mengdong 's comment?
@mengdong is this the criteo dataset gotten by running our download/convert notebook https://github.com/NVIDIA-Merlin/NVTabular/blob/main/examples/scaling-criteo/01-Download-Convert.ipynb ?  Or was this parquet dataset generated by another script like https://github.com/jarokaz/merlin-on-vertex/blob/main/src/training/dataprep/convert_to_parquet.py ?

@jperez999 @rjzamora can either of you reproduce this issue?
it is https://github.com/jarokaz/merlin-on-vertex/blob/main/src/training/dataprep/convert_to_parquet.py.
@mengdong can you link that script for conversion again, seems the links above are now reporting a 404.
Hi @jperez999,
We are using this code to convert to Parquet:
https://github.com/jarokaz/merlin-on-vertex/blob/7bea3758b8ca1166c8a45bd547080013b4078c9a/src/preprocessing/etl.py#L79

Please let me know if you have any question.
Thanks

Resolving patch:
diff --git a/nvtabular/io/dataset.py b/nvtabular/io/dataset.py
--- a/nvtabular/io/dataset.py
+++ b/nvtabular/io/dataset.py
@@ -861,7 +861,14 @@ def to_parquet(
                     fns = output_files[i : i + files_per_task]
                     start = i * split
                     stop = min(start + split * len(fns), ddf.npartitions)
-                    new[tuple(fns)] = np.arange(start, stop)
+                    if start < stop:
+                        new[tuple(fns)] = np.arange(start, stop)
+                # let user know they will not have expected number of output files.
+                if len(new.keys()) < len(output_files):
+                    warnings.warn(
+                        f"Only created {len(new.keys())} files did not have enough\n"
+                        f"partitions to create {len(output_files)} files."
+                    )
                 output_files = new
                 suffix = ""  # Don't add a suffix later - Names already include it
             if not isinstance(output_files, dict):


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
@@ -158,6 +158,36 @@ def test_dask_dataset_itr(tmpdir, datasets, engine, gpu_memory_frac):
     assert len(my_iter) == size
 
 
+def test_io_partitions_push(tmpdir):
+    os.makedirs(os.path.join(tmpdir, "csv"))
+
+    # Generate random csv files
+    files = [os.path.join(tmpdir, f"csv/day_{i}") for i in range(23)]
+    for file in files:
+        with open(file, "w") as f:
+            f.write("0,1,2,3,a,b,c\n" * 1000)
+
+    # Load csv files
+    label_columns = ["label"]
+    cont_columns = ["I1", "I2", "I3"]
+    cat_columns = ["C1", "C2", "C3"]
+    columns = label_columns + cont_columns + cat_columns
+    dataset = nvt.Dataset(files, engine="csv", names=columns)
+    print("npartitions of dataset:", dataset.npartitions)
+
+    for x in range(20):
+        dataset.to_parquet(
+            output_files=x,
+            output_path=os.path.join(tmpdir, f"parquet{x}"),
+            cats=cat_columns,
+            conts=cont_columns,
+            labels=label_columns,
+        )
+
+        df_lib = dispatch.get_lib()
+        df_lib.read_parquet(os.path.join(tmpdir, f"parquet{x}/part_0.parquet"))
+
+
 @pytest.mark.parametrize("engine", ["csv", "parquet", "csv-no-header"])
 @pytest.mark.parametrize("num_files", [1, 2])
 @pytest.mark.parametrize("cpu", [None, True])