Remove residual SA 2.0 pyre-ignores after Mapped[T] migration (#5206)

Summary:
Pull Request resolved: #5206

Builds on the Mapped[T] migration in the parent diff (D105247335) to remove inline `# pyre-ignore[6/7/8]: SA 2.0 Column[T] ...` comments that the Mapped[T] migration made redundant, plus remove the broad file-level `# pyre-ignore-all-errors[6, 8, 9]` directives in `decoder.py`, `save.py`, `load.py`, `delete.py` that were no longer needed once Mapped[T] resolved the underlying type mismatches.

Concrete changes:
- `fbcode/ax/storage/sqa_store/decoder.py`: removed file-level `# pyre-ignore-all-errors[6, 8, 9]` and 13 inline `# pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.` comments. With `Mapped[T]` annotations on the SQA columns, pyre now resolves `experiment_sqa.name`/etc. to plain `T` and the kwargs flow through `Experiment(...)`/`MultiTypeExperiment(...)` without complaint.
- `fbcode/ax/storage/sqa_store/save.py`, `load.py`, `delete.py`: removed file-level `# pyre-ignore-all-errors[...]` directives. Pyre clean after removal.
- `fbcode/ax/storage/sqa_store/load.py`: also fixed a type bug at the cast on line ~717: was `cast(SQAExperiment, ...)` (instance type) → `cast(type[SQAExperiment], ...)` (class type), matching the convention used elsewhere in the file.
- `fbcode/ax/storage/sqa_store/db.py`: added `_bound_engine(scoped_session) -> Engine` helper that narrows SA 2.0's `Union[None, Connection, Engine]` typing of `scoped_session.bind` via an `isinstance` check. Replaced 4 inline `# pyre-ignore[16]` / `# pyre-ignore[7]` sites with calls to the helper.
- `fbcode/ax/fb/storage/sqa_store/decoder.py`, `encoder.py`, `load_helper.py`: removed inline `# pyre-ignore[6/7/8]: SA 2.0 Column[T] ...` comments. Same Mapped[T]-driven cleanup.

NOT done (would have broken SA 1.4 consumers — see admarket/ranking/automl/gain/searcher pinning SA 1.4):
- `fbcode/ax/storage/sqa_store/json.py`: kept `JSONEncodedList: TypeDecorator = ...` rather than `TypeEngine[Any]`. SA 1.4's `TypeEngine` is not a Generic class and `TypeEngine[Any]` raises `type 'TypeEngine' is not subscriptable` at runtime. Inline `# pyre-ignore[9]` retained on each of the 4 alias lines instead.
- `fbcode/ax/storage/sqa_store/reduced_state.py`, `validation.py`: kept bare `InstrumentedAttribute` rather than `InstrumentedAttribute[Any]`. Same SA 1.4 runtime constraint — `InstrumentedAttribute` is not subscriptable in 1.4. File-level `# pyre-ignore-all-errors[24]` on `reduced_state.py` and inline `# pyre-ignore[24]` on `validation.py`.

Kept intentionally:
- FB `metadata.py` `hybrid_property.expression` inline pyre-ignores on `Column("...", String(...))` return values — hybrid SQL expressions don't migrate to Mapped[T].
- FB `sqa_classes.py` `# pyre-ignore[8]` on `association_proxy(...)` lines — association_proxy returns its own descriptor, not a Mapped attr.
- Pre-existing `# pyre-fixme[*]` comments unrelated to the SA 1.4 → 2.0 bump.

Net suppression count delta (this diff vs the prior D104875016 era): -4 file-level `# pyre-ignore-all-errors[...]` directives in `decoder.py`/`save.py`/`load.py`/`delete.py`, -18 inline `# pyre-ignore[6/7/8]` comments across `decoder.py` (OSS+FB), `encoder.py` (FB), `load_helper.py` (FB), and `db.py`. Total: -22 suppressions. The remaining 3 file-level `# pyre-ignore-all-errors[8/24]` in OSS `sqa_classes.py`, FB `metadata.py`, and `reduced_state.py`, plus a handful of inline ignores, are intentionally kept to preserve dual-version (SA 1.4 + SA 2.0) compatibility for OSS Ax consumers.

Reviewed By: yangjoanna

Differential Revision: D106016101

fbshipit-source-id: 7c49964654cd561ac952d72779c1700fdc332722

Author: Jiawei Yang

ax/storage/sqa_store/db.py

@@ -42,6 +42,22 @@
 T = TypeVar("T")
 
 
+def _bound_engine(session_factory: scoped_session) -> Engine:
+    """Return the ``Engine`` bound to ``session_factory``, raising if not bound.
+
+    SA 2.0 types ``scoped_session.bind`` as ``Connection | Engine | None``. In
+    this module we always bind a freshly-created ``Engine`` via ``sessionmaker``,
+    so the runtime value is always an ``Engine``. This narrows the type for both
+    pyre and runtime safety.
+    """
+    bind = session_factory.bind
+    if not isinstance(bind, Engine):
+        raise ValueError(
+            f"SESSION_FACTORY must be bound to an Engine, got {type(bind).__name__}."
+        )
+    return bind
+
+
 class SQABase:
     """Metaclass for SQLAlchemy classes corresponding to core Ax classes."""
 
@@ -164,8 +180,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()
+            _bound_engine(SESSION_FACTORY).dispose()
         else:
             return
     if url is not None:
@@ -205,8 +220,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()
+            _bound_engine(SESSION_FACTORY).dispose()
         else:
             return
     engine = create_test_engine(path=tier_or_path, echo=echo)
@@ -264,8 +278,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
+    return _bound_engine(SESSION_FACTORY)
 
 
 @contextmanager
@@ -333,6 +346,5 @@ 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()
+        _bound_engine(session_factory).dispose()
         SESSION_FACTORY = old_session

ax/storage/sqa_store/decoder.py

@@ -5,7 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-# pyre-ignore-all-errors[6, 8, 9]
 
 import re
 import warnings
@@ -195,10 +194,8 @@ 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,
                 )
