Support serving in container without tensorflow, torch, or cudf (…

…#326) * Enable ensemble to run in container without tensorflow installed * Check if cudf is available in convert_format * Support Triton runtime with torch op in env without torch package * Use tensorflow/torch imports from merlin.core.compat --------- Co-authored-by: Karl Higley <[email protected]>
NVIDIA-Merlin · Apr 13, 2023 · da739e2 · da739e2
1 parent 386ef22
commit da739e2
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 5 deletions.
diff --git a/merlin/systems/dag/ops/pytorch.py b/merlin/systems/dag/ops/pytorch.py
@@ -16,9 +16,8 @@
 import os
 
 import numpy as np
-import torch
-import torch.utils.dlpack
 
+from merlin.core.compat.torch import torch
 from merlin.core.protocols import Transformable
 from merlin.dag import ColumnSelector
 from merlin.schema import Schema

diff --git a/merlin/systems/dag/ops/tensorflow.py b/merlin/systems/dag/ops/tensorflow.py
@@ -20,8 +20,7 @@
 # this needs to be before any modules that import protobuf
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 
-import tensorflow as tf  # noqa
-
+from merlin.core.compat.tensorflow import tensorflow as tf  # noqa
 from merlin.core.protocols import Transformable  # noqa
 from merlin.dag import ColumnSelector  # noqa
 from merlin.schema import ColumnSchema, Schema  # noqa

diff --git a/merlin/systems/triton/conversions.py b/merlin/systems/triton/conversions.py
@@ -240,7 +240,7 @@ def convert_format(tensors, kind, target_kind):
         elif kind == Supports.GPU_DATAFRAME:
             return _cudf_to_array(tensors, True), Supports.CPU_DICT_ARRAY
 
-    elif target_kind & Supports.GPU_DATAFRAME:
+    elif cudf and target_kind & Supports.GPU_DATAFRAME:
         if kind == Supports.CPU_DATAFRAME:
             return cudf.DataFrame(tensors), Supports.GPU_DATAFRAME
         return _array_to_cudf(tensors), Supports.GPU_DATAFRAME