Repository: NVIDIA-Merlin/NVTabular

Issue:
[FEA] Triton inference support should be able to select running on cpu/gpu



Discussion:

Resolving patch:
diff --git a/nvtabular/inference/triton/__init__.py b/nvtabular/inference/triton/__init__.py
--- a/nvtabular/inference/triton/__init__.py
+++ b/nvtabular/inference/triton/__init__.py
@@ -18,6 +18,8 @@
 import os
 from shutil import copyfile
 
+import pandas as pd
+
 # this needs to be before any modules that import protobuf
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 
@@ -350,6 +352,8 @@ def convert_df_to_triton_input(column_names, batch, input_class=grpcclient.Infer
     inputs = []
     for i, (name, col) in enumerate(columns):
         if _is_list_dtype(col):
+            if isinstance(col, pd.Series):
+                raise ValueError("this function doesn't support CPU list values yet")
             inputs.append(
                 _convert_column_to_triton_input(
                     col._column.offsets.values_host.astype("int64"), name + "__nnzs", input_class
@@ -361,7 +365,8 @@ def convert_df_to_triton_input(column_names, batch, input_class=grpcclient.Infer
                 )
             )
         else:
-            inputs.append(_convert_column_to_triton_input(col.values_host, name, input_class))
+            values = col.values if isinstance(col, pd.Series) else col.values_host
+            inputs.append(_convert_column_to_triton_input(values, name, input_class))
     return inputs
 
 
diff --git a/nvtabular/inference/triton/model.py b/nvtabular/inference/triton/model.py
--- a/nvtabular/inference/triton/model.py
+++ b/nvtabular/inference/triton/model.py
@@ -24,6 +24,7 @@
 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import functools
 import json
 import logging
 import os
@@ -62,6 +63,18 @@ def _initialize_ops(self, column_group, visited=None):
             if inference_op:
                 column_group.op = inference_op
 
+            supported = column_group.op.supports
+
+            # if we're running on the CPU only, mask off support for GPU data formats
+            if self.kind == "CPU":
+                supported = functools.reduce(
+                    lambda a, b: a | b,
+                    (v for v in list(Supports) if v & supported and "CPU" in str(v)),
+                )
+            # the 'supports' property is readonly, and we can't always attach a new property
+            # to some of the operators (C++ categorify etc). set on the column_group instead
+            column_group.inference_supports = supported
+
         for parent in column_group.parents:
             if parent not in visited:
                 visited.add(parent)
@@ -72,6 +85,7 @@ def initialize(self, args):
             args["model_repository"], str(args["model_version"]), "workflow"
         )
         self.workflow = nvtabular.Workflow.load(workflow_path)
+        self.kind = args["model_instance_kind"]
         self.model_config = json.loads(args["model_config"])
         self.output_model = self.model_config["parameters"]["output_model"]["string_value"]
 
@@ -229,7 +243,7 @@ def _transform_tensors(input_tensors, column_group):
                     # we have multiple different kinds of data here (dataframe/array on cpu/gpu)
                     # we need to convert to a common format here first before concatentating.
                     op = column_group.op
-                    target_kind = op.supports if op else Supports.CPU_DICT_ARRAY
+                    target_kind = column_group.inference_supports if op else Supports.CPU_DICT_ARRAY
                     # note : the 2nd convert_format call needs to be stricter in what the kind is
                     # (exact match rather than a bitmask of values)
                     tensors, kind = convert_format(tensors, kind, target_kind)
@@ -244,8 +258,8 @@ def _transform_tensors(input_tensors, column_group):
     if column_group.op:
         try:
             # if the op doesn't support the current kind - we need to convert
-            if not column_group.op.supports & kind:
-                tensors, kind = convert_format(tensors, kind, column_group.op.supports)
+            if not column_group.inference_supports & kind:
+                tensors, kind = convert_format(tensors, kind, column_group.inference_supports)
 
             tensors = column_group.op.transform(column_group.columns, tensors)
 
diff --git a/nvtabular/workflow.py b/nvtabular/workflow.py
--- a/nvtabular/workflow.py
+++ b/nvtabular/workflow.py
@@ -258,9 +258,14 @@ def check_version(stored, current, name):
         lib = cudf if cudf else pd
         versions = meta["versions"]
         check_version(versions["nvtabular"], nvt_version, "nvtabular")
-        check_version(versions[lib.__name__], lib.__version__, lib.__name__)
         check_version(versions["python"], sys.version, "python")
 
+        if lib.__name__ in versions:
+            check_version(versions[lib.__name__], lib.__version__, lib.__name__)
+        else:
+            expected = "GPU" if "cudf" in versions else "CPU"
+            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.client = client


Test patch:
diff --git a/tests/unit/test_triton_inference.py b/tests/unit/test_triton_inference.py
--- a/tests/unit/test_triton_inference.py
+++ b/tests/unit/test_triton_inference.py
@@ -11,7 +11,7 @@
 
 import nvtabular as nvt
 import nvtabular.ops as ops
-from nvtabular.dispatch import HAS_GPU, _make_df
+from nvtabular.dispatch import HAS_GPU, _hash_series, _make_df
 from nvtabular.ops.operator import Supports
 from tests.conftest import assert_eq
 
@@ -141,7 +141,7 @@ def test_concatenate_dataframe(tmpdir):
         }
     )
     # this bug only happened with a dataframe representation: force this by using a lambda
-    cats = ["cat"] >> ops.LambdaOp(lambda col: col.hash_values() % 1000)
+    cats = ["cat"] >> ops.LambdaOp(lambda col: _hash_series(col) % 1000)
     conts = ["cont"] >> ops.Normalize() >> ops.FillMissing() >> ops.LogOp()
     workflow = nvt.Workflow(cats + conts)
     _verify_workflow_on_tritonserver(tmpdir, workflow, df, "test_concatenate_dataframe")