From 399ef2385e7f2fb6ab9101ed53ab9712cd49e894 Mon Sep 17 00:00:00 2001
From: Oliver Holworthy <oholworthy@nvidia.com>
Date: Thu, 8 Jun 2023 14:53:36 +0100
Subject: [PATCH] Fix Pytorch integration test postmerge test (#363)

* Extend LD_LIBRARY_PATH with pytorch backend path

* Add skipif for GPU in test_pytorch_backend

* Extend match_representations type coercion to support torch tensors

* Run gpu-postmerge on pull request temporarily

* Use `name` instead of `to` for numpy/cupy dtype conversion

* Remove pull_request from postmerge-gpu workflow
---
 merlin/systems/triton/conversions.py          | 30 +++++++++++++++++--
 tests/integration/t4r/test_pytorch_backend.py |  2 ++
 tox.ini                                       |  2 +-
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/merlin/systems/triton/conversions.py b/merlin/systems/triton/conversions.py
index dee1367cb..585118f75 100644
--- a/merlin/systems/triton/conversions.py
+++ b/merlin/systems/triton/conversions.py
@@ -105,9 +105,9 @@ def match_representations(schema: Schema, dict_array: Dict[str, Any]) -> Dict[st
                 aligned[offs_name] = offsets
 
             if dtype != md.unknown:
-                aligned[vals_name] = aligned[vals_name].astype(dtype.to_numpy)
+                aligned[vals_name] = _astype(aligned[vals_name], dtype)
 
-            aligned[offs_name] = aligned[offs_name].astype("int32")
+            aligned[offs_name] = _astype(aligned[offs_name], md.dtype("int32"))
         else:
             try:
                 # Look for values and offsets that already exist,
@@ -120,7 +120,7 @@ def match_representations(schema: Schema, dict_array: Dict[str, Any]) -> Dict[st
                 aligned[col_name] = dict_array[col_name]
 
             if dtype != md.unknown:
-                aligned[col_name] = aligned[col_name].astype(dtype.to_numpy)
+                aligned[col_name] = _astype(aligned[col_name], dtype)
 
     return aligned
 
@@ -139,6 +139,30 @@ def _from_values_offsets(values, offsets, shape):
     return values.reshape(new_shape)
 
 
+@singledispatch
+def _astype(value, target_dtype):
+    raise NotImplementedError(f"_to_dtype not implemented for {type(value)}")
+
+
+@_astype.register
+def _(array: np.ndarray, target_dtype: md.DType):
+    return array.astype(target_dtype.name)
+
+
+if cp:
+
+    @_astype.register
+    def _(array: cp.ndarray, target_dtype: md.DType):
+        return array.astype(target_dtype.name)
+
+
+if torch:
+
+    @_astype.register
+    def _(tensor: torch.Tensor, target_dtype: md.DType):
+        return tensor.to(target_dtype.to("torch"))
+
+
 @singledispatch
 def _to_values_offsets(values):
     """Convert array to values/offsets representation
diff --git a/tests/integration/t4r/test_pytorch_backend.py b/tests/integration/t4r/test_pytorch_backend.py
index bb23bdf51..d93dc6204 100644
--- a/tests/integration/t4r/test_pytorch_backend.py
+++ b/tests/integration/t4r/test_pytorch_backend.py
@@ -29,6 +29,7 @@
 tritonclient = pytest.importorskip("tritonclient")
 grpcclient = pytest.importorskip("tritonclient.grpc")
 
+from merlin.core.compat import HAS_GPU  # noqa
 from merlin.core.dispatch import make_df  # noqa
 from merlin.systems.dag import Ensemble  # noqa
 from merlin.systems.dag.ops.pytorch import PredictPyTorch  # noqa
@@ -39,6 +40,7 @@
 
 
 @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found")
+@pytest.mark.skipif(not HAS_GPU, reason="GPU Device required for test")
 def test_serve_t4r_with_torchscript(tmpdir):
     # ===========================================
     # Generate training data
diff --git a/tox.ini b/tox.ini
index 657eac65d..bba7db44c 100644
--- a/tox.ini
+++ b/tox.ini
@@ -73,7 +73,7 @@ sitepackages=true
 ; need to add some back.
 setenv = 
     TF_GPU_ALLOCATOR=cuda_malloc_async
-    LD_LIBRARY_PATH=/opt/tritonserver/backends/pytorch
+    LD_LIBRARY_PATH=/opt/tritonserver/backends/pytorch:{env:LD_LIBRARY_PATH}
 passenv =
     OPAL_PREFIX
 deps =