Remove pyre-fixme/pyre-ignore from ax/core/ source files (#4976)

Summary:
Pull Request resolved: #4976

Remove ~29 pyre-fixme/pyre-ignore suppression comments from 11 source files
in ax/core/ by applying proper type fixes:

- Use `none_throws()` for Optional unwrapping
- Use `cast()` for type narrowing
- Add proper type annotations (e.g., `partial[Any]`, `type[Any]`)
- Refactor pandas `itertuples()` access to use `Any`-typed row variable
- Make `Parameter.clone()` abstract with `abstractmethod`
- Use `float()` casts for numeric comparisons
- Remove explicit `: int` from enum members (pyre-fixme[35])

24 genuinely unfixable suppressions remain (documented):
- `copy_doc` with `property` interaction
- Property setter decorator type inference
- `np.floating`/`np.integer` generic params (runtime isinstance limitation)
- Intentional inconsistent overrides in OptimizationConfig subclasses
- Abstract class instantiation via `type(self)`

Reviewed By: dme65

Differential Revision: D95264859

fbshipit-source-id: 77326c274feb662a6e2c57be9347ec09ae005c0e

Author: saitcakmak

ax/core/base_trial.py

@@ -38,15 +38,12 @@
 MAX_ABANDONED_REASON_LENGTH = 1000
 
 
-def immutable_once_run(func: Callable) -> Callable:
+def immutable_once_run(func: Callable[..., Any]) -> Callable[..., Any]:
     """Decorator for methods that should throw Error when
     trial is running or has ever run and immutable.
     """
 
-    # no type annotation for now; breaks sphinx-autodoc-typehints
-    # pyre-fixme[3]: Return type must be annotated.
-    # pyre-fixme[2]: Parameter must be annotated.
-    def _immutable_once_run(self, *args, **kwargs):
+    def _immutable_once_run(self: Any, *args: Any, **kwargs: Any) -> Any:
         if self._status != TrialStatus.CANDIDATE:
             raise TrialMutationError(
                 "Cannot modify a trial that is running or has ever run. "

ax/core/batch_trial.py

@@ -396,7 +396,7 @@ def in_design_arms(self) -> list[Arm]:
             if self.experiment.search_space.check_membership(arm.parameters)
         ]
 
-    # pyre-ignore[6]: T77111662.
+    # pyre-ignore[6]: pyre does not understand @copy_doc with @property.
     @copy_doc(BaseTrial.generator_runs)
     @property
     def generator_runs(self) -> list[GeneratorRun]:

ax/core/data.py

@@ -16,7 +16,7 @@
 from functools import cached_property
 from io import StringIO
 from logging import Logger
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -177,29 +177,24 @@ def __init__(
                     f"Dataframe must contain required columns {list(missing_columns)}."
                 )
 
+            # Using itertuples() instead of iterrows() for speed.
+            # cast() suppresses pyre errors on namedtuple attribute access.
             self._data_rows = [
                 DataRow(
-                    # pyre-ignore[16] Intentional unsafe namedtuple access
                     # int() cast needed because pd.read_json with dtype=False
                     # can return string trial indices from storage
                     trial_index=int(row.trial_index),
-                    # pyre-ignore[16] Intentional unsafe namedtuple access
                     arm_name=row.arm_name,
-                    # pyre-ignore[16] Intentional unsafe namedtuple access
                     metric_name=row.metric_name,
-                    # pyre-ignore[16] Intentional unsafe namedtuple access
                     metric_signature=row.metric_signature,
-                    # pyre-ignore[16] Intentional unsafe namedtuple access
                     mean=row.mean,
-                    # pyre-ignore[16] Intentional unsafe namedtuple access
                     se=row.sem,
                     step=getattr(row, "step", None),
                     start_time=getattr(row, "start_time", None),
                     end_time=getattr(row, "end_time", None),
                     n=getattr(row, "n", None),
                 )
-                # Using itertuples() instead of iterrows() for speed
-                for row in df.itertuples(index=False)
+                for row in cast(Any, df.itertuples(index=False))
             ]
         else:
             self._data_rows = []

ax/core/evaluations_to_data.py

@@ -8,6 +8,7 @@
 
 from collections.abc import Mapping
 from enum import Enum
+from typing import cast
 
 from ax.core.data import Data, DataRow
 from ax.core.types import FloatLike, SingleMetricData, TEvaluationOutcome
@@ -117,9 +118,11 @@ def raw_evaluations_to_data(
                     "multiple metrics."
                 )
             metric_name = next(iter(metric_name_to_signature.keys()))
-            # pyre-fixme[6]: Incmopatible parameter type (Pyre doesn't know that
-            # this is in fact a SingleMetricData)
-            mean, sem = _validate_and_extract_single_metric_data(dat=evaluation)
+            # After eliminating dict and list cases above, evaluation is
+            # SingleMetricData, but pyre can't narrow union type aliases.
+            mean, sem = _validate_and_extract_single_metric_data(
+                dat=cast(SingleMetricData, evaluation)
+            )
             metric_names_seen.add(metric_name)
             data_rows.append(
                 DataRow(

ax/core/experiment.py

@@ -72,8 +72,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,
 )
@@ -135,9 +134,7 @@ def __init__(
                 DeprecationWarning,
                 stacklevel=2,
             )
-        # appease pyre
-        # pyre-fixme[13]: Attribute `_search_space` is never initialized.
-        self._search_space: SearchSpace
+        self._search_space: SearchSpace = search_space
         self._status_quo: Arm | None = None
 
         self._name = name
@@ -194,7 +191,6 @@ def __init__(
         self.add_tracking_metrics(tracking_metrics or [])
 
         # call setters defined below
-        self.search_space: SearchSpace = search_space
         self.status_quo = status_quo
         if optimization_config is not None:
             self.optimization_config = optimization_config

ax/core/metric.py

@@ -15,7 +15,7 @@
 from datetime import timedelta
 from functools import reduce
 from logging import Logger
-from typing import Any, TYPE_CHECKING
+from typing import Any, cast, TYPE_CHECKING
 
 from ax.core.data import Data
 from ax.utils.common.base import SortableBase
@@ -461,18 +461,16 @@ def maybe_raise_deprecation_warning_on_class_methods(self) -> None:
         # This is a temporary hack to allow us to deprecate old metric class method
         # implementations. There does not seem to be another way of checking whether
         # base class' classmethods are overridden in subclasses.
-        is_fetch_trial_data_multi_overriden = (
-            getattr(self.__class__.fetch_trial_data_multi, "__code__", "DEFAULT")
-            != Metric.fetch_trial_data_multi.__code__  # pyre-ignore[16]
-        )
-        is_fetch_experiment_data_multi_overriden = (
-            getattr(
-                self.__class__.fetch_experiment_data_multi,
-                "__code__",
-                "DEFAULT",
-            )
-            != Metric.fetch_experiment_data_multi.__code__  # pyre-ignore[16]
-        )
+        # Asymmetric defaults: if __code__ is missing on the subclass, "DEFAULT"
+        # != None is True, treating the method as overridden (safe default).
+        is_fetch_trial_data_multi_overriden = getattr(
+            self.__class__.fetch_trial_data_multi, "__code__", "DEFAULT"
+        ) != getattr(Metric.fetch_trial_data_multi, "__code__", None)
+        is_fetch_experiment_data_multi_overriden = getattr(
+            self.__class__.fetch_experiment_data_multi,
+            "__code__",
+            "DEFAULT",
+        ) != getattr(Metric.fetch_experiment_data_multi, "__code__", None)
         # Raise deprecation warning if this method from the base class is used (meaning
         # that it is not overridden and the classmethod is overridden instead), unless
         # the only overridden method is `fetch_trial_data` (in which case the setup is
@@ -623,16 +621,18 @@ def _wrap_trial_data_multi(cls, data: Data) -> dict[str, MetricFetchResult]:
     def _wrap_experiment_data_multi(
         cls, data: Data
     ) -> dict[int, dict[str, MetricFetchResult]]:
-        # pyre-fixme[7]
-        return {
-            trial_index: {
-                metric_signature: Ok(
-                    value=data.filter(
-                        trial_indices=[trial_index],
-                        metric_signatures=[metric_signature],
+        return cast(
+            dict[int, dict[str, MetricFetchResult]],
+            {
+                trial_index: {
+                    metric_signature: Ok(
+                        value=data.filter(
+                            trial_indices=[trial_index],
+                            metric_signatures=[metric_signature],
+                        )
                     )
-                )
-                for metric_signature in data.full_df["metric_signature"]
-            }
-            for trial_index in data.full_df["trial_index"]
-        }
+                    for metric_signature in data.full_df["metric_signature"]
+                }
+                for trial_index in data.full_df["trial_index"]
+            },
+        )

ax/core/multi_type_experiment.py

@@ -109,11 +109,12 @@ def add_trial_type(self, trial_type: str, runner: Runner) -> Self:
         self._trial_type_to_runner[trial_type] = runner
         return self
 
-    # pyre-fixme [56]: Pyre was not able to infer the type of the decorator
-    # `Experiment.optimization_config.setter`.
+    # pyre does not support inferring the type of property setter decorators
+    # or the `.fset` attribute on properties.
+    # pyre-fixme[56]: Pyre was not able to infer the type of the decorator.
     @Experiment.optimization_config.setter
     def optimization_config(self, optimization_config: OptimizationConfig) -> None:
-        # pyre-fixme [16]: `Optional` has no attribute `fset`.
+        # pyre-fixme[16]: `Optional` has no attribute `fset`.
         Experiment.optimization_config.fset(self, optimization_config)
         for metric_name in optimization_config.metrics.keys():
             # Optimization config metrics are required to be the default trial type

ax/core/parameter.py

@@ -55,14 +55,10 @@
 
 
 class ParameterType(Enum):
-    # pyre-fixme[35]: Target cannot be annotated.
-    BOOL: int = 0
-    # pyre-fixme[35]: Target cannot be annotated.
-    INT: int = 1
-    # pyre-fixme[35]: Target cannot be annotated.
-    FLOAT: int = 2
-    # pyre-fixme[35]: Target cannot be annotated.
-    STRING: int = 3
+    BOOL = 0
+    INT = 1
+    FLOAT = 2
+    STRING = 3
 
     @property
     def is_numeric(self) -> bool:
@@ -71,10 +67,7 @@ def is_numeric(self) -> bool:
 
 TParameterType = Union[type[int], type[float], type[str], type[bool]]
 
-# pyre: PARAMETER_PYTHON_TYPE_MAP is declared to have type
-# pyre: `Dict[ParameterType, Union[Type[bool], Type[float], Type[int],
-# pyre: Type[str]]]` but is used as type `Dict[ParameterType,
-# pyre-fixme[9]: Type[Union[float, str]]]`.
+# pyre-fixme[9]: Pyre collapses individual type[] values into Type[Union[...]].
 PARAMETER_PYTHON_TYPE_MAP: dict[ParameterType, TParameterType] = {
     ParameterType.INT: int,
     ParameterType.FLOAT: float,
@@ -87,9 +80,7 @@ def is_numeric(self) -> bool:
 ] = tuple(PARAMETER_PYTHON_TYPE_MAP.values())
 
 
-# pyre-fixme[24]: Generic type `type` expects 1 type parameter, use `typing.Type` to
-#  avoid runtime subscripting errors.
-def _get_parameter_type(python_type: type) -> ParameterType:
+def _get_parameter_type(python_type: type[Any]) -> ParameterType:
     """Given a Python type, retrieve corresponding Ax ``ParameterType``."""
     for param_type, py_type in PARAMETER_PYTHON_TYPE_MAP.items():
         if py_type == python_type:
@@ -236,9 +227,8 @@ def dependents(self) -> dict[TParamValue, list[str]]:
             "Only fixed and choice hierarchical parameters are currently supported."
         )
 
-    # pyre-fixme[7]: Expected `Parameter` but got implicit return value of `None`.
-    def clone(self) -> Self:
-        pass
+    @abstractmethod
+    def clone(self) -> Self: ...
 
     def disable(self, default_value: TParamValue) -> None:
         """
@@ -541,9 +531,7 @@ def set_digits(self, digits: int | None) -> RangeParameter:
         # Re-scale min and max to new digits definition
         cast_lower = self.cast(self._lower)
         cast_upper = self.cast(self._upper)
-        # `<=` is not supported for operand types `Union[float, int]` and `int`.
-        # pyre-ignore [58]
-        if cast_lower >= cast_upper:
+        if float(cast_lower) >= float(cast_upper):
             raise UserInputError(
                 f"Lower bound {cast_lower} is >= upper bound {cast_upper}."
             )
@@ -1304,13 +1292,7 @@ class DerivedParameter(Parameter):
     extendable to non-linear functions.
     """
 
-    # pyre-fixme [13]: Uninitialized attribute [13]: Attribute `_intercept` is
-    # declared in class `DerivedParameter` to have type `float` but is never
-    # initialized.
     _intercept: float
-    # pyre-fixme [13]: Uninitialized attribute [13]: Attribute
-    # `_parameter_names_to_weights` is declared in class `DerivedParameter` to#
-    # have type `typing.Dict[str, float]` but is never initialized.
     _parameter_names_to_weights: dict[str, float]
 
     def __init__(
@@ -1343,6 +1325,8 @@ def __init__(
         self._parameter_type = parameter_type  # Set first so validation works
         self._is_fidelity = is_fidelity
         self._target_value = target_value
+        self._intercept = 0.0
+        self._parameter_names_to_weights = {}
 
         # Parse expression and validate type constraint (reuses set_expression_str)
         self.set_expression_str(expression_str)

ax/core/search_space.py

@@ -13,6 +13,7 @@
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from logging import Logger
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -1079,7 +1080,7 @@ def cast_observation_features(
         parameter should not be in the arm within the search space due to its
         hierarchical structure.
         """
-        full_parameterization_md = {
+        full_parameterization_md: dict[str, Any] = {
             Keys.FULL_PARAMETERIZATION: observation_features.parameters.copy()
         }
         obs_feats = observation_features.clone(
@@ -1089,7 +1090,7 @@ def cast_observation_features(
             )
         )
         if not obs_feats.metadata:
-            obs_feats.metadata = full_parameterization_md  # pyre-ignore[8]
+            obs_feats.metadata = full_parameterization_md
         else:
             obs_feats.metadata = {**obs_feats.metadata, **full_parameterization_md}
 

ax/core/trial.py

@@ -25,8 +25,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,
 )
@@ -115,19 +114,15 @@ def add_arm(
             The trial instance.
         """
 
+        cand_metadata_by_sig: dict[str, TCandidateMetadata] | None = None
+        if candidate_metadata is not None:
+            cand_metadata: TCandidateMetadata = candidate_metadata.copy()
+            cand_metadata_by_sig = {arm.signature: cand_metadata}
         return self.add_generator_run(
             generator_run=GeneratorRun(
                 arms=[arm],
                 type=GeneratorRunType.MANUAL.name,
-                # pyre-ignore[6]: In call `GeneratorRun.__init__`, for 3rd parameter
-                # `candidate_metadata_by_arm_signature`
-                # expected `Optional[Dict[str, Optional[Dict[str, typing.Any]]]]`
-                # but got `Optional[Dict[str, Dict[str, typing.Any]]]`
-                candidate_metadata_by_arm_signature=(
-                    None
-                    if candidate_metadata is None
-                    else {arm.signature: candidate_metadata.copy()}
-                ),
+                candidate_metadata_by_arm_signature=cand_metadata_by_sig,
             )
         )