feat: Add multihost_gpu_node_count to Vertex SDK.

PiperOrigin-RevId: 739019987

Author: vertex-sdk-bot

google/cloud/aiplatform/preview/models.py

@@ -139,6 +139,7 @@ def create(
         sync=True,
         create_request_timeout: Optional[float] = None,
         required_replica_count: Optional[int] = 0,
+        multihost_gpu_node_count: Optional[int] = None,
     ) -> "DeploymentResourcePool":
         """Creates a new DeploymentResourcePool.
 
@@ -205,6 +206,9 @@ def create(
                 set, the model deploy/mutate operation will succeed once
                 available_replica_count reaches required_replica_count, and the
                 rest of the replicas will be retried.
+            multihost_gpu_node_count (int):
+                Optional. The number of nodes per replica for multihost GPU
+                deployments. Required for multihost GPU deployments.
 
         Returns:
             DeploymentResourcePool
@@ -232,6 +236,7 @@ def create(
             sync=sync,
             create_request_timeout=create_request_timeout,
             required_replica_count=required_replica_count,
+            multihost_gpu_node_count=multihost_gpu_node_count,
         )
 
     @classmethod
@@ -254,6 +259,7 @@ def _create(
         sync=True,
         create_request_timeout: Optional[float] = None,
         required_replica_count: Optional[int] = 0,
+        multihost_gpu_node_count: Optional[int] = None,
     ) -> "DeploymentResourcePool":
         """Creates a new DeploymentResourcePool.
 
