[FEATURE REQUEST]: Provide a minimal working example for client.predict()

[kurt-essential]

Motivation

I'm using Ax in a lightweight ask-tell setup and I struggled to figure out how to get client.predict to work. Eventually I figured out that running get_next_trials will progress to a predictive GenerationNode. (I looked into creating a runner but the example in the orchestration doc felt overly detailed). It would be helpful to have some minimal examples in the docs explaining concepts like this (i.e. I did not appreciate how get_next_trials changes state after trials are attached.)

Example:

from ax.api.client import Client
from ax.api.configs import RangeParameterConfig

client = Client()
parameters = [
    RangeParameterConfig(
        name=c, parameter_type="float", bounds=(0, 1)
        ) for c in ['x1','x2']
        ]
client.configure_experiment(name='test', parameters=parameters)
client.configure_optimization(objective='my_objective')
my_data = [
    ({"x1": 0.0, "x2": 0.0}, {"my_objective": 0.0}),
    ({"x1": 0.0, "x2": 1.0}, {"my_objective": 1.0}),
    ({"x1": 0.0, "x2": 0.5}, {"my_objective": 0.5}),
    ({"x1": 0.5, "x2": 0.5}, {"my_objective": 0.5}),
    ({"x1": 0.25, "x2": 0.25}, {"my_objective": 0.35}),
    ({"x1": 1.0, "x2": 1.0}, {"my_objective": 2.0}),
]

for my_parameters, raw_data in my_data:
    trial_index = client.attach_trial(
        parameters=my_parameters,
    )
    client.complete_trial(
        trial_index=trial_index, raw_data=raw_data,
    )

# predict works after getting 3 trials
trials = client.get_next_trials(max_trials=3)

# predict does not work after getting < 3 trials
# trials = client.get_next_trials(max_trials=2)

predict_data = [{"x1": 0.0, "x2": 0.0}, {"x1": 0.40, "x2": 0.30}]
client.predict(predict_data)

Describe the solution you'd like to see implemented in Ax.

add a minimal example to docs

Describe any alternatives you've considered to the above solution.

No response

Is this related to an existing issue in Ax or another repository? If so please include links to those Issues here.

No response

Code of Conduct

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

[mpolson64]

Hi @kurt-essential -- thanks for raising this with us; it is indeed a confusing user experience. As you discovered, under the hood the call to "predict" relies on using the surrogate model currently being used for Bayesian optimization. We use the GenerationStrategy to keep track of which surrogate model (and acquisition function, data preprocessing, etc.) Ax should be using, and its state changes as more trials are requested and data provided.

We should prioritize both adding a more informative error message here and adding a page on ax.dev, but longer term I believe it would be best to allow Ax to make predictions using some generic GP if a user calls predict and when the current GenerationNode does not have a suitable surrogate. This will take some work on our end, but I will do my best to get it prioritized and update this issue when it is ready.