Skip to content

Commit 8779d73

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Add service_account parameter to Model Garden deploy methods.
PiperOrigin-RevId: 1013562281
1 parent d62419a commit 8779d73

5 files changed

Lines changed: 189 additions & 0 deletions

File tree

‎agentplatform/model_garden/_model_garden.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -462,6 +462,7 @@ def deploy(
462462
serving_container_health_probe_timeout_seconds: Optional[int] = None,
463463
enable_private_service_connect: bool = False,
464464
psc_project_allow_list: Optional[Sequence[str]] = None,
465+
service_account: Optional[str] = None,
465466
) -> aiplatform.Endpoint:
466467
"""Deploys an Open Model to an endpoint.
467468
@@ -601,6 +602,9 @@ def deploy(
601602
connect.
602603
psc_project_allow_list (Sequence[str]): The list of projects that are
603604
allowed to access the endpoint over private service connect.
605+
service_account (str): Optional. The service account that the
606+
deployed model runs as. The caller must have the
607+
`iam.serviceAccounts.actAs` permission on this service account.
604608
605609
Returns:
606610
endpoint (aiplatform.Endpoint):
@@ -683,6 +687,9 @@ def deploy(
683687
if system_labels:
684688
request.deploy_config.system_labels = system_labels
685689

690+
if service_account:
691+
request.deploy_config.service_account = service_account
692+
686693
if serving_container_spec:
687694
if not serving_container_spec.image_uri:
688695
raise ValueError(
@@ -1126,6 +1133,7 @@ def deploy(
11261133
enable_private_service_connect: bool = False,
11271134
psc_project_allow_list: Optional[List[str]] = None,
11281135
deploy_request_timeout: Optional[float] = None,
1136+
service_account: Optional[str] = None,
11291137
) -> aiplatform.Endpoint:
11301138
"""Deploys a Custom Model to an endpoint.
11311139
@@ -1174,6 +1182,9 @@ def deploy(
11741182
access the endpoint over private service connect.
11751183
deploy_request_timeout: The timeout for the deploy request. Default is 2
11761184
hours.
1185+
service_account (str): Optional. The service account that the
1186+
deployed model runs as. The caller must have the
1187+
`iam.serviceAccounts.actAs` permission on this service account.
11771188
11781189
Returns:
11791190
endpoint (aiplatform.Endpoint):
@@ -1194,6 +1205,7 @@ def deploy(
11941205
enable_private_service_connect=enable_private_service_connect,
11951206
psc_project_allow_list=psc_project_allow_list,
11961207
deploy_request_timeout=deploy_request_timeout,
1208+
service_account=service_account,
11971209
)
11981210

11991211
def _deploy_model_registry_model(self) -> aiplatform.Endpoint:
@@ -1218,6 +1230,7 @@ def _deploy_gcs_uri(
12181230
endpoint_display_name: Optional[str] = None,
12191231
model_display_name: Optional[str] = None,
12201232
deploy_request_timeout: Optional[float] = None,
1233+
service_account: Optional[str] = None,
12211234
) -> aiplatform.Endpoint:
12221235
"""Deploys a Custom Model to an endpoint.
12231236
@@ -1268,6 +1281,9 @@ def _deploy_gcs_uri(
12681281
access the endpoint over private service connect.
12691282
deploy_request_timeout: The timeout for the deploy request. Default is 2
12701283
hours.
1284+
service_account (str): Optional. The service account that the
1285+
deployed model runs as. The caller must have the
1286+
`iam.serviceAccounts.actAs` permission on this service account.
12711287
12721288
Returns:
12731289
endpoint (aiplatform.Endpoint):
@@ -1296,6 +1312,8 @@ def has_all_or_none_values(var1, var2, var3) -> bool:
12961312
request.model_config.model_display_name = model_display_name
12971313
if system_labels:
12981314
request.deploy_config.system_labels = system_labels
1315+
if service_account:
1316+
request.deploy_config.service_account = service_account
12991317

13001318
if enable_private_service_connect and psc_project_allow_list:
13011319
request.endpoint_config.private_service_connect_config = (
@@ -1407,6 +1425,7 @@ def deploy(
14071425
endpoint_display_name: Optional[str] = None,
14081426
model_display_name: Optional[str] = None,
14091427
deploy_request_timeout: Optional[float] = None,
1428+
service_account: Optional[str] = None,
14101429
) -> aiplatform.Endpoint:
14111430
"""Deploys an Open Model to an endpoint.
14121431
@@ -1439,6 +1458,9 @@ def deploy(
14391458
model_display_name: The display name of the uploaded model.
14401459
deploy_request_timeout: The timeout for the deploy request. Default is 2
14411460
hours.
1461+
service_account (str): Optional. The service account that the
1462+
deployed model runs as. The caller must have the
1463+
`iam.serviceAccounts.actAs` permission on this service account.
14421464
14431465
Returns:
14441466
endpoint (aiplatform.Endpoint):
@@ -1469,6 +1491,9 @@ def deploy(
14691491
)
14701492
request.deploy_config.dedicated_resources = dedicated_resources
14711493

1494+
if service_account:
1495+
request.deploy_config.service_account = service_account
1496+
14721497
_LOGGER.info(f"Deploying model: {self._model_name}")
14731498

14741499
operation_future = self._model_garden_client.deploy(request)

‎google/cloud/aiplatform_v1beta1/types/model_garden_service.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -457,6 +457,17 @@ class DeployConfig(proto.Message):
457457
Optional. System labels for Model Garden
458458
deployments. These labels are managed by Google
459459
and for tracking purposes only.
460+
service_account (str):
461+
Optional. The service account that the
462+
container of the deployed model runs as. Specify
463+
the email address of the service account. If not
464+
specified, the container runs as a service
465+
account that has no access to the resource
466+
project.
467+
468+
Users deploying the model must have the
469+
``iam.serviceAccounts.actAs`` permission on this
470+
service account.
460471
"""
461472

462473
dedicated_resources: machine_resources.DedicatedResources = proto.Field(
@@ -473,6 +484,10 @@ class DeployConfig(proto.Message):
473484
proto.STRING,
474485
number=3,
475486
)
487+
service_account: str = proto.Field(
488+
proto.STRING,
489+
number=5,
490+
)
476491

477492
publisher_model_name: str = proto.Field(
478493
proto.STRING,

‎tests/unit/agentplatform/model_garden/test_model_garden.py‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070
_TEST_PARTNER_MODEL_SIMPLIFIED_RESOURCE_NAME = "ai21/jamba-large-1.6@001"
7171

7272
_TEST_GCS_URI = "gs://some-bucket/some-model"
73+
_TEST_SERVICE_ACCOUNT = "[email protected]"
7374
_TEST_ENDPOINT_NAME = "projects/test-project/locations/us-central1/endpoints/1234567890"
7475
_TEST_MODEL_NAME = "projects/test-project/locations/us-central1/models/9876543210"
7576
_TEST_IMAGE_URI = "us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00"
@@ -653,6 +654,26 @@ def test_deploy_with_display_names_success(self, deploy_mock):
653654
)
654655
)
655656

657+
def test_deploy_with_service_account_success(self, deploy_mock):
658+
"""Tests deploying a model with a service account."""
659+
aiplatform.init(
660+
project=_TEST_PROJECT,
661+
location=_TEST_LOCATION,
662+
)
663+
model = model_garden.PartnerModel(
664+
model_name=_TEST_PARTNER_MODEL_FULL_RESOURCE_NAME
665+
)
666+
model.deploy(service_account=_TEST_SERVICE_ACCOUNT)
667+
deploy_mock.assert_called_once_with(
668+
types.DeployRequest(
669+
publisher_model_name=_TEST_PARTNER_MODEL_FULL_RESOURCE_NAME,
670+
destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
671+
deploy_config=types.DeployRequest.DeployConfig(
672+
service_account=_TEST_SERVICE_ACCOUNT
673+
),
674+
)
675+
)
676+
656677

657678
@pytest.mark.usefixtures(
658679
"google_auth_mock",
@@ -985,6 +1006,24 @@ def test_deploy_with_system_labels_success(self, deploy_mock):
9851006
)
9861007
)
9871008

1009+
def test_deploy_with_service_account_success(self, deploy_mock):
1010+
"""Tests deploying a model with a service account."""
1011+
aiplatform.init(
1012+
project=_TEST_PROJECT,
1013+
location=_TEST_LOCATION,
1014+
)
1015+
model = model_garden.OpenModel(model_name=_TEST_MODEL_FULL_RESOURCE_NAME)
1016+
model.deploy(service_account=_TEST_SERVICE_ACCOUNT)
1017+
deploy_mock.assert_called_once_with(
1018+
types.DeployRequest(
1019+
publisher_model_name=_TEST_MODEL_FULL_RESOURCE_NAME,
1020+
destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
1021+
deploy_config=types.DeployRequest.DeployConfig(
1022+
service_account=_TEST_SERVICE_ACCOUNT
1023+
),
1024+
)
1025+
)
1026+
9881027
def test_deploy_with_fast_tryout_enabled_success(self, deploy_mock):
9891028
"""Tests deploying a model with fast tryout enabled."""
9901029
aiplatform.init(
@@ -1718,6 +1757,29 @@ def test_deploy_custom_model_with_system_labels_success(self, deploy_mock):
17181757
)
17191758
)
17201759

1760+
def test_deploy_custom_model_with_service_account_success(self, deploy_mock):
1761+
aiplatform.init(
1762+
project=_TEST_PROJECT,
1763+
location=_TEST_LOCATION,
1764+
)
1765+
model = model_garden_preview.CustomModel(gcs_uri=_TEST_GCS_URI)
1766+
model.deploy(service_account=_TEST_SERVICE_ACCOUNT)
1767+
deploy_mock.assert_called_once_with(
1768+
types.DeployRequest(
1769+
destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
1770+
custom_model=types.DeployRequest.CustomModel(
1771+
gcs_uri=_TEST_GCS_URI,
1772+
),
1773+
deploy_config=types.DeployRequest.DeployConfig(
1774+
dedicated_resources=types.DedicatedResources(
1775+
min_replica_count=1,
1776+
max_replica_count=1,
1777+
),
1778+
service_account=_TEST_SERVICE_ACCOUNT,
1779+
),
1780+
)
1781+
)
1782+
17211783
@pytest.mark.parametrize("filter_by_user_quota", [True, False])
17221784
def test_list_deploy_options_with_recommendations(self, filter_by_user_quota):
17231785
"""Tests list_deploy_options when recommend_spec returns recommendations."""

‎tests/unit/vertexai/model_garden/test_vertexai_model_garden.py‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070
_TEST_PARTNER_MODEL_SIMPLIFIED_RESOURCE_NAME = "ai21/jamba-large-1.6@001"
7171

7272
_TEST_GCS_URI = "gs://some-bucket/some-model"
73+
_TEST_SERVICE_ACCOUNT = "[email protected]"
7374
_TEST_ENDPOINT_NAME = "projects/test-project/locations/us-central1/endpoints/1234567890"
7475
_TEST_MODEL_NAME = "projects/test-project/locations/us-central1/models/9876543210"
7576
_TEST_IMAGE_URI = "us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00"
@@ -653,6 +654,26 @@ def test_deploy_with_display_names_success(self, deploy_mock):
653654
)
654655
)
655656

657+
def test_deploy_with_service_account_success(self, deploy_mock):
658+
"""Tests deploying a model with a service account."""
659+
aiplatform.init(
660+
project=_TEST_PROJECT,
661+
location=_TEST_LOCATION,
662+
)
663+
model = model_garden.PartnerModel(
664+
model_name=_TEST_PARTNER_MODEL_FULL_RESOURCE_NAME
665+
)
666+
model.deploy(service_account=_TEST_SERVICE_ACCOUNT)
667+
deploy_mock.assert_called_once_with(
668+
types.DeployRequest(
669+
publisher_model_name=_TEST_PARTNER_MODEL_FULL_RESOURCE_NAME,
670+
destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
671+
deploy_config=types.DeployRequest.DeployConfig(
672+
service_account=_TEST_SERVICE_ACCOUNT
673+
),
674+
)
675+
)
676+
656677

657678
@pytest.mark.usefixtures(
658679
"google_auth_mock",
@@ -985,6 +1006,24 @@ def test_deploy_with_system_labels_success(self, deploy_mock):
9851006
)
9861007
)
9871008

1009+
def test_deploy_with_service_account_success(self, deploy_mock):
1010+
"""Tests deploying a model with a service account."""
1011+
aiplatform.init(
1012+
project=_TEST_PROJECT,
1013+
location=_TEST_LOCATION,
1014+
)
1015+
model = model_garden.OpenModel(model_name=_TEST_MODEL_FULL_RESOURCE_NAME)
1016+
model.deploy(service_account=_TEST_SERVICE_ACCOUNT)
1017+
deploy_mock.assert_called_once_with(
1018+
types.DeployRequest(
1019+
publisher_model_name=_TEST_MODEL_FULL_RESOURCE_NAME,
1020+
destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
1021+
deploy_config=types.DeployRequest.DeployConfig(
1022+
service_account=_TEST_SERVICE_ACCOUNT
1023+
),
1024+
)
1025+
)
1026+
9881027
def test_deploy_with_fast_tryout_enabled_success(self, deploy_mock):
9891028
"""Tests deploying a model with fast tryout enabled."""
9901029
aiplatform.init(
@@ -1718,6 +1757,29 @@ def test_deploy_custom_model_with_system_labels_success(self, deploy_mock):
17181757
)
17191758
)
17201759

1760+
def test_deploy_custom_model_with_service_account_success(self, deploy_mock):
1761+
aiplatform.init(
1762+
project=_TEST_PROJECT,
1763+
location=_TEST_LOCATION,
1764+
)
1765+
model = model_garden_preview.CustomModel(gcs_uri=_TEST_GCS_URI)
1766+
model.deploy(service_account=_TEST_SERVICE_ACCOUNT)
1767+
deploy_mock.assert_called_once_with(
1768+
types.DeployRequest(
1769+
destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
1770+
custom_model=types.DeployRequest.CustomModel(
1771+
gcs_uri=_TEST_GCS_URI,
1772+
),
1773+
deploy_config=types.DeployRequest.DeployConfig(
1774+
dedicated_resources=types.DedicatedResources(
1775+
min_replica_count=1,
1776+
max_replica_count=1,
1777+
),
1778+
service_account=_TEST_SERVICE_ACCOUNT,
1779+
),
1780+
)
1781+
)
1782+
17211783
@pytest.mark.parametrize("filter_by_user_quota", [True, False])
17221784
def test_list_deploy_options_with_recommendations(self, filter_by_user_quota):
17231785
"""Tests list_deploy_options when recommend_spec returns recommendations."""

0 commit comments

Comments
 (0)