@@ -323,6 +329,9 @@ def _create(
                 set, the model deploy/mutate operation will succeed once
                 available_replica_count reaches required_replica_count, and the
                 rest of the replicas will be retried.
+            multihost_gpu_node_count (int):
+                Optional. The number of nodes per replica for multihost GPU
+                deployments. Required for multihost GPU deployments.
 
         Returns:
             DeploymentResourcePool
@@ -339,7 +348,8 @@ def _create(
         )
 
         machine_spec = gca_machine_resources_compat.MachineSpec(
-            machine_type=machine_type
+            machine_type=machine_type,
+            multihost_gpu_node_count=multihost_gpu_node_count,
         )
 
         if autoscaling_target_cpu_utilization:
@@ -369,6 +379,9 @@ def _create(
                     [autoscaling_metric_spec]
                 )
 
+        if multihost_gpu_node_count:
+            machine_spec.multihost_gpu_node_count = multihost_gpu_node_count
+
         dedicated_resources.machine_spec = machine_spec
 
         gapic_drp = gca_deployment_resource_pool_compat.DeploymentResourcePool(
@@ -691,6 +704,7 @@ def deploy(
         system_labels: Optional[Dict[str, str]] = None,
         required_replica_count: Optional[int] = 0,
         rollout_options: Optional[RolloutOptions] = None,
+        multihost_gpu_node_count: Optional[int] = None,
     ) -> None:
         """Deploys a Model to the Endpoint.
 
@@ -789,6 +803,9 @@ def deploy(
                 rest of the replicas will be retried.
             rollout_options (RolloutOptions):
                 Optional. Options to configure a rolling deployment.
+            multihost_gpu_node_count (int): Optional. The number of nodes per
+              replica for multihost GPU deployments. Required for multihost GPU
+              deployments.
 
         """
         self._sync_gca_resource_if_skipped()
@@ -832,6 +849,7 @@ def deploy(
             system_labels=system_labels,
             required_replica_count=required_replica_count,
             rollout_options=rollout_options,
+            multihost_gpu_node_count=multihost_gpu_node_count,
         )
 
     @base.optional_sync()
@@ -859,6 +877,7 @@ def _deploy(
         system_labels: Optional[Dict[str, str]] = None,
         required_replica_count: Optional[int] = 0,
         rollout_options: Optional[RolloutOptions] = None,
+        multihost_gpu_node_count: Optional[int] = None,
     ) -> None:
         """Deploys a Model to the Endpoint.
 
@@ -951,6 +970,10 @@ def _deploy(
               rest of the replicas will be retried.
             rollout_options (RolloutOptions): Optional.
               Options to configure a rolling deployment.
+            multihost_gpu_node_count (int): Optional. The number of nodes per
+              replica for multihost GPU deployments. Required for multihost
+              GPU deployments.
+
         """
         _LOGGER.log_action_start_against_resource(
             f"Deploying Model {model.resource_name} to", "", self
@@ -982,6 +1005,7 @@ def _deploy(
             system_labels=system_labels,
             required_replica_count=required_replica_count,
             rollout_options=rollout_options,
+            multihost_gpu_node_count=multihost_gpu_node_count,
         )
 
         _LOGGER.log_action_completed_against_resource("model", "deployed", self)
@@ -1016,6 +1040,7 @@ def _deploy_call(
         system_labels: Optional[Dict[str, str]] = None,
         required_replica_count: Optional[int] = 0,
         rollout_options: Optional[RolloutOptions] = None,
+        multihost_gpu_node_count: Optional[int] = None,
     ) -> None:
         """Helper method to deploy model to endpoint.
 
@@ -1115,6 +1140,9 @@ def _deploy_call(
               rest of the replicas will be retried.
             rollout_options (RolloutOptions): Optional. Options to configure a
               rolling deployment.
+            multihost_gpu_node_count (int):
+              Optional. The number of nodes per replica for multihost GPU
+              deployments. Required for multihost GPU deployments.
 
         Raises:
             ValueError: If only `accelerator_type` or `accelerator_count` is
@@ -1195,7 +1223,8 @@ def _deploy_call(
                 )
 
                 machine_spec = gca_machine_resources_compat.MachineSpec(
-                    machine_type=machine_type
+                    machine_type=machine_type,
+                    multihost_gpu_node_count=multihost_gpu_node_count,
                 )
 
                 if autoscaling_target_cpu_utilization:
@@ -1538,6 +1567,7 @@ def deploy(
         system_labels: Optional[Dict[str, str]] = None,
         required_replica_count: Optional[int] = 0,
         rollout_options: Optional[RolloutOptions] = None,
+        multihost_gpu_node_count: Optional[int] = None,
     ) -> Union[Endpoint, models.PrivateEndpoint]:
         """Deploys model to endpoint.
 
@@ -1657,6 +1687,9 @@ def deploy(
                 rest of the replicas will be retried.
             rollout_options (RolloutOptions):
               Optional. Options to configure a rolling deployment.
+            multihost_gpu_node_count (int):
+                Optional. The number of nodes per replica for multihost GPU
+                deployments. Required for multihost GPU deployments.
 
         Returns:
             endpoint (Union[Endpoint, models.PrivateEndpoint]):
@@ -1717,6 +1750,7 @@ def deploy(
             system_labels=system_labels,
             required_replica_count=required_replica_count,
             rollout_options=rollout_options,
+            multihost_gpu_node_count=multihost_gpu_node_count,
         )
 
     def _should_enable_dedicated_endpoint(self, fast_tryout_enabled: bool) -> bool:
@@ -1753,6 +1787,7 @@ def _deploy(
         system_labels: Optional[Dict[str, str]] = None,
         required_replica_count: Optional[int] = 0,
         rollout_options: Optional[RolloutOptions] = None,
+        multihost_gpu_node_count: Optional[int] = None,
     ) -> Union[Endpoint, models.PrivateEndpoint]:
         """Deploys model to endpoint.
 
@@ -1863,6 +1898,9 @@ def _deploy(
               rest of the replicas will be retried.
             rollout_options (RolloutOptions):
               Optional. Options to configure a rolling deployment.
+            multihost_gpu_node_count (int):
+              Optional. The number of nodes per replica for multihost GPU
+              deployments. Required for multihost GPU deployments.
 
         Returns:
             endpoint (Union[Endpoint, models.PrivateEndpoint]):
@@ -1928,6 +1966,7 @@ def _deploy(
             fast_tryout_enabled=fast_tryout_enabled,
             system_labels=system_labels,
             required_replica_count=required_replica_count,
+            multihost_gpu_node_count=multihost_gpu_node_count,
             **preview_kwargs,
         )
 

tests/unit/aiplatform/test_models.py

@@ -146,6 +146,11 @@
 _TEST_TPU_MACHINE_TYPE = "ct5lp-hightpu-4t"
 _TEST_TPU_TOPOLOGY = "2x2"
 
+_TEST_GPU_MACHINE_TYPE = "a3-highgpu-8g"
+_TEST_GPU_ACCELERATOR_TYPE = "NVIDIA_TESLA_A100"
+_TEST_GPU_ACCELERATOR_COUNT = 8
+_TEST_MULTIHOST_GPU_NODE_COUNT = 2
+
 _TEST_BATCH_SIZE = 16
 
 _TEST_PIPELINE_RESOURCE_NAME = (
@@ -2239,6 +2244,61 @@ def test_deploy_no_endpoint_with_tpu_topology(self, deploy_model_mock, sync):
             timeout=None,
         )
 
+    @pytest.mark.usefixtures(
+        "get_endpoint_mock",
+        "get_model_mock",
+        "create_endpoint_mock",
+        "preview_deploy_model_mock",
+    )
+    @pytest.mark.parametrize("sync", [True, False])
+    def test_deploy_no_endpoint_with_multihost_gpu_node_count(
+        self, preview_deploy_model_mock, sync
+    ):
+        test_model = preview_models.Model(_TEST_ID)
+        test_model._gca_resource.supported_deployment_resources_types.append(
+            aiplatform.gapic.Model.DeploymentResourcesType.DEDICATED_RESOURCES
+        )
+
+        test_endpoint = test_model.deploy(
+            machine_type=_TEST_GPU_MACHINE_TYPE,
+            accelerator_type=_TEST_GPU_ACCELERATOR_TYPE,
+            accelerator_count=_TEST_GPU_ACCELERATOR_COUNT,
+            multihost_gpu_node_count=_TEST_MULTIHOST_GPU_NODE_COUNT,
+            sync=sync,
+            deploy_request_timeout=None,
+        )
+
+        if not sync:
+            test_endpoint.wait()
+
+        expected_machine_spec = gca_machine_resources_v1beta1.MachineSpec(
+            machine_type=_TEST_GPU_MACHINE_TYPE,
+            accelerator_type=_TEST_GPU_ACCELERATOR_TYPE,
+            accelerator_count=_TEST_GPU_ACCELERATOR_COUNT,
+            multihost_gpu_node_count=_TEST_MULTIHOST_GPU_NODE_COUNT,
+        )
+        expected_dedicated_resources = gca_machine_resources_v1beta1.DedicatedResources(
+            machine_spec=expected_machine_spec,
+            min_replica_count=1,
+            max_replica_count=1,
+            spot=False,
+        )
+        expected_deployed_model = gca_endpoint_v1beta1.DeployedModel(
+            dedicated_resources=expected_dedicated_resources,
+            model=test_model.resource_name,
+            display_name=None,
+            enable_container_logging=True,
+            faster_deployment_config=gca_endpoint_v1beta1.FasterDeploymentConfig(),
+        )
+
+        preview_deploy_model_mock.assert_called_once_with(
+            endpoint=test_endpoint.resource_name,
+            deployed_model=expected_deployed_model,
+            traffic_split={"0": 100},
+            metadata=(),
+            timeout=None,
+        )
+
     @pytest.mark.usefixtures(
         "get_endpoint_mock", "get_model_mock", "create_endpoint_mock"
     )