Как я могу запустить параллельные испытания с использованием AX?

Question

В соответствии с документацией AX (https://ax.dev/docs/api.html), в API-интерфейсе службы) говорится, что «испытания могут оцениваться параллельно. Как я могу этого добиться?

У меня есть 4 графических процессора, и я хотел бы оценить 4 разных набора гиперпараметров одновременно (4 тренировки одновременно).

from ax.service.ax_client import AxClient
from ax.utils.measurement.synthetic_functions import branin

ax = AxClient()
ax.create_experiment(
    name="branin_test_experiment",
    parameters=[
        {
            "name": "x1",
            "type": "range",
            "bounds": [-5.0, 10.0],
            "value_type": "float",
        },
        {
            "name": "x2",
            "type": "range",
            "bounds": [0.0, 10.0],
        },
    ],
    objective_name="branin",
    minimize=True,
)

for _ in range(15):
    parameters, trial_index = ax.get_next_trial()
    ax.complete_trial(trial_index=trial_index, raw_data=branin(parameters["x1"], parameters["x2"]))

best_parameters, metrics = ax.get_best_parameters()