Remove arms per node override (#4822)

Summary:
Pull Request resolved: #4822

We initially provided this override during transition from legacy dispatch to GS, and at the time voted to persist the override in case folks wanted to provide this information manually. However, in the ~year+ since we added this, I'm not sure we've had any usage of this, and keeping it around increases the complexity of an already challenging part of the stack.

My suggestion is to remove this, and for advanced users to directly leverage nodes or models if custom modeling is needed

Reviewed By: saitcakmak, lena-kashtelyan

Differential Revision: D91513177

fbshipit-source-id: ac5ceaa324a20552574b2f23eb746973c611761c

Author: mgarrard

ax/generation_strategy/center_generation_node.py

@@ -71,7 +71,6 @@ def gen(
         skip_fit: bool = False,
         data: Data | None = None,
         n: int | None = None,
-        arms_per_node: dict[str, int] | None = None,
         **gs_gen_kwargs: Any,
     ) -> GeneratorRun | None:
         """Generate candidates or skip if search space is exhausted.
@@ -100,7 +99,6 @@ def gen(
             skip_fit=skip_fit,
             data=data,
             n=n,
-            arms_per_node=arms_per_node,
             **gs_gen_kwargs,
         )
 

ax/generation_strategy/generation_node.py

@@ -426,7 +426,6 @@ def gen(
         skip_fit: bool = False,
         data: Data | None = None,
         n: int | None = None,
-        arms_per_node: dict[str, int] | None = None,
         **gs_gen_kwargs: Any,
     ) -> GeneratorRun | None:
         """This method generates candidates using `self._gen` and handles deduplication
@@ -451,9 +450,6 @@ def gen(
             data: Optional override for the experiment data used to generate candidates;
                 if not specified, will use ``experiment.lookup_data()`` (extracted in
                 ``Adapter``).
-            arms_per_node: A manual override for users interacting with a gen. strategy
-                via a Python API; a mapping from node name to the specific number of
-                arms it should produce. Passed down here by `GenerationStrategy.gen`.
             gs_gen_kwargs: Keyword arguments, passed to ``GenerationStrategy.gen``.
                 These might be modified by this node's input constructors, before
                 being passed down to ``ModelSpec.gen``, where these will override any
@@ -485,13 +481,6 @@ def gen(
             logger.debug(f"Skipping generation for node {self.name}.")
             return None
 
-        if arms_per_node:
-            if self.name not in arms_per_node:
-                raise UnsupportedError(
-                    "If manually specifying arms per node, all nodes must be specified."
-                )
-            generator_gen_kwargs["n"] = arms_per_node[self.name]
-
         # TODO[drfreund]: Move this to `Adapter` or another more suitable place.
         # Keeping here for now to limit the scope of the current changeset.
         generator_gen_kwargs["fixed_features"] = (

ax/generation_strategy/generation_node_input_constructors.py

@@ -257,14 +257,10 @@ def _get_default_n(experiment: Experiment, next_node: GenerationNode) -> int:
         The default number of arms to generate from the next node, used if no n is
         provided to the ``GenerationStrategy``'s gen call.
     """
-    # If the generator spec contains `n` use that value first
-    # TODO #1 [drfreund, mgarrard]: Eliminate the need to do this; the order should be:
-    # `arms_per_node[node_name]` > `input_constuctors(n)` > `gen_spec...kwargs["n"]`
-    # NOTE: We might need to simply disallow `n` in `gen_spec...kwargs`: it should
-    # probably never be hardcoded there. Without it, we can just enforce that at a
-    # point within generation strategy, an `n` is passed down to `gen_spec.gen`.
-    # And if we keep it, we don't have a clear point in this stack at which we are
-    # "no longer allowed to have a null `n`."
+    # If the generator spec contains `n` use that value first.
+    # TODO [drfreund, mgarrard]: Consider disallowing `n` in `gen_spec...kwargs`:
+    # it should probably never be hardcoded there. This would enforce that `n`
+    # is always passed down through the generation strategy at runtime.
     if next_node.generator_spec_to_gen_from.generator_gen_kwargs.get("n") is not None:
         return next_node.generator_spec_to_gen_from.generator_gen_kwargs["n"]
 

ax/generation_strategy/generation_strategy.py

@@ -19,12 +19,7 @@
 from ax.core.generator_run import GeneratorRun
 from ax.core.observation import ObservationFeatures
 from ax.core.utils import extend_pending_observations, extract_pending_observations
-from ax.exceptions.core import (
-    AxError,
-    DataRequiredError,
-    UnsupportedError,
-    UserInputError,
-)
+from ax.exceptions.core import AxError, DataRequiredError, UnsupportedError
 from ax.exceptions.generation_strategy import (
     GenerationStrategyCompleted,
     GenerationStrategyMisconfiguredException,
@@ -247,7 +242,6 @@ def gen(
         n: int | None = None,
         fixed_features: ObservationFeatures | None = None,
         num_trials: int = 1,
-        arms_per_node: dict[str, int] | None = None,
     ) -> list[list[GeneratorRun]]:
         """Produce GeneratorRuns for multiple trials at once with the possibility of
         using multiple models per trial, getting multiple GeneratorRuns per trial.
@@ -275,12 +269,6 @@ def gen(
                 important to specify all necessary fixed features.
             num_trials: Number of trials to generate generator runs for in this call.
                 If not provided, defaults to 1.
-            arms_per_node: An optional map from node name to the number of arms to
-                generate from that node. If not provided, will default to the number
-                of arms specified in the node's ``InputConstructors`` or n if no
-                ``InputConstructors`` are defined on the node. We expect either n or
-                arms_per_node to be provided, but not both, and this is an advanced
-                argument that should only be used by advanced users.
 
         Returns:
             A list of lists of lists generator runs. Each outer list represents
@@ -306,7 +294,6 @@ def gen(
                     data=data,
                     n=n,
                     pending_observations=pending_observations,
-                    arms_per_node=arms_per_node,
                     fixed_features=fixed_features,
                     first_generation_in_multi=len(grs_for_multiple_trials) < 1,
                 )
@@ -467,24 +454,6 @@ def _validate_and_set_node_graph(self, nodes: list[GenerationNode]) -> None:
 
         self._curr = nodes[0]
 
-    def _validate_arms_per_node(self, arms_per_node: dict[str, int] | None) -> None:
-        """Validate that the arms_per_node argument is valid if it is provided.
-
-        Args:
-            arms_per_node: A map from node name to the number of arms to
-                generate from that node.
-        """
-        if arms_per_node is not None and not set(self.nodes_by_name).issubset(
-            arms_per_node
-        ):
-            raise UserInputError(
-                "Each node defined in the `GenerationStrategy` must have an "
-                "associated number of arms to generate from that node defined "
-                f"in `arms_per_node`. {arms_per_node} does not include all of "
-                f"{self.nodes_by_name.keys()}. "
-                "It may help to double-check the spelling."
-            )
-
     def _make_default_name(self) -> str:
         """Make a default name for this generation strategy; used when no name is passed
         to the constructor. For node-based generation strategies, the name is
@@ -515,10 +484,6 @@ def _gen_with_multiple_nodes(
         pending_observations: dict[str, list[ObservationFeatures]] | None = None,
         data: Data | None = None,
         fixed_features: ObservationFeatures | None = None,
-        # TODO: Consider naming `arms_per_node` smtg like `arms_per_node_override`,
-        # to convey its manually-specified nature (if it's not specified, GS selects
-        # what to do on its own).
-        arms_per_node: dict[str, int] | None = None,
         first_generation_in_multi: bool = True,
     ) -> list[GeneratorRun]:
         """Produces a List of GeneratorRuns for a single trial, either ``Trial`` or
@@ -548,12 +513,6 @@ def _gen_with_multiple_nodes(
                 passed down to the underlying nodes. Note: if provided this will
                 override any algorithmically determined fixed features so it is
                 important to specify all necessary fixed features.
-            arms_per_node: An optional map from node name to the number of arms to
-                generate from that node. If not provided, will default to the number
-                of arms specified in the node's ``InputConstructors`` or n if no
-                ``InputConstructors`` are defined on the node. We expect either n or
-                arms_per_node to be provided, but not both, and this is an advanced
-                argument that should only be used by advanced users.
 
         Returns:
             A list of ``GeneratorRuns`` for a single trial.
@@ -570,7 +529,6 @@ def _gen_with_multiple_nodes(
             pending_observations if pending_observations is not None else {}
         )
         self.experiment = experiment
-        self._validate_arms_per_node(arms_per_node=arms_per_node)
         pack_gs_gen_kwargs = {
             "grs_this_gen": grs_this_gen,
             "fixed_features": fixed_features,
@@ -596,7 +554,6 @@ def _gen_with_multiple_nodes(
                     pending_observations=pending_observations,
                     skip_fit=not (first_generation_in_multi or transitioned),
                     n=n,
-                    arms_per_node=arms_per_node,
                     **pack_gs_gen_kwargs,
                 )
             except DataRequiredError as err:

ax/generation_strategy/tests/test_generation_strategy.py

@@ -1448,17 +1448,12 @@ def test_gen_with_multiple_nodes_pending_points(self) -> None:
             ]
         )
         gs.experiment = exp
-        arms_per_node = {
-            "sobol_1": 2,
-            "sobol_2": 1,
-            "sobol_3": 3,
-        }
         with mock_patch_method_original(
             mock_path=f"{GeneratorSpec.__module__}.GeneratorSpec.gen",
             original_method=GeneratorSpec.gen,
         ) as gen_spec_gen_mock:
             # Generate a trial that should be composed of arms from 3 nodes
-            grs = gs.gen(experiment=exp, arms_per_node=arms_per_node)[0]
+            grs = gs.gen(experiment=exp, n=6)[0]
             self.assertEqual(len(grs), 3)  # len == 3 due to 3 nodes contributing
             self.assertEqual(gen_spec_gen_mock.call_count, 3)
             pending_in_each_gen = enumerate(
@@ -1491,15 +1486,17 @@ def test_gen_with_multiple_nodes_pending_points(self) -> None:
             # check that we can pass in pending points
             grs = gs.gen(
                 experiment=exp,
-                arms_per_node=arms_per_node,
+                n=3,
             )[0]
             self.assertEqual(len(grs), 3)  # len == 3 due to 3 nodes contributing
             pending_in_each_gen = enumerate(
                 call_kwargs.get("pending_observations")
                 for _, call_kwargs in gen_spec_gen_mock.call_args_list
             )
-            # check pending points is now 12 (from the previous trial having 6 arms)
-            self.assertEqual(len(list(pending_in_each_gen)[0][1]["m1"]), 12)
+            # check pending points is now 27 (18 arms from previous trial with 3 nodes
+            # each generating n=6 arms, plus 9 arms from the new generation with 3 nodes
+            # each generating n=3 arms)
+            self.assertEqual(len(list(pending_in_each_gen)[0][1]["m1"]), 27)
 
     def test_gs_initializes_default_props_correctly(self) -> None:
         """Test that all previous nodes are initialized to None"""
@@ -1707,68 +1704,6 @@ def test_transition_edges(self) -> None:
             },
         )
 
-    def test_multiple_arms_per_node(self) -> None:
-        """Test that a ``GenerationStrategy`` which expects some trials to be composed
-        of multiple nodes can generate multiple arms per node using `arms_per_node`.
-        """
-        exp = get_branin_experiment()
-        gs = self.complex_multinode_per_trial_gs
-        gs.experiment = exp
-        # first check that arms_per node validation works
-        arms_per_node = {
-            "sobol": 3,
-            "sobol_2": 2,
-            "sobol_3": 1,
-            "sobol_4": 4,
-        }
-        with self.assertRaisesRegex(UserInputError, "defined in `arms_per_node`"):
-            gs.gen(exp, arms_per_node=arms_per_node)
-
-        # now we will check that the first trial contains 3 arms, the second trial
-        # contains 6 arms (2 from mbm, 1 from sobol_2, 3 from sobol_3), and all
-        # remaining trials contain 4 arms
-        arms_per_node = {
-            "sobol": 3,
-            "mbm": 1,
-            "sobol_2": 2,
-            "sobol_3": 3,
-            "sobol_4": 4,
-        }
-        # for the first trial, we start on sobol, we generate the trial, but it hasn't
-        # been run yet, so we remain on sobol
-        trial0 = exp.new_batch_trial(
-            generator_runs=gs.gen(exp, arms_per_node=arms_per_node)[0]
-        )
-        self.assertEqual(len(trial0.arms_by_name), 3)
-        self.assertEqual(trial0.generator_runs[0]._generation_node_name, "sobol")
-        trial0.run()
-
-        # after trial 0 is run, we create a trial with nodes mbm, sobol_2, and sobol_3
-        # However, the sobol_3 criterion requires that we have two running trials. We
-        # don't move onto sobol_4 until we have two running trials, instead we reset
-        # to the last first node in a trial.
-        for _i in range(0, 2):
-            trial = exp.new_batch_trial(
-                generator_runs=gs.gen(exp, arms_per_node=arms_per_node)[0]
-            )
-            self.assertEqual(gs.current_node_name, "sobol_3")
-            self.assertEqual(len(trial.arms_by_name), 6)
-            self.assertEqual(len(trial.generator_runs), 3)
-            self.assertEqual(trial.generator_runs[0]._generation_node_name, "mbm")
-            self.assertEqual(len(trial.generator_runs[0].arms), 1)
-            self.assertEqual(trial.generator_runs[1]._generation_node_name, "sobol_2")
-            self.assertEqual(len(trial.generator_runs[1].arms), 2)
-            self.assertEqual(trial.generator_runs[2]._generation_node_name, "sobol_3")
-            self.assertEqual(len(trial.generator_runs[2].arms), 3)
-
-        # after running the next trial should be made from sobol 4
-        trial.run()
-        trial = exp.new_batch_trial(
-            generator_runs=gs.gen(exp, arms_per_node=arms_per_node)[0]
-        )
-        self.assertEqual(trial.generator_runs[0]._generation_node_name, "sobol_4")
-        self.assertEqual(len(trial.generator_runs[0].arms), 4)
-
     def test_gen_with_multiple_uses_total_concurrent_arms_for_a_default(self) -> None:
         exp = get_branin_experiment()
         self.sobol_node._input_constructors = {