Update / add task execution hooks

docs/reference/lifecycle-hooks/TaskExecutionHook.rst

@@ -0,0 +1,9 @@
+===================================
+lifecycle.api.TaskExecutionHook
+===================================
+
+
+.. autoclass:: hamilton.lifecycle.api.TaskExecutionHook
+   :special-members: __init__
+   :members:
+   :inherited-members:

docs/reference/lifecycle-hooks/TaskGroupingHook.rst

@@ -0,0 +1,9 @@
+===================================
+lifecycle.api.TaskGroupingHook
+===================================
+
+
+.. autoclass:: hamilton.lifecycle.api.TaskGroupingHook
+   :special-members: __init__
+   :members:
+   :inherited-members:

docs/reference/lifecycle-hooks/index.rst

@@ -24,7 +24,8 @@ looking forward.
    NodeExecutionMethod
    StaticValidator
    GraphConstructionHook
-
+   TaskExecutionHook
+   TaskGroupingHook
 
 Available Adapters
 -------------------

hamilton/driver.py

@@ -241,6 +241,14 @@ def execute(
         results_cache = state.DictBasedResultCache(prehydrated_results)
         # Create tasks from the grouped nodes, filtering/pruning as we go
         tasks = grouping.create_task_plan(grouped_nodes, final_vars, overrides, self.adapter)
+
+        if self.adapter.does_hook("post_task_group", is_async=False):
+            self.adapter.call_all_lifecycle_hooks_sync(
+                "post_task_group",
+                run_id=run_id,
+                tasks=tasks,
+            )
+
         # Create a task graph and execution state
         execution_state = state.ExecutionState(
             tasks, results_cache, run_id

hamilton/execution/executors.py

@@ -108,6 +108,8 @@ def base_execute_task(task: TaskImplementation) -> Dict[str, Any]:
             nodes=task.nodes,
             inputs=task.dynamic_inputs,
             overrides=task.overrides,
+            spawning_task_id=task.spawning_task_id,
+            purpose=task.purpose,
         )
     error = None
     success = True
@@ -139,6 +141,8 @@ def base_execute_task(task: TaskImplementation) -> Dict[str, Any]:
                 results=results,
                 success=success,
                 error=error,
+                spawning_task_id=task.spawning_task_id,
+                purpose=task.purpose,
             )
     # This selection is for GC
     # We also need to get the override values

hamilton/execution/state.py

