Automate selection of appropriate parameters for BoTorch components in Ax based on experiment and data size #674
Description
Hello again (I hope I am not causing too much trouble to the team :) ),
I am here to report a possible bug. Attaching many trials through the Service API consumes a lot of memory. Here you have an example: (Warning: A couple of machines started thrashing and then froze with the following code).
from ax.service.ax_client import AxClient
from ax.service.utils.instantiation import ObjectiveProperties
import itertools
def evaluate(args):
return {
'a': (5_000, 0.0),
}
ax_client = AxClient(random_seed=64)
ax_client.create_experiment(
name="ax_err",
parameters=[
{'name': 'p1', 'type': 'range', 'bounds': [0, 5000], 'value_type': 'int'},
{'name': 'p2', 'type': 'range', 'bounds': [0, 6000], 'value_type': 'int'},
{'name': 'p3', 'type': 'range', 'bounds': [0, 7000], 'value_type': 'int'},
],
objectives={
'a': ObjectiveProperties(minimize=True, threshold=10_000)
},
)
def range_float(stop, percent=10/100):
l = []
c = 0
while c < stop:
l.append(int(c))
c += percent * stop
return l
r_p1 = range_float(5000)
r_p2 = range_float(6000)
r_p3 = range_float(7000)
force_trials = []
for p1, p2, p3 in itertools.product(r_p1, r_p2, r_p3):
config, trial_index = ax_client.attach_trial({
'p1': p1,
'p2': p2,
'p3': p3
})
evaluations = evaluate(config)
ax_client.complete_trial(trial_index=trial_index, raw_data=evaluations)
for _ in range(15):
(config, trial_index) = ax_client.get_next_trial()
evaluations = evaluate(config)
ax_client.complete_trial(trial_index=trial_index, raw_data=evaluations)
I've replicated this issue with MOO. Reducing the percentage (e.g to 20/100) does not cause this behavior. If I had to guess, it does not appear to be a memory leak. I think that some internal operation of ax/botorch has very steep memory complexity.