Reloading experiment from the database in Service API does not reload the generation strategy if a custom one was set

[farrokhsiar]

I initiate a AxClient as it shown below:

_gs = GenerationStrategy( steps=[ GenerationStep( model=Models.SOBOL, num_trials=3, max_parallelism=max_parallelism, model_kwargs={"seed": 999}, model_gen_kwargs={}, ), GenerationStep( model=Models.GPEI, num_trials=-1, max_parallelism=self.max_parallelism, ), ] ) ax_client = AxClient( generation_strategy=_gs, db_settings=db_settings )

And then I create the experiment: ax_client.create_experiment(....), without any generation strategy arguments. However, after the first call to get_next_trial, I realized the generation strategy has gone back to default 6 Sobol iteration, and then GPEI. Am I missing anything in customizing the generation strategy?

[danielcohenlive]

Hi @farrokhsiar, thanks for the question. Happy to help! I tried to reproduce this problem, but I wasn't successful.

> [gr._model_key for trial in ax_client.experiment.trials.values() for gr in trial.generator_runs]
['Sobol', 'Sobol', 'Sobol', 'GPEI', 'GPEI']

If you sent me a more full reproduction of the problem I could take another look. I'm fudging some numbers like max_parallelism, so maybe there is a bug with some specific values, and otherwise basing my repro off the Service API Tutorial

[farrokhsiar]

Thanks, @danielcohenlive. Did you instantiate the same way I described? and if so, could you share the version of AX you used to reproduce the problem?

[danielcohenlive]

I'm on the latest. This is the complete code I ran. Try it out on the version you are running and let me know if you get different results.

from ax.service.ax_client import AxClient
from ax.modelbridge.generation_strategy import GenerationStrategy
from ax.modelbridge.generation_node import GenerationStep
from ax.utils.measurement.synthetic_functions import hartmann6
from ax.modelbridge.registry import Models

_gs = GenerationStrategy( steps=[ GenerationStep( model=Models.SOBOL, num_trials=3, model_kwargs={"seed": 999}, model_gen_kwargs={}, max_parallelism=3), GenerationStep( model=Models.GPEI, num_trials=-1, max_parallelism=3), ] ) 
ax_client = AxClient( generation_strategy=_gs )

ax_client.create_experiment(
    name="hartmann_test_experiment",
    parameters=[
        {
            "name": "x1",
            "type": "range",
            "bounds": [0.0, 1.0],
            "value_type": "float",  # Optional, defaults to inference from type of "bounds".
            "log_scale": False,  # Optional, defaults to False.
        },
        {
            "name": "x2",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "x3",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "x4",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "x5",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "x6",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
    ],
    objective_name="hartmann6",
    minimize=True,  # Optional, defaults to False.
    parameter_constraints=["x1 + x2 <= 2.0"],  # Optional.
    outcome_constraints=["l2norm <= 1.25"],  # Optional.
)
import numpy as np
def evaluate(parameters):
    x = np.array([parameters.get(f"x{i+1}") for i in range(6)])
    # In our case, standard error is 0, since we are computing a synthetic function.
    return {"hartmann6": (hartmann6(x), 0.0), "l2norm": (np.sqrt((x ** 2).sum()), 0.0)}

for i in range(5):
    parameters, trial_index = ax_client.get_next_trial()
    # Local evaluation here can be replaced with deployment to external system.
    ax_client.complete_trial(trial_index=trial_index, raw_data=evaluate(parameters))

[gr._model_key for trial in ax_client.experiment.trials.values() for gr in trial.generator_runs]

[farrokhsiar]

Yes! I ran and got the same results!

[danielcohenlive]

I'm going to this issue. Feel free to reopen if you still believe that the client is misbehaving in some situation.

[farrokhsiar]

I think I found the exact problem. I am using the AX in the Service mode with SQL backend. On the experiment retrieval, I use: ax_client = AxClient(db_settings=self.db_setting) ax_client.load_experiment_from_database(experiment_name=name)

The problem is the experiment name on the database has a generation strategy, but load_experiment_from_database does not retrieve any generation strategy, and therefore uses the default generation strategy. Unless I am saving the ax_clinet to SQL in the wrong way or retrieving it in the wrong way, this seems to be a bug.

[farrokhsiar]

To help with reproducing the issue, my current workflow is:

1. Create AxClient with a specific generation strategy with Postg storage
2. Retrieve the Client from DB
3. Perform a trial using get_next_trial So I do not ask for a trial and perform experiments before retrieval of the experiment.

[lena-kashtelyan]

Thank you very much for the repro steps, @farrokhsiar!

[danielcohenlive]

https://github.com/facebook/Ax/pull/931 should handle it.

[lena-kashtelyan]

This seems resolved; if it isn't, please feel free to reopen @farrokhsiar!