Remove pyre-fixme/pyre-ignore from ax/service/ and ax/api/ (#4979)

Summary:
Pull Request resolved: #4979

Remove ~43 pyre-fixme/pyre-ignore suppression comments from 12 files
in ax/service/ (source) and ax/api/ (source + tests):

- Use `cast()` for API TParameterization/TParamValue type mismatches
- Use `assert_is_instance()` and `assert_is_instance_optional()` for config values
- Add `partial[Any]` annotation for `round_floats_for_logging`
- Convert bare attribute declarations to `property abstractmethod`
- Use `npt.NDArray` for numpy array types
- Add `Literal` types for test parameterization

Reviewed By: dme65

Differential Revision: D95264965

fbshipit-source-id: 3291e2d9da3b9cf82bd87e626135b6d011690495

Author: saitcakmak

ax/api/client.py

@@ -8,7 +8,7 @@
 import json
 from collections.abc import Iterable, Sequence
 from logging import Logger
-from typing import Any, Literal, Self
+from typing import Any, cast, Literal, Self
 
 import numpy as np
 import pandas as pd
@@ -39,6 +39,7 @@
 from ax.core.runner import Runner
 from ax.core.trial import Trial
 from ax.core.trial_status import TrialStatus  # Used as a return type
+from ax.core.types import TParameterization as CoreTParameterization
 from ax.early_stopping.strategies import (
     BaseEarlyStoppingStrategy,
     PercentileEarlyStoppingStrategy,
@@ -183,8 +184,7 @@ def configure_optimization(
         pruning_target_arm: Arm | None = None
         if pruning_target_parameterization is not None:
             self._experiment.search_space.validate_membership(
-                # pyre-fixme[6]: Core Ax TParameterization is dict not Mapping
-                parameters=pruning_target_parameterization
+                parameters=cast(CoreTParameterization, pruning_target_parameterization)
             )
             pruning_target_arm = Arm(
                 parameters=pruning_target_parameterization, name="pruning_target"
@@ -442,9 +442,9 @@ def get_next_trials(
                 experiment=self._experiment,
                 n=1,
                 fixed_features=(
-                    # pyre-fixme[6]: Type narrowing broken because core Ax
-                    # TParameterization is dict not Mapping
-                    ObservationFeatures(parameters=fixed_parameters)
+                    ObservationFeatures(
+                        parameters=cast(CoreTParameterization, fixed_parameters)
+                    )
                     if fixed_parameters is not None
                     else None
                 ),
@@ -483,9 +483,10 @@ def get_next_trials(
             experiment=self._experiment, trials=trials
         )
 
-        # pyre-fixme[7]: Core Ax allows users to specify TParameterization values as
-        # None, but we do not allow this in the API.
-        return {trial.index: none_throws(trial.arm).parameters for trial in trials}
+        return {
+            trial.index: cast(TParameterization, none_throws(trial.arm).parameters)
+            for trial in trials
+        }
 
     def complete_trial(
         self,
@@ -573,9 +574,7 @@ def attach_trial(
             The index of the attached trial.
         """
         _, trial_index = self._experiment.attach_trial(
-            # pyre-fixme[6]: Type narrowing broken because core Ax TParameterization
-            # is dict not Mapping
-            parameterizations=[parameters],
+            parameterizations=[cast(CoreTParameterization, parameters)],
             arm_names=[arm_name] if arm_name else None,
         )
 
@@ -888,13 +887,14 @@ def get_best_parameterization(
             )
         )
 
-        # pyre-fixme[7]: Core Ax allows users to specify TParameterization values as
-        # None but we do not allow this in the API.
-        return BestPointMixin._to_best_point_tuple(
-            experiment=self._experiment,
-            trial_index=trial_index,
-            parameterization=parameterization,
-            model_prediction=model_prediction,
+        return cast(
+            tuple[TParameterization, TOutcome, int, str],
+            BestPointMixin._to_best_point_tuple(
+                experiment=self._experiment,
+                trial_index=trial_index,
+                parameterization=parameterization,
+                model_prediction=model_prediction,
+            ),
         )
 
     def get_pareto_frontier(
@@ -945,14 +945,15 @@ def get_pareto_frontier(
             use_model_predictions=use_model_predictions,
         )
 
-        # pyre-fixme[7]: Core Ax allows users to specify TParameterization values as
-        # None but we do not allow this in the API.
         return [
-            BestPointMixin._to_best_point_tuple(
-                experiment=self._experiment,
-                trial_index=trial_index,
-                parameterization=parameterization,
-                model_prediction=model_prediction,
+            cast(
+                tuple[TParameterization, TOutcome, int, str],
+                BestPointMixin._to_best_point_tuple(
+                    experiment=self._experiment,
+                    trial_index=trial_index,
+                    parameterization=parameterization,
+                    model_prediction=model_prediction,
+                ),
             )
             for trial_index, (parameterization, model_prediction) in frontier.items()
         ]
@@ -978,9 +979,9 @@ def predict(
         try:
             mean, covariance = none_throws(self._generation_strategy.adapter).predict(
                 observation_features=[
-                    # pyre-fixme[6]: Core Ax allows users to specify TParameterization
-                    # values as None but we do not allow this in the API.
-                    ObservationFeatures(parameters=parameters)
+                    ObservationFeatures(
+                        parameters=cast(CoreTParameterization, parameters)
+                    )
                     for parameters in points
                 ]
             )

ax/api/protocols/utils.py

@@ -10,7 +10,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Iterable, Mapping
-from typing import Any
+from typing import Any, cast
 
 import pandas as pd
 from ax.api.types import TParameterization
@@ -115,10 +115,10 @@ def run(self, trial: BaseTrial) -> dict[str, Any]:
         """
         metadata = self.run_trial(
             trial_index=trial.index,
-            # pyre-ignore[6] Arms in core Ax may have None in their parameters
-            parameterization=none_throws(
-                assert_is_instance(trial, Trial).arm
-            ).parameters,
+            parameterization=cast(
+                TParameterization,
+                none_throws(assert_is_instance(trial, Trial).arm).parameters,
+            ),
         )
 
         # Runtime validate metadata is JSON serializable to avoid issues when

ax/api/tests/test_client.py

@@ -1794,10 +1794,11 @@ def test_get_next_trials_with_derived_parameters(self) -> None:
             self.assertIn("x2", trial_params)
             self.assertIn("x3", trial_params)
             # Verify derived parameter is correctly computed
-            # pyre-fixme[58]: Arithmetic operations on TParameterValue
-            expected_x3 = 1.0 - trial_params["x1"] - trial_params["x2"]
-            # pyre-fixme[6]: Type mismatch on assertAlmostEqual
-            self.assertAlmostEqual(trial_params["x3"], expected_x3, places=6)
+            x1_val = assert_is_instance(trial_params["x1"], float)
+            x2_val = assert_is_instance(trial_params["x2"], float)
+            x3_val = assert_is_instance(trial_params["x3"], float)
+            expected_x3 = 1.0 - x1_val - x2_val
+            self.assertAlmostEqual(x3_val, expected_x3, places=6)
 
     def test_complete_trial_with_derived_parameters(self) -> None:
         # Setup: Configure experiment with derived parameter and generate trial

ax/api/utils/instantiation/from_config.py

@@ -6,6 +6,8 @@
 # pyre-strict
 
 
+from typing import cast
+
 import numpy as np
 from ax.api.configs import (
     ChoiceParameterConfig,
@@ -20,6 +22,7 @@
     ParameterType as CoreParameterType,
     RangeParameter,
 )
+from ax.core.types import TParamValue
 from ax.exceptions.core import UserInputError
 
 
@@ -79,21 +82,21 @@ def parameter_from_config(
                 name=config.name,
                 parameter_type=_parameter_type_converter(config.parameter_type),
                 value=config.values[0],
-                # pyre-fixme[6] Variance issue caused by FixedParameter.dependents
-                # using List instead of immutable container type.
-                dependents=config.dependent_parameters,
+                dependents=cast(
+                    dict[TParamValue, list[str]] | None,
+                    config.dependent_parameters,
+                ),
             )
 
         return ChoiceParameter(
             name=config.name,
             parameter_type=_parameter_type_converter(config.parameter_type),
-            # pyre-fixme[6] Variance issue caused by ChoiceParameter.value using List
-            # instead of immutable container type.
-            values=config.values,
+            values=cast(list[TParamValue], config.values),
             is_ordered=config.is_ordered,
-            # pyre-fixme[6] Variance issue caused by ChoiceParameter.dependents using
-            # List instead of immutable container type.
-            dependents=config.dependent_parameters,
+            dependents=cast(
+                dict[TParamValue, list[str]] | None,
+                config.dependent_parameters,
+            ),
             sort_values=config.parameter_type != "str",  # Matches default behavior.
         )
 

ax/api/utils/tests/test_generation_strategy_dispatch.py

@@ -8,7 +8,7 @@
 
 
 from itertools import product
-from typing import Any
+from typing import Any, Literal
 
 import torch
 from ax.adapter.registry import Generators
@@ -234,13 +234,9 @@ def test_choose_gs_single_sobol_initialization(self) -> None:
         self.assertEqual(mbm_node.name, "MBM")
 
     def test_gs_simplify_parameter_changes(self) -> None:
-        for simplify, method in product((True, False), ("fast", "quality")):
+        methods: list[Literal["fast", "quality"]] = ["fast", "quality"]
+        for simplify, method in product((True, False), methods):
             struct = GenerationStrategyDispatchStruct(
-                # pyre-fixme [6]: In call
-                # `GenerationStrategyDispatchStruct.__init__`, for argument
-                # `method`, expected `Union[typing_extensions.Literal['fast'],
-                # typing_extensions.Literal['quality'],
-                # typing_extensions.Literal['random_search']]` but got `str`
                 method=method,
                 simplify_parameter_changes=simplify,
             )
@@ -323,7 +319,13 @@ def test_choose_gs_with_custom_botorch_acqf_class(self) -> None:
         """Test that custom botorch_acqf_class is properly passed to generator kwargs
         and appended to the node name. Tests both fast and custom methods.
         """
-        for method, model_config, expected_name in [
+        test_cases: list[
+            tuple[
+                Literal["quality", "fast", "random_search", "custom"],
+                ModelConfig | None,
+                str,
+            ]
+        ] = [
             ("fast", None, "Sobol+MBM:fast+qLogNoisyExpectedImprovement"),
             (
                 "custom",
@@ -333,10 +335,11 @@ def test_choose_gs_with_custom_botorch_acqf_class(self) -> None:
                 ),
                 "Sobol+MBM:MAPSAAS+qLogNoisyExpectedImprovement",
             ),
-        ]:
+        ]
+        for method, model_config, expected_name in test_cases:
             with self.subTest(method=method):
                 struct = GenerationStrategyDispatchStruct(
-                    method=method,  # pyre-ignore [6]
+                    method=method,
                     initialization_budget=3,
                     initialize_with_center=False,
                 )

ax/service/ax_client.py

@@ -87,8 +87,7 @@
 
 ROUND_FLOATS_IN_LOGS_TO_DECIMAL_PLACES: int = 6
 
-# pyre-fixme[5]: Global expression must be annotated.
-round_floats_for_logging = partial(
+round_floats_for_logging: partial[Any] = partial(
     _round_floats_for_logging,
     decimal_places=ROUND_FLOATS_IN_LOGS_TO_DECIMAL_PLACES,
 )
@@ -1303,8 +1302,7 @@ def save_to_json_file(self, filepath: str = "ax_client_snapshot.json") -> None:
     def load_from_json_file(
         cls: type[AxClientSubclass],
         filepath: str = "ax_client_snapshot.json",
-        # pyre-fixme[2]: Parameter must be annotated.
-        **kwargs,
+        **kwargs: Any,
     ) -> AxClientSubclass:
         """Restore an `AxClient` and its state from a JSON-serialized snapshot,
         residing in a .json file by the given path.
@@ -1315,13 +1313,10 @@ def load_from_json_file(
 
     def to_json_snapshot(
         self,
-        # pyre-fixme[24]: Generic type `type` expects 1 type parameter, use
-        #  `typing.Type` to avoid runtime subscripting errors.
-        encoder_registry: dict[type, Callable[[Any], dict[str, Any]]] | None = None,
-        # pyre-fixme[24]: Generic type `type` expects 1 type parameter, use
-        #  `typing.Type` to avoid runtime subscripting errors.
-        class_encoder_registry: None
-        | (dict[type, Callable[[Any], dict[str, Any]]]) = None,
+        encoder_registry: dict[type[Any], Callable[[Any], dict[str, Any]]]
+        | None = None,
+        class_encoder_registry: dict[type[Any], Callable[[Any], dict[str, Any]]]
+        | None = None,
     ) -> dict[str, Any]:
         """Serialize this `AxClient` to JSON to be able to interrupt and restart
         optimization and save it to file by the provided path.
@@ -1357,8 +1352,7 @@ def from_json_snapshot(
         decoder_registry: TDecoderRegistry | None = None,
         class_decoder_registry: None
         | (dict[str, Callable[[dict[str, Any]], Any]]) = None,
-        # pyre-fixme[2]: Parameter must be annotated.
-        **kwargs,
+        **kwargs: Any,
     ) -> AxClientSubclass:
         """Recreate an `AxClient` from a JSON snapshot."""
         if decoder_registry is None:

ax/service/interactive_loop.py

@@ -28,7 +28,6 @@ def interactive_optimize(
     candidate_queue_maxsize: int,
     candidate_generator_function: Callable[..., None],
     data_attacher_function: Callable[..., None],
-    # pyre-ignore[2]: Missing parameter annotation
     elicitation_function: Callable[..., Any],
     candidate_generator_kwargs: dict[str, Any] | None = None,
     data_attacher_kwargs: dict[str, Any] | None = None,

ax/service/managed_loop.py

@@ -11,7 +11,8 @@
 import inspect
 import logging
 import warnings
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
+from typing import cast
 
 # Manual import to avoid strange error, see Diff for details.
 import ax.generation_strategy.generation_node_input_constructors  # noqa
@@ -73,12 +74,13 @@ def __init__(
         )
         self.experiment = experiment
         if generation_strategy is None:
-            # pyre-fixme[4]: Attribute must be annotated.
-            self.generation_strategy = choose_generation_strategy_legacy(
-                search_space=experiment.search_space,
-                use_batch_trials=self.arms_per_trial > 1,
-                random_seed=self.random_seed,
-                experiment=experiment,
+            self.generation_strategy: GenerationStrategy = (
+                choose_generation_strategy_legacy(
+                    search_space=experiment.search_space,
+                    use_batch_trials=self.arms_per_trial > 1,
+                    random_seed=self.random_seed,
+                    experiment=experiment,
+                )
             )
         else:
             self.generation_strategy = generation_strategy
@@ -147,11 +149,15 @@ def _call_evaluation_function(
         signature = inspect.signature(self.evaluation_function)
         num_evaluation_function_params = len(signature.parameters.items())
         if num_evaluation_function_params == 1:
-            # pyre-ignore [20]: Can't run instance checks on subscripted generics.
-            evaluation = self.evaluation_function(parameterization)
+            evaluation = cast(
+                Callable[[TParameterization], TEvaluationOutcome],
+                self.evaluation_function,
+            )(parameterization)
         elif num_evaluation_function_params == 2:
-            # pyre-ignore [19]: Can't run instance checks on subscripted generics.
-            evaluation = self.evaluation_function(parameterization, weight)
+            evaluation = cast(
+                Callable[[TParameterization, float | None], TEvaluationOutcome],
+                self.evaluation_function,
+            )(parameterization, weight)
         else:
             raise UserInputError(
                 "Evaluation function must take either one parameter "