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

Summary:
Pull Request resolved: facebook#5005

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

@@ -1138,6 +1138,37 @@ def fetch_data(
         Returns:
             Data for the experiment.
         """
+        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():
+                type_metrics = (
+                    metrics
+                    if metrics is not None
+                    else (
+                        self.metrics_for_trial_type(trial_type)
+                        if trial_type is not None
+                        else 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
@@ -1278,6 +1309,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

@@ -2772,3 +2772,114 @@ 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_filters_by_trial_type(self) -> None:
+        """fetch_data should return only the metrics mapped to each trial's
+        type when _trial_type_to_metric_names is populated."""
+        n = 10
+        exp = self._setup_multi_type_branin_experiment(n=n)
+
+        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)
+
+    def test_fetch_data_with_trial_indices_and_trial_types(self) -> None:
+        """fetch_data with trial_indices should respect trial type filtering."""
+        n = 10
+        exp = self._setup_multi_type_branin_experiment(n=n)
+
+        # 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())
+
+    def test_fetch_data_skips_non_expecting_trials_with_trial_types(self) -> None:
+        """fetch_data should skip trials not expecting data when
+        _trial_type_to_metric_names is populated."""
+        n = 10
+        exp = self._setup_multi_type_branin_experiment(n=n)
+
+        # 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_filters_metrics_by_trial_type(self) -> None:
+        """_fetch_trial_data should filter to only metrics relevant to the
+        trial's type when _trial_type_to_metric_names is populated."""
+        n = 10
+        exp = self._setup_multi_type_branin_experiment(n=n)
+
+        # 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)
+
+    def test_fetch_trial_data_filters_explicit_metrics_by_trial_type(self) -> None:
+        """_fetch_trial_data should filter even an explicit metrics list to
+        only those relevant to the trial's type."""
+        n = 10
+        exp = self._setup_multi_type_branin_experiment(n=n)
+
+        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)