[Bug]: cannot save client when metrics are tracked

[chanansh]

What happened?

Bug: Adding tracking metrics breaks saving Client state

Description

When using client.configure_metrics() with IMetric objects to add tracking metrics (as documented in https://ax.dev/docs/recipes/tracking-metrics/), the client.save_to_json_file() method fails with a serialization error because IMetric objects are not registered in ENCODER_REGISTRY.

This prevents checkpointing/saving the Client state when tracking metrics are configured, which is a critical issue for long-running optimization experiments.

Steps to Reproduce

Run the following minimal reproducible example:

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 FAILS)
client.save_to_json_file("ax_client_checkpoint.json")

Expected Behavior

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

Actual 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.

Environment

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

Additional Context

• The tracking metrics functionality works correctly for completing trials
• The issue only occurs when trying to save the Client state
• This is a blocker for production use cases where checkpointing is essential
• The documentation at https://ax.dev/docs/recipes/tracking-metrics/ does not mention this limitation

Related Documentation

• Tracking Metrics: https://ax.dev/docs/recipes/tracking-metrics/
• JSON Storage: https://ax.dev/docs/recipes/experiment-to-json

Please provide a minimal, reproducible example of the unexpected behavior.

#!/usr/bin/env python3
"""
Minimal reproducible example: Adding tracking metrics breaks saving Client state.

Issue: When using client.configure_metrics() with IMetric objects, client.save_to_json_file()
fails with a serialization error because IMetric objects are not registered in ENCODER_REGISTRY.

This follows the example from: https://ax.dev/docs/recipes/tracking-metrics/

To reproduce:
    python ax_issue_reproduce.py
"""

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"),
])

print("✓ Client configured with tracking metrics")

# 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,
        }
    )
    print(f"✓ Completed trial {trial_index} with tracking metrics")

# Step 4: Try to save Client state (THIS FAILS)
print("\nAttempting to save Client state...")
try:
    client.save_to_json_file("ax_client_checkpoint.json")
    print("✓ Successfully saved Client state")
except Exception as e:
    print(f"✗ FAILED to save Client state:")
    print(f"  Error: {type(e).__name__}: {e}")
    print("\nThis is the bug: IMetric objects cannot be serialized to JSON.")
    print("The tracking metrics configured via configure_metrics() prevent checkpoint saving.")

Please paste any relevant traceback/logs produced by the example provided.

tracking scalar metrics (e.g. string, int, float) should be easy

Ax Version

1.2.1

Python Version

3.12

Operating System

linux

(Optional) Describe any potential fixes you've considered to the issue outlined above.

Possible Solutions

1. Register IMetric in ENCODER_REGISTRY with appropriate encoder/decoder functions
2. Store tracking metrics configuration separately and reconfigure on load
3. Update documentation to warn about this limitation and provide workarounds

Pull Request

None

Code of Conduct

• I agree to follow Ax's Code of Conduct

[chanansh]

@Cesar-Cardoso

[chanansh]

Found it. There is a bug in the documentation recipe. One does not need to use IMetric but just Metric.

from ax.core.metric import Metric
        tracking_metrics = [
            Metric(name="metric_1"),
            Metric(name="metric_2"),
        ]
        client.configure_metrics(tracking_metrics) 

[mgarrard]

Hey @chanansh sorry for the delay getting back here, and thanks for flagging this. I think using Metric for now is fine, but it is not technically "api" meaning it is a bit out of the wheelhouse of support. I'll put up a PR this evening to ensure IMetric is in our storage layer :)