[Bug]: Custom metric issue

[leigedove]

What happened?

When using custom metrics, trial data for each Sobol trial is generated correctly and attached to the experiment, but the fetched data is still empty.

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

import os
import matplotlib.pyplot as plt
import numpy as np
import torch

from ax import Data, Experiment, ParameterType, RangeParameter, SearchSpace
from ax.core.metric import Metric
from ax.core.objective import Objective
from ax.core.optimization_config import OptimizationConfig
from ax.metrics.branin import BraninMetric
from ax.modelbridge.cross_validation import cross_validate
from ax.modelbridge.registry import Models
from ax.models.torch.botorch_modular.surrogate import Surrogate
from ax.runners.synthetic import SyntheticRunner
from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
import pandas as pd

SMOKE_TEST = os.environ.get("SMOKE_TEST")
torch.manual_seed(12345)  # To always get the same Sobol points
tkwargs = {
    "dtype": torch.double,
    "device": torch.device("cuda" if torch.cuda.is_available() else "cpu"),
}

# Define metric
class BoothMetric(Metric):
    def fetch_trial_data(self, trial):
        records = []
        for arm_name, arm in trial.arms_by_name.items():
            params = arm.parameters
            records.append(
                {
                    "arm_name": arm_name,
                    "metric_name": self.name,
                    "trial_index": trial.index,
                    # in practice, the mean and sem will be looked up based on trial metadata
                    # but for this tutorial we will calculate them
                    "mean": (params["x1"] + 2 * params["x2"] - 7) ** 2
                    + (2 * params["x1"] + params["x2"] - 5) ** 2,
                    "sem": 0.0,
                }
            )
        return Data(df=pd.DataFrame.from_records(records))

    def is_available_while_running(self) -> bool:
        return True

# Search space
search_space = SearchSpace(
    parameters=[
        RangeParameter(
            name=f"x{i}", parameter_type=ParameterType.FLOAT, lower=-5.0, upper=10.0
        )
        for i in range(25)
    ]
    + [
        RangeParameter(
            name=f"x{i + 25}",
            parameter_type=ParameterType.FLOAT,
            lower=0.0,
            upper=15.0,
        )
        for i in range(25)
    ]
)

# Optimization config
optimization_config = OptimizationConfig(
    objective=Objective(
        metric = BoothMetric(
            name="objective",
        ),
        minimize=True,
    )
)

N_INIT = 10
BATCH_SIZE = 3
N_BATCHES = 1 if SMOKE_TEST else 10

print(f"Doing {N_INIT + N_BATCHES * BATCH_SIZE} evaluations")

# Experiment
experiment = Experiment(
    name="saasbo_experiment",
    search_space=search_space,
    optimization_config=optimization_config,
    runner=SyntheticRunner(),
)

# Initial Sobol points
sobol = Models.SOBOL(search_space=experiment.search_space)
for _ in range(N_INIT):
    experiment.new_trial(sobol.gen(1)).run()

# Run SAASBO
data = experiment.fetch_data()
print(f"Initial data: {data.df}")

for i in range(N_BATCHES):
    model = Models.SAASBO(experiment=experiment, data=data)
    generator_run = model.gen(BATCH_SIZE)
    trial = experiment.new_batch_trial(generator_run=generator_run)
    trial.run()
    data = Data.from_multiple_data([data, trial.fetch_data()])

    new_value = trial.fetch_data().df["mean"].min()
    print(
        f"Iteration: {i}, Best in iteration {new_value:.3f}, Best so far: {data.df['mean'].min():.3f}"
    )

The error is:

Exception                                 Traceback (most recent call last)
Exception: []

The above exception was the direct cause of the following exception:

UnwrapError                               Traceback (most recent call last)

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

No response

Ax Version

0.4.0

Python Version

3.12.2

Operating System

linux

Code of Conduct

• [x] I agree to follow Ax's Code of Conduct

[mgrange1998]

Hi, thank you for opening the issue.

From the code snippet it looks like "mark_completed" is not being called prior to "fetch_data". In this section of the Building Blocks of Ax tutorial it says

"To fetch trial data, we need to run it and mark it completed. For most metrics in Ax, data is only available once the status of the trial is COMPLETED, since in real-worlds scenarios, metrics can typically only be fetched after the trial finished running."

Could you try calling "mark_completed", using the tutorial as a guide? If this doesn't work, please copy the full error trace as well as the printed statements so I can help debug further. Thanks!

[leigedove]

Hi, thank you for your response. However, after modifying my code based on the tutorial, there is an error: "AttributeError: 'Data' object has no attribute 'is_ok'" when running data = experiment.fetch_data().

And when I ran the code from the tutorial, the same error was shown.

It might be a problem with the ax_platform version. However, when I changed the version from 0.4.0 to 0.1.9, there were many environment issues, such as incompatibilities between gpytorch and botorch versions.

This is my code:

import os
import pandas as pd
import torch
from ax import Data, Experiment, ParameterType, RangeParameter, SearchSpace, Models
from ax.core.metric import Metric
from ax.core.objective import Objective
from ax.core.optimization_config import OptimizationConfig
from ax.runners.synthetic import SyntheticRunner

