Replace isinstance(MultiTypeExperiment) checks with feature checks (#5001)#5001
Open
mpolson64 wants to merge 6 commits intofacebook:mainfrom
Open
Replace isinstance(MultiTypeExperiment) checks with feature checks (#5001)#5001mpolson64 wants to merge 6 commits intofacebook:mainfrom
mpolson64 wants to merge 6 commits intofacebook:mainfrom
Conversation
meta-cla
Bot
added
the
CLA Signed
|
@mpolson64 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D94999020. |
mpolson64
added a commit
to mpolson64/Ax
that referenced
this pull request
Mar 11, 2026
…acebook#5001) Summary: Pull Request resolved: facebook#5001 Move `filter_trials_by_type` and `get_trial_indices_for_statuses` from `multi_type_experiment.py` to `experiment.py`, with backward-compatible re-exports from the old module. Replace all `isinstance(experiment, MultiTypeExperiment)` checks with `experiment.default_trial_type is not None` feature checks in: - `ax/orchestration/orchestrator.py` - `ax/fb/realtime/utils.py` - `ax/fb/axolotl/utils_chronos.py` Replace `assert_is_instance(experiment, MultiTypeExperiment).method()` calls with direct calls on base `Experiment` (which now has `runner_for_trial_type` and `metrics_for_trial_type`). This is Phase 6 of the MultiTypeExperiment pull-up plan. Differential Revision: D94999020
mpolson64
force-pushed
the
export-D94999020
branch
from
3086ccc to
f56e855
Compare
mpolson64
added a commit
to mpolson64/Ax
that referenced
this pull request
Mar 16, 2026
…acebook#5001) Summary: Pull Request resolved: facebook#5001 Move `filter_trials_by_type` and `get_trial_indices_for_statuses` from `multi_type_experiment.py` to `experiment.py`, with backward-compatible re-exports from the old module. Replace all `isinstance(experiment, MultiTypeExperiment)` checks with `experiment.default_trial_type is not None` feature checks in: - `ax/orchestration/orchestrator.py` - `ax/fb/realtime/utils.py` - `ax/fb/axolotl/utils_chronos.py` Replace `assert_is_instance(experiment, MultiTypeExperiment).method()` calls with direct calls on base `Experiment` (which now has `runner_for_trial_type` and `metrics_for_trial_type`). This is Phase 6 of the MultiTypeExperiment pull-up plan. Differential Revision: D94999020
meta-codesync
Bot
changed the title
mpolson64
force-pushed
the
export-D94999020
branch
from
f56e855 to
2c5e8a8
Compare
mpolson64
added a commit
to mpolson64/Ax
that referenced
this pull request
Mar 17, 2026
…acebook#5001) Summary: Pull Request resolved: facebook#5001 Move `filter_trials_by_type` and `get_trial_indices_for_statuses` from `multi_type_experiment.py` to `experiment.py`, with backward-compatible re-exports from the old module. Replace all `isinstance(experiment, MultiTypeExperiment)` checks with `experiment.default_trial_type is not None` feature checks in: - `ax/orchestration/orchestrator.py` - `ax/fb/realtime/utils.py` - `ax/fb/axolotl/utils_chronos.py` Replace `assert_is_instance(experiment, MultiTypeExperiment).method()` calls with direct calls on base `Experiment` (which now has `runner_for_trial_type` and `metrics_for_trial_type`). This is Phase 6 of the MultiTypeExperiment pull-up plan. Differential Revision: D94999020
mpolson64
force-pushed
the
export-D94999020
branch
from
2c5e8a8 to
b11ceca
Compare
mpolson64
added a commit
to mpolson64/Ax
that referenced
this pull request
Mar 17, 2026
…acebook#5001) Summary: Pull Request resolved: facebook#5001 Move `filter_trials_by_type` and `get_trial_indices_for_statuses` from `multi_type_experiment.py` to `experiment.py`, with backward-compatible re-exports from the old module. Replace all `isinstance(experiment, MultiTypeExperiment)` checks with `experiment.default_trial_type is not None` feature checks in: - `ax/orchestration/orchestrator.py` - `ax/fb/realtime/utils.py` - `ax/fb/axolotl/utils_chronos.py` Replace `assert_is_instance(experiment, MultiTypeExperiment).method()` calls with direct calls on base `Experiment` (which now has `runner_for_trial_type` and `metrics_for_trial_type`). This is Phase 6 of the MultiTypeExperiment pull-up plan. Differential Revision: D94999020
mpolson64
force-pushed
the
export-D94999020
branch
from
b11ceca to
27de1d6
Compare
Summary: This is Phase 1 of moving MultiTypeExperiment features into the base Experiment class, enabling eventual deprecation of MultiTypeExperiment. Adds `_trial_type_to_metric_names: dict[str, set[str]]` to Experiment — a mapping from trial type to the set of metric names relevant to that type. This is the natural complement to the existing `_trial_type_to_runner` dict. Along with it, adds the following properties and methods to Experiment: - `trial_type_to_metric_names`: read-only property (shallow copy) - `metric_to_trial_type`: computed inverse mapping, with optimization config metrics pinned to `default_trial_type` - `metrics_for_trial_type(trial_type)`: returns Metric objects for a given trial type - `default_trials`: returns trial indices matching the default type MultiTypeExperiment is updated to populate `_trial_type_to_metric_names` alongside `_metric_to_trial_type` in all mutation paths (init, optimization_config setter, add/update/remove tracking metric). The redundant MTE overrides for `metric_to_trial_type`, `metrics_for_trial_type`, `default_trials`, and `default_trial_type` are removed — they are now inherited from the base class. The JSON decoder is updated to rebuild `_trial_type_to_metric_names` from `_metric_to_trial_type` during deserialization for backward compatibility. Differential Revision: D94970662
Summary: Phase 2 of moving MultiTypeExperiment features into base Experiment. Updates the base Experiment metric management methods (`add_metric`, `update_metric`, `remove_metric`) to accept an optional `trial_type` parameter. When provided, metrics are associated with the specified trial type in `_trial_type_to_metric_names`. The `__init__` and `optimization_config` setter also now register metrics when `default_trial_type` is set. The deprecated wrappers (`add_tracking_metric`, `add_tracking_metrics`, `update_tracking_metric`) now accept and pass through `trial_type` and `canonical_name` parameters. On MultiTypeExperiment, overrides are simplified to delegate to the base class methods: - `add_tracking_metric` delegates to `self.add_metric()` - `add_tracking_metrics` override removed (inherited from base) - `update_tracking_metric` delegates to `self.update_metric()` - `remove_tracking_metric` replaced with `remove_metric` override Differential Revision: D94986440
…iment Summary: Phase 3 of moving MultiTypeExperiment features into base Experiment. Moves `add_trial_type` and `update_runner` from MultiTypeExperiment to the base Experiment class, making them available to all experiments. Updates `supports_trial_type` to unify the logic: for multi-type experiments (where `default_trial_type` is set), only trial types registered in `_trial_type_to_runner` are supported. For single-type experiments, `None` is supported along with `SHORT_RUN` and `LONG_RUN` for backward compatibility with generation strategies that use those trial types. Removes the corresponding overrides from MultiTypeExperiment — all three methods are now inherited from the base class. Differential Revision: D94988577
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
…ding Summary: Phase 5 of moving MultiTypeExperiment features into base Experiment. JSON storage: - `experiment_to_dict` now serializes `default_trial_type` and `_trial_type_to_metric_names` (sets converted to sorted lists for JSON determinism). - `experiment_from_json` pops and restores `_trial_type_to_metric_names` (lists back to sets), defaulting to None for backward compat with old JSON. - `multi_type_experiment_from_json` pops `_trial_type_to_metric_names` (now present via `experiment_to_dict`) to prevent it being passed as a constructor kwarg. SQA storage: - Encoder replaces the `isinstance(experiment, MultiTypeExperiment)` main branch with `experiment._default_trial_type is not None`, using the base class `metric_to_trial_type` computed property instead of MTE-specific `_metric_to_trial_type`. Keeps `isinstance` only for MTE-specific fields (`Keys.SUBCLASS` marker, `_metric_to_canonical_name`). - No decoder changes needed — `_trial_type_to_metric_names` is already populated correctly through the `__init__` + `add_tracking_metric` call chain. Differential Revision: D94994120
mpolson64
added a commit
to mpolson64/Ax
that referenced
this pull request
Mar 25, 2026
…acebook#5001) Summary: Move `filter_trials_by_type` and `get_trial_indices_for_statuses` from `multi_type_experiment.py` to `experiment.py`, with backward-compatible re-exports from the old module. Replace all `isinstance(experiment, MultiTypeExperiment)` checks with `experiment.default_trial_type is not None` feature checks in: - `ax/orchestration/orchestrator.py` - `ax/fb/realtime/utils.py` - `ax/fb/axolotl/utils_chronos.py` Replace `assert_is_instance(experiment, MultiTypeExperiment).method()` calls with direct calls on base `Experiment` (which now has `runner_for_trial_type` and `metrics_for_trial_type`). This is Phase 6 of the MultiTypeExperiment pull-up plan. Differential Revision: D94999020
mpolson64
force-pushed
the
export-D94999020
branch
from
27de1d6 to
74abf89
Compare
…acebook#5001) Summary: Pull Request resolved: facebook#5001 Move `filter_trials_by_type` and `get_trial_indices_for_statuses` from `multi_type_experiment.py` to `experiment.py`, with backward-compatible re-exports from the old module. Replace all `isinstance(experiment, MultiTypeExperiment)` checks with `experiment.default_trial_type is not None` feature checks in: - `ax/orchestration/orchestrator.py` - `ax/fb/realtime/utils.py` - `ax/fb/axolotl/utils_chronos.py` Replace `assert_is_instance(experiment, MultiTypeExperiment).method()` calls with direct calls on base `Experiment` (which now has `runner_for_trial_type` and `metrics_for_trial_type`). This is Phase 6 of the MultiTypeExperiment pull-up plan. Differential Revision: D94999020
mpolson64
force-pushed
the
export-D94999020
branch
from
74abf89 to
642178d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
Move
filter_trials_by_typeandget_trial_indices_for_statusesfrommulti_type_experiment.pytoexperiment.py, with backward-compatiblere-exports from the old module.
Replace all
isinstance(experiment, MultiTypeExperiment)checks withexperiment.default_trial_type is not Nonefeature checks in:ax/orchestration/orchestrator.pyax/fb/realtime/utils.pyax/fb/axolotl/utils_chronos.pyReplace
assert_is_instance(experiment, MultiTypeExperiment).method()calls with direct calls on base
Experiment(which now hasrunner_for_trial_typeandmetrics_for_trial_type).This is Phase 6 of the MultiTypeExperiment pull-up plan.
Differential Revision: D94999020