Parallel processing

miplearn/components/component.py

@@ -11,6 +11,8 @@ from miplearn.features import Sample
 from miplearn.instance.base import Instance
 from miplearn.types import LearningSolveStats
 
+from p_tqdm import p_umap
+
 if TYPE_CHECKING:
     from miplearn.solvers.learning import LearningSolver
 
@@ -159,7 +161,6 @@ class Component(EnforceOverrides):
         self,
         instance: Optional[Instance],
         sample: Sample,
-        pre: Optional[List[Any]] = None,
     ) -> Tuple[Dict, Dict]:
         """
         Returns a pair of x and y dictionaries containing, respectively, the matrices
@@ -168,6 +169,9 @@ class Component(EnforceOverrides):
         """
         pass
 
+    def pre_fit(self, pre: List[Any]):
+        pass
+
     def user_cut_cb(
         self,
         solver: "LearningSolver",
@@ -183,6 +187,7 @@ class Component(EnforceOverrides):
     def fit_multiple(
         components: Dict[str, "Component"],
         instances: List[Instance],
+        n_jobs: int = 1,
     ) -> None:
         def _pre_sample_xy(instance: Instance) -> Dict:
             pre_instance: Dict = {}
@@ -195,7 +200,17 @@ class Component(EnforceOverrides):
             instance.free()
             return pre_instance
 
-        def _sample_xy(instance: Instance, pre: Dict) -> Tuple[Dict, Dict]:
+        pre = p_umap(_pre_sample_xy, instances, num_cpus=n_jobs)
+        pre_combined: Dict = {}
+        for (cname, comp) in components.items():
+            pre_combined[cname] = []
+            for p in pre:
+                pre_combined[cname].extend(p[cname])
+
+        for (cname, comp) in components.items():
+            comp.pre_fit(pre_combined[cname])
+
+        def _sample_xy(instance: Instance) -> Tuple[Dict, Dict]:
             x_instance: Dict = {}
             y_instance: Dict = {}
             for (cname, comp) in components.items():
@@ -206,7 +221,7 @@ class Component(EnforceOverrides):
                 for (cname, comp) in components.items():
                     x = x_instance[cname]
                     y = y_instance[cname]
-                    x_sample, y_sample = comp.sample_xy(instance, sample, pre[cname])
+                    x_sample, y_sample = comp.sample_xy(instance, sample)
                     for cat in x_sample.keys():
                         if cat not in x:
                             x[cat] = []
@@ -216,15 +231,7 @@ class Component(EnforceOverrides):
             instance.free()
             return x_instance, y_instance
 
-        pre = [_pre_sample_xy(instance) for instance in instances]
+        xy = p_umap(_sample_xy, instances)
-
-        pre_combined: Dict = {}
-        for (cname, comp) in components.items():
-            pre_combined[cname] = []
-            for p in pre:
-                pre_combined[cname].extend(p[cname])
-
-        xy = [_sample_xy(instances, pre_combined) for instances in instances]
 
         for (cname, comp) in components.items():
             x_comp: Dict = {}

miplearn/components/dynamic_common.py

@@ -89,16 +89,18 @@ class DynamicConstraintsComponent(Component):
         self,
         instance: Optional[Instance],
         sample: Sample,
-        pre: Optional[List[Any]] = None,
     ) -> Tuple[Dict, Dict]:
+        x, y, _ = self.sample_xy_with_cids(instance, sample)
+        return x, y
+
+    @overrides
+    def pre_fit(self, pre: List[Any]) -> None:
         assert pre is not None
         known_cids: Set = set()
         for cids in pre:
             known_cids |= cids
         self.known_cids.clear()
         self.known_cids.extend(sorted(known_cids))
-        x, y, _ = self.sample_xy_with_cids(instance, sample)
-        return x, y
 
     def sample_predict(
         self,

miplearn/components/dynamic_lazy.py

@@ -108,9 +108,12 @@ class DynamicLazyConstraintsComponent(Component):
         self,
         instance: Optional[Instance],
         sample: Sample,
-        pre: Optional[List[Any]] = None,
     ) -> Tuple[Dict, Dict]:
-        return self.dynamic.sample_xy(instance, sample, pre=pre)
+        return self.dynamic.sample_xy(instance, sample)
+
+    @overrides
+    def pre_fit(self, pre: List[Any]) -> None:
+        self.dynamic.pre_fit(pre)
 
     def sample_predict(
         self,

miplearn/components/dynamic_user_cuts.py

@@ -101,9 +101,12 @@ class UserCutsComponent(Component):
         self,
         instance: "Instance",
         sample: Sample,
-        pre: Optional[List[Any]] = None,
     ) -> Tuple[Dict, Dict]:
-        return self.dynamic.sample_xy(instance, sample, pre=pre)
+        return self.dynamic.sample_xy(instance, sample)
+
+    @overrides
+    def pre_fit(self, pre: List[Any]) -> None:
+        self.dynamic.pre_fit(pre)
 
     def sample_predict(
         self,

miplearn/components/objective.py

@@ -76,7 +76,6 @@ class ObjectiveValueComponent(Component):
         self,
         _: Optional[Instance],
         sample: Sample,
-        pre: Optional[List[Any]] = None,
     ) -> Tuple[Dict[Hashable, List[List[float]]], Dict[Hashable, List[List[float]]]]:
         # Instance features
         assert sample.after_load is not None

miplearn/components/primal.py

@@ -145,7 +145,6 @@ class PrimalSolutionComponent(Component):
         self,
         _: Optional[Instance],
         sample: Sample,
-        pre: Optional[List[Any]] = None,
     ) -> Tuple[Dict[Category, List[List[float]]], Dict[Category, List[List[float]]]]:
         x: Dict = {}
         y: Dict = {}

miplearn/components/static_lazy.py

@@ -154,7 +154,6 @@ class StaticLazyConstraintsComponent(Component):
         self,
         _: Optional[Instance],
         sample: Sample,
-        pre: Optional[List[Any]] = None,
     ) -> Tuple[Dict[Hashable, List[List[float]]], Dict[Hashable, List[List[float]]]]:
         x, y, __ = self._sample_xy_with_cids(sample)
         return x, y

miplearn/solvers/learning.py

@@ -391,11 +391,19 @@ class LearningSolver:
             self._restore_miplearn_logger()
             return stats
 
-    def fit(self, training_instances: List[Instance]) -> None:
+    def fit(
+        self,
+        training_instances: List[Instance],
+        n_jobs: int = 1,
+    ) -> None:
         if len(training_instances) == 0:
             logger.warning("Empty list of training instances provided. Skipping.")
             return
-        Component.fit_multiple(self.components, training_instances)
+        Component.fit_multiple(
+            self.components,
+            training_instances,
+            n_jobs=n_jobs,
+        )
 
     def _add_component(self, component: Component) -> None:
         name = component.__class__.__name__

tests/components/test_dynamic_lazy.py

@@ -87,6 +87,7 @@ def training_instances() -> List[Instance]:
 
 def test_sample_xy(training_instances: List[Instance]) -> None:
     comp = DynamicLazyConstraintsComponent()
+    comp.pre_fit([{"c1", "c2", "c3", "c4"}])
     x_expected = {
         "type-a": [[5.0, 1.0, 2.0, 3.0], [5.0, 4.0, 5.0, 6.0]],
         "type-b": [[5.0, 1.0, 2.0], [5.0, 3.0, 4.0]],
@@ -98,7 +99,6 @@ def test_sample_xy(training_instances: List[Instance]) -> None:
     x_actual, y_actual = comp.sample_xy(
         training_instances[0],
         training_instances[0].samples[0],
-        pre=[{"c1", "c2", "c3", "c4"}],
     )
     assert_equals(x_actual, x_expected)
     assert_equals(y_actual, y_expected)