# Set seed and device
torch.manual_seed(12345)  # To always get the same Sobol points
tkwargs = {
    "dtype": torch.double,
    "device": torch.device("cuda" if torch.cuda.is_available() else "cpu"),
}

# Define BoothMetric
class BoothMetric(Metric):
    def fetch_trial_data(self, trial):
        records = []
        for arm_name, arm in trial.arms_by_name.items():
            params = arm.parameters
            records.append(
                {
                    "arm_name": arm_name,
                    "metric_name": self.name,
                    "trial_index": trial.index,
                    "mean": (params["x1"] + 2 * params["x2"] - 7) ** 2
                    + (2 * params["x1"] + params["x2"] - 5) ** 2,
                    "sem": 0.0,
                }
            )
        return Data(df=pd.DataFrame.from_records(records))

# Define the search space
search_space = SearchSpace(
    parameters=[
        RangeParameter(
            name=f"x{i}", parameter_type=ParameterType.FLOAT, lower=-5.0, upper=10.0
        )
        for i in range(25)
    ]
    + [
        RangeParameter(
            name=f"x{i + 25}",
            parameter_type=ParameterType.FLOAT,
            lower=0.0,
            upper=15.0,
        )
        for i in range(25)
    ]
)

# Define optimization configuration
optimization_config = OptimizationConfig(
    objective=Objective(
        metric=BoothMetric(name="objective"),
        minimize=True,
    )
)

# Initialize experiment parameters
N_INIT = 10
BATCH_SIZE = 3
N_BATCHES = 1 if os.environ.get("SMOKE_TEST") else 10

# Create experiment
experiment = Experiment(
    name="saasbo_experiment",
    search_space=search_space,
    optimization_config=optimization_config,
    runner=SyntheticRunner(),
)

# Generate initial Sobol points
sobol = Models.SOBOL(search_space=experiment.search_space)
for _ in range(N_INIT):
    trial = experiment.new_trial(sobol.gen(1))
    trial.run()
    trial.mark_completed()

# Check initial data
data = experiment.fetch_data()
if data.df.empty:
    print("Initial data from Sobol trials is empty. Check the objective function.")
else:
    print(f"Initial data: {data.df}")

# Run SAASBO optimization
for i in range(N_BATCHES):
    model = Models.SAASBO(experiment=experiment, data=data)
    generator_run = model.gen(BATCH_SIZE)
    trial = experiment.new_batch_trial(generator_run=generator_run)
    trial.run()
    trial.mark_completed()

    # Fetch new data and merge with existing data
    new_data = trial.fetch_data()
    data = Data.from_multiple_data([data, new_data])

    new_value = new_data.df["mean"].min()
    print(
        f"Iteration: {i}, Best in iteration {new_value:.3f}, Best so far: {data.df['mean'].min():.3f}"
    )

The full error trace: Traceback (most recent call last): File "/home/building_blocks.py", line 232, in data = experiment.fetch_data() ^^^^^^^^^^^^^^^^^^^^^^^ File "/home.../lib/python3.12/site-packages/ax/core/experiment.py", line 562, in fetch_data results = self._lookup_or_fetch_trials_results( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home...python3.12/site-packages/ax/core/experiment.py", line 649, in _lookup_or_fetch_trials_results ) = first_metric_of_group.fetch_data_prefer_lookup( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home...anaconda3/envs/charging/lib/python3.12/site-packages/ax/core/metric.py", line 352, in fetch_data_prefer_lookup contains_new_data = any( ^^^^ File "/home...anaconda3/envs/charging/lib/python3.12/site-packages/ax/core/metric.py", line 353, in result.is_ok() for result in fetched_trial_data.values() ^^^^^^^^^^^^ AttributeError: 'Data' object has no attribute 'is_ok'

[mgrange1998]

Thanks for the detailed explanation.

What is happening is that BoothMetric::fetch_trial_data is returning a Data object, when the Metric class interface expects that a MetricFetchResult object like Ok or Err is returned, which has the is_ok attribute.

Because the tutorial is for an old version, its BoothMetric implementation returns Data instead of a MetricFetchResult.

For your purposes, you can try fixing this by updating BoothMetric's fetch_trial_data to the following:

def fetch_trial_data(self, trial: BaseTrial) -> MetricFetchResult:
        records = []
        for arm_name, arm in trial.arms_by_name.items():
            params = arm.parameters
            records.append(
                {
                    "arm_name": arm_name,
                    "metric_name": self.name,
                    "trial_index": trial.index,
                    "mean": (params["x1"] + 2 * params["x2"] - 7) ** 2
                    + (2 * params["x1"] + params["x2"] - 5) ** 2,
                    "sem": 0.0,
                }
            )
        return Ok(value=Data(df=pd.DataFrame.from_records(records)))

You may also need to add error catching in case the data is not yet ready.

I'll make a note for us to update our tutorials and documentation with correct implementations of custom metrics. For now you can look off of current metric implementations like NoisyFunctionMetric.

Let me know if you encounter other blockers, thanks!

[mgrange1998]

I'll close out this issue for now- we've added the work to fix the tutorials to our backlog. Please open a new issue if any additional support is needed. Thank you!