Skip to content

Commit a8b01e0

Browse files
Ark-kuncopybara-github
authored andcommitted
fix: Fixed the PipelineJob bucket creation after a breaking change in Google Cloud Storage client library
Fixes #2936 PiperOrigin-RevId: 600929548
1 parent 59fae72 commit a8b01e0

File tree

2 files changed

+9
-11
lines changed

2 files changed

+9
-11
lines changed

google/cloud/aiplatform/utils/gcs_utils.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -234,10 +234,10 @@ def create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
234234
credentials=credentials,
235235
)
236236

237-
pipelines_bucket = storage.Blob.from_string(
237+
pipelines_bucket = storage.Bucket.from_string(
238238
uri=output_artifacts_gcs_dir,
239239
client=storage_client,
240-
).bucket
240+
)
241241

242242
if not pipelines_bucket.exists():
243243
_logger.info(

tests/unit/aiplatform/test_utils.py

+7-9
Original file line numberDiff line numberDiff line change
@@ -128,14 +128,6 @@ def mock_storage_blob_download_to_filename():
128128
yield mock_blob_download_to_filename
129129

130130

131-
@pytest.fixture()
132-
def mock_bucket_not_exist():
133-
with patch("google.cloud.storage.Blob.from_string") as mock_bucket_not_exist, patch(
134-
"google.cloud.storage.Bucket.exists", return_value=False
135-
):
136-
yield mock_bucket_not_exist
137-
138-
139131
def test_invalid_region_raises_with_invalid_region():
140132
with pytest.raises(ValueError):
141133
aiplatform.utils.validate_region(region="us-east5")
@@ -587,8 +579,13 @@ def test_generate_gcs_directory_for_pipeline_artifacts(self):
587579
)
588580
assert output == "gs://project-vertex-pipelines-us-central1/output_artifacts/"
589581

582+
@patch.object(storage.Bucket, "exists", return_value=False)
583+
@patch.object(storage, "Client")
584+
@patch.object(
585+
gcs_utils.resource_manager_utils, "get_project_number", return_value=12345
586+
)
590587
def test_create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
591-
self, mock_bucket_not_exist, mock_storage_client
588+
self, mock_bucket_not_exist, mock_storage_client, mock_get_project_number
592589
):
593590
output = (
594591
gcs_utils.create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
@@ -597,6 +594,7 @@ def test_create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
597594
)
598595
assert mock_storage_client.called
599596
assert mock_bucket_not_exist.called
597+
assert mock_get_project_number.called
600598
assert (
601599
output == "gs://test-project-vertex-pipelines-us-central1/output_artifacts/"
602600
)

0 commit comments

Comments
 (0)