pass datasets as shared memory instead of via default joblib method

Author: paulbkoch

python/interpret-core/interpret-core.pyproj

@@ -94,6 +94,7 @@
     <Compile Include="interpret\utils\_preprocessor.py" />
     <Compile Include="interpret\utils\_privacy.py" />
     <Compile Include="interpret\utils\_seed.py" />
+    <Compile Include="interpret\utils\_shared_dataset.py" />
     <Compile Include="interpret\utils\_synthetic.py" />
     <Compile Include="interpret\utils\_unify_predict.py" />
     <Compile Include="interpret\utils\_unify_data.py" />

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

@@ -7,6 +7,7 @@
 
 from ._clean_x import categorical_encode, unify_columns
 from ._native import Native
+from multiprocessing import shared_memory
 
 _log = logging.getLogger(__name__)
 
@@ -20,6 +21,7 @@ def bin_native(
     sample_weight,
     feature_names_in,
     feature_types_in,
+    shared,
 ):
     # called under: fit
 
@@ -99,7 +101,12 @@ def bin_native(
         _log.error(msg)
         raise ValueError(msg)
 
-    dataset = np.empty(n_bytes, np.ubyte)  # joblib loky doesn't support RawArray
+    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)
+    shared.dataset = dataset
 
     native.fill_dataset_header(len(feature_idxs), n_weights, 1, dataset)
 
