WIP: decorators should wrap functions with vanilla wrapper #1282

skrawcz · 2025-02-26T20:35:35Z

This is a prototype sketch.

each decorator wraps the function -- and returns that pointer.
we only wrap a function 1 level deep.
we validate on the original function.

This should then enable:

reusing base functions easily.
not mutating the original function.
also work in SERDE contexts -- but something to properly test.

Changes

all decorators now modify a wrapper
we special case mutate (somewhat hacky)
adds tests
fixes parameterize to work with async
fixes data validators to work with async

How I tested this

unit tests
test with parallel & collect for distributed stuff (seems to work)
run main examples
Test export
Validate "originating function" is as intended
Ask hamilton library people to test this change with their stuff

Notes

needs some more tests
regression found with ray + hamilton sdk when using parallel & collect.

Checklist

PR has an informative and human-readable title (this will be pulled into the release notes)
Changes are limited to a single goal (no scope creep)
Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
Any change in functionality is tested
New functions are documented (with a description, list of inputs, and expected output)
Placeholder code is flagged / future TODOs are captured in comments
Project documentation has been updated if adding/changing functionality.

This is a prototype sketch. 1. each decorator wraps the function -- and returns that pointer. 2. we only wrap a function 1 level deep. 3. we validate on the original function. This should then enable: 1. reusing base functions easily. 2. not mutating the original function. 3. also work in SERDE contexts -- but something to properly test.

elijahbenizzy · 2025-02-26T21:15:11Z

Copying the function might be cleaner -- if this works then great but I think there are a few assumptions that you'll have to check for (async, signature stuff), which functools.wraps doesn't cover.

Might be easier to use this utility function:

hamilton/hamilton/ad_hoc_utils.py

Line 15 in 1057a40

def _copy_func(f):

skrawcz · 2025-02-26T21:47:06Z

Copying the function might be cleaner -- if this works then great but I think there are a few assumptions that you'll have to check for (async, signature stuff), which functools.wraps doesn't cover.

Might be easier to use this utility function:

hamilton/hamilton/ad_hoc_utils.py

Line 15 in 1057a40

def _copy_func(f):

yep some testing to be done for sure.

We need to audit all decorators and see if they work with async or not...

skrawcz · 2025-02-27T17:06:32Z

@jernejfrank any thoughts on what assumption I broke with @mutate?

TODO: properly fix this that doesn't leak stuff

skrawcz · 2025-02-27T20:24:35Z

@jernejfrank any thoughts on what assumption I broke with @mutate?

figured it out. I have a hack that works I think. will need to refactor a few things in this PR for a proper change

jernejfrank · 2025-02-27T20:54:08Z

@jernejfrank any thoughts on what assumption I broke with @mutate?

figured it out. I have a hack that works I think. will need to refactor a few things in this PR for a proper change

Great, just to write down why I think it failed for easier lookup:

mutate looks at the functions you want to mutate and figures out if there's already a pipe_output decorator on it. If so, it appends the mutating function (the one decorated with mutate) as the next step in the already existing pipe_output sequence of steps. Otherwise, it will create a pipe_output decorator on the target function with the mutating function as the step.

I guess with the new wrapper for the decorators, this check if pipe_output exists doesn't work, because it looked at the wrong level?

Re: refactoring -- the current implementation is also not ideal when you use jupyter notebooks, because it could potentially stack the same mutating function twice if you don't restart the kernel. We were looking with @elijahbenizzy at how to change this at the notebook level, but it is a weird edge case rabbit hole that got labeled as TODO once somebody actually runs into the problem.

With 3.0 they've changed a few things. This pins that until someone complains / wants to fix / update the code for 3.0+.

skrawcz · 2025-02-27T22:16:46Z

Okay there is one bug with parallel & collect on ray using the tracker -- the parallelized nodes aren't being reflected in the Hamilton UI -- I suspect it could be something could be messing with the "source code" pointers the SDK captures or something.
NVM this bug exists on main - running python 3.11, running ray 2.42.0. -- It works without ray just fine on main though.

elijahbenizzy

If it works then great, but wow there's a lot of complexity here. I think this is a decent approach, but curious if there's a way to simplify accounting for the wrappers?

elijahbenizzy · 2025-02-27T23:37:39Z

hamilton/function_modifiers/base.py

-        self.validate(fn)
+        # # stop unwrapping if not a hamilton function
+        # # should only be one level of "hamilton wrapping" - and that's what we attach things to.
+        self.validate(unwrap(fn, stop=lambda f: not hasattr(f, "__hamilton__")))


Why do we only do one level deep? Aren't there some weird cases with multiple imports, imports of imports, etc...?

before we attached to the function itself. What I'm trying to do here is replicate that by using one wrapper layer around a function.

E.g.

@tag(...) @extract_columns(...) @parameterize(...) def some_func(...) -> ... :

Should only have one "wrapper" around some_func - which is what is being modified by the decorators.

This then enables someone to do this:

def a(..) -> ...: # base func q = @parameterize(x=...)(a) a_validated = @check_output(range=...)(a) a_validated.__name__ = "a_validated" # need to rename to enable this to be stand alone...

This above would enable a, x, and a_validated to be in the DAG. Before it wasn't possible.

Yes, I got that. Why just one layer was the question?

