Unable to attach trials to HierarchicalSearchSpace using service API

[ailitw]

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'}.

[lena-kashtelyan]

Hi @ailitw! You definitely can attach parameters, they'd just need to abide by the hierarchical structure of your search space. So for example, if your "cat1" parameter evaluates to "x1", the parameterization you attach as a trial should have the parameters that depend on that value –– {"x1"}, but not the ones that depend on other values –– {"x2", "x3"}. So the parameterization that will succeed will be: {"root": True, "cat1": "x1", "cat2": "x4", "x1": 0.5, "x4": 0.5,}. Does that make sense?

In other words, a parameter "x2" depending on "cat1" having the value of "x2" means that parameter "x2" is not expected to appear in the parameterization unless "cat1" takes on the value of "x2" in that parameterization.

Let me know if that does not resolve the issue for you!

[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]

Oh I'm so sorry I misunderstood you. This is a bug and should be an easy fix, we'll take care of it shortly : )

[lena-kashtelyan]

The fix for this (and another issue with using HSS in combination with attaching trials) has been shipped; we'll close this issue once a stable version of Ax with the fix has been released.

[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)!