Skip to content

Serialize _trial_type_to_metric_names and use base class for SQA encoding (#4999)#4999

Open
mpolson64 wants to merge 6 commits intofacebook:mainfrom
mpolson64:export-D94994120
Open

Serialize _trial_type_to_metric_names and use base class for SQA encoding (#4999)#4999
mpolson64 wants to merge 6 commits intofacebook:mainfrom
mpolson64:export-D94994120

Conversation

@mpolson64
Copy link
Copy Markdown
Contributor

@mpolson64 mpolson64 commented Mar 9, 2026

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

@meta-cla meta-cla Bot added the CLA Signed Do not delete this pull request or issue due to inactivity. label Mar 9, 2026
@meta-codesync
Copy link
Copy Markdown

meta-codesync Bot commented Mar 9, 2026

@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
@meta-codesync meta-codesync Bot changed the title Serialize _trial_type_to_metric_names and use base class for SQA encoding Serialize _trial_type_to_metric_names and use base class for SQA encoding (#4999) Mar 13, 2026
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
…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
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
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed Do not delete this pull request or issue due to inactivity. fb-exported meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant