feat: add custom and hp tuning

google/cloud/aiplatform/__init__.py

@@ -26,9 +26,15 @@
     TimeSeriesDataset,
     VideoDataset,
 )
+from google.cloud.aiplatform import hyperparameter_tuning
+from google.cloud.aiplatform.metadata import metadata
 from google.cloud.aiplatform.models import Endpoint
 from google.cloud.aiplatform.models import Model
-from google.cloud.aiplatform.jobs import BatchPredictionJob
+from google.cloud.aiplatform.jobs import (
+    BatchPredictionJob,
+    CustomJob,
+    HyperparameterTuningJob,
+)
 from google.cloud.aiplatform.training_jobs import (
     CustomTrainingJob,
     CustomContainerTrainingJob,
@@ -39,7 +45,6 @@
     AutoMLTextTrainingJob,
     AutoMLVideoTrainingJob,
 )
-from google.cloud.aiplatform.metadata import metadata
 
 """
 Usage:
@@ -60,6 +65,7 @@
     "explain",
     "gapic",
     "init",
+    "hyperparameter_tuning",
     "log_params",
     "log_metrics",
     "get_experiment_df",
@@ -71,11 +77,13 @@
     "AutoMLTextTrainingJob",
     "AutoMLVideoTrainingJob",
     "BatchPredictionJob",
+    "CustomJob",
     "CustomTrainingJob",
     "CustomContainerTrainingJob",
     "CustomPythonPackageTrainingJob",
     "Endpoint",
     "ImageDataset",
+    "HyperparameterTuningJob",
     "Model",
     "TabularDataset",
     "TextDataset",

google/cloud/aiplatform/base.py

@@ -101,6 +101,29 @@ def log_create_complete(
             f"{variable_name} = aiplatform.{cls.__name__}('{resource.name}')"
         )
 
+    def log_create_complete_with_getter(
+        self,
+        cls: Type["AiPlatformResourceNoun"],
+        resource: proto.Message,
+        variable_name: str,
+    ):
+        """Logs create event is complete.
+
+        Will also include code snippet to instantiate resource in SDK.
+
+        Args:
+            cls (AiPlatformResourceNoun):
+                AI Platform Resource Noun class that is being created.
+            resource (proto.Message):
+                AI Platform Resourc proto.Message
+            variable_name (str): Name of variable to use for code snippet
+        """
+        self._logger.info(f"{cls.__name__} created. Resource name: {resource.name}")
+        self._logger.info(f"To use this {cls.__name__} in another session:")
+        self._logger.info(
+            f"{variable_name} = aiplatform.{cls.__name__}.get('{resource.name}')"
+        )
+
     def log_action_start_against_resource(
         self, action: str, noun: str, resource_noun_obj: "AiPlatformResourceNoun"
     ):
@@ -543,6 +566,11 @@ def update_time(self) -> datetime.datetime:
         self._sync_gca_resource()
         return self._gca_resource.update_time
 
+    @property
+    def gca_resource(self) -> proto.Message:
+        """The underlying resource proto represenation."""
+        return self._gca_resource
+
     def __repr__(self) -> str:
         return f"{object.__repr__(self)} \nresource name: {self.resource_name}"
 

google/cloud/aiplatform/compat/__init__.py

@@ -70,6 +70,7 @@
     types.prediction_service = types.prediction_service_v1beta1
     types.specialist_pool = types.specialist_pool_v1beta1
     types.specialist_pool_service = types.specialist_pool_service_v1beta1
+    types.study = types.study_v1beta1
     types.training_pipeline = types.training_pipeline_v1beta1
     types.metadata_service = types.metadata_service_v1beta1
     types.tensorboard_service = types.tensorboard_service_v1beta1
@@ -120,6 +121,7 @@
     types.prediction_service = types.prediction_service_v1
     types.specialist_pool = types.specialist_pool_v1
     types.specialist_pool_service = types.specialist_pool_service_v1
+    types.study = types.study_v1
     types.training_pipeline = types.training_pipeline_v1
 
 __all__ = (

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

@@ -49,6 +49,7 @@
     prediction_service as prediction_service_v1beta1,
     specialist_pool as specialist_pool_v1beta1,
     specialist_pool_service as specialist_pool_service_v1beta1,
+    study as study_v1beta1,
     training_pipeline as training_pipeline_v1beta1,
     metadata_service as metadata_service_v1beta1,
     tensorboard_service as tensorboard_service_v1beta1,
@@ -90,6 +91,7 @@
     prediction_service as prediction_service_v1,
     specialist_pool as specialist_pool_v1,
     specialist_pool_service as specialist_pool_service_v1,
+    study as study_v1,
     training_pipeline as training_pipeline_v1,
 )
 

google/cloud/aiplatform/hyperparameter_tuning.py

@@ -0,0 +1,215 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2021 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 abc
+from typing import Dict, List, Optional, Sequence, Tuple, Union
+
+import proto
+
+from google.cloud.aiplatform.compat.types import study as gca_study_compat
+
+_SCALE_TYPE_MAP = {
+    "linear": gca_study_compat.StudySpec.ParameterSpec.ScaleType.UNIT_LINEAR_SCALE,
+    "log": gca_study_compat.StudySpec.ParameterSpec.ScaleType.UNIT_LOG_SCALE,
+    "reverse_log": gca_study_compat.StudySpec.ParameterSpec.ScaleType.UNIT_REVERSE_LOG_SCALE,
+    "unspecified": gca_study_compat.StudySpec.ParameterSpec.ScaleType.SCALE_TYPE_UNSPECIFIED,
+}
+
+
+class _ParameterSpec(metaclass=abc.ABCMeta):
+    """Base class represents a single parameter to optimize."""
+
+    def __init__(
+        self,
+        conditional_parameter_spec: Optional[Dict[str, "_ParameterSpec"]] = None,
+        parent_values: Optional[List[Union[float, int, str]]] = None,
+    ):
+
+        self.conditional_parameter_spec = conditional_parameter_spec
+        self.parent_values = parent_values
+
+    @property
+    @classmethod
+    @abc.abstractmethod
+    def _proto_parameter_value_class(self) -> proto.Message:
+        """The proto represenation of this parameter."""
+        pass
+
+    @property
+    @classmethod
+    @abc.abstractmethod
+    def _parameter_value_map(self) -> Tuple[Tuple[str, str]]:
+        """A Tuple map of parameter key to underlying proto key."""
+        pass
+
+    @property
+    @classmethod
+    @abc.abstractmethod
+    def _parameter_spec_value_key(self) -> Tuple[Tuple[str, str]]:
+        """The ParameterSpec key this parameter should be assigned."""
+        pass
+
+    @property
+    def _proto_parameter_value_spec(self) -> proto.Message:
+        """Converts this parameter to it's parameter value representation."""
+        proto_parameter_value_spec = self._proto_parameter_value_class()
+        for self_attr_key, proto_attr_key in self._parameter_value_map:
+            setattr(
+                proto_parameter_value_spec, proto_attr_key, getattr(self, self_attr_key)
+            )
+        return proto_parameter_value_spec
+
+    def _to_parameter_spec(
+        self, parameter_id: str
+    ) -> gca_study_compat.StudySpec.ParameterSpec:
+        """Converts this parameter to ParameterSpec."""
+        # TODO: Conditional parameters
+        parameter_spec = gca_study_compat.StudySpec.ParameterSpec(
+            parameter_id=parameter_id,
+            scale_type=_SCALE_TYPE_MAP.get(getattr(self, "scale", "unspecified")),
+        )
+
+        setattr(
+            parameter_spec,
+            self._parameter_spec_value_key,
+            self._proto_parameter_value_spec,
+        )
+
+        return parameter_spec
+
+
+class DoubleParameterSpec(_ParameterSpec):
+
+    _proto_parameter_value_class = (
+        gca_study_compat.StudySpec.ParameterSpec.DoubleValueSpec
+    )
+    _parameter_value_map = (("min", "min_value"), ("max", "max_value"))
+    _parameter_spec_value_key = "double_value_spec"
+
+    def __init__(
+        self, min: float, max: float, scale: str,
+    ):
+        """
+        Value specification for a parameter in ``DOUBLE`` type.
+
+        Args:
+            min (float):
+                Required. Inclusive minimum value of the
+                parameter.
+            max (float):
+                Required. Inclusive maximum value of the
+                parameter.
+            scale (str): 
+                Required. The type of scaling that should be applied to this parameter.
+
+                Accepts: 'linear', 'log', 'reverse_log'
+        """
+
+        super().__init__()
+
+        self.min = min
+        self.max = max
+        self.scale = scale
+
+
+class IntegerParameterSpec(_ParameterSpec):
+
+    _proto_parameter_value_class = (
+        gca_study_compat.StudySpec.ParameterSpec.IntegerValueSpec
+    )
+    _parameter_value_map = (("min", "min_value"), ("max", "max_value"))
+    _parameter_spec_value_key = "integer_value_spec"
+
+    def __init__(
+        self, min: int, max: int, scale: str,
+    ):
+        """
+        Value specification for a parameter in ``INTEGER`` type.
+
+        Args:
+            min (float):
+                Required. Inclusive minimum value of the
+                parameter.
+            max (float):
+                Required. Inclusive maximum value of the
+                parameter.
+            scale (str): 
+                Required. The type of scaling that should be applied to this parameter.
+
+                Accepts: 'linear', 'log', 'reverse_log'
+        """
+
+        super().__init__()
+
+        self.min = min
+        self.max = max
+        self.scale = scale
+
+
+class CategoricalParameterSpec(_ParameterSpec):
+
+    _proto_parameter_value_class = (
+        gca_study_compat.StudySpec.ParameterSpec.CategoricalValueSpec
+    )
+    _parameter_value_map = (("values", "values"),)
+    _parameter_spec_value_key = "categorical_value_spec"
+
+    def __init__(
+        self, values: Sequence[str],
+    ):
+        """Value specification for a parameter in ``CATEGORICAL`` type.
+
+        Args:
+            values (Sequence[str]): 
+                Required. The list of possible categories.
+        """
+
+        super().__init__()
+
+        self.values = values
+
+
+class DiscreteParameterSpec(_ParameterSpec):
+
+    _proto_parameter_value_class = (
+        gca_study_compat.StudySpec.ParameterSpec.DiscreteValueSpec
+    )
+    _parameter_value_map = (("values", "values"),)
+    _parameter_spec_value_key = "discrete_value_spec"
+
+    def __init__(
+        self, values: Sequence[float], scale: str,
+    ):
+        """Value specification for a parameter in ``DISCRETE`` type.
+
+        values (Sequence[float]):
+            Required. A list of possible values.
+            The list should be in increasing order and at
+            least 1e-10 apart. For instance, this parameter
+            might have possible settings of 1.5, 2.5, and
+            4.0. This list should not contain more than
+            1,000 values.
+        scale (str): 
+                Required. The type of scaling that should be applied to this parameter.
+
+                Accepts: 'linear', 'log', 'reverse_log'
+        """
+
+        super().__init__()
+
+        self.values = values
+        self.scale = scale