feat: [vertexai] Added concise option name to OpenModel.list_deploy_options()

vertex-sdk-bot · copybara-github · commit 9a0eec645446 · 2025-06-23T18:04:20.000-07:00
PiperOrigin-RevId: 774997367
diff --git a/tests/unit/vertexai/model_garden/test_model_garden.py b/tests/unit/vertexai/model_garden/test_model_garden.py
@@ -181,6 +181,7 @@ def get_publisher_model_mock():
                     multi_deploy_vertex=types.PublisherModel.CallToAction.DeployVertex(
                         multi_deploy_vertex=[
                             types.PublisherModel.CallToAction.Deploy(
+                                deploy_task_name="vLLM 32K context",
                                 container_spec=types.ModelContainerSpec(
                                     image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
                                     command=["python", "main.py"],
@@ -198,6 +199,7 @@ def get_publisher_model_mock():
                                 ),
                             ),
                             types.PublisherModel.CallToAction.Deploy(
+                                deploy_task_name="vLLM 128K context",
                                 container_spec=types.ModelContainerSpec(
                                     image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
                                     command=["python", "main.py"],
@@ -1032,17 +1034,17 @@ def test_list_deploy_options_concise(self, get_publisher_model_mock):
         result = model.list_deploy_options(concise=True)
         expected_result = textwrap.dedent(
             """\
-        [Option 1]
-          serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
-          machine_type="g2-standard-16",
-          accelerator_type="NVIDIA_L4",
-          accelerator_count=1,
+        [Option 1: vLLM 32K context]
+            serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
+            machine_type="g2-standard-16",
+            accelerator_type="NVIDIA_L4",
+            accelerator_count=1,
 
-        [Option 2]
-          serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
-          machine_type="g2-standard-32",
-          accelerator_type="NVIDIA_L4",
-          accelerator_count=4,"""
+        [Option 2: vLLM 128K context]
+            serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
+            machine_type="g2-standard-32",
+            accelerator_type="NVIDIA_L4",
+            accelerator_count=4,"""
         )
         assert result == expected_result
         get_publisher_model_mock.assert_called_with(
@@ -1058,16 +1060,16 @@ def test_list_deploy_options_concise(self, get_publisher_model_mock):
         expected_hf_result = textwrap.dedent(
             """\
         [Option 1]
-          serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
-          machine_type="g2-standard-16",
-          accelerator_type="NVIDIA_L4",
-          accelerator_count=1,
+            serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00",
+            machine_type="g2-standard-16",
+            accelerator_type="NVIDIA_L4",
+            accelerator_count=1,
 
         [Option 2]
-          serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
-          machine_type="g2-standard-32",
-          accelerator_type="NVIDIA_L4",
-          accelerator_count=4,"""
+            serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/text-generation-inference-cu121.2-1.py310:latest",
+            machine_type="g2-standard-32",
+            accelerator_type="NVIDIA_L4",
+            accelerator_count=4,"""
         )
         assert hf_result == expected_hf_result
         get_publisher_model_mock.assert_called_with(
diff --git a/vertexai/model_garden/_model_garden.py b/vertexai/model_garden/_model_garden.py
@@ -661,7 +661,7 @@ def list_deploy_options(
 
         Args:
             concise: If true, returns a human-readable string with container and
-              machine specs.
+            machine specs.
 
         Returns:
             A list of deploy options or a concise formatted string.
@@ -694,8 +694,10 @@ def _extract_config(option):
                 if option.dedicated_resources
                 else None
             )
+            option_name = getattr(option, "deploy_task_name", None)
 
             return {
+                "option_name": option_name,
                 "serving_container_image_uri": container,
                 "machine_type": getattr(machine, "machine_type", None),
                 "accelerator_type": getattr(
@@ -706,11 +708,15 @@ def _extract_config(option):
 
         concise_deploy_options = [_extract_config(opt) for opt in deploy_options]
         return "\n\n".join(
-            f"[Option {i + 1}]\n"
+            (
+                f"[Option {i + 1}: {config['option_name']}]\n"
+                if config.get("option_name")
+                else f"[Option {i + 1}]\n"
+            )
             + "\n".join(
-                f'  {k}="{v}",' if k != "accelerator_count" else f"  {k}={v},"
+                f'    {k}="{v}",' if k != "accelerator_count" else f"    {k}={v},"
                 for k, v in config.items()
-                if v is not None
+                if v is not None and k != "option_name"
             )
             for i, config in enumerate(concise_deploy_options)
         )
