Skip to content

Commit acc301a

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: [vertexai] Added system labels to OpenModel.deploy()
feat: [vertexai] Added a concise option to `OpenModel.list_deploy_options()` PiperOrigin-RevId: 755048831
1 parent bf79bdf commit acc301a

2 files changed

Lines changed: 322 additions & 143 deletions

File tree

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

Lines changed: 142 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
"""Unit tests for ModelGarden class."""
1515

1616
import importlib
17+
import textwrap
1718
from unittest import mock
1819

1920
from google import auth
@@ -180,14 +181,39 @@ def get_publisher_model_mock():
180181
multi_deploy_vertex=types.PublisherModel.CallToAction.DeployVertex(
181182
multi_deploy_vertex=[
182183
types.PublisherModel.CallToAction.Deploy(
184+
container_spec=types.ModelContainerSpec(
185+
image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
186+
command=["python", "main.py"],
187+
args=["--model-id=gemma-2b"],
188+
env=[
189+
types.EnvVar(name="MODEL_ID", value="gemma-2b")
190+
],
191+
),
183192
dedicated_resources=types.DedicatedResources(
184193
machine_spec=types.MachineSpec(
185194
machine_type="g2-standard-16",
186195
accelerator_type="NVIDIA_L4",
187196
accelerator_count=1,
188197
)
189-
)
190-
)
198+
),
199+
),
200+
types.PublisherModel.CallToAction.Deploy(
201+
container_spec=types.ModelContainerSpec(
202+
image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
203+
command=["python", "main.py"],
204+
args=["--model-id=gemma-2b"],
205+
env=[
206+
types.EnvVar(name="MODEL_ID", value="gemma-2b")
207+
],
208+
),
209+
dedicated_resources=types.DedicatedResources(
210+
machine_spec=types.MachineSpec(
211+
machine_type="g2-standard-32",
212+
accelerator_type="NVIDIA_L4",
213+
accelerator_count=4,
214+
)
215+
),
216+
),
191217
]
192218
)
193219
),
@@ -198,14 +224,39 @@ def get_publisher_model_mock():
198224
multi_deploy_vertex=types.PublisherModel.CallToAction.DeployVertex(
199225
multi_deploy_vertex=[
200226
types.PublisherModel.CallToAction.Deploy(
227+
container_spec=types.ModelContainerSpec(
228+
image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
229+
command=["python", "main.py"],
230+
args=["--model-id=gemma-2b"],
231+
env=[
232+
types.EnvVar(name="MODEL_ID", value="gemma-2b")
233+
],
234+
),
201235
dedicated_resources=types.DedicatedResources(
202236
machine_spec=types.MachineSpec(
203237
machine_type="g2-standard-16",
204238
accelerator_type="NVIDIA_L4",
205239
accelerator_count=1,
206240
)
207-
)
208-
)
241+
),
242+
),
243+
types.PublisherModel.CallToAction.Deploy(
244+
container_spec=types.ModelContainerSpec(
245+
image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
246+
command=["python", "main.py"],
247+
args=["--model-id=gemma-2b"],
248+
env=[
249+
types.EnvVar(name="MODEL_ID", value="gemma-2b")
250+
],
251+
),
252+
dedicated_resources=types.DedicatedResources(
253+
machine_spec=types.MachineSpec(
254+
machine_type="g2-standard-32",
255+
accelerator_type="NVIDIA_L4",
256+
accelerator_count=4,
257+
)
258+
),
259+
),
209260
]
210261
)
211262
),
@@ -444,8 +495,8 @@ def accept_model_license_agreement_mock():
444495
"check_license_agreement_status_mock",
445496
"accept_model_license_agreement_mock",
446497
)
447-
class TestModelGarden:
448-
"""Test cases for ModelGarden class."""
498+
class TestModelGardenOpenModel:
499+
"""Test cases for Model Garden OpenModel class."""
449500

450501
def setup_method(self):
451502
importlib.reload(aiplatform.initializer)
@@ -746,6 +797,24 @@ def test_deploy_with_dedicated_endpoint_success(self, deploy_mock):
746797
)
747798
)
748799

