Use Runtimes and Executors to serve NVT Workflows

merlin/systems/dag/ensemble.py

@@ -23,6 +23,7 @@
 
 from merlin.core.protocols import Transformable
 from merlin.dag import Graph
+from merlin.systems.dag.runtimes.base_runtime import Runtime
 from merlin.systems.dag.runtimes.triton import TritonExecutorRuntime
 
 
@@ -74,7 +75,7 @@ def transform(self, transformable: Transformable, runtime=None):
         Transformable
             transformed data
         """
-        runtime = runtime or TritonExecutorRuntime()
+        runtime = runtime or Runtime()
         return runtime.transform(self.graph, transformable)
 
     def save(self, path):

merlin/systems/dag/runtimes/__init__.py

@@ -13,6 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-# flake8: noqa
 
+# flake8: noqa
 from .base_runtime import Runtime

merlin/systems/dag/runtimes/base_runtime.py

@@ -14,12 +14,14 @@
 # limitations under the License.
 #
 from merlin.core.protocols import Transformable
-from merlin.dag import Graph
+from merlin.dag import Graph, postorder_iter_nodes
 from merlin.dag.executors import LocalExecutor
+from merlin.systems.dag.runtimes.op_table import OpTable
 
 
 class Runtime:
-    """A Systems Graph Runtime.
+    """
+    A runtime for Merlin DAGs that supports using custom implementations of existing operators
 
     This class can be used as a base class for custom runtimes.
     """
@@ -33,9 +35,32 @@ def __init__(self, executor=None):
             The Graph Executor to use to use for the transform, by default None
         """
         self.executor = executor or LocalExecutor()
-        self.op_table = {}
+        self.op_table = OpTable()
+
+    def convert(self, graph: Graph):
+        """
+        Replace the operators in the supplied graph with ops from this runtime's op table
+
+        Parameters
+        ----------
+        graph : Graph
+            Graph of nodes container operator chains for data manipulation.
 
-    def transform(self, graph: Graph, transformable: Transformable):
+        Returns
+        -------
+        Graph
+            Copy of the graph with operators converted to this runtime's versions
+        """
+        if not self.op_table.empty:
+            nodes = list(postorder_iter_nodes(graph.output_node))
+
+            for node in nodes:
+                if self.op_table.has_impl(node.op):
+                    node.op = self.op_table.replace(node.op)
+
+        return graph
+
+    def transform(self, graph: Graph, transformable: Transformable, convert=True):
         """Run the graph with the input data.
 
         Parameters
@@ -44,12 +69,17 @@ def transform(self, graph: Graph, transformable: Transformable):
             Graph of nodes container operator chains for data manipulation.
         transformable : Transformable
             Input data to transform in graph.
+        convert: bool
+            If True, converts the operators in the graph to this runtime's versions
 
         Returns
         -------
         Transformable
             Input data after it has been transformed via graph.
         """
+        if convert:
+            graph = self.convert(graph)
+
         return self.executor.transform(transformable, [graph.output_node])
 
     def export(self):

merlin/systems/dag/runtimes/nvtabular/__init__.py

@@ -0,0 +1,15 @@
+#
+# Copyright (c) 2023, 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.
+#