Refactor to enable RayGraphAdapter and HamiltonTracker to work well together

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
       args: [ --fix ]
     # Run the formatter.
     - id: ruff-format
-#      args: [ --diff ]  # Use for previewing changes
+      # args: [ --diff ]  # Use for previewing changes
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.6.0
     hooks:

hamilton/dev_utils/deprecation.py

@@ -48,8 +48,8 @@ class deprecated:
     @deprecate(
         warn_starting=(1,10,0)
         fail_starting=(2,0,0),
-        use_instead=parameterize_values,
-        reason='We have redefined the parameterization decorators to consist of `parametrize`, `parametrize_inputs`, and `parametrize_values`
+        use_this=parameterize_values,
+        explanation='We have redefined the parameterization decorators to consist of `parametrize`, `parametrize_inputs`, and `parametrize_values`
         migration_guide="https://github.com/dagworks-inc/hamilton/..."
     )
     class parameterized(...):
@@ -66,7 +66,7 @@ class parameterized(...):
     explanation: str
     migration_guide: Optional[
         str
-    ]  # If this is None, this means that the use_instead is a drop in replacement
+    ]  # If this is None, this means that the use_this is a drop in replacement
     current_version: Union[Tuple[int, int, int], Version] = dataclasses.field(
         default_factory=lambda: CURRENT_VERSION
     )

hamilton/driver.py

@@ -19,6 +19,7 @@
 import pandas as pd
 
 from hamilton import common, graph_types, htypes
+from hamilton.dev_utils import deprecation
 from hamilton.execution import executors, graph_functions, grouping, state
 from hamilton.graph_types import HamiltonNode
 from hamilton.io import materialization
@@ -580,10 +581,12 @@ def execute(
             )
         start_time = time.time()
         run_successful = True
-        error = None
+        telemetry_error = None
+        execution_error = None
+        outputs = None
         _final_vars = self._create_final_vars(final_vars)
         try:
-            outputs = self.raw_execute(_final_vars, overrides, display_graph, inputs=inputs)
+            outputs = self.__raw_execute(_final_vars, overrides, display_graph, inputs=inputs)
             if self.adapter.does_method("do_build_result", is_async=False):
                 # Build the result if we have a result builder
                 return self.adapter.call_lifecycle_method_sync("do_build_result", outputs=outputs)
