Fix bugs in JSON and SQA storage modules (#4893)

Summary:
Pull Request resolved: #4893

Fix multiple bugs in ax/storage:

- `int(None)` crash in SQA encoder: `metric_registry.get()` can return None, but the result was immediately passed to `int()` before the None check could run (sqa_store/encoder.py:393)
- `AttributeError` on None kwargs in JSON decoder: `kwargs.pop()` was called before the `if kwargs is not None` guard (json_store/decoder.py:1092)
- `decode_args_list` length mismatch in SQA save: used `len(trials)` instead of `len(trials_to_reduce_state)` and `[...] * len(trials)` instead of a single-element list (sqa_store/save.py:245,254)
- SQLAlchemy filter no-op: `SQAExperiment.id is not None` uses Python identity check instead of SQLAlchemy's `.isnot(None)` (sqa_store/load.py:760)
- Inconsistent threshold in error message: says "> 10" but the actual check is `> 15` (sqa_store/utils.py:77)
- Garbled error message in with_db_settings_base.py: sentence fragments in wrong order (sqa_store/with_db_settings_base.py:106)
- Doubled path `ax/ax/storage` in error message (sqa_store/encoder.py:131)
- Backslash line-continuation inside string literal causes 16 extra spaces in error message (sqa_store/encoder.py:654)
- Missing space between concatenated strings (sqa_store/db.py)
- Typo "SCalarized" → "Scalarized" (sqa_store/encoder.py)
- Missing space in SKIP_ATTRS_ERROR_SUFFIX (sqa_store/utils.py)

Reviewed By: bernardbeckerman

Differential Revision: D92879499

Author: saitcakmak

ax/storage/json_store/decoder.py

@@ -1098,10 +1098,10 @@ def generator_spec_from_json(
         kwargs = generator_spec_json.pop("model_kwargs", None)
     else:
         kwargs = generator_spec_json.pop("generator_kwargs", None)
-    for k in _DEPRECATED_GENERATOR_KWARGS:
-        # Remove deprecated model kwargs.
-        kwargs.pop(k, None)
     if kwargs is not None:
+        for k in _DEPRECATED_GENERATOR_KWARGS:
+            # Remove deprecated model kwargs.
+            kwargs.pop(k, None)
         kwargs = _sanitize_surrogate_spec_input(object_json=kwargs)
     if "model_gen_kwargs" in generator_spec_json:
         gen_kwargs = generator_spec_json.pop("model_gen_kwargs", None)

ax/storage/sqa_store/db.py

@@ -218,8 +218,8 @@ def create_all_tables(engine: Engine) -> None:
     """
     if engine.dialect.name == "mysql" and engine.dialect.default_schema_name == "ax":
         raise ValueError(
-            "The open-source Ax table creation is likely not applicable in this case,"
-            + "please contact the Adaptive Experimentation team if you need help."
+            "The open-source Ax table creation is likely not applicable in this case, "
+            "please contact the Adaptive Experimentation team if you need help."
         )
     Base.metadata.create_all(engine)
 

ax/storage/sqa_store/encoder.py

@@ -128,7 +128,7 @@ def validate_experiment_metadata(
                 raise ValueError(
                     f"An experiment already exists with the name {experiment.name}. "
                     "If you need to override this existing experiment, first delete it "
-                    "via `delete_experiment` in ax/ax/storage/sqa_store/delete.py, "
+                    "via `delete_experiment` in ax/storage/sqa_store/delete.py, "
                     "and then resave."
                 )
 
@@ -391,14 +391,15 @@ def get_metric_type_and_properties(
         json blob.
         """
         metric_class = type(metric)
-        metric_type = int(self.config.metric_registry.get(metric_class))
-        if metric_type is None:
+        metric_type_or_none = self.config.metric_registry.get(metric_class)
+        if metric_type_or_none is None:
             raise SQAEncodeError(
                 "Cannot encode metric to SQLAlchemy because metric's "
                 f"subclass ({metric_class}) is missing from the registry. "
                 "The metric registry currently contains the following: "
                 f"{','.join(map(str, self.config.metric_registry.keys()))} "
             )
+        metric_type = int(metric_type_or_none)
 
         properties = metric_class.serialize_init_args(obj=metric)
         return metric_type, object_to_json(
@@ -646,13 +647,13 @@ def outcome_constraint_to_sqa(
     def scalarized_outcome_constraint_to_sqa(
         self, outcome_constraint: ScalarizedOutcomeConstraint
     ) -> SQAMetric:
-        """Convert Ax SCalarized OutcomeConstraint to SQLAlchemy."""
+        """Convert Ax Scalarized OutcomeConstraint to SQLAlchemy."""
         metrics, weights = outcome_constraint.metrics, outcome_constraint.weights
 
         if metrics is None or weights is None or len(metrics) != len(weights):
             raise SQAEncodeError(
-                "Metrics and weights in scalarized OutcomeConstraint \
-                must be lists of equal length."
+                "Metrics and weights in scalarized OutcomeConstraint "
+                "must be lists of equal length."
             )
 
         metrics_by_name = self.get_children_metrics_by_name(

ax/storage/sqa_store/load.py

@@ -768,7 +768,7 @@ def _query_historical_experiments_given_parameters(
                 )
             )
             .filter(SQAExperiment.is_test == False)  # noqa E712 `is` won't work for SQA
-            .filter(SQAExperiment.id is not None)
+            .filter(SQAExperiment.id.isnot(None))
             # Experiments with some data
             .join(SQAData, SQAParameter.experiment_id == SQAData.experiment_id)
         )

ax/storage/sqa_store/save.py

@@ -260,7 +260,9 @@ def trial_to_reduced_state_sqa_encoder(t: BaseTrial) -> SQATrial:
             objs=trials_to_reduce_state,
             encode_func=trial_to_reduced_state_sqa_encoder,
             decode_func=decoder.trial_from_sqa,
-            decode_args_list=[{"experiment": experiment} for _ in range(len(trials))],
+            decode_args_list=[
+                {"experiment": experiment} for _ in range(len(trials_to_reduce_state))
+            ],
             modify_sqa=add_experiment_id,
             batch_size=batch_size,
         )
@@ -269,7 +271,7 @@ def trial_to_reduced_state_sqa_encoder(t: BaseTrial) -> SQATrial:
             objs=[latest_trial],
             encode_func=encoder.trial_to_sqa,
             decode_func=decoder.trial_from_sqa,
-            decode_args_list=[{"experiment": experiment} for _ in range(len(trials))],
+            decode_args_list=[{"experiment": experiment}],
             modify_sqa=add_experiment_id,
             batch_size=batch_size,
         )

ax/storage/sqa_store/utils.py

@@ -47,7 +47,9 @@
     "_metric_fetching_errors",
     "_data_rows",
 }
