[core] weak reference observers to prevent leaks

Author: terhorst

smcpp/_smcpp.pyx

@@ -123,6 +123,8 @@ cdef class _PyInferenceManager:
     cdef InferenceManager* _im
     cdef vector[int*] _obs_ptrs
 
+    cdef object __weakref__
+
     def __my_cinit__(self, observations, hidden_states, im_id=None):
         self._im_id = im_id
         self.seed = 1

smcpp/analysis.py

@@ -59,7 +59,7 @@ def _init_penalty(self):
     def _init_optimizer(self, outdir, algorithm, xtol, ftol):
         self._optimizer = self._OPTIMIZER_CLS(self, algorithm, xtol, ftol)
         if outdir:
-            self._optimizer.register(analysis_saver.AnalysisSaver(outdir))
+            self._optimizer.register_plugin(analysis_saver.AnalysisSaver(outdir))
 
     def rescale(self, x):
         return x / (2. * self._N0)
@@ -361,7 +361,7 @@ def _init_optimizer(self, outdir, algorithm, xtol, ftol, learn_rho=False):
         super()._init_optimizer(outdir, algorithm, xtol, ftol)
         if learn_rho:
             rho_bounds = 2. * self._N0 * np.array([1e-10, 1e-5])
-            self._optimizer.register(
+            self._optimizer.register_plugin(
                     parameter_optimizer.ParameterOptimizer("rho", tuple(rho_bounds)))
 
 
@@ -391,7 +391,7 @@ def _validate_data(self):
 
     def _init_optimizer(self, outdir, algorithm, xtol, ftol):
         super()._init_optimizer(outdir, algorithm, xtol, ftol)
-        self._optimizer.register(parameter_optimizer.ParameterOptimizer("split",
+        self._optimizer.register_plugin(parameter_optimizer.ParameterOptimizer("split",
                                                    (0., self._max_split),
                                                    "model"))
 

smcpp/observe.py

@@ -1,6 +1,7 @@
 from __future__ import absolute_import
 from abc import ABCMeta, abstractmethod
 import wrapt
+import weakref
 
 
 # Decorator to target specific messages.
@@ -27,19 +28,16 @@ def update(self, *args, **kwargs):
 
 class Observable(object):
     def __init__(self):
-        self.observers = []
+        self.observers = weakref.WeakSet()
 
     def register(self, observer):
-        if not observer in self.observers:
-            self.observers.append(observer)
+        self.observers.add(observer)
 
     def unregister(self, observer):
-        if observer in self.observers:
-            self.observers.remove(observer)
+        self.observers.discard(observer)
 
     def unregister_all(self):
-        if self.observers:
-            del self.observers[:]
+        self.observers.clear()
 
     def update_observers(self, *args, **kwargs):
         for observer in self.observers:

smcpp/optimize/optimizers.py

@@ -21,6 +21,7 @@ class AbstractOptimizer(Observable):
     '''
     def __init__(self, analysis, algorithm, xtol, ftol):
         Observable.__init__(self)
+        self._plugins = []
         self._analysis = analysis
         self._algorithm = algorithm
         self._ftol = ftol
@@ -174,6 +175,10 @@ def _callback(self, xk):
         if self._delta < self._xtol:
             raise ConvergedException("delta=%f < xtol=%f" % (self._delta, self._xtol))
 
+    def register_plugin(self, p):
+        self._plugins.append(p)
+        self.register(p)
+
     def update_observers(self, *args, **kwargs):
         kwargs.update({
             'optimizer': self,
@@ -190,7 +195,7 @@ def __init__(self, analysis, algorithm, xtol, ftol):
         for cls in OptimizerPlugin.__subclasses__():
             try:
                 if not cls.DISABLED:
-                    self.register(cls())
+                    self.register_plugin(cls())
             except TypeError:
                 # Only register listeners with null constructor
                 pass