Repository: NVIDIA-Merlin/NVTabular

Issue:
[FEA] Add DistributedClient API and use global client objects in Workflow and Dataset
**Is your feature request related to a problem? Please describe.**
In order to perform distributed ETL and/or spill data between device/host/disk, NVTabular currently requires the user to provide a Dask-Distributed (or Dask-CUDA) cluster.  Furthermore, the `nvt.Workflow` (and sometimes `nvt.Dataset`) need to be defined with an explicit `client=` argument in order for the distributed cluster to be used.  Although I feel strongly that it would be dangerous and messy for NVTabular to automatically spin up a distributed cluster by default, I do suspect that the user experience could be much better. 

**Describe the solution you'd like**
To improve the user experience of distributed ETL with NVTabular, I propose:

1. Simple `LocalCluster`/`LocalCUDACluster`-wrapper APIs be added to NVTabular so that users can enable multi-GPU processing and/or spilling without interacting with distributed/dask_cuda. I am not yet sure of the ideal API to expose in NVTabular, but perhaps something like `DistributedClient` (wrapping `distributed.Client`). This API could be used to automatically generate a local cluster (if the address of an existing cluster is not provided), and we could add a `cpu=False` kwarg to toggle between gpu and cpu mode.
2. [**DONE** in #1318] Automatically detect and **use** an existing Dask client object. NVTabular already [checks for a global dask client](https://github.com/NVIDIA-Merlin/NVTabular/blob/34d01d7e6090d6029ac40010ed79e1558f18759c/nvtabular/workflow/workflow.py#L88) in both `Workflow ` and `Dataset`.  However, these checks result in a UserWarning whenever a global client is detected.  Instead of warning the user, I propose that NVTabular automatically **use** the client object when it is detected (which is actually the defult behavior in `dask.dataframe` and `dask.array` anyway).

**Describe alternatives you've considered**
The alternative is to continue requiring the user to deploy Dask clusters/clients with `distributed`/`dask_cuda`, and require them to include an explicit `client` argument when defining an `Workflow` objects.



Discussion:
I think this is a great idea, and it would be very valuable. All our examples/notebooks would be automatically running in multi-gpu and customers would be able to use multi-gpu with minimum effort and knowledge.

Happy to work into this, unless Rick wants to tackle it.
>Happy to work into this, unless Rick wants to tackle it.

Happy to let you tackle it. We could also split the work if you want. I imagine that it would be best to add features in two distinct PRs: (1) automatically use a global client object, and (2) introduce the `DistributedClient` API (or some other NVT-based mechanism for users to say: "I want distributed execution")
We now have the second part of this issue resolved (automatic client detection/handling).  In resolving this, we added a `nvtabular.utils.set_dask_client(client="auto")` utility.

Perhaps we could use the same machinery to solve the remaining challenge of this issue. That is: Add a `client="new"` option to `set_dask_client`, and suport `cluster_type=` and `cluster_options=` arguments. In the case that the user specifies `client="new"`, we would use `cluster_type` and `cluster_options` to automatically deploy a local cluster and define a new global client.

By default, using `client="new"`would result in a default `LocalCUDACluster` if gpus are detected on the machine (and a default `LocalCluster` otherwise).

Resolving patch:
diff --git a/nvtabular/__init__.py b/nvtabular/__init__.py
--- a/nvtabular/__init__.py
+++ b/nvtabular/__init__.py
@@ -18,7 +18,7 @@
 from merlin.graph import ColumnSelector
 from merlin.schema import ColumnSchema, Schema
 
-from . import graph, io, workflow  # noqa
+from . import graph, io, utils, workflow  # noqa
 from ._version import get_versions
 
 # suppress some warnings with cudf warning about column ordering with dlpack
@@ -30,6 +30,8 @@
 WorkflowNode = workflow.WorkflowNode
 Workflow = workflow.Workflow
 Dataset = io.dataset.Dataset
+Distributed = utils.Distributed
+Serial = utils.Serial
 
 
 # Provides an alias of ColumnSelector so that old usages of ColumnGroup to
@@ -47,6 +49,8 @@ def ColumnGroup(columns):
     "ColumnSelector",
     "ColumnSchema",
     "Schema",
+    "Distributed",
+    "Serial",
 ]
 
 # cudf warns about column ordering with dlpack methods, ignore it


Test patch:
diff --git a/tests/unit/test_dask_nvt.py b/tests/unit/test_dask_nvt.py
--- a/tests/unit/test_dask_nvt.py
+++ b/tests/unit/test_dask_nvt.py
@@ -26,7 +26,7 @@
 
 from nvtabular import ColumnSelector, Dataset, Workflow, ops
 from nvtabular.io import Shuffle
-from nvtabular.utils import set_dask_client
+from nvtabular.utils import global_dask_client, set_dask_client
 from tests.conftest import allcols_csv, mycols_csv, mycols_pq
 
 cudf = pytest.importorskip("cudf")
@@ -295,3 +295,21 @@ def test_filtered_partition(tmpdir, cpu):
 
     # Write result to disk
     workflow.transform(dataset).to_parquet(str(tmpdir))
+
+
+def test_merlin_core_execution_managers(client):
+    # This functionality lives in merlin-core,
+    # but we are testing that we can import
+    # from NVTabular with proper behavior.
+    from nvtabular import Distributed, Serial
+
+    # Set distributed client
+    with Distributed(client=client):
+        assert global_dask_client() == client
+
+        # Check that the global dask client
+        # becomes None in a `with Serial()` block
+        with Serial():
+            assert global_dask_client() is None
+
+        assert global_dask_client() == client