roualdes / bridgestan

BridgeStan provides efficient in-memory access through Python, Julia, and R to the methods of a Stan model.
https://roualdes.github.io/bridgestan
BSD 3-Clause "New" or "Revised" License
87 stars 12 forks source link

Fix memory leak in log_density(propto=True) #180

Closed WardBrian closed 8 months ago

WardBrian commented 8 months ago

When I submitted #165 to fix the unnecessary computation in log_density(propto=True) I didn't consider that removing the call to grad also meant the var stack was no longer being cleared, so the memory would build up over time.

You can see this pretty easily in the current version, e.g.:

import numpy as np
import bridgestan

m = bridgestan.StanModel("./test_models/gaussian/gaussian_model.so", "./test_models/gaussian/gaussian.data.json")

def print_lp():
    p = np.random.randn(2)
    lp = m.log_density(p, propto=True)
    print(lp)

for _ in range(100_000):
    print_lp()

Running this in a repl, you can see that the memory used by the process grows throughout the loop and stays high. After this PR, it does not.

WardBrian commented 8 months ago

Good point. We do catch the top level exceptions somewhere else, so we can catch, recover memory, and re-raise

bob-carpenter commented 8 months ago

That's a better idea (to re-raise C++ exception). Just be sure to catch both ... (which you can only regrow generically) and also catch std::exception, from which you can extract message to rethrow.

WardBrian commented 8 months ago

... includes std::exception, so there's no need to separate the cases here (they are separate in the top-level function that calls this, though!). We have a test for the log density failing with an error message which is still passing

WardBrian commented 8 months ago

The other option we could take is using what gradient itself uses which is the nested_rev_autodiff RAII class which cleans up memory in its destructor. We don't need nested autodiff (since we don't need any autodiff) but I believe it would be functionally equivalent in this case, and perhaps a bit cleaner.

bob-carpenter commented 8 months ago

Thanks. Happy to review another PR if you think it'll be cleaner.