fix: Fixed the PipelineJob bucket creation after a breaking change in…

… Google Cloud Storage client library

Fixes #2936

PiperOrigin-RevId: 600653928

google/cloud/aiplatform/utils/gcs_utils.py

@@ -234,10 +234,10 @@ def create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
         credentials=credentials,
     )
 
-    pipelines_bucket = storage.Blob.from_string(
+    pipelines_bucket = storage.Bucket.from_string(
         uri=output_artifacts_gcs_dir,
         client=storage_client,
-    ).bucket
+    )
 
     if not pipelines_bucket.exists():
         _logger.info(

tests/unit/aiplatform/test_utils.py

@@ -128,14 +128,6 @@ def mock_storage_blob_download_to_filename():
         yield mock_blob_download_to_filename
 
 
-@pytest.fixture()
-def mock_bucket_not_exist():
-    with patch("google.cloud.storage.Blob.from_string") as mock_bucket_not_exist, patch(
-        "google.cloud.storage.Bucket.exists", return_value=False
-    ):
-        yield mock_bucket_not_exist
-
-
 def test_invalid_region_raises_with_invalid_region():
     with pytest.raises(ValueError):
         aiplatform.utils.validate_region(region="us-east5")
@@ -587,8 +579,11 @@ def test_generate_gcs_directory_for_pipeline_artifacts(self):
         )
         assert output == "gs://project-vertex-pipelines-us-central1/output_artifacts/"
 
+    @patch.object(storage.Bucket, "exists", return_value=False)
+    @patch.object(storage, "Client")
+    @patch.object(gcs_utils.resource_manager_utils, "get_project_number", return_value=12345)
     def test_create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
-        self, mock_bucket_not_exist, mock_storage_client
+        self, mock_bucket_not_exist, mock_storage_client, mock_get_project_number
     ):
         output = (
             gcs_utils.create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
@@ -597,6 +592,7 @@ def test_create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
         )
         assert mock_storage_client.called
         assert mock_bucket_not_exist.called
+        assert mock_get_project_number.called
         assert (
             output == "gs://test-project-vertex-pipelines-us-central1/output_artifacts/"
         )