Skip to content

Commit b2af60d

Browse files
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>
1 parent 2d68307 commit b2af60d

File tree

1 file changed

+41
-0
lines changed

1 file changed

+41
-0
lines changed

merlin/dag/node.py

+41
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
# limitations under the License.
1515
#
1616
import collections.abc
17+
import os
1718
from typing import List, Union
1819

1920
from merlin.dag.operator import Operator
@@ -446,6 +447,46 @@ def exportable(self, backend: str = None):
446447
backends = getattr(self.op, "exportable_backends", [])
447448
return hasattr(self.op, "export") and backend in backends
448449

450+
def export(
451+
self,
452+
output_path: Union[str, os.PathLike],
453+
node_id: int = None,
454+
version: int = 1,
455+
):
456+
"""
457+
Export a directory for this node, containing the required artifacts
458+
to run in the target context.
459+
460+
Parameters
461+
----------
462+
output_path : Union[str, os.PathLike]
463+
The base path to write this node's export directory.
464+
node_id : int, optional
465+
The id of this node in a larger graph (for disambiguation), by default None.
466+
version : int, optional
467+
The version of the node to use for this export, by default 1.
468+
469+
"""
470+
return self.op.export(
471+
output_path,
472+
self.input_schema,
473+
self.output_schema,
474+
node_id=node_id,
475+
version=version,
476+
)
477+
478+
@property
479+
def export_name(self):
480+
"""
481+
Name for the exported node directory.
482+
483+
Returns
484+
-------
485+
str
486+
Name supplied by this node's operator.
487+
"""
488+
return self.op.export_name
489+
449490
@property
450491
def parents_with_dependencies(self):
451492
nodes = []

0 commit comments

Comments
 (0)