[tune]Can't use AxClient with 1.0.0 - TypeError: 'AxClient' object is not iterable

[nscotto]

When using AxClient with ray 1.0.0 I get the error TypeError: 'AxClient' object is not iterable, this make the Ax unusable with ray at the moment (it might be an Ax problem though). I have seen this behavior in my project (too big to be posted), as well as using the example from Ax documentation. You can run the example from the Ax documentation to reproduce the error, even the documentation (an executed notebook) contains the error.

Here is the code from Ax documentation:

import logging
from ray import tune
from ray.tune import track
from ray.tune.suggest.ax import AxSearch
logger = logging.getLogger(tune.__name__)  
logger.setLevel(level=logging.CRITICAL)  # Reduce the number of Ray warnings that are not relevant here.

import torch
import numpy as np

from ax.plot.contour import plot_contour
from ax.plot.trace import optimization_trace_single_method
from ax.service.ax_client import AxClient
from ax.utils.notebook.plotting import render, init_notebook_plotting
from ax.utils.tutorials.cnn_utils import CNN, load_mnist, train, evaluate

init_notebook_plotting()

ax = AxClient(enforce_sequential_optimization=False)

ax.create_experiment(
    name="mnist_experiment",
    parameters=[
        {"name": "lr", "type": "range", "bounds": [1e-6, 0.4], "log_scale": True},
        {"name": "momentum", "type": "range", "bounds": [0.0, 1.0]},
    ],
    objective_name="mean_accuracy",
)

def train_evaluate(parameterization):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    train_loader, valid_loader, test_loader = load_mnist(data_path='~/.data')
    net = train(net=CNN(), train_loader=train_loader, parameters=parameterization, dtype=torch.float, device=device)
    track.log(
        mean_accuracy=evaluate(
            net=net,
            data_loader=valid_loader,
            dtype=torch.float,
            device=device,
        )
    )

tune.run(
    train_evaluate, 
    num_samples=30, 
    search_alg=AxSearch(ax),  # Note that the argument here is the `AxClient`.
    verbose=0,  # Set this level to 1 to see status updates and to 2 to also see trial results.
    # To use GPU, specify: resources_per_trial={"gpu": 1}.
)

best_parameters, values = ax.get_best_parameters()
best_parameters

means, covariances = values
means

render(
    plot_contour(
        model=ax.generation_strategy.model, param_x='lr', param_y='momentum', metric_name='mean_accuracy'
    )
)

# `plot_single_method` expects a 2-d array of means, because it expects to average means from multiple 
# optimization runs, so we wrap out best objectives array in another array.
best_objectives = np.array([[trial.objective_mean * 100 for trial in ax.experiment.trials.values()]])
best_objective_plot = optimization_trace_single_method(
    y=np.maximum.accumulate(best_objectives, axis=1),
    title="Model performance vs. # of iterations",
    ylabel="Accuracy",
)
render(best_objective_plot)

The full stack trace:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/AutoML/raytune_pytorch_cnn.py in <module>
     45     train_evaluate,
     46     num_samples=30,
---> 47     search_alg=AxSearch(ax),  # Note that the argument here is the `AxClient`.
     48     verbose=0,  # Set this level to 1 to see status updates and to 2 to also see trial results.
     49     # To use GPU, specify: resources_per_trial={"gpu": 1}.

~/anaconda3/envs/automl/lib/python3.7/site-packages/ray/tune/suggest/ax.py in __init__(self, space, metric, mode, parameter_constraints, outcome_constraints, ax_client, use_early_stopped_trials, max_concurrent)
    134
    135         if self._ax or self._space:
--> 136             self.setup_experiment()
    137
    138     def setup_experiment(self):

~/anaconda3/envs/automl/lib/python3.7/site-packages/ray/tune/suggest/ax.py in setup_experiment(self)
    158                 parameter_constraints=self._parameter_constraints,
    159                 outcome_constraints=self._outcome_constraints,
--> 160                 minimize=self._mode != "max")
    161         else:
    162             if any([

~/anaconda3/envs/automl/lib/python3.7/site-packages/ax/service/ax_client.py in create_experiment(self, parameters, name, objective_name, minimize, parameter_constraints, outcome_constraints, status_quo, overwrite_existing_experiment, experiment_type, choose_generation_strategy_kwargs)
    239             outcome_constraints=outcome_constraints,
    240             status_quo=status_quo,
--> 241             experiment_type=experiment_type,
    242         )
    243

~/anaconda3/envs/automl/lib/python3.7/site-packages/ax/service/utils/instantiation.py in make_experiment(parameters, name, objective_name, minimize, parameter_constraints, outcome_constraints, status_quo, experiment_type)
    292     without importing or instantiating any Ax classes."""
    293
--> 294     exp_parameters: List[Parameter] = [parameter_from_json(p) for p in parameters]
    295     status_quo_arm = None if status_quo is None else Arm(parameters=status_quo)
    296     parameter_map = {p.name: p for p in exp_parameters}

TypeError: 'AxClient' object is not iterable

[richardliaw]

@nscotto thanks a bunch for this bug report! Seems like we just need to add AxClient(ax_client=ax).

[nscotto]

You mean AxSearch not AxClient, but thanks that solves the issue :)