From 634f513e31fe884993b5d1869b9f1efa4bbbf4d8 Mon Sep 17 00:00:00 2001 From: juacrumar Date: Mon, 27 Jul 2026 11:05:24 +0200 Subject: [PATCH 1/5] fix hyperopt leak by running single trials in a separate subprocess --- .../backends/keras_backend/internal_state.py | 23 +++++++ .../n3fit/hyper_optimization/hyper_scan.py | 64 ++++++++++++++++++- n3fit/src/n3fit/model_trainer.py | 1 - 3 files changed, 85 insertions(+), 3 deletions(-) diff --git a/n3fit/src/n3fit/backends/keras_backend/internal_state.py b/n3fit/src/n3fit/backends/keras_backend/internal_state.py index 2404aa2ca7..93c026d192 100644 --- a/n3fit/src/n3fit/backends/keras_backend/internal_state.py +++ b/n3fit/src/n3fit/backends/keras_backend/internal_state.py @@ -18,8 +18,16 @@ from keras import backend as K import numpy as np +from validphys.convolution import central_predictions, predictions + log = logging.getLogger(__name__) + +def _internal_backend_clear(): + """Keras might need extra help to clear some backends.""" + pass + + # Prepare Keras-backend dependent functions if (kback := K.backend()) == "torch": @@ -36,6 +44,7 @@ def set_threading(threads, core, double_precision=False): elif K.backend() == "tensorflow": import tensorflow as tf + from tensorflow.python.framework import ops def set_eager(flag=True): """Set eager mode on or off @@ -62,6 +71,14 @@ def set_threading(threads, cores, double_precision=False): "Could not set tensorflow parallelism settings from n3fit, maybe tensorflow is already initialized by a third program" ) + def _internal_backend_clear(): + """TensorFlow saves the gradient of the custom models we create as custom gradients. + These are not followed by keras and are thus not cleared during the clear_backend call.""" + registry = ops.gradient_registry._registry + for name in tuple(registry): # so python don't complain about the waning dictionary + if name.startswith("CustomGradient-"): + registry.pop(name) + elif K.backend() == "jax": import jax @@ -120,6 +137,12 @@ def clear_backend_state(): and unused memory. """ log.info("Clearing session") + # The cache of validphys' predictions might be keeping a reference to the N3PDF model, remove + central_predictions.cache_clear() + predictions.cache_clear() + # Then clear the backend + _internal_backend_clear() + # and finally Keras K.clear_session() diff --git a/n3fit/src/n3fit/hyper_optimization/hyper_scan.py b/n3fit/src/n3fit/hyper_optimization/hyper_scan.py index 0e2229d102..753f1be16d 100644 --- a/n3fit/src/n3fit/hyper_optimization/hyper_scan.py +++ b/n3fit/src/n3fit/hyper_optimization/hyper_scan.py @@ -16,6 +16,9 @@ import contextlib import copy import logging +import multiprocessing +import os +import traceback try: import hyperopt @@ -47,7 +50,58 @@ HYPEROPT_STATUSES = {True: "ok", False: "fail"} -HYPEROPT_SEED = 42 +HYPEROPT_SEED = int(os.environ.get("HYPEROPT_SEED", 42)) +print(f"Running with hyperopt seed: {HYPEROPT_SEED}") + + +def _run_trial_in_subprocess(objective, params): + """Running a hyperparameter scan leaks memory trial by trial. + In order to ensure that the memory associated to one trial is eliminated once it has finished, + run it as a separate (forked) process that will die at the end and its memories gone. + Parallel hyperopt will skip this. + """ + # Note: this strategy seems to work, but I haven't found the actual source of the leak + + if not hasattr(os, "fork"): + return objective(params) + + context = multiprocessing.get_context("fork") + receiver, sender = context.Pipe(duplex=False) + + def _run_objective(): + # The child will ever send info, close the receiver + receiver.close() + try: + sender.send((True, objective(params))) + except Exception: + # On any failure, propagate it back to the parent process + sender.send((False, traceback.format_exc())) + finally: + sender.close() + + process = context.Process(target=_run_objective) + process.start() + # The parent only receives, close its sender + sender.close() + try: + # Wait for the child to finish running + success, ret = receiver.recv() + except EOFError as error: + process.join() + exitcode = process.exitcode + process.close() + raise RuntimeError(f"Hyperopt subprocess failed with error= {exitcode}") from error + finally: + receiver.close() + process.join() + exitcode = process.exitcode + log.debug(f"Hyperopt subprocess finished with {exitcode=}") + process.close() + + if not success: + raise RuntimeError(f"Hyperopt trial subprocess failed:\n{ret}") + + return ret # These are just wrapper around some hyperopt's sampling expresions defined in here @@ -172,8 +226,14 @@ def hyper_scan_wrapper(replica_path_set, model_trainer, hyperscanner, max_evals= # Initialize seed for hyperopt trials.rstate = np.random.default_rng(HYPEROPT_SEED) # And prepare the generic arguments to fmin + objective = model_trainer.hyperparametrizable + if not hyperscanner.parallel_hyperopt and hasattr(os, "fork"): + objective = lambda params: _run_trial_in_subprocess( + model_trainer.hyperparametrizable, params + ) + fmin_args = { - "fn": model_trainer.hyperparametrizable, + "fn": objective, "space": hyperscanner.as_dict(), "algo": hyperopt.tpe.suggest, "max_evals": max_evals, diff --git a/n3fit/src/n3fit/model_trainer.py b/n3fit/src/n3fit/model_trainer.py index 9454f895db..44619adb1f 100644 --- a/n3fit/src/n3fit/model_trainer.py +++ b/n3fit/src/n3fit/model_trainer.py @@ -857,7 +857,6 @@ def hyperparametrizable(self, params): All other parameters are passed to the corresponding functions """ - # Reset the internal state of the backend every time this function is called print("") clear_backend_state() From 7cb12d42e92ce7482f5986cabf99bf51af6e8fba Mon Sep 17 00:00:00 2001 From: juacrumar Date: Mon, 27 Jul 2026 21:30:42 +0200 Subject: [PATCH 2/5] fork only under linux --- n3fit/src/n3fit/hyper_optimization/hyper_scan.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/n3fit/src/n3fit/hyper_optimization/hyper_scan.py b/n3fit/src/n3fit/hyper_optimization/hyper_scan.py index 753f1be16d..9c6373feef 100644 --- a/n3fit/src/n3fit/hyper_optimization/hyper_scan.py +++ b/n3fit/src/n3fit/hyper_optimization/hyper_scan.py @@ -58,11 +58,10 @@ def _run_trial_in_subprocess(objective, params): """Running a hyperparameter scan leaks memory trial by trial. In order to ensure that the memory associated to one trial is eliminated once it has finished, run it as a separate (forked) process that will die at the end and its memories gone. - Parallel hyperopt will skip this. + It should only be used in Linux and sequential runs. """ # Note: this strategy seems to work, but I haven't found the actual source of the leak - - if not hasattr(os, "fork"): + if not hasattr(os, "fork") or os.uname().sysname == "Darwin": return objective(params) context = multiprocessing.get_context("fork") From 359c274ef869f4da67f7639d99388a851c75cde4 Mon Sep 17 00:00:00 2001 From: "Juan M. Cruz-Martinez" Date: Tue, 1 Sep 2026 09:29:26 +0200 Subject: [PATCH 3/5] remove print to stdout --- n3fit/src/n3fit/hyper_optimization/hyper_scan.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/n3fit/src/n3fit/hyper_optimization/hyper_scan.py b/n3fit/src/n3fit/hyper_optimization/hyper_scan.py index 9c6373feef..ad62ed4ace 100644 --- a/n3fit/src/n3fit/hyper_optimization/hyper_scan.py +++ b/n3fit/src/n3fit/hyper_optimization/hyper_scan.py @@ -48,11 +48,7 @@ # Hyperopt uses these strings for a passed and failed run # it also has statuses "new", "running" and "suspended", but we don't use them HYPEROPT_STATUSES = {True: "ok", False: "fail"} - - HYPEROPT_SEED = int(os.environ.get("HYPEROPT_SEED", 42)) -print(f"Running with hyperopt seed: {HYPEROPT_SEED}") - def _run_trial_in_subprocess(objective, params): """Running a hyperparameter scan leaks memory trial by trial. From 170822336c3c0f62d225abe31efba4434e3505e4 Mon Sep 17 00:00:00 2001 From: juacrumar Date: Mon, 7 Sep 2026 19:04:21 +0200 Subject: [PATCH 4/5] add review comments --- .../backends/keras_backend/internal_state.py | 19 ++++++++++--------- n3fit/src/n3fit/model_trainer.py | 5 +++++ 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/n3fit/src/n3fit/backends/keras_backend/internal_state.py b/n3fit/src/n3fit/backends/keras_backend/internal_state.py index 93c026d192..8886509f71 100644 --- a/n3fit/src/n3fit/backends/keras_backend/internal_state.py +++ b/n3fit/src/n3fit/backends/keras_backend/internal_state.py @@ -18,8 +18,6 @@ from keras import backend as K import numpy as np -from validphys.convolution import central_predictions, predictions - log = logging.getLogger(__name__) @@ -74,10 +72,16 @@ def set_threading(threads, cores, double_precision=False): def _internal_backend_clear(): """TensorFlow saves the gradient of the custom models we create as custom gradients. These are not followed by keras and are thus not cleared during the clear_backend call.""" - registry = ops.gradient_registry._registry - for name in tuple(registry): # so python don't complain about the waning dictionary - if name.startswith("CustomGradient-"): - registry.pop(name) + try: + registry = ops.gradient_registry._registry + for name in tuple(registry): # so python don't complain about the waning dictionary + if name.startswith("CustomGradient-"): + registry.pop(name) + except Exception as e: + log.error( + "Error found when trying to use internal Tensorflow features to clean memory. Please report this issue." + ) + raise e elif K.backend() == "jax": @@ -137,9 +141,6 @@ def clear_backend_state(): and unused memory. """ log.info("Clearing session") - # The cache of validphys' predictions might be keeping a reference to the N3PDF model, remove - central_predictions.cache_clear() - predictions.cache_clear() # Then clear the backend _internal_backend_clear() # and finally Keras diff --git a/n3fit/src/n3fit/model_trainer.py b/n3fit/src/n3fit/model_trainer.py index 44619adb1f..5ad6550ff1 100644 --- a/n3fit/src/n3fit/model_trainer.py +++ b/n3fit/src/n3fit/model_trainer.py @@ -26,6 +26,7 @@ from n3fit.scaler import generate_scaler from n3fit.stopping import Stopping from n3fit.vpinterface import N3PDF, compute_hyperopt_metrics +from validphys.convolution import central_predictions, predictions from validphys.core import DataGroupSpec from validphys.loader import Loader from validphys.photon.compute import Photon @@ -860,6 +861,10 @@ def hyperparametrizable(self, params): # Reset the internal state of the backend every time this function is called print("") clear_backend_state() + # Clean also validphys' internal caches which keep references to n3fit models + central_predictions.cache_clear() + predictions.cache_clear() + # When doing hyperopt some entries in the params dictionary # can bring with them overriding arguments if self.mode_hyperopt: From d7bf7c2134b17f6ff9a8f1c49f0b7941921d2073 Mon Sep 17 00:00:00 2001 From: juacrumar Date: Tue, 8 Sep 2026 10:10:16 +0200 Subject: [PATCH 5/5] add tests for forking doing the expected thing --- n3fit/src/n3fit/tests/test_hyperopt.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/n3fit/src/n3fit/tests/test_hyperopt.py b/n3fit/src/n3fit/tests/test_hyperopt.py index ad4a83f595..d12aaa441c 100644 --- a/n3fit/src/n3fit/tests/test_hyperopt.py +++ b/n3fit/src/n3fit/tests/test_hyperopt.py @@ -13,6 +13,7 @@ from numpy.testing import assert_approx_equal import pytest +from n3fit.hyper_optimization.hyper_scan import _run_trial_in_subprocess from n3fit.hyper_optimization.rewards import HyperLoss from n3fit.model_gen import ReplicaSettings, generate_pdf_model from n3fit.tests.helpers import run_n3fit, run_setupfit @@ -291,3 +292,21 @@ def test_parallel_restart(tmp_path): assert initial_json[i]['state'] == final_json[i]['state'] assert initial_json[i]['tid'] == final_json[i]['tid'] assert initial_json[i]['result'] == final_json[i]['result'] + + +@pytest.mark.darwin +def test_hyperopt_runs_normally_in_macos(): + """Checks that hyperopt doesn't try to fork out in macos.""" + current_pid = os.getpid() + not_forked_pid = _run_trial_in_subprocess(lambda _: os.getpid(), None) + assert current_pid == not_forked_pid + + +@pytest.mark.linux +def test_hyperopt_runs_in_subprocess(): + """Check that, in linux, hyperopt is able to fork out. + In a linux system with no fork, this test will fail. Report this situation. + """ + current_pid = os.getpid() + forked_pid = _run_trial_in_subprocess(lambda _: os.getpid(), None) + assert current_pid != forked_pid