diff --git a/agentplatform/model_garden/_model_garden.py b/agentplatform/model_garden/_model_garden.py index 5b80f7e71d..f9f292d3dc 100644 --- a/agentplatform/model_garden/_model_garden.py +++ b/agentplatform/model_garden/_model_garden.py @@ -462,6 +462,7 @@ def deploy( serving_container_health_probe_timeout_seconds: Optional[int] = None, enable_private_service_connect: bool = False, psc_project_allow_list: Optional[Sequence[str]] = None, + service_account: Optional[str] = None, ) -> aiplatform.Endpoint: """Deploys an Open Model to an endpoint. @@ -601,6 +602,9 @@ def deploy( connect. psc_project_allow_list (Sequence[str]): The list of projects that are allowed to access the endpoint over private service connect. + service_account (str): Optional. The service account that the + deployed model runs as. The caller must have the + `iam.serviceAccounts.actAs` permission on this service account. Returns: endpoint (aiplatform.Endpoint): @@ -683,6 +687,9 @@ def deploy( if system_labels: request.deploy_config.system_labels = system_labels + if service_account: + request.deploy_config.service_account = service_account + if serving_container_spec: if not serving_container_spec.image_uri: raise ValueError( @@ -1126,6 +1133,7 @@ def deploy( enable_private_service_connect: bool = False, psc_project_allow_list: Optional[List[str]] = None, deploy_request_timeout: Optional[float] = None, + service_account: Optional[str] = None, ) -> aiplatform.Endpoint: """Deploys a Custom Model to an endpoint. @@ -1174,6 +1182,9 @@ def deploy( access the endpoint over private service connect. deploy_request_timeout: The timeout for the deploy request. Default is 2 hours. + service_account (str): Optional. The service account that the + deployed model runs as. The caller must have the + `iam.serviceAccounts.actAs` permission on this service account. Returns: endpoint (aiplatform.Endpoint): @@ -1194,6 +1205,7 @@ def deploy( enable_private_service_connect=enable_private_service_connect, psc_project_allow_list=psc_project_allow_list, deploy_request_timeout=deploy_request_timeout, + service_account=service_account, ) def _deploy_model_registry_model(self) -> aiplatform.Endpoint: @@ -1218,6 +1230,7 @@ def _deploy_gcs_uri( endpoint_display_name: Optional[str] = None, model_display_name: Optional[str] = None, deploy_request_timeout: Optional[float] = None, + service_account: Optional[str] = None, ) -> aiplatform.Endpoint: """Deploys a Custom Model to an endpoint. @@ -1268,6 +1281,9 @@ def _deploy_gcs_uri( access the endpoint over private service connect. deploy_request_timeout: The timeout for the deploy request. Default is 2 hours. + service_account (str): Optional. The service account that the + deployed model runs as. The caller must have the + `iam.serviceAccounts.actAs` permission on this service account. Returns: endpoint (aiplatform.Endpoint): @@ -1296,6 +1312,8 @@ def has_all_or_none_values(var1, var2, var3) -> bool: request.model_config.model_display_name = model_display_name if system_labels: request.deploy_config.system_labels = system_labels + if service_account: + request.deploy_config.service_account = service_account if enable_private_service_connect and psc_project_allow_list: request.endpoint_config.private_service_connect_config = ( @@ -1407,6 +1425,7 @@ def deploy( endpoint_display_name: Optional[str] = None, model_display_name: Optional[str] = None, deploy_request_timeout: Optional[float] = None, + service_account: Optional[str] = None, ) -> aiplatform.Endpoint: """Deploys an Open Model to an endpoint. @@ -1439,6 +1458,9 @@ def deploy( model_display_name: The display name of the uploaded model. deploy_request_timeout: The timeout for the deploy request. Default is 2 hours. + service_account (str): Optional. The service account that the + deployed model runs as. The caller must have the + `iam.serviceAccounts.actAs` permission on this service account. Returns: endpoint (aiplatform.Endpoint): @@ -1469,6 +1491,9 @@ def deploy( ) request.deploy_config.dedicated_resources = dedicated_resources + if service_account: + request.deploy_config.service_account = service_account + _LOGGER.info(f"Deploying model: {self._model_name}") operation_future = self._model_garden_client.deploy(request) diff --git a/google/cloud/aiplatform_v1beta1/types/model_garden_service.py b/google/cloud/aiplatform_v1beta1/types/model_garden_service.py index fc9a97cea3..091d256429 100644 --- a/google/cloud/aiplatform_v1beta1/types/model_garden_service.py +++ b/google/cloud/aiplatform_v1beta1/types/model_garden_service.py @@ -457,6 +457,17 @@ class DeployConfig(proto.Message): Optional. System labels for Model Garden deployments. These labels are managed by Google and for tracking purposes only. + service_account (str): + Optional. The service account that the + container of the deployed model runs as. Specify + the email address of the service account. If not + specified, the container runs as a service + account that has no access to the resource + project. + + Users deploying the model must have the + ``iam.serviceAccounts.actAs`` permission on this + service account. """ dedicated_resources: machine_resources.DedicatedResources = proto.Field( @@ -473,6 +484,10 @@ class DeployConfig(proto.Message): proto.STRING, number=3, ) + service_account: str = proto.Field( + proto.STRING, + number=5, + ) publisher_model_name: str = proto.Field( proto.STRING, diff --git a/tests/unit/agentplatform/model_garden/test_model_garden.py b/tests/unit/agentplatform/model_garden/test_model_garden.py index 23c18d095f..ab913bc2f1 100644 --- a/tests/unit/agentplatform/model_garden/test_model_garden.py +++ b/tests/unit/agentplatform/model_garden/test_model_garden.py @@ -70,6 +70,7 @@ _TEST_PARTNER_MODEL_SIMPLIFIED_RESOURCE_NAME = "ai21/jamba-large-1.6@001" _TEST_GCS_URI = "gs://some-bucket/some-model" +_TEST_SERVICE_ACCOUNT = "deployed-model@test-project.iam.gserviceaccount.com" _TEST_ENDPOINT_NAME = "projects/test-project/locations/us-central1/endpoints/1234567890" _TEST_MODEL_NAME = "projects/test-project/locations/us-central1/models/9876543210" _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): ) ) + def test_deploy_with_service_account_success(self, deploy_mock): + """Tests deploying a model with a service account.""" + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + ) + model = model_garden.PartnerModel( + model_name=_TEST_PARTNER_MODEL_FULL_RESOURCE_NAME + ) + model.deploy(service_account=_TEST_SERVICE_ACCOUNT) + deploy_mock.assert_called_once_with( + types.DeployRequest( + publisher_model_name=_TEST_PARTNER_MODEL_FULL_RESOURCE_NAME, + destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}", + deploy_config=types.DeployRequest.DeployConfig( + service_account=_TEST_SERVICE_ACCOUNT + ), + ) + ) + @pytest.mark.usefixtures( "google_auth_mock", @@ -985,6 +1006,24 @@ def test_deploy_with_system_labels_success(self, deploy_mock): ) ) + def test_deploy_with_service_account_success(self, deploy_mock): + """Tests deploying a model with a service account.""" + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + ) + model = model_garden.OpenModel(model_name=_TEST_MODEL_FULL_RESOURCE_NAME) + model.deploy(service_account=_TEST_SERVICE_ACCOUNT) + deploy_mock.assert_called_once_with( + types.DeployRequest( + publisher_model_name=_TEST_MODEL_FULL_RESOURCE_NAME, + destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}", + deploy_config=types.DeployRequest.DeployConfig( + service_account=_TEST_SERVICE_ACCOUNT + ), + ) + ) + def test_deploy_with_fast_tryout_enabled_success(self, deploy_mock): """Tests deploying a model with fast tryout enabled.""" aiplatform.init( @@ -1718,6 +1757,29 @@ def test_deploy_custom_model_with_system_labels_success(self, deploy_mock): ) ) + def test_deploy_custom_model_with_service_account_success(self, deploy_mock): + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + ) + model = model_garden_preview.CustomModel(gcs_uri=_TEST_GCS_URI) + model.deploy(service_account=_TEST_SERVICE_ACCOUNT) + deploy_mock.assert_called_once_with( + types.DeployRequest( + destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}", + custom_model=types.DeployRequest.CustomModel( + gcs_uri=_TEST_GCS_URI, + ), + deploy_config=types.DeployRequest.DeployConfig( + dedicated_resources=types.DedicatedResources( + min_replica_count=1, + max_replica_count=1, + ), + service_account=_TEST_SERVICE_ACCOUNT, + ), + ) + ) + @pytest.mark.parametrize("filter_by_user_quota", [True, False]) def test_list_deploy_options_with_recommendations(self, filter_by_user_quota): """Tests list_deploy_options when recommend_spec returns recommendations.""" diff --git a/tests/unit/vertexai/model_garden/test_vertexai_model_garden.py b/tests/unit/vertexai/model_garden/test_vertexai_model_garden.py index 7150043f08..b474ed103c 100644 --- a/tests/unit/vertexai/model_garden/test_vertexai_model_garden.py +++ b/tests/unit/vertexai/model_garden/test_vertexai_model_garden.py @@ -70,6 +70,7 @@ _TEST_PARTNER_MODEL_SIMPLIFIED_RESOURCE_NAME = "ai21/jamba-large-1.6@001" _TEST_GCS_URI = "gs://some-bucket/some-model" +_TEST_SERVICE_ACCOUNT = "deployed-model@test-project.iam.gserviceaccount.com" _TEST_ENDPOINT_NAME = "projects/test-project/locations/us-central1/endpoints/1234567890" _TEST_MODEL_NAME = "projects/test-project/locations/us-central1/models/9876543210" _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): ) ) + def test_deploy_with_service_account_success(self, deploy_mock): + """Tests deploying a model with a service account.""" + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + ) + model = model_garden.PartnerModel( + model_name=_TEST_PARTNER_MODEL_FULL_RESOURCE_NAME + ) + model.deploy(service_account=_TEST_SERVICE_ACCOUNT) + deploy_mock.assert_called_once_with( + types.DeployRequest( + publisher_model_name=_TEST_PARTNER_MODEL_FULL_RESOURCE_NAME, + destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}", + deploy_config=types.DeployRequest.DeployConfig( + service_account=_TEST_SERVICE_ACCOUNT + ), + ) + ) + @pytest.mark.usefixtures( "google_auth_mock", @@ -985,6 +1006,24 @@ def test_deploy_with_system_labels_success(self, deploy_mock): ) ) + def test_deploy_with_service_account_success(self, deploy_mock): + """Tests deploying a model with a service account.""" + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + ) + model = model_garden.OpenModel(model_name=_TEST_MODEL_FULL_RESOURCE_NAME) + model.deploy(service_account=_TEST_SERVICE_ACCOUNT) + deploy_mock.assert_called_once_with( + types.DeployRequest( + publisher_model_name=_TEST_MODEL_FULL_RESOURCE_NAME, + destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}", + deploy_config=types.DeployRequest.DeployConfig( + service_account=_TEST_SERVICE_ACCOUNT + ), + ) + ) + def test_deploy_with_fast_tryout_enabled_success(self, deploy_mock): """Tests deploying a model with fast tryout enabled.""" aiplatform.init( @@ -1718,6 +1757,29 @@ def test_deploy_custom_model_with_system_labels_success(self, deploy_mock): ) ) + def test_deploy_custom_model_with_service_account_success(self, deploy_mock): + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + ) + model = model_garden_preview.CustomModel(gcs_uri=_TEST_GCS_URI) + model.deploy(service_account=_TEST_SERVICE_ACCOUNT) + deploy_mock.assert_called_once_with( + types.DeployRequest( + destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}", + custom_model=types.DeployRequest.CustomModel( + gcs_uri=_TEST_GCS_URI, + ), + deploy_config=types.DeployRequest.DeployConfig( + dedicated_resources=types.DedicatedResources( + min_replica_count=1, + max_replica_count=1, + ), + service_account=_TEST_SERVICE_ACCOUNT, + ), + ) + ) + @pytest.mark.parametrize("filter_by_user_quota", [True, False]) def test_list_deploy_options_with_recommendations(self, filter_by_user_quota): """Tests list_deploy_options when recommend_spec returns recommendations.""" diff --git a/vertexai/model_garden/_model_garden.py b/vertexai/model_garden/_model_garden.py index 98e828bf40..474781133a 100644 --- a/vertexai/model_garden/_model_garden.py +++ b/vertexai/model_garden/_model_garden.py @@ -463,6 +463,7 @@ def deploy( serving_container_health_probe_timeout_seconds: Optional[int] = None, enable_private_service_connect: bool = False, psc_project_allow_list: Optional[Sequence[str]] = None, + service_account: Optional[str] = None, ) -> aiplatform.Endpoint: """Deploys an Open Model to an endpoint. @@ -601,6 +602,9 @@ def deploy( connect. psc_project_allow_list (Sequence[str]): The list of projects that are allowed to access the endpoint over private service connect. + service_account (str): Optional. The service account that the + deployed model runs as. The caller must have the + `iam.serviceAccounts.actAs` permission on this service account. Returns: endpoint (aiplatform.Endpoint): @@ -683,6 +687,9 @@ def deploy( if system_labels: request.deploy_config.system_labels = system_labels + if service_account: + request.deploy_config.service_account = service_account + if serving_container_spec: if not serving_container_spec.image_uri: raise ValueError( @@ -1126,6 +1133,7 @@ def deploy( enable_private_service_connect: bool = False, psc_project_allow_list: Optional[List[str]] = None, deploy_request_timeout: Optional[float] = None, + service_account: Optional[str] = None, ) -> aiplatform.Endpoint: """Deploys a Custom Model to an endpoint. @@ -1174,6 +1182,9 @@ def deploy( access the endpoint over private service connect. deploy_request_timeout: The timeout for the deploy request. Default is 2 hours. + service_account (str): Optional. The service account that the + deployed model runs as. The caller must have the + `iam.serviceAccounts.actAs` permission on this service account. Returns: endpoint (aiplatform.Endpoint): @@ -1194,6 +1205,7 @@ def deploy( enable_private_service_connect=enable_private_service_connect, psc_project_allow_list=psc_project_allow_list, deploy_request_timeout=deploy_request_timeout, + service_account=service_account, ) def _deploy_model_registry_model(self) -> aiplatform.Endpoint: @@ -1218,6 +1230,7 @@ def _deploy_gcs_uri( endpoint_display_name: Optional[str] = None, model_display_name: Optional[str] = None, deploy_request_timeout: Optional[float] = None, + service_account: Optional[str] = None, ) -> aiplatform.Endpoint: """Deploys a Custom Model to an endpoint. @@ -1268,6 +1281,9 @@ def _deploy_gcs_uri( access the endpoint over private service connect. deploy_request_timeout: The timeout for the deploy request. Default is 2 hours. + service_account (str): Optional. The service account that the + deployed model runs as. The caller must have the + `iam.serviceAccounts.actAs` permission on this service account. Returns: endpoint (aiplatform.Endpoint): @@ -1296,6 +1312,8 @@ def has_all_or_none_values(var1, var2, var3) -> bool: request.model_config.model_display_name = model_display_name if system_labels: request.deploy_config.system_labels = system_labels + if service_account: + request.deploy_config.service_account = service_account if enable_private_service_connect and psc_project_allow_list: request.endpoint_config.private_service_connect_config = ( @@ -1407,6 +1425,7 @@ def deploy( endpoint_display_name: Optional[str] = None, model_display_name: Optional[str] = None, deploy_request_timeout: Optional[float] = None, + service_account: Optional[str] = None, ) -> aiplatform.Endpoint: """Deploys an Open Model to an endpoint. @@ -1439,6 +1458,9 @@ def deploy( model_display_name: The display name of the uploaded model. deploy_request_timeout: The timeout for the deploy request. Default is 2 hours. + service_account (str): Optional. The service account that the + deployed model runs as. The caller must have the + `iam.serviceAccounts.actAs` permission on this service account. Returns: endpoint (aiplatform.Endpoint): @@ -1469,6 +1491,9 @@ def deploy( ) request.deploy_config.dedicated_resources = dedicated_resources + if service_account: + request.deploy_config.service_account = service_account + _LOGGER.info(f"Deploying model: {self._model_name}") operation_future = self._model_garden_client.deploy(request)