Unable to attach trials to HierarchicalSearchSpace using service API

Hey, I am testing the recently added HierarchicalSearchSpace (https://github.com/facebook/Ax/issues/140) for my use case. However, I discovered that it is not possible to attach trials using the service API. Below is a simple example and all the formats I tried. I believe the issue lies in AxClient()._validate_search_space_membership.

import numpy as np

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


def evaluate(parameters):
    x = np.array([parameters.get(f"x{i+1}", 0) for i in range(6)])
    return {"hartmann6": (hartmann6(x), 0.0)}


parameters = [
    {
        "name": "root",
        "type": "fixed",
        "value": True,
        "dependents": {True: ["cat1", "cat2"]},
    },
    {
        "name": "cat1",
        "type": "choice",
        "values": ["x1", "x2", "x3"],
        "dependents": {"x1": ["x1"], "x2": ["x2"], "x3": ["x3"]},
    },
    {
        "name": "cat2",
        "type": "choice",
        "values": ["x4", "x5", "x6"],
        "dependents": {"x4": ["x4"], "x5": ["x5"], "x6": ["x6"]},
    },
] + [
    {
        "name": f"x{i+1}",
        "type": "range",
        "bounds": [0.0, 1.0],
        "value_type": "float",
    }
    for i in range(6)
]

ax_client = AxClient()
ax_client.create_experiment(
    name="test_experiment",
    parameters=parameters,
    objective_name="hartmann6",
    minimize=True,
)

ax_client.attach_trial({"root": True, "cat1": "x1", "cat2": "x4", "x1": 0.5, "x4": 0.5})
# KeyError: 'x2'

ax_client.attach_trial({"root": True, "cat1": "x1", "cat2": "x4", 
                                      "x1": 0.5, "x2": None, "x3": None, "x4": 0.5, "x5": None, "x6": None})
# ValueError: None is not a valid value for parameter 
# RangeParameter(name='x2', parameter_type=FLOAT, range=[0.0, 1.0])

ax_client.attach_trial({"root": True, "cat1": "x1", "cat2": "x4", 
                                      "x1": 0.5, "x2": 0.0, "x3": 0.0, "x4": 0.5, "x5": 0.0, "x6": 0.0,})
# ValueError: Parameterization violates the hierarchical structure of the searchspace; 
# cast version would have parameters: {'cat2', 'cat1', 'x4', 'root', 'x1'}, but full version contains 
# parameters: {'cat2', 'cat1', 'x4', 'root', 'x6', 'x1', 'x5', 'x2', 'x3'}.

[ailitw]

Hi @lena-kashtelyan , thanks for the response! That is exactly how I expected attaching trials to work. Unfortunately, I’m getting a KeyError for that parametrization

ax_client.attach_trial({"root": True, "cat1": "x1", "cat2": "x4", "x1": 0.5, "x4": 0.5})
# KeyError: 'x2'

The error comes from this line in _validate_search_space_membership https://github.com/facebook/Ax/blob/aad4cda670f2947fb5169df91082e81a25bff328/ax/service/ax_client.py#L1697 It seems to loop through ALL parameters (not just the ones that match the hierarchical structure) in the search space to check that the supplied parameter values match.

[lena-kashtelyan]

New stable version of Ax, 0.2.7, is now out, and it should resolve this –– example in https://github.com/facebook/Ax/issues/1025#issuecomment-1188850721 should be working now. @ailitw, please let us know if it still doesn’t (and please reopen the issue in that case, so we see your response)!