Parallelization Enhancement Ideas

[mitstake]

Hi, as I've been working with the Parallelizable/Collect capabilities of Hamilton I've come across a couple limitations that I am having trouble finding adequate workarounds for. I have some ideas that could address these that I would like to discuss with the Hamilton team.

The limitations:

1. Nested Parallelizable - something like this does not work:

def product_line() -> Parallelizable[str]:
    return ['Cars', 'Trucks']

def product(product_line: str) -> Parallelizable[str]:
    return {'Cars': ['Camry', 'Corolla'], 'Trucks': ['Tacoma', 'Tundra']}[product_line]

def product_data(product_line: str, product: str) -> int:
    return len(product_line + product)

def product_line_data(product_data: Collect[int]) -> int:
    return sum(product_data)

def all_data(product_line_data: Collect[int]) -> int:
    return sum(product_line_data)

The workaround that I considered was flattening to just one Parallelizable that returns a tuple of (product_lineX, productY), but this is no good because then combining/processing the data will need to be done for both levels in one function (all_data). This means you lose the parallelization benefits for the product_line level (in this simple case it's negligible, but not in all real world cases), and also the logic in all_data becomes more complicated as it must now be responsible for parsing/processing multiple levels.

2. Isolate Individual Parallelizable Path - previously discussed here as "not planned". However, I would request reopening discussion on this given the context of the suggestions in this thread.
   Parallel Execution: KeyError: 'key df not found in cache' #1029

The workaround I considered here was something like this:

def product() -> Parallelizable[str]:
    return ['Hat', 'Shirt', 'Shoes']

def product_data(product: str) -> Tuple[str, int]:
    return (product, len(product))

def all_data_raw(product_data: Collect[Tuple[str, int]]) -> Dict[str, int]:
    return {product: pd for product, pd in product_data}

def all_data(all_data_raw: Dict[str, int]) -> int:
    return sum(all_data_raw.values())

def main(all_data: int, all_data_raw: Dict[str, int]) -> Tuple[int, int]:
    return (all_data, all_data_raw['Hat'])

This pattern will give you the desired combined result (all_data), and also a result indexed by individual path result (all_data_raw) that you can slice from. However, this breaks down in a few ways. (1) You need to pass around a tuple containing the info that the path is parameterized on, through each function. (2) This will not allow you to access intermediate results in the path, only the terminal node result, which is not necessarily the specific path result you want. (3) If you only ever want a specific path result and not the others, you'll still need to compute all the others, resulting in wasted compute (or alternatively you could reconfigure your driver to avoid this, but that also adds complexity).

Proposal:

The idea is still rough and not fully fleshed out, but what I would propose is a new additional syntax for defining parallizable paths. I think it's best if I start with a simple example to illustrate the idea:

#################################################

# Existing Hamilton API
 
def product() -> Parallelizable[str]:
    return ['Hat', 'Coat', ...]
 
def product_cost_raw(product: str) -> int:
    return len(product)
 
def product_cost(product_cost_raw: int) -> int:
    return product_cost_raw + 1
 
def total_cost_raw(product_cost: Collect[int]) -> int:
    return sum(product_cost)
 
def total_cost(total_cost_raw: int) -> int:
    return total_cost_raw + 1
 
driver = driver.Builder()…
driver.execute(['total_cost'])

#################################################

# Proposed Additional API - add driver and parallelized parameter as argument for parallelized path functions
 
def product_cost_raw(driver, product: str) -> int:
    return len(product)
 
def product_cost(driver, product: str) -> int:
    return driver.call('product_cost_raw', product=product) + 1
 
def total_cost_raw(driver) -> int:
    # parallelization happens here as a regular for loop
    product_costs = [driver.call('product_cost', product=product) for product in ['Hat', 'Coat', ...]]
    product_costs = Collect[int](product_costs)
    return sum(product_costs)
 
def total_cost(total_cost_raw: int) -> int:  # remains the same
    return total_cost_raw + 1
 
driver = driver.Builder()…
driver.execute('total_cost')
driver.execute('product_cost', product='Hat')  # can now run individual parallelized path in isolation

#################################################

The idea here is that when you want to do some branching you can provide to your hamilton function a reference to the hamilton driver itself. This would be used instead of the function arguments to encode the function dependencies by "calling" them through the driver. And using Collect as a delimiting point to say stop defining dependencies and collect their resolved values now. These "calls" would be used during the graph build, so that instead of inspecting the function arguments you run the function and follow its driver calls to the dependencies, which can then potentially either be another "driver" function or standard hamilton function. And what this means for parallelization is that it's simply represented with a regular python for loop!

This unlocks a few benefits:
(1) similar to plain python - this follows extremely closely to how a naive plain python implementation would be coded. Just replace driver.call with the direct python call and remove Collect calls and you have your standard python code. Plain for loops look the same, args passed the same way, etc, so it's easy for users to understand what going on.

(2) nested parallelization possible - you could simply do something like this:

def C(driver, x):
    return ...

def B(driver, x):
    cs = [driver.call('C', i) for i in [1, 2, 3]]
    cs = Collect[int](cs)
    return sum(cs) + x

def A(driver):
    bs = [driver.call('B', i) for i in [1, 2, 3]]
    bs = Collect[int](bs)
    return sum(bs)

(3) can call an individual path - e.g. can call this driver.execute('B', x=1) (or driver.call within another hamilton function) without needing to build all the other possible paths and without needing to alter/reconfigure the driver.

(4) can get a static graph - with Parallelizable the parallelized branches are necessarily only known at run time. With this approach you can get the paths in a static graph during the graph build, which opens up further optimizations/inspections that can be done.

Would be happy to hear the teams thoughts and discuss further. Thanks!

[skrawcz]

Interesting!

Will need to think about this some more.

But some intial questions:

1. You're only doing single node parallelization with your proposal, right? You're not parallelizing a whole sub-DAG, IIUC?
2. driver would need a type, and would need to be passed in as input to execute to conform to how Hamilton normally behaves, right?

[skrawcz]

Yep I see now, this is just to build the graph. I think we can think more about the API, like what is the object passed in, what does it have reference to ?

Otherwise I want to improve the flexibility in creating graphs in general, e.g. configuration driven -> classes + functions -> nodes... and this fits into that effort.

To help me think through your proposal, I want to have a clear mental model of "when" things would occur to help design the API.

E.g. do I start with a basic DAG that I then want to parameterize more? Should the build process be split into two: define atomic functions, then define parameterizations (I'm thinking build basic driver, then pass that driver into another driver that has functions that are parameterizations, or should functions that take a driver as argument are built in a second pass...)?
E.g. how do we catch errors? When we would know it's an error?

Thoughts? What would match your ideal development process mental model?

[mitstake]

To help me think through your proposal, I want to have a clear mental model of "when" things would occur to help design the API.

You can take a look at how it works in my darl library. That’s a fully “production” ready implementation and shouldn’t be too tricky to decipher.
https://github.com/mitstake/darl/blob/main/src/darl/graph_build/sequential.py

The basic idea is that the standard asset style functions and the driver/parameter aware (I’ll call them darl functions) can all be resolved in the same single pass through the graph build. I don’t have any insight into how Hamilton builds the graph so it might not be as straightforward without some rearchitecting there. In fact the darl functions are just a generalization of the asset functions. Darl also allows asset functions and just wraps them to convert them to the darl style under the hood: https://github.com/mitstake/darl/blob/112ff936045e23f7921129aa97348e893d66b85a/src/darl/special_providers.py#L57
But you don’t even need to do that honestly, in your graph build you can just treat them separate.

A pseudo code example of the graph build looks like this (sorry on mobile so this might not be great)

class FakeDriver:
    def __init__(self):
        self. called_dependencies = []
    def __getattr__(self, item):
        def caller(*args):
            self. called_dependencies.append((item, args))
            return Proxy()
        return caller

graph = Graph()
def build_graph(func_name, args):
    node = (func_name, args)
    fake_driver = FakeDriver()
    func = real_driver.get_func(func_name)
    if is_darl_func(func):
        try:
            func(fake_driver, *args)
        except CollectExitException:      # raised by ngn.collect()
            pass
        deps = fake_driver.called_dependencies
    else:
        deps = parse_names_from_signature(func)
    for dep_func_name, dep_args in deps:
        dep_node = (dep_func_name, dep_args)
        graph.add_edge(dep_node, node)
        build_graph(dep_func_name, dep_args)    # recursive graph build

Hopefully that mostly gets the idea across. Regarding what I mentioned about 1 graph per driver, a darl function which is parameterized can technically result in infinitely many different graphs if the args aren’t constrained (remember you can also have a non parameterized darl style function, see snippet below). So one potential thing you could do to limit it to one graph still is require that a root execution on the driver be a regular asset style function or non-parameterized darl function. This way the set of parameters to get passed to a parameterized darl function will be known at graph build time, instead of driver invocation time.

def parameterized_darl_func(ngn, x):
    …

def non_parameterized_darl_func(ngn):
    …

ngn.parameterized_darl_func(99)
ngn.non_parameterized_darl_func()

classes + functions -> nodes

Regarding your point about defining nodes with classes, darl allows this too. Technically it’s just any callable works, so you can do something like this:

# asset style callable
class MyResult1:
    def __init__(self, val):
        self.val = val

    def __call__(self, MyDep):
        return MyDep + self.val

# darl style callable
class MyResult2:
    def __init__(self, val):
        self.val = val

    def __call__(self, ngn):
        md = ngn.MyDep()
        ngn.collect()
        return md + self.val

how do we catch errors?

check out the way it works in darl here: https://github.com/mitstake/darl/tree/main?tab=readme-ov-file#error-handling

What would match your ideal development process mental model?

Honestly the way things work in darl exactly matches my mental model. It’s a model I’ve worked with for several years with great success. I would look through that/the readme and see how you can borrow ideas from that for Hamilton. I think some of those ideas would make Hamilton a lot better/more powerful, especially the ones that I linked from the other discussions (tracing, updating/scoped updating, more targeted error handling). All of those have sections about them in the readme.

Let me know if there’s anything you want to discuss/dive into on this. Id be happy to discuss further.

[skrawcz]

@mitstake sorry I meant at a higher level. What would we tell a new user to do, to utilize this feature? What is the way the code is being developed and used? I would want the API to naturally guide / fit into that pattern.

Is your pattern what most people would do? or are there alternative "patterns" to consider?

In general the goal of the API UX is to make certain things easier (or harder) to accomplish and trading off what is required, how configurable is it, etc.

[mitstake]

ok so knowing a little bit more about the philosophy of the UX this might not be a great fit for hamilton. I'll get to why in a second. but to answer the question: "is your pattern what most people would do?" I think the answer is yes. the darl style allows you to write code how you would if you were writing naive imperative python with modular nested functions and still unlock the power and benefits of compiling to a dag, which I think is a compelling proposition on the surface.
E.g. writing code like this can become dagified with some restrictions.

def A():
    b = B()
    return b + 1

The reason I think this might not be a good fit for the Hamilton philosophy is what else this unleashes on the environment (note I don't think all these things are bad in general, I like some of them personally, but I do think it's at odds with the Hamilton way of doing things, and including it could just lead to confusion/frustration among users). It might create a lower barrier to entry by looking like standard naive python code, but it also includes several things that could be considered too flexible. Some of those things are:

1. configuring routing logic in functions rather than externally, e.g. doing this instead of configuring with @config decorators

def PizzaSpecial(ngn):  # ngn stands for engine and is the darl equivalent of driver
    if is_vegetarian:
        pizza = ngn.VeggiePizza()
    else:
        pizza = ngn.MeatLoversPizza()
    ngn.collect()
    return pizza

2. parameterizing on unreasonable data - probably the biggest frustration I've experienced with parameterizing nodes is people passing in large objects into the arguments. this has caused all sorts of issues with the memory footprint of the graph blowing up because function args are stored on the graph object. and if you're passing around 1gb dataframes through all your nodes that balloons very quickly. Also function arguments necessarily need to be hashable which can be annoying. Also if users modify function argument objects in place that leads to issues so you need to either do deepcopys of the arguments or depend on user education.
3. chaining results from function calls - if you allow parameterized functions one of the first things people are going to want to do is pass the result of one function call to another, e.g:

def A(ngn):
    b = ngn.B()
    c = ngn.C(b)  # <--- result of B passed into C
    return b + c

this is possible to do and implemented in darl, however it significantly complicates the graph building logic (and often the workflow business logic too). And to your point of the UX should make it harder to accomplish certain things, I'd say this is a good example. Instead of encouraging users to do the above, you'd prefer them to do something like this:

def A(B, C):
    return B + C

def C(B):
    return B + 1

4. Your parameterization needs to propagate through every function in the chain - one nice thing about Hamilton Parallelizable is that the functions in parallelized branch can live outside the context of the variable they're being parallelized on. With the darl style for loop everything in the parallelized branch needs to be parameterized to take that variable as well.

There's probably some other things I'm forgetting too.

Update functionality

So in the same vein regarding my suggestion on the update function (#1397, https://github.com/mitstake/darl/tree/main?tab=readme-ov-file#updates-and-shocks), I personally think this is a good feature for enabling flexible scenario configuration. But the more I'm picking up on the hamilton philosophy, I think it is potentially at odds there too. This feature would factor configuration of your driver out of explicitly defined and visible @config decorated workflows and open up the ability to configure your driver "anywhere" and "anyhow" which means you have a lot less visibility/lineage into what configuration constitutes your driver. Again I think it has upsides, but potentially heavier downsides from the hamilton pov?

Tracing

One thing I think would definitely be useful for Hamilton and would not conflict with any UX philosophy is the darl trace. https://github.com/mitstake/darl/tree/main?tab=readme-ov-file#debugging-tracing-and-replaying. This makes debugging and navigating through a workflow execution very easy and pleasant.

Anyway, I may be overthinking all of this and both the darl style functions and the update functionality make sense to you within hamilton. Interested to hear your thoughts on this, and happy to continue the conversation.

[mitstake]

One more point against updates functionality in hamilton potentially. The updates almost always without fail end up getting used to configure the standard base case rather than used for one off scenario experiements. These are updates that are not tracked through explicitly defined workflow configurations. Scoped updates end up getting used in this way too, and those can be very daunting to deal with in base case scenarios where they often get over-employed. Potentially falls under your umbrella of the ux should make certain things harder to accomplish.