Serialize _trial_type_to_metric_names and use base class for SQA encoding (#4999)#4999
Open
mpolson64 wants to merge 6 commits intofacebook:mainfrom
Open
Serialize _trial_type_to_metric_names and use base class for SQA encoding (#4999)#4999mpolson64 wants to merge 6 commits intofacebook:mainfrom
mpolson64 wants to merge 6 commits intofacebook:mainfrom
Conversation
|
@mpolson64 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D94994120. |
mpolson64
added a commit
to mpolson64/Ax
that referenced
this pull request
Mar 11, 2026
…ding (facebook#4999) Summary: Pull Request resolved: facebook#4999 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
ee9e15e to
c0fd160
Compare
c0fd160 to
9ffe05a
Compare
mpolson64
added a commit
to mpolson64/Ax
that referenced
this pull request
Mar 13, 2026
…ding (facebook#4999) Summary: Pull Request resolved: facebook#4999 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 16, 2026
…ding (facebook#4999) Summary: Pull Request resolved: facebook#4999 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. Reviewed By: lena-kashtelyan Differential Revision: D94994120
9ffe05a to
108da89
Compare
…pter Summary: D93520819 removed Metric instances from the OptimizationConfig which made it more verbose to get any individual Metric's signature, which we use for bookkeeping throughout the adapter stack. This diff ... Differential Revision: D96855090
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
108da89 to
ea4a184
Compare
mpolson64
added a commit
to mpolson64/Ax
that referenced
this pull request
Mar 17, 2026
…ding (facebook#4999) 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. Reviewed By: lena-kashtelyan Differential Revision: D94994120
…ding (facebook#4999) Summary: Pull Request resolved: facebook#4999 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. Reviewed By: lena-kashtelyan Differential Revision: D94994120
ea4a184 to
5b0ae76
Compare
mpolson64
added a commit
to mpolson64/Ax
that referenced
this pull request
Mar 25, 2026
…ding (facebook#4999) 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. Reviewed By: lena-kashtelyan Differential Revision: D94994120
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
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:
Phase 5 of moving MultiTypeExperiment features into base Experiment.
JSON storage:
experiment_to_dictnow serializesdefault_trial_typeand_trial_type_to_metric_names(sets converted to sorted lists for JSON determinism).experiment_from_jsonpops and restores_trial_type_to_metric_names(lists back to sets), defaulting to None for backward compat with old JSON.multi_type_experiment_from_jsonpops_trial_type_to_metric_names(now present viaexperiment_to_dict) to prevent it being passed as a constructor kwarg.SQA storage:
isinstance(experiment, MultiTypeExperiment)main branch withexperiment._default_trial_type is not None, using the base classmetric_to_trial_typecomputed property instead of MTE-specific_metric_to_trial_type. Keepsisinstanceonly for MTE-specific fields (Keys.SUBCLASSmarker,_metric_to_canonical_name)._trial_type_to_metric_namesis already populated correctly through the__init__+add_tracking_metriccall chain.Reviewed By: lena-kashtelyan
Differential Revision: D94994120