Skip to content

Commit 27ee0dc

Browse files
yinghsienwucopybara-github
authored andcommitted
fix: warn when the pipeline output bucket is not owned by the project, and add aiplatform.init(enforce_bucket_ownership=True) to raise an error instead
PiperOrigin-RevId: 1017744214
1 parent 87a7380 commit 27ee0dc

7 files changed

Lines changed: 376 additions & 23 deletions

File tree

‎google/cloud/aiplatform/initializer.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@
6262

6363
_TOP_GOOGLE_CONSTRUCTOR_METHOD_TAG = "top_google_constructor_method"
6464

65+
_ENFORCE_BUCKET_OWNERSHIP_ENV_VAR = "GOOGLE_CLOUD_AIPLATFORM_ENFORCE_BUCKET_OWNERSHIP"
66+
6567

6668
class _Product(enum.Enum):
6769
"""Notebook product types."""
@@ -133,6 +135,7 @@ def __init__(self):
133135
self._request_metadata = None
134136
self._resource_type = None
135137
self._async_rest_credentials = None
138+
self._enforce_bucket_ownership = None
136139

137140
def init(
138141
self,
@@ -153,6 +156,7 @@ def init(
153156
api_key: Optional[str] = None,
154157
api_transport: Optional[str] = None,
155158
request_metadata: Optional[Sequence[Tuple[str, str]]] = None,
159+
enforce_bucket_ownership: Optional[bool] = None,
156160
):
157161
"""Updates common initialization parameters with provided options.
158162
@@ -220,6 +224,13 @@ def init(
220224
beta state (preview).
221225
request_metadata:
222226
Optional. Additional gRPC metadata to send with every client request.
227+
enforce_bucket_ownership (bool):
228+
Optional. If True, raise an error when an existing GCS bucket
229+
that the SDK writes pipeline artifacts to is not owned by the
230+
project, or when its owner cannot be verified. If False or not
231+
set, only a warning is shown. Can also be enabled with the
232+
GOOGLE_CLOUD_AIPLATFORM_ENFORCE_BUCKET_OWNERSHIP=true
233+
environment variable.
223234
Raises:
224235
ValueError:
225236
If experiment_description is provided but experiment is not.
@@ -287,6 +298,8 @@ def init(
287298
self._request_metadata = request_metadata
288299
if api_key is not None:
289300
self._api_key = api_key
301+
if enforce_bucket_ownership is not None:
302+
self._enforce_bucket_ownership = enforce_bucket_ownership
290303
self._resource_type = None
291304

292305
# Finally, perform secondary state updates
@@ -388,6 +401,15 @@ def staging_bucket(self) -> Optional[str]:
388401
"""Default staging bucket, if provided."""
389402
return self._staging_bucket
390403

404+
@property
405+
def enforce_bucket_ownership(self) -> bool:
406+
"""Whether to raise an error for a GCS bucket owned by another project."""
407+
if self._enforce_bucket_ownership is not None:
408+
return self._enforce_bucket_ownership
409+
return os.environ.get(
410+
_ENFORCE_BUCKET_OWNERSHIP_ENV_VAR, ""
411+
).strip().lower() in ("1", "true", "yes")
412+
391413
@property
392414
def credentials(self) -> Optional[auth_credentials.Credentials]:
393415
"""Default credentials."""

‎google/cloud/aiplatform/pipeline_jobs.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -468,6 +468,9 @@ def submit(
468468
location=self.location,
469469
credentials=self.credentials,
470470
)
471+
except gcs_utils.BucketOwnershipError:
472+
# Raised only in strict mode: aiplatform.init(enforce_bucket_ownership=True).
473+
raise
471474
except: # noqa: E722
472475
_LOGGER.exception(
473476
"Error when trying to get or create a GCS bucket for the pipeline output artifacts"

‎google/cloud/aiplatform/utils/gcs_utils.py‎

Lines changed: 123 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
import os
2626
import pathlib
2727
import tempfile
28-
from typing import Optional, TYPE_CHECKING
28+
from typing import Optional, Tuple, TYPE_CHECKING
2929
import warnings
3030

3131
from google.auth import credentials as auth_credentials
@@ -63,6 +63,45 @@
6363
_DEFAULT_STAGING_BUCKET_SALT = str(uuid.uuid4())
6464

6565

66+
class BucketOwnershipError(ValueError):
67+
"""Raised when a GCS bucket does not belong to the expected project."""
68+
69+
70+
def _get_bucket_project_numbers(
71+
bucket: storage.Bucket,
72+
expected_project: str,
73+
client: storage.Client,
74+
) -> Tuple[Optional[str], Optional[str]]:
75+
"""Gets the project number of a GCS bucket and of the expected project.
76+
77+
Args:
78+
bucket: The GCS bucket to check.
79+
expected_project: The project ID that should own the bucket.
80+
client: Storage client instance.
81+
82+
Returns:
83+
A tuple of the bucket's project number and the expected project's
84+
number. Each value is None if it could not be retrieved.
85+
"""
86+
bucket_project_number = None
87+
expected_project_number = None
88+
try:
89+
bucket.reload(client=client)
90+
if bucket.project_number is not None:
91+
bucket_project_number = str(bucket.project_number)
92+
except Exception:
93+
pass
94+
try:
95+
expected_project_number = str(
96+
resource_manager_utils.get_project_number(
97+
expected_project, credentials=getattr(client, "_credentials", None)
98+
)
99+
)
100+
except Exception:
101+
pass
102+
return bucket_project_number, expected_project_number
103+
104+
66105
def _verify_bucket_ownership(
67106
bucket: storage.Bucket,
68107
expected_project: str,
@@ -81,15 +120,70 @@ def _verify_bucket_ownership(
81120
Returns:
82121
True if the bucket belongs to the expected project, False otherwise.
83122
"""
84-
try:
85-
bucket.reload(client=client)
86-
bucket_project_number = str(bucket.project_number)
87-
expected_project_number = str(
88-
resource_manager_utils.get_project_number(expected_project)
123+
bucket_project_number, expected_project_number = _get_bucket_project_numbers(
124+
bucket, expected_project, client
125+
)
126+
return (
127+
bucket_project_number is not None
128+
and bucket_project_number == expected_project_number
129+
)
130+
131+
132+
def _check_pipeline_bucket_ownership(
133+
bucket: storage.Bucket,
134+
expected_project: str,
135+
client: storage.Client,
136+
is_sdk_generated: bool = False,
137+
) -> None:
138+
"""Warns, or raises, if a bucket is not owned by the project.
139+
140+
Args:
141+
bucket: The existing GCS bucket for pipeline artifacts.
142+
expected_project: The project ID that should own the bucket.
143+
client: Storage client instance.
144+
is_sdk_generated: Whether the SDK generated the bucket name. A
145+
generated name can be predictable, so a mismatch always raises.
146+
147+
Raises:
148+
BucketOwnershipError: If the bucket is not owned by `expected_project`,
149+
or its owner cannot be verified, and either `is_sdk_generated` is
150+
True or `aiplatform.init(enforce_bucket_ownership=True)` is set.
151+
"""
152+
bucket_project_number, expected_project_number = _get_bucket_project_numbers(
153+
bucket, expected_project, client
154+
)
155+
if bucket_project_number is not None and (
156+
bucket_project_number == expected_project_number
157+
):
158+
return
159+
160+
if bucket_project_number is None or expected_project_number is None:
161+
message = (
162+
f'Could not verify that GCS bucket "{bucket.name}" is owned by '
163+
f'project "{expected_project}".'
89164
)
90-
return bucket_project_number == expected_project_number
91-
except Exception:
92-
return False
165+
else:
166+
message = (
167+
f'GCS bucket "{bucket.name}" is owned by project number '
168+
f'"{bucket_project_number}", not by project "{expected_project}" '
169+
f'(project number "{expected_project_number}").'
170+
)
171+
message += (
172+
" Pipeline artifacts will be written to this bucket. If you do not "
173+
"own this bucket, other users may be able to read or modify the "
174+
"artifacts. Use a bucket owned by your project, for example by "
175+
"passing pipeline_root or by setting "
176+
"aiplatform.init(staging_bucket='gs://your-bucket')."
177+
)
178+
179+
if is_sdk_generated or initializer.global_config.enforce_bucket_ownership:
180+
raise BucketOwnershipError(message)
181+
warnings.warn(
182+
f"{message} To raise an error instead, set "
183+
"aiplatform.init(enforce_bucket_ownership=True).",
184+
UserWarning,
185+
stacklevel=2,
186+
)
93187

94188

95189
def blob_from_uri(uri: str, client: storage.Client) -> storage.Blob:
@@ -335,12 +429,12 @@ def create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
335429
service_account = service_account or initializer.global_config.service_account
336430
credentials = credentials or initializer.global_config.credentials
337431

432+
default_output_artifacts_gcs_dir = generate_gcs_directory_for_pipeline_artifacts(
433+
project=project,
434+
location=location,
435+
)
338436
output_artifacts_gcs_dir = (
339-
output_artifacts_gcs_dir
340-
or generate_gcs_directory_for_pipeline_artifacts(
341-
project=project,
342-
location=location,
343-
)
437+
output_artifacts_gcs_dir or default_output_artifacts_gcs_dir
344438
)
345439

346440
# Creating the bucket if needed
@@ -377,13 +471,21 @@ def create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
377471
f"serviceAccount:{service_account}"
378472
)
379473
pipelines_bucket.set_iam_policy(bucket_iam_policy)
380-
elif not _verify_bucket_ownership(pipelines_bucket, project, storage_client):
381-
raise ValueError(
382-
f'Output artifacts bucket "{pipelines_bucket.name}" exists but does '
383-
f'not belong to project "{project}". This may indicate a '
384-
f"bucket squatting attack. Please provide an explicit "
385-
f"output_artifacts_gcs_dir parameter or configure one via "
386-
f"aiplatform.init(output_artifacts_gcs_dir='gs://your-bucket')."
474+
else:
475+
# PipelineJob generates its default pipeline_root before calling this
476+
# function, so compare bucket names instead of checking for None.
477+
is_sdk_generated = (
478+
pipelines_bucket.name
479+
== bucket_from_uri(
480+
uri=default_output_artifacts_gcs_dir,
481+
client=storage_client,
482+
).name
483+
)
484+
_check_pipeline_bucket_ownership(
485+
pipelines_bucket,
486+
project,
487+
storage_client,
488+
is_sdk_generated=is_sdk_generated,
387489
)
388490
return output_artifacts_gcs_dir
389491

‎tests/unit/aiplatform/test_initializer.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,24 @@ def test_init_service_account_sets_service_account(self):
138138
initializer.global_config.init(service_account=_TEST_SERVICE_ACCOUNT)
139139
assert initializer.global_config.service_account == _TEST_SERVICE_ACCOUNT
140140

141+
def test_enforce_bucket_ownership_defaults_to_false(self):
142+
with patch.dict(
143+
os.environ, {initializer._ENFORCE_BUCKET_OWNERSHIP_ENV_VAR: ""}
144+
):
145+
assert initializer.global_config.enforce_bucket_ownership is False
146+
147+
def test_init_enforce_bucket_ownership_sets_enforce_bucket_ownership(self):
148+
initializer.global_config.init(enforce_bucket_ownership=True)
149+
assert initializer.global_config.enforce_bucket_ownership is True
150+
151+
def test_enforce_bucket_ownership_from_env_var(self):
152+
with patch.dict(
153+
os.environ, {initializer._ENFORCE_BUCKET_OWNERSHIP_ENV_VAR: "true"}
154+
):
155+
assert initializer.global_config.enforce_bucket_ownership is True
156+
initializer.global_config.init(enforce_bucket_ownership=False)
157+
assert initializer.global_config.enforce_bucket_ownership is False
158+
141159
@patch.object(_experiment_tracker, "set_experiment")
142160
def test_init_experiment_sets_experiment(self, set_experiment_mock):
143161
initializer.global_config.init(experiment=_TEST_EXPERIMENT)

‎tests/unit/aiplatform/test_pipeline_jobs.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1627,6 +1627,42 @@ def test_submit_call_gcs_utils_get_or_create_with_correct_arguments(
16271627
credentials=_TEST_CREDENTIALS,
16281628
)
16291629

1630+
@pytest.mark.parametrize(
1631+
"job_spec",
1632+
[_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_SPEC_YAML, _TEST_PIPELINE_JOB],
1633+
)
1634+
def test_submit_raises_when_pipeline_bucket_belongs_to_another_project(
1635+
self,
1636+
mock_pipeline_service_create,
1637+
mock_pipeline_service_get,
1638+
job_spec,
1639+
mock_load_yaml_and_json,
1640+
):
1641+
job = pipeline_jobs.PipelineJob(
1642+
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
1643+
template_path=_TEST_TEMPLATE_PATH,
1644+
job_id=_TEST_PIPELINE_JOB_ID,
1645+
parameter_values=_TEST_PIPELINE_PARAMETER_VALUES,
1646+
enable_caching=True,
1647+
project=_TEST_PROJECT,
1648+
pipeline_root=_TEST_GCS_OUTPUT_DIRECTORY,
1649+
location=_TEST_LOCATION,
1650+
credentials=_TEST_CREDENTIALS,
1651+
)
1652+
1653+
with mock.patch.object(
1654+
gcs_utils,
1655+
"create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist",
1656+
side_effect=gcs_utils.BucketOwnershipError("not owned"),
1657+
):
1658+
with pytest.raises(gcs_utils.BucketOwnershipError):
1659+
job.submit(
1660+
service_account=_TEST_SERVICE_ACCOUNT,
1661+
network=_TEST_NETWORK,
1662+
)
1663+
1664+
mock_pipeline_service_create.assert_not_called()
1665+
16301666
@pytest.mark.parametrize(
16311667
"job_spec",
16321668
[_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_SPEC_YAML, _TEST_PIPELINE_JOB],

0 commit comments

Comments
 (0)