Ax: Loop API Batch Trial in parallel

[dqii]

I tried to run the loop API with batch trial to speed things up. I have multiple CPUs. I noticed it took about 90 seconds to run each iteration with a standard trial, but with I did a batch trial with 5 arms, it took about 7 minutes. This suggests to me that the 5 arms are not running in parallel.

Is it possible to get them to run in parallel with the Loop API? If not, is it possible with the service API or the developer API?

[lena-kashtelyan]

Hey, @dqii! You are right, the Loop API does not run arms in a batch in parallel, which is definitely something that needs fixing; for Loop API with batched trials, the evaluation function should be not for each separate arm (aka parameter configuration), but it should evaluate them all together –– this would allow you to distribute the evaluation across your CPUs, correct?

To do this via the Service API:
just request 5 1-arm trials, then evaluate them in parallel and log back data, then repeat until done –– something like this for one iteration of this repetition:

# Generate 5 trials.
trials_to_evaluate = {}
for _ in range(5):
    parameterization, trial_index = ax_client.get_next_trial()
    trials_to_evaluate[trial_index] = parameterization

# Computing results would actually evaluate your trials in parallel,
# you'd implement it in whatever way works for your setup.
results = { 
    trial_index: evaluate(parameterization)
    for trial_index, parameterization in trials_to_evaluate
}

# Update them all with data at once.
for trial_index in results:
    ax_client.complete_trial(trial_index, results.get(trial_index)

Let me know if that helps!

[dqii]

This helps, thank you! I will take the first approach.

for Loop API with batched trials, the evaluation function should be not for each separate arm (aka parameter configuration), but it should evaluate them all together –– this would allow you to distribute the evaluation across your CPUs, correct?

I wasn't 100% sure what you meant by this -- just to clarify, I intended for my evaluation function f to only accept a single set of parameters p, but I thought that if arms=5, then f(p) would be called 5 times, once in each of 5 processes, with 5 different sets of parameters.

[lena-kashtelyan]

@dqii, I think this issue stems from the conceptual confusion between 1-arm Trial-s and BatchTrial-s. Batched trials are not just a set or arms evaluated in parallel; a set of arms constitutes a batch when it need to be deployed and evaluated together (for example, in A/B testing [1], when you are exposing users to multiple parameter configurations at the same time). Therefore in this case, an evaluation function that evaluates all arms in the batch simultaneously should be required.

The case of parallel execution of multiple arms, like the one you are describing, is still a case for 1-arm Trials, since each one is evaluated separately. So the first approach I suggested is right for you and aligns with the concepts of Ax correctly.

I will edit my comment above to not include the second suggested approach, since while it's possible to use it, it violates the concept of batched trials and we should therefore not suggest it.

Let me know if any confusion remains!

---

[1] A little more detail: batched trials are important is when non-stationarity of the data is a problem. That is, if your evaluation can be biased by temporal effects, novelty effects, etc. and you have to evaluate all your arms at exactly the same time for the same duration, which often occurs in A/B tests, but can occur in other situations too.

[lena-kashtelyan]

Closing this as it seems that we found the solution, but will note the batched trial evaluation in the Loop API as something we will need to take another pass over.