feat: GenAI - Tuning - Added support for CMEK

PiperOrigin-RevId: 642696819

Author: Ark-kun

tests/unit/vertexai/test_tuning.py

@@ -19,10 +19,12 @@
 
 import copy
 import datetime
+import importlib
 from typing import Dict, Iterable
 from unittest import mock
 import uuid
 
+from google.cloud import aiplatform
 import vertexai
 from google.cloud.aiplatform import compat
 from google.cloud.aiplatform import initializer
@@ -150,6 +152,10 @@ class TestgenerativeModelTuning:
     """Unit tests for generative model tuning."""
 
     def setup_method(self):
+        importlib.reload(initializer)
+        importlib.reload(aiplatform)
+        importlib.reload(vertexai)
+
         vertexai.init(
             project=_TEST_PROJECT,
             location=_TEST_LOCATION,
@@ -197,3 +203,18 @@ def test_genai_tuning_service_supervised_tuning_tune_model(self):
         assert sft_tuning_job._experiment_name
         assert sft_tuning_job.tuned_model_name
         assert sft_tuning_job.tuned_model_endpoint_name
+
+    @mock.patch.object(
+        target=tuning.TuningJob,
+        attribute="client_class",
+        new=MockTuningJobClientWithOverride,
+    )
+    def test_genai_tuning_service_encryption_spec(self):
+        """Test that the global encryption spec propagates to the tuning job."""
+        vertexai.init(encryption_spec_key_name="test-key")
+
+        sft_tuning_job = supervised_tuning.train(
+            source_model="gemini-1.0-pro-001",
+            train_dataset="gs://some-bucket/some_dataset.jsonl",
+        )
+        assert sft_tuning_job.encryption_spec.kms_key_name == "test-key"

vertexai/tuning/_tuning.py

@@ -195,6 +195,11 @@ def _create(
         else:
             raise RuntimeError(f"Unsupported tuning_spec kind: {tuning_spec}")
 
+        if aiplatform_initializer.global_config.encryption_spec_key_name:
+            gca_tuning_job.encryption_spec.kms_key_name = (
+                aiplatform_initializer.global_config.encryption_spec_key_name
+            )
+
         tuning_job: TuningJob = cls._construct_sdk_resource_from_gapic(
             gapic_resource=gca_tuning_job,
             project=project,