Repository: NVIDIA-Merlin/NVTabular

Issue:
[BUG] NVTabular: Support for GCS paths: read + write (Dataset, Workflow classes)
**Describe the bug**
Dataset: Only supports reading from GCS, not writing.
Workflow: Does not support GCS paths for reading and writing.

**Steps/Code to reproduce bug**
follow
https://github.com/jarokaz/merlin-on-vertex/blob/main/01-dataset-preprocessing.ipynb

Details will be added by Vertex team.

**Expected behavior**
as description

GCP - n1-highmem-32+4xT4 - nvcr.io/nvidia/merlin/merlin-training:21.09
Driver: NVIDIA-SMI 460.73.01 Driver Version: 460.73.01 CUDA Version: 11.2



Discussion:
In the past (Maybe 6 months ago), I have been able to run an NVTabular workflow reading/writing data from/in cloud(gcs and S3) using this script: https://github.com/NVIDIA-Merlin/NVTabular/blob/main/bench/examples/dask-nvtabular-criteo-benchmark.py 

I actually got some perf numbers and was working in how to better configure GCP/AWS VMs storage to get better I/O access.

Not sure if this is exactly the same, but just adding the information.


>I have been able to run an NVTabular workflow reading/writing data from/in cloud(gcs and S3) using this script: https://github.com/NVIDIA-Merlin/NVTabular/blob/main/bench/examples/dask-nvtabular-criteo-benchmark.py

@albert17 - Are you saying you were able to run your script with `--out-path` pointing to a location in both s3 and gcs? I have not spent much time trying to write to remote storage, but I do recall having problems writing to gcs. I believe I was able to write directly to s3, but needed to copy (`put`) local files to get them into gcs.
@rjzamora Yes, I meant that. I added a PR at that time to add the support to the script and to update several components from NVTabular. #540 
@mengdong for writing the dataset, are you using the latest code? There was a fix pushed about a week ago that at least fixed this for writing the dataset to S3: https://github.com/NVIDIA-Merlin/NVTabular/issues/1214


Resolving patch:
diff --git a/nvtabular/ops/categorify.py b/nvtabular/ops/categorify.py
--- a/nvtabular/ops/categorify.py
+++ b/nvtabular/ops/categorify.py
@@ -1345,14 +1345,15 @@ def _hash_bucket(df, num_buckets, col, encode_type="joint"):
 
 def _copy_storage(existing_stats, existing_path, new_path, copy):
     """helper function to copy files to a new storage location"""
-    from shutil import copyfile
-
+    existing_fs = get_fs_token_paths(existing_path)[0]
+    new_fs = get_fs_token_paths(new_path)[0]
     new_locations = {}
     for column, existing_file in existing_stats.items():
         new_file = existing_file.replace(str(existing_path), str(new_path))
         if copy and new_file != existing_file:
-            os.makedirs(os.path.dirname(new_file), exist_ok=True)
-            copyfile(existing_file, new_file)
+            new_fs.makedirs(os.path.dirname(new_file), exist_ok=True)
+            with new_fs.open(new_file, "wb") as output:
+                output.write(existing_fs.open(existing_file, "rb").read())
 
         new_locations[column] = new_file
 
diff --git a/nvtabular/worker.py b/nvtabular/worker.py
--- a/nvtabular/worker.py
+++ b/nvtabular/worker.py
@@ -23,7 +23,6 @@
     import cudf
 except ImportError:
     cudf = None
-import fsspec
 import pyarrow as pa
 from dask.distributed import get_worker
 
