moved all triton tests into runtime triton folder

tests/unit/systems/dag/runtimes/local_runtime/__init__.py → tests/unit/systems/dag/runtimes/triton/ops/faiss/__init__.py

@@ -13,3 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import pytest
+
+pytest.importorskip("faiss")

tests/unit/systems/dag/runtimes/local_runtime/ops/__init__.py → tests/unit/systems/dag/runtimes/triton/ops/implicit/__init__.py

@@ -13,3 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+
+import pytest
+
+pytest.importorskip("implicit")

tests/unit/systems/dag/runtimes/triton/ops/implicit/test_op.py

@@ -0,0 +1,85 @@
+#
+# Copyright (c) 2022, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from distutils.spawn import find_executable
+
+import implicit
+import numpy as np
+import pytest
+from scipy.sparse import csr_matrix
+from tritonclient import grpc as grpcclient
+
+from merlin.schema import ColumnSchema, Schema
+from merlin.systems.dag.ensemble import Ensemble
+from merlin.systems.dag.ops.implicit import PredictImplicit
+from merlin.systems.dag.runtimes.triton import TritonExecutorRuntime
+from merlin.systems.triton.utils import run_triton_server
+
+TRITON_SERVER_PATH = find_executable("tritonserver")
+
+triton = pytest.importorskip("merlin.systems.triton")
+
+
+@pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found")
+@pytest.mark.parametrize("runtime", [None, TritonExecutorRuntime()])
+@pytest.mark.parametrize(
+    "model_cls",
+    [
+        implicit.bpr.BayesianPersonalizedRanking,
+        implicit.als.AlternatingLeastSquares,
+        implicit.lmf.LogisticMatrixFactorization,
+    ],
+)
+def test_ensemble(model_cls, runtime, tmpdir):
+    model = model_cls()
+    n = 100
+    user_items = csr_matrix(np.random.choice([0, 1], size=n * n, p=[0.9, 0.1]).reshape(n, n))
+    model.fit(user_items)
+
+    num_to_recommend = np.random.randint(1, n)
+
+    user_items = None
+    ids, scores = model.recommend(
+        [0, 1], user_items, N=num_to_recommend, filter_already_liked_items=False
+    )
+
+    implicit_op = PredictImplicit(model, num_to_recommend=num_to_recommend)
+
+    input_schema = Schema([ColumnSchema("user_id", dtype="int64")])
+
+    triton_chain = input_schema.column_names >> implicit_op
+
+    triton_ens = Ensemble(triton_chain, input_schema)
+    ensemble_config, _ = triton_ens.export(tmpdir, runtime=runtime)
+
+    input_user_id = np.array([[0], [1]], dtype=np.int64)
+    inputs = [
+        grpcclient.InferInput(
+            "user_id", input_user_id.shape, triton.np_to_triton_dtype(input_user_id.dtype)
+        ),
+    ]
+    inputs[0].set_data_from_numpy(input_user_id)
+    outputs = [grpcclient.InferRequestedOutput("scores"), grpcclient.InferRequestedOutput("ids")]
+
+    response = None
+
+    with run_triton_server(tmpdir) as client:
+        response = client.infer(ensemble_config.name, inputs, outputs=outputs)
+
+    response_ids = response.as_numpy("ids")
+    response_scores = response.as_numpy("scores")
+
+    np.testing.assert_array_equal(ids, response_ids)
+    np.testing.assert_array_equal(scores, response_scores)

tests/unit/systems/dag/runtimes/local_runtime/ops/torch/__init__.py → tests/unit/systems/dag/runtimes/triton/ops/tf/__init__.py

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2021, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,3 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+
+import pytest
+
+pytest.importorskip("tensorflow")

tests/unit/systems/dag/test_executors.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 #
 import random
-from distutils.spawn import find_executable
 
 import numpy as np
 import pandas as pd
@@ -28,8 +27,6 @@
 from merlin.systems.dag.ensemble import Ensemble
 from merlin.systems.dag.ops.session_filter import FilterCandidates
 
-TRITON_SERVER_PATH = find_executable("tritonserver")
-
 
 def test_run_dag_on_dictarray_with_local_executor():
     request_schema = Schema(

tests/unit/systems/ops/implicit/test_op.py

@@ -14,23 +14,14 @@
 # limitations under the License.
 #
 import json
-from distutils.spawn import find_executable
 
 import implicit
 import numpy as np
 import pytest
 from scipy.sparse import csr_matrix
-from tritonclient import grpc as grpcclient
 
-from merlin.schema import ColumnSchema, Schema
-from merlin.systems.dag.ensemble import Ensemble
+from merlin.schema import Schema
 from merlin.systems.dag.ops.implicit import PredictImplicit
-from merlin.systems.dag.runtimes.triton import TritonExecutorRuntime
-from merlin.systems.triton.utils import run_triton_server
-
-TRITON_SERVER_PATH = find_executable("tritonserver")
-
-triton = pytest.importorskip("merlin.systems.triton")
 
 
 @pytest.mark.parametrize(
@@ -73,56 +64,3 @@ def test_reload_from_config(model_cls, tmpdir):
 
     np.testing.assert_array_equal(ids, reloaded_ids)
     np.testing.assert_array_equal(scores, reloaded_scores)
-
-
-@pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found")
-@pytest.mark.parametrize("runtime", [None, TritonExecutorRuntime()])
-@pytest.mark.parametrize(
-    "model_cls",
-    [
-        implicit.bpr.BayesianPersonalizedRanking,
-        implicit.als.AlternatingLeastSquares,
-        implicit.lmf.LogisticMatrixFactorization,
-    ],
-)
-def test_ensemble(model_cls, runtime, tmpdir):
-    model = model_cls()
-    n = 100
-    user_items = csr_matrix(np.random.choice([0, 1], size=n * n, p=[0.9, 0.1]).reshape(n, n))
-    model.fit(user_items)
-
-    num_to_recommend = np.random.randint(1, n)
-
-    user_items = None
-    ids, scores = model.recommend(
-        [0, 1], user_items, N=num_to_recommend, filter_already_liked_items=False
-    )
-
-    implicit_op = PredictImplicit(model, num_to_recommend=num_to_recommend)
-
-    input_schema = Schema([ColumnSchema("user_id", dtype="int64")])
-
-    triton_chain = input_schema.column_names >> implicit_op
-
-    triton_ens = Ensemble(triton_chain, input_schema)
-    ensemble_config, _ = triton_ens.export(tmpdir, runtime=runtime)
-
-    input_user_id = np.array([[0], [1]], dtype=np.int64)
-    inputs = [
-        grpcclient.InferInput(
-            "user_id", input_user_id.shape, triton.np_to_triton_dtype(input_user_id.dtype)
-        ),
-    ]
-    inputs[0].set_data_from_numpy(input_user_id)
-    outputs = [grpcclient.InferRequestedOutput("scores"), grpcclient.InferRequestedOutput("ids")]
-
-    response = None
-
-    with run_triton_server(tmpdir) as client:
-        response = client.infer(ensemble_config.name, inputs, outputs=outputs)
-
-    response_ids = response.as_numpy("ids")
-    response_scores = response.as_numpy("scores")
-
-    np.testing.assert_array_equal(ids, response_ids)
-    np.testing.assert_array_equal(scores, response_scores)