[FEATURE REQUEST]: Enable requests of batches of experiments in Service API

[CompRhys]

Motivation

Currently the service API doesn't allow for batches of experiments to be requested unless they're sequential greedy strategies in which case it grabs experiments and feeds them into the X_pending state. There are some methods such as recently released qPO https://arxiv.org/abs/2410.06333 that are not compatible with this formulation. Otherwise (our current usecase) it would just be nice to be able to replicate the testing setup on the dev API where different arms are grouped into the same trial if they were requested at the same time.

Our lab does batch processing of samples.

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

add an n argument to AxClient.get_next_trial or introduce an equivalent get_next_trial_batch that enables this functionality.

Happy to do this myself assuming that an interface/API acceptable to Ax maintainers is agreed.

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

Just using the service API's serial trials and assigning other metadata to indicate that these candidates formed part of the same batch.

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

• I agree to follow Ax's Code of Conduct

[Balandat]

This is something we generally want to support in our new APIs - @mpolson64 knows more about the plans and timelines.

[lena-kashtelyan]

Hi @CompRhys! Are you sure you need batches or would AxClient.get_next_trials do? What is the reason why you want to use BatchTrial-s specifically?

[CompRhys]

We're designing in batches and so given that it exists in the Ax schema in dev API it would be nice to match that in the service API. We can do it in service API by just leaving X_pending and then defining the batch idx in our own data schema.

So it's not needed but would be a nice to have.

[lena-kashtelyan]

In Ax, BatchTrial-s have a pretty specific meaning, you can read more about it here: https://ax.dev/docs/core#trial-vs-batch-trial. My sense is that your experiment doesn't match that definition, because the evaluation of each arm is in fact independent from one another, they are just done in parallel, is that right?

If so, my suggestion would be to:

1. Use get_next_trials (not get_next_trial) to get a "batch" of trials all at once (we might need to correctly configure max_parallelism of the Ax' GenerationStrategy so it returns as many trials at once as you want –– let me know if you would like help with that),
2. use the metadata field of complete_trial (

   Ax/ax/service/ax_client.py

   Line 769 in 0f8660a

   metadata: dict[str, str | int] | None = None,

   ) to track the index of the batch on your end.

What do you think of this setup?

We can do it in service API by just leaving X_pending

I'm not sure I understood this bit: why are you finding yourself needing to interact with X_pending? : ) AxClient should be handling the passing of all pending features, for you.

[CompRhys]

I can't remember the details as I haven't looked at what my coworker worked up in a while but I believe we we're handling X_pending as we're not keeping the process alive or using the Ax backend. When we trigger we were manually recreating the object and then needed to do something to reassign X_pending. If we were to call get_next_trial many times in a row in the same process then that would work as you say. I am not sure whether get_next_trials was already in the API when I posted this.

[lena-kashtelyan]

get_next_trials was in the API for a long time, but not featured in tutorials unfortunately, so it makes only the most sense that you didn't stumble into it! LMK if it doesn't solve the issue, I'll be glad to help more.