feat: Add preview HyperparameterTuningJob which can be run on persist…

…ent resource PiperOrigin-RevId: 582032717
googleapis · Nov 13, 2023 · 0da8c53 · 0da8c53
1 parent 9f46f7c
commit 0da8c53
Show file tree

Hide file tree

Showing 4 changed files with 776 additions and 3 deletions.
diff --git a/google/cloud/aiplatform/hyperparameter_tuning.py b/google/cloud/aiplatform/hyperparameter_tuning.py
@@ -20,7 +20,10 @@
 
 import proto
 
-from google.cloud.aiplatform.compat.types import study as gca_study_compat
+from google.cloud.aiplatform.compat.types import (
+    study_v1beta1 as gca_study_compat_v1beta1,
+    study as gca_study_compat,
+)
 
 SEARCH_ALGORITHM_TO_PROTO_VALUE = {
     "random": gca_study_compat.StudySpec.Algorithm.RANDOM_SEARCH,
@@ -91,6 +94,46 @@ def _proto_parameter_value_spec(self) -> proto.Message:
             )
         return proto_parameter_value_spec
 
+    @property
+    def _proto_parameter_value_spec_v1beta1(self) -> proto.Message:
+        """Converts this parameter to it's parameter value representation."""
+        if isinstance(
+            self._proto_parameter_value_class(),
+            gca_study_compat.StudySpec.ParameterSpec.DoubleValueSpec,
+        ):
+            proto_parameter_value_spec = (
+                gca_study_compat_v1beta1.StudySpec.ParameterSpec.DoubleValueSpec()
+            )
+        elif isinstance(
+            self._proto_parameter_value_class(),
+            gca_study_compat.StudySpec.ParameterSpec.IntegerValueSpec,
+        ):
+            proto_parameter_value_spec = (
+                gca_study_compat_v1beta1.StudySpec.ParameterSpec.IntegerValueSpec()
+            )
+        elif isinstance(
+            self._proto_parameter_value_class(),
+            gca_study_compat.StudySpec.ParameterSpec.CategoricalValueSpec,
+        ):
+            proto_parameter_value_spec = (
+                gca_study_compat_v1beta1.StudySpec.ParameterSpec.CategoricalValueSpec()
+            )
+        elif isinstance(
+            self._proto_parameter_value_class(),
+            gca_study_compat.StudySpec.ParameterSpec.DiscreteValueSpec,
+        ):
+            proto_parameter_value_spec = (
+                gca_study_compat_v1beta1.StudySpec.ParameterSpec.DiscreteValueSpec()
+            )
+        else:
+            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:
@@ -129,6 +172,46 @@ def _to_parameter_spec(
 
         return parameter_spec
 
+    def _to_parameter_spec_v1beta1(
+        self, parameter_id: str
+    ) -> gca_study_compat_v1beta1.StudySpec.ParameterSpec:
+        """Converts this parameter to ParameterSpec."""
+        conditions = []
+        if self.conditional_parameter_spec is not None:
+            for (conditional_param_id, spec) in self.conditional_parameter_spec.items():
+                condition = (
+                    gca_study_compat_v1beta1.StudySpec.ParameterSpec.ConditionalParameterSpec()
+                )
+                if self._parameter_spec_value_key == _INT_VALUE_SPEC:
+                    condition.parent_int_values = gca_study_compat_v1beta1.StudySpec.ParameterSpec.ConditionalParameterSpec.IntValueCondition(
+                        values=spec.parent_values
+                    )
+                elif self._parameter_spec_value_key == _CATEGORICAL_VALUE_SPEC:
+                    condition.parent_categorical_values = gca_study_compat_v1beta1.StudySpec.ParameterSpec.ConditionalParameterSpec.CategoricalValueCondition(
+                        values=spec.parent_values
+                    )
+                elif self._parameter_spec_value_key == _DISCRETE_VALUE_SPEC:
+                    condition.parent_discrete_values = gca_study_compat_v1beta1.StudySpec.ParameterSpec.ConditionalParameterSpec.DiscreteValueCondition(
+                        values=spec.parent_values
+                    )
+                condition.parameter_spec = spec._to_parameter_spec_v1beta1(
+                    conditional_param_id
+                )
+                conditions.append(condition)
+        parameter_spec = gca_study_compat_v1beta1.StudySpec.ParameterSpec(
+            parameter_id=parameter_id,
+            scale_type=_SCALE_TYPE_MAP.get(getattr(self, "scale", "unspecified")),
+            conditional_parameter_specs=conditions,
+        )
+
+        setattr(
+            parameter_spec,
+            self._parameter_spec_value_key,
+            self._proto_parameter_value_spec_v1beta1,
+        )
+
+        return parameter_spec
+
 
 class DoubleParameterSpec(_ParameterSpec):
 

diff --git a/google/cloud/aiplatform/jobs.py b/google/cloud/aiplatform/jobs.py
@@ -2036,7 +2036,7 @@ def job_spec(self):
         return self._gca_resource.job_spec
 
 
-class HyperparameterTuningJob(_RunnableJob):
+class HyperparameterTuningJob(_RunnableJob, base.PreviewMixin):
     """Vertex AI Hyperparameter Tuning Job."""
 
     _resource_noun = "hyperparameterTuningJobs"
@@ -2047,6 +2047,9 @@ class HyperparameterTuningJob(_RunnableJob):
     _parse_resource_name_method = "parse_hyperparameter_tuning_job_path"
     _format_resource_name_method = "hyperparameter_tuning_job_path"
     _job_type = "training"
+    _preview_class = (
+        "google.cloud.aiplatform.aiplatform.preview.jobs.HyperparameterTuningJob"
+    )
 
     def __init__(
         self,