From b2af60d0bbf8a0f3ce3e8ed38065daadc301d10d Mon Sep 17 00:00:00 2001
From: Julio Perez <37191411+jperez999@users.noreply.github.com>
Date: Thu, 13 Jul 2023 11:54:57 -0400
Subject: [PATCH] Merge `InferenceNode` functionality from Systems into the
 base `Node` class (#357)

* migrate inf operator to base operator

* migrate inf node to node

* remove model registry from core and make export a no op by default

* remove mention of triton in base op and node code

* change model config to config.

* removed the return for export docs

* remove unnecessary import

---------

Co-authored-by: Karl Higley <karlb@nvidia.com>
---
 merlin/dag/node.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/merlin/dag/node.py b/merlin/dag/node.py
index e99046143..153d0ab3f 100644
--- a/merlin/dag/node.py
+++ b/merlin/dag/node.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 #
 import collections.abc
+import os
 from typing import List, Union
 
 from merlin.dag.operator import Operator
@@ -446,6 +447,46 @@ def exportable(self, backend: str = None):
         backends = getattr(self.op, "exportable_backends", [])
         return hasattr(self.op, "export") and backend in backends
 
+    def export(
+        self,
+        output_path: Union[str, os.PathLike],
+        node_id: int = None,
+        version: int = 1,
+    ):
+        """
+        Export a directory for this node, containing the required artifacts
+        to run in the target context.
+
+        Parameters
+        ----------
+        output_path : Union[str, os.PathLike]
+            The base path to write this node's export directory.
+        node_id : int, optional
+            The id of this node in a larger graph (for disambiguation), by default None.
+        version : int, optional
+            The version of the node to use for this export, by default 1.
+
+        """
+        return self.op.export(
+            output_path,
+            self.input_schema,
+            self.output_schema,
+            node_id=node_id,
+            version=version,
+        )
+
+    @property
+    def export_name(self):
+        """
+        Name for the exported node directory.
+
+        Returns
+        -------
+        str
+            Name supplied by this node's operator.
+        """
+        return self.op.export_name
+
     @property
     def parents_with_dependencies(self):
         nodes = []