Repository: NVIDIA-Merlin/NVTabular

Issue:
[BUG] NVTabular workflow writes more input columns for triton inference server
**Describe the bug**
Using `generate_triton_model` generates the config file for triton inference. It depends on workflow.input_dtypes.
If my workflow does not use all available columns, then the workflow.input_dtypes contains more input columns and triton will fail loading the config.

NVTabular workflow is:
![image](https://user-images.githubusercontent.com/19567274/109032917-c0140000-768b-11eb-8dc6-5c6e50332fb0.png)

Current (fails):
![image](https://user-images.githubusercontent.com/19567274/109032961-cc985880-768b-11eb-8108-2c441fef9c9b.png)

Error:
```
E0224 15:58:10.330248 178 model_repository_manager.cc:963] failed to load 'amazonreview_tf' version 1: Internal: unable to create stream: the provided PTX was compiled with an unsupported toolchain.
/nvtabular/nvtabular/workflow.py:236: UserWarning: Loading workflow generated with cudf version 0+untagged.1.gbd321d1 - but we are running cudf 0.18.0a+253.g53ed28e91c. This might cause issues
  warnings.warn(
E0224 15:58:20.534884 178 model_repository_manager.cc:963] failed to load 'amazonreview_nvt' version 1: Internal: Traceback (most recent call last):
  File "/opt/tritonserver/backends/python/startup.py", line 197, in Init
    self.backend.initialize(args)
  File "/models/models/amazonreview_nvt/1/model.py", line 57, in initialize
    self.output_dtypes[name] = triton_string_to_numpy(conf["data_type"])
TypeError: 'NoneType' object is not subscriptable
I0224 15:58:20.535093 178 server.cc:490]
```

This error occurs, when triton tries to load the config for the column `brand`, which is not used in the workflow

How to fix it:
![image](https://user-images.githubusercontent.com/19567274/109033003-d4f09380-768b-11eb-874e-953beecbdc8a.png)

**Steps/Code to reproduce bug**
Use MovieLens example and do not drop the timestamp column


Discussion:

Resolving patch:
diff --git a/nvtabular/workflow.py b/nvtabular/workflow.py
--- a/nvtabular/workflow.py
+++ b/nvtabular/workflow.py
@@ -151,7 +151,7 @@ def fit(self, dataset: Dataset):
         # hack: store input/output dtypes here. We should have complete dtype
         # information for each operator (like we do for column names), but as
         # an interim solution this gets us what we need.
-        input_dtypes = dataset.to_ddf().dtypes
+        input_dtypes = dataset.to_ddf()[self._input_columns()].dtypes
         self.input_dtypes = dict(zip(input_dtypes.index, input_dtypes))
         output_dtypes = self.transform(dataset).to_ddf().head(1).dtypes
         self.output_dtypes = dict(zip(output_dtypes.index, output_dtypes))


Test patch:
diff --git a/tests/unit/test_workflow.py b/tests/unit/test_workflow.py
--- a/tests/unit/test_workflow.py
+++ b/tests/unit/test_workflow.py
@@ -580,3 +580,14 @@ def test_workflow_move_saved(tmpdir):
     # also check that when transforming our input we get the same results after loading
     transformed = workflow2.transform(Dataset(data)).to_ddf().compute()
     assert_eq(expected, transformed)
+
+
+def test_workflow_input_output_dtypes():
+    df = cudf.DataFrame({"genre": ["drama", "comedy"], "user": ["a", "b"], "unneeded": [1, 2]})
+    features = [["genre", "user"], "genre"] >> ops.Categorify(encode_type="combo")
+    workflow = Workflow(features)
+    workflow.fit(Dataset(df))
+
+    assert "unneeded" not in workflow.input_dtypes
+    assert set(workflow.input_dtypes.keys()) == {"genre", "user"}
+    assert set(workflow.output_dtypes.keys()) == {"genre_user", "genre"}