@@ -83,8 +82,7 @@ def fetch_table_data(
             if reader == _lib.read_parquet:  # pylint: disable=comparison-with-callable
                 # Using cudf-backed data with "host" caching.
                 # Cache as an Arrow table.
-                with contextlib.closing(fsspec.open(path, "rb")) as f:
-                    table = reader(f, **use_kwargs)
+                table = reader(path, **use_kwargs)
                 if cudf:
                     table_cache[path] = table.to_arrow()
                 else:
diff --git a/nvtabular/workflow/workflow.py b/nvtabular/workflow/workflow.py
--- a/nvtabular/workflow/workflow.py
+++ b/nvtabular/workflow/workflow.py
@@ -15,13 +15,13 @@
 #
 import json
 import logging
-import os
 import sys
 import time
 import warnings
 from typing import TYPE_CHECKING, Optional
 
 import cloudpickle
+import fsspec
 
 try:
     import cudf
@@ -283,7 +283,9 @@ def save(self, path):
         # avoid a circular import getting the version
         from nvtabular import __version__ as nvt_version
 
-        os.makedirs(path, exist_ok=True)
+        fs = fsspec.get_fs_token_paths(path)[0]
+
+        fs.makedirs(path, exist_ok=True)
 
         # point all stat ops to store intermediate output (parquet etc) at the path
         # this lets us easily bundle
@@ -292,7 +294,7 @@ def save(self, path):
 
         # generate a file of all versions used to generate this bundle
         lib = cudf if cudf else pd
-        with open(os.path.join(path, "metadata.json"), "w") as o:
+        with fs.open(fs.sep.join([path, "metadata.json"]), "w") as o:
             json.dump(
                 {
                     "versions": {
@@ -306,7 +308,7 @@ def save(self, path):
             )
 
         # dump out the full workflow (graph/stats/operators etc) using cloudpickle
-        with open(os.path.join(path, "workflow.pkl"), "wb") as o:
+        with fs.open(fs.sep.join([path, "workflow.pkl"]), "wb") as o:
             cloudpickle.dump(self, o)
 
     @classmethod
@@ -327,8 +329,10 @@ def load(cls, path, client=None):
         # avoid a circular import getting the version
         from nvtabular import __version__ as nvt_version
 
+        fs = fsspec.get_fs_token_paths(path)[0]
+
         # check version information from the metadata blob, and warn if we have a mismatch
-        meta = json.load(open(os.path.join(path, "metadata.json")))
+        meta = json.load(fs.open(fs.sep.join([path, "metadata.json"])))
 
         def parse_version(version):
             return version.split(".")[:2]
@@ -354,7 +358,7 @@ def check_version(stored, current, name):
             warnings.warn(f"Loading workflow generated on {expected}")
 
         # load up the workflow object di
-        workflow = cloudpickle.load(open(os.path.join(path, "workflow.pkl"), "rb"))
+        workflow = cloudpickle.load(fs.open(fs.sep.join([path, "workflow.pkl"]), "rb"))
         workflow.client = client
 
         # we might have been copied since saving, update all the stat ops


Test patch:
diff --git a/tests/unit/test_s3.py b/tests/unit/test_s3.py
--- a/tests/unit/test_s3.py
+++ b/tests/unit/test_s3.py
@@ -85,7 +85,7 @@ def test_s3_dataset(s3_base, s3so, paths, datasets, engine, df, patch_aiobotocor
         meta.write_metadata_file(files[fn])
         files[fn].seek(0)
 
-    with s3_context(s3_base=s3_base, bucket=engine, files=files):
+    with s3_context(s3_base=s3_base, bucket=engine, files=files) as s3fs:
         # Create nvt.Dataset from mock s3 paths
         url = f"s3://{engine}" if engine == "parquet" else f"s3://{engine}/*"
         dataset = nvt.Dataset(url, engine=engine, storage_options=s3so)
@@ -108,3 +108,17 @@ def test_s3_dataset(s3_base, s3so, paths, datasets, engine, df, patch_aiobotocor
         # make sure we can write out the dataset back to S3
         # (https://github.com/NVIDIA-Merlin/NVTabular/issues/1214)
         processor.transform(dataset).to_parquet(f"s3://{engine}/output")
+        expected = processor.transform(dataset).to_ddf().compute()
+
+        # make sure we can write out the workflow to s3
+        processor.save(f"s3://{engine}/saved_workflow/")
+
+        # make sure the workflow got saved to the right spot in S3
+        workflow_files = s3fs.ls(f"/{engine}/saved_workflow/")
+        assert workflow_files
+
+        # finally make sure we can read in the workflow from S3, and use it
+        # to transform values and get the same result as on the local fs
+        reloaded = nvt.Workflow.load(f"s3://{engine}/saved_workflow/")
+        from_s3 = reloaded.transform(dataset).to_ddf().compute()
+        assert_eq(expected, from_s3)