PickleGzInstance: Replace implicit load by load/free methods

miplearn/instance/base.py

@@ -158,6 +158,12 @@ class Instance(ABC):
     def build_user_cut(self, model: Any, violation: Hashable) -> Any:
         return None
 
+    def load(self) -> None:
+        pass
+
+    def free(self) -> None:
+        pass
+
     def flush(self) -> None:
         """
         Save any pending changes made to the instance to the underlying data store.

miplearn/instance/picklegz.py

@@ -5,23 +5,13 @@
 import gzip
 import os
 import pickle
+import gc
 from typing import Optional, Any, List, Hashable, cast, IO, Callable
 
 from miplearn.instance.base import logger, Instance
 from miplearn.types import VarIndex
 
 
-def lazy_load(func: Callable) -> Callable:
-    def inner(self: Any, *args: Any) -> Any:
-        if self.instance is None:
-            self.instance = self._load()
-            self.features = self.instance.features
-            self.training_data = self.instance.training_data
-        return func(self, *args)
-
-    return inner
-
-
 class PickleGzInstance(Instance):
     """
     An instance backed by a gzipped pickle file.
@@ -41,22 +31,18 @@ class PickleGzInstance(Instance):
         self.instance: Optional[Instance] = None
         self.filename: str = filename
 
-    @lazy_load
     def to_model(self) -> Any:
         assert self.instance is not None
         return self.instance.to_model()
 
-    @lazy_load
     def get_instance_features(self) -> List[float]:
         assert self.instance is not None
         return self.instance.get_instance_features()
 
-    @lazy_load
     def get_variable_features(self, var_name: str, index: VarIndex) -> List[float]:
         assert self.instance is not None
         return self.instance.get_variable_features(var_name, index)
 
-    @lazy_load
     def get_variable_category(
         self,
         var_name: str,
@@ -65,55 +51,55 @@ class PickleGzInstance(Instance):
         assert self.instance is not None
         return self.instance.get_variable_category(var_name, index)
 
-    @lazy_load
     def get_constraint_features(self, cid: str) -> Optional[List[float]]:
         assert self.instance is not None
         return self.instance.get_constraint_features(cid)
 
-    @lazy_load
     def get_constraint_category(self, cid: str) -> Optional[Hashable]:
         assert self.instance is not None
         return self.instance.get_constraint_category(cid)
 
-    @lazy_load
     def has_static_lazy_constraints(self) -> bool:
         assert self.instance is not None
         return self.instance.has_static_lazy_constraints()
 
-    @lazy_load
     def has_dynamic_lazy_constraints(self) -> bool:
         assert self.instance is not None
         return self.instance.has_dynamic_lazy_constraints()
 
-    @lazy_load
     def is_constraint_lazy(self, cid: str) -> bool:
         assert self.instance is not None
         return self.instance.is_constraint_lazy(cid)
 
-    @lazy_load
     def find_violated_lazy_constraints(self, model: Any) -> List[Hashable]:
         assert self.instance is not None
         return self.instance.find_violated_lazy_constraints(model)
 
-    @lazy_load
     def build_lazy_constraint(self, model: Any, violation: Hashable) -> Any:
         assert self.instance is not None
         return self.instance.build_lazy_constraint(model, violation)
 
-    @lazy_load
     def find_violated_user_cuts(self, model: Any) -> List[Hashable]:
         assert self.instance is not None
         return self.instance.find_violated_user_cuts(model)
 
-    @lazy_load
     def build_user_cut(self, model: Any, violation: Hashable) -> Any:
         assert self.instance is not None
         return self.instance.build_user_cut(model, violation)
 
-    def _load(self) -> Instance:
-        obj = read_pickle_gz(self.filename)
-        assert isinstance(obj, Instance)
-        return obj
+    def load(self) -> None:
+        if self.instance is None:
+            obj = read_pickle_gz(self.filename)
+            assert isinstance(obj, Instance)
+            self.instance = obj
+            self.features = self.instance.features
+            self.training_data = self.instance.training_data
+
+    def free(self) -> None:
+        self.instance = None  # type: ignore
+        self.features = None  # type: ignore
+        self.training_data = None  # type: ignore
+        gc.collect()
 
     def flush(self) -> None:
         write_pickle_gz(self.instance, self.filename)