feat: Add explicit constraints for update_ray_cluster

PiperOrigin-RevId: 589973886

google/cloud/aiplatform/preview/vertex_ray/cluster_init.py

@@ -124,6 +124,14 @@ def create_ray_cluster(
                 "[Ray on Vertex AI]: For head_node_type, "
                 + "Resources.node_count must be 1."
             )
+        if (
+            head_node_type.accelerator_type is None
+            and head_node_type.accelerator_count > 0
+        ):
+            raise ValueError(
+                "[Ray on Vertex]: accelerator_type must be specified when"
+                + " accelerator_count is set to a value other than 0."
+            )
 
     resource_pool_images = {}
 
@@ -147,6 +155,14 @@ def create_ray_cluster(
     i = 0
     if worker_node_types:
         for worker_node_type in worker_node_types:
+            if (
+                worker_node_type.accelerator_type is None
+                and worker_node_type.accelerator_count > 0
+            ):
+                raise ValueError(
+                    "[Ray on Vertex]: accelerator_type must be specified when"
+                    + " accelerator_count is set to a value other than 0."
+                )
             # Worker and head share the same MachineSpec, merge them into the
             # same ResourcePool
             additional_replica_count = resources._check_machine_spec_identical(
@@ -327,31 +343,64 @@ def update_ray_cluster(
     Returns:
         The cluster_resource_name of the Ray cluster on Vertex.
     """
+    # worker_node_types should not be duplicated.
+    for i in range(len(worker_node_types)):
+        for j in range(len(worker_node_types)):
+            additional_replica_count = resources._check_machine_spec_identical(
+                worker_node_types[i], worker_node_types[j]
+            )
+            if additional_replica_count > 0 and i != j:
+                raise ValueError(
+                    "[Ray on Vertex AI]: Worker_node_types have duplicate machine specs: ",
+                    worker_node_types[i],
+                    "and ",
+                    worker_node_types[j],
+                )
+
     persistent_resource = _gapic_utils.get_persistent_resource(
         persistent_resource_name=cluster_resource_name
     )
 
     current_persistent_resource = copy.deepcopy(persistent_resource)
-    head_node_type = get_ray_cluster(cluster_resource_name).head_node_type
     current_persistent_resource.resource_pools[0].replica_count = 1
-    # TODO(b/300146407): Raise ValueError for duplicate resource pools
+
+    previous_ray_cluster = get_ray_cluster(cluster_resource_name)
+    head_node_type = previous_ray_cluster.head_node_type
+    previous_worker_node_types = previous_ray_cluster.worker_node_types
+
+    # new worker_node_types and previous_worker_node_types should be the same length.
+    if len(worker_node_types) != len(previous_worker_node_types):
+        raise ValueError(
+            f"[Ray on Vertex AI]: Desired number of worker_node_types ({len(worker_node_types)}) does not match the number of the existing worker_node_type({len(previous_worker_node_types)}).",
+        )
+
+    # Merge worker_node_type and head_node_type if the share
+    # the same machine spec.
     not_merged = 1
     for i in range(len(worker_node_types)):
         additional_replica_count = resources._check_machine_spec_identical(
             head_node_type, worker_node_types[i]
         )
-        if additional_replica_count != 0:
-            # merge the 1st duplicated worker with head
+        if additional_replica_count != 0 or (
+            additional_replica_count == 0 and worker_node_types[i].node_count == 0
+        ):
+            # Merge the 1st duplicated worker with head, allow scale down to 0 worker
             current_persistent_resource.resource_pools[0].replica_count = (
                 1 + additional_replica_count
             )
-            # reset not_merged
+            # Reset not_merged
             not_merged = 0
         else:
             # No duplication w/ head node, write the 2nd worker node to the 2nd resource pool.
             current_persistent_resource.resource_pools[
                 i + not_merged
             ].replica_count = worker_node_types[i].node_count
+            # New worker_node_type.node_count should be >=1 unless the worker_node_type
+            # and head_node_type are merged due to the same machine specs.
+            if worker_node_types[i].node_count == 0:
+                raise ValueError(
+                    f"[Ray on Vertex AI]: Worker_node_type ({worker_node_types[i]}) must update to >= 1 nodes",
+                )
 
     request = persistent_resource_service.UpdatePersistentResourceRequest(
         persistent_resource=current_persistent_resource,

google/cloud/aiplatform/preview/vertex_ray/util/resources.py

@@ -19,6 +19,7 @@
 from google.cloud.aiplatform_v1beta1.types import PersistentResource
 
 
+@dataclasses.dataclass
 class Resources:
     """Resources for a ray cluster node.
 
@@ -38,28 +39,12 @@ class Resources:
             be either unspecified or within the range of [100, 64000].
     """
 
-    def __init__(
-        self,
-        machine_type: Optional[str] = "n1-standard-4",
-        node_count: Optional[int] = 1,
-        accelerator_type: Optional[str] = None,
-        accelerator_count: Optional[int] = 0,
-        boot_disk_type: Optional[str] = "pd-ssd",
-        boot_disk_size_gb: Optional[int] = 100,
-    ):
-
-        self.machine_type = machine_type
-        self.node_count = node_count
-        self.accelerator_type = accelerator_type
-        self.accelerator_count = accelerator_count
-        self.boot_disk_type = boot_disk_type
-        self.boot_disk_size_gb = boot_disk_size_gb
-
-        if accelerator_type is None and accelerator_count > 0:
-            raise ValueError(
-                "[Ray on Vertex]: accelerator_type must be specified when"
-                + " accelerator_count is set to a value other than 0."
-            )
+    machine_type: Optional[str] = "n1-standard-4"
+    node_count: Optional[int] = 1
+    accelerator_type: Optional[str] = None
+    accelerator_count: Optional[int] = 0
+    boot_disk_type: Optional[str] = "pd-ssd"
+    boot_disk_size_gb: Optional[int] = 100
 
 
 @dataclasses.dataclass

tests/unit/vertex_ray/test_cluster_init.py

@@ -46,6 +46,11 @@
 )
 _TEST_RESPONSE_RUNNING_2_POOLS_RESIZE.resource_pools[1].replica_count = 1
 
+_TEST_RESPONSE_RUNNING_1_POOL_RESIZE_0_WORKER = copy.deepcopy(
+    tc.ClusterConstants._TEST_RESPONSE_RUNNING_1_POOL
+)
+_TEST_RESPONSE_RUNNING_1_POOL_RESIZE_0_WORKER.resource_pools[0].replica_count = 1
+
 
 @pytest.fixture
 def create_persistent_resource_1_pool_mock():
@@ -163,6 +168,22 @@ def update_persistent_resource_1_pool_mock():
         yield update_persistent_resource_1_pool_mock
 
 
+@pytest.fixture
+def update_persistent_resource_1_pool_0_worker_mock():
+    with mock.patch.object(
+        PersistentResourceServiceClient,
+        "update_persistent_resource",
+    ) as update_persistent_resource_1_pool_0_worker_mock:
+        update_persistent_resource_lro_mock = mock.Mock(ga_operation.Operation)
+        update_persistent_resource_lro_mock.result.return_value = (
+            _TEST_RESPONSE_RUNNING_1_POOL_RESIZE_0_WORKER
+        )
+        update_persistent_resource_1_pool_0_worker_mock.return_value = (
+            update_persistent_resource_lro_mock
+        )
+        yield update_persistent_resource_1_pool_0_worker_mock
+
+
 @pytest.fixture
 def update_persistent_resource_2_pools_mock():
     with mock.patch.object(
@@ -472,6 +493,30 @@ def test_update_ray_cluster_1_pool(self, update_persistent_resource_1_pool_mock)
 
         assert returned_name == tc.ClusterConstants._TEST_VERTEX_RAY_PR_ADDRESS
 
+    @pytest.mark.usefixtures("get_persistent_resource_1_pool_mock")
+    def test_update_ray_cluster_1_pool_to_0_worker(
+        self, update_persistent_resource_1_pool_mock
+    ):
+
+        new_worker_node_types = []
+        for worker_node_type in tc.ClusterConstants._TEST_CLUSTER.worker_node_types:
+            # resize worker node to node_count = 0
+            worker_node_type.node_count = 0
+            new_worker_node_types.append(worker_node_type)
+
+        returned_name = vertex_ray.update_ray_cluster(
+            cluster_resource_name=tc.ClusterConstants._TEST_VERTEX_RAY_PR_ADDRESS,
+            worker_node_types=new_worker_node_types,
+        )
+
+        request = persistent_resource_service.UpdatePersistentResourceRequest(
+            persistent_resource=_TEST_RESPONSE_RUNNING_1_POOL_RESIZE_0_WORKER,
+            update_mask=_EXPECTED_MASK,
+        )
+        update_persistent_resource_1_pool_mock.assert_called_once_with(request)
+
+        assert returned_name == tc.ClusterConstants._TEST_VERTEX_RAY_PR_ADDRESS
+
     @pytest.mark.usefixtures("get_persistent_resource_2_pools_mock")
     def test_update_ray_cluster_2_pools(self, update_persistent_resource_2_pools_mock):
 
@@ -493,3 +538,49 @@ def test_update_ray_cluster_2_pools(self, update_persistent_resource_2_pools_moc
         update_persistent_resource_2_pools_mock.assert_called_once_with(request)
 
         assert returned_name == tc.ClusterConstants._TEST_VERTEX_RAY_PR_ADDRESS
+
+    @pytest.mark.usefixtures("get_persistent_resource_2_pools_mock")
+    def test_update_ray_cluster_2_pools_0_worker_fail(self):
+
+        new_worker_node_types = []
+        for worker_node_type in tc.ClusterConstants._TEST_CLUSTER_2.worker_node_types:
+            # resize worker node to node_count = 0
+            worker_node_type.node_count = 0
+            new_worker_node_types.append(worker_node_type)
+
+        with pytest.raises(ValueError) as e:
+            vertex_ray.update_ray_cluster(
+                cluster_resource_name=tc.ClusterConstants._TEST_VERTEX_RAY_PR_ADDRESS,
+                worker_node_types=new_worker_node_types,
+            )
+
+            e.match(regexp=r"must update to >= 1 nodes.")
+
+    @pytest.mark.usefixtures("get_persistent_resource_1_pool_mock")
+    def test_update_ray_cluster_duplicate_worker_node_types_error(self):
+        new_worker_node_types = (
+            tc.ClusterConstants._TEST_CLUSTER_2.worker_node_types
+            + tc.ClusterConstants._TEST_CLUSTER_2.worker_node_types
+        )
+        with pytest.raises(ValueError) as e:
+            vertex_ray.update_ray_cluster(
+                cluster_resource_name=tc.ClusterConstants._TEST_VERTEX_RAY_PR_ADDRESS,
+                worker_node_types=new_worker_node_types,
+            )
+
+            e.match(regexp=r"Worker_node_types have duplicate machine specs")
+
+    @pytest.mark.usefixtures("get_persistent_resource_1_pool_mock")
+    def test_update_ray_cluster_mismatch_worker_node_types_count_error(self):
+        with pytest.raises(ValueError) as e:
+            new_worker_node_types = (
+                tc.ClusterConstants._TEST_CLUSTER_2.worker_node_types
+            )
+            vertex_ray.update_ray_cluster(
+                cluster_resource_name=tc.ClusterConstants._TEST_VERTEX_RAY_PR_ADDRESS,
+                worker_node_types=new_worker_node_types,
+            )
+
+            e.match(
+                regexp=r"does not match the number of the existing worker_node_type"
+            )