Dynamic DAG build fails #1409
Description
Current behavior
My project uses Hamilton to manage a DAG-based processing pipeline that mainly manipulates xarray datasets. The DAG terminal nodes change based on user configuration. To implement this, I use the resolve decorator, combined with an extract_field wrapper that implements field selection based on configuration. The following reproducer illustrates this:
definitions.py
import hamilton
from hamilton.function_modifiers import ResolveAt, extract_fields, resolve
hamilton.enable_power_user_mode = True
def _extract_fields(field_name: str):
return extract_fields({field_name: int})
@resolve(when=ResolveAt.CONFIG_AVAILABLE, decorate_with=_extract_fields)
def a() -> dict:
return {"b": 1, "c": 2}main.py (I actually run this in a notebook)
import definitions
from hamilton import driver
drv = (
driver.Builder()
.with_config({"hamilton.enable_power_user_mode": True, "field_name": "b"})
.with_modules(definitions)
.build()
)
drv.display_all_functions()With Hamilton v1.88 and old, that runs just fine. With Hamilton v1.89.0, it raises an AttributeError with the following stack trace:
Stack trace
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[1], line 6
1 import definitions
3 from hamilton import driver
5 drv = (
----> 6 driver.Builder()
7 .with_config({"hamilton.enable_power_user_mode": True, "field_name": "b"})
8 .with_modules(definitions)
9 .build()
10 )
12 drv.display_all_functions()
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/driver.py:2189, in Builder.build(self)
2182 grouping_strategy = self.grouping_strategy or grouping.GroupByRepeatableBlocks()
2183 graph_executor = TaskBasedGraphExecutor(
2184 execution_manager=execution_manager,
2185 grouping_strategy=grouping_strategy,
2186 adapter=lifecycle_base.LifecycleAdapterSet(*adapter),
2187 )
-> 2189 return Driver(
2190 self.config,
2191 *self.modules,
2192 adapter=adapter,
2193 _materializers=self.materializers,
2194 _graph_executor=graph_executor,
2195 _use_legacy_adapter=False,
2196 allow_module_overrides=self._allow_module_overrides,
2197 )
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/driver.py:492, in Driver.__init__(self, config, adapter, allow_module_overrides, _materializers, _graph_executor, _use_legacy_adapter, *modules)
490 error = telemetry.sanitize_error(*sys.exc_info())
491 logger.error(SLACK_ERROR_MESSAGE)
--> 492 raise e
493 finally:
494 # TODO -- update this to use the lifecycle methods
495 self.capture_constructor_telemetry(error, modules, config, adapter)
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/driver.py:466, in Driver.__init__(self, config, adapter, allow_module_overrides, _materializers, _graph_executor, _use_legacy_adapter, *modules)
464 self.graph_modules = modules
465 try:
--> 466 self.graph = graph.FunctionGraph.from_modules(
467 *modules,
468 config=config,
469 adapter=adapter,
470 allow_module_overrides=allow_module_overrides,
471 )
472 if _materializers:
473 materializer_factories, extractor_factories = self._process_materializers(
474 _materializers
475 )
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/graph.py:753, in FunctionGraph.from_modules(config, adapter, allow_module_overrides, *modules)
734 @staticmethod
735 def from_modules(
736 *modules: ModuleType,
(...)
739 allow_module_overrides: bool = False,
740 ):
741 """Initializes a function graph from the specified modules. Note that this was the old
742 way we constructed FunctionGraph -- this is not a public-facing API, so we replaced it
743 with a constructor that takes in nodes directly. If you hacked in something using
(...)
750 :return: a function graph.
751 """
--> 753 nodes = create_function_graph(
754 *modules, config=config, adapter=adapter, allow_module_overrides=allow_module_overrides
755 )
756 return FunctionGraph(nodes, config, adapter)
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/graph.py:188, in create_function_graph(config, adapter, fg, allow_module_overrides, *modules)
186 # create non-input nodes -- easier to just create this in one loop
187 for _func_name, f in functions:
--> 188 for n in fm_base.resolve_nodes(f, config):
189 if n.name in config:
190 continue # This makes sure we overwrite things if they're in the config...
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/function_modifiers/base.py:850, in resolve_nodes(fn, config)
848 except Exception as e:
849 logger.exception(_resolve_nodes_error(fn))
--> 850 raise e
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/function_modifiers/base.py:843, in resolve_nodes(fn, config)
841 node_transformers = function_decorators[NodeTransformer.get_lifecycle_name()]
842 for dag_modifier in node_transformers:
--> 843 nodes = dag_modifier.transform_dag(nodes, filter_config(config, dag_modifier), fn)
844 function_decorators = function_decorators[NodeDecorator.get_lifecycle_name()]
845 for node_decorator in function_decorators:
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/function_modifiers/base.py:532, in NodeTransformer.transform_dag(self, nodes, config, fn)
530 nodes_to_keep = self.compliment(nodes, nodes_to_transform)
531 out = list(nodes_to_keep)
--> 532 out += self.transform_targets(nodes_to_transform, config, fn)
533 return out
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/function_modifiers/base.py:584, in SingleNodeNodeTransformer.transform_targets(self, targets, config, fn)
579 if len(targets) != 1:
580 raise InvalidDecoratorException(
581 f"Expected a single node to transform, but got {len(targets)}. {self.__class__} "
582 f" can only operate on a single node, but multiple nodes were created by {fn.__qualname__}"
583 )
--> 584 return super().transform_targets(targets, config, fn)
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/function_modifiers/base.py:515, in NodeTransformer.transform_targets(self, targets, config, fn)
513 out = []
514 for node_to_transform in targets:
--> 515 out += list(self.transform_node(node_to_transform, config, fn))
516 return out
File ~/Documents/src/rayference/rtm/eradiate/.pixi/envs/dev/lib/python3.9/site-packages/hamilton/function_modifiers/expanders.py:918, in extract_fields.transform_node(self, node_, config, fn)
914 return dict_generated
916 output_nodes = [node_.copy_with(callabl=dict_generator)]
--> 918 for field, field_type in self.resolved_fields.items():
919 doc_string = base_doc # default doc string of base function.
921 # This extractor is constructed to avoid closure issues.
AttributeError: 'extract_fields' object has no attribute 'resolved_fields'
Is this expected?
Workaround
Based on a Claude Code suggestion
Is the following, updated code fine?
definitions.py
import hamilton
from hamilton.function_modifiers import ResolveAt, extract_fields, resolve
hamilton.enable_power_user_mode = True
def _extract_fields(field_name: str):
# Create the decorator
fields = {field_name: int}
decorator = extract_fields(fields)
# WORKAROUND for Hamilton v1.89:
# The extract_fields decorator needs its resolved_fields and output_type
# attributes set before it can be used with @resolve(decorate_with=...).
# Normally these are set in validate(), but that hasn't been called yet.
# We set them manually here by mimicking what validate() does.
if hamilton.__version__ >= (1, 89, 0):
from hamilton.function_modifiers.expanders import _determine_fields_to_extract
decorator.output_type = dict
decorator.resolved_fields = _determine_fields_to_extract(fields, dict)
return decorator
@resolve(when=ResolveAt.CONFIG_AVAILABLE, decorate_with=_extract_fields)
def a() -> dict:
return {"b": 1, "c": 2}Library & System Information
- Python: 3.9
- Hamilton: 1.89.0
- System: macOS 14.7 and Ubuntu 24.04
Expected behavior
With Hamilton 1.88, this runs without problem and I get the following graph:
