Garbage collection/memory optimization

[elijahbenizzy]

This deletes unused references in parallelism/dynamism. That said, the python garbage collector occasionally has to be told what to do, so the user should be aware.

Fixes #373

[Short description explaining the high-level reason for the pull request]

Changes

How I tested this

Notes

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.

[elijahbenizzy]

This works nicely. See this script:

import numpy as np
import pandas as pd

from hamilton import driver
from hamilton.ad_hoc_utils import create_temporary_module
from hamilton.function_modifiers import parameterize, source

NUM_ITERS = 100

import pandas as pd
import numpy as np


def foo_0(memory_size: int = 100_000_000) -> pd.DataFrame:
    """
    Generates a large DataFrame with memory size close to the specified memory_size_gb.

    Parameters:
    memory_size_gb (float): Desired memory size of the DataFrame in GB. Default is 1 GB.

    Returns:
    pd.DataFrame: Generated DataFrame with approximate memory usage of memory_size_gb.
    """
    # Number of rows in the DataFrame
    num_rows = 10 ** 6

    # Calculate the number of columns required to make a DataFrame close to memory_size_gb
    # Assuming float64 type which takes 8 bytes
    bytes_per_row = 8 * num_rows
    target_bytes = memory_size
    num_cols = target_bytes // bytes_per_row

    # Create a DataFrame with random data
    data = {f'col_{i}': np.random.random(num_rows) for i in range(int(num_cols))}
    df = pd.DataFrame(data)

    # Print DataFrame info, including memory usage
    print(df.info(memory_usage='deep'))
    return df


@parameterize(
    **{f"foo_{i}" : {"foo_i_minus_one" : source(f"foo_{i-1}")}  for i in range(1,NUM_ITERS)}
)
def foo_i(foo_i_minus_one: pd.DataFrame) -> pd.DataFrame:
    print(f"foo_{i}")
    return foo_i_minus_one * 1.01

if __name__ == '__main__':
    mod = create_temporary_module(foo_i, foo_0)
    dr = driver.Builder().with_modules(mod).build()
    output = dr.execute(
        [f"foo_{NUM_ITERS-1}"],
        inputs=dict(memory_size=1_000_000_000)
    )

Using mprof, we can easily run it, and get the following plot. My computer dies without that code:

[elijahbenizzy]

If you comment out the new code, you'll notice that things act strangely. The caveat here is that mprof doesn't record swap space. So, you see a large alocation, then things start to act quite junky (my computer's music skips 😆). Note this is 100mb increments instead of 1gb. 1gb will just give up.

Now, with the code added, it caps out at 500mb before GC -- I think this is storing a few pointers to the node + general python memory usage. Its keeping a few extra ones around -- Hard to tell why, as there are likely more than one reference being kept around due to the way the del function works with dicts/when GC is happening:

[skrawcz]

I think we have enough pushed that 1.32.0 is warrented.