Deprecate deduplicate setting for SOBOL Generator; Set GenerationNode.should_deduplicate=True by default

Summary:
Previously, we had dedup settings on both `Generator` and `GenerationNode` level, which is unnecessary and error-prone for accidentally using 2 different dedup settings at the 2 levels (see D92094386)

Given that only `SobolGenerator` has special dedup logic, we've decided to deprecate `deduplicate` setting in `Generator` and only rely on `GenerationNode` level deduplication.

Furthermore, since we generally expect dedup to be True for OSS use cases and AutoML use cases, we set `GenerationNode.should_deduplicate=True` by default while explicitly overwriting GNode dedup to False for online use cases.


**One Potential Issue**
Currently `RandomAdapter._gen` solely relies on SOBOL generator to dedup (by passing in `generated_points` to sobol generator and then `rejection_sample`- https://fburl.com/code/zmqfhn8g) and if we remove that, anyone that's using standalone `RandomAdapter` (without a GNode) may get duplicated trials if they call `RandomAdapter.gen` multiple times. This is partially mitigated by advancing `init_position` on sobol generator -- there shouldn't be duplicated points for continuous parameters, but if there are discrete parameters we can be producing duplicated points after rounding. The recommended path is through GenerationStrategy/GenerationNode, where dedup is properly handled. But there's no enforcement on never using standalone RandomAdapter. (ps. `TorchAdapter` passes pending_observations through to the BoTorch acquisition function as X_pending so it doesn't have the same problem)

Differential Revision: D92884352

Author: Sunny Shen

ax/adapter/factory.py

@@ -34,7 +34,6 @@
 def get_sobol(
     search_space: SearchSpace,
     seed: int | None = None,
-    deduplicate: bool = False,
     init_position: int = 0,
     scramble: bool = True,
     fallback_to_sample_polytope: bool = False,
@@ -52,7 +51,6 @@ def get_sobol(
         Generators.SOBOL(
             experiment=Experiment(search_space=search_space),
             seed=seed,
-            deduplicate=deduplicate,
             init_position=init_position,
             scramble=scramble,
             fallback_to_sample_polytope=fallback_to_sample_polytope,

ax/adapter/random.py

@@ -9,7 +9,6 @@
 
 from collections.abc import Mapping, Sequence
 
-import numpy as np
 from ax.adapter.adapter_utils import (
     extract_parameter_constraints,
     extract_search_space_digest,
@@ -93,35 +92,6 @@ def _gen(
         linear_constraints = extract_parameter_constraints(
             search_space.parameter_constraints, self.parameters
         )
-        # Extract generated points to deduplicate against.
-        # Exclude out-of-design arms (which can only be manual arms
-        # instead of adapter-generated arms).
-        generated_points = None
-        if self.generator.deduplicate:
-            arms_to_deduplicate = self._experiment.arms_by_signature_for_deduplication
-            generated_obs = [
-                ObservationFeatures.from_arm(arm=arm)
-                for arm in arms_to_deduplicate.values()
-                if self._search_space.check_membership(parameterization=arm.parameters)
-            ]
-            # Transform
-            for t in self.transforms.values():
-                generated_obs = t.transform_observation_features(generated_obs)
-            # Add pending observations -- already transformed.
-            generated_obs.extend(
-                [obs for obs_list in pending_observations.values() for obs in obs_list]
-            )
-            if len(generated_obs) > 0:
-                # Extract generated points array (n x d).
-                generated_points = np.array(
-                    [
-                        [obs.parameters[p] for p in self.parameters]
-                        for obs in generated_obs
-                    ]
-                )
-                # Take unique points only, since there may be duplicates coming
-                # from pending observations for different metrics.
-                generated_points = np.unique(generated_points, axis=0)
 
         # Generate the candidates
         X, w = self.generator.gen(
@@ -131,7 +101,6 @@ def _gen(
             fixed_features=fixed_features_dict,
             model_gen_options=model_gen_options,
             rounding_func=transform_callback(self.parameters, self.transforms),
-            generated_points=generated_points,
         )
         observation_features = parse_observation_features(X, self.parameters)
         return GenResults(

ax/adapter/registry.py

@@ -173,6 +173,9 @@ class GeneratorSetup(NamedTuple):
     default_generator_kwargs: Mapping[str, Any] | None = None
     standard_adapter_kwargs: Mapping[str, Any] | None = None
     not_saved_generator_kwargs: Sequence[str] | None = None
+    # Kwargs that were removed from the generator but may still exist in
+    # serialized data. These are silently filtered out before validation.
+    deprecated_generator_kwargs: Sequence[str] | None = None
 
 
 """A mapping of string keys that indicate a generator, to the corresponding
@@ -209,6 +212,8 @@ class GeneratorSetup(NamedTuple):
         adapter_class=RandomAdapter,
         generator_class=SobolGenerator,
         transforms=Cont_X_trans,
+        # These kwargs were removed but may exist in old serialized data.
+        deprecated_generator_kwargs=["deduplicate"],
     ),
     "Uniform": GeneratorSetup(
         adapter_class=RandomAdapter,
@@ -291,6 +296,11 @@ def __call__(
         adapter_class = model_setup_info.adapter_class
         search_space = experiment.search_space
 
+        # Filter out deprecated kwargs that may exist in old serialized data.
+        if model_setup_info.deprecated_generator_kwargs:
+            for deprecated_key in model_setup_info.deprecated_generator_kwargs:
+                kwargs.pop(deprecated_key, None)
+
         if not silently_filter_kwargs:
             # Check correct kwargs are present
             callables = (generator_class, adapter_class)

ax/adapter/tests/test_random_adapter.py

@@ -6,7 +6,6 @@
 
 # pyre-strict
 
-import dataclasses
 from unittest import mock
 
 import numpy as np
@@ -20,15 +19,10 @@
 from ax.core.parameter import ParameterType, RangeParameter
 from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.search_space import SearchSpace
-from ax.exceptions.core import SearchSpaceExhausted
 from ax.generators.random.base import RandomGenerator
 from ax.generators.random.sobol import SobolGenerator
 from ax.utils.common.testutils import TestCase
-from ax.utils.testing.core_stubs import (
-    get_data,
-    get_search_space_for_range_values,
-    get_small_discrete_search_space,
-)
+from ax.utils.testing.core_stubs import get_data
 from ax.utils.testing.modeling_stubs import get_experiment_for_value
 
 
@@ -147,22 +141,6 @@ def test_gen_simple(self) -> None:
         self.assertIsNone(gen_args["linear_constraints"])
         self.assertIsNone(gen_args["fixed_features"])
 
-    def test_deduplicate(self) -> None:
-        exp = Experiment(search_space=get_small_discrete_search_space())
-        sobol = RandomAdapter(
-            experiment=exp,
-            generator=SobolGenerator(deduplicate=True),
-            transforms=Cont_X_trans,
-        )
-        for _ in range(4):  # Search space is {[0, 1], {"red", "panda"}}
-            # Generate & attach trials to the experiment so that the
-            # generated points are used for deduplication.
-            gr = sobol.gen(1)
-            exp.new_trial(generator_run=gr).mark_running(no_runner_required=True)
-            self.assertEqual(len(gr.arms), 1)
-        with self.assertRaises(SearchSpaceExhausted):
-            sobol.gen(1)
-
     def test_search_space_not_expanded(self) -> None:
         data = get_data(num_non_sq_arms=0)
         sq_arm = Arm(name="status_quo", parameters={"x": 10.0, "y": 1.0, "z": 1.0})
@@ -186,124 +164,6 @@ def test_search_space_not_expanded(self) -> None:
         sobol.gen(1)
         self.assertEqual(sobol._model_space, sobol._search_space)
 
-    def test_generated_points(self) -> None:
-        # Checks for generated points argument passed to Generator.gen.
-        # Search space has two range parameters in [0, 5].
-        exp = Experiment(
-            search_space=get_search_space_for_range_values(min=0.0, max=5.0)
-        )
-        ssd = extract_search_space_digest(
-            search_space=exp.search_space,
-            param_names=list(exp.search_space.parameters.keys()),
-        )
-        ssd = dataclasses.replace(ssd, bounds=[(0.0, 1.0), (0.0, 1.0)])
-        generator = SobolGenerator(deduplicate=True)
-        gen_res = generator.gen(n=1, search_space_digest=ssd, rounding_func=lambda x: x)
-        # Using Cont_X_trans, particularly UnitX here to test transform application.
-        adapter = RandomAdapter(
-            experiment=exp, generator=generator, transforms=Cont_X_trans
-        )
-
-        # No pending points or previous trials on the experiment.
-        with mock.patch.object(generator, "gen", return_value=gen_res) as mock_gen:
-            adapter.gen(n=1)
-        self.assertIsNone(mock_gen.call_args.kwargs["generated_points"])
-
-        # Attach two trials to the experiment.
-        exp.new_trial().add_arm(Arm(parameters={"x": 0.0, "y": 0.0})).mark_running(
-            no_runner_required=True
-        )
-        exp.new_trial().add_arm(Arm(parameters={"x": 2.0, "y": 2.0})).mark_running(
-            no_runner_required=True
-        )
-        with mock.patch.object(generator, "gen", return_value=gen_res) as mock_gen:
-            adapter.gen(n=1)
-        self.assertEqual(
-            mock_gen.call_args.kwargs["generated_points"].tolist(),
-            [[0.0, 0.0], [0.4, 0.4]],
-        )
-
-        # Add pending points -- only unique ones should be passed down.
-        pending_observations = {
-            m: [ObservationFeatures(parameters={"x": 3.0, "y": 3.0})]
-            for m in ("m1", "m2")
-        }
-        with mock.patch.object(generator, "gen", return_value=gen_res) as mock_gen:
-            adapter.gen(n=1, pending_observations=pending_observations)
-        self.assertEqual(
-            mock_gen.call_args.kwargs["generated_points"].tolist(),
-            [[0.0, 0.0], [0.4, 0.4], [0.6, 0.6]],
-        )
-
-        # Turn off deduplicate, nothing should be passed down.
-        generator.deduplicate = False
-        with mock.patch.object(generator, "gen", return_value=gen_res) as mock_gen:
-            adapter.gen(n=1, pending_observations=pending_observations)
-        self.assertIsNone(mock_gen.call_args.kwargs["generated_points"])
-
-        # Test filtering out-of-design arms during deduplication
-        # Create experiment with in-design and out-of-design arms
-        exp_with_ood_arms = Experiment(
-            search_space=get_search_space_for_range_values(min=0.0, max=5.0)
-        )
-        in_design_arm = Arm(
-            name="in_design", parameters={"x": 2.0, "y": 3.0}
-        )  # Within [0, 5]
-        out_of_design_arm = Arm(
-            name="out_of_design", parameters={"x": 6.0, "y": 7.0}
-        )  # Outside [0, 5]
-
-        exp_with_ood_arms.new_trial().add_arm(in_design_arm).mark_running(
-            no_runner_required=True
-        )
-        exp_with_ood_arms.new_trial().add_arm(out_of_design_arm).mark_running(
-            no_runner_required=True
-        )
-
-        generator = SobolGenerator(deduplicate=True)
-        adapter_mixed = RandomAdapter(
-            experiment=exp_with_ood_arms,
-            generator=generator,
-            transforms=Cont_X_trans,
-        )
-
-        # Only the in-design arm should be included in generated_points
-        with mock.patch.object(generator, "gen", return_value=gen_res) as mock_gen:
-            adapter_mixed.gen(n=1)
-
-        generated_points = mock_gen.call_args.kwargs["generated_points"]
-        self.assertEqual(len(generated_points), 1)
-
-        # Test case where all arms are out-of-design
-        exp_all_out_of_design = Experiment(
-            search_space=get_search_space_for_range_values(min=0.0, max=5.0)
-        )
-        out_of_design_arm1 = Arm(name="out1", parameters={"x": 6.0, "y": 7.0})
-        out_of_design_arm2 = Arm(name="out2", parameters={"x": -1.0, "y": 8.0})
-
-        exp_all_out_of_design.new_trial().add_arm(out_of_design_arm1).mark_running(
-            no_runner_required=True
-        )
-        exp_all_out_of_design.new_trial().add_arm(out_of_design_arm2).mark_running(
-            no_runner_required=True
-        )
-
-        generator_all_out = SobolGenerator(deduplicate=True)
-        adapter_all_out = RandomAdapter(
-            experiment=exp_all_out_of_design,
-            generator=generator_all_out,
-            transforms=Cont_X_trans,
-        )
-
-        # When all arms are out-of-design, generated_points should be empty
-        with mock.patch.object(
-            generator_all_out, "gen", return_value=gen_res
-        ) as mock_gen:
-            adapter_all_out.gen(n=1)
-
-        generated_points_all_out = mock_gen.call_args.kwargs["generated_points"]
-        self.assertIsNone(generated_points_all_out)
-
     def test_generation_with_all_fixed(self) -> None:
         # Make sure candidate generation succeeds and returns correct parameters
         # when all parameters are fixed.

ax/adapter/tests/test_registry.py

@@ -174,7 +174,6 @@ def test_view_defaults(self) -> None:
             (
                 {
                     "seed": None,
-                    "deduplicate": True,
                     "init_position": 0,
                     "scramble": True,
                     "generated_points": None,
@@ -194,7 +193,7 @@ def test_view_defaults(self) -> None:
         self.assertTrue(
             all(
                 kw in Generators.SOBOL.view_kwargs()[0]
-                for kw in ["seed", "deduplicate", "init_position", "scramble"]
+                for kw in ["seed", "init_position", "scramble"]
             ),
             all(
                 kw in Generators.SOBOL.view_kwargs()[1]

ax/benchmark/tests/test_benchmark.py

@@ -934,11 +934,7 @@ def test_replication_with_generation_node(self) -> None:
                 nodes=[
                     GenerationNode(
                         name="Sobol",
-                        generator_specs=[
-                            GeneratorSpec(
-                                Generators.SOBOL, generator_kwargs={"deduplicate": True}
-                            )
-                        ],
+                        generator_specs=[GeneratorSpec(Generators.SOBOL)],
                     )
                 ]
             ),

ax/generation_strategy/dispatch_utils.py

@@ -51,7 +51,7 @@ def _make_sobol_step(
     enforce_num_trials: bool = True,
     max_parallelism: int | None = None,
     seed: int | None = None,
-    should_deduplicate: bool = False,
+    should_deduplicate: bool = True,
 ) -> GenerationStep:
     """Shortcut for creating a Sobol generation step."""
     return GenerationStep(
@@ -61,7 +61,7 @@ def _make_sobol_step(
         min_trials_observed=min_trials_observed or ceil(num_trials / 2),
         enforce_num_trials=enforce_num_trials,
         max_parallelism=max_parallelism,
-        generator_kwargs={"deduplicate": True, "seed": seed},
+        generator_kwargs={"seed": seed},
         should_deduplicate=should_deduplicate,
         use_all_trials_in_exp=True,
     )
@@ -76,7 +76,7 @@ def _make_botorch_step(
     generator_kwargs: dict[str, Any] | None = None,
     winsorization_config: None
     | (WinsorizationConfig | dict[str, WinsorizationConfig]) = None,
-    should_deduplicate: bool = False,
+    should_deduplicate: bool = True,
     disable_progbar: bool | None = None,
     jit_compile: bool | None = None,
     derelativize_with_raw_status_quo: bool = False,
@@ -299,7 +299,7 @@ def choose_generation_strategy_legacy(
     max_parallelism_cap: int | None = None,
     max_parallelism_override: int | None = None,
     optimization_config: OptimizationConfig | None = None,
-    should_deduplicate: bool = False,
+    should_deduplicate: bool = True,
     use_saasbo: bool = False,
     disable_progbar: bool | None = None,
     jit_compile: bool | None = None,

ax/generation_strategy/generation_node.py

@@ -122,7 +122,6 @@ class GenerationNode(SerializationMixin, SortableBase):
 
     # Required options:
     generator_specs: list[GeneratorSpec]
-    # TODO: Move `should_deduplicate` to `GeneratorSpec` if possible, and make optional
     should_deduplicate: bool
     _name: str
 
@@ -150,7 +149,7 @@ def __init__(
         generator_specs: list[GeneratorSpec],
         transition_criteria: Sequence[TransitionCriterion] | None = None,
         best_model_selector: BestModelSelector | None = None,
-        should_deduplicate: bool = False,
+        should_deduplicate: bool = True,
         input_constructors: TInputConstructorsByPurpose | None = None,
         previous_node_name: str | None = None,
         trial_type: str | None = None,
@@ -1014,7 +1013,7 @@ def __new__(
         min_trials_observed: int = 0,
         max_parallelism: int | None = None,
         enforce_num_trials: bool = True,
-        should_deduplicate: bool = False,
+        should_deduplicate: bool = True,
         generator_name: str | None = None,
         use_all_trials_in_exp: bool = False,
         use_update: bool = False,  # DEPRECATED.

ax/generation_strategy/tests/test_dispatch_utils.py

@@ -860,22 +860,24 @@ def test_max_parallelism_adjustments(self) -> None:
         self.assertEqual(self._get_max_parallelism(sobol_gpei._nodes[1]), 10)
 
     def test_set_should_deduplicate(self) -> None:
+        # Default is now should_deduplicate=True
         sobol_gpei = choose_generation_strategy_legacy(
             search_space=get_branin_search_space(),
             use_batch_trials=True,
             num_initialization_trials=3,
         )
         self.assertListEqual(
-            [s.should_deduplicate for s in sobol_gpei._nodes], [False] * 2
+            [s.should_deduplicate for s in sobol_gpei._nodes], [True] * 2
         )
+        # Explicitly set should_deduplicate=False
         sobol_gpei = choose_generation_strategy_legacy(
             search_space=get_branin_search_space(),
             use_batch_trials=True,
             num_initialization_trials=3,
-            should_deduplicate=True,
+            should_deduplicate=False,
         )
         self.assertListEqual(
-            [s.should_deduplicate for s in sobol_gpei._nodes], [True] * 2
+            [s.should_deduplicate for s in sobol_gpei._nodes], [False] * 2
         )
 
     def test_setting_experiment_attribute(self) -> None: