fix: add v1 conversion value rule

google/cloud/aiplatform/helpers/_decorators.py

@@ -68,3 +68,5 @@ def _from_map(map_):
 
 marshal = Marshal(name="google.cloud.aiplatform.v1beta1")
 marshal.register(Value, ConversionValueRule(marshal=marshal))
+marshal = Marshal(name="google.cloud.aiplatform.v1")
+marshal.register(Value, ConversionValueRule(marshal=marshal))

google/cloud/aiplatform/initializer.py

@@ -38,6 +38,7 @@
     encryption_spec_v1beta1 as gca_encryption_spec_v1beta1,
 )
 
+
 class _Config:
     """Stores common parameters and options for API calls."""
 
@@ -100,7 +101,12 @@ def get_encryption_spec(
         self,
         encryption_spec_key_name: Optional[str],
         select_version: Optional[str] = compat.DEFAULT_VERSION,
-    ) -> Optional[Union[gca_encryption_spec_v1.EncryptionSpec, gca_encryption_spec_v1beta1.EncryptionSpec]]:
+    ) -> Optional[
+        Union[
+            gca_encryption_spec_v1.EncryptionSpec,
+            gca_encryption_spec_v1beta1.EncryptionSpec,
+        ]
+    ]:
         """Creates a gca_encryption_spec.EncryptionSpec instance from the given key name.
         If the provided key name is None, it uses the default key name if provided.
 

google/cloud/aiplatform/jobs.py

@@ -487,7 +487,7 @@ def create(
         # Optional Fields
         gapic_batch_prediction_job.encryption_spec = initializer.global_config.get_encryption_spec(
             encryption_spec_key_name=encryption_spec_key_name,
-            select_version=select_version
+            select_version=select_version,
         )
 
         if model_parameters:

tests/unit/aiplatform/test_end_to_end.py

@@ -25,15 +25,6 @@
 from google.cloud.aiplatform import schema
 from google.cloud.aiplatform import training_jobs
 
-from google.cloud.aiplatform_v1beta1.types import (
-    dataset as gca_dataset_v1beta1,
-    encryption_spec as gca_encryption_spec_v1beta1,
-    io as gca_io_v1beta1,
-    model as gca_model_v1beta1,
-    pipeline_state as gca_pipeline_state_v1beta1,
-    training_pipeline as gca_training_pipeline_v1beta1,
-)
-
 from google.cloud.aiplatform_v1.types import (
     dataset as gca_dataset,
     encryption_spec as gca_encryption_spec,

tests/unit/aiplatform/test_endpoints.py

@@ -29,9 +29,6 @@
 from google.cloud.aiplatform import models
 from google.cloud.aiplatform import utils
 
-from google.cloud.aiplatform_v1beta1.services.model_service import (
-    client as model_service_client_v1beta1,
-)
 from google.cloud.aiplatform_v1beta1.services.endpoint_service import (
     client as endpoint_service_client_v1beta1,
 )
@@ -40,11 +37,9 @@
 )
 from google.cloud.aiplatform_v1beta1.types import (
     endpoint as gca_endpoint_v1beta1,
-    model as gca_model_v1beta1,
     machine_resources as gca_machine_resources_v1beta1,
     prediction_service as gca_prediction_service_v1beta1,
     endpoint_service as gca_endpoint_service_v1beta1,
-    encryption_spec as gca_encryption_spec_v1beta1,
 )
 
 from google.cloud.aiplatform_v1.services.model_service import (
@@ -99,7 +94,9 @@
 _TEST_ACCELERATOR_TYPE = "NVIDIA_TESLA_P100"
 _TEST_ACCELERATOR_COUNT = 2
 
-_TEST_EXPLANATIONS = [gca_prediction_service_v1beta1.explanation.Explanation(attributions=[])]
+_TEST_EXPLANATIONS = [
+    gca_prediction_service_v1beta1.explanation.Explanation(attributions=[])
+]
 
 _TEST_ATTRIBUTIONS = [
     gca_prediction_service_v1beta1.explanation.Attribution(

tests/unit/aiplatform/test_jobs.py

@@ -39,19 +39,15 @@
     batch_prediction_job as gca_batch_prediction_job_v1beta1,
     explanation as gca_explanation_v1beta1,
     io as gca_io_v1beta1,
-    job_state as gca_job_state_v1beta1,
     machine_resources as gca_machine_resources_v1beta1,
 )
 
-from google.cloud.aiplatform_v1.services.job_service import (
-    client as job_service_client,
-)
+from google.cloud.aiplatform_v1.services.job_service import client as job_service_client
 
 from google.cloud.aiplatform_v1.types import (
     batch_prediction_job as gca_batch_prediction_job,
     io as gca_io,
     job_state as gca_job_state,
-    machine_resources as gca_machine_resources,
 )
 
 _TEST_PROJECT = "test-project"
@@ -485,7 +481,9 @@ def test_batch_predict_gcs_source_bq_dest(
 
     @pytest.mark.parametrize("sync", [True, False])
     @pytest.mark.usefixtures("get_batch_prediction_job_mock")
-    def test_batch_predict_with_all_args(self, create_batch_prediction_job_with_explanations_mock, sync):
+    def test_batch_predict_with_all_args(
+        self, create_batch_prediction_job_with_explanations_mock, sync
+    ):
         aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
         creds = auth_credentials.AnonymousCredentials()
 
@@ -518,7 +516,9 @@ def test_batch_predict_with_all_args(self, create_batch_prediction_job_with_expl
             model=_TEST_MODEL_NAME,
             input_config=gca_batch_prediction_job_v1beta1.BatchPredictionJob.InputConfig(
                 instances_format="jsonl",
-                gcs_source=gca_io_v1beta1.GcsSource(uris=[_TEST_BATCH_PREDICTION_GCS_SOURCE]),
+                gcs_source=gca_io_v1beta1.GcsSource(
+                    uris=[_TEST_BATCH_PREDICTION_GCS_SOURCE]
+                ),
             ),
             output_config=gca_batch_prediction_job_v1beta1.BatchPredictionJob.OutputConfig(
                 gcs_destination=gca_io_v1beta1.GcsDestination(

tests/unit/aiplatform/test_models.py

@@ -43,7 +43,6 @@
     env_var as gca_env_var_v1beta1,
     explanation as gca_explanation_v1beta1,
     io as gca_io_v1beta1,
-    job_state as gca_job_state_v1beta1,
     model as gca_model_v1beta1,
     endpoint as gca_endpoint_v1beta1,
     machine_resources as gca_machine_resources_v1beta1,
@@ -55,15 +54,12 @@
 from google.cloud.aiplatform_v1.services.endpoint_service import (
     client as endpoint_service_client,
 )
-from google.cloud.aiplatform_v1.services.job_service import (
-    client as job_service_client,
-)
+from google.cloud.aiplatform_v1.services.job_service import client as job_service_client
 from google.cloud.aiplatform_v1.services.model_service import (
     client as model_service_client,
 )
 from google.cloud.aiplatform_v1.types import (
     batch_prediction_job as gca_batch_prediction_job,
-    env_var as gca_env_var,
     io as gca_io,
     job_state as gca_job_state,
     model as gca_model,
@@ -184,6 +180,7 @@ def get_model_mock():
         )
         yield get_model_mock
 
+
 @pytest.fixture
 def get_model_with_explanations_mock():
     with mock.patch.object(
@@ -194,6 +191,7 @@ def get_model_with_explanations_mock():
         )
         yield get_model_mock
 
+
 @pytest.fixture
 def get_model_with_custom_location_mock():
     with mock.patch.object(
@@ -244,7 +242,6 @@ def upload_model_with_explanations_mock():
         yield upload_model_mock
 
 
-
 @pytest.fixture
 def upload_model_with_custom_project_mock():
     with mock.patch.object(
@@ -300,7 +297,6 @@ def deploy_model_mock():
         yield deploy_model_mock
 
 
-
 @pytest.fixture
 def deploy_model_with_explanations_mock():
     with mock.patch.object(
@@ -343,6 +339,7 @@ def create_batch_prediction_job_mock():
         create_batch_prediction_job_mock.return_value = batch_prediction_job_mock
         yield create_batch_prediction_job_mock
 
+
 @pytest.fixture
 def create_batch_prediction_job_with_explanations_mock():
     with mock.patch.object(
@@ -355,6 +352,7 @@ def create_batch_prediction_job_with_explanations_mock():
         create_batch_prediction_job_mock.return_value = batch_prediction_job_mock
         yield create_batch_prediction_job_mock
 
+
 @pytest.fixture
 def create_client_mock():
     with mock.patch.object(
@@ -746,7 +744,9 @@ def test_deploy_no_endpoint_dedicated_resources(self, deploy_model_mock, sync):
         "get_endpoint_mock", "get_model_mock", "create_endpoint_mock"
     )
     @pytest.mark.parametrize("sync", [True, False])
-    def test_deploy_no_endpoint_with_explanations(self, deploy_model_with_explanations_mock, sync):
+    def test_deploy_no_endpoint_with_explanations(
+        self, deploy_model_with_explanations_mock, sync
+    ):
         aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
         test_model = models.Model(_TEST_ID)
         test_endpoint = test_model.deploy(
@@ -834,9 +834,7 @@ def test_init_aiplatform_with_encryption_key_name_and_batch_predict_gcs_source_a
             ),
             input_config=gca_batch_prediction_job.BatchPredictionJob.InputConfig(
                 instances_format="jsonl",
-                gcs_source=gca_io.GcsSource(
-                    uris=[_TEST_BATCH_PREDICTION_GCS_SOURCE]
-                ),
+                gcs_source=gca_io.GcsSource(uris=[_TEST_BATCH_PREDICTION_GCS_SOURCE]),
             ),
             output_config=gca_batch_prediction_job.BatchPredictionJob.OutputConfig(
                 gcs_destination=gca_io.GcsDestination(
@@ -940,7 +938,9 @@ def test_batch_predict_gcs_source_bq_dest(
 
     @pytest.mark.parametrize("sync", [True, False])
     @pytest.mark.usefixtures("get_model_mock", "get_batch_prediction_job_mock")
-    def test_batch_predict_with_all_args(self, create_batch_prediction_job_with_explanations_mock, sync):
+    def test_batch_predict_with_all_args(
+        self, create_batch_prediction_job_with_explanations_mock, sync
+    ):
         aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
         test_model = models.Model(_TEST_ID)
         creds = auth_credentials.AnonymousCredentials()
@@ -977,7 +977,9 @@ def test_batch_predict_with_all_args(self, create_batch_prediction_job_with_expl
             ),
             input_config=gca_batch_prediction_job_v1beta1.BatchPredictionJob.InputConfig(
                 instances_format="jsonl",
-                gcs_source=gca_io_v1beta1.GcsSource(uris=[_TEST_BATCH_PREDICTION_GCS_SOURCE]),
+                gcs_source=gca_io_v1beta1.GcsSource(
+                    uris=[_TEST_BATCH_PREDICTION_GCS_SOURCE]
+                ),
             ),
             output_config=gca_batch_prediction_job_v1beta1.BatchPredictionJob.OutputConfig(
                 gcs_destination=gca_io_v1beta1.GcsDestination(