Why just one layer

Doesn't seem like we need more?

This mirrors the prior behavior of directly attaching to the function - there were no layers before.

elijahbenizzy · 2025-02-27T23:38:11Z

hamilton/function_modifiers/base.py

@@ -805,6 +829,9 @@ def resolve_nodes(fn: Callable, config: Dict[str, Any]) -> Collection[node.Node]
    which configuration they need.
    :return: A list of nodes into which this function transforms.
    """
+    # check for mutate...
+    fn = handle_mutate_hack(fn)


Take away the name "hack"

elijahbenizzy · 2025-02-27T23:38:34Z

hamilton/function_modifiers/base.py

+    :param fn: Function to check
+    :return: Function or wrapped function to use if applicable.
+    """
+    if hasattr(fn, "__hamilton_wrappers__"):


Is this worth the complexity it adds? Is there a simpler way to approach it? Need to grok better.

The problem is that with mutate the current design was to attach to the function and it would be resolved properly via that way.

I think the better approach would be to register mutates via a registry that is applied to functions that are found ... and that way we wouldn't need to attach it to the function directly...

Curious, if we use a registry for mutate, would it be sensible to just have one global registry for all decorators instead of the additional 1-layer wrapper?

elijahbenizzy · 2025-02-27T23:39:09Z

hamilton/function_modifiers/expanders.py

@@ -231,6 +231,65 @@ def replacement_function(
                    # now the error will be clear enough
                return node_.callable(*args, **new_kwargs)

+            async def async_replacement_function(


note -- merge with the above when you get to polishing

skrawcz mentioned this pull request Feb 26, 2025

Allow decorator use for functions defined in other modules #1281

Open

Adds some async changes

d2052b1

We need to audit all decorators and see if they work with async or not...

skrawcz force-pushed the prototype_decorator_change branch from 2b655fb to d2052b1 Compare February 27, 2025 00:36

skrawcz added 2 commits February 27, 2025 10:54

WIP - trying to fix mutate assumptions

cd1f317

hack for fixing mutate

1d057a6

TODO: properly fix this that doesn't leak stuff

skrawcz added 3 commits February 27, 2025 13:22

Updated hack fix for mutate

448ce0b

Adds doc strings to help explain hack

17c298e

Pins ddtrace < 3.0 due to backwards incompatible change

22509fd

With 3.0 they've changed a few things. This pins that until someone complains / wants to fix / update the code for 3.0+.

elijahbenizzy reviewed Feb 27, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

WIP: decorators should wrap functions with vanilla wrapper #1282

WIP: decorators should wrap functions with vanilla wrapper #1282

skrawcz commented Feb 26, 2025 •

edited

Loading

elijahbenizzy commented Feb 26, 2025 •

edited

Loading

skrawcz commented Feb 26, 2025

skrawcz commented Feb 27, 2025

skrawcz commented Feb 27, 2025

jernejfrank commented Feb 27, 2025 •

edited

Loading

skrawcz commented Feb 27, 2025 •

edited

Loading

elijahbenizzy left a comment

elijahbenizzy Feb 27, 2025

skrawcz Feb 28, 2025 •

edited

Loading

elijahbenizzy Feb 28, 2025

skrawcz Feb 28, 2025

elijahbenizzy Feb 27, 2025

elijahbenizzy Feb 27, 2025

skrawcz Feb 28, 2025 •

edited

Loading

jernejfrank Feb 28, 2025

elijahbenizzy Feb 27, 2025

WIP: decorators should wrap functions with vanilla wrapper #1282

Are you sure you want to change the base?

WIP: decorators should wrap functions with vanilla wrapper #1282

Conversation

skrawcz commented Feb 26, 2025 • edited Loading

Changes

How I tested this

Notes

Checklist

elijahbenizzy commented Feb 26, 2025 • edited Loading

skrawcz commented Feb 26, 2025

skrawcz commented Feb 27, 2025

skrawcz commented Feb 27, 2025

jernejfrank commented Feb 27, 2025 • edited Loading

skrawcz commented Feb 27, 2025 • edited Loading

elijahbenizzy left a comment

Choose a reason for hiding this comment

elijahbenizzy Feb 27, 2025

Choose a reason for hiding this comment

skrawcz Feb 28, 2025 • edited Loading

Choose a reason for hiding this comment

elijahbenizzy Feb 28, 2025

Choose a reason for hiding this comment

skrawcz Feb 28, 2025

Choose a reason for hiding this comment

elijahbenizzy Feb 27, 2025

Choose a reason for hiding this comment

elijahbenizzy Feb 27, 2025

Choose a reason for hiding this comment

skrawcz Feb 28, 2025 • edited Loading

Choose a reason for hiding this comment

jernejfrank Feb 28, 2025

Choose a reason for hiding this comment

elijahbenizzy Feb 27, 2025

Choose a reason for hiding this comment

skrawcz commented Feb 26, 2025 •

edited

Loading

elijahbenizzy commented Feb 26, 2025 •

edited

Loading

jernejfrank commented Feb 27, 2025 •

edited

Loading

skrawcz commented Feb 27, 2025 •

edited

Loading

skrawcz Feb 28, 2025 •

edited

Loading

skrawcz Feb 28, 2025 •

edited

Loading