@@ -384,6 +384,13 @@ def update_task_state(
                 self.realize_parameterized_group(
                     completed_task.task_id, parameterization_values, input_to_parameterize
                 )
+                if completed_task.adapter.does_hook("post_task_expand", is_async=False):
+                    completed_task.adapter.call_all_lifecycle_hooks_sync(
+                        "post_task_expand",
+                        run_id=completed_task.run_id,
+                        task_id=completed_task.task_id,
+                        parameters=parameterization_values,
+                    )
             else:
                 for candidate_task in self.base_reverse_dependencies[completed_task.base_id]:
                     # This means its not spawned by another task, or a node spawning group itself

hamilton/lifecycle/api.py

@@ -15,7 +15,11 @@
 # To really fix this we should move everything user-facing out of base, which is a pretty sloppy name for a package anyway
 # And put it where it belongs. For now we're OK with the TYPE_CHECKING hack
 if TYPE_CHECKING:
+    from hamilton.execution.grouping import NodeGroupPurpose, TaskSpec
     from hamilton.graph import FunctionGraph
+else:
+    NodeGroupPurpose = None
+    TaskSpec = None
 
 from hamilton.graph_types import HamiltonGraph, HamiltonNode
 from hamilton.lifecycle.base import (
@@ -27,6 +31,8 @@
     BasePostGraphExecute,
     BasePostNodeExecute,
     BasePostTaskExecute,
+    BasePostTaskExpand,
+    BasePostTaskGroup,
     BasePreGraphExecute,
     BasePreNodeExecute,
     BasePreTaskExecute,
@@ -379,13 +385,17 @@ def pre_task_execute(
         nodes: List["node.Node"],
         inputs: Dict[str, Any],
         overrides: Dict[str, Any],
+        spawning_task_id: Optional[str],
+        purpose: NodeGroupPurpose,
     ):
         self.run_before_task_execution(
             run_id=run_id,
             task_id=task_id,
             nodes=[HamiltonNode.from_node(n) for n in nodes],
             inputs=inputs,
             overrides=overrides,
+            spawning_task_id=spawning_task_id,
+            purpose=purpose,
         )
 
     def post_task_execute(
@@ -397,6 +407,8 @@ def post_task_execute(
         results: Optional[Dict[str, Any]],
         success: bool,
         error: Exception,
+        spawning_task_id: Optional[str],
+        purpose: NodeGroupPurpose,
     ):
         self.run_after_task_execution(
             run_id=run_id,
@@ -405,6 +417,8 @@ def post_task_execute(
             results=results,
             success=success,
             error=error,
+            spawning_task_id=spawning_task_id,
+            purpose=purpose,
         )
 
     @abc.abstractmethod
@@ -416,6 +430,8 @@ def run_before_task_execution(
         nodes: List[HamiltonNode],
         inputs: Dict[str, Any],
         overrides: Dict[str, Any],
+        spawning_task_id: Optional[str],
+        purpose: NodeGroupPurpose,
         **future_kwargs,
     ):
         """Implement this to run something after task execution. Tasks are tols used to group nodes.
@@ -428,6 +444,8 @@ def run_before_task_execution(
         :param inputs: Inputs to the task
         :param overrides: Overrides passed to the task
         :param future_kwargs: Reserved for backwards compatibility.
+        :param spawning_task_id: ID of the task that spawned this task
+        :param purpose: Purpose of the current task group
         """
         pass
 
@@ -441,6 +459,8 @@ def run_after_task_execution(
         results: Optional[Dict[str, Any]],
         success: bool,
         error: Exception,
+        spawning_task_id: Optional[str],
+        purpose: NodeGroupPurpose,
         **future_kwargs,
     ):
         """Implement this to run something after task execution. See note in run_before_task_execution.
@@ -452,6 +472,8 @@ def run_after_task_execution(
         :param success: Whether the task was successful
         :param error: The error the task threw, if any
         :param future_kwargs: Reserved for backwards compatibility.
+        :param spawning_task_id: ID of the task that spawned this task
+        :param purpose: Purpose of the current task group
         """
         pass
 
@@ -614,6 +636,42 @@ def validate_graph(
         return self.run_to_validate_graph(graph=HamiltonGraph.from_graph(graph))
 
 
+class TaskGroupingHook(BasePostTaskGroup, BasePostTaskExpand):
+    """Implement this to run something after task grouping or task expansion. This will allow you to
+    capture information about the tasks during `Parallelize`/`Collect` blocks in dynamic DAG execution."""
+
+    @override
+    @final
+    def post_task_group(self, *, run_id: str, tasks: List[TaskSpec]):
+        return self.run_after_task_grouping(run_id=run_id, tasks=tasks)
+
+    @override
+    @final
+    def post_task_expand(self, *, run_id: str, task_id: str, parameters: Dict[str, Any]):
+        return self.run_after_task_expansion(run_id=run_id, task_id=task_id, parameters=parameters)
+
+    @abc.abstractmethod
+    def run_after_task_grouping(self, *, run_id: str, tasks: List[TaskSpec], **future_kwargs):
+        """Hook that is called after task grouping.
+        :param run_id: ID of the run, unique in scope of the driver.
+        :param tasks: List of tasks that were grouped together.
+        :param future_kwargs: Additional keyword arguments -- this is kept for backwards compatibility.
+        """
+        pass
+
+    @abc.abstractmethod
+    def run_after_task_expansion(
+        self, *, run_id: str, task_id: str, parameters: Dict[str, Any], **future_kwargs
+    ):
+        """Hook that is called after task expansion.
+        :param run_id: ID of the run, unique in scope of the driver.
+        :param task_id: ID of the task that was expanded.
+        :param parameters: Parameters that were passed to the task.
+        :param future_kwargs: Additional keyword arguments -- this is kept for backwards compatibility.
+        """
+        pass
+
+
 class GraphConstructionHook(BasePostGraphConstruct, abc.ABC):
     """Hook that is run after graph construction. This allows you to register/capture info on the graph.
     Note that, in the case of materialization, this may be called multiple times (once when we create the graph,

hamilton/lifecycle/base.py

@@ -43,6 +43,10 @@
 # python, which (our usage of) leans type-hinting trigger-happy, this will suffice.
 if TYPE_CHECKING:
     from hamilton import graph, node
+    from hamilton.execution.grouping import NodeGroupPurpose, TaskSpec
+else:
+    NodeGroupPurpose = None
+    TaskSpec = None
 
 # All of these are internal APIs. Specifically, structure required to manage a set of
 # hooks/methods/validators that we will likely expand. We store them in constants (rather than, say, a more complex single object)
@@ -418,6 +422,8 @@ def pre_task_execute(
         nodes: List["node.Node"],
         inputs: Dict[str, Any],
         overrides: Dict[str, Any],
+        spawning_task_id: Optional[str],
+        purpose: NodeGroupPurpose,
     ):
         """Hook that is called immediately prior to task execution. Note that this is only useful in dynamic
         execution, although we reserve the right to add this back into the standard hamilton execution pattern.
@@ -427,6 +433,8 @@ def pre_task_execute(
         :param nodes: Nodes that are being executed
         :param inputs: Inputs to the task
         :param overrides: Overrides to task execution
+        :param spawning_task_id: ID of the task that spawned this task
+        :param purpose: Purpose of the current task group
         """
         pass
 
@@ -442,6 +450,8 @@ async def pre_task_execute(
         nodes: List["node.Node"],
         inputs: Dict[str, Any],
         overrides: Dict[str, Any],
+        spawning_task_id: Optional[str],
+        purpose: NodeGroupPurpose,
     ):
         """Hook that is called immediately prior to task execution. Note that this is only useful in dynamic
         execution, although we reserve the right to add this back into the standard hamilton execution pattern.
@@ -451,6 +461,8 @@ async def pre_task_execute(
         :param nodes: Nodes that are being executed
         :param inputs: Inputs to the task
         :param overrides: Overrides to task execution
+        :param spawning_task_id: ID of the task that spawned this task
+        :param purpose: Purpose of the current task group
         """
         pass
 
@@ -615,6 +627,31 @@ async def post_node_execute(
         pass
 
 
+@lifecycle.base_hook("post_task_group")
+class BasePostTaskGroup(abc.ABC):
+    @abc.abstractmethod
+    def post_task_group(self, *, run_id: str, tasks: List[TaskSpec]):
+        """Hook that is called immediately after a task group is created. Note that this is only useful in dynamic
+        execution, although we reserve the right to add this back into the standard hamilton execution pattern.
+
+        :param run_id: ID of the run, unique in scope of the driver.
+        :param tasks: Tasks specs that are in the group."""
+        pass
+
+
+@lifecycle.base_hook("post_task_expand")
+class BasePostTaskExpand(abc.ABC):
+    @abc.abstractmethod
+    def post_task_expand(self, *, run_id: str, task_id: str, parameters: Dict[str, Any]):
+        """Hook that is called immediately after a task is expanded into separate task. Note that this is only useful
+        in dynamic execution, although we reserve the right to add this back into the standard hamilton execution pattern.
+
+        :param run_id: ID of the run, unique in scope of the driver.
+        :param task_id: ID of the task.
+        :param parameters: Parameters that are being passed to each of the expanded tasks."""
+        pass
+
+
 @lifecycle.base_hook("post_task_execute")
 class BasePostTaskExecute(abc.ABC):
     @abc.abstractmethod
@@ -627,6 +664,8 @@ def post_task_execute(
         results: Optional[Dict[str, Any]],
         success: bool,
         error: Exception,
+        spawning_task_id: Optional[str],
+        purpose: NodeGroupPurpose,
     ):
         """Hook called immediately after task execution. Note that this is only useful in dynamic
         execution, although we reserve the right to add this back into the standard hamilton execution pattern.
@@ -637,6 +676,8 @@ def post_task_execute(
         :param results: Results of the task
         :param success: Whether or not the task executed successfully
         :param error: The error that was raised, if any
+        :param spawning_task_id: ID of the task that spawned this task
+        :param purpose: Purpose of the current task group
         """
         pass
 
@@ -653,6 +694,8 @@ async def post_task_execute(
         results: Optional[Dict[str, Any]],
         success: bool,
         error: Exception,
+        spawning_task_id: Optional[str],
+        purpose: NodeGroupPurpose,
     ):
         """Asynchronous Hook called immediately after task execution. Note that this is only useful in dynamic
         execution, although we reserve the right to add this back into the standard hamilton execution pattern.
@@ -663,6 +706,8 @@ async def post_task_execute(
         :param results: Results of the task
         :param success: Whether or not the task executed successfully
         :param error: The error that was raised, if any
+        :param spawning_task_id: ID of the task that spawned this task
+        :param purpose: Purpose of the current task group
         """
         pass
 
@@ -737,6 +782,8 @@ def do_build_result(self, *, outputs: Any) -> Any:
     BasePostGraphConstructAsync,
     BasePreGraphExecute,
     BasePreGraphExecuteAsync,
+    BasePostTaskGroup,
+    BasePostTaskExpand,
     BasePreTaskExecute,
     BasePreTaskExecuteAsync,
     BasePreNodeExecute,

tests/lifecycle/lifecycle_adapters_for_testing.py

@@ -5,6 +5,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 from hamilton import node
+from hamilton.execution.grouping import NodeGroupPurpose, TaskSpec
 from hamilton.graph import FunctionGraph
 from hamilton.lifecycle.base import (
     BaseDoBuildResult,
@@ -15,6 +16,8 @@
     BasePostGraphExecute,
     BasePostNodeExecute,
     BasePostTaskExecute,
+    BasePostTaskExpand,
+    BasePostTaskGroup,
     BasePreDoAnythingHook,
     BasePreGraphExecute,
     BasePreNodeExecute,
@@ -108,6 +111,11 @@ def pre_graph_execute(
         pass
 
 
+class TrackingPostTaskGroupHook(ExtendToTrackCalls, BasePostTaskGroup):
+    def post_task_group(self, run_id: str, tasks: List[TaskSpec]):
+        pass
+
+
 class TrackingPreTaskExecuteHook(ExtendToTrackCalls, BasePreTaskExecute):
     def pre_task_execute(
         self,
@@ -116,6 +124,8 @@ def pre_task_execute(
         nodes: List[node.Node],
         inputs: Dict[str, Any],
         overrides: Dict[str, Any],
+        spawning_task_id: Optional[str],
+        purpose: NodeGroupPurpose,
     ):
         pass
 
@@ -150,10 +160,17 @@ def post_task_execute(
         results: Optional[Dict[str, Any]],
         success: bool,
         error: Exception,
+        spawning_task_id: Optional[str],
+        purpose: NodeGroupPurpose,
     ):
         pass
 
 
+class TrackingPostTaskExpandHook(ExtendToTrackCalls, BasePostTaskExpand):
+    def post_task_expand(self, run_id: str, task_id: str, parameters: Dict[str, Any]):
+        pass
+
+
 class TrackingPostGraphExecuteHook(ExtendToTrackCalls, BasePostGraphExecute):
     def post_graph_execute(
         self,