Make SQA store SQLAlchemy 2.0-compatible (#5201)

Summary:
Pull Request resolved: #5201

The OSS Ax SQA store has a hard guard in with_db_settings_base.py that raises IncompatibleDependencyVersion when SQLAlchemy major > 1, disabling SQL storage entirely. This blocks SA 2.0 adoption (T163607006) and Python 3.13/3.14 (which auto-select SA 2.0.48 from third-party). Two additional SA 2.0 incompatibilities exist in OSS: defer("col_name") in load.py and reduced_state.py, which SA 2.0 rejects in favor of class-bound attribute references.

This diff removes the guard and converts the string-based loader options to attribute references. Adds a dual-version Buck test target tests_sa2 via constraint_overrides plus a self-proving TestSQLAlchemyDualVersionCompat class so each target proves its constraint took effect (EXPECTED_SA_MAJOR env var asserted at runtime).

Reviewed By: yangjoanna, andycylmeta

Differential Revision: D104875017

fbshipit-source-id: 6835c4753e251161c13eff2a452aeb50ea439001

Author: Jiawei Yang

ax/storage/sqa_store/db.py

@@ -164,6 +164,7 @@ def init_engine_and_session_factory(
 
     if SESSION_FACTORY is not None:
         if force_init:
+            # pyre-ignore[16]: SA 2.0 bind is Union; runtime Engine.
             SESSION_FACTORY.bind.dispose()
         else:
             return
@@ -204,6 +205,7 @@ def init_test_engine_and_session_factory(
 
     if SESSION_FACTORY is not None:
         if force_init:
+            # pyre-ignore[16]: SA 2.0 bind is Union; runtime Engine.
             SESSION_FACTORY.bind.dispose()
         else:
             return
@@ -262,6 +264,7 @@ def get_engine() -> Engine:
     global SESSION_FACTORY
     if SESSION_FACTORY is None:
         raise ValueError("Engine must be initialized first.")
+    # pyre-ignore[7]: SA 2.0 bind is Union; runtime Engine.
     return SESSION_FACTORY.bind
 
 
@@ -330,5 +333,6 @@ def session_context(
     finally:
         # Restore the old session factory
         session_factory.close()
+        # pyre-ignore[16]: SA 2.0 bind is Union; runtime Engine.
         session_factory.bind.dispose()
         SESSION_FACTORY = old_session

ax/storage/sqa_store/decoder.py

@@ -194,8 +194,10 @@ def _auxiliary_experiments_by_purpose_from_experiment_sqa(
                 ):
                     continue
                 aux_experiment = auxiliary_experiment_from_name(
+                    # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
                     experiment_name=auxiliary_experiment_sqa.source_experiment.name,
                     config=self.config,
+                    # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
                     is_active=auxiliary_experiment_sqa.is_active,
                     reduced_state=reduced_state,
                 )
@@ -227,6 +229,7 @@ def _init_experiment_from_sqa(
         # so need to convert it to regular dict.
         properties = dict(experiment_sqa.properties or {})
         if Keys.LLM_MESSAGES in properties:
+            # pyre-ignore[6]: SA 2.0 properties values are ColumnElement; runtime list.
             properties[Keys.LLM_MESSAGES] = [
                 LLMMessage(**m) for m in properties[Keys.LLM_MESSAGES]
             ]
@@ -249,7 +252,9 @@ def _init_experiment_from_sqa(
             raise SQADecodeError("Experiment SearchSpace cannot be None.")
         status_quo = (
             Arm(
+                # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
                 parameters=experiment_sqa.status_quo_parameters,
+                # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
                 name=experiment_sqa.status_quo_name,
             )
             if experiment_sqa.status_quo_parameters is not None
@@ -295,6 +300,7 @@ def _init_mt_experiment_from_sqa(
         """First step of conversion within experiment_from_sqa."""
         properties = dict(experiment_sqa.properties or {})
         if Keys.LLM_MESSAGES in properties:
+            # pyre-ignore[6]: SA 2.0 properties values are ColumnElement; runtime list.
             properties[Keys.LLM_MESSAGES] = [
                 LLMMessage(**m) for m in properties[Keys.LLM_MESSAGES]
             ]
@@ -317,7 +323,9 @@ def _init_mt_experiment_from_sqa(
             raise SQADecodeError("Experiment SearchSpace cannot be None.")
         status_quo = (
             Arm(
+                # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
                 parameters=experiment_sqa.status_quo_parameters,
+                # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
                 name=experiment_sqa.status_quo_name,
             )
             if experiment_sqa.status_quo_parameters is not None

ax/storage/sqa_store/json.py

@@ -93,6 +93,7 @@ class JSONEncodedLongText(JSONEncodedObject):
     impl = Text(LONGTEXT_BYTES)
 
 
+# pyre-ignore[9]: SA 2.0 typed as_mutable returns TypeEngine; runtime TypeDecorator.
 JSONEncodedList: TypeDecorator = MutableList.as_mutable(JSONEncodedObject)
 JSONEncodedDict: TypeDecorator = MutableDict.as_mutable(JSONEncodedObject)
 JSONEncodedTextDict: TypeDecorator = MutableDict.as_mutable(JSONEncodedText)

ax/storage/sqa_store/load.py

@@ -614,10 +614,12 @@ def get_generation_strategy_sqa_reduced_state(
                 gr_sqa_class.parameter_constraints
             ),
             defaultload(gs_sqa_class.generator_runs).lazyload(gr_sqa_class.metrics),
-            defaultload(gs_sqa_class.generator_runs).defer("model_kwargs"),
-            defaultload(gs_sqa_class.generator_runs).defer("bridge_kwargs"),
-            defaultload(gs_sqa_class.generator_runs).defer("model_state_after_gen"),
-            defaultload(gs_sqa_class.generator_runs).defer("gen_metadata"),
+            defaultload(gs_sqa_class.generator_runs).defer(gr_sqa_class.model_kwargs),
+            defaultload(gs_sqa_class.generator_runs).defer(gr_sqa_class.bridge_kwargs),
+            defaultload(gs_sqa_class.generator_runs).defer(
+                gr_sqa_class.model_state_after_gen
+            ),
+            defaultload(gs_sqa_class.generator_runs).defer(gr_sqa_class.gen_metadata),
         ],
     )
 

ax/storage/sqa_store/reduced_state.py

@@ -58,6 +58,5 @@ def get_query_options_to_defer_large_model_cols() -> list[strategy_options.Load]
     when loading experiment and generation strategy in reduced state.
     """
     return [
-        defaultload(SQATrial.generator_runs).defer(col.key)
-        for col in GR_LARGE_MODEL_ATTRS
+        defaultload(SQATrial.generator_runs).defer(col) for col in GR_LARGE_MODEL_ATTRS
     ]

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -240,9 +240,12 @@ def test_connection_to_db_with_creator(self) -> None:
         )
         with session_scope() as session:
             engine = session.bind
+            # pyre-ignore[16]: SA 2.0 bind is Union; runtime Engine.
             engine.connect()
             self.assertEqual(mocked_dbapi.connect.call_count, 1)
+            # pyre-ignore[16]: SA 2.0 bind is Union; runtime Engine.
             self.assertTrue(engine.echo)
+            # pyre-ignore[16]: SA 2.0 bind is Union; runtime Engine.
             self.assertEqual(engine.pool.size(), 2)
 
     def test_connection_to_db_with_session_context(self) -> None:
@@ -279,10 +282,12 @@ def test_generator_run_type_validation(self) -> None:
 
         generator_run._generator_run_type = "STATUS_QUO"
         generator_run_sqa = self.encoder.generator_run_to_sqa(generator_run)
+        # pyre-ignore[8]: SA 2.0 Column[T] attr; runtime assign is fine.
         generator_run_sqa.generator_run_type = 2
         with self.assertRaises(SQADecodeError):
             self.decoder.generator_run_from_sqa(generator_run_sqa, False, False)
 
+        # pyre-ignore[8]: SA 2.0 Column[T] attr; runtime assign is fine.
         generator_run_sqa.generator_run_type = 0
         self.decoder.generator_run_from_sqa(generator_run_sqa, False, False)
 

ax/storage/sqa_store/tests/test_with_db_settings_base.py

@@ -6,15 +6,18 @@
 
 # pyre-strict
 import logging
+import os
 import random
 import string
 from unittest.mock import patch
 
+import sqlalchemy
 from ax.core.base_trial import BaseTrial
 from ax.core.experiment import Experiment
 from ax.core.trial import Trial
 from ax.core.trial_status import TrialStatus
 from ax.generation_strategy.generation_strategy import GenerationStrategy
+from ax.storage.sqa_store import with_db_settings_base as _wdb_module
 from ax.storage.sqa_store.db import init_test_engine_and_session_factory
 from ax.storage.sqa_store.load import (
     _load_experiment,
@@ -396,3 +399,47 @@ def test_try_load_generation_strategy(self) -> None:
                 lg.output[0],
             )
         self.assertEqual(output, generation_strategy)
+
+
+class TestSQLAlchemyDualVersionCompat(TestCase):
+    """Self-proving checks that the dual-version SA 2.0 BUCK targets actually
+    resolved their constraint_overrides and that the SA 2.0 hard guard is gone.
+
+    Part of the SA 2.0 dual-version migration (T163607006).
+    """
+
+    def test_module_level_dbsettings_is_defined(self) -> None:
+        """The SA 2.0 hard guard previously set DBSettings = None at module level
+        when SA major > 1, breaking WithDBSettingsBase.__init__ type checks. Now
+        that the guard is removed, DBSettings must always resolve to the real
+        type when SQLAlchemy is importable. Uses getattr because DBSettings is
+        conditionally defined in a try/except in with_db_settings_base.
+        """
+        # pyre-ignore[16]: DBSettings is conditionally defined in with_db_settings_base.
+        module_dbsettings = getattr(_wdb_module, "DBSettings", None)
+        self.assertIsNotNone(
+            module_dbsettings,
+            "with_db_settings_base.DBSettings is None -- guard removal regressed",
+        )
+        self.assertIs(module_dbsettings, DBSettings)
+
+    def test_sa_major_matches_buck_target(self) -> None:
+        """When the BUCK target sets EXPECTED_SA_MAJOR, assert the runtime
+        SQLAlchemy major matches. Makes :tests vs :tests_sa2 self-proving.
+        Skipped when EXPECTED_SA_MAJOR is unset (e.g., local one-off invocations).
+        """
+        expected_major_str = os.environ.get("EXPECTED_SA_MAJOR")
+        if expected_major_str is None:
+            self.skipTest(
+                "EXPECTED_SA_MAJOR not set; only enforced under the dual-version "
+                "BUCK targets that pin SQLAlchemy via constraint_overrides"
+            )
+        # pyre-ignore[16]: Module `sqlalchemy` has no attribute `__version__`.
+        actual_version = sqlalchemy.__version__
+        actual_major = int(actual_version.split(".")[0])
+        self.assertEqual(
+            actual_major,
+            int(expected_major_str),
+            f"BUCK target expected SQLAlchemy major {expected_major_str}, "
+            f"got {actual_version}",
+        )

ax/storage/sqa_store/with_db_settings_base.py

@@ -6,7 +6,6 @@
 
 # pyre-strict
 
-import re
 import time
 from collections.abc import Sequence
 from logging import INFO, Logger
@@ -16,11 +15,7 @@
 from ax.core.experiment import Experiment
 from ax.core.generator_run import GeneratorRun
 from ax.core.runner import Runner
-from ax.exceptions.core import (
-    IncompatibleDependencyVersion,
-    ObjectNotFoundError,
-    UnsupportedError,
-)
+from ax.exceptions.core import ObjectNotFoundError, UnsupportedError
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.utils.common.executils import retry_on_exception
 from ax.utils.common.logger import _round_floats_for_logging, get_logger
@@ -37,18 +32,7 @@
     from sqlalchemy import __version__ as sqa_version
 
     # pyre-fixme[16]: Module `sqlalchemy` has no attribute `__version__`.
-    sqa_major_version = int(none_throws(re.match(r"^\d*", sqa_version))[0])
-    if sqa_major_version > 1:
-        msg = (
-            "Ax currently requires a sqlalchemy version below 2.0. This will be "
-            "addressed in a future release. Disabling SQL storage in Ax for now, if "
-            "you would like to use SQL storage please install Ax with mysql extras "
-            "via `pip install ax-platform[mysql]`."
-        )
-
-        logger.warning(msg)
-
-        raise IncompatibleDependencyVersion(msg)
+    logger.info(f"Ax SQL storage initialized with SQLAlchemy {sqa_version}")
 
     from ax.storage.sqa_store.db import init_engine_and_session_factory
     from ax.storage.sqa_store.decoder import Decoder
@@ -78,7 +62,7 @@
 
     # We retry on `OperationalError` if saving to DB.
     RETRY_EXCEPTION_TYPES = (OperationalError, StaleDataError)
-except (ModuleNotFoundError, IncompatibleDependencyVersion, TypeError):
+except (ModuleNotFoundError, TypeError):
     DBSettings = None
     TDBSettings = None
     Decoder = None