context_minimizer.py

import zfit
import tempfile
from zfit                          import dill                     
from pathlib                       import Path
from multiprocessing               import get_context

class ContextMinimizer:
    '''
    Class meant to run minimization in a separate process
    This is needed to deal with the memory leak described here:

    https://github.com/zfit/zfit/issues/665

    And could be removed once issue is fixed
    '''
    def __init__(self):
        '''
        '''
        self._ctx = get_context(method = 'spawn')
    # ----------------------
    @staticmethod
    def _minimize(obj : bytes, path : Path):
        '''
        Parameters
        ----------------
        obj : Bytes representation of likelihood
        path: Path where the result object will be saved
        '''
        import zfit
        from zfit import dill

        nll = dill.loads(obj)
        min = zfit.minimize.Minuit()
        res = min.minimize(loss = nll)

        with open(path, 'w+b') as ofile:
            dill.dump(res, ofile)
    # ----------------------
    def minimize(self, nll : zlos) -> zres:
        '''
        Parameters
        -------------
        nll: zfit Negative log likelihood

        Returns
        -------------
        Zfit result object
        '''
        nll_byt = dill.dumps(nll)
        with tempfile.NamedTemporaryFile(mode='w+b', delete=True) as temp:
            path = Path(temp.name)

            proc = self._ctx.Process(target = self._minimize, args = (nll_byt,path))
            proc.start()
            proc.join()

            with open(path, 'rb') as ifile:
                res = dill.load(ifile)

        return res