Repository navigation
DeploymentResourcePool.query_deployed_models returns page containers instead of DeployedModelRef objects #7230
Description
Environment details
- OS type and version: Linux
- Python version: 3.11.16
google-cloud-aiplatformversion: 2.5.0 (main at faed3f4)
Steps to reproduce
- Call
query_deployed_models()on aDeploymentResourcePoolwhere the response has one page holding two deployed model refs. - Read
deployed_model_idfrom the first item of the returned list.
query_deployed_models documents its return value as a list of DeployedModelRef objects, but it builds the list with itertools.chain(page.deployed_model_refs for page in response.pages). Called with a single generator, chain does not flatten anything, so each item in the result is one page's RepeatedComposite of refs. With one page holding two refs the method returns a list of length 1, and every item is a container, not a ref. The same line is in google/cloud/aiplatform/models.py and google/cloud/aiplatform/preview/models.py.
The existing unit test uses two pages with one ref each and only checks len(dm_refs) == 2, so it passes with the bug.
Code example
Only the two client calls are stubbed, everything else is the real SDK code path:
from unittest import mock
from google.auth.credentials import AnonymousCredentials
from google.cloud import aiplatform
from google.cloud.aiplatform.compat.services import deployment_resource_pool_service_client as svc
from google.cloud.aiplatform.compat.types import (
deployed_model_ref as ref_types,
deployment_resource_pool as drp_types,
deployment_resource_pool_service as svc_types,
)
NAME = "projects/p/locations/us-central1/deploymentResourcePools/pool"
aiplatform.init(project="p", location="us-central1", credentials=AnonymousCredentials())
one_page = svc_types.QueryDeployedModelsResponse(deployed_model_refs=[
ref_types.DeployedModelRef(endpoint="e1", deployed_model_id="m1"),
ref_types.DeployedModelRef(endpoint="e1", deployed_model_id="m2"),
])
pager = mock.Mock(pages=[one_page])
Client = svc.DeploymentResourcePoolServiceClient
with mock.patch.object(Client, "get_deployment_resource_pool", return_value=drp_types.DeploymentResourcePool(name=NAME)), \
mock.patch.object(Client, "query_deployed_models", return_value=pager):
refs = aiplatform.DeploymentResourcePool(NAME).query_deployed_models()
print("len:", len(refs), "| item types:", [type(r).__name__ for r in refs])
print(refs[0].deployed_model_id)Stack trace
len: 1 | item types: ['RepeatedComposite']
AttributeError: 'RepeatedComposite' object has no attribute 'deployed_model_id'
The fix is to use itertools.chain.from_iterable(...) in both classes. I have a PR with tests ready and will link it here.
Activity