Skip to content

Commit caa416a

Browse files
racinmatcopybara-github
authored andcommitted
fix: using location and project parsed from resource name if not provided
test: versioning registry location COPYBARA_INTEGRATE_REVIEW=#6641 from racinmat:versioning-registry 9e3c763 PiperOrigin-RevId: 925456884
1 parent be945df commit caa416a

2 files changed

Lines changed: 29 additions & 2 deletions

File tree

google/cloud/aiplatform/models.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5312,8 +5312,8 @@ def __init__(
53125312
# Create ModelRegistry with the unversioned resource name
53135313
self._registry = ModelRegistry(
53145314
self.resource_name,
5315-
location=location,
5316-
project=project,
5315+
location=location or self.location,
5316+
project=project or self.project,
53175317
credentials=credentials,
53185318
)
53195319

tests/unit/aiplatform/test_models.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4701,6 +4701,33 @@ def test_init_with_version_arg(self, get_model_with_version):
47014701
# The Model yielded from upload SHOULD have a version in the versioned resource name
47024702
assert model.versioned_resource_name.endswith(f"@{_TEST_VERSION_ID}")
47034703

4704+
def test_versioning_registry_uses_location_from_resource_name(
4705+
self, create_client_mock
4706+
):
4707+
# Regression test for https://github.com/googleapis/python-aiplatform/issues/2608:
4708+
# When a Model is initialized with a fully-qualified resource name that encodes a
4709+
# non-default location, the versioning registry client must use that location, not
4710+
# the global default from aiplatform.init().
4711+
models.Model(_TEST_MODEL_RESOURCE_NAME_CUSTOM_LOCATION)
4712+
create_client_mock.assert_any_call(
4713+
client_class=utils.ModelClientWithOverride,
4714+
credentials=initializer.global_config.credentials,
4715+
location_override=_TEST_LOCATION_2,
4716+
appended_user_agent=None,
4717+
)
4718+
4719+
def test_versioning_registry_uses_project_from_resource_name(
4720+
self, get_model_with_custom_project_mock
4721+
):
4722+
# Regression test for https://github.com/googleapis/python-aiplatform/issues/2608:
4723+
# When a Model is initialized with a fully-qualified resource name that encodes a
4724+
# non-default project, the versioning registry must use that project, not the
4725+
# global default from aiplatform.init().
4726+
model = models.Model(_TEST_MODEL_RESOURCE_NAME_CUSTOM_PROJECT)
4727+
assert model._registry.model_resource_name.startswith(
4728+
f"projects/{_TEST_PROJECT_2}/"
4729+
)
4730+
47044731
@pytest.mark.parametrize(
47054732
"parent,location,project",
47064733
[

0 commit comments

Comments
 (0)