fix(algos): Fixes infinite loop in mutate_parents when mutations are …

archai/discrete_search/algos/bananas.py

@@ -167,6 +167,7 @@ def mutate_parents(
             nb_tries = 0
 
             while len(candidates) < mutations_per_parent and nb_tries < patience:
+                nb_tries += 1
                 mutated_model = self.search_space.mutate(p)
                 mutated_model.metadata["parent"] = p.archid
 
@@ -176,8 +177,6 @@ def mutate_parents(
                 if mutated_model.archid not in self.seen_archs:
                     candidates[mutated_model.archid] = mutated_model
 
-                nb_tries += 1
-
             mutations.update(candidates)
 
         if len(mutations) == 0:

archai/discrete_search/algos/local_search.py

@@ -122,6 +122,7 @@ def mutate_parents(
             nb_tries = 0
 
             while len(candidates) < mutations_per_parent and nb_tries < patience:
+                nb_tries += 1
                 mutated_model = self.search_space.mutate(p)
                 mutated_model.metadata["parent"] = p.archid
 
@@ -131,7 +132,6 @@ def mutate_parents(
                 if mutated_model.archid not in self.seen_archs:
                     mutated_model.metadata["generation"] = self.iter_num
                     candidates[mutated_model.archid] = mutated_model
-                nb_tries += 1
             mutations.update(candidates)
 
         return list(mutations.values())

archai/discrete_search/algos/regularized_evolution.py

@@ -135,6 +135,7 @@ def mutate_parents(
             nb_tries = 0
 
             while len(candidates) < mutations_per_parent and nb_tries < patience:
+                nb_tries += 1
                 mutated_model = self.search_space.mutate(p)
                 mutated_model.metadata["parent"] = p.archid
 
@@ -144,7 +145,6 @@ def mutate_parents(
                 if mutated_model.archid not in self.seen_archs:
                     mutated_model.metadata["generation"] = self.iter_num
                     candidates[mutated_model.archid] = mutated_model
-                nb_tries += 1
             mutations.update(candidates)
 
         return list(mutations.values())

tests/discrete_search/algos/test_bananas.py

@@ -37,3 +37,32 @@ def test_bananas(output_dir, search_space, search_objectives, surrogate_model):
     # Checks if all registered models satisfy constraints
     _, valid_models = search_objectives.validate_constraints(all_models)
     assert len(valid_models) == len(all_models)
+
+
+def test_bananas_mutate_parents_stops_on_invalid_mutations(
+    output_dir, search_space, search_objectives, surrogate_model, monkeypatch
+):
+    algo = MoBananasSearch(
+        search_space=search_space,
+        search_objectives=search_objectives,
+        output_dir=output_dir,
+        surrogate_model=surrogate_model,
+        num_iters=2,
+        init_num_models=5,
+    )
+
+    mutate_calls = []
+
+    def counting_mutate(model):
+        mutate_calls.append(model.archid)
+        # Fails instead of hanging if `patience` is not honored
+        assert len(mutate_calls) <= 3
+        return mutate(model)
+
+    mutate = search_space.mutate
+    monkeypatch.setattr(search_space, "mutate", counting_mutate)
+    monkeypatch.setattr(algo.so, "is_model_valid", lambda model: False)
+
+    parent = search_space.random_sample()
+    assert algo.mutate_parents([parent], mutations_per_parent=1, patience=3) == []
+    assert len(mutate_calls) == 3

tests/discrete_search/algos/test_local_search.py

@@ -35,3 +35,32 @@ def test_local_search(output_dir, search_space, search_objectives):
     # Checks if all registered models satisfy constraints
     _, valid_models = search_objectives.validate_constraints(all_models)
     assert len(valid_models) == len(all_models)
+
+
+def test_local_search_mutate_parents_stops_on_invalid_mutations(
+    output_dir, search_space, search_objectives, monkeypatch
+):
+    algo = LocalSearch(
+        search_space,
+        search_objectives,
+        output_dir,
+        num_iters=2,
+        init_num_models=5,
+        seed=1,
+    )
+
+    mutate_calls = []
+
+    def counting_mutate(model):
+        mutate_calls.append(model.archid)
+        # Fails instead of hanging if `patience` is not honored
+        assert len(mutate_calls) <= 3
+        return mutate(model)
+
+    mutate = search_space.mutate
+    monkeypatch.setattr(search_space, "mutate", counting_mutate)
+    monkeypatch.setattr(algo.so, "is_model_valid", lambda model: False)
+
+    parent = search_space.random_sample()
+    assert algo.mutate_parents([parent], mutations_per_parent=1, patience=3) == []
+    assert len(mutate_calls) == 3

tests/discrete_search/algos/test_regularized_evolution.py

@@ -38,3 +38,34 @@ def test_regularized_evolution(output_dir, search_space, search_objectives, surr
     # Checks if all registered models satisfy constraints
     _, valid_models = search_objectives.validate_constraints(all_models)
     assert len(valid_models) == len(all_models)
+
+
+def test_regularized_evolution_mutate_parents_stops_on_invalid_mutations(
+    output_dir, search_space, search_objectives, monkeypatch
+):
+    algo = RegularizedEvolutionSearch(
+        search_space=search_space,
+        search_objectives=search_objectives,
+        output_dir=output_dir,
+        num_iters=2,
+        init_num_models=4,
+        pareto_sample_size=4,
+        history_size=10,
+        seed=1,
+    )
+
+    mutate_calls = []
+
+    def counting_mutate(model):
+        mutate_calls.append(model.archid)
+        # Fails instead of hanging if `patience` is not honored
+        assert len(mutate_calls) <= 3
+        return mutate(model)
+
+    mutate = search_space.mutate
+    monkeypatch.setattr(search_space, "mutate", counting_mutate)
+    monkeypatch.setattr(algo.so, "is_model_valid", lambda model: False)
+
+    parent = search_space.random_sample()
+    assert algo.mutate_parents([parent], mutations_per_parent=1, patience=3) == []
+    assert len(mutate_calls) == 3