feat: Add an arg to turn off Ray metrics collection during cluster creation

yinghsienwu · copybara-github · commit e33d11fa02eb · 2024-03-20T13:42:14.000-07:00
PiperOrigin-RevId: 617612703
diff --git a/google/cloud/aiplatform/preview/vertex_ray/cluster_init.py b/google/cloud/aiplatform/preview/vertex_ray/cluster_init.py
@@ -28,6 +28,7 @@
 from google.cloud.aiplatform_v1beta1.types.persistent_resource import (
     PersistentResource,
     RaySpec,
+    RayMetricSpec,
     ResourcePool,
     ResourceRuntimeSpec,
 )
@@ -49,6 +50,7 @@ def create_ray_cluster(
     cluster_name: Optional[str] = None,
     worker_node_types: Optional[List[resources.Resources]] = None,
     custom_images: Optional[resources.NodeImages] = None,
+    enable_metrics_collection: Optional[bool] = True,
     labels: Optional[Dict[str, str]] = None,
 ) -> str:
     """Create a ray cluster on the Vertex AI.
@@ -107,6 +109,7 @@ def create_ray_cluster(
             has a specific custom image, use `Resources.custom_image` for
             head/worker_node_type(s). Note that configuring `Resources.custom_image`
             will override `custom_images` here. Allowlist only.
+        enable_metrics_collection: Enable Ray metrics collection for visualization.
         labels:
             The labels with user-defined metadata to organize Ray cluster.
 
@@ -244,8 +247,11 @@ def create_ray_cluster(
             i += 1
 
     resource_pools = [resource_pool_0] + worker_pools
-
-    ray_spec = RaySpec(resource_pool_images=resource_pool_images)
+    disabled = not enable_metrics_collection
+    ray_metric_spec = RayMetricSpec(disabled=disabled)
+    ray_spec = RaySpec(
+        resource_pool_images=resource_pool_images, ray_metric_spec=ray_metric_spec
+    )
     resource_runtime_spec = ResourceRuntimeSpec(ray_spec=ray_spec)
     persistent_resource = PersistentResource(
         resource_pools=resource_pools,
diff --git a/tests/unit/vertex_ray/test_cluster_init.py b/tests/unit/vertex_ray/test_cluster_init.py
@@ -34,7 +34,6 @@
 
 
 # -*- coding: utf-8 -*-
-# TODO(b/328684671)
 _EXPECTED_MASK = field_mask_pb2.FieldMask(paths=["resource_pools.replica_count"])
 
 # for manual scaling
@@ -241,6 +240,22 @@ def update_persistent_resource_2_pools_mock():
         yield update_persistent_resource_2_pools_mock
 
 
+def cluster_eq(returned_cluster, expected_cluster):
+    assert vars(returned_cluster.head_node_type) == vars(
+        expected_cluster.head_node_type
+    )
+    assert vars(returned_cluster.worker_node_types[0]) == vars(
+        expected_cluster.worker_node_types[0]
+    )
+    assert (
+        returned_cluster.cluster_resource_name == expected_cluster.cluster_resource_name
+    )
+    assert returned_cluster.python_version == expected_cluster.python_version
+    assert returned_cluster.ray_version == expected_cluster.ray_version
+    assert returned_cluster.network == expected_cluster.network
+    assert returned_cluster.state == expected_cluster.state
+
+
 @pytest.mark.usefixtures("google_auth_mock", "get_project_number_mock")
 class TestClusterManagement:
     def setup_method(self):
@@ -315,6 +330,7 @@ def test_create_ray_cluster_1_pool_gpu_with_labels_success(
             network=tc.ProjectConstants.TEST_VPC_NETWORK,
             cluster_name=tc.ClusterConstants.TEST_VERTEX_RAY_PR_ID,
             labels=tc.ClusterConstants.TEST_LABELS,
+            enable_metrics_collection=False,
         )
 
         assert tc.ClusterConstants.TEST_VERTEX_RAY_PR_ADDRESS == cluster_name
@@ -465,21 +481,7 @@ def test_get_ray_cluster_success(self, get_persistent_resource_1_pool_mock):
         )
 
         get_persistent_resource_1_pool_mock.assert_called_once()
-
-        assert vars(cluster.head_node_type) == vars(
-            tc.ClusterConstants.TEST_CLUSTER.head_node_type
-        )
-        assert vars(cluster.worker_node_types[0]) == vars(
-            tc.ClusterConstants.TEST_CLUSTER.worker_node_types[0]
-        )
-        assert (
-            cluster.cluster_resource_name
-            == tc.ClusterConstants.TEST_CLUSTER.cluster_resource_name
-        )
-        assert cluster.python_version == tc.ClusterConstants.TEST_CLUSTER.python_version
-        assert cluster.ray_version == tc.ClusterConstants.TEST_CLUSTER.ray_version
-        assert cluster.network == tc.ClusterConstants.TEST_CLUSTER.network
-        assert cluster.state == tc.ClusterConstants.TEST_CLUSTER.state
+        cluster_eq(cluster, tc.ClusterConstants.TEST_CLUSTER)
 
     def test_get_ray_cluster_with_custom_image_success(
         self, get_persistent_resource_2_pools_custom_image_mock
@@ -489,27 +491,7 @@ def test_get_ray_cluster_with_custom_image_success(
         )
 
         get_persistent_resource_2_pools_custom_image_mock.assert_called_once()
-
-        assert vars(cluster.head_node_type) == vars(
-            tc.ClusterConstants.TEST_CLUSTER_CUSTOM_IMAGE.head_node_type
-        )
-        assert vars(cluster.worker_node_types[0]) == vars(
-            tc.ClusterConstants.TEST_CLUSTER_CUSTOM_IMAGE.worker_node_types[0]
-        )
-        assert (
-            cluster.cluster_resource_name
-            == tc.ClusterConstants.TEST_CLUSTER_CUSTOM_IMAGE.cluster_resource_name
-        )
-        assert (
-            cluster.python_version
-            == tc.ClusterConstants.TEST_CLUSTER_CUSTOM_IMAGE.python_version
-        )
-        assert (
-            cluster.ray_version
-            == tc.ClusterConstants.TEST_CLUSTER_CUSTOM_IMAGE.ray_version
-        )
-        assert cluster.network == tc.ClusterConstants.TEST_CLUSTER_CUSTOM_IMAGE.network
-        assert cluster.state == tc.ClusterConstants.TEST_CLUSTER_CUSTOM_IMAGE.state
+        cluster_eq(cluster, tc.ClusterConstants.TEST_CLUSTER_CUSTOM_IMAGE)
 
     @pytest.mark.usefixtures("get_persistent_resource_exception_mock")
     def test_get_ray_cluster_error(self):
@@ -526,42 +508,9 @@ def test_list_ray_clusters_success(self, list_persistent_resources_mock):
         list_persistent_resources_mock.assert_called_once()
 
         # first ray cluster
-        assert vars(clusters[0].head_node_type) == vars(
-            tc.ClusterConstants.TEST_CLUSTER.head_node_type
-        )
-        assert vars(clusters[0].worker_node_types[0]) == vars(
-            tc.ClusterConstants.TEST_CLUSTER.worker_node_types[0]
-        )
-        assert (
-            clusters[0].cluster_resource_name
-            == tc.ClusterConstants.TEST_CLUSTER.cluster_resource_name
-        )
-        assert (
-            clusters[0].python_version
-            == tc.ClusterConstants.TEST_CLUSTER.python_version
-        )
-        assert clusters[0].ray_version == tc.ClusterConstants.TEST_CLUSTER.ray_version
-        assert clusters[0].network == tc.ClusterConstants.TEST_CLUSTER.network
-        assert clusters[0].state == tc.ClusterConstants.TEST_CLUSTER.state
-
+        cluster_eq(clusters[0], tc.ClusterConstants.TEST_CLUSTER)
         # second ray cluster
-        assert vars(clusters[1].head_node_type) == vars(
-            tc.ClusterConstants.TEST_CLUSTER_2.head_node_type
-        )
-        assert vars(clusters[1].worker_node_types[0]) == vars(
-            tc.ClusterConstants.TEST_CLUSTER_2.worker_node_types[0]
-        )
-        assert (
-            clusters[1].cluster_resource_name
-            == tc.ClusterConstants.TEST_CLUSTER_2.cluster_resource_name
-        )
-        assert (
-            clusters[1].python_version
-            == tc.ClusterConstants.TEST_CLUSTER_2.python_version
-        )
-        assert clusters[1].ray_version == tc.ClusterConstants.TEST_CLUSTER_2.ray_version
-        assert clusters[1].network == tc.ClusterConstants.TEST_CLUSTER_2.network
-        assert clusters[1].state == tc.ClusterConstants.TEST_CLUSTER_2.state
+        cluster_eq(clusters[1], tc.ClusterConstants.TEST_CLUSTER_2)
 
     def test_list_ray_clusters_initialized_success(
         self, get_project_number_mock, list_persistent_resources_mock
diff --git a/tests/unit/vertex_ray/test_constants.py b/tests/unit/vertex_ray/test_constants.py
@@ -29,6 +29,9 @@
     PersistentResource,
 )
 from google.cloud.aiplatform_v1beta1.types.persistent_resource import RaySpec
+from google.cloud.aiplatform_v1beta1.types.persistent_resource import (
+    RayMetricSpec,
+)
 from google.cloud.aiplatform_v1beta1.types.persistent_resource import (
     ResourcePool,
 )
@@ -116,22 +119,31 @@ class ClusterConstants:
     TEST_REQUEST_RUNNING_1_POOL = PersistentResource(
         resource_pools=[TEST_RESOURCE_POOL_0],
         resource_runtime_spec=ResourceRuntimeSpec(
-            ray_spec=RaySpec(resource_pool_images={"head-node": TEST_GPU_IMAGE}),
+            ray_spec=RaySpec(
+                resource_pool_images={"head-node": TEST_GPU_IMAGE},
+                ray_metric_spec=RayMetricSpec(disabled=False),
+            ),
         ),
         network=ProjectConstants.TEST_VPC_NETWORK,
     )
     TEST_REQUEST_RUNNING_1_POOL_WITH_LABELS = PersistentResource(
         resource_pools=[TEST_RESOURCE_POOL_0],
         resource_runtime_spec=ResourceRuntimeSpec(
-            ray_spec=RaySpec(resource_pool_images={"head-node": TEST_GPU_IMAGE}),
+            ray_spec=RaySpec(
+                resource_pool_images={"head-node": TEST_GPU_IMAGE},
+                ray_metric_spec=RayMetricSpec(disabled=True),
+            ),
         ),
         network=ProjectConstants.TEST_VPC_NETWORK,
         labels=TEST_LABELS,
     )
     TEST_REQUEST_RUNNING_1_POOL_CUSTOM_IMAGES = PersistentResource(
         resource_pools=[TEST_RESOURCE_POOL_0],
         resource_runtime_spec=ResourceRuntimeSpec(
-            ray_spec=RaySpec(resource_pool_images={"head-node": TEST_CUSTOM_IMAGE}),
+            ray_spec=RaySpec(
+                resource_pool_images={"head-node": TEST_CUSTOM_IMAGE},
+                ray_metric_spec=RayMetricSpec(disabled=False),
+            ),
         ),
         network=ProjectConstants.TEST_VPC_NETWORK,
     )
@@ -140,7 +152,10 @@ class ClusterConstants:
         name=TEST_VERTEX_RAY_PR_ADDRESS,
         resource_pools=[TEST_RESOURCE_POOL_0],
         resource_runtime_spec=ResourceRuntimeSpec(
-            ray_spec=RaySpec(resource_pool_images={"head-node": TEST_GPU_IMAGE}),
+            ray_spec=RaySpec(
+                resource_pool_images={"head-node": TEST_GPU_IMAGE},
+                ray_metric_spec=RayMetricSpec(disabled=False),
+            ),
         ),
         network=ProjectConstants.TEST_VPC_NETWORK,
         resource_runtime=ResourceRuntime(
@@ -156,7 +171,10 @@ class ClusterConstants:
         name=TEST_VERTEX_RAY_PR_ADDRESS,
         resource_pools=[TEST_RESOURCE_POOL_0],
         resource_runtime_spec=ResourceRuntimeSpec(
-            ray_spec=RaySpec(resource_pool_images={"head-node": TEST_CUSTOM_IMAGE}),
+            ray_spec=RaySpec(
+                resource_pool_images={"head-node": TEST_CUSTOM_IMAGE},
+                ray_metric_spec=RayMetricSpec(disabled=False),
+            ),
         ),
         network=ProjectConstants.TEST_VPC_NETWORK,
         resource_runtime=ResourceRuntime(
@@ -218,7 +236,8 @@ class ClusterConstants:
                 resource_pool_images={
                     "head-node": TEST_CPU_IMAGE,
                     "worker-pool1": TEST_GPU_IMAGE,
-                }
+                },
+                ray_metric_spec=RayMetricSpec(disabled=False),
             ),
         ),
         network=ProjectConstants.TEST_VPC_NETWORK,
@@ -230,7 +249,8 @@ class ClusterConstants:
                 resource_pool_images={
                     "head-node": TEST_CUSTOM_IMAGE,
                     "worker-pool1": TEST_CUSTOM_IMAGE,
-                }
+                },
+                ray_metric_spec=RayMetricSpec(disabled=False),
             ),
         ),
         network=ProjectConstants.TEST_VPC_NETWORK,
@@ -243,7 +263,8 @@ class ClusterConstants:
                 resource_pool_images={
                     "head-node": TEST_CPU_IMAGE,
                     "worker-pool1": TEST_GPU_IMAGE,
-                }
+                },
+                ray_metric_spec=RayMetricSpec(disabled=False),
             ),
         ),
         network=ProjectConstants.TEST_VPC_NETWORK,
@@ -263,7 +284,8 @@ class ClusterConstants:
                 resource_pool_images={
                     "head-node": TEST_CUSTOM_IMAGE,
                     "worker-pool1": TEST_CUSTOM_IMAGE,
-                }
+                },
+                ray_metric_spec=RayMetricSpec(disabled=False),
             ),
         ),
         network=ProjectConstants.TEST_VPC_NETWORK,
