diff --git a/.coveragerc b/.coveragerc index 01d28d4b2c..e0d73b97a8 100644 --- a/.coveragerc +++ b/.coveragerc @@ -5,6 +5,7 @@ branch = True show_missing = True omit = google/cloud/aiplatform/v1/schema/trainingjob/definition/__init__.py + google/cloud/aiplatform/vizier/pyvizier/* exclude_lines = # Re-enable the standard pragma pragma: NO COVER diff --git a/google/cloud/aiplatform/compat/__init__.py b/google/cloud/aiplatform/compat/__init__.py index 98ecacc5b5..6b21fd94bf 100644 --- a/google/cloud/aiplatform/compat/__init__.py +++ b/google/cloud/aiplatform/compat/__init__.py @@ -44,6 +44,7 @@ services.index_endpoint_service_client = ( services.index_endpoint_service_client_v1beta1 ) + services.vizier_service_client = services.vizier_service_client_v1beta1 types.accelerator_type = types.accelerator_type_v1beta1 types.annotation = types.annotation_v1beta1 @@ -117,6 +118,7 @@ types.tensorboard_time_series = types.tensorboard_time_series_v1beta1 types.training_pipeline = types.training_pipeline_v1beta1 types.types = types.types_v1beta1 + types.vizier_service = types.vizier_service_v1beta1 if DEFAULT_VERSION == V1: @@ -134,6 +136,7 @@ services.tensorboard_service_client = services.tensorboard_service_client_v1 services.index_service_client = services.index_service_client_v1 services.index_endpoint_service_client = services.index_endpoint_service_client_v1 + services.vizier_service_client = services.vizier_service_client_v1 types.accelerator_type = types.accelerator_type_v1 types.annotation = types.annotation_v1 @@ -204,6 +207,7 @@ types.tensorboard_time_series = types.tensorboard_time_series_v1 types.training_pipeline = types.training_pipeline_v1 types.types = types.types_v1 + types.vizier_service = types.vizier_service_v1 __all__ = ( DEFAULT_VERSION, diff --git a/google/cloud/aiplatform/compat/services/__init__.py b/google/cloud/aiplatform/compat/services/__init__.py index 68440de4c5..627c77b258 100644 --- a/google/cloud/aiplatform/compat/services/__init__.py +++ b/google/cloud/aiplatform/compat/services/__init__.py @@ -54,6 +54,9 @@ from google.cloud.aiplatform_v1beta1.services.tensorboard_service import ( client as tensorboard_service_client_v1beta1, ) +from google.cloud.aiplatform_v1beta1.services.vizier_service import ( + client as vizier_service_client_v1beta1, +) from google.cloud.aiplatform_v1.services.dataset_service import ( client as dataset_service_client_v1, @@ -94,6 +97,9 @@ from google.cloud.aiplatform_v1.services.tensorboard_service import ( client as tensorboard_service_client_v1, ) +from google.cloud.aiplatform_v1.services.vizier_service import ( + client as vizier_service_client_v1, +) __all__ = ( # v1 @@ -110,6 +116,7 @@ prediction_service_client_v1, specialist_pool_service_client_v1, tensorboard_service_client_v1, + vizier_service_client_v1, # v1beta1 dataset_service_client_v1beta1, endpoint_service_client_v1beta1, @@ -124,4 +131,5 @@ specialist_pool_service_client_v1beta1, metadata_service_client_v1beta1, tensorboard_service_client_v1beta1, + vizier_service_client_v1beta1, ) diff --git a/google/cloud/aiplatform/compat/types/__init__.py b/google/cloud/aiplatform/compat/types/__init__.py index d07f86fcb0..ceb9d84157 100644 --- a/google/cloud/aiplatform/compat/types/__init__.py +++ b/google/cloud/aiplatform/compat/types/__init__.py @@ -81,6 +81,7 @@ tensorboard_time_series as tensorboard_time_series_v1beta1, training_pipeline as training_pipeline_v1beta1, types as types_v1beta1, + vizier_service as vizier_service_v1beta1, ) from google.cloud.aiplatform_v1.types import ( accelerator_type as accelerator_type_v1, @@ -147,6 +148,7 @@ tensorboard_time_series as tensorboard_time_series_v1, training_pipeline as training_pipeline_v1, types as types_v1, + vizier_service as vizier_service_v1, ) __all__ = ( diff --git a/google/cloud/aiplatform/utils/__init__.py b/google/cloud/aiplatform/utils/__init__.py index 9ec2f27779..8049c9376b 100644 --- a/google/cloud/aiplatform/utils/__init__.py +++ b/google/cloud/aiplatform/utils/__init__.py @@ -47,6 +47,7 @@ pipeline_service_client_v1beta1, prediction_service_client_v1beta1, tensorboard_service_client_v1beta1, + vizier_service_client_v1beta1, ) from google.cloud.aiplatform.compat.services import ( dataset_service_client_v1, @@ -61,6 +62,7 @@ pipeline_service_client_v1, prediction_service_client_v1, tensorboard_service_client_v1, + vizier_service_client_v1, ) from google.cloud.aiplatform.compat.types import ( @@ -82,6 +84,7 @@ job_service_client_v1beta1.JobServiceClient, metadata_service_client_v1beta1.MetadataServiceClient, tensorboard_service_client_v1beta1.TensorboardServiceClient, + vizier_service_client_v1beta1.VizierServiceClient, # v1 dataset_service_client_v1.DatasetServiceClient, endpoint_service_client_v1.EndpointServiceClient, @@ -93,6 +96,7 @@ pipeline_service_client_v1.PipelineServiceClient, job_service_client_v1.JobServiceClient, tensorboard_service_client_v1.TensorboardServiceClient, + vizier_service_client_v1.VizierServiceClient, ) @@ -570,6 +574,15 @@ class TensorboardClientWithOverride(ClientWithOverride): ) +class VizierClientWithOverride(ClientWithOverride): + _is_temporary = True + _default_version = compat.DEFAULT_VERSION + _version_map = ( + (compat.V1, vizier_service_client_v1.VizierServiceClient), + (compat.V1BETA1, vizier_service_client_v1beta1.VizierServiceClient), + ) + + VertexAiServiceClientWithOverride = TypeVar( "VertexAiServiceClientWithOverride", DatasetClientWithOverride, @@ -582,6 +595,7 @@ class TensorboardClientWithOverride(ClientWithOverride): PredictionClientWithOverride, MetadataClientWithOverride, TensorboardClientWithOverride, + VizierClientWithOverride, ) diff --git a/google/cloud/aiplatform/vizier/__init__.py b/google/cloud/aiplatform/vizier/__init__.py new file mode 100644 index 0000000000..5fab17a0bb --- /dev/null +++ b/google/cloud/aiplatform/vizier/__init__.py @@ -0,0 +1,22 @@ +# -*- coding: utf-8 -*- + +# Copyright 2022 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from google.cloud.aiplatform.vizier.study import Study +from google.cloud.aiplatform.vizier.trial import Trial + +__all__ = ( + "Study", + "Trial", +) diff --git a/google/cloud/aiplatform/vizier/client_abc.py b/google/cloud/aiplatform/vizier/client_abc.py new file mode 100644 index 0000000000..bb1c637fe2 --- /dev/null +++ b/google/cloud/aiplatform/vizier/client_abc.py @@ -0,0 +1,192 @@ +# -*- coding: utf-8 -*- + +# Copyright 2022 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Cross-platform Vizier client interfaces. + +Aside from "materialize_" methods, code written using these interfaces are +compatible with OSS and Cloud Vertex Vizier. Note importantly that subclasses +may have more methods than what is required by interfaces, and such methods +are not cross compatible. Our recommendation is to explicitly type your objects +to be `StudyInterface` or `TrialInterface` when you want to guarantee that +a code block is cross-platform. + +Keywords: + +#Materialize: The method returns a deep copy of the underlying pyvizier object. +Modifying the returned object does not update the Vizier service. +""" + +from __future__ import annotations + +from typing import Optional, Collection, Type, TypeVar, Mapping, Any +import abc + +from google.cloud.aiplatform.vizier import pyvizier as vz + +_T = TypeVar("_T") + + +class ResourceNotFoundError(LookupError): + """Error raised by Vizier clients when resource is not found.""" + + pass + + +class TrialInterface(abc.ABC): + """Responsible for trial-level operations.""" + + @property + @abc.abstractmethod + def uid(self) -> int: + """Unique identifier of the trial.""" + + @property + @abc.abstractmethod + def parameters(self) -> Mapping[str, Any]: + """Parameters of the trial.""" + + @property + @abc.abstractmethod + def status(self) -> vz.TrialStatus: + """Trial's status.""" + + @abc.abstractmethod + def delete(self) -> None: + """Delete the Trial in Vizier service. + + There is currently no promise on how this object behaves after `delete()`. + If you are sharing a Trial object in parallel processes, proceed with + caution. + """ + + @abc.abstractmethod + def complete( + self, + measurement: Optional[vz.Measurement] = None, + *, + infeasible_reason: Optional[str] = None, + ) -> Optional[vz.Measurement]: + """Completes the trial and #materializes the measurement. + + * If `measurement` is provided, then Vizier writes it as the trial's final + measurement and returns it. + * If `infeasible_reason` is provided, `measurement` is not needed. + * If neither is provided, then Vizier selects an existing (intermediate) + measurement to be the final measurement and returns it. + + Args: + measurement: Final measurement. + infeasible_reason: Infeasible reason for missing final measurement. + + Returns: + The final measurement of the trial, or None if the trial is marked + infeasible. + + Raises: + ValueError: If neither `measurement` nor `infeasible_reason` is provided + but the trial does not contain any intermediate measurements. + """ + + @abc.abstractmethod + def should_stop(self) -> bool: + """Returns true if the trial should stop.""" + + @abc.abstractmethod + def add_measurement(self, measurement: vz.Measurement) -> None: + """Adds an intermediate measurement.""" + + @abc.abstractmethod + def materialize(self, *, include_all_measurements: bool = True) -> vz.Trial: + """#Materializes the Trial. + + Args: + include_all_measurements: If True, returned Trial includes all + intermediate measurements. The final measurement is always provided. + + Returns: + Trial object. + """ + + +class StudyInterface(abc.ABC): + """Responsible for study-level operations.""" + + @abc.abstractmethod + def create_or_load( + self, display_name: str, problem: vz.ProblemStatement + ) -> StudyInterface: + """ """ + + @abc.abstractmethod + def suggest( + self, *, count: Optional[int] = None, worker: str = "" + ) -> Collection[TrialInterface]: + """Returns Trials to be evaluated by worker. + + Args: + count: Number of suggestions. + worker: When new Trials are generated, their `assigned_worker` field is + populated with this worker. suggest() first looks for existing Trials + that are assigned to `worker`, before generating new ones. + + Returns: + Trials. + """ + + @abc.abstractmethod + def delete(self) -> None: + """Deletes the study.""" + + @abc.abstractmethod + def trials( + self, trial_filter: Optional[vz.TrialFilter] = None + ) -> Collection[TrialInterface]: + """Fetches a collection of trials.""" + + @abc.abstractmethod + def get_trial(self, uid: int) -> TrialInterface: + """Fetches a single trial. + + Args: + uid: Unique identifier of the trial within study. + + Returns: + Trial. + + Raises: + ResourceNotFoundError: If trial does not exist. + """ + + @abc.abstractmethod + def optimal_trials(self) -> Collection[TrialInterface]: + """Returns optimal trial(s).""" + + @abc.abstractmethod + def materialize_study_config(self) -> vz.StudyConfig: + """#Materializes the study config.""" + + @abc.abstractclassmethod + def from_uid(cls: Type[_T], uid: str) -> _T: + """Fetches an existing study from the Vizier service. + + Args: + uid: Unique identifier of the study. + + Returns: + Study. + + Raises: + ResourceNotFoundError: If study does not exist. + """ diff --git a/google/cloud/aiplatform/vizier/pyvizier/__init__.py b/google/cloud/aiplatform/vizier/pyvizier/__init__.py new file mode 100644 index 0000000000..6c5dfc412b --- /dev/null +++ b/google/cloud/aiplatform/vizier/pyvizier/__init__.py @@ -0,0 +1,84 @@ +"""PyVizier classes for Pythia policies.""" + +try: + from vizier.pyvizier import MetricInformation + from vizier.pyvizier import MetricsConfig + from vizier.pyvizier import MetricType + from vizier.pyvizier import ( + ObjectiveMetricGoal, + ) + from vizier.pyvizier import ProblemStatement + from vizier.pyvizier import SearchSpace + from vizier.pyvizier import ( + SearchSpaceSelector, + ) + from vizier.pyvizier import Metadata + from vizier.pyvizier import MetadataValue + from vizier.pyvizier import Namespace + from vizier.pyvizier import ExternalType + from vizier.pyvizier import ParameterConfig + from vizier.pyvizier import ParameterType + from vizier.pyvizier import ScaleType + from vizier.pyvizier import CompletedTrial + from vizier.pyvizier import Measurement + from vizier.pyvizier import MonotypeParameterSequence + from vizier.pyvizier import Metric + from vizier.pyvizier import ParameterDict + from vizier.pyvizier import ParameterValue + from vizier.pyvizier import Trial + from vizier.pyvizier import ParameterValueTypes + from vizier.pyvizier import TrialFilter + from vizier.pyvizier import TrialStatus + from vizier.pyvizier import TrialSuggestion +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.3a"' + ) + +from google.cloud.aiplatform.vizier.pyvizier.proto_converters import TrialConverter +from google.cloud.aiplatform.vizier.pyvizier.proto_converters import ( + ParameterConfigConverter, +) +from google.cloud.aiplatform.vizier.pyvizier.proto_converters import ( + MeasurementConverter, +) +from google.cloud.aiplatform.vizier.pyvizier.study_config import StudyConfig +from google.cloud.aiplatform.vizier.pyvizier.study_config import Algorithm +from google.cloud.aiplatform.vizier.pyvizier.automated_stopping import ( + AutomatedStoppingConfig, +) + +__all__ = ( + "MetricInformation", + "MetricsConfig", + "MetricType", + "ObjectiveMetricGoal", + "ProblemStatement", + "SearchSpace", + "SearchSpaceSelector", + "Metadata", + "MetadataValue", + "Namespace", + "ParameterConfigConverter", + "ParameterValueTypes", + "MeasurementConverter", + "MonotypeParameterSequence", + "TrialConverter", + "StudyConfig", + "Algorithm", + "AutomatedStoppingConfig", + "ExternalType", + "ParameterConfig", + "ParameterType", + "ScaleType", + "CompletedTrial", + "Measurement", + "Metric", + "ParameterDict", + "ParameterValue", + "Trial", + "TrialFilter", + "TrialStatus", + "TrialSuggestion", +) diff --git a/google/cloud/aiplatform/vizier/pyvizier/automated_stopping.py b/google/cloud/aiplatform/vizier/pyvizier/automated_stopping.py new file mode 100644 index 0000000000..b46f2b12ec --- /dev/null +++ b/google/cloud/aiplatform/vizier/pyvizier/automated_stopping.py @@ -0,0 +1,79 @@ +"""Convenience classes for configuring Vizier Early-Stopping Configs.""" +import copy +from typing import Union + +import attr + +from google.cloud.aiplatform.compat.types import study as study_pb2 + +AutomatedStoppingConfigProto = Union[ + study_pb2.StudySpec.DecayCurveAutomatedStoppingSpec, + study_pb2.StudySpec.MedianAutomatedStoppingSpec, +] + + +@attr.s(frozen=True, init=True, slots=True, kw_only=True) +class AutomatedStoppingConfig: + """A wrapper for study_pb2.automated_stopping_spec.""" + + _proto: AutomatedStoppingConfigProto = attr.ib(init=True, kw_only=True) + + @classmethod + def decay_curve_stopping_config(cls, use_steps: bool) -> "AutomatedStoppingConfig": + """Create a DecayCurve automated stopping config. + + Vizier will early stop the Trial if it predicts the Trial objective value + will not be better than previous Trials. + + Args: + use_steps: Bool. If set, use Measurement.step_count as the measure of + training progress. Otherwise, use Measurement.elapsed_duration. + + Returns: + AutomatedStoppingConfig object. + + Raises: + ValueError: If more than one metric is configured. + Note that Vizier Early Stopping currently only supports single-objective + studies. + """ + config = study_pb2.StudySpec.DecayCurveAutomatedStoppingSpec( + use_elapsed_duration=not use_steps + ) + return cls(proto=config) + + @classmethod + def median_automated_stopping_config( + cls, use_steps: bool + ) -> "AutomatedStoppingConfig": + """Create a Median automated stopping config. + + Vizier will early stop the Trial if it predicts the Trial objective value + will not be better than previous Trials. + + Args: + use_steps: Bool. If set, use Measurement.step_count as the measure of + training progress. Otherwise, use Measurement.elapsed_duration. + + Returns: + AutomatedStoppingConfig object. + + Raises: + ValueError: If more than one metric is configured. + Note that Vizier Early Stopping currently only supports single-objective + studies. + """ + config = study_pb2.StudySpec.MedianAutomatedStoppingSpec( + use_elapsed_duration=not use_steps + ) + return cls(proto=config) + + @classmethod + def from_proto( + cls, proto: AutomatedStoppingConfigProto + ) -> "AutomatedStoppingConfig": + return cls(proto=proto) + + def to_proto(self) -> AutomatedStoppingConfigProto: + """Returns this object as a proto.""" + return copy.deepcopy(self._proto) diff --git a/google/cloud/aiplatform/vizier/pyvizier/proto_converters.py b/google/cloud/aiplatform/vizier/pyvizier/proto_converters.py new file mode 100644 index 0000000000..86ce9df770 --- /dev/null +++ b/google/cloud/aiplatform/vizier/pyvizier/proto_converters.py @@ -0,0 +1,502 @@ +"""Converters for OSS Vizier's protos from/to PyVizier's classes.""" +import datetime +import logging +from typing import List, Optional, Sequence, Tuple, Union + +from google.protobuf import duration_pb2 +from google.cloud.aiplatform.compat.types import study as study_pb2 +from google.cloud.aiplatform.vizier.pyvizier import ScaleType +from google.cloud.aiplatform.vizier.pyvizier import ParameterType +from google.cloud.aiplatform.vizier.pyvizier import ParameterValue +from google.cloud.aiplatform.vizier.pyvizier import MonotypeParameterSequence +from google.cloud.aiplatform.vizier.pyvizier import ParameterConfig +from google.cloud.aiplatform.vizier.pyvizier import Measurement +from google.cloud.aiplatform.vizier.pyvizier import Metric +from google.cloud.aiplatform.vizier.pyvizier import TrialStatus +from google.cloud.aiplatform.vizier.pyvizier import Trial + +_ScaleTypePb2 = study_pb2.StudySpec.ParameterSpec.ScaleType + + +class _ScaleTypeMap: + """Proto converter for scale type.""" + + _pyvizier_to_proto = { + ScaleType.LINEAR: _ScaleTypePb2.UNIT_LINEAR_SCALE, + ScaleType.LOG: _ScaleTypePb2.UNIT_LOG_SCALE, + ScaleType.REVERSE_LOG: _ScaleTypePb2.UNIT_REVERSE_LOG_SCALE, + } + _proto_to_pyvizier = {v: k for k, v in _pyvizier_to_proto.items()} + + @classmethod + def to_proto(cls, pyvizier: ScaleType) -> _ScaleTypePb2: + return cls._pyvizier_to_proto[pyvizier] + + @classmethod + def from_proto(cls, proto: _ScaleTypePb2) -> ScaleType: + return cls._proto_to_pyvizier[proto] + + +class ParameterConfigConverter: + """Converter for ParameterConfig.""" + + @classmethod + def _set_bounds( + cls, + proto: study_pb2.StudySpec.ParameterSpec, + lower: float, + upper: float, + parameter_type: ParameterType, + ): + """Sets the proto's min_value and max_value fields.""" + if parameter_type == ParameterType.INTEGER: + proto.integer_value_spec.min_value = lower + proto.integer_value_spec.max_value = upper + elif parameter_type == ParameterType.DOUBLE: + proto.double_value_spec.min_value = lower + proto.double_value_spec.max_value = upper + + @classmethod + def _set_feasible_points( + cls, proto: study_pb2.StudySpec.ParameterSpec, feasible_points: Sequence[float] + ): + """Sets the proto's feasible_points field.""" + feasible_points = sorted(feasible_points) + proto.discrete_value_spec.values.clear() + proto.discrete_value_spec.values.extend(feasible_points) + + @classmethod + def _set_categories( + cls, proto: study_pb2.StudySpec.ParameterSpec, categories: Sequence[str] + ): + """Sets the protos' categories field.""" + proto.categorical_value_spec.values.clear() + proto.categorical_value_spec.values.extend(categories) + + @classmethod + def _set_default_value( + cls, + proto: study_pb2.StudySpec.ParameterSpec, + default_value: Union[float, int, str], + ): + """Sets the protos' default_value field.""" + which_pv_spec = proto.WhichOneof("parameter_value_spec") + getattr(proto, which_pv_spec).default_value.value = default_value + + @classmethod + def _matching_parent_values( + cls, proto: study_pb2.StudySpec.ParameterSpec.ConditionalParameterSpec + ) -> MonotypeParameterSequence: + """Returns the matching parent values, if set.""" + oneof_name = proto.WhichOneof("parent_value_condition") + if not oneof_name: + return [] + if oneof_name in ( + "parent_discrete_values", + "parent_int_values", + "parent_categorical_values", + ): + return list(getattr(getattr(proto, oneof_name), "values")) + raise ValueError("Unknown matching_parent_vals: {}".format(oneof_name)) + + @classmethod + def from_proto( + cls, + proto: study_pb2.StudySpec.ParameterSpec, + *, + strict_validation: bool = False + ) -> ParameterConfig: + """Creates a ParameterConfig. + + Args: + proto: + strict_validation: If True, raise ValueError to enforce that + from_proto(proto).to_proto == proto. + + Returns: + ParameterConfig object + + Raises: + ValueError: See the "strict_validtion" arg documentation. + """ + feasible_values = [] + 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), + ) + elif oneof_name == "double_value_spec": + bounds = ( + proto.double_value_spec.min_value, + proto.double_value_spec.max_value, + ) + elif oneof_name == "discrete_value_spec": + bounds = None + feasible_values = proto.discrete_value_spec.values + elif oneof_name == "categorical_value_spec": + bounds = None + feasible_values = proto.categorical_value_spec.values + + default_value = None + if getattr(proto, oneof_name).default_value: + default_value = getattr(proto, oneof_name).default_value + + if proto.conditional_parameter_specs: + children = [] + for conditional_ps in proto.conditional_parameter_specs: + parent_values = cls._matching_parent_values(conditional_ps) + children.append( + (parent_values, cls.from_proto(conditional_ps.parameter_spec)) + ) + else: + children = None + + scale_type = None + if proto.scale_type: + scale_type = _ScaleTypeMap.from_proto(proto.scale_type) + + try: + config = ParameterConfig.factory( + name=proto.parameter_id, + feasible_values=feasible_values, + bounds=bounds, + children=children, + scale_type=scale_type, + default_value=default_value, + ) + except ValueError as e: + raise ValueError( + "The provided proto was misconfigured. {}".format(proto) + ) from e + + if strict_validation and cls.to_proto(config) != proto: + raise ValueError( + "The provided proto was misconfigured. Expected: {} Given: {}".format( + cls.to_proto(config), proto + ) + ) + return config + + @classmethod + def _set_child_parameter_configs( + cls, + parent_proto: study_pb2.StudySpec.ParameterSpec, + pc: ParameterConfig, + ): + """Sets the parent_proto's conditional_parameter_specs field. + + Args: + parent_proto: Modified in place. + pc: Parent ParameterConfig to copy children from. + + Raises: + ValueError: If the child configs are invalid + """ + children: List[Tuple[MonotypeParameterSequence, ParameterConfig]] = [] + for child in pc.child_parameter_configs: + children.append((child.matching_parent_values, child)) + if not children: + return + parent_proto.conditional_parameter_specs.clear() + 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)""" + ) + + logging.debug( + "_set_child_parameter_configs: parent_proto=%s, children=%s", + parent_proto, + children, + ) + for unsorted_parent_values, child in children: + parent_values = sorted(unsorted_parent_values) + child_proto = cls.to_proto(child.clone_without_children) + conditional_parameter_spec = ( + study_pb2.StudySpec.ParameterSpec.ConditionalParameterSpec( + parameter_spec=child_proto + ) + ) + + if parent_proto.HasField("discrete_value_spec"): + conditional_parameter_spec.parent_discrete_values.values[ + : + ] = parent_values + elif parent_proto.HasField("categorical_value_spec"): + conditional_parameter_spec.parent_categorical_values.values[ + : + ] = parent_values + elif parent_proto.HasField("integer_value_spec"): + conditional_parameter_spec.parent_int_values.values[:] = parent_values + else: + raise ValueError("DOUBLE type cannot have child parameters") + if child.child_parameter_configs: + cls._set_child_parameter_configs(child_proto, child) + parent_proto.conditional_parameter_specs.extend( + [conditional_parameter_spec] + ) + + @classmethod + def to_proto(cls, pc: ParameterConfig) -> study_pb2.StudySpec.ParameterSpec: + """Returns a ParameterConfig Proto.""" + proto = study_pb2.StudySpec.ParameterSpec(parameter_id=pc.name) + if pc.type == ParameterType.DISCRETE: + cls._set_feasible_points(proto, [float(v) for v in pc.feasible_values]) + elif pc.type == ParameterType.CATEGORICAL: + cls._set_categories(proto, pc.feasible_values) + elif pc.type in (ParameterType.INTEGER, ParameterType.DOUBLE): + cls._set_bounds(proto, pc.bounds[0], pc.bounds[1], pc.type) + else: + raise ValueError("Invalid ParameterConfig: {}".format(pc)) + if pc.scale_type is not None and pc.scale_type != ScaleType.UNIFORM_DISCRETE: + proto.scale_type = _ScaleTypeMap.to_proto(pc.scale_type) + if pc.default_value is not None: + cls._set_default_value(proto, pc.default_value) + + cls._set_child_parameter_configs(proto, pc) + return proto + + +class ParameterValueConverter: + """Converter for ParameterValue.""" + + @classmethod + def from_proto(cls, proto: study_pb2.Trial.Parameter) -> Optional[ParameterValue]: + """Returns whichever value that is populated, or None.""" + potential_value = proto.value + if ( + isinstance(potential_value, float) + or isinstance(potential_value, str) + or isinstance(potential_value, bool) + ): + return ParameterValue(potential_value) + else: + return None + + @classmethod + def to_proto( + cls, parameter_value: ParameterValue, name: str + ) -> study_pb2.Trial.Parameter: + """Returns Parameter Proto.""" + proto = study_pb2.Trial.Parameter(parameter_id=name) + + if isinstance(parameter_value.value, int): + proto.value.number_value = parameter_value.value + elif isinstance(parameter_value.value, bool): + proto.value.bool_value = parameter_value.value + elif isinstance(parameter_value.value, float): + proto.value.number_value = parameter_value.value + elif isinstance(parameter_value.value, str): + proto.value.string_value = parameter_value.value + + return proto + + +class MeasurementConverter: + """Converter for MeasurementConverter.""" + + @classmethod + def from_proto(cls, proto: study_pb2.Measurement) -> Measurement: + """Creates a valid instance from proto. + + Args: + proto: Measurement proto. + + Returns: + A valid instance of Measurement object. Metrics with invalid values + are automatically filtered out. + """ + + metrics = dict() + + for metric in proto.metrics: + if ( + metric.metric_id in metrics + and metrics[metric.metric_id].value != metric.value + ): + logging.log_first_n( + logging.ERROR, + 'Duplicate metric of name "%s".' + "The newly found value %s will be used and " + "the previously found value %s will be discarded." + "This always happens if the proto has an empty-named metric.", + 5, + metric.metric_id, + metric.value, + metrics[metric.metric_id].value, + ) + try: + metrics[metric.metric_id] = Metric(value=metric.value) + except ValueError: + pass + return Measurement( + metrics=metrics, + elapsed_secs=proto.elapsed_duration.seconds, + steps=proto.step_count, + ) + + @classmethod + def to_proto(cls, measurement: Measurement) -> study_pb2.Measurement: + """Converts to Measurement proto.""" + proto = study_pb2.Measurement() + for name, metric in measurement.metrics.items(): + proto.metrics.append( + study_pb2.Measurement.Metric(metric_id=name, value=metric.value) + ) + + proto.step_count = measurement.steps + int_seconds = int(measurement.elapsed_secs) + proto.elapsed_duration = duration_pb2.Duration( + seconds=int_seconds, + nanos=int(1e9 * (measurement.elapsed_secs - int_seconds)), + ) + return proto + + +def _to_pyvizier_trial_status(proto_state: study_pb2.Trial.State) -> TrialStatus: + """from_proto conversion for Trial statuses.""" + if proto_state == study_pb2.Trial.State.REQUESTED: + return TrialStatus.REQUESTED + elif proto_state == study_pb2.Trial.State.ACTIVE: + return TrialStatus.ACTIVE + if proto_state == study_pb2.Trial.State.STOPPING: + return TrialStatus.STOPPING + if proto_state == study_pb2.Trial.State.SUCCEEDED: + return TrialStatus.COMPLETED + elif proto_state == study_pb2.Trial.State.INFEASIBLE: + return TrialStatus.COMPLETED + else: + return TrialStatus.UNKNOWN + + +def _from_pyvizier_trial_status( + status: TrialStatus, infeasible: bool +) -> study_pb2.Trial.State: + """to_proto conversion for Trial states.""" + if status == TrialStatus.REQUESTED: + return study_pb2.Trial.State.REQUESTED + elif status == TrialStatus.ACTIVE: + return study_pb2.Trial.State.ACTIVE + elif status == TrialStatus.STOPPING: + return study_pb2.Trial.State.STOPPING + elif status == TrialStatus.COMPLETED: + if infeasible: + return study_pb2.Trial.State.INFEASIBLE + else: + return study_pb2.Trial.State.SUCCEEDED + else: + return study_pb2.Trial.State.STATE_UNSPECIFIED + + +class TrialConverter: + """Converter for TrialConverter.""" + + @classmethod + def from_proto(cls, proto: study_pb2.Trial) -> Trial: + """Converts from Trial proto to object. + + Args: + proto: Trial proto. + + Returns: + A Trial object. + """ + parameters = {} + for parameter in proto.parameters: + value = ParameterValueConverter.from_proto(parameter) + if value is not None: + if parameter.parameter_id in parameters: + raise ValueError( + "Invalid trial proto contains duplicate parameter {}" + ": {}".format(parameter.parameter_id, proto) + ) + parameters[parameter.parameter_id] = value + else: + logging.warning( + "A parameter without a value will be dropped: %s", parameter + ) + + final_measurement = None + if proto.final_measurement: + final_measurement = MeasurementConverter.from_proto(proto.final_measurement) + + completion_time = None + infeasibility_reason = None + if proto.state == study_pb2.Trial.State.SUCCEEDED: + if proto.end_time: + completion_ts = proto.end_time.nanosecond / 1e9 + completion_time = datetime.datetime.fromtimestamp(completion_ts) + elif proto.state == study_pb2.Trial.State.INFEASIBLE: + infeasibility_reason = proto.infeasible_reason + + measurements = [] + for measure in proto.measurements: + measurements.append(MeasurementConverter.from_proto(measure)) + + creation_time = None + if proto.start_time: + creation_ts = proto.start_time.nanosecond / 1e9 + creation_time = datetime.datetime.fromtimestamp(creation_ts) + return Trial( + id=int(proto.name.split("/")[-1]), + description=proto.name, + assigned_worker=proto.client_id or None, + is_requested=proto.state == study_pb2.Trial.State.REQUESTED, + stopping_reason=( + "stopping reason not supported yet" + if proto.state == study_pb2.Trial.State.STOPPING + else None + ), + parameters=parameters, + creation_time=creation_time, + completion_time=completion_time, + infeasibility_reason=infeasibility_reason, + final_measurement=final_measurement, + measurements=measurements, + ) # pytype: disable=wrong-arg-types + + @classmethod + def from_protos(cls, protos: Sequence[study_pb2.Trial]) -> List[Trial]: + """Convenience wrapper for from_proto.""" + return [TrialConverter.from_proto(proto) for proto in protos] + + @classmethod + def to_protos(cls, pytrials: Sequence[Trial]) -> List[study_pb2.Trial]: + return [TrialConverter.to_proto(pytrial) for pytrial in pytrials] + + @classmethod + def to_proto(cls, pytrial: Trial) -> study_pb2.Trial: + """Converts a pyvizier Trial to a Trial proto.""" + proto = study_pb2.Trial() + if pytrial.description is not None: + proto.name = pytrial.description + proto.id = str(pytrial.id) + proto.state = _from_pyvizier_trial_status(pytrial.status, pytrial.infeasible) + proto.client_id = pytrial.assigned_worker or "" + + for name, value in pytrial.parameters.items(): + proto.parameters.append(ParameterValueConverter.to_proto(value, name)) + + # pytrial always adds an empty metric. Ideally, we should remove it if the + # metric does not exist in the study config. + if pytrial.final_measurement is not None: + proto.final_measurement.CopyFrom( + MeasurementConverter.to_proto(pytrial.final_measurement) + ) + + for measurement in pytrial.measurements: + proto.measurements.append(MeasurementConverter.to_proto(measurement)) + + if pytrial.creation_time is not None: + creation_secs = datetime.datetime.timestamp(pytrial.creation_time) + proto.start_time.seconds = int(creation_secs) + proto.start_time.nanos = int(1e9 * (creation_secs - int(creation_secs))) + if pytrial.completion_time is not None: + completion_secs = datetime.datetime.timestamp(pytrial.completion_time) + proto.end_time.seconds = int(completion_secs) + proto.end_time.nanos = int(1e9 * (completion_secs - int(completion_secs))) + if pytrial.infeasibility_reason is not None: + proto.infeasible_reason = pytrial.infeasibility_reason + return proto diff --git a/google/cloud/aiplatform/vizier/pyvizier/study_config.py b/google/cloud/aiplatform/vizier/pyvizier/study_config.py new file mode 100644 index 0000000000..0314e1442f --- /dev/null +++ b/google/cloud/aiplatform/vizier/pyvizier/study_config.py @@ -0,0 +1,453 @@ +"""Convenience classes for configuring Vizier Study Configs and Search Spaces. + +This module contains several classes, used to access/build Vizier StudyConfig +protos: + * `StudyConfig` class is the main class, which: + 1) Allows to easily build Vizier StudyConfig protos via a convenient + Python API. + 2) Can be initialized from an existing StudyConfig proto, to enable easy + Pythonic accessors to information contained in StudyConfig protos, + and easy field editing capabilities. + + * `SearchSpace` and `SearchSpaceSelector` classes deals with Vizier search + spaces. Both flat spaces and conditional parameters are supported. +""" +import collections +import copy +import enum +from typing import Dict, Iterable, List, Optional, Sequence, Tuple, Union + +import attr +from google.cloud.aiplatform.vizier.pyvizier.automated_stopping import ( + AutomatedStoppingConfig, +) +from google.cloud.aiplatform.vizier.pyvizier import proto_converters +from google.cloud.aiplatform.vizier.pyvizier import SearchSpace +from google.cloud.aiplatform.vizier.pyvizier import ProblemStatement +from google.cloud.aiplatform.vizier.pyvizier import ObjectiveMetricGoal +from google.cloud.aiplatform.vizier.pyvizier import SearchSpaceSelector +from google.cloud.aiplatform.vizier.pyvizier import MetricsConfig +from google.cloud.aiplatform.vizier.pyvizier import MetricInformation +from google.cloud.aiplatform.vizier.pyvizier import Trial +from google.cloud.aiplatform.vizier.pyvizier import ParameterValueTypes +from google.cloud.aiplatform.vizier.pyvizier import ParameterConfig +from google.cloud.aiplatform.compat.types import study as study_pb2 + +################### PyTypes ################### +# A sequence of possible internal parameter values. +# Possible types for trial parameter values after cast to external types. +ParameterValueSequence = Union[ + ParameterValueTypes, + Sequence[int], + Sequence[float], + Sequence[str], + Sequence[bool], +] + +################### Enums ################### + + +class Algorithm(enum.Enum): + """Valid Values for StudyConfig.Algorithm.""" + + ALGORITHM_UNSPECIFIED = study_pb2.StudySpec.Algorithm.ALGORITHM_UNSPECIFIED + # GAUSSIAN_PROCESS_BANDIT = study_pb2.StudySpec.Algorithm.GAUSSIAN_PROCESS_BANDIT + GRID_SEARCH = study_pb2.StudySpec.Algorithm.GRID_SEARCH + RANDOM_SEARCH = study_pb2.StudySpec.Algorithm.RANDOM_SEARCH + # NSGA2 = study_pb2.StudySpec.Algorithm.NSGA2 + # EMUKIT_GP_EI = study_pb2.StudySpec.Algorithm.EMUKIT_GP_EI + + +class ObservationNoise(enum.Enum): + """Valid Values for StudyConfig.ObservationNoise.""" + + OBSERVATION_NOISE_UNSPECIFIED = ( + study_pb2.StudySpec.ObservationNoise.OBSERVATION_NOISE_UNSPECIFIED + ) + LOW = study_pb2.StudySpec.ObservationNoise.LOW + HIGH = study_pb2.StudySpec.ObservationNoise.HIGH + + +################### Classes For Various Config Protos ################### +@attr.define(frozen=False, init=True, slots=True, kw_only=True) +class MetricInformationConverter: + """A wrapper for vizier_pb2.MetricInformation.""" + + @classmethod + def from_proto(cls, proto: study_pb2.StudySpec.MetricSpec) -> MetricInformation: + """Converts a MetricInformation proto to a MetricInformation object.""" + if proto.goal not in list(ObjectiveMetricGoal): + raise ValueError("Unknown MetricInformation.goal: {}".format(proto.goal)) + + return MetricInformation( + name=proto.metric_id, + goal=proto.goal, + safety_threshold=None, + safety_std_threshold=None, + min_value=None, + max_value=None, + ) + + @classmethod + def to_proto(cls, obj: MetricInformation) -> study_pb2.StudySpec.MetricSpec: + """Returns this object as a proto.""" + return study_pb2.StudySpec.MetricSpec(metric_id=obj.name, goal=obj.goal.value) + + +class MetricsConfig(MetricsConfig): + """Metrics config.""" + + @classmethod + def from_proto( + cls, protos: Iterable[study_pb2.StudySpec.MetricSpec] + ) -> "MetricsConfig": + return cls(MetricInformationConverter.from_proto(m) for m in protos) + + def to_proto(self) -> List[study_pb2.StudySpec.MetricSpec]: + return [MetricInformationConverter.to_proto(metric) for metric in self] + + +SearchSpaceSelector = SearchSpaceSelector + + +@attr.define(frozen=True, init=True, slots=True, kw_only=True) +class SearchSpace(SearchSpace): + """A Selector for all, or part of a SearchSpace.""" + + @classmethod + def from_proto(cls, proto: study_pb2.StudySpec) -> "SearchSpace": + """Extracts a SearchSpace object from a StudyConfig proto.""" + parameter_configs = [] + for pc in proto.parameters: + parameter_configs.append( + proto_converters.ParameterConfigConverter.from_proto(pc) + ) + return cls._factory(parameter_configs=parameter_configs) + + @property + def parameter_protos(self) -> List[study_pb2.StudySpec.ParameterSpec]: + """Returns the search space as a List of ParameterConfig protos.""" + return [ + proto_converters.ParameterConfigConverter.to_proto(pc) + for pc in self._parameter_configs + ] + + +################### Main Class ################### +# +# A StudyConfig object can be initialized: +# (1) From a StudyConfig proto using StudyConfig.from_proto(): +# study_config_proto = study_pb2.StudySpec(...) +# study_config = pyvizier.StudyConfig.from_proto(study_config_proto) +# # Attributes can be modified. +# new_proto = study_config.to_proto() +# +# (2) By directly calling __init__ and setting attributes: +# study_config = pyvizier.StudyConfig( +# metric_information=[pyvizier.MetricInformation( +# name='accuracy', goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE)], +# search_space=SearchSpace.from_proto(proto), +# ) +# # OR: +# study_config = pyvizier.StudyConfig() +# study_config.metric_information.append( +# pyvizier.MetricInformation( +# name='accuracy', goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE)) +# +# # Since building a search space is more involved, get a reference to the +# # search space, and add parameters to it. +# root = study_config.search_space.select_root() +# root.add_float_param('learning_rate', 0.001, 1.0, +# scale_type=pyvizier.ScaleType.LOG) +# +@attr.define(frozen=False, init=True, slots=True, kw_only=True) +class StudyConfig(ProblemStatement): + """A builder and wrapper for study_pb2.StudySpec proto.""" + + search_space: SearchSpace = attr.field( + init=True, + factory=SearchSpace, + validator=attr.validators.instance_of(SearchSpace), + on_setattr=attr.setters.validate, + ) + + algorithm: Algorithm = attr.field( + init=True, + validator=attr.validators.instance_of(Algorithm), + on_setattr=[attr.setters.convert, attr.setters.validate], + default=Algorithm.ALGORITHM_UNSPECIFIED, + kw_only=True, + ) + + metric_information: MetricsConfig = attr.field( + init=True, + factory=MetricsConfig, + converter=MetricsConfig, + validator=attr.validators.instance_of(MetricsConfig), + kw_only=True, + ) + + observation_noise: ObservationNoise = attr.field( + init=True, + validator=attr.validators.instance_of(ObservationNoise), + on_setattr=attr.setters.validate, + default=ObservationNoise.OBSERVATION_NOISE_UNSPECIFIED, + kw_only=True, + ) + + automated_stopping_config: Optional[AutomatedStoppingConfig] = attr.field( + init=True, + default=None, + validator=attr.validators.optional( + attr.validators.instance_of(AutomatedStoppingConfig) + ), + on_setattr=attr.setters.validate, + kw_only=True, + ) + + # An internal representation as a StudyConfig proto. + # If this object was created from a StudyConfig proto, a copy of the original + # proto is kept, to make sure that unknown proto fields are preserved in + # round trip serialization. + # TODO: Fix the broken proto validation. + _study_config: study_pb2.StudySpec = attr.field( + init=True, factory=study_pb2.StudySpec, kw_only=True + ) + + # Public attributes, methods and properties. + @classmethod + def from_proto(cls, proto: study_pb2.StudySpec) -> "StudyConfig": + """Converts a StudyConfig proto to a StudyConfig object. + + Args: + proto: StudyConfig proto. + + Returns: + A StudyConfig object. + """ + metric_information = MetricsConfig( + sorted( + [MetricInformationConverter.from_proto(m) for m in proto.metrics], + key=lambda x: x.name, + ) + ) + + oneof_name = proto._pb.WhichOneof("automated_stopping_spec") + if not oneof_name: + automated_stopping_config = None + else: + automated_stopping_config = AutomatedStoppingConfig.from_proto( + getattr(proto, oneof_name) + ) + + return cls( + search_space=SearchSpace.from_proto(proto), + algorithm=Algorithm(proto.algorithm), + metric_information=metric_information, + observation_noise=ObservationNoise(proto.observation_noise), + automated_stopping_config=automated_stopping_config, + study_config=copy.deepcopy(proto), + ) + + def to_proto(self) -> study_pb2.StudySpec: + """Serializes this object to a StudyConfig proto.""" + proto = copy.deepcopy(self._study_config) + proto.algorithm = self.algorithm.value + proto.observation_noise = self.observation_noise.value + + del proto.metrics[:] + proto.metrics.extend(self.metric_information.to_proto()) + + del proto.parameters[:] + proto.parameters.extend(self.search_space.parameter_protos) + + if self.automated_stopping_config is not None: + auto_stop_proto = self.automated_stopping_config.to_proto() + if isinstance( + auto_stop_proto, study_pb2.StudySpec.DecayCurveAutomatedStoppingSpec + ): + proto.decay_curve_stopping_spec = copy.deepcopy(auto_stop_proto) + elif isinstance( + auto_stop_proto, study_pb2.StudySpec.DecayCurveAutomatedStoppingSpec + ): + for method_name in dir(proto.decay_curve_stopping_spec): + if callable( + getattr(proto.median_automated_stopping_spec, method_name) + ): + print(method_name) + proto.median_automated_stopping_spec = copy.deepcopy(auto_stop_proto) + + return proto + + @property + def is_single_objective(self) -> bool: + """Returns True if only one objective metric is configured.""" + return len(self.metric_information) == 1 + + @property + def single_objective_metric_name(self) -> Optional[str]: + """Returns the name of the single-objective metric, if set. + + Returns: + String: name of the single-objective metric. + None: if this is not a single-objective study. + """ + if len(self.metric_information) == 1: + return list(self.metric_information)[0].name + return None + + def _trial_to_external_values( + self, pytrial: Trial + ) -> Dict[str, Union[float, int, str, bool]]: + """Returns the trial paremeter values cast to external types.""" + parameter_values: Dict[str, Union[float, int, str]] = {} + external_values: Dict[str, Union[float, int, str, bool]] = {} + # parameter_configs is a list of Tuple[parent_name, ParameterConfig]. + parameter_configs: List[Tuple[Optional[str], ParameterConfig]] = [ + (None, p) for p in self.search_space.parameters + ] + remaining_parameters = copy.deepcopy(pytrial.parameters) + # Traverse the conditional tree using a BFS. + while parameter_configs and remaining_parameters: + parent_name, pc = parameter_configs.pop(0) + parameter_configs.extend( + (pc.name, child) for child in pc.child_parameter_configs + ) + if pc.name not in remaining_parameters: + continue + if parent_name is not None: + # This is a child parameter. If the parent was not seen, + # skip this parameter config. + if parent_name not in parameter_values: + continue + parent_value = parameter_values[parent_name] + if parent_value not in pc.matching_parent_values: + continue + parameter_values[pc.name] = remaining_parameters[pc.name].value + if pc.external_type is None: + external_value = remaining_parameters[pc.name].value + else: + external_value = remaining_parameters[pc.name].cast( + pc.external_type + ) # pytype: disable=wrong-arg-types + external_values[pc.name] = external_value + remaining_parameters.pop(pc.name) + return external_values + + def trial_parameters( + self, proto: study_pb2.Trial + ) -> Dict[str, ParameterValueSequence]: + """Returns the trial values, cast to external types, if they exist. + + Args: + proto: + + Returns: + Parameter values dict: cast to each parameter's external_type, if exists. + NOTE that the values in the dict may be a Sequence as opposed to a single + element. + + Raises: + ValueError: If the trial parameters do not exist in this search space. + ValueError: If the trial contains duplicate parameters. + """ + pytrial = proto_converters.TrialConverter.from_proto(proto) + return self._pytrial_parameters(pytrial) + + def _pytrial_parameters(self, pytrial: Trial) -> Dict[str, ParameterValueSequence]: + """Returns the trial values, cast to external types, if they exist. + + Args: + pytrial: + + Returns: + Parameter values dict: cast to each parameter's external_type, if exists. + NOTE that the values in the dict may be a Sequence as opposed to a single + element. + + Raises: + ValueError: If the trial parameters do not exist in this search space. + ValueError: If the trial contains duplicate parameters. + """ + trial_external_values: Dict[ + str, Union[float, int, str, bool] + ] = self._trial_to_external_values(pytrial) + if len(trial_external_values) != len(pytrial.parameters): + raise ValueError( + "Invalid trial for this search space: failed to convert " + "all trial parameters: {}".format(pytrial) + ) + + # Combine multi-dimensional parameter values to a list of values. + trial_final_values: Dict[str, ParameterValueSequence] = {} + # multi_dim_params: Dict[str, List[Tuple[int, ParameterValueSequence]]] + multi_dim_params = collections.defaultdict(list) + for name in trial_external_values: + base_index = SearchSpaceSelector.parse_multi_dimensional_parameter_name( + name + ) + if base_index is None: + trial_final_values[name] = trial_external_values[name] + else: + base_name, index = base_index + multi_dim_params[base_name].append((index, trial_external_values[name])) + for name in multi_dim_params: + multi_dim_params[name].sort(key=lambda x: x[0]) + trial_final_values[name] = [x[1] for x in multi_dim_params[name]] + + return trial_final_values + + def trial_metrics( + self, proto: study_pb2.Trial, *, include_all_metrics=False + ) -> Dict[str, float]: + """Returns the trial's final measurement metric values. + + If the trial is not completed, or infeasible, no metrics are returned. + By default, only metrics configured in the StudyConfig are returned + (e.g. only objective and safety metrics). + + Args: + proto: + include_all_metrics: If True, all metrics in the final measurements are + returned. If False, only metrics configured in the StudyConfig are + returned. + + Returns: + Dict[metric name, metric value] + """ + pytrial = proto_converters.TrialConverter.from_proto(proto) + return self._pytrial_metrics(pytrial, include_all_metrics=include_all_metrics) + + def _pytrial_metrics( + self, pytrial: Trial, *, include_all_metrics=False + ) -> Dict[str, float]: + """Returns the trial's final measurement metric values. + + If the trial is not completed, or infeasible, no metrics are returned. + By default, only metrics configured in the StudyConfig are returned + (e.g. only objective and safety metrics). + + Args: + pytrial: + include_all_metrics: If True, all metrics in the final measurements are + returned. If False, only metrics configured in the StudyConfig are + returned. + + Returns: + Dict[metric name, metric value] + """ + configured_metrics = [m.name for m in self.metric_information] + + metrics: Dict[str, float] = {} + if pytrial.is_completed and not pytrial.infeasible: + for name in pytrial.final_measurement.metrics: + if include_all_metrics or ( + not include_all_metrics and name in configured_metrics + ): + # Special case: Measurement always adds an empty metric by default. + # If there is a named single objective in study_config, drop the empty + # metric. + if not name and self.single_objective_metric_name != name: + continue + metrics[name] = pytrial.final_measurement.metrics[name].value + return metrics diff --git a/google/cloud/aiplatform/vizier/study.py b/google/cloud/aiplatform/vizier/study.py new file mode 100644 index 0000000000..79ad3c2961 --- /dev/null +++ b/google/cloud/aiplatform/vizier/study.py @@ -0,0 +1,300 @@ +# -*- coding: utf-8 -*- + +# Copyright 2022 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import copy +from typing import Optional, Collection, Type, TypeVar + +from google.api_core import exceptions +from google.auth import credentials as auth_credentials +from google.cloud.aiplatform import base +from google.cloud.aiplatform import utils +from google.cloud.aiplatform import initializer +from google.cloud.aiplatform.vizier import client_abc +from google.cloud.aiplatform.vizier import pyvizier as vz +from google.cloud.aiplatform.vizier.trial import Trial + + +from google.cloud.aiplatform.compat.types import study as gca_study + + +_T = TypeVar("_T") +_LOGGER = base.Logger(__name__) + + +class Study(base.VertexAiResourceNounWithFutureManager, client_abc.StudyInterface): + """Manage Study resource for Vertex Vizier.""" + + client_class = utils.VizierClientWithOverride + + _resource_noun = "study" + _getter_method = "get_study" + _list_method = "list_studies" + _delete_method = "delete_study" + _parse_resource_name_method = "parse_study_path" + _format_resource_name_method = "study_path" + + def __init__( + self, + study_id: str, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """Retrieves an existing managed study given a study resource name or a study id. + + Example Usage: + study = aiplatform.Study(study_id = '12345678') + or + study = aiplatform.Study(study_id = 'projects/123/locations/us-central1/studies/12345678') + + Args: + study_id (str): + Required. A fully-qualified study resource name or a study ID. + Example: "projects/123/locations/us-central1/studies/12345678" or "12345678" when + project and location are initialized or passed. + project (str): + Optional. Project to retrieve study from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve study from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Feature. Overrides + credentials set in aiplatform.init. + """ + base.VertexAiResourceNounWithFutureManager.__init__( + self, + project=project, + location=location, + credentials=credentials, + resource_name=study_id, + ) + self._gca_resource = self._get_gca_resource(resource_name=study_id) + + @classmethod + @base.optional_sync() + def create_or_load( + cls, + display_name: str, + problem: vz.ProblemStatement, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> client_abc.StudyInterface: + """Creates a Study resource. + + Example Usage: + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name='pr-auc', goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE)) + root = sc.search_space.select_root() + root.add_float_param( + 'learning_rate', 0.00001, 1.0, scale_type=pyvizier.ScaleType.LINEAR) + root.add_categorical_param('optimizer', ['adagrad', 'adam', 'experimental']) + study = aiplatform.Study.create_or_load(display_name='tuning_study', problem=sc) + + Args: + display_name (str): + Required. A name to describe the Study. It's unique per study. An existing study + will be returned if the study has the same display name. + problem (vz.ProblemStatement): + Required. Configurations of the study. It defines the problem to create the study. + project (str): + Optional. Project to retrieve study from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve study from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Feature. Overrides + credentials set in aiplatform.init. + Returns: + StudyInterface - The created study resource object. + """ + project = initializer.global_config.project if not project else project + location = initializer.global_config.location if not location else location + credentials = ( + initializer.global_config.credentials if not credentials else credentials + ) + + api_client = cls._instantiate_client( + location=location, + credentials=credentials, + ) + study = gca_study.Study( + display_name=display_name, study_spec=problem.to_proto() + ) + + try: + study = api_client.create_study( + parent=initializer.global_config.common_location_path( + project, + location, + ), + study=study, + ) + except exceptions.AlreadyExists: + _LOGGER.info("The study is already created. Using existing study.") + study = api_client.lookup_study( + request={ + "parent": initializer.global_config.common_location_path( + project, + location, + ), + "display_name": display_name, + }, + ) + + return Study(study.name) + + def get_trial(self, uid: int) -> client_abc.TrialInterface: + """Retrieves the trial under the study by given trial id. + + Args: + uid (int): Required. Unique identifier of the trial to search. + Returns: + TrialInterface - The trial resource object. + """ + study_path_components = self._parse_resource_name(self.resource_name) + return Trial( + Trial._format_resource_name( + project=study_path_components["project"], + location=study_path_components["location"], + study=study_path_components["study"], + trial=uid, + ), + credentials=self.credentials, + ) + + def trials( + self, trial_filter: Optional[vz.TrialFilter] = None + ) -> Collection[client_abc.TrialInterface]: + """Fetches a collection of trials. + + Args: + trial_filter (int): Optional. A filter for the trials. + Returns: + Collection[TrialInterface] - A list of trials resource object belonging + to the study. + """ + list_trials_request = {"parent": self.resource_name} + trials_response = self.api_client.list_trials(request=list_trials_request) + return [ + Trial._construct_sdk_resource_from_gapic( + trial, + project=self.project, + location=self.location, + credentials=self.credentials, + ) + for trial in trials_response.trials + ] + + def optimal_trials(self) -> Collection[client_abc.TrialInterface]: + """Returns optimal trial(s). + + Returns: + Collection[TrialInterface] - A list of optimal trials resource object. + """ + list_optimal_trials_request = {"parent": self.resource_name} + optimal_trials_response = self.api_client.list_optimal_trials( + request=list_optimal_trials_request + ) + return [ + Trial._construct_sdk_resource_from_gapic( + trial, + project=self.project, + location=self.location, + credentials=self.credentials, + ) + for trial in optimal_trials_response.optimal_trials + ] + + def materialize_study_config(self) -> vz.StudyConfig: + """#Materializes the study config. + + Returns: + StudyConfig - A deepcopy of StudyConfig from the study. + """ + study = self.api_client.get_study( + name=self.resource_name, credentials=self.credentials + ) + return copy.deepcopy(vz.StudyConfig.from_proto(study.study_spec)) + + @classmethod + def from_uid( + cls: Type[_T], + uid: str, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> _T: + """Fetches an existing study from the Vizier service. + + Args: + uid (str): Required. Unique identifier of the study. + Returns: + StudyInterface - The study resource object. + """ + project = initializer.global_config.project if not project else project + location = initializer.global_config.location if not location else location + credentials = ( + initializer.global_config.credentials if not credentials else credentials + ) + + return Study( + study_id=uid, project=project, location=location, credentials=credentials + ) + + def suggest( + self, *, count: Optional[int] = None, worker: str = "" + ) -> Collection[client_abc.TrialInterface]: + """Returns Trials to be evaluated by worker. + + Args: + count (int): Optional. Number of suggestions. + worker (str): When new Trials are generated, their `assigned_worker` field is + populated with this worker. suggest() first looks for existing Trials + that are assigned to `worker`, before generating new ones. + Returns: + Collection[TrialInterface] - A list of suggested trial resource objects. + """ + suggest_trials_lro = self.api_client.suggest_trials( + request={ + "parent": self.resource_name, + "suggestion_count": count, + "client_id": worker, + }, + ) + _LOGGER.log_action_started_against_resource_with_lro( + "Suggest", "study", self.__class__, suggest_trials_lro + ) + _LOGGER.info(self.client_class.get_gapic_client_class()) + trials = suggest_trials_lro.result() + _LOGGER.log_action_completed_against_resource("study", "suggested", self) + return [ + Trial._construct_sdk_resource_from_gapic( + trial, + project=self.project, + location=self.location, + credentials=self.credentials, + ) + for trial in trials.trials + ] + + def delete(self) -> None: + """Deletes the study.""" + self.api_client.delete_study(name=self.resource_name) diff --git a/google/cloud/aiplatform/vizier/trial.py b/google/cloud/aiplatform/vizier/trial.py new file mode 100644 index 0000000000..310f08e36b --- /dev/null +++ b/google/cloud/aiplatform/vizier/trial.py @@ -0,0 +1,180 @@ +# -*- coding: utf-8 -*- + +# Copyright 2022 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import copy + +from typing import Optional, TypeVar, Mapping, Any +from google.cloud.aiplatform.vizier.client_abc import TrialInterface + +from google.auth import credentials as auth_credentials +from google.cloud.aiplatform import base +from google.cloud.aiplatform import utils +from google.cloud.aiplatform.vizier import study +from google.cloud.aiplatform.vizier import pyvizier as vz + +_T = TypeVar("_T") +_LOGGER = base.Logger(__name__) + + +class Trial(base.VertexAiResourceNounWithFutureManager, TrialInterface): + """Manage Trial resource for Vertex Vizier.""" + + client_class = utils.VizierClientWithOverride + + _resource_noun = "trial" + _getter_method = "get_trial" + _list_method = "list_trials" + _delete_method = "delete_trial" + _parse_resource_name_method = "parse_trial_path" + _format_resource_name_method = "trial_path" + + def __init__( + self, + trial_name: str, + study_id: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """Retrieves an existing managed trial given a trial resource name or a study id. + + Example Usage: + trial = aiplatform.Trial(trial_name = 'projects/123/locations/us-central1/studies/12345678/trials/1') + or + trial = aiplatform.Trial(trial_name = '1', study_id = '12345678') + + Args: + trial_name (str): + Required. A fully-qualified trial resource name or a trial ID. + Example: "projects/123/locations/us-central1/studies/12345678/trials/1" or "12345678" when + project and location are initialized or passed. + study_id (str): + Optional. A fully-qualified study resource name or a study ID. + Example: "projects/123/locations/us-central1/studies/12345678" or "12345678" when + project and location are initialized or passed. + project (str): + Optional. Project to retrieve trial from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve trial from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Feature. Overrides + credentials set in aiplatform.init. + """ + + base.VertexAiResourceNounWithFutureManager.__init__( + self, + project=project, + location=location, + credentials=credentials, + resource_name=trial_name, + ) + self._gca_resource = self._get_gca_resource( + resource_name=trial_name, + parent_resource_name_fields={ + study.Study._resource_noun: study_id, + } + if study_id + else study_id, + ) + + @property + def uid(self) -> int: + """Unique identifier of the trial.""" + trial_path_components = self._parse_resource_name(self.resource_name) + return int(trial_path_components["trial"]) + + @property + def parameters(self) -> Mapping[str, Any]: + """Parameters of the trial.""" + trial = self.api_client.get_trial(name=self.resource_name) + return vz.TrialConverter.from_proto(trial).parameters + + @property + def status(self) -> vz.TrialStatus: + """Status of the Trial.""" + trial = self.api_client.get_trial(name=self.resource_name) + return vz.TrialConverter.from_proto(trial).status + + def delete(self) -> None: + """Deletes the Trial in Vizier service.""" + self.api_client.delete_trial(name=self.resource_name) + + def complete( + self, + measurement: Optional[vz.Measurement] = None, + *, + infeasible_reason: Optional[str] = None + ) -> Optional[vz.Measurement]: + """Completes the trial and #materializes the measurement. + + * If `measurement` is provided, then Vizier writes it as the trial's final + measurement and returns it. + * If `infeasible_reason` is provided, `measurement` is not needed. + * If neither is provided, then Vizier selects an existing (intermediate) + measurement to be the final measurement and returns it. + + Args: + measurement: Final measurement. + infeasible_reason: Indefeasibly reason for missing final measurement. + """ + complete_trial_request = {"name": self.resource_name} + if infeasible_reason is not None: + complete_trial_request["infeasible_reason"] = infeasible_reason + complete_trial_request["trial_infeasible"] = True + if measurement is not None: + complete_trial_request[ + "final_measurement" + ] = vz.MeasurementConverter.to_proto(measurement) + trial = self.api_client.complete_trial(request=complete_trial_request) + return ( + vz.MeasurementConverter.from_proto(trial.final_measurement) + if trial.final_measurement + else None + ) + + def should_stop(self) -> bool: + """Returns true if the Trial should stop.""" + check_trial_early_stopping_state_request = {"trial_name": self.resource_name} + should_stop_lro = self.api_client.check_trial_early_stopping_state( + request=check_trial_early_stopping_state_request + ) + _LOGGER.log_action_started_against_resource_with_lro( + "ShouldStop", "trial", self.__class__, should_stop_lro + ) + should_stop_lro.result() + _LOGGER.log_action_completed_against_resource("trial", "should_stop", self) + return should_stop_lro.result().should_stop + + def add_measurement(self, measurement: vz.Measurement) -> None: + """Adds an intermediate measurement.""" + add_trial_measurement_request = { + "trial_name": self.resource_name, + } + add_trial_measurement_request["measurement"] = vz.MeasurementConverter.to_proto( + measurement + ) + self.api_client.add_trial_measurement(request=add_trial_measurement_request) + + def materialize(self, *, include_all_measurements: bool = True) -> vz.Trial: + """#Materializes the Trial. + + Args: + include_all_measurements: If True, returned Trial includes all + intermediate measurements. The final measurement is always provided. + """ + trial = self.api_client.get_trial(name=self.resource_name) + return copy.deepcopy(vz.TrialConverter.from_proto(trial)) diff --git a/samples/model-builder/requirements.txt b/samples/model-builder/requirements.txt index 983ca6b4c6..d09caa45e0 100644 --- a/samples/model-builder/requirements.txt +++ b/samples/model-builder/requirements.txt @@ -1 +1 @@ -google-cloud-aiplatform \ No newline at end of file +google-cloud-aiplatform diff --git a/setup.py b/setup.py index 85231d8d7f..4ac1acf68a 100644 --- a/setup.py +++ b/setup.py @@ -58,6 +58,19 @@ datasets_extra_require = [ "pyarrow >= 3.0.0, < 8.0dev", ] + +vizier_extra_require = [ + "attrs==21.4.0", + "absl-py>=0.7", + "numpy>=1.19.0", + "protobuf>=3.6,<4.0", + "keras-tuner>=1.0,<2.0", + "portpicker==1.3.1", + "googleapis-common-protos==1.56.0", + "google-api-python-client==1.12.8", + "google-vizier==0.0.3a", +] + prediction_extra_require = [ "docker >= 5.0.3", # TODO: remove the upper bound after a new version is released. @@ -66,6 +79,7 @@ "starlette >= 0.17.1", "uvicorn >= 0.16.0", ] + private_endpoints_extra_require = [ "urllib3 >=1.21.1, <1.27", ] @@ -78,6 +92,7 @@ + featurestore_extra_require + pipelines_extra_require + datasets_extra_require + + vizier_extra_require + prediction_extra_require + private_endpoints_extra_require ) @@ -129,6 +144,7 @@ "lit": lit_extra_require, "cloud_profiler": profiler_extra_require, "pipelines": pipelines_extra_require, + "vizier": vizier_extra_require, "prediction": prediction_extra_require, "datasets": datasets_extra_require, "private_endpoints": private_endpoints_extra_require, diff --git a/tests/system/aiplatform/test_vizier.py b/tests/system/aiplatform/test_vizier.py new file mode 100644 index 0000000000..0f592cc7ce --- /dev/null +++ b/tests/system/aiplatform/test_vizier.py @@ -0,0 +1,106 @@ +import pytest + +from google.api_core import exceptions +from google.cloud import aiplatform +from google.cloud.aiplatform.vizier import Study +from google.cloud.aiplatform.vizier import Trial +from tests.system.aiplatform import e2e_base +from google.cloud.aiplatform.vizier import pyvizier + +_TEST_STUDY_ID = 123 + + +@pytest.mark.usefixtures("tear_down_resources") +class TestVizier(e2e_base.TestEndToEnd): + _temp_prefix = "temp_vertex_sdk_e2e_vizier_test" + + def test_vizier_lifecycle(self, shared_state): + aiplatform.init( + project=e2e_base._PROJECT, + location=e2e_base._LOCATION, + ) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name="pr-auc", goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + "learning_rate", 0.00001, 1.0, scale_type=pyvizier.ScaleType.LINEAR + ) + root.add_categorical_param("optimizer", ["adagrad", "adam", "experimental"]) + sc.automated_stopping_config = ( + pyvizier.AutomatedStoppingConfig.decay_curve_stopping_config(use_steps=True) + ) + + study = Study.create_or_load(display_name=self._temp_prefix, problem=sc) + shared_state["resources"] = [study] + trials = study.suggest(count=3, worker="halio_test_worker") + for trial in trials: + if not trial.should_stop(): + measurement = pyvizier.Measurement() + measurement.metrics["pr-auc"] = 0.4 + trial.add_measurement(measurement=measurement) + trial.complete(measurement=measurement) + optimal_trials = study.optimal_trials() + + for trial in study.trials(): + assert trial.status == pyvizier.TrialStatus.COMPLETED + assert optimal_trials[0].status == pyvizier.TrialStatus.COMPLETED + + def test_vizier_study_deletion(self, shared_state): + aiplatform.init( + project=e2e_base._PROJECT, + location=e2e_base._LOCATION, + ) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name="pr-auc", goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + "learning_rate", 0.00001, 1.0, scale_type=pyvizier.ScaleType.LINEAR + ) + root.add_categorical_param("optimizer", ["adagrad", "adam", "experimental"]) + sc.automated_stopping_config = ( + pyvizier.AutomatedStoppingConfig.decay_curve_stopping_config(use_steps=True) + ) + + study = Study.create_or_load(display_name=self._temp_prefix, problem=sc) + study.delete() + + with pytest.raises(exceptions.NotFound): + study = Study(study_id=study.name) + + def test_vizier_trial_deletion(self, shared_state): + aiplatform.init( + project=e2e_base._PROJECT, + location=e2e_base._LOCATION, + ) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name="pr-auc", goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + "learning_rate", 0.00001, 1.0, scale_type=pyvizier.ScaleType.LINEAR + ) + root.add_categorical_param("optimizer", ["adagrad", "adam", "experimental"]) + sc.automated_stopping_config = ( + pyvizier.AutomatedStoppingConfig.decay_curve_stopping_config(use_steps=True) + ) + + study = Study.create_or_load(display_name=self._temp_prefix, problem=sc) + trials = study.suggest(count=1, worker="halio_test_worker") + trials[0].delete() + + with pytest.raises(exceptions.NotFound): + study = Trial(study_id=study.name, trial_name=trials[0].name) diff --git a/tests/unit/aiplatform/test_vizier.py b/tests/unit/aiplatform/test_vizier.py new file mode 100644 index 0000000000..9b47761368 --- /dev/null +++ b/tests/unit/aiplatform/test_vizier.py @@ -0,0 +1,621 @@ +# -*- coding: utf-8 -*- + +# Copyright 2022 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + + +from unittest import mock +from importlib import reload +from unittest.mock import patch +from unittest.mock import ANY + +from google.api_core import exceptions +from google.api_core import operation + +from google.cloud import aiplatform +from google.cloud.aiplatform.vizier import Study +from google.cloud.aiplatform.vizier import Trial +from google.cloud.aiplatform import initializer +from google.cloud.aiplatform.vizier import pyvizier + +from google.cloud.aiplatform.compat.services import vizier_service_client +from google.cloud.aiplatform.compat.types import ( + study as gca_study, + vizier_service as gca_vizier_service, +) +from google.protobuf import duration_pb2 + + +# project +_TEST_PROJECT = "test-project" +_TEST_LOCATION = "us-central1" +_TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" + +# study +_TEST_STUDY_ID = "12345" +_TEST_STUDY_NAME = f"{_TEST_PARENT}/studies/{_TEST_STUDY_ID}" + +# trial +_TEST_TRIAL_ID = "1" +_TEST_TRIAL_NAME = f"{_TEST_STUDY_NAME}/trials/{_TEST_TRIAL_ID}" + +_TEST_METRIC_ID = "pr-auc" +_TEST_DISPLAY_NAME = "test_study_python_aiplatform" + +_TEST_PARAMETER_ID_1 = "learning_rate" +_TEST_PARAMETER_ID_MIN_VALUE_1 = 1e-05 +_TEST_PARAMETER_ID_MAX_VALUE_1 = 1.0 + +_TEST_PARAMETER_ID_2 = "optimizer" +_TEST_PARAMETER_VALUE_2 = ["adagrad", "adam", "experimental"] + +_TEST_STUDY = gca_study.Study( + display_name=_TEST_DISPLAY_NAME, + study_spec=gca_study.StudySpec( + algorithm=gca_study.StudySpec.Algorithm.RANDOM_SEARCH, + metrics=[ + gca_study.StudySpec.MetricSpec( + metric_id=_TEST_METRIC_ID, + goal=gca_study.StudySpec.MetricSpec.GoalType.MAXIMIZE, + ) + ], + parameters=[ + gca_study.StudySpec.ParameterSpec( + parameter_id=_TEST_PARAMETER_ID_1, + scale_type=gca_study.StudySpec.ParameterSpec.ScaleType.UNIT_LINEAR_SCALE, + double_value_spec=gca_study.StudySpec.ParameterSpec.DoubleValueSpec( + min_value=_TEST_PARAMETER_ID_MIN_VALUE_1, + max_value=_TEST_PARAMETER_ID_MAX_VALUE_1, + ), + ), + gca_study.StudySpec.ParameterSpec( + parameter_id=_TEST_PARAMETER_ID_2, + categorical_value_spec=gca_study.StudySpec.ParameterSpec.CategoricalValueSpec( + values=_TEST_PARAMETER_VALUE_2 + ), + ), + ], + ), +) + + +@pytest.fixture +def get_study_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "get_study" + ) as get_study_mock: + get_study_mock.return_value = gca_study.Study(name=_TEST_STUDY_NAME) + yield get_study_mock + + +@pytest.fixture +def get_trial_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "get_trial" + ) as get_trial_mock: + get_trial_mock.return_value = gca_study.Trial( + name=_TEST_TRIAL_NAME, + state=gca_study.Trial.State.ACTIVE, + parameters=[ + gca_study.Trial.Parameter( + parameter_id=_TEST_PARAMETER_ID_1, + value=_TEST_PARAMETER_ID_MIN_VALUE_1, + ) + ], + ) + yield get_trial_mock + + +@pytest.fixture +def create_study_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "create_study" + ) as create_study_mock: + create_study_mock.return_value = gca_study.Study( + name=_TEST_STUDY_NAME, + ) + yield create_study_mock + + +@pytest.fixture +def lookup_study_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "lookup_study" + ) as lookup_study_mock: + lookup_study_mock.return_value = gca_study.Study( + name=_TEST_STUDY_NAME, + ) + yield lookup_study_mock + + +@pytest.fixture +def suggest_trials_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "suggest_trials" + ) as suggest_trials_mock: + suggest_trials_lro_mock = mock.Mock(operation.Operation) + suggest_trials_lro_mock.result.return_value = ( + gca_vizier_service.SuggestTrialsResponse( + trials=[gca_study.Trial(name=_TEST_TRIAL_NAME)] + ) + ) + suggest_trials_mock.return_value = suggest_trials_lro_mock + yield suggest_trials_mock + + +@pytest.fixture +def list_optimal_trials_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "list_optimal_trials" + ) as list_optimal_trials_mock: + list_optimal_trials_mock.return_value = ( + gca_vizier_service.ListOptimalTrialsResponse( + optimal_trials=[gca_study.Trial(name=_TEST_TRIAL_NAME)] + ) + ) + yield list_optimal_trials_mock + + +@pytest.fixture +def list_trials_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "list_trials" + ) as list_trials_mock: + list_trials_mock.return_value = gca_vizier_service.ListTrialsResponse( + trials=[gca_study.Trial(name=_TEST_TRIAL_NAME)] + ) + yield list_trials_mock + + +@pytest.fixture +def delete_study_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "delete_study" + ) as delete_study_mock: + yield delete_study_mock + + +@pytest.fixture +def delete_trial_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "delete_trial" + ) as delete_trial_mock: + yield delete_trial_mock + + +@pytest.fixture +def complete_trial_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "complete_trial" + ) as complete_trial_mock: + complete_trial_mock.return_value = gca_study.Trial( + name=_TEST_TRIAL_NAME, + final_measurement=gca_study.Measurement( + step_count=3, + metrics=[gca_study.Measurement.Metric(metric_id="y", value=5)], + ), + ) + yield complete_trial_mock + + +@pytest.fixture +def complete_trial_empty_measurement_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "complete_trial" + ) as complete_trial_empty_measurement_mock: + complete_trial_empty_measurement_mock.return_value = gca_study.Trial( + name=_TEST_TRIAL_NAME + ) + yield complete_trial_empty_measurement_mock + + +@pytest.fixture +def should_stop_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "check_trial_early_stopping_state" + ) as should_stop_mock: + should_stop_lro_mock = mock.Mock(operation.Operation) + should_stop_lro_mock.result.return_value = ( + gca_vizier_service.CheckTrialEarlyStoppingStateResponse(should_stop=True) + ) + should_stop_mock.return_value = should_stop_lro_mock + yield should_stop_mock + + +@pytest.fixture +def create_study_mock_already_exists(): + with patch.object( + vizier_service_client.VizierServiceClient, "create_study" + ) as create_study_mock: + create_study_mock.side_effect = [ + exceptions.AlreadyExists("Study already exists."), + gca_study.Study( + name=_TEST_STUDY_NAME, + ), + ] + yield create_study_mock + + +@pytest.fixture +def add_measurement_mock(): + with patch.object( + vizier_service_client.VizierServiceClient, "add_trial_measurement" + ) as add_measurement_mock: + yield add_measurement_mock + + +@pytest.mark.usefixtures("google_auth_mock") +class TestStudy: + def setup_method(self): + reload(initializer) + reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + @pytest.mark.usefixtures("get_study_mock") + def test_create_study(self, create_study_mock): + aiplatform.init(project=_TEST_PROJECT) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name=_TEST_METRIC_ID, goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + _TEST_PARAMETER_ID_1, + _TEST_PARAMETER_ID_MIN_VALUE_1, + _TEST_PARAMETER_ID_MAX_VALUE_1, + scale_type=pyvizier.ScaleType.LINEAR, + ) + root.add_categorical_param(_TEST_PARAMETER_ID_2, _TEST_PARAMETER_VALUE_2) + + study = Study.create_or_load(display_name=_TEST_DISPLAY_NAME, problem=sc) + + create_study_mock.assert_called_once_with( + parent=_TEST_PARENT, study=_TEST_STUDY + ) + assert type(study) == Study + + @pytest.mark.usefixtures("get_study_mock") + def test_create_study_already_exists( + self, create_study_mock_already_exists, lookup_study_mock + ): + aiplatform.init(project=_TEST_PROJECT) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name=_TEST_METRIC_ID, goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + _TEST_PARAMETER_ID_1, + _TEST_PARAMETER_ID_MIN_VALUE_1, + _TEST_PARAMETER_ID_MAX_VALUE_1, + scale_type=pyvizier.ScaleType.LINEAR, + ) + root.add_categorical_param(_TEST_PARAMETER_ID_2, _TEST_PARAMETER_VALUE_2) + + study = Study.create_or_load(display_name=_TEST_DISPLAY_NAME, problem=sc) + + lookup_study_mock.assert_called_once_with( + request={"parent": _TEST_PARENT, "display_name": _TEST_DISPLAY_NAME} + ) + assert type(study) == Study + + @pytest.mark.usefixtures("get_study_mock") + def test_materialize_study_config(self, create_study_mock): + aiplatform.init(project=_TEST_PROJECT) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name=_TEST_METRIC_ID, goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + _TEST_PARAMETER_ID_1, + _TEST_PARAMETER_ID_MIN_VALUE_1, + _TEST_PARAMETER_ID_MAX_VALUE_1, + scale_type=pyvizier.ScaleType.LINEAR, + ) + root.add_categorical_param(_TEST_PARAMETER_ID_2, _TEST_PARAMETER_VALUE_2) + study = Study.create_or_load(display_name=_TEST_DISPLAY_NAME, problem=sc) + + study_config = study.materialize_study_config() + + create_study_mock.assert_called_once_with( + parent=_TEST_PARENT, study=_TEST_STUDY + ) + assert type(study_config) == pyvizier.StudyConfig + + @pytest.mark.usefixtures("get_study_mock", "get_trial_mock") + def test_suggest(self, create_study_mock, suggest_trials_mock): + aiplatform.init(project=_TEST_PROJECT) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name=_TEST_METRIC_ID, goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + _TEST_PARAMETER_ID_1, + _TEST_PARAMETER_ID_MIN_VALUE_1, + _TEST_PARAMETER_ID_MAX_VALUE_1, + scale_type=pyvizier.ScaleType.LINEAR, + ) + root.add_categorical_param(_TEST_PARAMETER_ID_2, _TEST_PARAMETER_VALUE_2) + study = Study.create_or_load(display_name=_TEST_DISPLAY_NAME, problem=sc) + + trials = study.suggest(count=5, worker="test_worker") + + suggest_trials_mock.assert_called_once_with( + request={ + "parent": _TEST_STUDY_NAME, + "suggestion_count": 5, + "client_id": "test_worker", + } + ) + assert type(trials[0]) == Trial + + @pytest.mark.usefixtures("get_study_mock") + def test_from_uid(self): + aiplatform.init(project=_TEST_PROJECT) + + study = Study.from_uid(uid=_TEST_STUDY_ID) + + assert type(study) == Study + assert study.name == _TEST_STUDY_ID + + @pytest.mark.usefixtures("get_study_mock") + def test_delete(self, create_study_mock, delete_study_mock): + aiplatform.init(project=_TEST_PROJECT) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name=_TEST_METRIC_ID, goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + _TEST_PARAMETER_ID_1, + _TEST_PARAMETER_ID_MIN_VALUE_1, + _TEST_PARAMETER_ID_MAX_VALUE_1, + scale_type=pyvizier.ScaleType.LINEAR, + ) + root.add_categorical_param(_TEST_PARAMETER_ID_2, _TEST_PARAMETER_VALUE_2) + study = Study.create_or_load(display_name=_TEST_DISPLAY_NAME, problem=sc) + + study.delete() + + delete_study_mock.assert_called_once_with(name=_TEST_STUDY_NAME) + + @pytest.mark.usefixtures("get_study_mock", "create_study_mock", "get_trial_mock") + def test_optimal_trials(self, list_optimal_trials_mock): + aiplatform.init(project=_TEST_PROJECT) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name=_TEST_METRIC_ID, goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + _TEST_PARAMETER_ID_1, + _TEST_PARAMETER_ID_MIN_VALUE_1, + _TEST_PARAMETER_ID_MAX_VALUE_1, + scale_type=pyvizier.ScaleType.LINEAR, + ) + root.add_categorical_param(_TEST_PARAMETER_ID_2, _TEST_PARAMETER_VALUE_2) + study = Study.create_or_load(display_name=_TEST_DISPLAY_NAME, problem=sc) + + trials = study.optimal_trials() + + list_optimal_trials_mock.assert_called_once_with( + request={"parent": _TEST_STUDY_NAME} + ) + assert type(trials[0]) == Trial + + @pytest.mark.usefixtures("get_study_mock", "create_study_mock", "get_trial_mock") + def test_list_trials(self, list_trials_mock): + aiplatform.init(project=_TEST_PROJECT) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name=_TEST_METRIC_ID, goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + _TEST_PARAMETER_ID_1, + _TEST_PARAMETER_ID_MIN_VALUE_1, + _TEST_PARAMETER_ID_MAX_VALUE_1, + scale_type=pyvizier.ScaleType.LINEAR, + ) + root.add_categorical_param(_TEST_PARAMETER_ID_2, _TEST_PARAMETER_VALUE_2) + study = Study.create_or_load(display_name=_TEST_DISPLAY_NAME, problem=sc) + + trials = study.trials() + + list_trials_mock.assert_called_once_with(request={"parent": _TEST_STUDY_NAME}) + assert type(trials[0]) == Trial + + @pytest.mark.usefixtures("get_study_mock", "create_study_mock") + def test_get_trial(self, get_trial_mock): + aiplatform.init(project=_TEST_PROJECT) + sc = pyvizier.StudyConfig() + sc.algorithm = pyvizier.Algorithm.RANDOM_SEARCH + sc.metric_information.append( + pyvizier.MetricInformation( + name=_TEST_METRIC_ID, goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE + ) + ) + root = sc.search_space.select_root() + root.add_float_param( + _TEST_PARAMETER_ID_1, + _TEST_PARAMETER_ID_MIN_VALUE_1, + _TEST_PARAMETER_ID_MAX_VALUE_1, + scale_type=pyvizier.ScaleType.LINEAR, + ) + root.add_categorical_param(_TEST_PARAMETER_ID_2, _TEST_PARAMETER_VALUE_2) + study = Study.create_or_load(display_name=_TEST_DISPLAY_NAME, problem=sc) + + trial = study.get_trial(1) + + get_trial_mock.assert_called_once_with(name=_TEST_TRIAL_NAME, retry=ANY) + assert type(trial) == Trial + + +@pytest.mark.usefixtures("google_auth_mock") +class TestTrial: + def setup_method(self): + reload(initializer) + reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + @pytest.mark.usefixtures("get_trial_mock") + def test_delete(self, delete_trial_mock): + aiplatform.init(project=_TEST_PROJECT) + trial = Trial(trial_name=_TEST_TRIAL_NAME) + + trial.delete() + + delete_trial_mock.assert_called_once_with(name=_TEST_TRIAL_NAME) + assert type(trial) == Trial + + @pytest.mark.usefixtures("get_trial_mock") + def test_complete(self, complete_trial_mock): + aiplatform.init(project=_TEST_PROJECT) + trial = Trial(trial_name=_TEST_TRIAL_NAME) + measurement = pyvizier.Measurement() + measurement.metrics["y"] = 4 + + measurement = trial.complete( + measurement=measurement, infeasible_reason="infeasible" + ) + + complete_trial_mock.assert_called_once_with( + request={ + "name": _TEST_TRIAL_NAME, + "infeasible_reason": "infeasible", + "trial_infeasible": True, + "final_measurement": gca_study.Measurement( + elapsed_duration=duration_pb2.Duration(), + metrics=[gca_study.Measurement.Metric(metric_id="y", value=4)], + ), + } + ) + assert type(measurement) == pyvizier.Measurement + + @pytest.mark.usefixtures("get_trial_mock") + def test_complete_empty_measurement(self, complete_trial_empty_measurement_mock): + aiplatform.init(project=_TEST_PROJECT) + trial = Trial(trial_name=_TEST_TRIAL_NAME) + measurement = pyvizier.Measurement() + measurement.metrics["y"] = 4 + + measurement = trial.complete( + measurement=measurement, infeasible_reason="infeasible" + ) + + complete_trial_empty_measurement_mock.assert_called_once_with( + request={ + "name": _TEST_TRIAL_NAME, + "infeasible_reason": "infeasible", + "trial_infeasible": True, + "final_measurement": gca_study.Measurement( + elapsed_duration=duration_pb2.Duration(), + metrics=[gca_study.Measurement.Metric(metric_id="y", value=4)], + ), + } + ) + assert measurement is None + + @pytest.mark.usefixtures("get_trial_mock") + def test_should_stop(self, should_stop_mock): + aiplatform.init(project=_TEST_PROJECT) + trial = Trial(trial_name=_TEST_TRIAL_NAME) + + should_stop = trial.should_stop() + + should_stop_mock.assert_called_once_with( + request={"trial_name": _TEST_TRIAL_NAME} + ) + assert should_stop is True + + @pytest.mark.usefixtures("get_trial_mock") + def test_add_measurement(self, add_measurement_mock): + aiplatform.init(project=_TEST_PROJECT) + trial = Trial(trial_name=_TEST_TRIAL_NAME) + measurement = pyvizier.Measurement() + measurement.metrics["y"] = 4 + + add_measurement = trial.add_measurement(measurement=measurement) + + add_measurement_mock.assert_called_once_with( + request={ + "trial_name": _TEST_TRIAL_NAME, + "measurement": gca_study.Measurement( + elapsed_duration=duration_pb2.Duration(), + metrics=[gca_study.Measurement.Metric(metric_id="y", value=4)], + ), + } + ) + assert add_measurement is None + + @pytest.mark.usefixtures("get_trial_mock") + def test_properties(self): + aiplatform.init(project=_TEST_PROJECT) + trial = Trial(trial_name=_TEST_TRIAL_NAME) + measurement = pyvizier.Measurement() + measurement.metrics["y"] = 4 + + uid = trial.uid + status = trial.status + parameters = trial.parameters + + assert uid == 1 + assert status == pyvizier.TrialStatus.ACTIVE + assert ( + parameters.get_value(_TEST_PARAMETER_ID_1) == _TEST_PARAMETER_ID_MIN_VALUE_1 + ) + + @pytest.mark.usefixtures("get_trial_mock") + def test_materialize(self): + aiplatform.init(project=_TEST_PROJECT) + trial = Trial(trial_name=_TEST_TRIAL_NAME) + measurement = pyvizier.Measurement() + measurement.metrics["y"] = 4 + + materialize_trial = trial.materialize() + + assert materialize_trial.id == 1 + assert ( + materialize_trial.parameters.get_value(_TEST_PARAMETER_ID_1) + == _TEST_PARAMETER_ID_MIN_VALUE_1 + )