Add AsExternalType to OSS Vizier proto following #1185

PiperOrigin-RevId: 691261698

Author: xingyousong

vizier/_src/pyvizier/oss/proto_converters.py

@@ -35,6 +35,8 @@
 
 ScaleType = parameter_config.ScaleType
 _ScaleTypePb2 = study_pb2.StudySpec.ParameterSpec.ScaleType
+ExternalType = parameter_config.ExternalType
+_ExternalTypePb2 = study_pb2.StudySpec.ParameterSpec.ExternalType
 ParameterType = parameter_config.ParameterType
 MonotypeParameterSequence = parameter_config.MonotypeParameterSequence
 
@@ -89,6 +91,26 @@ def from_proto(cls, proto: _ScaleTypePb2) -> ScaleType:
     return cls._proto_to_pyvizier[proto]
 
 
+class _ExternalTypeMap:
+  """Proto converter for external type."""
+
+  _pyvizier_to_proto = dict([
+      (ExternalType.INTERNAL, _ExternalTypePb2.AS_INTERNAL),
+      (ExternalType.BOOLEAN, _ExternalTypePb2.AS_BOOLEAN),
+      (ExternalType.INTEGER, _ExternalTypePb2.AS_INTEGER),
+      (ExternalType.FLOAT, _ExternalTypePb2.AS_FLOAT),
+  ])
+  _proto_to_pyvizier = {v: k for k, v in _pyvizier_to_proto.items()}
+
+  @classmethod
+  def to_proto(cls, pyvizier: ExternalType) -> _ExternalTypePb2:
+    return cls._pyvizier_to_proto[pyvizier]
+
+  @classmethod
+  def from_proto(cls, proto: _ExternalTypePb2) -> ExternalType:
+    return cls._proto_to_pyvizier[proto]
+
+
 class ParameterConfigConverter:
   """Converter for ParameterConfig."""
 
