Move WithDBSettingsBase to sqa_store (#4700)

Summary:
Pull Request resolved: #4700

Same reasoning as in D89737933: `ax/service` was always a misnomer (of my creation, no denying here), time for it to be ready to die once `AxClient` is deprecated!

Reviewed By: esantorella

Differential Revision:
D89749717

Privacy Context Container: L1307644

fbshipit-source-id: 38f37fc3837c58997c5c1146c2249272f4bc12cb

Author: Lena Kashtelyan

ax/api/client.py

@@ -46,7 +46,6 @@
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.orchestration.orchestrator import Orchestrator, OrchestratorOptions
 from ax.service.utils.best_point_mixin import BestPointMixin
-from ax.service.utils.with_db_settings_base import WithDBSettingsBase
 from ax.storage.json_store.decoder import (
     generation_strategy_from_json,
     object_from_json,
@@ -58,6 +57,7 @@
     CORE_DECODER_REGISTRY,
     CORE_ENCODER_REGISTRY,
 )
+from ax.storage.sqa_store.with_db_settings_base import WithDBSettingsBase
 from ax.utils.common.logger import _round_floats_for_logging, get_logger
 from ax.utils.common.random import with_rng_seed
 

ax/api/tests/test_client.py

@@ -40,10 +40,10 @@
 from ax.core.trial_status import TrialStatus
 from ax.early_stopping.strategies import PercentileEarlyStoppingStrategy
 from ax.exceptions.core import UnsupportedError, UserInputError
-from ax.service.utils.with_db_settings_base import (
+from ax.storage.sqa_store.db import init_test_engine_and_session_factory
+from ax.storage.sqa_store.with_db_settings_base import (
     _save_generation_strategy_to_db_if_possible,
 )
-from ax.storage.sqa_store.db import init_test_engine_and_session_factory
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
     get_branin_experiment,
@@ -392,7 +392,7 @@ def test_get_next_trials(self) -> None:
 
         # Test respects fixed features
         with mock.patch(
-            "ax.service.utils.with_db_settings_base"
+            "ax.storage.sqa_store.with_db_settings_base"
             "._save_generation_strategy_to_db_if_possible"
         ) as mock_save:
             trials = client.get_next_trials(max_trials=1, fixed_parameters={"x1": 0.5})
@@ -430,7 +430,7 @@ def test_get_next_trials_with_db(self) -> None:
         )
         # Generate one more trial, so that GS transitions to BO.
         with mock.patch(
-            "ax.service.utils.with_db_settings_base"
+            "ax.storage.sqa_store.with_db_settings_base"
             "._save_generation_strategy_to_db_if_possible",
             wraps=_save_generation_strategy_to_db_if_possible,
         ) as mock_save:

ax/orchestration/orchestrator.py

@@ -48,7 +48,7 @@
 from ax.orchestration.orchestrator_options import OrchestratorOptions, TrialType
 from ax.service.utils.best_point import derelativize_opt_config, is_row_feasible
 from ax.service.utils.best_point_mixin import BestPointMixin
-from ax.service.utils.with_db_settings_base import DBSettings, WithDBSettingsBase
+from ax.storage.sqa_store.with_db_settings_base import DBSettings, WithDBSettingsBase
 from ax.utils.common.constants import Keys
 from ax.utils.common.executils import retry_on_exception
 from ax.utils.common.logger import (

ax/orchestration/tests/test_orchestrator.py

@@ -78,7 +78,6 @@
     TEST_MEAN,
     TestOrchestrator,
 )
-from ax.service.utils.with_db_settings_base import WithDBSettingsBase
 from ax.storage.json_store.encoders import runner_to_dict
 from ax.storage.json_store.registry import CORE_DECODER_REGISTRY, CORE_ENCODER_REGISTRY
 from ax.storage.metric_registry import CORE_METRIC_REGISTRY
@@ -89,6 +88,7 @@
 from ax.storage.sqa_store.save import save_experiment
 from ax.storage.sqa_store.sqa_config import SQAConfig
 from ax.storage.sqa_store.structs import DBSettings
+from ax.storage.sqa_store.with_db_settings_base import WithDBSettingsBase
 from ax.utils.common.constants import Keys
 from ax.utils.common.logger import AX_ROOT_LOGGER_NAME
 from ax.utils.common.testutils import TestCase

ax/service/ax_client.py

@@ -61,7 +61,6 @@
     InstantiationBase,
     ObjectiveProperties,
 )
-from ax.service.utils.with_db_settings_base import TDBSettings
 from ax.storage.json_store.decoder import (
     generation_strategy_from_json,
     object_from_json,
@@ -74,6 +73,7 @@
     CORE_ENCODER_REGISTRY,
     TDecoderRegistry,
 )
