chore: GenAI - Tuning - Switched tuning to use the v1beta1

PiperOrigin-RevId: 642382810

Author: Ark-kun

tests/unit/vertexai/test_tuning.py

@@ -28,9 +28,9 @@
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform import utils as aiplatform_utils
 from google.cloud.aiplatform.metadata import experiment_resources
-from google.cloud.aiplatform_v1.services import gen_ai_tuning_service
-from google.cloud.aiplatform_v1.types import job_state
-from google.cloud.aiplatform_v1.types import tuning_job as gca_tuning_job
+from google.cloud.aiplatform_v1beta1.services import gen_ai_tuning_service
+from google.cloud.aiplatform_v1beta1.types import job_state
+from google.cloud.aiplatform_v1beta1.types import tuning_job as gca_tuning_job
 from vertexai.preview import tuning
 from vertexai.preview.tuning import sft as supervised_tuning
 
@@ -131,11 +131,10 @@ def cancel_tuning_job(self, *, name: str, **_) -> None:
 
 class MockTuningJobClientWithOverride(aiplatform_utils.ClientWithOverride):
     _is_temporary = False
-    _default_version = compat.V1
+    _default_version = compat.V1BETA1
     _version_map = (
         (compat.V1, MockGenAiTuningServiceClient),
-        # v1beta1 version does not exist
-        # (compat.V1BETA1, gen_ai_tuning_service_v1beta1.client.JobServiceClient),
+        (compat.V1BETA1, MockGenAiTuningServiceClient),
     )
 
 

vertexai/tuning/_supervised_tuning.py

@@ -15,7 +15,7 @@
 
 from typing import Dict, Literal, Optional, Union
 
-from google.cloud.aiplatform_v1.types import (
+from google.cloud.aiplatform_v1beta1.types import (
     tuning_job as gca_tuning_job_types,
 )
 from vertexai import generative_models

vertexai/tuning/_tuning.py

@@ -29,10 +29,13 @@
 from google.cloud.aiplatform_v1.services import (
     gen_ai_tuning_service as gen_ai_tuning_service_v1,
 )
-from google.cloud.aiplatform_v1.types import (
+from google.cloud.aiplatform_v1beta1.services import (
+    gen_ai_tuning_service as gen_ai_tuning_service_v1beta1,
+)
+from google.cloud.aiplatform_v1beta1.types import (
     tuning_job as gca_tuning_job_types,
 )
-from google.cloud.aiplatform_v1 import types as gca_types
+from google.cloud.aiplatform_v1beta1 import types as gca_types
 
 from google.rpc import status_pb2  # type: ignore
 
@@ -42,11 +45,10 @@
 
 class TuningJobClientWithOverride(aiplatform_utils.ClientWithOverride):
     _is_temporary = True
-    _default_version = compat.V1
+    _default_version = compat.V1BETA1
     _version_map = (
         (compat.V1, gen_ai_tuning_service_v1.client.GenAiTuningServiceClient),
-        # v1beta1 version does not exist
-        # (compat.V1BETA1, gen_ai_tuning_service_v1beta1.client.GenAiTuningServiceClient),
+        (compat.V1BETA1, gen_ai_tuning_service_v1beta1.client.GenAiTuningServiceClient),
     )
 
 
@@ -66,7 +68,7 @@ class TuningJob(aiplatform_base._VertexAiResourceNounPlus):
     client_class = TuningJobClientWithOverride
 
     _gca_resource: gca_tuning_job_types.TuningJob
-    api_client: gen_ai_tuning_service_v1.client.GenAiTuningServiceClient
+    api_client: gen_ai_tuning_service_v1beta1.client.GenAiTuningServiceClient
 
     def __init__(self, tuning_job_name: str):
         super().__init__(resource_name=tuning_job_name)