Ax: [Question] TrialStatus not marked as completed when using fetch data. `trial.status.expecting_data` is True

[BCJuan]

Hello,

I have some doubts about the BO procedure of appending new trials, running them, fetching the data and updating.

First, after doing experiment.fetch_data(), the trial.status.expecting_data value is still True, which makes the trial run once again every time experiment.fetch_data() is ran.

In my case I have the following procedure:

1. Define quasirandom model sobol = get_sobol(exp.search_space)

2. I have this function

def model_loop(model, batch_size, experiment):
    """
    Makes a loop of random generation and fetching
    """
    if batch_size > 1:
        n = batch_size
        _ = experiment.new_batch_trial(
                generator_run=model.gen(n=n)).run()
    else:
        n = 1
        _ = experiment.new_trial(
                generator_run=model.gen(n=n)).run()

    data = experiment.fetch_data()
    return experiment, model, data

Make a pass of sobol model of model_loop function:

        exp, sobol, data = model_loop(
            sobol,
            batch,
            experiment)

1. Define botorch model and use botorch model the same way, in a loop like

for i in range(n):
        exp, botorch, trial_number, data = model_loop(
            get_botorch(experiment=exp, data=data),
            batch_size,
            exp)

The thing is that the trials are not marked as TrialStatus.COMPLETED after calling experiment.fetch_data(). Should I mark them myself? Also, although I mark them as completed, its status of expecting_data is True, hence every time fetch_data is run ALL, and literally all trials are ran once again. How can this be solved?

I see in class TrialStatus that the property expecting_datais True for both TrialStatus RUNNINGand COMPLETED. Shouldn't it be false for COMPLETED?

Maybe I have missed something because I do not understand the current behaviour.

Thanks

[ldworkin]

Hi @BCJuan ! A few things:

1) You're correct that experiment.fetch_data() doesn't mark trials as completed. That's because the Developer API is built to be as flexible as possible, and so you have to explicitly manage the lifecycle of your trials yourself. So yup, go ahead and mark them completed yourself.

2) The behavior of expecting_data is correct. Again, for maximum flexibility in the developer API, we want it to be possible to fetch data for completed trials as well as running trials. This is desired for several applications, like A/B testing.

Can you give us more information about how your data fetching works? That might help us figure out how best to support your use case. You might also want to take a look at the Service or Loop APIs, which seem like it might be better suited to you, but we'll have a better sense of that when we know how your metrics are implemented.

cc @lena-kashtelyan

[BCJuan]

Hi @ldworkin,

I am using metrics as in the tutorial Building Blocks of Ax,

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,
                "mean": (params["x1"] + 2*params["x2"] - 7)**2 + (2*params["x1"] + params["x2"] - 5)**2,
                "sem": 0.0,
                "trial_index": trial.index,
            })
        return Data(df=pd.DataFrame.from_records(records))

The thing is that I have multiple metrics, which I think is the worst obstacle to try out those APIs. I change experiment and metric parameters on the fly but I thing this could be managed in the service APIs. However, If I am not wrong the function optimize in managed_loop and the method create_experiment of 'AxClient work with single objectives.

I have been looking at the implementation of AxClient and maybe could pass my experiment (of class Experiment) and avoid creating an experiment. I'll try it out.

Any recommendation is totally appreciated. Also, if you would like to know more information, please, do not doubt to ask.

[ldworkin]

Hi @BCJuan ,

A couple thoughts:

Even if you can't use the Loop API exactly as is, you might want to take a look at the code here, since it seems pretty similar to what you're trying to do, and might be helpful.

In particular, I think you could fix the issue you're running into by just fetching data for the most recent trial using trial.fetch_data(), rather than experiment.fetch_data(). That will avoid the problem of fetching data for all trials every time.

Also I'm curious -- can you elaborate on what you mean about having multiple objectives? Ax in general doesn't really support that right now, so I'm curious how you're making that work :)

[BCJuan]

Yes, I have turned to experiment._fetch_trial_data(trial_index). I am thinking now if it has the same effect on the experiment than trial.fetch_data(). I do something similar to the sample link you provided.

Regarding the multiple objective concept, I mean having different ouputs/performance metrics for the same system. I am using ScalarizedObjective which allows for having different metrics and making a sole objective from the composition of several. I thought it allowed fusing several metrics into one. Doesn't it work that way? IF not I am completely lost.

For example, in the function get_NEI from botorch_defaults, you have LinearMCObjective which makes a linear combination of objectives.

[ldworkin]

Oh yes! If you're using ScalarizedObjective then that's exactly right -- sorry for the confusion!

Yes, experiment._fetch_trial_data should be exactly the same as trial.fetch_data, and seems like the right thing to use in this case.

So it sounds like you're unblocked for now? Let us know if you have any further questions!

[BCJuan]

Yes, I am unblocked for now. Thank you very much. I think that is appropriate to close the issue.