+from ax.storage.sqa_store.with_db_settings_base import TDBSettings
 from ax.utils.common.docutils import copy_doc
 from ax.utils.common.executils import retry_on_exception
 from ax.utils.common.logger import _round_floats_for_logging, get_logger

ax/service/utils/analysis_base.py

@@ -12,7 +12,7 @@
 from ax.core.analysis_card import AnalysisCardBase
 from ax.core.experiment import Experiment
 from ax.generation_strategy.generation_strategy import GenerationStrategy
-from ax.service.utils.with_db_settings_base import WithDBSettingsBase
+from ax.storage.sqa_store.with_db_settings_base import WithDBSettingsBase
 
 
 class AnalysisBase(WithDBSettingsBase):

…vice/tests/test_with_db_settings_base.py …tore/tests/test_with_db_settings_base.pyax/service/tests/test_with_db_settings_base.py renamed to ax/storage/sqa_store/tests/test_with_db_settings_base.py ax/service/tests/test_with_db_settings_base.py renamed to ax/storage/sqa_store/tests/test_with_db_settings_base.py

@@ -13,10 +13,6 @@
 from ax.core.experiment import Experiment
 from ax.core.trial_status import TrialStatus
 from ax.generation_strategy.generation_strategy import GenerationStrategy
-from ax.service.utils.with_db_settings_base import (
-    try_load_generation_strategy,
-    WithDBSettingsBase,
-)
 from ax.storage.sqa_store.db import init_test_engine_and_session_factory
 from ax.storage.sqa_store.load import (
     _load_experiment,
@@ -29,6 +25,10 @@
     _save_or_update_trials,
 )
 from ax.storage.sqa_store.structs import DBSettings
+from ax.storage.sqa_store.with_db_settings_base import (
+    try_load_generation_strategy,
+    WithDBSettingsBase,
+)
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import DEFAULT_USER, get_experiment, get_generator_run
 from ax.utils.testing.modeling_stubs import get_generation_strategy
@@ -362,7 +362,7 @@ def test_try_load_generation_strategy(self) -> None:
         )
         # test logging with no experiment/gs saved
         with self.assertLogs(
-            logger="ax.service.utils.with_db_settings_base", level=logging.DEBUG
+            logger="ax.storage.sqa_store.with_db_settings_base", level=logging.DEBUG
         ) as lg:
             output = try_load_generation_strategy(
                 experiment_name=experiment.name,
@@ -384,7 +384,7 @@ def test_try_load_generation_strategy(self) -> None:
         self.assertTrue(exp_saved)
         self.assertTrue(gs_saved)
         with self.assertLogs(
-            logger="ax.service.utils.with_db_settings_base", level=logging.DEBUG
+            logger="ax.storage.sqa_store.with_db_settings_base", level=logging.DEBUG
         ) as lg:
             output = try_load_generation_strategy(
                 experiment_name=experiment.name,

…db_settings_base_minimal_dependencies.py …db_settings_base_minimal_dependencies.pyax/service/tests/test_with_db_settings_base_minimal_dependencies.py renamed to ax/storage/sqa_store/tests/test_with_db_settings_base_minimal_dependencies.py ax/service/tests/test_with_db_settings_base_minimal_dependencies.py renamed to ax/storage/sqa_store/tests/test_with_db_settings_base_minimal_dependencies.py

@@ -13,4 +13,6 @@
 class TestWithDbSetingsBaseMinimalDependencies(TestCase):
     @patch.dict("sys.modules", {"sqlalchemy": None})
     def test_with_db_settings_base_no_sql_alchemy(self) -> None:
-        from ax.service.utils.with_db_settings_base import WithDBSettingsBase  # noqa
+        from ax.storage.sqa_store.with_db_settings_base import (  # noqa
+            WithDBSettingsBase,
+        )

ax/service/utils/with_db_settings_base.py …orage/sqa_store/with_db_settings_base.pyax/service/utils/with_db_settings_base.py renamed to ax/storage/sqa_store/with_db_settings_base.py ax/service/utils/with_db_settings_base.py renamed to ax/storage/sqa_store/with_db_settings_base.py

@@ -83,15 +83,6 @@ Reporting
     :show-inheritance:
 
 
-WithDBSettingsBase
-~~~~~~~~~~~~~~~~~~
-
-.. automodule:: ax.service.utils.with_db_settings_base
-    :members:
-    :undoc-members:
-    :show-inheritance:
-
-
 EarlyStopping
 ~~~~~~~~~~~~~