Fix InSampleUniformGenerator selecting arms from CANDIDATE trials (#5109)

Summary:
Pull Request resolved: #5109

When LILO labeling generates pairwise comparisons, `InSampleUniformGenerator` selects 2 arms uniformly at random from `generated_points` built by `RandomAdapter`. Previously, this pool included arms from CANDIDATE trials (which have no observed data) via two sources:

1. `arms_by_signature_for_deduplication` (all non-FAILED arms)
2. `pending_observations` (CANDIDATE/STAGED/RUNNING arms appended for dedup)

When a CANDIDATE trial arm was selected, `LILOPairwiseMetric` could not find source metric data for it, causing `fetch_data` to raise an `ExceptionGroup`.

**Fix:** In `RandomAdapter._gen()`, when the generator is `InSampleUniformGenerator`:
- Filter `arms_to_deduplicate` to only arms from `expecting_data` trials (COMPLETED, EARLY_STOPPED, RUNNING)
- Skip appending `pending_observations` (which re-adds CANDIDATE arms)

This is correct for all current and foreseeable use cases of `InSampleUniformGenerator`. The generator is used in two scenarios: LILO labeling (active) and potentially bandit candidate generation (not currently using it). In both cases, we need arms with observed data -- selecting an arm from a CANDIDATE or STAGED trial that has never been evaluated is never meaningful. Excluding pending observations is similarly safe: pending points exist to prevent regular generators (Sobol, BO) from re-suggesting in-flight arms, but for in-sample selection the entire point is to pick from already-observed arms.

The change is gated behind an `isinstance` check -- no effect on other generators (Sobol, Uniform, etc.).

Reviewed By: saitcakmak

Differential Revision: D98627073

fbshipit-source-id: b1bf9f378aac1eb56907d32da687e5874d22b632

Author: ItsMrLin

ax/adapter/random.py

@@ -26,6 +26,7 @@
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.search_space import SearchSpace
 from ax.generators.random.base import RandomGenerator
+from ax.generators.random.in_sample import InSampleUniformGenerator
 from ax.generators.types import TConfig
 
 
@@ -97,8 +98,25 @@ def _gen(
         # Exclude out-of-design arms (which can only be manual arms
         # instead of adapter-generated arms).
         generated_points = None
+        is_in_sample = isinstance(self.generator, InSampleUniformGenerator)
         if self.generator.deduplicate:
             arms_to_deduplicate = self._experiment.arms_by_signature_for_deduplication
+            # For in-sample generators, restrict to arms from trials that
+            # have or expect observed data (COMPLETED, EARLY_STOPPED,
+            # RUNNING). This prevents selecting arms from CANDIDATE/STAGED
+            # trials that have never been evaluated.
+            if is_in_sample:
+                expecting_sigs = {
+                    arm.signature
+                    for trial in self._experiment.trials.values()
+                    if trial.status.expecting_data
+                    for arm in trial.arms
+                }
+                arms_to_deduplicate = {
+                    sig: arm
+                    for sig, arm in arms_to_deduplicate.items()
+                    if sig in expecting_sigs
+                }
             generated_obs = [
                 ObservationFeatures.from_arm(arm=arm)
                 for arm in arms_to_deduplicate.values()
@@ -108,9 +126,16 @@ def _gen(
             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]
-            )
+            # Skipped for in-sample generators: pending observations include
+            # CANDIDATE arms that should not enter the selection pool.
+            if not is_in_sample:
+                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(

ax/adapter/tests/test_random_adapter.py

@@ -22,6 +22,7 @@
 from ax.core.search_space import SearchSpace
 from ax.exceptions.core import SearchSpaceExhausted
 from ax.generators.random.base import RandomGenerator
+from ax.generators.random.in_sample import InSampleUniformGenerator
 from ax.generators.random.sobol import SobolGenerator
 from ax.generators.types import TConfig
 from ax.utils.common.testutils import TestCase
@@ -297,6 +298,60 @@ def test_generated_points(self) -> None:
         generated_points_all_out = mock_gen.call_args.kwargs["generated_points"]
         self.assertIsNone(generated_points_all_out)
 
+    def test_in_sample_excludes_non_data_bearing_trial_arms(self) -> None:
+        """For InSampleUniformGenerator, generated_points should only contain
+        arms from trials with expecting_data status (COMPLETED, RUNNING,
+        EARLY_STOPPED). Arms from CANDIDATE, FAILED, and ABANDONED trials
+        must be excluded even though they exist on the experiment."""
+        search_space = SearchSpace(self.parameters[:2])
+        exp = Experiment(search_space=search_space)
+
+        # Trial 0: COMPLETED -- should be included.
+        exp.new_trial().add_arm(Arm(parameters={"x": 0.5, "y": 1.5})).mark_running(
+            no_runner_required=True
+        )
+        exp.trials[0].mark_completed()
+
+        # Trial 1: CANDIDATE -- should be excluded.
+        exp.new_trial().add_arm(Arm(parameters={"x": 0.8, "y": 1.8}))
+
+        # Trial 2: RUNNING -- should be included.
+        exp.new_trial().add_arm(Arm(parameters={"x": 0.3, "y": 1.3})).mark_running(
+            no_runner_required=True
+        )
+
+        # Trial 3: FAILED -- should be excluded.
+        exp.new_trial().add_arm(Arm(parameters={"x": 0.1, "y": 1.1})).mark_running(
+            no_runner_required=True
+        )
+        exp.trials[3].mark_failed()
+
+        # Trial 4: ABANDONED -- should be excluded.
+        exp.new_trial().add_arm(Arm(parameters={"x": 0.9, "y": 1.9})).mark_running(
+            no_runner_required=True
+        )
+        exp.trials[4].mark_abandoned()
+
+        generator = InSampleUniformGenerator(seed=0)
+        adapter = RandomAdapter(
+            experiment=exp,
+            generator=generator,
+            transforms=Cont_X_trans,
+        )
+
+        with mock.patch.object(
+            generator,
+            "gen",
+            wraps=generator.gen,
+        ) as mock_gen:
+            adapter.gen(n=2)
+
+        # generated_points should have exactly 2 points (COMPLETED + RUNNING).
+        # CANDIDATE, FAILED, and ABANDONED arms must be excluded.
+        generated_points = mock_gen.call_args.kwargs["generated_points"]
+        assert generated_points is not None
+        self.assertEqual(len(generated_points), 2)
+
     def test_generation_with_all_fixed(self) -> None:
         # Make sure candidate generation succeeds and returns correct parameters
         # when all parameters are fixed.

ax/generators/random/in_sample.py

@@ -19,8 +19,10 @@ class InSampleUniformGenerator(RandomGenerator):
     """Randomly select candidates from existing experiment arms.
 
     Selects n arms uniformly at random without replacement from the
-    ``generated_points`` array passed by the adapter. This array contains
-    the in-design, non-failed arms on the experiment (deduplicated).
+    ``generated_points`` array passed by the adapter. For this generator,
+    the adapter restricts ``generated_points`` to arms from trials that
+    have or expect observed data (``status.expecting_data``), excluding
+    arms from CANDIDATE or STAGED trials that have never been evaluated.
 
     Used for model-free candidate selection in use cases like LILO
     (Language-in-the-Loop Optimization), where a labeling node needs
@@ -51,8 +53,8 @@ def gen(
             model_gen_options: Not used. Accepted for interface compatibility.
             rounding_func: Not used. Accepted for interface compatibility.
             generated_points: A numpy array of shape ``(num_arms, d)`` containing
-                the existing experiment arms to select from. Constructed by the
-                adapter from in-design, non-failed arms (deduplicated).
+                the existing experiment arms to select from. The adapter
+                filters this to arms from trials with observed data.
 
         Returns:
             2-element tuple containing