add ability to handle direct memory allocated datasets in addition to shared memory

Author: paulbkoch

docs/benchmarks/ebm-benchmark.ipynb

@@ -17,7 +17,7 @@
     "force_recreate = False\n",
     "exist_ok = True\n",
     "TIMEOUT_SEC = 60 * 60 * 24 * 180  # 180 days\n",
-    "wheel_filepaths = ['interpret_core-0.6.13-py3-none-any.whl', 'powerlift-0.1.12-py3-none-any.whl']\n",
+    "wheel_filepaths = ['interpret_core-0.6.15-py3-none-any.whl', 'powerlift-0.1.12-py3-none-any.whl']\n",
     "\n",
     "import datetime\n",
     "experiment_name = datetime.datetime.now().strftime('%Y_%m_%d_%H%M__') + 'myexperiment'\n",

python/interpret-core/interpret/glassbox/_ebm/_boost.py

@@ -17,7 +17,7 @@ def boost(
     shm_name,
     bag_idx,
     callback,
-    dataset_name,
+    dataset,
     intercept_rounds,
     intercept_learning_rate,
     intercept,
@@ -56,10 +56,12 @@ def boost(
 ):
     try:
         develop._develop_options = develop_options  # restore these in this process
+        shared_dataset = None
         try:
-            shared_dataset = shared_memory.SharedMemory(name=dataset_name)
-            # we do not know the length of the dataset, so we create a 1-element array
-            dataset = np.ndarray(1, dtype=np.ubyte, buffer=shared_dataset.buf)
+            if isinstance(dataset, str):  # if str it is shared memory
+                shared_dataset = shared_memory.SharedMemory(name=dataset)
+                # we do not know the length of the dataset, so we create a 1-element array
+                dataset = np.ndarray(1, dtype=np.ubyte, buffer=shared_dataset.buf)
 
             shm = None
             try:
@@ -372,6 +374,7 @@ def boost(
                 if shm is not None:
                     shm.close()
         finally:
-            shared_dataset.close()
+            if shared_dataset is not None:
+                shared_dataset.close()
     except Exception as e:
         return e, None, None, None, None

python/interpret-core/interpret/glassbox/_ebm/_ebm.py

@@ -1166,7 +1166,7 @@ def fit(self, X, y, sample_weight=None, bags=None, init_score=None):
                             shm_name,
                             idx,
                             callback,
-                            shared.name,
+                            shared.name if shared.name is not None else shared.dataset,
                             n_intercept_rounds,
                             develop.get_option("intercept_learning_rate"),
                             bagged_intercept[idx],
@@ -1292,7 +1292,9 @@ def fit(self, X, y, sample_weight=None, bags=None, init_score=None):
                                 (
                                     shm_name,
                                     idx,
-                                    shared.name,
+                                    shared.name
+                                    if shared.name is not None
+                                    else shared.dataset,
                                     bagged_intercept[idx],
                                     internal_bags[idx],
                                     scores_bags[idx],
@@ -1434,7 +1436,9 @@ def fit(self, X, y, sample_weight=None, bags=None, init_score=None):
                                 shm_name,
                                 idx,
                                 callback,
-                                shared.name,
+                                shared.name
+                                if shared.name is not None
+                                else shared.dataset,
                                 0,  # intercept should already be close for pairs
                                 0.0,  # intercept should already be close for pairs
                                 bagged_intercept[idx],
@@ -1570,7 +1574,7 @@ def fit(self, X, y, sample_weight=None, bags=None, init_score=None):
                         None,
                         0,
                         None,
-                        shared.name,
+                        shared.dataset,
                         develop.get_option("n_intercept_rounds_final"),
                         develop.get_option("intercept_learning_rate"),
                         np.zeros(n_scores, np.float64),
@@ -1741,8 +1745,9 @@ def estimate_mem(self, X):
         # One shared memory copy of the data mapped into all processes, plus a copy of
         # the test and train data for each outer bag. Assume all processes are started
         # at some point and are eating up memory.
-
-        max_bytes = n_bytes_mains + n_bytes_mains * self.outer_bags
+        # When we cannot use shared memory the parent has a copy of the dataset and
+        # all the children share one copy.
+        max_bytes = n_bytes_mains + n_bytes_mains + n_bytes_mains * self.outer_bags
 
         n_features_in = len(bins)
 
@@ -1766,7 +1771,7 @@ def estimate_mem(self, X):
             # each outer bag makes a copy of the features. Only the training features
             # are kept for interaction detection, but don't estimate that for now.
             interaction_detection_bytes = (
-                n_bytes_pairs + n_bytes_pairs * self.outer_bags
+                n_bytes_pairs + n_bytes_pairs + n_bytes_pairs * self.outer_bags
             )
 
             max_bytes = max(max_bytes, interaction_detection_bytes)
@@ -1777,8 +1782,10 @@ def estimate_mem(self, X):
             # then the resulting data storage should take approx 14 bits in total,
             # so as a loose approximation we can add the bits in a pair.
             interaction_multiple *= 2.0
-            interaction_boosting_bytes = n_bytes_pairs + int(
-                n_bytes_pairs * interaction_multiple * self.outer_bags
+            interaction_boosting_bytes = (
+                n_bytes_pairs
+                + n_bytes_pairs
+                + int(n_bytes_pairs * interaction_multiple * self.outer_bags)
             )
 
             max_bytes = max(max_bytes, interaction_boosting_bytes)

python/interpret-core/interpret/utils/_compressed_dataset.py

@@ -102,11 +102,15 @@ def bin_native(
         raise ValueError(msg)
 
     if shared is not None:
-        shared_mem = shared_memory.SharedMemory(create=True, size=n_bytes, name=None)
-        shared.shared_memory = shared_mem
-        shared.name = shared_mem.name
+        # shared_mem = shared_memory.SharedMemory(create=True, size=n_bytes, name=None)
+        # shared.shared_memory = shared_mem
+        # shared.name = shared_mem.name
+        # dataset = np.ndarray(n_bytes, dtype=np.ubyte, buffer=shared_mem.buf)
+
+        # Large amounts of shared memory cannot be allocated inside docker images
+        # so for now allocate a normal numpy array.
+        dataset = np.empty(n_bytes, np.ubyte)
 
-        dataset = np.ndarray(n_bytes, dtype=np.ubyte, buffer=shared_mem.buf)
         shared.dataset = dataset
 
         native.fill_dataset_header(len(feature_idxs), n_weights, 1, dataset)

python/interpret-core/interpret/utils/_measure_interactions.py

@@ -270,7 +270,7 @@ def measure_interactions(
         ranked_interactions = rank_interactions(
             None,
             0,
-            dataset_name=shared.name,
+            dataset=shared.name if shared.name is not None else shared.dataset,
             intercept=None,
             bag=None,
             init_scores=init_score,

python/interpret-core/interpret/utils/_rank_interactions.py

@@ -20,7 +20,7 @@
 def rank_interactions(
     shm_name,
     bag_idx,
-    dataset_name,
+    dataset,
     intercept,
     bag,
     init_scores,
@@ -44,10 +44,12 @@ def rank_interactions(
     try:
         develop._develop_options = develop_options  # restore these in this process
 
+        shared_dataset = None
         try:
-            shared_dataset = shared_memory.SharedMemory(name=dataset_name)
-            # we do not know the length of the dataset, so we create a 1-element array
-            dataset = np.ndarray(1, dtype=np.ubyte, buffer=shared_dataset.buf)
+            if isinstance(dataset, str):  # if str it is shared memory
+                shared_dataset = shared_memory.SharedMemory(name=dataset)
+                # we do not know the length of the dataset, so we create a 1-element array
+                dataset = np.ndarray(1, dtype=np.ubyte, buffer=shared_dataset.buf)
 
             shm = None
             try:
@@ -100,6 +102,7 @@ def rank_interactions(
                 if shm is not None:
                     shm.close()
         finally:
-            shared_dataset.close()
+            if shared_dataset is not None:
+                shared_dataset.close()
     except Exception as e:
         return e