Do not persist entire AutoMLState in Searcher (#870)

* Do not persist entire AutoMLState in Searcher Signed-off-by: Antoni Baum <antoni.baum@protonmail.com> * Fix tests Signed-off-by: Antoni Baum <antoni.baum@protonmail.com> Signed-off-by: Antoni Baum <antoni.baum@protonmail.com>
microsoft · Jan 6, 2023 · 5f67c0a · 5f67c0a
1 parent 90aea9c
commit 5f67c0a
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 7 deletions.
diff --git a/flaml/automl/automl.py b/flaml/automl/automl.py
@@ -491,15 +491,15 @@ def _train_with_config(
         return estimator, train_time
 
 
-def size(state: AutoMLState, config: dict) -> float:
+def size(learner_classes: dict, config: dict) -> float:
     """Size function.
 
     Returns:
         The mem size in bytes for a config.
     """
     config = config.get("ml", config)
     estimator = config["learner"]
-    learner_class = state.learner_classes.get(estimator)
+    learner_class = learner_classes.get(estimator)
     return learner_class.size(config)
 
 
@@ -3125,7 +3125,7 @@ def _search_parallel(self):
                 min_resource=min_resource_all_estimator,
                 max_resource=self.max_resource,
                 config_constraints=[
-                    (partial(size, self._state), "<=", self._mem_thres)
+                    (partial(size, self._state.learner_classes), "<=", self._mem_thres)
                 ],
                 metric_constraints=self.metric_constraints,
                 seed=self._seed,

diff --git a/test/automl/test_constraints.py b/test/automl/test_constraints.py
@@ -55,7 +55,9 @@ def test_metric_constraints():
         min_resource=automl.min_resource,
         max_resource=automl.max_resource,
         time_budget_s=automl._state.time_budget,
-        config_constraints=[(partial(size, automl._state), "<=", automl._mem_thres)],
+        config_constraints=[
+            (partial(size, automl._state.learner_classes), "<=", automl._mem_thres)
+        ],
         metric_constraints=automl.metric_constraints,
         num_samples=5,
     )
@@ -159,7 +161,9 @@ def test_metric_constraints_custom():
         min_resource=automl.min_resource,
         max_resource=automl.max_resource,
         time_budget_s=automl._state.time_budget,
-        config_constraints=[(partial(size, automl._state), "<=", automl._mem_thres)],
+        config_constraints=[
+            (partial(size, automl._state.learner_classes), "<=", automl._mem_thres)
+        ],
         metric_constraints=automl.metric_constraints,
         num_samples=5,
     )

diff --git a/test/automl/test_python_log.py b/test/automl/test_python_log.py
@@ -84,7 +84,11 @@ def test_logging_level(self):
                 min_resource=automl.min_resource,
                 max_resource=automl.max_resource,
                 config_constraints=[
-                    (partial(size, automl._state), "<=", automl._mem_thres)
+                    (
+                        partial(size, automl._state.learner_classes),
+                        "<=",
+                        automl._mem_thres,
+                    )
                 ],
                 metric_constraints=automl.metric_constraints,
             )

diff --git a/test/automl/test_xgboost2d.py b/test/automl/test_xgboost2d.py
@@ -77,7 +77,9 @@ def test_simple(method=None):
         min_resource=automl.min_resource,
         max_resource=automl.max_resource,
         time_budget_s=automl._state.time_budget,
-        config_constraints=[(partial(size, automl._state), "<=", automl._mem_thres)],
+        config_constraints=[
+            (partial(size, automl._state.learner_classes), "<=", automl._mem_thres)
+        ],
         metric_constraints=automl.metric_constraints,
         num_samples=5,
     )