-SKIP_ATTRS_ERROR_SUFFIX = "Consider adding to COPY_DB_IDS_ATTRS_TO_SKIP if appropriate."
+SKIP_ATTRS_ERROR_SUFFIX = (
+    " Consider adding to COPY_DB_IDS_ATTRS_TO_SKIP if appropriate."
+)
 
 
 def is_foreign_key_field(field: str) -> bool:
@@ -75,7 +77,7 @@ def copy_db_ids(source: Any, target: Any, path: list[str] | None = None) -> None
     if len(path) > 15:
         # This shouldn't happen, but is a precaution against accidentally
         # introducing infinite loops
-        raise SQADecodeError(error_message_prefix + "Encountered path of length > 10.")
+        raise SQADecodeError(error_message_prefix + "Encountered path of length > 15.")
 
     if type(source) is not type(target):
         if not issubclass(type(target), type(source)):

ax/storage/sqa_store/with_db_settings_base.py

@@ -106,8 +106,9 @@ def __init__(
         if db_settings and (not DBSettings or not isinstance(db_settings, DBSettings)):
             raise ValueError(
                 "`db_settings` argument should be of type ax.storage.sqa_store."
-                f"(Got: {db_settings} of type {type(db_settings)}. "
-                "structs.DBSettings. To use `DBSettings`, you will need SQLAlchemy "
+                "structs.DBSettings. "
+                f"(Got: {db_settings} of type {type(db_settings)}). "
+                "To use `DBSettings`, you will need SQLAlchemy "
                 "installed in your environment (can be installed through pip)."
             )
         self._db_settings = db_settings or self._get_default_db_settings()