@@ -592,12 +595,22 @@ def execute(
         except Exception as e:
             run_successful = False
             logger.error(SLACK_ERROR_MESSAGE)
-            error = telemetry.sanitize_error(*sys.exc_info())
+            execution_error = e
+            telemetry_error = telemetry.sanitize_error(*sys.exc_info())
             raise e
         finally:
+            if self.adapter.does_hook("post_graph_execute", is_async=False):
+                self.adapter.call_all_lifecycle_hooks_sync(
+                    "post_graph_execute",
+                    run_id=self.run_id,
+                    graph=self.function_graph,
+                    success=run_successful,
+                    error=execution_error,
+                    results=outputs,
+                )
             duration = time.time() - start_time
             self.capture_execute_telemetry(
-                error, _final_vars, inputs, overrides, run_successful, duration
+                telemetry_error, _final_vars, inputs, overrides, run_successful, duration
             )
 
     def _create_final_vars(self, final_vars: List[Union[str, Callable, Variable]]) -> List[str]:
@@ -649,6 +662,13 @@ def capture_execute_telemetry(
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(f"Error caught in processing telemetry: \n{e}")
 
+    @deprecation.deprecated(
+        warn_starting=(1, 75, 0),
+        fail_starting=(2, 0, 0),
+        use_this=None,
+        explanation="This has become a private method and does not guarantee that all the adapters work correctly.",
+        migration_guide="Don't use this entry point for execution directly. Always go through `.execute()`.",
+    )
     def raw_execute(
         self,
         final_vars: List[str],
@@ -657,24 +677,55 @@ def raw_execute(
         inputs: Dict[str, Any] = None,
         _fn_graph: graph.FunctionGraph = None,
     ) -> Dict[str, Any]:
-        """Raw execute function that does the meat of execute.
-
-        Don't use this entry point for execution directly. Always go through `.execute()`.
+        """Don't use this entry point for execution directly. Always go through `.execute()`.
         In case you are using `.raw_execute()` directly, please switch to `.execute()` using a
         `base.DictResult()`. Note: `base.DictResult()` is the default return of execute if you are
         using the `driver.Builder()` class to create a `Driver()` object.
+        """
+        success = True
+        error = None
+        results = None
+        try:
+            return self.__raw_execute(final_vars, overrides, display_graph, inputs=inputs)
+        except Exception as e:
+            success = False
+            logger.error(SLACK_ERROR_MESSAGE)
+            error = e
+            raise e
+        finally:
+            if self.adapter.does_hook("post_graph_execute", is_async=False):
+                self.adapter.call_all_lifecycle_hooks_sync(
+                    "post_graph_execute",
+                    run_id=self.run_id,
+                    graph=self.function_graph,
+                    success=success,
+                    error=error,
+                    results=results,
+                )
+
+    def __raw_execute(
+        self,
+        final_vars: List[str],
+        overrides: Dict[str, Any] = None,
+        display_graph: bool = False,
+        inputs: Dict[str, Any] = None,
+        _fn_graph: graph.FunctionGraph = None,
+    ) -> Dict[str, Any]:
+        """Raw execute function that does the meat of execute.
+
+        Private method since the result building and post_graph_execute lifecycle hooks are performed outside and so this returns an incomplete result.
 
         :param final_vars: Final variables to compute
         :param overrides: Overrides to run.
         :param display_graph: DEPRECATED. DO NOT USE. Whether or not to display the graph when running it
         :param inputs: Runtime inputs to the DAG
         :return:
         """
-        function_graph = _fn_graph if _fn_graph is not None else self.graph
-        run_id = str(uuid.uuid4())
-        nodes, user_nodes = function_graph.get_upstream_nodes(final_vars, inputs, overrides)
+        self.function_graph = _fn_graph if _fn_graph is not None else self.graph
+        self.run_id = str(uuid.uuid4())
+        nodes, user_nodes = self.function_graph.get_upstream_nodes(final_vars, inputs, overrides)
         Driver.validate_inputs(
-            function_graph, self.adapter, user_nodes, inputs, nodes
+            self.function_graph, self.adapter, user_nodes, inputs, nodes
         )  # TODO -- validate within the function graph itself
         if display_graph:  # deprecated flow.
             logger.warning(
@@ -683,46 +734,31 @@ def raw_execute(
             )
             self.visualize_execution(final_vars, "test-output/execute.gv", {"view": True})
             if self.has_cycles(
-                final_vars, function_graph
+                final_vars, self.function_graph
             ):  # here for backwards compatible driver behavior.
                 raise ValueError("Error: cycles detected in your graph.")
         all_nodes = nodes | user_nodes
         self.graph_executor.validate(list(all_nodes))
         if self.adapter.does_hook("pre_graph_execute", is_async=False):
             self.adapter.call_all_lifecycle_hooks_sync(
                 "pre_graph_execute",
-                run_id=run_id,
-                graph=function_graph,
+                run_id=self.run_id,
+                graph=self.function_graph,
                 final_vars=final_vars,
                 inputs=inputs,
                 overrides=overrides,
             )
         results = None
-        error = None
-        success = False
         try:
             results = self.graph_executor.execute(
-                function_graph,
+                self.function_graph,
                 final_vars,
                 overrides if overrides is not None else {},
                 inputs if inputs is not None else {},
-                run_id,
+                self.run_id,
             )
-            success = True
         except Exception as e:
-            error = e
-            success = False
             raise e
-        finally:
-            if self.adapter.does_hook("post_graph_execute", is_async=False):
-                self.adapter.call_all_lifecycle_hooks_sync(
-                    "post_graph_execute",
-                    run_id=run_id,
-                    graph=function_graph,
-                    success=success,
-                    error=error,
-                    results=results,
-                )
         return results
 
     @capture_function_usage
@@ -1517,6 +1553,8 @@ def materialize(
         start_time = time.time()
         run_successful = True
         error = None
+        execution_error = None
+        raw_results_output = None
 
         final_vars = self._create_final_vars(additional_vars)
         # This is so the finally logging statement does not accidentally die
@@ -1550,7 +1588,7 @@ def materialize(
             Driver.validate_inputs(function_graph, self.adapter, user_nodes, inputs, nodes)
             all_nodes = nodes | user_nodes
             self.graph_executor.validate(list(all_nodes))
-            raw_results = self.raw_execute(
+            raw_results = self.__raw_execute(
                 final_vars=final_vars + materializer_vars,
                 inputs=inputs,
                 overrides=overrides,
@@ -1563,9 +1601,19 @@ def materialize(
         except Exception as e:
             run_successful = False
             logger.error(SLACK_ERROR_MESSAGE)
+            execution_error = e
             error = telemetry.sanitize_error(*sys.exc_info())
             raise e
         finally:
+            if self.adapter.does_hook("post_graph_execute", is_async=False):
+                self.adapter.call_all_lifecycle_hooks_sync(
+                    "post_graph_execute",
+                    run_id=self.run_id,
+                    graph=self.function_graph,
+                    success=run_successful,
+                    error=execution_error,
+                    results=raw_results_output,
+                )
             duration = time.time() - start_time
             self.capture_execute_telemetry(
                 error, final_vars + materializer_vars, inputs, overrides, run_successful, duration