Add hamilton.packaging for reusable, validatable dataflow subgraphs

docs/index.md

@@ -71,6 +71,7 @@ tryhamilton.dev <https://www.tryhamilton.dev/>
 reference/decorators/index
 reference/drivers/index
 reference/caching/index
+reference/packaging/index
 reference/graph-adapters/index
 reference/lifecycle-hooks/index
 reference/result-builders/index

docs/reference/packaging/index.rst

@@ -0,0 +1,81 @@
+==============
+Packaging
+==============
+
+Subgraph packaging for Hamilton dataflows. A :class:`hamilton.packaging.NodePackage`
+carves a group of nodes out of a driver's built internal graph (``Driver.graph``) and
+treats it as a first-class, reusable unit that records the nodes it provides, the
+external dependencies it requires, the configuration keys it needs, and the dependency
+edges among them. Packages serialize to JSON manifests for export and import, validate
+for compatibility against a host dataflow's graph, detect conflicts against other
+packages, and compose deterministically into larger packages. Packages supplied with
+``Builder.with_packages(...)`` are enforced when the driver is built.
+
+NodePackage
+-----------
+
+.. autoclass:: hamilton.packaging.NodePackage
+   :members:
+
+PackageManifest
+---------------
+
+.. autoclass:: hamilton.packaging.PackageManifest
+   :members:
+
+Validation results
+------------------
+
+.. autoclass:: hamilton.packaging.ValidationReport
+   :members:
+
+.. autoclass:: hamilton.packaging.ValidationProblem
+   :members:
+
+Conflicts
+---------
+
+.. autoclass:: hamilton.packaging.PackageConflict
+   :members:
+
+Diffs
+-----
+
+.. autoclass:: hamilton.packaging.ManifestDiff
+   :members:
+
+Errors
+------
+
+.. autoclass:: hamilton.packaging.PackagingError
+
+.. autoclass:: hamilton.packaging.PackageValidationError
+
+.. autoclass:: hamilton.packaging.PackageConflictError
+
+Module-level helpers
+--------------------
+
+.. autofunction:: hamilton.packaging.validate_package
+
+.. autofunction:: hamilton.packaging.find_conflicts
+
+.. autofunction:: hamilton.packaging.compose_packages
+
+.. autofunction:: hamilton.packaging.diff_manifests
+
+.. autofunction:: hamilton.packaging.export_package
+
+.. autofunction:: hamilton.packaging.import_manifest
+
+Driver integration
+------------------
+
+Packages integrate with the driver lifecycle:
+
+* ``Builder.with_packages(*node_packages)`` declares the packages a dataflow must be
+  compatible with.
+* Building a driver validates its graph against every declared package, raising
+  :class:`hamilton.packaging.PackageValidationError` when one is incompatible.
+* ``Driver.list_packages()`` returns the declared packages and
+  ``Driver.validate_packages()`` returns one report per package.

hamilton/async_driver.py

@@ -23,7 +23,7 @@
 from typing import Any
 
 import hamilton.lifecycle.base as lifecycle_base
-from hamilton import base, driver, graph, lifecycle, node
+from hamilton import base, driver, graph, lifecycle, node, packaging
 from hamilton.execution.graph_functions import create_error_message
 from hamilton.io.materialization import ExtractorFactory, MaterializerFactory
 
@@ -213,6 +213,7 @@ def __init__(
         result_builder: base.ResultMixin | None = None,
         adapters: list[lifecycle.LifecycleAdapter] = None,
         allow_module_overrides: bool = False,
+        _node_packages: typing.Sequence["packaging.NodePackage"] = None,
     ):
         """Instantiates an asynchronous driver.
 
@@ -231,6 +232,7 @@ def __init__(
             The order of listing the modules is important, since later ones will overwrite the previous ones.
             This is a global call affecting all imported modules.
             See https://github.com/apache/hamilton/tree/main/examples/module_overrides for more info.
+        :param _node_packages: Not public facing, do not use this parameter. This is injected by the builder.
         """
         if adapters is None:
             adapters = []
@@ -263,6 +265,7 @@ def __init__(
                 *async_adapters,  # note async adapters will not be called during synchronous execution -- this is for access later
             ],
             allow_module_overrides=allow_module_overrides,
+            _node_packages=_node_packages,
         )
         self.initialized = False
 
@@ -451,6 +454,7 @@ def build_without_init(self) -> AsyncDriver:
             adapters=self.adapters,
             result_builder=specified_result_builder,
             allow_module_overrides=self._allow_module_overrides,
+            _node_packages=self.node_packages,
         )
 
     async def build(self) -> AsyncDriver:

hamilton/driver.py

@@ -40,7 +40,7 @@
 import pandas as pd
 from typing_extensions import Self
 
-from hamilton import common, graph_types, htypes
+from hamilton import common, graph_types, htypes, packaging
 from hamilton.caching.adapter import HamiltonCacheAdapter
 from hamilton.caching.stores.base import MetadataStore, ResultStore
 from hamilton.dev_utils import deprecation
@@ -393,6 +393,22 @@ def _perform_graph_validations(
                 )
                 raise lifecycle_base.ValidationException(error_str)
 
