Fix-ish hyperopt leak

n3fit/src/n3fit/backends/keras_backend/internal_state.py

@@ -20,6 +20,12 @@
 
 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 +42,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 +69,20 @@ 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."""
+        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":
 
     import jax
@@ -120,6 +141,9 @@ def clear_backend_state():
     and unused memory.
     """
     log.info("Clearing session")
+    # Then clear the backend
+    _internal_backend_clear()
+    # and finally Keras
     K.clear_session()
 
 

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
@@ -45,9 +48,55 @@
 # 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))
 
+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.
+    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") or os.uname().sysname == "Darwin":
+        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}")
 
-HYPEROPT_SEED = 42
+    return ret
 
 
 # These are just wrapper around some hyperopt's sampling expresions defined in here
@@ -172,8 +221,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,

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
@@ -857,10 +858,13 @@ 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()
+        # 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:

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