Fix minor bugs across api, benchmark, generators, metrics, runners, and other modules (#4899)

Summary:

Fix bugs in smaller modules with 1-2 files each:

- Client JSON snapshot crash when no GenerationStrategy exists: `"generation_strategy" in snapshot` is True for the key with a None value, causing a decode crash. Changed to `snapshot.get("generation_strategy") is not None` (api/client.py:1351)
- Benchmark ternary discarding input in both branches: `{} if target_fidelity_and_task is None else {}` always produces an empty dict regardless of input (benchmark/benchmark.py:701)
- String `"False"` default (truthy) preventing trials from being marked CANDIDATE: `t.run_metadata.get(STARTED_KEY, "False")` returns the string `"False"` which is truthy, so `not "False"` is always False (runners/map_replay.py:44)
- Operator precedence making `noisy=False` set all means to 0.0: `item["mean"] + noise if noisy else 0.0` parses as `(item["mean"] + noise) if noisy else 0.0`, replacing the entire mean with 0.0 when `noisy=False` (metrics/noisy_function_map.py:86, metrics/branin_map.py:119)
- Misleading "total number of trials" vs actual trial index (global_stopping/strategies/improvement.py:132) + test update
- Typos: "trial rial" → "trial" (fb/tutorials/auto_tuning.py), "prefererd" → "preferred" and stray `f` in f-string (fb/utils/preference/preference.py), "acquistion" → "acquisition" (generators/torch/botorch_modular/utils.py + test)

Reviewed By: ItsMrLin

Differential Revision: D92879402

Author: saitcakmak

ax/api/client.py

@@ -1338,7 +1338,7 @@ def _from_json_snapshot(
                 decoder_registry=decoder_registry,
                 class_decoder_registry=class_decoder_registry,
             )
-            if "generation_strategy" in snapshot
+            if snapshot.get("generation_strategy") is not None
             else None
         )
 

ax/benchmark/benchmark.py

@@ -701,7 +701,9 @@ def compute_baseline_value_from_sobol(
         n_repeats: Number of times to repeat the five Sobol trials.
     """
     method = get_sobol_benchmark_method()
-    target_fidelity_and_task = {} if target_fidelity_and_task is None else {}
+    target_fidelity_and_task = (
+        {} if target_fidelity_and_task is None else target_fidelity_and_task
+    )
 
     # set up a dummy problem so we can use `benchmark_replication`
     # MOO problems are always higher-is-better because they use hypervolume

ax/early_stopping/strategies/base.py

@@ -8,7 +8,7 @@
 
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import Iterable, Sequence
+from collections.abc import Iterable
 from logging import Logger
 from typing import cast
 
@@ -665,19 +665,3 @@ def _lookup_and_validate_data(
             full_df = full_df[full_df[MAP_KEY] >= self.min_progression_modeling]
             map_data = Data(df=full_df)
         return map_data
-
-    def get_training_data(
-        self,
-        experiment: Experiment,
-        map_data: Data,
-        max_training_size: int | None = None,
-        outcomes: Sequence[str] | None = None,
-        parameters: list[str] | None = None,
-    ) -> None:
-        # Deprecated in Ax 1.1.0, so should be removed in Ax 1.2.0+.
-        raise DeprecationWarning(
-            "`ModelBasedEarlyStoppingStrategy.get_training_data` is deprecated. "
-            "Subclasses should either extract the training data manually, "
-            "or rely on the fitted surrogates available in the current generation "
-            "node that is passed into `should_stop_trials_early`."
-        )

ax/generators/torch/botorch_modular/utils.py

@@ -539,7 +539,7 @@ def construct_acquisition_and_optimizer_options(
             if len(botorch_acqf_classes_with_options) > 1:
                 warnings.warn(
                     message="botorch_acqf_options are being ignored, due to using "
-                    "MultiAcquisition. Specify options for each acquistion function"
+                    "MultiAcquisition. Specify options for each acquisition function "
                     "via botorch_acqf_classes_with_options.",
                     category=AxWarning,
                     stacklevel=4,

ax/generators/torch/tests/test_utils.py

@@ -693,7 +693,7 @@ def test_construct_acquisition_and_optimizer_options(self) -> None:
             self.assertEqual(
                 str(warning.message),
                 "botorch_acqf_options are being ignored, due to using "
-                "MultiAcquisition. Specify options for each acquistion function"
+                "MultiAcquisition. Specify options for each acquisition function "
                 "via botorch_acqf_classes_with_options.",
             )
 

ax/global_stopping/strategies/improvement.py

@@ -130,8 +130,8 @@ def _should_stop_optimization(
             trial_to_check = max_completed_trial
         elif trial_to_check > max_completed_trial:
             raise ValueError(
-                "trial_to_check is larger than the total number of "
-                f"trials (={max_completed_trial})."
+                "trial_to_check is larger than the maximum completed "
+                f"trial index (={max_completed_trial})."
             )
 
         # Only counting the trials up to trial_to_check.

ax/global_stopping/tests/test_strategies.py

@@ -85,7 +85,7 @@ def test_base_cases(self) -> None:
         # Should raise ValueError if trying to check an invalid trial
         with self.assertRaisesRegex(
             ValueError,
-            r"trial_to_check is larger than the total number of trials \(=4\).",
+            r"trial_to_check is larger than the maximum completed trial index \(=4\).",
         ):
             stop, message = gss.should_stop_optimization(
                 experiment=exp, trial_to_check=5

ax/metrics/branin_map.py

@@ -117,9 +117,8 @@ def fetch_trial_data(
                         "sem": self.noise_sd if noisy else 0.0,
                         "trial_index": trial.index,
                         "mean": [
-                            item["mean"] + self.noise_sd * np.random.randn()
-                            if noisy
-                            else 0.0
+                            item["mean"]
+                            + (self.noise_sd * np.random.randn() if noisy else 0.0)
                             for item in res
                         ],
                         "metric_signature": self.signature,

ax/metrics/noisy_function_map.py

@@ -84,9 +84,8 @@ def fetch_trial_data(
                     "sem": self.noise_sd if noisy else 0.0,
                     "trial_index": trial.index,
                     "mean": [
-                        item["mean"] + self.noise_sd * np.random.randn()
-                        if noisy
-                        else 0.0
+                        item["mean"]
+                        + (self.noise_sd * np.random.randn() if noisy else 0.0)
                         for item in res
                     ],
                     "metric_signature": self.signature,

ax/runners/map_replay.py

@@ -41,7 +41,7 @@ def poll_trial_status(
         # depending on whether or not there is more data available,
         # mark it either RUNNING or COMPLETED.
         for t in trials:
-            if not t.run_metadata.get(STARTED_KEY, "False"):
+            if not t.run_metadata.get(STARTED_KEY, False):
                 result[TrialStatus.CANDIDATE].add(t.index)
             elif not self.replay_metric.has_trial_data(t.index):
                 result[TrialStatus.ABANDONED].add(t.index)