Filter fetch_data by trial type in base Experiment (#5005)

Summary:

Phase 4 of moving MultiTypeExperiment features into base Experiment.

Updates the base Experiment's `_fetch_trial_data` to filter metrics by trial type when `_trial_type_to_metric_names` is populated. Only metrics whose names appear in the set for the trial's type are fetched, preventing metrics from being evaluated against trials of the wrong type.

Updates `fetch_data` to iterate trial-by-trial when `_trial_type_to_metric_names` is populated (so each trial fetches only its relevant metrics), while preserving the existing bulk fetch path for single-type experiments.

Removes the `fetch_data` and `_fetch_trial_data` overrides from MultiTypeExperiment, along with now-unused imports (Iterable, Data, MetricFetchResult).

Differential Revision: D94990429

Author: mpolson64

ax/core/experiment.py

@@ -1180,6 +1180,41 @@ def fetch_data(
         Returns:
             Data for the experiment.
         """
+        # Only use trial-type-aware fetching for multi-type experiments;
+        # single-type experiments have an empty mapping.
+        if self._trial_type_to_metric_names:
+            # When metrics are mapped to trial types, group trials by type
+            # and bulk-fetch per group so each group only fetches its
+            # relevant metrics.
+            all_trials = (
+                list(self.trials.values())
+                if trial_indices is None
+                else self.get_trials_by_indices(trial_indices)
+            )
+            trials_by_type: dict[str | None, list[BaseTrial]] = defaultdict(list)
+            for trial in all_trials:
+                if trial.status.expecting_data:
+                    trials_by_type[trial.trial_type].append(trial)
+            all_results: dict[int, dict[str, MetricFetchResult]] = {}
+            for trial_type, type_trials in trials_by_type.items():
+                if metrics is not None and trial_type is not None:
+                    valid_names = self._trial_type_to_metric_names.get(
+                        trial_type, set()
+                    )
+                    type_metrics = [m for m in metrics if m.name in valid_names]
+                elif metrics is not None:
+                    type_metrics = metrics
+                elif trial_type is not None:
+                    type_metrics = self.metrics_for_trial_type(trial_type)
+                else:
+                    type_metrics = list(self.metrics.values())
+                results = self._lookup_or_fetch_trials_results(
+                    trials=type_trials,
+                    metrics=type_metrics,
+                    **kwargs,
+                )
+                all_results.update(results)
+            return Metric._unwrap_experiment_data_multi(results=all_results)
         results = self._lookup_or_fetch_trials_results(
             trials=list(self.trials.values())
             if trial_indices is None
@@ -1320,6 +1355,14 @@ def _fetch_trial_data(
     ) -> dict[str, MetricFetchResult]:
         trial = self.trials[trial_index]
 
+        # When metrics are mapped to trial types, filter to only the
+        # metrics relevant to this trial's type.
+        trial_type = trial.trial_type
+        if self._trial_type_to_metric_names and trial_type is not None:
+            valid_names = self._trial_type_to_metric_names.get(trial_type, set())
+            all_metrics = list(metrics or self.metrics.values())
+            metrics = [m for m in all_metrics if m.name in valid_names]
+
         trial_data = self._lookup_or_fetch_trials_results(
             trials=[trial], metrics=metrics, **kwargs
         )

ax/core/multi_type_experiment.py

@@ -6,14 +6,13 @@
 
 # pyre-strict
 
-from collections.abc import Iterable, Sequence
+from collections.abc import Sequence
 from typing import Any, Self
 
 from ax.core.arm import Arm
 from ax.core.base_trial import BaseTrial, TrialStatus
-from ax.core.data import Data
 from ax.core.experiment import Experiment
-from ax.core.metric import Metric, MetricFetchResult
+from ax.core.metric import Metric
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.runner import Runner
 from ax.core.search_space import SearchSpace
@@ -171,40 +170,6 @@ def remove_metric(self, metric_name: str) -> Self:
         self._metric_to_canonical_name.pop(metric_name, None)
         return self
 
-    @copy_doc(Experiment.fetch_data)
-    def fetch_data(
-        self,
-        trial_indices: Iterable[int] | None = None,
-        metrics: list[Metric] | None = None,
-        **kwargs: Any,
-    ) -> Data:
-        # TODO: make this more efficient for fetching
-        # data for multiple trials of the same type
-        # by overriding Experiment._lookup_or_fetch_trials_results
-        return Data.from_multiple_data(
-            [
-                (
-                    trial.fetch_data(**kwargs, metrics=metrics)
-                    if trial.status.expecting_data
-                    else Data()
-                )
-                for trial in self.trials.values()
-            ]
-        )
-
-    @copy_doc(Experiment._fetch_trial_data)
-    def _fetch_trial_data(
-        self, trial_index: int, metrics: list[Metric] | None = None, **kwargs: Any
-    ) -> dict[str, MetricFetchResult]:
-        trial = self.trials[trial_index]
-        metrics = [
-            metric
-            for metric in (metrics or self.metrics.values())
-            if self.metric_to_trial_type[metric.name] == trial.trial_type
-        ]
-        # Invoke parent's fetch method using only metrics for this trial_type
-        return super()._fetch_trial_data(trial.index, metrics=metrics, **kwargs)
-
 
 def filter_trials_by_type(
     trials: Sequence[BaseTrial], trial_type: str | None

ax/core/tests/test_experiment.py

@@ -2816,3 +2816,100 @@ def test_extract_relevant_trials(self) -> None:
             )
             self.assertEqual(len(trials), 1)
             self.assertEqual(trials[0].index, 0)
+
+    def _setup_multi_type_branin_experiment(self, n: int = 10) -> Experiment:
+        """Create a base Experiment with two trial types and metrics mapped
+        to each, mimicking a multi-type setup without using
+        MultiTypeExperiment.
+        """
+        exp = Experiment(
+            name="multi_type_test",
+            search_space=get_branin_search_space(),
+            default_trial_type="type1",
+            tracking_metrics=[
+                BraninMetric(name="m1", param_names=["x1", "x2"]),
+            ],
+            runner=SyntheticRunner(),
+        )
+        # Register a second trial type with its own runner and metric.
+        exp._trial_type_to_runner["type2"] = SyntheticRunner()
+        exp.add_tracking_metric(
+            BraninMetric(name="m2", param_names=["x2", "x1"]),
+            trial_type="type2",
+        )
+
+        # Create one batch per trial type and run them.
+        b1 = exp.new_batch_trial(trial_type="type1")
+        b1.add_arms_and_weights(arms=get_branin_arms(n=n, seed=0))
+        b1.run()
+
+        b2 = exp.new_batch_trial(trial_type="type2")
+        b2.add_arms_and_weights(arms=get_branin_arms(n=n, seed=0))
+        b2.run()
+
+        return exp
+
+    def test_fetch_data_with_trial_types(self) -> None:
+        """fetch_data should correctly filter metrics by trial type."""
+        n = 10
+        exp = self._setup_multi_type_branin_experiment(n=n)
+
+        with self.subTest("filters_by_trial_type"):
+            df = exp.fetch_data().df
+            # Each trial should have n rows (one per arm), for a total of 2*n.
+            self.assertEqual(len(df), 2 * n)
+
+            # Trial 0 (type1) should only have metric "m1".
+            trial_0_df = df[df["trial_index"] == 0]
+            self.assertEqual(set(trial_0_df["metric_name"]), {"m1"})
+            self.assertEqual(len(trial_0_df), n)
+
+            # Trial 1 (type2) should only have metric "m2".
+            trial_1_df = df[df["trial_index"] == 1]
+            self.assertEqual(set(trial_1_df["metric_name"]), {"m2"})
+            self.assertEqual(len(trial_1_df), n)
+
+        with self.subTest("with_trial_indices"):
+            # Fetch only trial 1 (type2).
+            df = exp.fetch_data(trial_indices=[1]).df
+            self.assertEqual(len(df), n)
+            self.assertEqual(set(df["metric_name"]), {"m2"})
+            self.assertTrue((df["trial_index"] == 1).all())
+
+        with self.subTest("skips_non_expecting_trials"):
+            # Mark trial 0 as abandoned so it doesn't expect data.
+            exp.trials[0].mark_abandoned()
+
+            df = exp.fetch_data().df
+            # Only trial 1 should have data.
+            self.assertEqual(len(df), n)
+            self.assertTrue((df["trial_index"] == 1).all())
+            self.assertEqual(set(df["metric_name"]), {"m2"})
+
+    def test_fetch_trial_data_with_trial_types(self) -> None:
+        """_fetch_trial_data should filter metrics by trial type."""
+        n = 10
+        exp = self._setup_multi_type_branin_experiment(n=n)
+
+        with self.subTest("filters_metrics_by_trial_type"):
+            # Fetch data for trial 0 (type1) -- should only contain "m1".
+            results_0 = exp._fetch_trial_data(trial_index=0)
+            self.assertIn("m1", results_0)
+            self.assertNotIn("m2", results_0)
+
+            # Fetch data for trial 1 (type2) -- should only contain "m2".
+            results_1 = exp._fetch_trial_data(trial_index=1)
+            self.assertIn("m2", results_1)
+            self.assertNotIn("m1", results_1)
+
+        with self.subTest("filters_explicit_metrics_by_trial_type"):
+            both_metrics = list(exp.metrics.values())
+            # Passing both metrics to a type1 trial should still only return m1.
+            results_0 = exp._fetch_trial_data(trial_index=0, metrics=both_metrics)
+            self.assertIn("m1", results_0)
+            self.assertNotIn("m2", results_0)
+
+            # Passing both metrics to a type2 trial should still only return m2.
+            results_1 = exp._fetch_trial_data(trial_index=1, metrics=both_metrics)
+            self.assertIn("m2", results_1)
+            self.assertNotIn("m1", results_1)