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 (https://github.com/DAGWorks-Inc/hamilton/blob/046ed599e00ad186c361e2ce94192b9f0688e8d1/hamilton/lifecycle/api.py#L271).

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?