Update reset_gs code to orphan old gs if not indicated for deletion (facebook#4735)

Summary:
Pull Request resolved: facebook#4735

V2:
After review we decided to allow for multiple GS to be associated with an experiment in the db and only load the most recent one.

V1:
I was looking at this method and realized that although it will work within a single instance of client, upon saving/re-loading the experiment load will fail because there will be two instances of gs associated with an experiment and this will cause a merge error.

This implements:
1. an orphan gs method -- allows the legacy gs to still live in the db and just sets expeirment column to none, essentially orphaning it, but we coud go look it up in the future if we wanted
2. cleans up some pyres taht weren't doing anything i saw along the way
3. extends tests to account for this behavior

Reviewed By: mpolson64

Differential Revision: D90064851

fbshipit-source-id: a5083dbadb4d51e26b546ec1ec690f1ea9daa965

Author: mgarrard

ax/storage/sqa_store/load.py

@@ -6,6 +6,8 @@
 
 # pyre-strict
 
+import logging
+
 from math import ceil
 from typing import Any, cast, Mapping
 
@@ -53,6 +55,8 @@
 from sqlalchemy.orm import defaultload, joinedload, lazyload, noload
 from sqlalchemy.orm.exc import DetachedInstanceError
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 # ---------------------------- Loading `Experiment`. ---------------------------
 
@@ -537,21 +541,32 @@ def _load_generation_strategy_by_id(
 def get_generation_strategy_id(experiment_name: str, decoder: Decoder) -> int | None:
     """Get DB ID of the generation strategy, associated with the experiment
     with the given name if its in DB, return None otherwise.
+
+    If multiple generation strategies are associated with the experiment,
+    returns the latest one (highest DB ID).
     """
     exp_sqa_class = decoder.config.class_to_sqa_class[Experiment]
     gs_sqa_class = decoder.config.class_to_sqa_class[GenerationStrategy]
     with session_scope() as session:
-        sqa_gs_id = (
+        sqa_gs_ids = (
             session.query(gs_sqa_class.id)  # pyre-ignore[16]
             .join(exp_sqa_class.generation_strategy)  # pyre-ignore[16]
             # pyre-fixme[16]: `SQABase` has no attribute `name`.
             .filter(exp_sqa_class.name == experiment_name)
-            .one_or_none()
+            .order_by(gs_sqa_class.id.desc())
+            .all()
         )
 
-    if sqa_gs_id is None:
+    if not sqa_gs_ids:
         return None
-    return sqa_gs_id[0]
+
+    if len(sqa_gs_ids) > 1:
+        logger.warning(
+            f"Found {len(sqa_gs_ids)} generation strategies for experiment "
+            f"{experiment_name}. Loading the latest one (id={sqa_gs_ids[0][0]})."
+        )
+
+    return sqa_gs_ids[0][0]
 
 
 def get_generation_strategy_sqa(

ax/storage/sqa_store/sqa_enum.py

@@ -16,16 +16,13 @@
 class BaseNullableEnum(types.TypeDecorator):
     cache_ok = True
 
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
     def __init__(self, enum: Any, *arg: list[Any], **kw: dict[Any, Any]) -> None:
         types.TypeDecorator.__init__(self, *arg, **kw)
         # pyre-fixme[4]: Attribute must be annotated.
         self._member_map = enum._member_map_
         # pyre-fixme[4]: Attribute must be annotated.
         self._value2member_map = enum._value2member_map_
 
-    # pyre-fixme[3]: Return annotation cannot be `Any`.
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
     def process_bind_param(self, value: Any, dialect: Any) -> Any:
         if value is None:
             return value
@@ -40,8 +37,6 @@ def process_bind_param(self, value: Any, dialect: Any) -> Any:
             )
         return val._value_
 
-    # pyre-fixme[3]: Return annotation cannot be `Any`.
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
     def process_result_value(self, value: Any, dialect: Any) -> Any:
         if value is None:
             return value

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -2910,6 +2910,34 @@ def test_delete_generation_strategy_max_gs_to_delete(self) -> None:
         # Full GS fails the equality check
         self.assertEqual(str(generation_strategy), str(loaded_generation_strategy))
 
+    def test_load_latest_generation_strategy_when_multiple_exist(self) -> None:
+        experiment = get_branin_experiment()
+        gs1 = choose_generation_strategy_legacy(experiment.search_space)
+        gs1.experiment = experiment
+        save_experiment(experiment)
+        save_generation_strategy(generation_strategy=gs1)
+        self.assertEqual(
+            gs1.db_id,
+            load_generation_strategy_by_experiment_name(experiment.name).db_id,
+        )
+
+        # create a second generation strategy for the experiment
+        gs2 = choose_generation_strategy_legacy(experiment.search_space)
+        gs2._name = "second_gs"
+        gs2.experiment = experiment
+        save_generation_strategy(generation_strategy=gs2)
+
+        # check that the latest generation stragey is loaded
+        with self.assertLogs(
+            "ax.storage.sqa_store.load", level=logging.WARNING
+        ) as logs:
+            loaded_gs = load_generation_strategy_by_experiment_name(experiment.name)
+            self.assertEqual(loaded_gs.db_id, gs2.db_id)
+            self.assertEqual(loaded_gs.name, gs2.name)
+            self.assertTrue(
+                any("Found 2 generation strategies" in log for log in logs.output)
+            )
+
     def test_query_historical_experiments_given_parameters(self) -> None:
         # This test validates the query behavior for historical experiments.
         config = SQAConfig(experiment_type_enum=TestExperimentTypeEnum)

ax/storage/sqa_store/validation.py

@@ -51,7 +51,6 @@ def wrapper(fn: Callable) -> Callable:
     return wrapper
 
 
-# pyre-fixme[3]: Return annotation cannot be `Any`.
 def consistency_exactly_one(instance: SQABase, exactly_one_fields: list[str]) -> Any:
     """Ensure that exactly one of `exactly_one_fields` has a value set."""
     values = [getattr(instance, field) is not None for field in exactly_one_fields]