feat: Vertex Vizier support in SDK. (#1434)

* Added the Vizier client in the aiplatform folder.

* Copied the pyvizier from the open source vizier.

* Added the Trial in the aiplatform folder.

* Forked the pyvizier from the pyVizier.

* imported the objects from the pyvizier.

* Added the Vizier client and types in the aiplatform init file.

* Made the pyvizier converters compatible with the Vertex SDK proto.

* Forked the framework for the Vertex Vizier from the Open Source Vizier.

* Implemented the interfaces for Study and Trial. Added the unit tests and system tests for them.

* Setup the dependencies for the Vizier.

* Fix the lint error by running the nox -s blacken.

* Fixed the lint errors for the Vizier.

* Made the unit test import the google credentials.

* Disable the coverage check for the pyvizier. It will be imported from the open source vizier.

* Remove the converage dependency to avoid the conflicts.

* Fixing the py-3.9 issue in the sample/module-builder

* Convert the lambda function to avoid the import numpy be called in the confest.py test

* Revert the requirements file.

* Fix the lint error by running the nox -s blacken.

* Fixed the syntax issue in the setup.py

* Setup the local package

* scrube the TODO since it's the documentation ticket.

* Addresses the comments.

* Ran blacken on the test_vizier file.

* Import the OSS in the Vertex SDK.

* Already imported the package from open source vizier. Removing the code copied from oss.

* Import the google-vizier and fix the dependencies for Vertex Vizier.

* Configured the dependency of the google-vizier.

* Ran the nox -s blacken.

* Fixed the lint issue.

* Clean the debugging logs.

* Decouple the Study, Trial and the aiplatform to make the sample test pass.

* Fixed the issue in the system test got an unexpected keyword argument 'credentials'

* Ran the nox -s blacken to format the python file.

* Fixed the unit test failure.

* Add the wrapper to give more error information about the vizier import error.

Co-authored-by: Rosie Zou <rosiezou@users.noreply.github.com>

.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/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,

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

@@ -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__ = (

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