@@ -253,9 +250,7 @@ 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
@@ -282,16 +277,13 @@ def _init_experiment_from_sqa(
         )
 
         return Experiment(
-            # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
             name=experiment_sqa.name,
-            # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
             description=experiment_sqa.description,
             search_space=search_space,
             optimization_config=opt_config,
             tracking_metrics=all_metrics,
             runner=runner,
             status_quo=status_quo,
-            # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
             is_test=experiment_sqa.is_test,
             properties=properties,
             auxiliary_experiments_by_purpose=auxiliary_experiments_by_purpose,
@@ -327,9 +319,7 @@ 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
@@ -357,14 +347,10 @@ def _init_mt_experiment_from_sqa(
             trial_type_to_runner.update(dict.fromkeys(trial_types_with_metrics))
 
         experiment = MultiTypeExperiment(
-            # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
             name=experiment_sqa.name,
-            # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
             description=experiment_sqa.description,
             search_space=search_space,
-            # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
             default_trial_type=default_trial_type,
-            # pyre-ignore[6]: SA 2.0 Column[T] vs plain T param.
             default_runner=trial_type_to_runner.get(default_trial_type),
             optimization_config=opt_config,
             status_quo=status_quo,

ax/storage/sqa_store/delete.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-# pyre-ignore-all-errors[6]
 
 from logging import Logger
 from typing import cast

ax/storage/sqa_store/json.py

@@ -93,6 +93,10 @@ class JSONEncodedLongText(JSONEncodedObject):
     impl = Text(LONGTEXT_BYTES)
 
 
+# `Mutable*.as_mutable()` returns a `TypeEngine` subclass per SA 2.0 stubs.
+# Cannot annotate as `TypeEngine[Any]` because SA 1.4's `TypeEngine` is not a
+# Generic class (`type 'TypeEngine' is not subscriptable` at runtime under 1.4).
+# Keep `TypeDecorator` and suppress the SA 2.0 type-stub mismatch on each line.
 # pyre-ignore[9]: SA 2.0 typed as_mutable returns TypeEngine; runtime TypeDecorator.
 JSONEncodedList: TypeDecorator = MutableList.as_mutable(JSONEncodedObject)
 # pyre-ignore[9]: SA 2.0 typed as_mutable returns TypeEngine; runtime TypeDecorator.

ax/storage/sqa_store/load.py

@@ -5,7 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-# pyre-ignore-all-errors[6, 8]
 
 import logging
 from collections.abc import Mapping
@@ -715,8 +714,8 @@ def load_analysis_cards_by_experiment_name(
             analysis_card_sqa_class.children
         )
 
-    exp_sqa_class: SQAExperiment = cast(
-        SQAExperiment, decoder.config.class_to_sqa_class[Experiment]
+    exp_sqa_class = cast(
+        type[SQAExperiment], decoder.config.class_to_sqa_class[Experiment]
     )
 
     with session_scope() as session:

ax/storage/sqa_store/reduced_state.py

@@ -6,16 +6,17 @@
 
 # pyre-strict
 # pyre-ignore-all-errors[24]
+#
+# SA 2.0 requires a type param on InstrumentedAttribute, but SA 1.4
+# InstrumentedAttribute is not subscriptable at runtime (`type is not a generic
+# class`), so we keep the bare form to preserve dual-version compatibility.
 
 
 from ax.storage.sqa_store.sqa_classes import SQAGeneratorRun, SQATrial
 from sqlalchemy.orm import defaultload, strategy_options
 from sqlalchemy.orm.attributes import InstrumentedAttribute
 
 
-# pyre-fixme[9]: `GR_LARGE_MODEL_ATTRS` is declared as `List[InstrumentedAttribute]`
-#  but SQLAlchemy class attributes are typed as `Column` in stubs; they are
-#  `InstrumentedAttribute` instances at runtime.
 GR_LARGE_MODEL_ATTRS: list[InstrumentedAttribute] = [
     SQAGeneratorRun.model_kwargs,
     SQAGeneratorRun.bridge_kwargs,

ax/storage/sqa_store/save.py

@@ -5,7 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-# pyre-ignore-all-errors[6, 8, 9]
 
 import os
 from collections.abc import Callable, Generator, Sequence

ax/storage/sqa_store/validation.py

@@ -34,7 +34,8 @@
 
 
 def listens_for_multiple(
-    # pyre-ignore[24]: SA 2.0 requires a type param on InstrumentedAttribute.
+    # pyre-ignore[24]: SA 2.0 requires a type param; SA 1.4 InstrumentedAttribute
+    # is not subscriptable at runtime, so we keep the bare form.
     targets: list[InstrumentedAttribute],
     identifier: str,
     *args: Any,