800+
def test_deploy_with_system_labels_success(self, deploy_mock):
801+
"""Tests deploying a model with system labels."""
802+
aiplatform.init(
803+
project=_TEST_PROJECT,
804+
location=_TEST_LOCATION,
805+
)
806+
model = model_garden.OpenModel(model_name=_TEST_MODEL_FULL_RESOURCE_NAME)
807+
model.deploy(system_labels={"test-key": "test-value"})
808+
deploy_mock.assert_called_once_with(
809+
types.DeployRequest(
810+
publisher_model_name=_TEST_MODEL_FULL_RESOURCE_NAME,
811+
destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
812+
deploy_config=types.DeployRequest.DeployConfig(
813+
system_labels={"test-key": "test-value"}
814+
),
815+
)
816+
)
817+
749818
def test_deploy_with_fast_tryout_enabled_success(self, deploy_mock):
750819
"""Tests deploying a model with fast tryout enabled."""
751820
aiplatform.init(
@@ -917,9 +986,8 @@ def test_list_deploy_options(self, get_publisher_model_mock):
917986
)
918987

919988
expected_message = (
920-
"Model does not support deployment, please use a deploy-able model"
921-
" instead. You can use the list_deployable_models() method to find out"
922-
" which ones currently support deployment."
989+
"Model does not support deployment. "
990+
"Use `list_deployable_models()` to find supported models."
923991
)
924992
with pytest.raises(ValueError) as exception:
925993
model = model_garden.OpenModel(model_name=_TEST_MODEL_FULL_RESOURCE_NAME)
@@ -945,6 +1013,71 @@ def test_list_deploy_options(self, get_publisher_model_mock):
9451013
)
9461014
)
9471015

1016+
def test_list_deploy_options_concise(self, get_publisher_model_mock):
1017+
"""Tests getting the supported deploy options for a model."""
1018+
aiplatform.init(
1019+
project=_TEST_PROJECT,
1020+
location=_TEST_LOCATION,
1021+
)
1022+
1023+
expected_message = (
1024+
"Model does not support deployment. "
1025+
"Use `list_deployable_models()` to find supported models."
1026+
)
1027+
with pytest.raises(ValueError) as exception:
1028+
model = model_garden.OpenModel(model_name=_TEST_MODEL_FULL_RESOURCE_NAME)
1029+
_ = model.list_deploy_options(concise=True)
1030+
assert str(exception.value) == expected_message
1031+
1032+
result = model.list_deploy_options(concise=True)
1033+
expected_result = textwrap.dedent(
1034+
"""\
1035+
[Option 1]
1036+
serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
1037+
machine_type="g2-standard-16",
1038+
accelerator_type="NVIDIA_L4",
1039+
accelerator_count=1,
1040+
1041+
[Option 2]
1042+
serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
1043+
machine_type="g2-standard-32",
1044+
accelerator_type="NVIDIA_L4",
1045+
accelerator_count=4,"""
1046+
)
1047+
assert result == expected_result
1048+
get_publisher_model_mock.assert_called_with(
1049+
types.GetPublisherModelRequest(
1050+
name=_TEST_MODEL_FULL_RESOURCE_NAME,
1051+
is_hugging_face_model=False,
1052+
include_equivalent_model_garden_model_deployment_configs=True,
1053+
)
1054+
)
1055+
1056+
hf_model = model_garden.OpenModel(_TEST_MODEL_HUGGING_FACE_ID)
1057+
hf_result = hf_model.list_deploy_options(concise=True)
1058+
expected_hf_result = textwrap.dedent(
1059+
"""\
1060+
[Option 1]
1061+
serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
1062+
machine_type="g2-standard-16",
1063+
accelerator_type="NVIDIA_L4",
1064+
accelerator_count=1,
1065+
1066+
[Option 2]
1067+
serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
1068+
machine_type="g2-standard-32",
1069+
accelerator_type="NVIDIA_L4",
1070+
accelerator_count=4,"""
1071+
)
1072+
assert hf_result == expected_hf_result
1073+
get_publisher_model_mock.assert_called_with(
1074+
types.GetPublisherModelRequest(
1075+
name=_TEST_HUGGING_FACE_MODEL_FULL_RESOURCE_NAME,
1076+
is_hugging_face_model=True,
1077+
include_equivalent_model_garden_model_deployment_configs=True,
1078+
)
1079+
)
1080+
9481081
def test_list_deployable_models(self, list_publisher_models_mock):
9491082
"""Tests getting the supported deploy options for a model."""
9501083
aiplatform.init(

0 commit comments

Comments
 (0)