@@ -153,8 +160,6 @@ def bin_native(
         _log.error(msg)
         raise ValueError(msg)
 
-    return dataset
-
 
 def bin_native_by_dimension(
     n_classes,
@@ -165,6 +170,7 @@ def bin_native_by_dimension(
     sample_weight,
     feature_names_in,
     feature_types_in,
+    shared,
 ):
     # called under: fit
 
@@ -175,7 +181,7 @@ def bin_native_by_dimension(
         feature_bins = bin_levels[min(len(bin_levels), n_dimensions) - 1]
         bins_iter.append(feature_bins)
 
-    return bin_native(
+    bin_native(
         n_classes,
         feature_idxs,
         bins_iter,
@@ -184,4 +190,5 @@ def bin_native_by_dimension(
         sample_weight,
         feature_names_in,
         feature_types_in,
+        shared,
     )

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

@@ -26,6 +26,7 @@
 from ._native import Native
 from ._preprocessor import construct_bins
 from ._rank_interactions import rank_interactions
+from ._shared_dataset import SharedDataset
 
 _log = logging.getLogger(__name__)
 
@@ -239,59 +240,61 @@ def measure_interactions(
     bins = binning_result[2]
     n_features_in = len(bins)
 
-    dataset = bin_native_by_dimension(
-        n_classes=n_classes,
-        n_dimensions=2,
-        bins=bins,
-        X=X,
-        y=y,
-        sample_weight=sample_weight,
-        feature_names_in=feature_names_in,
-        feature_types_in=feature_types_in,
-    )
-
-    interaction_flags = Native.CalcInteractionFlags_Default
-    if develop.get_option("full_interaction"):
-        interaction_flags |= Native.CalcInteractionFlags_Full
-
-    if isinstance(interactions, int):
-        n_output_interactions = interactions
-        iter_term_features = combinations(range(n_features_in), 2)
-    elif interactions is None:
-        n_output_interactions = 0
-        iter_term_features = combinations(range(n_features_in), 2)
-    else:
-        n_output_interactions = 0
-        iter_term_features = interactions
-
-    ranked_interactions = rank_interactions(
-        None,
-        0,
-        dataset=dataset,
-        intercept=None,
-        bag=None,
-        init_scores=init_score,
-        iter_term_features=iter_term_features,
-        exclude=set(),
-        exclude_features=set(),
-        calc_interaction_flags=interaction_flags,
-        max_cardinality=max_cardinality,
-        min_samples_leaf=min_samples_leaf,
-        min_hessian=min_hessian,
-        reg_alpha=reg_alpha,
-        reg_lambda=reg_lambda,
-        max_delta_step=max_delta_step,
-        create_interaction_flags=(
-            Native.CreateInteractionFlags_DifferentialPrivacy
-            if is_differential_privacy
-            else Native.CreateInteractionFlags_Default
-        ),
-        objective=objective,
-        acceleration=develop.get_option("acceleration"),
-        experimental_params=None,
-        n_output_interactions=n_output_interactions,
-        develop_options=develop._develop_options,
-    )
+    with SharedDataset() as shared:
+        bin_native_by_dimension(
+            n_classes=n_classes,
+            n_dimensions=2,
+            bins=bins,
+            X=X,
+            y=y,
+            sample_weight=sample_weight,
+            feature_names_in=feature_names_in,
+            feature_types_in=feature_types_in,
+            shared=shared,
+        )
+
+        interaction_flags = Native.CalcInteractionFlags_Default
+        if develop.get_option("full_interaction"):
+            interaction_flags |= Native.CalcInteractionFlags_Full
+
+        if isinstance(interactions, int):
+            n_output_interactions = interactions
+            iter_term_features = combinations(range(n_features_in), 2)
+        elif interactions is None:
+            n_output_interactions = 0
+            iter_term_features = combinations(range(n_features_in), 2)
+        else:
+            n_output_interactions = 0
+            iter_term_features = interactions
+
+        ranked_interactions = rank_interactions(
+            None,
+            0,
+            dataset_name=shared.name,
+            intercept=None,
+            bag=None,
+            init_scores=init_score,
+            iter_term_features=iter_term_features,
+            exclude=set(),
+            exclude_features=set(),
+            calc_interaction_flags=interaction_flags,
+            max_cardinality=max_cardinality,
+            min_samples_leaf=min_samples_leaf,
+            min_hessian=min_hessian,
+            reg_alpha=reg_alpha,
+            reg_lambda=reg_lambda,
+            max_delta_step=max_delta_step,
+            create_interaction_flags=(
+                Native.CreateInteractionFlags_DifferentialPrivacy
+                if is_differential_privacy
+                else Native.CreateInteractionFlags_Default
+            ),
+            objective=objective,
+            acceleration=develop.get_option("acceleration"),
+            experimental_params=None,
+            n_output_interactions=n_output_interactions,
+            develop_options=develop._develop_options,
+        )
 
     if isinstance(ranked_interactions, Exception):
         raise ranked_interactions

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

@@ -20,7 +20,7 @@
 def rank_interactions(
     shm_name,
     bag_idx,
-    dataset,
+    dataset_name,
     intercept,
     bag,
     init_scores,
@@ -44,55 +44,62 @@ def rank_interactions(
     try:
         develop._develop_options = develop_options  # restore these in this process
 
-        shm = None
         try:
-            stop_flag = None
-            if shm_name is not None:
-                shm = shared_memory.SharedMemory(name=shm_name)
-                stop_flag = np.ndarray((1,), dtype=np.bool_, buffer=shm.buf)
+            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)
 
-            interaction_strengths = []
-            with InteractionDetector(
-                dataset,
-                intercept,
-                bag,
-                init_scores,
-                create_interaction_flags,
-                objective,
-                acceleration,
-                experimental_params,
-            ) as interaction_detector:
-                for feature_idxs in iter_term_features:
-                    if tuple(sorted(feature_idxs)) in exclude:
-                        continue
-                    if any(i in exclude_features for i in feature_idxs):
-                        continue
+            shm = None
+            try:
+                stop_flag = None
+                if shm_name is not None:
+                    shm = shared_memory.SharedMemory(name=shm_name)
+                    stop_flag = np.ndarray(1, dtype=np.bool_, buffer=shm.buf)
 
-                    strength = interaction_detector.calc_interaction_strength(
-                        feature_idxs,
-                        calc_interaction_flags,
-                        max_cardinality,
-                        min_samples_leaf,
-                        min_hessian,
-                        reg_alpha,
-                        reg_lambda,
-                        max_delta_step,
-                    )
-                    item = (strength, feature_idxs)
-                    if n_output_interactions <= 0:
-                        interaction_strengths.append(item)
-                    elif len(interaction_strengths) == n_output_interactions:
-                        heapq.heappushpop(interaction_strengths, item)
-                    else:
-                        heapq.heappush(interaction_strengths, item)
+                interaction_strengths = []
+                with InteractionDetector(
+                    dataset,
+                    intercept,
+                    bag,
+                    init_scores,
+                    create_interaction_flags,
+                    objective,
+                    acceleration,
+                    experimental_params,
+                ) as interaction_detector:
+                    for feature_idxs in iter_term_features:
+                        if tuple(sorted(feature_idxs)) in exclude:
+                            continue
+                        if any(i in exclude_features for i in feature_idxs):
+                            continue
 
-                    if stop_flag is not None and stop_flag[0]:
-                        break
+                        strength = interaction_detector.calc_interaction_strength(
+                            feature_idxs,
+                            calc_interaction_flags,
+                            max_cardinality,
+                            min_samples_leaf,
+                            min_hessian,
+                            reg_alpha,
+                            reg_lambda,
+                            max_delta_step,
+                        )
+                        item = (strength, feature_idxs)
+                        if n_output_interactions <= 0:
+                            interaction_strengths.append(item)
+                        elif len(interaction_strengths) == n_output_interactions:
+                            heapq.heappushpop(interaction_strengths, item)
+                        else:
+                            heapq.heappush(interaction_strengths, item)
 
-            interaction_strengths.sort(reverse=True)
-            return interaction_strengths
+                        if stop_flag is not None and stop_flag[0]:
+                            break
+
+                interaction_strengths.sort(reverse=True)
+                return interaction_strengths
+            finally:
+                if shm is not None:
+                    shm.close()
         finally:
-            if shm is not None:
-                shm.close()
+            shared_dataset.close()
     except Exception as e:
         return e

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

@@ -0,0 +1,25 @@
+# Copyright (c) 2023 The InterpretML Contributors
+# Distributed under the MIT software license
+
+from contextlib import AbstractContextManager
+
+class SharedDataset(AbstractContextManager):
+    def __init__(self):
+        self.shared_memory = None
+        self.dataset = None
+        self.name = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.reset()
+
+    def reset(self):
+        shared_memory = self.shared_memory
+        self.name = None
+        self.dataset = None
+        self.shared_memory = None
+        if shared_memory is not None:
+            shared_memory.close()
+            shared_memory.unlink()

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

@@ -9,7 +9,7 @@
 from interpret.utils._clean_x import preclean_X
 from interpret.utils._compressed_dataset import bin_native, bin_native_by_dimension
 from interpret.utils._preprocessor import construct_bins
-
+from interpret.utils._shared_dataset import SharedDataset
 
 @pytest.mark.skip(reason="skip this until we have support for missing values")
 def test_bin_native():
@@ -79,19 +79,27 @@ def test_bin_native():
         feature_idxs.append(feature_idx)
         bins_iter.append(feature_bins)
 
-    shared_dataset = bin_native(
-        n_classes,
-        feature_idxs,
-        bins_iter,
-        X,
-        y,
-        sample_weight,
-        feature_names_in,
-        feature_types_in,
-    )
-    assert shared_dataset is not None
+    with SharedDataset() as shared:
+        bin_native(
+            n_classes,
+            feature_idxs,
+            bins_iter,
+            X,
+            y,
+            sample_weight,
+            feature_names_in,
+            feature_types_in,
+            shared,
+        )
+        assert shared.shared_memory is not None
+        assert shared.dataset is not None
+        assert shared.name is not None
 
-    shared_dataset = bin_native_by_dimension(
-        n_classes, 1, bins, X, y, sample_weight, feature_names_in, feature_types_in
-    )
-    assert shared_dataset is not None
+    with SharedDataset() as shared:
+        bin_native_by_dimension(
+            n_classes, 1, bins, X, y, sample_weight, feature_names_in, feature_types_in,
+                shared,
+        )
+        assert shared.shared_memory is not None
+        assert shared.dataset is not None
+        assert shared.name is not None