Skip to content

Commit 9a0eec6

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: [vertexai] Added concise option name to OpenModel.list_deploy_options()
PiperOrigin-RevId: 774997367
1 parent 7b51d9e commit 9a0eec6

File tree

2 files changed

+30
-22
lines changed

2 files changed

+30
-22
lines changed

tests/unit/vertexai/model_garden/test_model_garden.py

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,7 @@ def get_publisher_model_mock():
181181
multi_deploy_vertex=types.PublisherModel.CallToAction.DeployVertex(
182182
multi_deploy_vertex=[
183183
types.PublisherModel.CallToAction.Deploy(
184+
deploy_task_name="vLLM 32K context",
184185
container_spec=types.ModelContainerSpec(
185186
image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
186187
command=["python", "main.py"],
@@ -198,6 +199,7 @@ def get_publisher_model_mock():
198199
),
199200
),
200201
types.PublisherModel.CallToAction.Deploy(
202+
deploy_task_name="vLLM 128K context",
201203
container_spec=types.ModelContainerSpec(
202204
image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
203205
command=["python", "main.py"],
@@ -1032,17 +1034,17 @@ def test_list_deploy_options_concise(self, get_publisher_model_mock):
10321034
result = model.list_deploy_options(concise=True)
10331035
expected_result = textwrap.dedent(
10341036
"""\
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,
1037+
[Option 1: vLLM 32K context]
1038+
serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
1039+
machine_type="g2-standard-16",
1040+
accelerator_type="NVIDIA_L4",
1041+
accelerator_count=1,
10401042
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,"""
1043+
[Option 2: vLLM 128K context]
1044+
serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
1045+
machine_type="g2-standard-32",
1046+
accelerator_type="NVIDIA_L4",
1047+
accelerator_count=4,"""
10461048
)
10471049
assert result == expected_result
10481050
get_publisher_model_mock.assert_called_with(
@@ -1058,16 +1060,16 @@ def test_list_deploy_options_concise(self, get_publisher_model_mock):
10581060
expected_hf_result = textwrap.dedent(
10591061
"""\
10601062
[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,
1063+
serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
1064+
machine_type="g2-standard-16",
1065+
accelerator_type="NVIDIA_L4",
1066+
accelerator_count=1,
10651067
10661068
[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,"""
1069+
serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
1070+
machine_type="g2-standard-32",
1071+
accelerator_type="NVIDIA_L4",
1072+
accelerator_count=4,"""
10711073
)
10721074
assert hf_result == expected_hf_result
10731075
get_publisher_model_mock.assert_called_with(

vertexai/model_garden/_model_garden.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -661,7 +661,7 @@ def list_deploy_options(
661661
662662
Args:
663663
concise: If true, returns a human-readable string with container and
664-
machine specs.
664+
machine specs.
665665
666666
Returns:
667667
A list of deploy options or a concise formatted string.
@@ -694,8 +694,10 @@ def _extract_config(option):
694694
if option.dedicated_resources
695695
else None
696696
)
697+
option_name = getattr(option, "deploy_task_name", None)
697698

698699
return {
700+
"option_name": option_name,
699701
"serving_container_image_uri": container,
700702
"machine_type": getattr(machine, "machine_type", None),
701703
"accelerator_type": getattr(
@@ -706,11 +708,15 @@ def _extract_config(option):
706708

707709
concise_deploy_options = [_extract_config(opt) for opt in deploy_options]
708710
return "\n\n".join(
709-
f"[Option {i + 1}]\n"
711+
(
712+
f"[Option {i + 1}: {config['option_name']}]\n"
713+
if config.get("option_name")
714+
else f"[Option {i + 1}]\n"
715+
)
710716
+ "\n".join(
711-
f' {k}="{v}",' if k != "accelerator_count" else f" {k}={v},"
717+
f' {k}="{v}",' if k != "accelerator_count" else f" {k}={v},"
712718
for k, v in config.items()
713-
if v is not None
719+
if v is not None and k != "option_name"
714720
)
715721
for i, config in enumerate(concise_deploy_options)
716722
)

0 commit comments

Comments
 (0)