From ed8cb966c01b65c7d8015adb91d3da13fc323e32 Mon Sep 17 00:00:00 2001 From: Xingyou Song Date: Tue, 29 Oct 2024 20:24:21 -0700 Subject: [PATCH] Add AsExternalType to OSS Vizier proto following https://github.com/google/vizier/issues/1185 PiperOrigin-RevId: 691261698 --- vizier/_src/pyvizier/oss/proto_converters.py | 38 ++++++++++++-- .../pyvizier/oss/proto_converters_test.py | 7 ++- vizier/_src/pyvizier/oss/study_config_test.py | 51 +++++++++++++++++++ vizier/_src/service/study.proto | 20 +++++++- 4 files changed, 108 insertions(+), 8 deletions(-) diff --git a/vizier/_src/pyvizier/oss/proto_converters.py b/vizier/_src/pyvizier/oss/proto_converters.py index af6ffac05..59d0139ee 100644 --- a/vizier/_src/pyvizier/oss/proto_converters.py +++ b/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 diff --git a/vizier/_src/pyvizier/oss/proto_converters_test.py b/vizier/_src/pyvizier/oss/proto_converters_test.py index bb7ba9e17..59615cee8 100644 --- a/vizier/_src/pyvizier/oss/proto_converters_test.py +++ b/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): diff --git a/vizier/_src/pyvizier/oss/study_config_test.py b/vizier/_src/pyvizier/oss/study_config_test.py index 7196a6da4..39aa1bac8 100644 --- a/vizier/_src/pyvizier/oss/study_config_test.py +++ b/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() diff --git a/vizier/_src/service/study.proto b/vizier/_src/service/study.proto index da65c7b90..de11dfd8d 100644 --- a/vizier/_src/service/study.proto +++ b/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 {}