TDQM hook with overrides

[zilto]

Current behavior

When using hamilton.plugins.h_tdqm.ProgressBar(), the progress bar uses "the number of nodes to return" as denominator and "the number of nodes executed" as numerator. However, passing values to .execute(overrides=...) counts towrds the first, but not the latter. Code runs to completion but the progress bar completes without reaching 100%

Screenshots

Dataflow

# joke.py

def joke_prompt(topic: str) -> str:
    return f"Knock, knock. Who's there? {topic}"

def reply(joke_prompt: str) -> str:
    _, _, right = joke_prompt.partition("? ")
    return f"{right} who?"

def punchline(reply: str) -> str:
    left, _, _ = reply.partition(" ")
    return f"No, {left} MooOOooo"

Execution

import joke
from hamilton import driver
from hamilton.plugins.h_tqdm import ProgressBar
dr = (
    driver.Builder()
    .with_adapters(ProgressBar())
    .with_modules(joke)
    .build()
)

dr.execute(["joke_prompt", "reply"], overrides=dict(punchline="..."), inputs=topic("hello")).

[elijahbenizzy]

Hmm this might be a bit tricky -- we'll probably want to add a get_execution_plan() for the HamiltonGraph object. Otherwise we won't have access to the number. Otherwise we'll also want a overrides in the GraphExecutionHook (

hamilton/hamilton/lifecycle/api.py

Line 271 in 046ed59

class GraphExecutionHook(BasePreGraphExecute, BasePostGraphExecute):

).

Everything will be backwards compatible, but not super simple.

[zilto]

This seems low priority, but good to have it documented! Thanks for the details

[skrawcz]

We could just make sure there's a post graph hook implementation that makes sure the bar hits 100% if there was no error?

[rajatkriplani]

Hi @skrawcz following up from issue #1186
and @zilto I have come up with the following:

• adjusting pre_graph_execute by adding the overridden node to execution path to reflect progress:

nodes_to_execute = set(all_nodes) - set(user_defined_nodes) | set(overrides.keys())

OR

nodes_to_execute = set(all_nodes) 

• in post_graph_execute force progress bar to 100% if the execution is successful:

if success and hasattr(self, 'progress_bar'):
       self.progress_bar.update_to_completion()

[skrawcz]

@rajatkriplani Cool feel free to create a PR with it. If there isn't a test we'll need to add one.