Adjust Node.exportable() to catch up to Merlin Core (#209)

* Adjust `Node.exportable()` to catch up to Merlin Core

As we prepare to export and/or run Merlin graphs in multiple contexts, the signature of `Node.exportable()` has changed a little. Further related changes are coming in #204, but these small changes make the tests pass with the latest `main` version of Core.

* Adjust `tox.ini` to run the tests against latest version of Merlin Core

Author: karlhigley

merlin/systems/dag/ensemble.py

@@ -66,10 +66,12 @@ def export(self, export_path, version=1):
         Write out an ensemble model configuration directory. The exported
         ensemble is designed for use with Triton Inference Server.
         """
+        backend = "ensemble"
+
         # Create ensemble config
         ensemble_config = model_config.ModelConfig(
             name=self.name,
-            platform="ensemble",
+            platform=backend,
             # max_batch_size=configs[0].max_batch_size
         )
 
@@ -95,14 +97,14 @@ def export(self, export_path, version=1):
         node_idx = 0
         node_id_lookup = {}
         for node in postorder_nodes:
-            if node.exportable:
+            if node.exportable(backend):
                 node_id_lookup[node] = node_idx
                 node_idx += 1
 
         node_configs = []
         # Export node configs and add ensemble steps
         for node in postorder_nodes:
-            if node.exportable:
+            if node.exportable(backend):
                 node_id = node_id_lookup.get(node, None)
                 node_name = f"{node_id}_{node.export_name}"
 
@@ -120,7 +122,9 @@ def export(self, export_path, version=1):
                 )
 
                 for input_col_name, input_col_schema in node.input_schema.column_schemas.items():
-                    source = _find_column_source(node.parents_with_dependencies, input_col_name)
+                    source = _find_column_source(
+                        node.parents_with_dependencies, input_col_name, backend
+                    )
                     source_id = node_id_lookup.get(source, None)
                     in_suffix = f"_{source_id}" if source_id is not None else ""
                     if input_col_schema.is_list and input_col_schema.is_ragged:
@@ -163,14 +167,14 @@ def export(self, export_path, version=1):
         return (ensemble_config, node_configs)
 
 
-def _find_column_source(upstream_nodes, column_name):
+def _find_column_source(upstream_nodes, column_name, backend):
     source_node = None
     for upstream_node in upstream_nodes:
         if column_name in upstream_node.output_columns.names:
             source_node = upstream_node
             break
 
-    if source_node and not source_node.exportable:
-        return _find_column_source(source_node.parents_with_dependencies, column_name)
+    if source_node and not source_node.exportable(backend):
+        return _find_column_source(source_node.parents_with_dependencies, column_name, backend)
     else:
         return source_node

merlin/systems/dag/node.py

@@ -23,6 +23,25 @@
 class InferenceNode(Node):
     """Specialized node class used in Triton Ensemble DAGs"""
 
+    def exportable(self, backend: str = None):
+        """
+        Determine whether the current node's operator is exportable for a given back-end
+
+        Parameters
+        ----------
+        backend : str, optional
+            The Merlin Systems (not Triton) back-end to use,
+            either "ensemble" or "executor", by default None
+
+        Returns
+        -------
+        bool
+            True if the node's operator is exportable for the supplied back-end
+        """
+        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 Triton config directory for this node.
@@ -58,6 +77,24 @@ def export_name(self):
         return self.op.export_name
 
     def validate_schemas(self, root_schema, strict_dtypes=False):
+        """
+        Checks that the output schema is valid given the previous
+        nodes in the graph and following nodes in the graph, as
+        well as any additional root inputs.
+
+        Parameters
+        ----------
+        root_schema : Schema
+            Schema of selection from the original data supplied
+        strict_dtypes : bool, optional
+            If True, raises an error when the dtypes in the input data
+            do not match the dtypes in the schema, by default False
+
+        Raises
+        ------
+        ValueError
+            If an output column is produced but not used by child nodes
+        """
         super().validate_schemas(root_schema, strict_dtypes)
 
         if self.children:

merlin/systems/dag/ops/operator.py

@@ -71,6 +71,10 @@ def export_name(self):
         """
         return self.__class__.__name__.lower()
 
+    @property
+    def exportable_backends(self):
+        return ["ensemble"]
+
     @abstractmethod
     def export(
         self,

tox.ini

@@ -14,7 +14,9 @@ commands =
 ; Runs all CPU-based tests. NOTE: if you are using an M1 mac, this will fail. You need to
 ; change the tensorflow dependency to `tensorflow-macos` in requirements/test-cpu.txt.
 deps = -rrequirements/test-cpu.txt
-commands = python -m pytest --cov-report term --cov=merlin -rxs tests/unit
+commands =
+    python -m pip install --upgrade git+https://github.com/NVIDIA-Merlin/core.git
+    python -m pytest --cov-report term --cov=merlin -rxs tests/unit
 
 [testenv:test-gpu]
 sitepackages=true
@@ -28,6 +30,7 @@ deps =
     pytest
     pytest-cov
 commands =
+    python -m pip install --upgrade git+https://github.com/NVIDIA-Merlin/core.git
     python -m pytest --cov-report term --cov merlin -rxs tests/unit
 
 [testenv:test-merlin]
@@ -57,6 +60,7 @@ commands =
 ; Install pre-commit-hooks to run these tests during development.
 deps = -rrequirements/dev.txt
 commands =
+    python -m pip install --upgrade git+https://github.com/NVIDIA-Merlin/core.git
     flake8 setup.py merlin/ tests/
     black --check --diff merlin tests
     pylint merlin