Move api module out of ax.preview (#3523)

Summary: Pull Request resolved: #3523

Reviewed By: bernardbeckerman

Differential Revision: D71401836

fbshipit-source-id: de26a92fa66d0c111c704502c399aa3196c4eb7c

Author: mpolson64

ax/analysis/tests/test_metric_summary.py

@@ -8,10 +8,10 @@
 import pandas as pd
 from ax.analysis.analysis import AnalysisCardCategory, AnalysisCardLevel
 from ax.analysis.metric_summary import MetricSummary
+from ax.api.client import Client
+from ax.api.configs import ExperimentConfig
 from ax.core.metric import Metric
 from ax.exceptions.core import UserInputError
-from ax.preview.api.client import Client
-from ax.preview.api.configs import ExperimentConfig
 from ax.utils.common.testutils import TestCase
 
 

ax/analysis/tests/test_search_space_summary.py

@@ -8,15 +8,15 @@
 import pandas as pd
 from ax.analysis.analysis import AnalysisCardCategory, AnalysisCardLevel
 from ax.analysis.search_space_summary import SearchSpaceSummary
-from ax.exceptions.core import UserInputError
-from ax.preview.api.client import Client
-from ax.preview.api.configs import (
+from ax.api.client import Client
+from ax.api.configs import (
     ChoiceParameterConfig,
     ExperimentConfig,
     ParameterScaling,
     ParameterType,
     RangeParameterConfig,
 )
+from ax.exceptions.core import UserInputError
 from ax.utils.common.testutils import TestCase
 
 

ax/analysis/tests/test_summary.py

@@ -9,10 +9,10 @@
 import pandas as pd
 from ax.analysis.analysis import AnalysisCardCategory, AnalysisCardLevel
 from ax.analysis.summary import Summary
+from ax.api.client import Client
+from ax.api.configs import ExperimentConfig, ParameterType, RangeParameterConfig
 from ax.core.trial import Trial
 from ax.exceptions.core import UserInputError
-from ax.preview.api.client import Client
-from ax.preview.api.configs import ExperimentConfig, ParameterType, RangeParameterConfig
 from ax.utils.common.testutils import TestCase
 from pyre_extensions import assert_is_instance, none_throws
 

ax/preview/api/__init__.py ax/api/__init__.py

@@ -6,8 +6,8 @@
 
 # pyre-strict
 
-from ax.preview.api.client import Client
-from ax.preview.api.configs import (
+from ax.api.client import Client
+from ax.api.configs import (
     ChoiceParameterConfig,
     ExperimentConfig,
     GenerationStrategyConfig,
@@ -17,7 +17,7 @@
     RangeParameterConfig,
     StorageConfig,
 )
-from ax.preview.api.types import TOutcome, TParameterization
+from ax.api.types import TOutcome, TParameterization
 
 __all__ = [
     "Client",

ax/preview/api/client.py ax/api/client.py

@@ -21,6 +21,19 @@
     markdown_analysis_card_from_analysis_e,
 )
 from ax.analysis.utils import choose_analyses
+from ax.api.configs import (
+    ExperimentConfig,
+    GenerationStrategyConfig,
+    OrchestrationConfig,
+    StorageConfig,
+)
+from ax.api.protocols.metric import IMetric
+from ax.api.protocols.runner import IRunner
+from ax.api.types import TOutcome, TParameterization
+from ax.api.utils.generation_strategy_dispatch import choose_generation_strategy
+from ax.api.utils.instantiation.from_config import experiment_from_config
+from ax.api.utils.instantiation.from_string import optimization_config_from_string
+from ax.api.utils.storage import db_settings_from_storage_config
 from ax.core.experiment import Experiment
 from ax.core.metric import Metric
 from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
@@ -36,21 +49,6 @@
 )
 from ax.exceptions.core import ObjectNotFoundError, UnsupportedError
 from ax.generation_strategy.generation_strategy import GenerationStrategy
-from ax.preview.api.configs import (
-    ExperimentConfig,
-    GenerationStrategyConfig,
-    OrchestrationConfig,
-    StorageConfig,
-)
-from ax.preview.api.protocols.metric import IMetric
-from ax.preview.api.protocols.runner import IRunner
-from ax.preview.api.types import TOutcome, TParameterization
-from ax.preview.api.utils.generation_strategy_dispatch import choose_generation_strategy
-from ax.preview.api.utils.instantiation.from_config import experiment_from_config
-from ax.preview.api.utils.instantiation.from_string import (
-    optimization_config_from_string,
-)
-from ax.preview.api.utils.storage import db_settings_from_storage_config
 from ax.service.scheduler import Scheduler, SchedulerOptions
 from ax.service.utils.best_point_mixin import BestPointMixin
 from ax.service.utils.with_db_settings_base import WithDBSettingsBase

ax/preview/api/configs.py ax/api/configs.py

@@ -10,7 +10,7 @@
 from enum import Enum
 from typing import Any
 
-from ax.preview.api.types import TParameterValue
+from ax.api.types import TParameterValue
 from ax.storage.registry_bundle import RegistryBundleBase
 
 

ax/preview/api/protocols/__init__.py ax/api/protocols/__init__.py

@@ -6,8 +6,8 @@
 
 # pyre-strict
 
-from ax.preview.api.protocols.metric import IMetric
-from ax.preview.api.protocols.runner import IRunner
+from ax.api.protocols.metric import IMetric
+from ax.api.protocols.runner import IRunner
 
 __all__ = [
     "IMetric",

ax/preview/api/protocols/metric.py ax/api/protocols/metric.py

@@ -9,7 +9,7 @@
 from collections.abc import Mapping
 from typing import Any
 
-from ax.preview.api.protocols.utils import _APIMetric
+from ax.api.protocols.utils import _APIMetric
 from pyre_extensions import override
 
 

ax/preview/api/protocols/runner.py ax/api/protocols/runner.py

@@ -9,9 +9,10 @@
 from collections.abc import Mapping
 from typing import Any
 
+from ax.api.protocols.utils import _APIRunner
+from ax.api.types import TParameterization
+
 from ax.core.trial_status import TrialStatus
-from ax.preview.api.protocols.utils import _APIRunner
-from ax.preview.api.types import TParameterization
 from pyre_extensions import override
 
 

ax/preview/api/protocols/utils.py ax/api/protocols/utils.py

@@ -13,6 +13,7 @@
 from typing import Any
 
 import pandas as pd
+from ax.api.types import TParameterization
 
 from ax.core.base_trial import BaseTrial, TrialStatus
 from ax.core.map_data import MapData, MapKeyInfo
@@ -21,7 +22,6 @@
 from ax.core.runner import Runner
 from ax.core.trial import Trial
 from ax.exceptions.storage import JSONEncodeError
-from ax.preview.api.types import TParameterization
 from ax.utils.common.result import Err, Ok
 from pyre_extensions import assert_is_instance, none_throws, override
 

ax/preview/api/tests/test_client.py ax/api/tests/test_client.py

@@ -12,6 +12,19 @@
 
 import pandas as pd
 from ax.analysis.plotly.parallel_coordinates import ParallelCoordinatesPlot
+from ax.api.client import Client
+from ax.api.configs import (
+    ChoiceParameterConfig,
+    ExperimentConfig,
+    GenerationStrategyConfig,
+    OrchestrationConfig,
+    ParameterType,
+    RangeParameterConfig,
+    StorageConfig,
+)
+from ax.api.protocols.metric import IMetric
+from ax.api.protocols.runner import IRunner
+from ax.api.types import TParameterization
 
 from ax.core.experiment import Experiment
 from ax.core.formatting_utils import DataType
@@ -31,19 +44,6 @@
 from ax.core.trial_status import TrialStatus
 from ax.early_stopping.strategies import PercentileEarlyStoppingStrategy
 from ax.exceptions.core import UnsupportedError
-from ax.preview.api.client import Client
-from ax.preview.api.configs import (
-    ChoiceParameterConfig,
-    ExperimentConfig,
-    GenerationStrategyConfig,
-    OrchestrationConfig,
-    ParameterType,
-    RangeParameterConfig,
-    StorageConfig,
-)
-from ax.preview.api.protocols.metric import IMetric
-from ax.preview.api.protocols.runner import IRunner
-from ax.preview.api.types import TParameterization
 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 (

ax/preview/api/types.py ax/api/types.py

@@ -8,6 +8,7 @@
 
 
 import torch
+from ax.api.configs import GenerationMethod, GenerationStrategyConfig
 from ax.core.trial_status import TrialStatus
 from ax.exceptions.core import UnsupportedError
 from ax.generation_strategy.generation_strategy import (
@@ -18,7 +19,6 @@
 from ax.generation_strategy.transition_criterion import MinTrials
 from ax.modelbridge.registry import Generators
 from ax.models.torch.botorch_modular.surrogate import ModelConfig, SurrogateSpec
-from ax.preview.api.configs import GenerationMethod, GenerationStrategyConfig
 from botorch.models.transforms.input import Normalize, Warp
 from gpytorch.kernels.linear_kernel import LinearKernel
 

ax/preview/__init__.py ax/api/utils/__init__.py

@@ -7,6 +7,14 @@
 
 
 import numpy as np
+from ax.api.configs import (
+    ChoiceParameterConfig,
+    ExperimentConfig,
+    ParameterScaling,
+    ParameterType,
+    RangeParameterConfig,
+)
+from ax.api.utils.instantiation.from_string import parse_parameter_constraint
 
 from ax.core.experiment import Experiment
 
@@ -21,14 +29,6 @@
 from ax.core.parameter_constraint import validate_constraint_parameters
 from ax.core.search_space import HierarchicalSearchSpace, SearchSpace
 from ax.exceptions.core import UserInputError
-from ax.preview.api.configs import (
-    ChoiceParameterConfig,
-    ExperimentConfig,
-    ParameterScaling,
-    ParameterType,
-    RangeParameterConfig,
-)
-from ax.preview.api.utils.instantiation.from_string import parse_parameter_constraint
 
 
 def parameter_from_config(

ax/preview/api/utils/generation_strategy_dispatch.py ax/api/utils/generation_strategy_dispatch.py

@@ -5,6 +5,18 @@
 
 # pyre-strict
 
+from ax.api.configs import (
+    ChoiceParameterConfig,
+    ExperimentConfig,
+    ParameterScaling,
+    ParameterType,
+    RangeParameterConfig,
+)
+from ax.api.utils.instantiation.from_config import (
+    _parameter_type_converter,
+    experiment_from_config,
+    parameter_from_config,
+)
 from ax.core.experiment import Experiment
 from ax.core.formatting_utils import DataType
 from ax.core.parameter import (
@@ -16,18 +28,6 @@
 from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.search_space import HierarchicalSearchSpace, SearchSpace
 from ax.exceptions.core import UserInputError
-from ax.preview.api.configs import (
-    ChoiceParameterConfig,
-    ExperimentConfig,
-    ParameterScaling,
-    ParameterType,
-    RangeParameterConfig,
-)
-from ax.preview.api.utils.instantiation.from_config import (
-    _parameter_type_converter,
-    experiment_from_config,
-    parameter_from_config,
-)
 from ax.utils.common.testutils import TestCase
 
 

ax/preview/api/utils/__init__.py ax/api/utils/instantiation/__init__.py

@@ -5,6 +5,13 @@
 
 # pyre-strict
 
+from ax.api.utils.instantiation.from_string import (
+    _sanitize_dot,
+    optimization_config_from_string,
+    parse_objective,
+    parse_outcome_constraint,
+    parse_parameter_constraint,
+)
 from ax.core.map_metric import MapMetric
 from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
 from ax.core.optimization_config import (
@@ -19,13 +26,6 @@
 )
 from ax.core.parameter_constraint import ParameterConstraint
 from ax.exceptions.core import UserInputError
-from ax.preview.api.utils.instantiation.from_string import (
-    _sanitize_dot,
-    optimization_config_from_string,
-    parse_objective,
-    parse_outcome_constraint,
-    parse_parameter_constraint,
-)
 from ax.utils.common.testutils import TestCase
 
 

ax/preview/api/utils/instantiation/from_config.py ax/api/utils/instantiation/from_config.py

@@ -3,7 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from ax.preview.api.configs import StorageConfig
+from ax.api.configs import StorageConfig
 from ax.storage.sqa_store.decoder import Decoder
 from ax.storage.sqa_store.encoder import Encoder
 from ax.storage.sqa_store.sqa_config import SQAConfig

ax/preview/api/utils/instantiation/from_string.py ax/api/utils/instantiation/from_string.py

@@ -6,13 +6,13 @@
 
 
 import torch
+from ax.api.configs import GenerationMethod, GenerationStrategyConfig
+from ax.api.utils.generation_strategy_dispatch import choose_generation_strategy
 from ax.core.trial import Trial
 from ax.core.trial_status import TrialStatus
 from ax.generation_strategy.transition_criterion import MinTrials
 from ax.modelbridge.registry import Generators
 from ax.models.torch.botorch_modular.surrogate import ModelConfig, SurrogateSpec
-from ax.preview.api.configs import GenerationMethod, GenerationStrategyConfig
-from ax.preview.api.utils.generation_strategy_dispatch import choose_generation_strategy
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
     get_branin_experiment,