Added IMetric and IRunner to JSON storage registries | Fix: Issue #4671

[sswez02]

Fixes #4671

Problem

When using client.configure_metrics() with IMetric objects to add tracking metrics, client.save_to_json_file() fails with a JSONEncodeError because IMetric and IRunner are not registered in the encoder/decoder registries.

Solution

Added IMetric and IRunner to both registries in ax/storage/json_store/registry.py:

• Added imports for IMetric and IRunner from ax.api.protocols
• Added IMetric: metric_to_dict to CORE_ENCODER_REGISTRY
• Added IRunner: runner_to_dict to CORE_ENCODER_REGISTRY
• Added "IMetric": IMetric to CORE_DECODER_REGISTRY
• Added "IRunner": IRunner to CORE_DECODER_REGISTRY

Testing

Verified using the reproduction script from the issue - save and load now work correctly with IMetric tracking metrics.

Run the following file:

from ax.api.client import Client
from ax.api.configs import RangeParameterConfig
from ax.api.protocols.metric import IMetric

# Step 1: Setup Client and configure experiment (as per Ax docs)
client = Client()

client.configure_experiment(
    parameters=[
        RangeParameterConfig(name="x1", parameter_type="float", bounds=(0.0, 1.0)),
        RangeParameterConfig(name="x2", parameter_type="float", bounds=(0.0, 1.0)),
    ]
)
client.configure_optimization(objective="my_objective")

# Step 2: Add tracking metrics (as per Ax docs)
# https://ax.dev/docs/recipes/tracking-metrics/
client.configure_metrics(
    [
        IMetric(name="my_tracking_metric_1"),
        IMetric(name="my_tracking_metric_2"),
    ]
)

# Step 3: Complete a trial with tracking metrics (as per Ax docs)
trials = client.get_next_trials(max_trials=1)
for trial_index, parameters in trials.items():
    client.complete_trial(
        trial_index=trial_index,
        raw_data={
            "my_objective": 0.5,
            "my_tracking_metric_1": 1.0,
            "my_tracking_metric_2": 2.0,
        },
    )

# Step 4: Try to save Client state (THIS WOULD FAIL, NOW WORKS)
client.save_to_json_file("ax_client_checkpoint.json")

# Step 5: Load it back to verify full round-trip works
loaded_client = Client.load_from_json_file("ax_client_checkpoint.json")
print("Load successful!")

Previous Behavior

The save operation fails with the following error:

ax.exceptions.storage.JSONEncodeError: Object IMetric('my_tracking_metric_1') passed to object_to_json (of type <class 'ax.api.protocols.metric.IMetric'>, module: ax.api.protocols.metric) is not registered with a corresponding encoder in ENCODER_REGISTRY. Please see our JSON storage tutorial (https://ax.dev/docs/recipes/experiment-to-json) for more details. The 'Customizing the Serialization Process' section will be relevant for saving Ax object subclasses.

New Behaviour

The client.save_to_json_file() successfully saves the Client state, including the configured tracking metrics, allowing the experiment to be resumed later.

Environment

Ax version: 1.2.1 (or latest)
Python version: 3.12
OS: Linux

[sswez02]

Hi @mpolson64, I noticed you were assigned to #4671 and mentioned putting up a fix. I submitted this PR to address it. Would you be able to take a look when you have a chance? Happy to make any changes needed. Thanks!