Unify the node and dataflow versioning API

hamilton/cli/logic.py

@@ -54,22 +54,15 @@ def version_hamilton_functions(module: ModuleType) -> Dict[str, str]:
 
     for origin_name, _ in graph_utils.find_functions(module):
         origin_callable = getattr(module, origin_name)
-        origins_version[origin_name] = graph_utils.hash_source_code(origin_callable, strip=True)
+        origins_version[origin_name] = graph_types.hash_source_code(origin_callable, strip=True)
 
     return origins_version
 
 
 def hash_hamilton_nodes(dr: driver.Driver) -> Dict[str, str]:
     """Hash the source code of Hamilton functions from nodes in a Driver"""
     graph = graph_types.HamiltonGraph.from_graph(dr.graph)
-
-    nodes_version = dict()
-    for n in graph.nodes:
-        node_origin = n.originating_functions[0]
-        origin_hash = graph_utils.hash_source_code(node_origin, strip=True)
-        nodes_version[n.name] = origin_hash
-
-    return nodes_version
+    return {n.name: n.version for n in graph.nodes}
 
 
 def map_nodes_to_functions(dr: driver.Driver) -> Dict[str, str]:

hamilton/graph_types.py

@@ -1,5 +1,7 @@
 """Module for external-facing graph constructs. These help the user navigate/manage the graph as needed."""
 
+import ast
+import hashlib
 import inspect
 import typing
 from dataclasses import dataclass
@@ -16,6 +18,70 @@
     from hamilton import graph
 
 
+def _remove_docs_and_comments(source: str) -> str:
+    """Remove the docs and comments from a source code string.
+
+    The use of `ast.unparse()` requires Python 3.9
+
+    1. Parsing then unparsing the AST of the source code will
+    create a code object and convert it back to a string. In the
+    process, comments are stripped.
+
+    2. walk the AST to check if first element after `def` is a
+    docstring. If so, edit AST to skip the docstring
+
+    NOTE. The ast parsing will fail if `source` has syntax errors. For the
+    majority of cases this is caught upstream (e.g., by calling `import`).
+    The foreseeable edge case is if `source` is the result of `inspect.getsource`
+    on a nested function, method, or callable where `def` isn't at column 0.
+    Standard usage of Hamilton requires users to define functions/nodes at the top
+    level of a module, and therefore no issues should arise.
+    """
+    parsed = ast.parse(source)
+    for n in ast.walk(parsed):
+        if not isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            continue
+
+        if not len(n.body):
+            continue
+
+        # check if 1st node is a docstring
+        if not isinstance(n.body[0], ast.Expr):
+            continue
+
+        if not hasattr(n.body[0], "value") or not isinstance(n.body[0].value, ast.Str):
+            continue
+
+        # skip docstring
+        n.body = n.body[1:]
+
+    return ast.unparse(parsed)
+
+
+def hash_source_code(source: typing.Union[str, typing.Callable], strip: bool = False) -> str:
+    """Hashes the source code of a function (str).
+
+    The `strip` parameter requires Python 3.9
+
+    If strip, try to remove docs and comments from source code string. Since
+    they don't impact function behavior, they shouldn't influence the hash.
+    """
+    if isinstance(source, typing.Callable):
+        source = inspect.getsource(source)
+
+    source = source.strip()
+
+    if strip:
+        try:
+            # could fail if source is indented code.
+            # see `remove_docs_and_comments` docstring for details.
+            source = _remove_docs_and_comments(source)
+        except Exception:
+            pass
+
+    return hashlib.sha256(source.encode()).hexdigest()
+
+
 @dataclass
 class HamiltonNode:
     """External facing API for hamilton Nodes. Having this as a dataclass allows us
@@ -45,6 +111,7 @@ def as_dict(self):
                 else None
             ),
             "documentation": self.documentation,
+            "version": self.version,
         }
 
     @staticmethod
@@ -73,6 +140,15 @@ def from_node(n: node.Node) -> "HamiltonNode":
             },
         )
 
+    @property
+    def version(self) -> str:
+        """Generate a hash of the node originating function source code.
+
+        The option `strip=True` means docstring and comments are ignored
+        when hashing the function.
+        """
+        return hash_source_code(self.originating_functions[0], strip=True)
+
     def __repr__(self):
         return f"{self.name}: {htypes.get_type_as_string(self.type)}"
 
@@ -101,3 +177,13 @@ def from_graph(fn_graph: "graph.FunctionGraph") -> "HamiltonGraph":
         return HamiltonGraph(
             nodes=[HamiltonNode.from_node(n) for n in fn_graph.nodes.values()],
         )
+
+    @property
+    def version(self) -> str:
+        """Generate a hash of the dataflow based on the collection of node hashes.
+
+        Node hashes are in a sorted list, then concatenated as a string before hashing.
+        To find differences between dataflows, you need to inspect the node level.
+        """
+        sorted_node_versions = sorted([n.version for n in self.nodes])
+        return hashlib.sha256(str(sorted_node_versions).encode()).hexdigest()

hamilton/graph_utils.py

@@ -1,8 +1,6 @@
-import ast
-import hashlib
 import inspect
 from types import ModuleType
-from typing import Callable, List, Tuple, Union
+from typing import Callable, List, Tuple
 
 
 def is_submodule(child: ModuleType, parent: ModuleType):
@@ -24,67 +22,3 @@ def valid_fn(fn):
         )
 
     return [f for f in inspect.getmembers(function_module, predicate=valid_fn)]
-
-
-def hash_source_code(source: Union[str, Callable], strip: bool = False) -> str:
-    """Hashes the source code of a function (str).
-
-    The `strip` parameter requires Python 3.9
-
-    If strip, try to remove docs and comments from source code string. Since
-    they don't impact function behavior, they shouldn't influence the hash.
-    """
-    if isinstance(source, Callable):
-        source = inspect.getsource(source)
-
-    source = source.strip()
-
-    if strip:
-        try:
-            # could fail if source is indented code.
-            # see `remove_docs_and_comments` docstring for details.
-            source = remove_docs_and_comments(source)
-        except Exception:
-            pass
-
-    return hashlib.sha256(source.encode()).hexdigest()
-
-
-def remove_docs_and_comments(source: str) -> str:
-    """Remove the docs and comments from a source code string.
-
-    The use of `ast.unparse()` requires Python 3.9
-
-    1. Parsing then unparsing the AST of the source code will
-    create a code object and convert it back to a string. In the
-    process, comments are stripped.
-
-    2. walk the AST to check if first element after `def` is a
-    docstring. If so, edit AST to skip the docstring
-
-    NOTE. The ast parsing will fail if `source` has syntax errors. For the
-    majority of cases this is caught upstream (e.g., by calling `import`).
-    The foreseeable edge case is if `source` is the result of `inspect.getsource`
-    on a nested function, method, or callable where `def` isn't at column 0.
-    Standard usage of Hamilton requires users to define functions/nodes at the top
-    level of a module, and therefore no issues should arise.
-    """
-    parsed = ast.parse(source)
-    for node in ast.walk(parsed):
-        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
-            continue
-
-        if not len(node.body):
-            continue
-
-        # check if 1st node is a docstring
-        if not isinstance(node.body[0], ast.Expr):
-            continue
-
-        if not hasattr(node.body[0], "value") or not isinstance(node.body[0].value, ast.Str):
-            continue
-
-        # skip docstring
-        node.body = node.body[1:]
-
-    return ast.unparse(parsed)

hamilton/plugins/h_diskcache.py

@@ -3,7 +3,7 @@
 
 import diskcache
 
-from hamilton import driver, graph_types, graph_utils, lifecycle, node
+from hamilton import driver, graph_types, lifecycle, node
 
 logger = logging.getLogger(__name__)
 
@@ -26,7 +26,7 @@ def evict_all_except(nodes_to_keep: Dict[str, node.Node], cache: diskcache.Cache
 
         if node_name in nodes_to_keep.keys():
             node_to_keep = nodes_to_keep[node_name]
-            hash_to_keep = graph_utils.hash_source_code(node_to_keep.callable, strip=True)
+            hash_to_keep = graph_types.hash_source_code(node_to_keep.callable, strip=True)
             history.remove(hash_to_keep)
             new_nodes_history[node_name] = [hash_to_keep]
 
@@ -99,7 +99,7 @@ def run_to_execute_node(
         if node_name not in self.cache_vars:
             return node_callable(**node_kwargs)
 
-        node_hash = graph_utils.hash_source_code(node_callable, strip=True)
+        node_hash = graph_types.hash_source_code(node_callable, strip=True)
         self.used_nodes_hash[node_name] = node_hash
         cache_key = (node_hash, *node_kwargs.values())
 

hamilton/plugins/h_experiments/hook.py

@@ -50,21 +50,6 @@ def json_encoder(obj: Any):
     return serialized
 
 
-def graph_hash(graph: graph_types.HamiltonGraph) -> str:
-    """Create a single hash (str) from the bytecode of all sorted functions"""
-    nodes_data = []
-    for node in graph.nodes:
-        source_code = ""
-        if node.originating_functions is not None:
-            source_code = inspect.getsource(node.originating_functions[0])
-
-        nodes_data.append(dict(name=node.name, source_code=source_code))
-
-    digest = hashlib.sha256()
-    digest.update(json.dumps(nodes_data, default=json_encoder, sort_keys=True).encode())
-    return digest.hexdigest()
-
-
 @dataclass
 class NodeImplementation:
     name: str
@@ -147,7 +132,7 @@ def run_before_graph_execution(
         **kwargs,
     ):
         """Store execution metadata: graph hash, inputs, overrides"""
-        self.graph_hash = graph_hash(graph)
+        self.graph_hash = graph.version
 
         for node in graph.nodes:
             if node.tags.get("module"):

tests/plugins/test_h_diskcache.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from hamilton import ad_hoc_utils, driver, graph_utils, node
+from hamilton import ad_hoc_utils, driver, graph_types, node
 from hamilton.plugins import h_diskcache
 
 
@@ -53,7 +53,7 @@ def A(external_input: int) -> int:
 
 def test_set_result(hook: h_diskcache.DiskCacheAdapter, node_a: node.Node):
     """Hook sets value and assert value in cache"""
-    node_hash = graph_utils.hash_source_code(node_a.callable, strip=True)
+    node_hash = graph_types.hash_source_code(node_a.callable, strip=True)
     node_kwargs = dict(external_input=7)
     cache_key = (node_hash, *node_kwargs.values())
     result = 2
@@ -73,7 +73,7 @@ def test_set_result(hook: h_diskcache.DiskCacheAdapter, node_a: node.Node):
 
 def test_get_result(hook: h_diskcache.DiskCacheAdapter, node_a: node.Node):
     """Hooks get value and assert cache hit"""
-    node_hash = graph_utils.hash_source_code(node_a.callable, strip=True)
+    node_hash = graph_types.hash_source_code(node_a.callable, strip=True)
     node_kwargs = dict(external_input=7)
     result = 2
     cache_key = (node_hash, *node_kwargs.values())
@@ -105,7 +105,7 @@ def test_append_nodes_history(
     hook.cache_vars = [node_a.name]
 
     # node version 1
-    hook.used_nodes_hash[node_name] = graph_utils.hash_source_code(node_a.callable, strip=True)
+    hook.used_nodes_hash[node_name] = graph_types.hash_source_code(node_a.callable, strip=True)
     hook.run_to_execute_node(
         node_name=node_name,
         node_kwargs=node_kwargs,
@@ -116,7 +116,7 @@ def test_append_nodes_history(
     assert len(hook.nodes_history.get(node_name, [])) == 1
 
     # node version 2
-    hook.used_nodes_hash[node_name] = graph_utils.hash_source_code(node_a_body.callable, strip=True)
+    hook.used_nodes_hash[node_name] = graph_types.hash_source_code(node_a_body.callable, strip=True)
     hook.run_to_execute_node(
         node_name=node_name,
         node_kwargs=node_kwargs,
@@ -141,8 +141,8 @@ def test_evict_all_except(
     node_a_body: node.Node,
 ):
     """Check utility function to evict all except passed nodes"""
-    node_a_hash = graph_utils.hash_source_code(node_a.callable, strip=True)
-    node_a_body_hash = graph_utils.hash_source_code(node_a_body.callable, strip=True)
+    node_a_hash = graph_types.hash_source_code(node_a.callable, strip=True)
+    node_a_body_hash = graph_types.hash_source_code(node_a_body.callable, strip=True)
     hook.cache[h_diskcache.DiskCacheAdapter.nodes_history_key] = dict(
         A=[node_a_hash, node_a_body_hash]
     )
@@ -159,8 +159,8 @@ def test_evict_from_driver(
     node_a_body: node.Node,
 ):
     """Check utility function to evict all except driver"""
-    node_a_hash = graph_utils.hash_source_code(node_a.callable, strip=True)
-    node_a_body_hash = graph_utils.hash_source_code(node_a_body.callable, strip=True)
+    node_a_hash = graph_types.hash_source_code(node_a.callable, strip=True)
+    node_a_body_hash = graph_types.hash_source_code(node_a_body.callable, strip=True)
     hook.cache[h_diskcache.DiskCacheAdapter.nodes_history_key] = dict(
         A=[node_a_hash, node_a_body_hash]
     )