+    @staticmethod
+    def _enforce_node_packages(
+        graph: graph.FunctionGraph,
+        node_packages: typing.Sequence["packaging.NodePackage"],
+    ):
+        """Validates the constructed graph against each supplied node package.
+
+        :param graph: Graph to validate against.
+        :param node_packages: Packages the driver should be compatible with.
+        :raises PackageValidationError: If a package is incompatible with the graph.
+        """
+        for node_package in node_packages:
+            report = packaging.validate_package(node_package, graph)
+            if not report.is_valid:
+                raise packaging.PackageValidationError(report)
+
     def __init__(
         self,
         config: dict[str, Any],
@@ -404,6 +420,7 @@ def __init__(
         _materializers: typing.Sequence[ExtractorFactory | MaterializerFactory] = None,
         _graph_executor: GraphExecutor = None,
         _use_legacy_adapter: bool = True,
+        _node_packages: typing.Sequence["packaging.NodePackage"] = None,
     ):
         """Constructor: creates a DAG given the configuration & modules to crawl.
 
@@ -422,6 +439,7 @@ def __init__(
         :param _use_legacy_adapter: Not public facing, do not use this parameter.
             This represents whether or not to use the legacy adapter. Defaults to True, as this should be
             backwards compatible. In Hamilton 2.0.0, this will be removed.
+        :param _node_packages: Not public facing, do not use this parameter. This is injected by the builder.
 
         """
 
@@ -445,6 +463,8 @@ def __init__(
                     self.graph, materializer_factories, extractor_factories
                 )
             Driver._perform_graph_validations(adapter, graph=self.graph, graph_modules=modules)
+            self._node_packages = list(_node_packages) if _node_packages else []
+            Driver._enforce_node_packages(self.graph, self._node_packages)
             if adapter.does_hook("post_graph_construct", is_async=False):
                 adapter.call_all_lifecycle_hooks_sync(
                     "post_graph_construct", graph=self.graph, modules=modules, config=config
@@ -793,6 +813,27 @@ def list_available_variables(
             results = [Variable.from_node(n) for n in all_nodes]
         return results
 
+    def list_packages(self) -> list["packaging.NodePackage"]:
+        """Returns the node packages this driver was built with.
+
+        :return: The packages supplied via ``Builder.with_packages`` (empty if none).
+        """
+        return list(self._node_packages)
+
+    def validate_packages(self) -> list["packaging.ValidationReport"]:
+        """Validates each of this driver's node packages against its built graph.
+
+        Since a driver only builds when its packages are compatible, this returns
+        problem-free reports for a successfully built driver. It is primarily useful
+        for inspecting the per-package validation result programmatically.
+
+        :return: One report per package, in the order the packages were supplied.
+        """
+        return [
+            packaging.validate_package(node_package, self.graph)
+            for node_package in self._node_packages
+        ]
+
     def get_variable(self, name: str) -> Variable:
         """Returns a variable by name.
 
@@ -1773,6 +1814,7 @@ def __init__(self):
         self.config = {}
         self.modules = []
         self.materializers = []
+        self.node_packages = []
 
         # Allow later modules to override nodes of the same name
         self._allow_module_overrides = False
@@ -1892,6 +1934,29 @@ def with_materializers(self, *materializers: ExtractorFactory | MaterializerFact
         self.materializers.extend(materializers)
         return self
 
+    def with_packages(self, *node_packages: "packaging.NodePackage") -> Self:
+        """Declares node packages the built dataflow must be compatible with.
+
+        Each package is validated against the driver's built graph when the driver is
+        built; building fails with a ``PackageValidationError`` if a package is
+        incompatible.
+
+        :param node_packages: packages the dataflow should satisfy
+        :return: self
+        """
+        if any(not isinstance(p, packaging.NodePackage) for p in node_packages):
+            if len(node_packages) == 1 and isinstance(node_packages[0], Sequence):
+                raise ValueError(
+                    "`.with_packages()` received a sequence. Unpack it by prepending `*` e.g., `*[package_a, package_b]`"
+                )
+            else:
+                raise ValueError(
+                    f"`.with_packages()` only accepts packages. Received instead: {node_packages}"
+                )
+
+        self.node_packages.extend(node_packages)
+        return self
+
     def with_cache(
         self,
         path: str | pathlib.Path = ".hamilton_cache",
@@ -2085,6 +2150,7 @@ def build(self) -> Driver:
             _graph_executor=graph_executor,
             _use_legacy_adapter=False,
             allow_module_overrides=self._allow_module_overrides,
+            _node_packages=self.node_packages,
         )
 
     def copy(self) -> "Builder":
@@ -2099,6 +2165,7 @@ def copy(self) -> "Builder":
         new_builder.legacy_graph_adapter = self.legacy_graph_adapter
         new_builder.adapters = self.adapters.copy()
         new_builder.materializers = self.materializers.copy()
+        new_builder.node_packages = self.node_packages.copy()
         new_builder.execution_manager = self.execution_manager
         new_builder.local_executor = self.local_executor
         new_builder.remote_executor = self.remote_executor