This will hopefully bring a little speed improvement to batch upscales

"Hopefully"? Did... did you not measure whether your optimization actually made things faster?

And on that note: if you do an optimization, please also report how much faster this actually makes things. As a reviewer, I also have to weigh whether the speed-up is worth the tradeoff in code simplicity/quality/maintainability. I can of course measure things myself, but having you give a reference case/number is useful.

Not sure if this is the best api for it, but packages can now define on_chain_finish callbacks to be run when the chain is finished, to do things like resource clearing.

I would suggest the following API: we have a def add_cleanup(callback: Callable[[], None]) -> None method to NodeContext. This will allow nodes to add cleanup functions that run when the chain is done. Cleanup callbacks are stored in a set, so nodes can decide to run a cleanup action once or once per node execution by doing:

def do_clean_up(): ...

@register(...)
def my_node(context: NodeContext):
    # run once per node execution
    context.add_cleanup(lambda: do_clean_up())
    # run once per chain
    context.add_cleanup(do_clean_up)

This fixes the following issue with the current set_on_chain_finish API:

• set_on_chain_finish always runs cleanup actions even if no nodes from that package were executed.
• Cleanup actions were not allowed to assume the dependencies of the package.

"Hopefully"? Did... did you not measure whether your optimization actually made things faster?

I did, after I made this PR, but that was with all of my changes together. I'll test this individually.

Turns out this is doing most of the heavy lifting for the optimizations I tested. a 2m20s processing time went down to 2m01s with just this change alone.

Probably needs an from __future__ import annotations somewhere.