@@ -173,7 +195,8 @@ def from_proto(
     Raises:
       ValueError: See the "strict_validtion" arg documentation.
     """
-    feasible_values = []
+    bounds = None
+    feasible_values = None
     oneof_name = proto.WhichOneof('parameter_value_spec')
     if oneof_name == 'integer_value_spec':
       bounds = (
@@ -210,6 +233,10 @@ def from_proto(
     if proto.scale_type:
       scale_type = _ScaleTypeMap.from_proto(proto.scale_type)
 
+    external_type = None
+    if proto.external_type:
+      external_type = _ExternalTypeMap.from_proto(proto.external_type)
+
     try:
       config = parameter_config.ParameterConfig.factory(
           name=proto.parameter_id,
@@ -218,6 +245,7 @@ def from_proto(
           children=children,
           scale_type=scale_type,
           default_value=default_value,
+          external_type=external_type,
       )
     except ValueError as e:
       raise ValueError(
@@ -258,10 +286,8 @@ def _set_child_parameter_configs(
     parent_proto.ClearField('conditional_parameter_specs')
     for child_pair in children:
       if len(child_pair) != 2:
-        raise ValueError(
-            """Each element in children must be a tuple of
-            (Sequence of valid parent values,  ParameterConfig)"""
-        )
+        raise ValueError("""Each element in children must be a tuple of
+            (Sequence of valid parent values,  ParameterConfig)""")
 
     logging.debug(
         '_set_child_parameter_configs: parent_proto=%s, children=%s',
@@ -316,6 +342,8 @@ def to_proto(
       proto.scale_type = _ScaleTypeMap.to_proto(pc.scale_type)
     if pc.default_value is not None:
       cls._set_default_value(proto, pc.default_value)
+    if pc.external_type is not None:
+      proto.external_type = _ExternalTypeMap.to_proto(pc.external_type)
 
     cls._set_child_parameter_configs(proto, pc)
     return proto

vizier/_src/pyvizier/oss/proto_converters_test.py

@@ -17,9 +17,7 @@
 """Tests for proto_converters."""
 
 from absl import logging
-
 import attr
-
 from vizier._src.pyvizier.oss import proto_converters
 from vizier._src.pyvizier.pythia import study
 from vizier._src.pyvizier.shared import parameter_config as pc
@@ -365,6 +363,7 @@ def testDiscreteConfigToProto(self):
         'name',
         feasible_values=feasible_values,
         scale_type=pc.ScaleType.LOG,
+        external_type=pc.ExternalType.INTEGER,
         default_value=2,
     )
 
@@ -376,6 +375,10 @@ def testDiscreteConfigToProto(self):
         proto.scale_type,
         study_pb2.StudySpec.ParameterSpec.ScaleType.UNIT_LOG_SCALE,
     )
+    self.assertEqual(
+        proto.external_type,
+        study_pb2.StudySpec.ParameterSpec.ExternalType.AS_INTEGER,
+    )
 
 
 class ParameterConfigConverterFromProtoTest(absltest.TestCase):

vizier/_src/pyvizier/oss/study_config_test.py

@@ -384,6 +384,57 @@ def testPyTrialToDict(self):
     self.assertIsInstance(parameters['batch_size'], int)
     self.assertIsInstance(parameters['floating_point_param'], float)
 
+  def testTrialToDictWithExternalType(self):
+    """Test conversion when external types are not specified."""
+    proto = study_pb2.StudySpec()
+    proto.parameters.add(
+        parameter_id='learning_rate',
+        double_value_spec=study_pb2.StudySpec.ParameterSpec.DoubleValueSpec(
+            min_value=1e-4, max_value=0.1),
+        scale_type=study_pb2.StudySpec.ParameterSpec.ScaleType.UNIT_LOG_SCALE)
+    proto.parameters.add(
+        parameter_id='batch_size',
+        discrete_value_spec=study_pb2.StudySpec.ParameterSpec.DiscreteValueSpec(
+            values=[1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0]),
+        external_type=study_pb2.StudySpec.ParameterSpec.ExternalType.AS_INTEGER)
+    proto.parameters.add(
+        parameter_id='training_steps',
+        discrete_value_spec=study_pb2.StudySpec.ParameterSpec.DiscreteValueSpec(
+            values=[1000.0, 10000.0]),
+        external_type=study_pb2.StudySpec.ParameterSpec.ExternalType.AS_INTEGER)
+    proto.observation_noise = study_pb2.StudySpec.ObservationNoise.HIGH
+    proto.metrics.add(
+        metric_id='loss', goal=study_pb2.StudySpec.MetricSpec.MINIMIZE)
+
+    trial_proto = study_pb2.Trial()
+    trial_proto.id = str(1)
+    trial_proto.parameters.add(
+        parameter_id='batch_size', value=struct_pb2.Value(number_value=128.0))
+    trial_proto.parameters.add(
+        parameter_id='learning_rate',
+        value=struct_pb2.Value(number_value=1.2137854406366652E-4))
+    trial_proto.parameters.add(
+        parameter_id='training_steps',
+        value=struct_pb2.Value(number_value=10000.0))
+
+    py_study_config = vz.StudyConfig.from_proto(proto)
+    self.assertEqual(
+        py_study_config.observation_noise, vz.ObservationNoise.HIGH
+    )
+    parameters = py_study_config.trial_parameters(trial_proto)
+    self.assertEqual(
+        py_study_config.observation_noise, vz.ObservationNoise.HIGH
+    )
+    expected = {
+        'batch_size': 128,
+        'learning_rate': 1.2137854406366652E-4,
+        'training_steps': 10000.0
+    }
+    self.assertEqual(expected, parameters)
+    self.assertIsInstance(parameters['learning_rate'], float)
+    self.assertIsInstance(parameters['batch_size'], int)
+    self.assertIsInstance(parameters['training_steps'], int)
+
   def testTrialToDictWithoutExternalType(self):
     """Test conversion when external types are not specified."""
     proto = study_pb2.StudySpec()

vizier/_src/service/study.proto

@@ -289,6 +289,24 @@ message StudySpec {
     // Leave unset for `CATEGORICAL` parameters.
     ScaleType scale_type = 6;
 
+    // This is a place where the Vizier client can note the representation it
+    // presents to its callers.
+    // e.g. boolean can be represented inside Vizier in several ways (e.g.
+    // CATEGORICAL, INTEGER, or DOUBLE). Or, to represent a python range like
+    // range(10, 100, 10), you need to use an internal DOUBLE representation and
+    // use the external AS_INTEGER representation.
+    //
+    // NOTE: This field is not examined or modified by the Vizier service.
+    // NOTE: Not all combinations of ExternalType and ParameterType make sense.
+    enum ExternalType {
+      AS_INTERNAL = 0;
+      AS_BOOLEAN = 1;
+      AS_INTEGER = 2;
+      AS_FLOAT = 3;
+    }
+
+    ExternalType external_type = 7;
+
     // Represents a parameter spec with condition from its parent parameter.
     message ConditionalParameterSpec {
       // The spec for a conditional parameter.
@@ -351,7 +369,7 @@ message StudySpec {
   // See pyvizier.Algorithm for a list of native Vizier algorithms.
   string algorithm = 3;
 
-  // Use the deault early stopping policy.
+  // Use the default early stopping policy.
   reserved 4, 5, 8;
   message DefaultEarlyStoppingSpec {}