From 82c2ed7774d5c632068f451c5baacc244b7e6caf Mon Sep 17 00:00:00 2001 From: Roger Pros Rius Date: Wed, 13 May 2026 09:08:13 +0200 Subject: [PATCH] Fix memory blowup in GAM.gridsearch over large lambda grids --- pygam/pygam.py | 67 +++++++----- pygam/tests/test_memory_leak_gridsearch.py | 114 +++++++++++++++++++++ pyproject.toml | 1 + 3 files changed, 155 insertions(+), 27 deletions(-) create mode 100644 pygam/tests/test_memory_leak_gridsearch.py diff --git a/pygam/pygam.py b/pygam/pygam.py index 6f86ae43..afccfe7a 100644 --- a/pygam/pygam.py +++ b/pygam/pygam.py @@ -2016,41 +2016,47 @@ def gridsearch( param_grid_list.append(dict(zip(params, candidate))) # set up data collection - best_model = None # keep the best model + best_model = None best_score = np.inf scores = [] - models = [] + last_model = None + + # Store the class of the model to instantiate fresh copies later + ModelClass = self.__class__ + # Get the base parameters of the current model + base_params = self.get_params() + + # Only create the list if strictly needed for return_scores + models = [] if return_scores else None + last_model = None - # check if our model has been fitted already and store it if self._is_fitted: - models.append(self) + if return_scores: + models.append(self) + last_model = self scores.append(self.statistics_[objective]) - - # our model is currently the best - best_model = models[-1] + best_model = self best_score = scores[-1] - # make progressbar optional if progress: pbar = ProgressBar() + else: def pbar(x): return x - # loop through candidate model params for param_grid in pbar(param_grid_list): try: - # try fitting - # define new model - gam = deepcopy(self) - gam.set_params(self.get_params()) + gam = ModelClass(**base_params) + gam.set_params(**param_grid) # warm start with parameters from previous build - if models: - coef = models[-1].coef_ + if last_model is not None: + coef = last_model.coef_.copy() gam.set_params(coef_=coef, force=True, verbose=False) + gam.fit(X, y, weights) except ValueError as error: @@ -2060,29 +2066,36 @@ def pbar(x): warnings.warn(msg) continue - # record results - models.append(gam) - scores.append(gam.statistics_[objective]) + current_score = float(gam.statistics_[objective]) + scores.append(current_score) + + last_model = gam - # track best - if scores[-1] < best_score: - best_model = models[-1] - best_score = scores[-1] + if return_scores: + models.append(gam) - # problems - if len(models) == 0: + if current_score < best_score: + best_model = gam + best_score = current_score + + if gam is not best_model and gam is not last_model: + del gam + + if best_model is None: msg = "No models were fitted." if self.verbose: warnings.warn(msg) return self - # copy over the best if keep_best: self.set_params(deep=True, force=True, **best_model.get_params(deep=True)) + + del best_model + del last_model + if return_scores: return OrderedDict(zip(models, scores)) - else: - return self + return self def sample( self, diff --git a/pygam/tests/test_memory_leak_gridsearch.py b/pygam/tests/test_memory_leak_gridsearch.py new file mode 100644 index 00000000..3392c967 --- /dev/null +++ b/pygam/tests/test_memory_leak_gridsearch.py @@ -0,0 +1,114 @@ +import gc +import os +import threading +import time + +import numpy as np +import psutil +import pytest + +from pygam import LinearGAM, s + +# --- Configuration --- +# LEAK TOLERANCE: How much extra memory (MB) can be left over? +# We allow a small buffer. +LEAK_TOLERANCE_MB = 50 + +# PEAK THRESHOLD: If memory usage exceeds 'X' times the starting memory, fail. +PEAK_MEMORY_MULTIPLIER = 4.0 + + +class PeakMemoryMonitor: + def __init__(self, pid): + self.process = psutil.Process(pid) + self.stop_event = threading.Event() + self.peak_memory = 0.0 + self.start_memory = 0.0 + self.thread = threading.Thread(target=self._monitor) + + def _monitor(self): + while not self.stop_event.is_set(): + # Get current memory in MB + current_mem = self.process.memory_info().rss / 1024 / 1024 + print(f"Current Memory: {current_mem:.2f} MB") + + if current_mem > self.peak_memory: + self.peak_memory = current_mem + + if self.peak_memory > self.start_memory * PEAK_MEMORY_MULTIPLIER: + pytest.fail( + f"Peak memory usage ({self.peak_memory:.2f} MB) exceeded safety limit" + ) + + time.sleep(0.5) + + def start(self): + self.start_memory = self.process.memory_info().rss / 1024 / 1024 + self.peak_memory = self.start_memory + self.thread.start() + + def stop(self): + self.stop_event.set() + self.thread.join() + + +def test_gridsearch_memory_robust(): + """ + Robustly tests for memory leaks by comparing Start vs. End memory + and ensuring peak memory doesn't explode. + """ + N_SAMPLES = 1000 + N_GRID_POINTS = 100 + + rng = np.random.default_rng(42) + X = rng.standard_normal((N_SAMPLES, 10)) + y = np.sin(X[:, 0]) + X[:, 1] ** 2 + rng.standard_normal(N_SAMPLES) * 0.1 + + gam = LinearGAM(s(0) + s(1) + s(2)) + lam_grid = np.logspace(-3, 3, N_GRID_POINTS) + + # Force GC to ensure we start at a true baseline + gc.collect() + time.sleep(0.5) + + monitor = PeakMemoryMonitor(os.getpid()) + monitor.start() + + print(f"\n[Test Info] Starting Memory: {monitor.start_memory:.2f} MB") + + try: + gam.gridsearch(X, y, lam=lam_grid, progress=False) + except Exception as e: + pytest.fail(f"Gridsearch failed with error: {e}") + finally: + monitor.stop() + + gc.collect() + + process = psutil.Process(os.getpid()) + end_memory = process.memory_info().rss / 1024 / 1024 + + print(f"[Test Info] Peak Memory: {monitor.peak_memory:.2f} MB") + print(f"[Test Info] Ending Memory: {end_memory:.2f} MB") + + max_allowed_peak = monitor.start_memory * PEAK_MEMORY_MULTIPLIER + if monitor.peak_memory > max_allowed_peak: + pytest.fail( + f"Peak memory usage ({monitor.peak_memory:.2f} MB) exceeded safety limit " + f"({max_allowed_peak:.2f} MB). This suggests runaway memory usage." + ) + + diff = end_memory - monitor.start_memory + + if diff > LEAK_TOLERANCE_MB: + pytest.fail( + f"Memory Leak Detected! Memory grew by {diff:.2f} MB (Threshold: {LEAK_TOLERANCE_MB} MB). " + f"Start: {monitor.start_memory:.2f} MB, End: {end_memory:.2f} MB." + ) + + print(f"[Test Info] Memory Difference: {diff:.2f} MB (PASSED)") + + +if __name__ == "__main__": + test_gridsearch_memory_robust() + print("Test finished successfully.") diff --git a/pyproject.toml b/pyproject.toml index cc726ee5..399a521f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,6 +46,7 @@ dev = [ "nbsphinx>=0.9.0,<1", "numpydoc>=1.8.0", "pandas>=2.0", + "psutil>=5.9.0", "pytest>=9.0.0", "pydata-sphinx-theme>=0.15.0", "pytest-cov>=7.0.0",