diff --git a/ax/core/experiment.py b/ax/core/experiment.py index 1f13f6181aa..36623948d50 100644 --- a/ax/core/experiment.py +++ b/ax/core/experiment.py @@ -161,6 +161,11 @@ def __init__( self._trial_type_to_runner: dict[str | None, Runner | None] = { default_trial_type: runner } + # Maps each trial type to the set of metric names relevant to that type. + # This is the complement of _trial_type_to_runner and is used for + # multi-type experiments where different metrics apply to different + # trial types. + self._trial_type_to_metric_names: dict[str, set[str]] = {} # Used to keep track of whether any trials on the experiment # specify a TTL. Since trials need to be checked for their TTL's # expiration often, having this attribute helps avoid unnecessary @@ -197,6 +202,10 @@ def __init__( # a naming collision occurs. for m in [*(tracking_metrics or []), *(metrics or [])]: self._metrics[m.name] = m + if self._default_trial_type is not None: + self._trial_type_to_metric_names.setdefault( + self._default_trial_type, set() + ).add(m.name) # call setters defined below self.status_quo = status_quo @@ -585,6 +594,12 @@ def optimization_config(self, optimization_config: OptimizationConfig) -> None: "but not found on experiment. Add it first with add_metric()." ) self._optimization_config = optimization_config + resolved_trial_type = self._resolve_trial_type(None) + if resolved_trial_type is not None: + for metric_name in optimization_config.metric_names: + self._trial_type_to_metric_names.setdefault( + resolved_trial_type, set() + ).add(metric_name) @property def is_moo_problem(self) -> bool: @@ -837,7 +852,41 @@ def get_metric(self, name: str) -> Metric: ) return self._metrics[name] - def add_metric(self, metric: Metric) -> Self: + def _resolve_trial_type(self, trial_type: str | None) -> str | None: + """Resolve an explicit or default trial type and validate it. + + Returns ``trial_type`` if explicitly provided (after validating via + ``supports_trial_type``), falls back to ``_default_trial_type`` when + available, and raises ``ValueError`` if this experiment uses trial types + (``_trial_type_to_metric_names`` is non-empty) but none could be + resolved. + + Args: + trial_type: The explicitly provided trial type, or ``None``. + + Returns: + The resolved trial type, which may be ``None`` for single-type + experiments. + + Raises: + ValueError: If ``trial_type`` is provided but not supported, or if + no trial type could be resolved for a multi-type experiment. + """ + if trial_type is not None: + if not self.supports_trial_type(trial_type): + raise ValueError(f"`{trial_type}` is not a supported trial type.") + return trial_type + if self._default_trial_type is not None: + return self._default_trial_type + if self._trial_type_to_metric_names: + raise ValueError( + "This experiment has trial-type-aware metrics but no " + "`trial_type` was specified and no `default_trial_type` is set. " + "Please specify a `trial_type`." + ) + return None + + def add_metric(self, metric: Metric, trial_type: str | None = None) -> Self: """Add a new metric to the experiment. Metrics that are not referenced by the experiment's optimization config @@ -846,54 +895,98 @@ def add_metric(self, metric: Metric) -> Self: Args: metric: Metric to be added. + trial_type: If provided, associates the metric with this trial type. + When ``None`` and a ``default_trial_type`` is set, defaults to + the default trial type. + + Raises: + ValueError: If the metric already exists, the trial type is not + supported, or trial types are in use but none could be resolved. """ if metric.name in self._metrics: raise ValueError( f"Metric `{metric.name}` already defined on experiment. " "Use `update_metric` to update an existing metric definition." ) + trial_type = self._resolve_trial_type(trial_type) + if trial_type is not None: + self._trial_type_to_metric_names.setdefault(trial_type, set()).add( + metric.name + ) self._metrics[metric.name] = metric return self - def add_tracking_metric(self, metric: Metric) -> Self: + def add_tracking_metric( + self, + metric: Metric, + trial_type: str | None = None, + canonical_name: str | None = None, + ) -> Self: """*Deprecated.* Use ``add_metric`` instead.""" warnings.warn( "add_tracking_metric is deprecated. Use add_metric instead.", DeprecationWarning, stacklevel=2, ) - return self.add_metric(metric) + return self.add_metric(metric, trial_type=trial_type) - def add_tracking_metrics(self, metrics: list[Metric]) -> Experiment: + def add_tracking_metrics( + self, + metrics: list[Metric], + metrics_to_trial_types: dict[str, str] | None = None, + canonical_names: dict[str, str] | None = None, + ) -> Experiment: """*Deprecated.* Use ``add_metric`` instead.""" warnings.warn( "add_tracking_metrics is deprecated. Use add_metric instead.", DeprecationWarning, stacklevel=2, ) + metrics_to_trial_types = metrics_to_trial_types or {} for metric in metrics: - self.add_metric(metric) + canonical_name = (canonical_names or {}).get(metric.name) + self.add_tracking_metric( + metric=metric, + trial_type=metrics_to_trial_types.get(metric.name), + canonical_name=canonical_name, + ) return self - def update_metric(self, metric: Metric) -> Self: + def update_metric(self, metric: Metric, trial_type: str | None = None) -> Self: """Redefine a metric that already exists on the experiment. Args: metric: New metric definition. + trial_type: If provided, reassociates the metric with this trial + type. When ``None``, keeps the metric's existing trial type. """ if metric.name not in self._metrics: raise ValueError(f"Metric `{metric.name}` doesn't exist on experiment.") + if trial_type is not None: + trial_type = self._resolve_trial_type(trial_type) + # Remove from any existing trial type set + for names in self._trial_type_to_metric_names.values(): + names.discard(metric.name) + # Add to new trial type set + self._trial_type_to_metric_names.setdefault(trial_type, set()).add( + metric.name + ) self._metrics[metric.name] = metric return self - def update_tracking_metric(self, metric: Metric) -> Experiment: + def update_tracking_metric( + self, + metric: Metric, + trial_type: str | None = None, + canonical_name: str | None = None, + ) -> Experiment: """*Deprecated.* Use ``update_metric`` instead.""" warnings.warn( "update_tracking_metric is deprecated. Use update_metric instead.", DeprecationWarning, stacklevel=2, ) - return self.update_metric(metric) + return self.update_metric(metric, trial_type=trial_type) def remove_metric(self, metric_name: str) -> Self: """Remove a metric from the experiment. @@ -914,6 +1007,9 @@ def remove_metric(self, metric_name: str) -> Self: f"Metric `{metric_name}` is referenced by the optimization config " "and cannot be removed. Update the optimization config first." ) + # Clean up _trial_type_to_metric_names + for names in self._trial_type_to_metric_names.values(): + names.discard(metric_name) del self._metrics[metric_name] return self @@ -1084,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 @@ -1224,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 ) @@ -1928,6 +2067,85 @@ def default_trial_type(self) -> str | None: """ return self._default_trial_type + @property + def trial_type_to_metric_names(self) -> dict[str, set[str]]: + """Map from trial type to the set of metric names relevant to that + type. + + Returns a shallow copy of the internal mapping. + """ + return dict(self._trial_type_to_metric_names) + + @property + def metric_to_trial_type(self) -> dict[str, str]: + """Map each metric name to its associated trial type. + + Computed from ``_trial_type_to_metric_names``. When a + ``default_trial_type`` is set and an ``optimization_config`` exists, + optimization config metrics are pinned to the default trial type. + """ + result: dict[str, str] = {} + for trial_type, metric_names in self._trial_type_to_metric_names.items(): + for name in metric_names: + result[name] = trial_type + opt_config = self._optimization_config + default_trial_type = self._default_trial_type + if default_trial_type is not None and opt_config is not None: + for metric_name in opt_config.metric_names: + result[metric_name] = default_trial_type + return result + + def metrics_for_trial_type(self, trial_type: str) -> list[Metric]: + """Return the metrics associated with a given trial type. + + Args: + trial_type: The trial type to look up metrics for. + + Raises: + ValueError: If the trial type is not supported. + """ + if not self.supports_trial_type(trial_type): + raise ValueError(f"Trial type `{trial_type}` is not supported.") + valid_names = self._trial_type_to_metric_names.get(trial_type, set()) + return [self._metrics[name] for name in valid_names if name in self._metrics] + + @property + def default_trials(self) -> set[int]: + """Return the indices for trials of the default type.""" + return { + idx + for idx, trial in self.trials.items() + if trial.trial_type == self.default_trial_type + } + + def add_trial_type(self, trial_type: str, runner: Runner | None = None) -> Self: + """Add a new trial type to be supported by this experiment. + + Args: + trial_type: The new trial type to be added. + runner: The default runner for trials of this type. + """ + if self.supports_trial_type(trial_type): + raise ValueError(f"Experiment already contains trial_type `{trial_type}`") + + if runner is not None: + self._trial_type_to_runner[trial_type] = runner + + return self + + def update_runner(self, trial_type: str, runner: Runner) -> Self: + """Update the default runner for an existing trial type. + + Args: + trial_type: The trial type whose runner should be updated. + runner: The new runner for trials of this type. + """ + if not self.supports_trial_type(trial_type): + raise ValueError(f"Experiment does not contain trial_type `{trial_type}`") + self._trial_type_to_runner[trial_type] = runner + self._runner = runner + return self + def runner_for_trial_type(self, trial_type: str | None) -> Runner | None: """The default runner to use for a given trial type. @@ -1942,14 +2160,20 @@ def runner_for_trial_type(self, trial_type: str | None) -> Runner | None: def supports_trial_type(self, trial_type: str | None) -> bool: """Whether this experiment allows trials of the given type. - The base experiment class only supports None. For experiments - with multiple trial types, use the MultiTypeExperiment class. + For experiments with a ``default_trial_type`` (multi-type experiments), + only trial types registered in ``_trial_type_to_runner`` are supported. + For single-type experiments, ``None`` is always supported, along with + ``SHORT_RUN`` and ``LONG_RUN`` for backward compatibility with + generation strategies that use those trial types. """ + if self._default_trial_type is not None: + return trial_type in self._trial_type_to_runner return ( trial_type is None or trial_type == Keys.SHORT_RUN or trial_type == Keys.LONG_RUN or trial_type == Keys.LILO_LABELING + or trial_type in self._trial_type_to_runner ) def attach_trial( diff --git a/ax/core/multi_type_experiment.py b/ax/core/multi_type_experiment.py index f6a1ce91add..ca8b67270da 100644 --- a/ax/core/multi_type_experiment.py +++ b/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 @@ -97,28 +96,13 @@ def __init__( ) # Ensure tracking metrics are registered in _metric_to_trial_type. - # super().__init__ sets self._metrics directly, bypassing - # add_tracking_metric, so tracking metrics won't be in - # _metric_to_trial_type yet. + # The base __init__ handles _trial_type_to_metric_names. for m in tracking_metrics or []: if m.name not in self._metric_to_trial_type: self._metric_to_trial_type[m.name] = none_throws( self._default_trial_type ) - def add_trial_type(self, trial_type: str, runner: Runner) -> Self: - """Add a new trial_type to be supported by this experiment. - - Args: - trial_type: The new trial_type to be added. - runner: The default runner for trials of this type. - """ - if self.supports_trial_type(trial_type): - raise ValueError(f"Experiment already contains trial_type `{trial_type}`") - - self._trial_type_to_runner[trial_type] = runner - return self - # pyre does not support inferring the type of property setter decorators # or the `.fset` attribute on properties. # pyre-fixme[56]: Pyre was not able to infer the type of the decorator. @@ -126,27 +110,12 @@ def add_trial_type(self, trial_type: str, runner: Runner) -> Self: def optimization_config(self, optimization_config: OptimizationConfig) -> None: # pyre-fixme[16]: `Optional` has no attribute `fset`. Experiment.optimization_config.fset(self, optimization_config) + # Base setter handles _trial_type_to_metric_names; update legacy dict. for metric_name in optimization_config.metric_names: - # Optimization config metrics are required to be the default trial type - # currently. TODO: remove that restriction (T202797235) self._metric_to_trial_type[metric_name] = none_throws( self.default_trial_type ) - def update_runner(self, trial_type: str, runner: Runner) -> Self: - """Update the default runner for an existing trial_type. - - Args: - trial_type: The new trial_type to be added. - runner: The new runner for trials of this type. - """ - if not self.supports_trial_type(trial_type): - raise ValueError(f"Experiment does not contain trial_type `{trial_type}`") - - self._trial_type_to_runner[trial_type] = runner - self._runner = runner - return self - def add_tracking_metric( self, metric: Metric, @@ -162,54 +131,12 @@ def add_tracking_metric( """ if trial_type is None: trial_type = self._default_trial_type - if not self.supports_trial_type(trial_type): - raise ValueError(f"`{trial_type}` is not a supported trial type.") - - super().add_tracking_metric(metric) + self.add_metric(metric, trial_type=trial_type) self._metric_to_trial_type[metric.name] = none_throws(trial_type) if canonical_name is not None: self._metric_to_canonical_name[metric.name] = canonical_name return self - def add_tracking_metrics( - self, - metrics: list[Metric], - metrics_to_trial_types: dict[str, str] | None = None, - canonical_names: dict[str, str] | None = None, - ) -> Experiment: - """Add a list of new metrics to the experiment. - - If any of the metrics are already defined on the experiment, - we raise an error and don't add any of them to the experiment - - Args: - metrics: Metrics to be added. - metrics_to_trial_types: The mapping from metric names to corresponding - trial types for each metric. If provided, the metrics will be - added to their trial types. If not provided, then the default - trial type will be used. - canonical_names: A mapping of metric names to their - canonical names(The default metrics for which the metrics are - proxies.) - - Returns: - The experiment with the added metrics. - """ - metrics_to_trial_types = metrics_to_trial_types or {} - canonical_name = None - for metric in metrics: - if canonical_names is not None: - canonical_name = none_throws(canonical_names).get(metric.name, None) - - self.add_tracking_metric( - metric=metric, - trial_type=metrics_to_trial_types.get( - metric.name, self._default_trial_type - ), - canonical_name=canonical_name, - ) - return self - def update_tracking_metric( self, metric: Metric, @@ -230,118 +157,19 @@ def update_tracking_metric( trial_type = self._metric_to_trial_type.get( metric.name, self._default_trial_type ) - oc = self.optimization_config - oc_metric_names = oc.metric_names if oc else set() - if metric.name in oc_metric_names and trial_type != self._default_trial_type: - raise ValueError( - f"Metric `{metric.name}` must remain a " - f"`{self._default_trial_type}` metric because it is part of the " - "optimization_config." - ) - elif not self.supports_trial_type(trial_type): - raise ValueError(f"`{trial_type}` is not a supported trial type.") - - super().update_tracking_metric(metric) + self.update_metric(metric, trial_type=trial_type) self._metric_to_trial_type[metric.name] = none_throws(trial_type) if canonical_name is not None: self._metric_to_canonical_name[metric.name] = canonical_name return self - @copy_doc(Experiment.remove_tracking_metric) - def remove_tracking_metric(self, metric_name: str) -> Self: - if metric_name not in self._metrics: - raise ValueError(f"Metric `{metric_name}` doesn't exist on experiment.") - - # Required fields - del self._metrics[metric_name] - del self._metric_to_trial_type[metric_name] - - # Optional - if metric_name in self._metric_to_canonical_name: - del self._metric_to_canonical_name[metric_name] + @copy_doc(Experiment.remove_metric) + def remove_metric(self, metric_name: str) -> Self: + super().remove_metric(metric_name) + self._metric_to_trial_type.pop(metric_name, None) + 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) - - @property - def default_trials(self) -> set[int]: - """Return the indicies for trials of the default type.""" - return { - idx - for idx, trial in self.trials.items() - if trial.trial_type == self.default_trial_type - } - - @property - def metric_to_trial_type(self) -> dict[str, str]: - """Map metrics to trial types. - - Adds in default trial type for OC metrics to custom defined trial types.. - """ - opt_config_types = { - metric_name: self.default_trial_type - for metric_name in self.optimization_config.metric_names - } - return {**opt_config_types, **self._metric_to_trial_type} - - # -- Overridden functions from Base Experiment Class -- - @property - def default_trial_type(self) -> str | None: - """Default trial type assigned to trials in this experiment.""" - return self._default_trial_type - - def metrics_for_trial_type(self, trial_type: str) -> list[Metric]: - """The default runner to use for a given trial type. - - Looks up the appropriate runner for this trial type in the trial_type_to_runner. - """ - if not self.supports_trial_type(trial_type): - raise ValueError(f"Trial type `{trial_type}` is not supported.") - return [ - self.metrics[metric_name] - for metric_name, metric_trial_type in self._metric_to_trial_type.items() - if metric_trial_type == trial_type - ] - - def supports_trial_type(self, trial_type: str | None) -> bool: - """Whether this experiment allows trials of the given type. - - Only trial types defined in the trial_type_to_runner are allowed. - """ - return trial_type in self._trial_type_to_runner.keys() - def filter_trials_by_type( trials: Sequence[BaseTrial], trial_type: str | None diff --git a/ax/core/tests/test_experiment.py b/ax/core/tests/test_experiment.py index 80e5de75e29..8eac8d0620a 100644 --- a/ax/core/tests/test_experiment.py +++ b/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) diff --git a/ax/core/tests/test_multi_type_experiment.py b/ax/core/tests/test_multi_type_experiment.py index b314cfd5b75..5418d5e73e8 100644 --- a/ax/core/tests/test_multi_type_experiment.py +++ b/ax/core/tests/test_multi_type_experiment.py @@ -125,12 +125,6 @@ def test_BadBehavior(self) -> None: with self.assertRaises(ValueError): self.experiment.remove_tracking_metric("m3") - # Try to change optimization metric to non-primary trial type - with self.assertRaises(ValueError): - self.experiment.update_tracking_metric( - BraninMetric("m1", ["x1", "x2"]), "type2" - ) - # Update metric definition for trial_type that doesn't exist with self.assertRaises(ValueError): self.experiment.update_tracking_metric( diff --git a/ax/storage/json_store/decoder.py b/ax/storage/json_store/decoder.py index fc5827a89a1..1c7cf70ffdc 100644 --- a/ax/storage/json_store/decoder.py +++ b/ax/storage/json_store/decoder.py @@ -720,6 +720,12 @@ def multi_type_experiment_from_json( experiment._metric_to_trial_type = _metric_to_trial_type experiment._trial_type_to_runner = _trial_type_to_runner + # Rebuild _trial_type_to_metric_names from _metric_to_trial_type + trial_type_to_metric_names: dict[str, set[str]] = {} + for metric_name, trial_type in _metric_to_trial_type.items(): + trial_type_to_metric_names.setdefault(trial_type, set()).add(metric_name) + experiment._trial_type_to_metric_names = trial_type_to_metric_names + _load_experiment_info( exp=experiment, exp_info=experiment_info,