Remove pyre-fixme/pyre-ignore from ax/service, ax/storage, ax/utils test files (#4990)

Summary:
Pull Request resolved: #4990

Remove pyre-fixme and pyre-ignore type suppression comments from test files in
ax/service/tests, ax/storage/*/tests, and ax/utils/*/tests. Uses proper type narrowing
via none_throws, assert_is_instance, cast, and explicit type annotations instead of
suppression comments.

Reviewed By: dme65

Differential Revision: D95273568

fbshipit-source-id: c8a716f5558dcd443692b9b53050d02fefa0c8e1

Author: saitcakmak

ax/service/tests/test_ax_client.py

@@ -48,9 +48,7 @@ def get_ax_client_for_branin(
     def evaluate(self, parameters: TParameterization) -> dict[str, tuple[float, float]]:
         """Evaluates the parameters for branin experiment."""
         x = np.array([parameters.get(f"x{i + 1}") for i in range(2)])
-        # pyre-fixme[7]: Expected `Dict[str, Tuple[float, float]]` but got
-        #  `Dict[str, Tuple[Union[float, ndarray], float]]`.
-        return {"branin": (branin(x), 0.0)}
+        return {"branin": (float(branin(x)), 0.0)}
 
     def test_global_stopping_integration(self) -> None:
         """
@@ -69,7 +67,6 @@ def test_global_stopping_integration(self) -> None:
             parameters, trial_index = ax_client.get_next_trial()
             ax_client.complete_trial(
                 trial_index=trial_index,
-                # pyre-fixme[6]: For 2nd param expected `Union[Dict[str, Union[Tuple[...
                 raw_data=self.evaluate(parameters),
             )
 
@@ -109,15 +106,13 @@ def test_min_trials(self) -> None:
             parameters, trial_index = ax_client.get_next_trial()
             ax_client.complete_trial(
                 trial_index=trial_index,
-                # pyre-fixme[6]: For 2nd param expected `Union[Dict[str, Union[Tuple[...
                 raw_data=self.evaluate(parameters),
             )
 
         # Since min_trials=3, GSS should not stop creating the 3rd iteration.
         parameters, trial_index = ax_client.get_next_trial()
         ax_client.complete_trial(
             trial_index=trial_index,
-            # pyre-fixme[6]: For 2nd param expected `Union[Dict[str, Union[Tuple[Unio...
             raw_data=self.evaluate(parameters),
         )
         self.assertIsNotNone(parameters)

ax/service/tests/test_global_stopping.py

@@ -6,6 +6,7 @@
 
 # pyre-strict
 
+from collections.abc import Sequence
 from typing import Any
 
 from ax.core.metric import Metric
@@ -16,6 +17,7 @@
     ParameterType,
     RangeParameter,
 )
+from ax.core.types import TParamValue
 from ax.runners.synthetic import SyntheticRunner
 from ax.service.utils.instantiation import InstantiationBase
 from ax.utils.common.testutils import TestCase
@@ -332,7 +334,9 @@ def test_choice_with_is_sorted(self) -> None:
             _ = InstantiationBase.parameter_from_json(representation)
 
     def test_hss(self) -> None:
-        parameter_dicts = [
+        parameter_dicts: list[
+            dict[str, TParamValue | Sequence[TParamValue] | dict[str, list[str]]]
+        ] = [
             {
                 "name": "root",
                 "type": "fixed",
@@ -361,7 +365,6 @@ def test_hss(self) -> None:
             {"name": "another_int", "type": "fixed", "value": "2"},
         ]
         search_space = InstantiationBase.make_search_space(
-            # pyre-fixme[6]: For 1st param expected `List[Dict[str, Union[None, Dict[...
             parameters=parameter_dicts,
             parameter_constraints=[],
         )

ax/service/tests/test_instantiation_utils.py

@@ -9,13 +9,15 @@
 
 import functools
 import time
+from collections.abc import Callable, Sequence
 from logging import WARN
 from queue import Queue
 from threading import Event, Lock
+from typing import cast
 
 import numpy as np
 from ax.adapter.registry import Generators
-from ax.core.types import TEvaluationOutcome, TParameterization
+from ax.core.types import TEvaluationOutcome, TParameterization, TParamValue
 from ax.generation_strategy.generation_strategy import (
     GenerationStep,
     GenerationStrategy,
@@ -41,17 +43,20 @@ def setUp(self) -> None:
             ]
         )
         self.ax_client = AxClient(generation_strategy=generation_strategy)
-        self.ax_client.create_experiment(
-            name="hartmann_test_experiment",
-            # pyre-fixme[6]
-            parameters=[
+        parameters = cast(
+            list[dict[str, TParamValue | Sequence[TParamValue] | dict[str, list[str]]]],
+            [
                 {
                     "name": f"x{i}",
                     "type": "range",
                     "bounds": [0.0, 1.0],
                 }
                 for i in range(1, 7)
             ],
+        )
+        self.ax_client.create_experiment(
+            name="hartmann_test_experiment",
+            parameters=parameters,
             objectives={"hartmann6": ObjectiveProperties(minimize=True)},
             tracking_metric_names=["l2norm"],
         )
@@ -76,8 +81,10 @@ def test_interactive_loop(self) -> None:
             ax_client=self.ax_client,
             num_trials=15,
             candidate_queue_maxsize=3,
-            # pyre-fixme[6]
-            elicitation_function=self._elicit,
+            elicitation_function=cast(
+                Callable[[tuple[TParameterization, int]], TEvaluationOutcome],
+                self._elicit,
+            ),
         )
 
         self.assertTrue(optimization_completed)
@@ -94,8 +101,10 @@ def _aborted_elicit(
             ax_client=self.ax_client,
             num_trials=15,
             candidate_queue_maxsize=3,
-            # pyre-fixme[6]
-            elicitation_function=_aborted_elicit,
+            elicitation_function=cast(
+                Callable[[tuple[TParameterization, int]], TEvaluationOutcome],
+                _aborted_elicit,
+            ),
         )
         self.assertFalse(optimization_completed)
 
@@ -144,8 +153,10 @@ def _sleep_elicit(
                 ax_client=self.ax_client,
                 num_trials=3,
                 candidate_queue_maxsize=3,
-                # pyre-fixme[6]
-                elicitation_function=_sleep_elicit,
+                elicitation_function=cast(
+                    Callable[[tuple[TParameterization, int]], TEvaluationOutcome],
+                    _sleep_elicit,
+                ),
             )
 
             # Assert sleep and retry warning is somewhere in the logs

ax/service/tests/test_interactive_loop.py

@@ -6,63 +6,63 @@
 
 # pyre-strict
 
+from typing import Any
 from unittest.mock import Mock, patch
 
 import numpy as np
-import numpy.typing as npt
 from ax.adapter.registry import Generators
+from ax.core.types import TParameterization
 from ax.exceptions.core import UserInputError
 from ax.generation_strategy.generation_strategy import (
     GenerationStep,
     GenerationStrategy,
 )
+from ax.generators.random.sobol import SobolGenerator
 from ax.metrics.branin import branin
 from ax.service.managed_loop import OptimizationLoop, optimize
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.mock import mock_botorch_optimize
+from pyre_extensions import assert_is_instance, none_throws
 
 
 def _branin_evaluation_function(
-    # pyre-fixme[2]: Parameter must be annotated.
-    parameterization,
-    weight=None,  # pyre-fixme[2]: Parameter must be annotated.
-) -> dict[str, tuple[float | npt.NDArray, float]]:
+    parameterization: TParameterization,
+    weight: float | None = None,
+) -> dict[str, tuple[float, float]]:
     if any(param_name not in parameterization.keys() for param_name in ["x1", "x2"]):
         raise ValueError("Parametrization does not contain x1 or x2")
-    x1, x2 = parameterization["x1"], parameterization["x2"]
+    x1, x2 = float(parameterization["x1"]), float(parameterization["x2"])
     return {
-        "branin": (branin(x1, x2), 0.0),
-        "constrained_metric": (-branin(x1, x2), 0.0),
+        "branin": (float(branin(x1, x2)), 0.0),
+        "constrained_metric": (float(-branin(x1, x2)), 0.0),
     }
 
 
 def _branin_evaluation_function_v2(
-    # pyre-fixme[2]: Parameter must be annotated.
-    parameterization,
-    weight=None,  # pyre-fixme[2]: Parameter must be annotated.
-) -> tuple[float | npt.NDArray, float]:
+    parameterization: TParameterization,
+    weight: float | None = None,
+) -> tuple[float, float]:
     if any(param_name not in parameterization.keys() for param_name in ["x1", "x2"]):
         raise ValueError("Parametrization does not contain x1 or x2")
-    x1, x2 = parameterization["x1"], parameterization["x2"]
-    return (branin(x1, x2), 0.0)
+    x1, x2 = float(parameterization["x1"]), float(parameterization["x2"])
+    return (float(branin(x1, x2)), 0.0)
 
 
 def _branin_evaluation_function_with_unknown_sem(
-    # pyre-fixme[2]: Parameter must be annotated.
-    parameterization,
-    weight=None,  # pyre-fixme[2]: Parameter must be annotated.
-) -> tuple[float | npt.NDArray, None]:
+    parameterization: TParameterization,
+    weight: float | None = None,
+) -> tuple[float, None]:
     if any(param_name not in parameterization.keys() for param_name in ["x1", "x2"]):
         raise ValueError("Parametrization does not contain x1 or x2")
-    x1, x2 = parameterization["x1"], parameterization["x2"]
-    return (branin(x1, x2), None)
+    x1, x2 = float(parameterization["x1"]), float(parameterization["x2"])
+    return (float(branin(x1, x2)), None)
 
 
 class TestManagedLoop(TestCase):
     """Check functionality of optimization loop."""
 
     def test_with_evaluation_function_propagates_parameter_constraints(self) -> None:
-        kwargs = {
+        kwargs: dict[str, Any] = {
             "parameters": [
                 {
                     "name": "x1",
@@ -151,9 +151,7 @@ def test_branin_with_active_parameter_constraints(self) -> None:
         bp, _ = loop.full_run().get_best_point()
         self.assertIn("x1", bp)
         self.assertIn("x2", bp)
-        # pyre-fixme[58]: `+` is not supported for operand types `Union[None, bool,
-        #  float, int, str]` and `Union[None, bool, float, int, str]`.
-        self.assertLessEqual(bp["x1"] + bp["x2"], 1.0 + 1e-8)
+        self.assertLessEqual(float(bp["x1"]) + float(bp["x2"]), 1.0 + 1e-8)
         with self.assertRaisesRegex(ValueError, "Optimization is complete"):
             loop.run_trial()
 
@@ -241,11 +239,8 @@ def test_branin_batch(self) -> None:
         self.assertIn("x2", bp)
         assert vals is not None
         self.assertIn("branin", vals[0])
-        # pyre-fixme[6]: For 2nd param expected `Union[Container[typing.Any],
-        #  Iterable[typing.Any]]` but got `Optional[Dict[str, Dict[str, float]]]`.
-        self.assertIn("branin", vals[1])
-        # pyre-fixme[16]: Optional type has no attribute `__getitem__`.
-        self.assertIn("branin", vals[1]["branin"])
+        self.assertIn("branin", none_throws(vals[1]))
+        self.assertIn("branin", none_throws(vals[1])["branin"])
         # Check that all total_trials * arms_per_trial * 2 metrics evaluations
         # are present in the dataframe.
         self.assertEqual(len(loop.experiment.fetch_data().df.index), 12)
@@ -270,11 +265,8 @@ def test_optimize(self) -> None:
         self.assertIn("x2", best)
         assert vals is not None
         self.assertIn("objective", vals[0])
-        # pyre-fixme[6]: For 2nd param expected `Union[Container[typing.Any],
-        #  Iterable[typing.Any]]` but got `Optional[Dict[str, Dict[str, float]]]`.
-        self.assertIn("objective", vals[1])
-        # pyre-fixme[16]: Optional type has no attribute `__getitem__`.
-        self.assertIn("objective", vals[1]["objective"])
+        self.assertIn("objective", none_throws(vals[1]))
+        self.assertIn("objective", none_throws(vals[1])["objective"])
 
     @patch(
         "ax.service.managed_loop."
@@ -301,11 +293,8 @@ def test_optimize_with_predictions(self, _) -> None:
         self.assertIn("x2", best)
         assert vals is not None
         self.assertIn("a", vals[0])
-        # pyre-fixme[6]: For 2nd param expected `Union[Container[typing.Any],
-        #  Iterable[typing.Any]]` but got `Optional[Dict[str, Dict[str, float]]]`.
-        self.assertIn("a", vals[1])
-        # pyre-fixme[16]: Optional type has no attribute `__getitem__`.
-        self.assertIn("a", vals[1]["a"])
+        self.assertIn("a", none_throws(vals[1]))
+        self.assertIn("a", none_throws(vals[1])["a"])
 
     @mock_botorch_optimize
     def test_optimize_unknown_sem(self) -> None:
@@ -327,11 +316,8 @@ def test_optimize_unknown_sem(self) -> None:
         self.assertIn("x2", best)
         self.assertIsNotNone(vals)
         self.assertIn("objective", vals[0])
-        # pyre-fixme[6]: For 2nd param expected `Union[Container[typing.Any],
-        #  Iterable[typing.Any]]` but got `Optional[Dict[str, Dict[str, float]]]`.
-        self.assertIn("objective", vals[1])
-        # pyre-fixme[16]: Optional type has no attribute `__getitem__`.
-        self.assertIn("objective", vals[1]["objective"])
+        self.assertIn("objective", none_throws(vals[1]))
+        self.assertIn("objective", none_throws(vals[1])["objective"])
 
     def test_optimize_propagates_random_seed(self) -> None:
         """Tests optimization as a single call."""
@@ -347,8 +333,8 @@ def test_optimize_propagates_random_seed(self) -> None:
             total_trials=5,
             random_seed=12345,
         )
-        # pyre-fixme[16]: Optional type has no attribute `model`.
-        self.assertEqual(12345, model.generator.seed)
+        generator = assert_is_instance(none_throws(model).generator, SobolGenerator)
+        self.assertEqual(12345, generator.seed)
 
     def test_optimize_search_space_exhausted(self) -> None:
         """Tests optimization as a single call."""
@@ -370,11 +356,8 @@ def test_optimize_search_space_exhausted(self) -> None:
         self.assertIn("x2", best)
         self.assertIsNotNone(vals)
         self.assertIn("objective", vals[0])
-        # pyre-fixme[6]: For 2nd param expected `Union[Container[typing.Any],
-        #  Iterable[typing.Any]]` but got `Optional[Dict[str, Dict[str, float]]]`.
-        self.assertIn("objective", vals[1])
-        # pyre-fixme[16]: Optional type has no attribute `__getitem__`.
-        self.assertIn("objective", vals[1]["objective"])
+        self.assertIn("objective", none_throws(vals[1]))
+        self.assertIn("objective", none_throws(vals[1])["objective"])
 
     def test_custom_gs(self) -> None:
         """Managed loop with custom generation strategy"""
@@ -432,18 +415,14 @@ def test_optimize_graceful_exit_on_exception(self) -> None:
         self.assertIn("x2", best)
         self.assertIsNotNone(vals)
         self.assertIn("objective", vals[0])
-        # pyre-fixme[6]: For 2nd param expected `Union[Container[typing.Any],
-        #  Iterable[typing.Any]]` but got `Optional[Dict[str, Dict[str, float]]]`.
-        self.assertIn("objective", vals[1])
-        # pyre-fixme[16]: Optional type has no attribute `__getitem__`.
-        self.assertIn("objective", vals[1]["objective"])
+        self.assertIn("objective", none_throws(vals[1]))
+        self.assertIn("objective", none_throws(vals[1])["objective"])
 
     @patch(
         "ax.core.experiment.Experiment.new_trial",
         side_effect=RuntimeError("cholesky_cpu error - bad matrix"),
     )
-    # pyre-fixme[3]: Return type must be annotated.
-    def test_annotate_exception(self, _):
+    def test_annotate_exception(self, _: Mock) -> None:
         strategy0 = GenerationStrategy(
             name="Sobol",
             steps=[GenerationStep(generator=Generators.SOBOL, num_trials=-1)],

ax/service/tests/test_managed_loop.py

@@ -9,7 +9,9 @@
 import itertools
 import logging
 from collections import namedtuple
+from collections.abc import Callable
 from logging import DEBUG, INFO, WARN
+from typing import Any
 from unittest import mock
 from unittest.mock import patch
 
@@ -265,8 +267,7 @@ def test_exp_to_df(self) -> None:
         with patch.object(Experiment, "lookup_data", lambda self: mock_results):
             df = exp_to_df(exp=exp)
         # all but two rows should have a metric value of NaN
-        # pyre-fixme[16]: `bool` has no attribute `sum`.
-        self.assertEqual(pd.isna(df[OBJECTIVE_NAME]).sum(), len(df.index) - 2)
+        self.assertEqual(df[OBJECTIVE_NAME].isna().sum(), len(df.index) - 2)
 
         # an experiment with more results than arms raises an error
         with (
@@ -369,16 +370,16 @@ def test_get_standard_plots(self) -> None:
                 self.assertTrue(all(isinstance(plot, go.Figure) for plot in plots))
 
         # Raise an exception in one plot and make sure we generate the others
-        for plot_function, num_expected_plots in [
-            [_get_curve_plot_dropdown, 8],  # Not used
-            [_get_objective_trace_plot, 6],
-            [_objective_vs_true_objective_scatter, 7],
-            [_get_objective_v_param_plots, 6],
-            [_get_cross_validation_plots, 7],
-            [plot_feature_importance_by_feature_plotly, 6],
-        ]:
+        plot_test_cases: list[tuple[Callable[..., Any], int]] = [
+            (_get_curve_plot_dropdown, 8),  # Not used
+            (_get_objective_trace_plot, 6),
+            (_objective_vs_true_objective_scatter, 7),
+            (_get_objective_v_param_plots, 6),
+            (_get_cross_validation_plots, 7),
+            (plot_feature_importance_by_feature_plotly, 6),
+        ]
+        for plot_function, num_expected_plots in plot_test_cases:
             with mock.patch(
-                # pyre-ignore
                 f"ax.service.utils.report_utils.{plot_function.__name__}",
                 side_effect=Exception(),
             ):