diff --git a/google/cloud/aiplatform/vizier/pyvizier/proto_converters.py b/google/cloud/aiplatform/vizier/pyvizier/proto_converters.py index 49862182c6..9072041e5d 100644 --- a/google/cloud/aiplatform/vizier/pyvizier/proto_converters.py +++ b/google/cloud/aiplatform/vizier/pyvizier/proto_converters.py @@ -222,15 +222,15 @@ def _set_child_parameter_configs( ) ) - if parent_proto.HasField("discrete_value_spec"): + if "discrete_value_spec" in parent_proto: conditional_parameter_spec.parent_discrete_values.values[ : ] = parent_values - elif parent_proto.HasField("categorical_value_spec"): + elif "categorical_value_spec" in parent_proto: conditional_parameter_spec.parent_categorical_values.values[ : ] = parent_values - elif parent_proto.HasField("integer_value_spec"): + elif "integer_value_spec" in parent_proto: conditional_parameter_spec.parent_int_values.values[:] = parent_values else: raise ValueError("DOUBLE type cannot have child parameters") diff --git a/tests/unit/aiplatform/test_vizier.py b/tests/unit/aiplatform/test_vizier.py index 60c95a5baf..b8070e0bd1 100644 --- a/tests/unit/aiplatform/test_vizier.py +++ b/tests/unit/aiplatform/test_vizier.py @@ -987,11 +987,15 @@ def test_measurement_back_to_back_conversion(self): class TestParameterConfigConverterToProto: def test_discrete_config_to_proto(self): feasible_values = (-1, 3, 2) + child_parameter_config = pyvizier.ParameterConfig.factory( + "child", bounds=(-1.0, 1.0) + ) parameter_config = pyvizier.ParameterConfig.factory( "name", feasible_values=feasible_values, scale_type=pyvizier.ScaleType.LOG, default_value=2, + children=[([-1], child_parameter_config)], ) proto = proto_converters.ParameterConfigConverter.to_proto(parameter_config) @@ -1002,6 +1006,43 @@ def test_discrete_config_to_proto(self): proto.scale_type == study_pb2.StudySpec.ParameterSpec.ScaleType.UNIT_LOG_SCALE ) + assert len(proto.conditional_parameter_specs) == 1 + + spec = proto.conditional_parameter_specs[0] + assert spec.parameter_spec.parameter_id == "child" + assert spec.parameter_spec.double_value_spec.min_value == -1.0 + assert spec.parameter_spec.double_value_spec.max_value == 1.0 + assert len(spec.parent_discrete_values.values) == 1 + assert spec.parent_discrete_values.values[0] == -1 + + def test_categorical_config_to_proto_with_children(self): + feasible_values = ("option_a", "option_b") + child_parameter_config = pyvizier.ParameterConfig.factory( + "child", bounds=(-1.0, 1.0) + ) + parameter_config = pyvizier.ParameterConfig.factory( + "name", + feasible_values=feasible_values, + children=[(["option_a"], child_parameter_config)], + ) + proto = proto_converters.ParameterConfigConverter.to_proto(parameter_config) + assert len(proto.conditional_parameter_specs) == 1 + spec = proto.conditional_parameter_specs[0] + assert len(spec.parent_categorical_values.values) == 1 + assert spec.parent_categorical_values.values[0] == "option_a" + + def test_integer_config_to_proto_with_children(self): + child_parameter_config = pyvizier.ParameterConfig.factory( + "child", bounds=(-1.0, 1.0) + ) + parameter_config = pyvizier.ParameterConfig.factory( + "name", bounds=(1, 10), children=[([6], child_parameter_config)] + ) + proto = proto_converters.ParameterConfigConverter.to_proto(parameter_config) + assert len(proto.conditional_parameter_specs) == 1 + spec = proto.conditional_parameter_specs[0] + assert len(spec.parent_int_values.values) == 1 + assert spec.parent_int_values.values[0] == 6 class TestParameterConfigConverterFromProto: