feat: Vertex Vizier support in SDK.

.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

google/cloud/aiplatform/__init__.py

@@ -52,6 +52,8 @@
     CustomJob,
     HyperparameterTuningJob,
 )
+from google.cloud.aiplatform.vizier import Study
+from google.cloud.aiplatform.vizier import Trial
 from google.cloud.aiplatform.pipeline_jobs import PipelineJob
 from google.cloud.aiplatform.tensorboard import (
     Tensorboard,
@@ -139,12 +141,14 @@
     "PipelineJob",
     "PrivateEndpoint",
     "SequenceToSequencePlusForecastingTrainingJob",
+    "Study",
     "TabularDataset",
     "Tensorboard",
     "TensorboardExperiment",
     "TensorboardRun",
     "TensorboardTimeSeries",
     "TextDataset",
     "TimeSeriesDataset",
+    "Trial",
     "VideoDataset",
 )

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
@@ -113,6 +114,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:
 
@@ -130,6 +132,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
@@ -198,6 +201,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,

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,
 )

google/cloud/aiplatform/compat/types/__init__.py

@@ -79,6 +79,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,
@@ -143,6 +144,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__ = (

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,
 )
 
 

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",
+)

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.
+        """