fix: Fixed INTEGER and BOOL parameters casting issue. Fix conversion bug where INTEGER and BOOL parameters are not cast to the correct type.

PiperOrigin-RevId: 586483856

Author: vertex-sdk-bot

google/cloud/aiplatform/vizier/pyvizier/__init__.py

@@ -33,7 +33,7 @@
 except ImportError:
     raise ImportError(
         "Google-vizier is not installed, and is required to use Vizier client."
-        'Please install the SDK using "pip install google-vizier==0.0.4"'
+        'Please install the SDK using "pip install google-vizier"'
     )
 
 from google.cloud.aiplatform.vizier.pyvizier.proto_converters import TrialConverter

google/cloud/aiplatform/vizier/pyvizier/proto_converters.py

@@ -7,6 +7,7 @@
 from google.protobuf import struct_pb2
 from google.protobuf import timestamp_pb2
 from google.cloud.aiplatform.compat.types import study as study_pb2
+from google.cloud.aiplatform.vizier.pyvizier import ExternalType
 from google.cloud.aiplatform.vizier.pyvizier import ScaleType
 from google.cloud.aiplatform.vizier.pyvizier import ParameterType
 from google.cloud.aiplatform.vizier.pyvizier import ParameterValue
@@ -122,12 +123,14 @@ def from_proto(
           ValueError: See the "strict_validtion" arg documentation.
         """
         feasible_values = []
+        external_type = ExternalType.INTERNAL
         oneof_name = proto._pb.WhichOneof("parameter_value_spec")
         if oneof_name == "integer_value_spec":
             bounds = (
                 int(proto.integer_value_spec.min_value),
                 int(proto.integer_value_spec.max_value),
             )
+            external_type = ExternalType.INTEGER
         elif oneof_name == "double_value_spec":
             bounds = (
                 proto.double_value_spec.min_value,
@@ -139,10 +142,17 @@ def from_proto(
         elif oneof_name == "categorical_value_spec":
             bounds = None
             feasible_values = proto.categorical_value_spec.values
+            # Boolean values are encoded as categoricals, check for the special
+            # hard-coded values.
+            boolean_values = ["False", "True"]
+            if sorted(list(feasible_values)) == boolean_values:
+                external_type = ExternalType.BOOLEAN
 
         default_value = None
         if getattr(proto, oneof_name).default_value:
             default_value = getattr(proto, oneof_name).default_value
+            if external_type == ExternalType.INTEGER:
+                default_value = int(default_value)
 
         if proto.conditional_parameter_specs:
             children = []
@@ -166,6 +176,7 @@ def from_proto(
                 children=children,
                 scale_type=scale_type,
                 default_value=default_value,
+                external_type=external_type,
             )
         except ValueError as e:
             raise ValueError(

tests/unit/aiplatform/test_vizier.py

@@ -1046,7 +1046,10 @@ def test_integer_config_to_proto_with_children(self):
 
 
 class TestParameterConfigConverterFromProto:
-    def test_creates_from_good_proto(self):
+    """Test ParameterConfigConverter.from_proto."""
+
+    def test_from_proto_discrete(self):
+        """Test from_proto."""
         proto = study_pb2.StudySpec.ParameterSpec(
             parameter_id="name",
             discrete_value_spec=study_pb2.StudySpec.ParameterSpec.DiscreteValueSpec(
@@ -1061,3 +1064,38 @@ def test_creates_from_good_proto(self):
         assert parameter_config.bounds == (1.0, 3.0)
         assert parameter_config.feasible_values == [1.0, 2.0, 3.0]
         assert parameter_config.default_value == 2.0
+        assert parameter_config.external_type == pyvizier.ExternalType.INTERNAL
+
+    def test_from_proto_integer(self):
+        """Test from_proto."""
+        proto = study_pb2.StudySpec.ParameterSpec(
+            parameter_id="name",
+            integer_value_spec=study_pb2.StudySpec.ParameterSpec.IntegerValueSpec(
+                default_value=2, min_value=1, max_value=3
+            ),
+        )
+
+        parameter_config = proto_converters.ParameterConfigConverter.from_proto(proto)
+
+        assert parameter_config.name == proto.parameter_id
+        assert parameter_config.type == pyvizier.ParameterType.INTEGER
+        assert parameter_config.bounds == (1, 3)
+        assert parameter_config.default_value == 2
+        assert parameter_config.external_type == pyvizier.ExternalType.INTEGER
+
+    def test_from_proto_bool(self):
+        """Test from_proto."""
+        proto = study_pb2.StudySpec.ParameterSpec(
+            parameter_id="name",
+            categorical_value_spec=study_pb2.StudySpec.ParameterSpec.CategoricalValueSpec(
+                default_value="True", values=["True", "False"]
+            ),
+        )
+
+        parameter_config = proto_converters.ParameterConfigConverter.from_proto(proto)
+
+        assert parameter_config.name == proto.parameter_id
+        assert parameter_config.type == pyvizier.ParameterType.CATEGORICAL
+        assert parameter_config.feasible_values == ["False", "True"]
+        assert parameter_config.default_value == "True"
+        assert parameter_config.external_type == pyvizier.ExternalType.BOOLEAN