GridSearchOracle: persist _ordered_ids and _populate_next across reloads (#1055)

keras_tuner/tuners/gridsearch.py

@@ -322,6 +322,36 @@ def end_trial(self, trial):
         # For not blocking _populate_space, we push it regardless of the status.
         self._populate_next.append(trial.trial_id)
 
+    def get_state(self):
+        # Persist the GridSearch-specific bookkeeping in addition to the base
+        # Oracle state so that a fresh process resuming a search (e.g.
+        # `GridSearch(..., overwrite=False)` after a kernel restart) keeps
+        # working without `KeyError` from `_ordered_ids.next()` (#1055).
+        state = super().get_state()
+        state["ordered_ids"] = list(self._ordered_ids._memory)
+        state["populate_next"] = list(self._populate_next)
+        return state
+
+    def set_state(self, state):
+        super().set_state(state)
+        # Older state files (saved by versions <= 1.4.8) did not include the
+        # new keys. Lazily rebuild `_ordered_ids` from `start_order` so we
+        # remain compatible with checkpoints written before this fix.
+        self._ordered_ids = LinkedList()
+        ordered_ids = state.get("ordered_ids")
+        if ordered_ids is None:
+            ordered_ids = list(self.start_order)
+        prev_id = None
+        for trial_id in ordered_ids:
+            self._ordered_ids.insert(trial_id, prev_id)
+            prev_id = trial_id
+        populate_next = state.get("populate_next")
+        if populate_next is None:
+            # Match the workaround in #1055: seed with the most recent
+            # completed trial so end_trial-driven progression resumes.
+            populate_next = [self.end_order[-1]] if self.end_order else []
+        self._populate_next = list(populate_next)
+
 
 @keras_tuner_export(["keras_tuner.GridSearch", "keras_tuner.tuners.GridSearch"])
 class GridSearch(tuner_module.Tuner):

keras_tuner/tuners/gridsearch_test.py

@@ -225,3 +225,88 @@ def test_linked_list():
     assert linked_list.next("1") == "3"
     assert linked_list.next("3") == "4"
     assert linked_list.next("4") is None
+
+
+def test_grid_search_oracle_state_round_trip_resumes_search(tmp_path):
+    """Regression test for #1055: GridSearchOracle.get_state / set_state must
+    persist `_ordered_ids` and `_populate_next` so a fresh-process resume
+    (``GridSearch(..., overwrite=False)`` after a kernel restart) doesn't
+    KeyError on the first completed trial."""
+    from keras_tuner.engine import hyperparameters as hp_module
+    from keras_tuner.tuners.gridsearch import GridSearchOracle
+
+    hps = hp_module.HyperParameters()
+    hps.Boolean("b1")
+    hps.Boolean("b2")
+
+    def make_oracle():
+        return GridSearchOracle(
+            objective="val_loss",
+            max_trials=10,
+            hyperparameters=hps,
+        )
+
+    # Drive the search through a couple of completed trials so the LinkedList
+    # and the populate-next queue actually have non-empty state.
+    oracle = make_oracle()
+    trial_1 = oracle.create_trial(tuner_id="1")
+    trial_2 = oracle.create_trial(tuner_id="2")
+    trial_1.status = trial_module.TrialStatus.COMPLETED
+    trial_2.status = trial_module.TrialStatus.COMPLETED
+    oracle.end_trial(trial_1)
+    oracle.end_trial(trial_2)
+    assert len(oracle._ordered_ids._memory) >= 2
+
+    # Round-trip through get_state / set_state on a brand-new oracle (mimics
+    # process restart + reload-from-disk).
+    state = oracle.get_state()
+    fresh_oracle = make_oracle()
+    fresh_oracle.trials = oracle.trials
+    fresh_oracle.set_state(state)
+
+    # Resumed oracle's LinkedList must contain the same trial ids.
+    assert fresh_oracle._ordered_ids._memory == oracle._ordered_ids._memory
+    assert fresh_oracle._populate_next == oracle._populate_next
+
+    # And the next create_trial must NOT raise KeyError from _ordered_ids.next.
+    trial_3 = fresh_oracle.create_trial(tuner_id="3")
+    assert trial_3.status in (
+        trial_module.TrialStatus.RUNNING,
+        trial_module.TrialStatus.IDLE,
+        trial_module.TrialStatus.STOPPED,
+    )
+
+
+def test_grid_search_oracle_set_state_recovers_from_legacy_state(tmp_path):
+    """A state dict written by an older keras-tuner that did not persist the
+    GridSearch bookkeeping must still rehydrate without KeyError — the new
+    set_state lazily rebuilds `_ordered_ids` from `start_order`."""
+    from keras_tuner.engine import hyperparameters as hp_module
+    from keras_tuner.tuners.gridsearch import GridSearchOracle
+
+    hps = hp_module.HyperParameters()
+    hps.Boolean("b1")
+
+    oracle = GridSearchOracle(
+        objective="val_loss",
+        max_trials=5,
+        hyperparameters=hps,
+    )
+    trial_1 = oracle.create_trial(tuner_id="1")
+    trial_1.status = trial_module.TrialStatus.COMPLETED
+    oracle.end_trial(trial_1)
+
+    # Build a "legacy" state by dropping the new keys, then round-trip.
+    state = oracle.get_state()
+    state.pop("ordered_ids", None)
+    state.pop("populate_next", None)
+
+    fresh_oracle = GridSearchOracle(
+        objective="val_loss",
+        max_trials=5,
+        hyperparameters=hps,
+    )
+    fresh_oracle.trials = oracle.trials
+    fresh_oracle.set_state(state)
+    # Rebuilt from start_order rather than directly from the missing key.
+    assert fresh_oracle._ordered_ids._memory == list(oracle.start_order)