fix: Increase memory for Ray on Vertex default machine type

matthew29tang · copybara-github · commit e6bcb1757645 · 2024-01-22T12:31:03.000-08:00
PiperOrigin-RevId: 600536040
diff --git a/google/cloud/aiplatform/preview/vertex_ray/cluster_init.py b/google/cloud/aiplatform/preview/vertex_ray/cluster_init.py
@@ -56,14 +56,14 @@ def create_ray_cluster(
     from vertex_ray import Resources
 
     head_node_type = Resources(
-        machine_type="n1-standard-4",
+        machine_type="n1-standard-8",
         node_count=1,
         accelerator_type="NVIDIA_TESLA_K80",
         accelerator_count=1,
     )
 
     worker_node_types = [Resources(
-        machine_type="n1-standard-4",
+        machine_type="n1-standard-8",
         node_count=2,
         accelerator_type="NVIDIA_TESLA_K80",
         accelerator_count=1,
diff --git a/google/cloud/aiplatform/preview/vertex_ray/util/resources.py b/google/cloud/aiplatform/preview/vertex_ray/util/resources.py
@@ -39,7 +39,7 @@ class Resources:
             be either unspecified or within the range of [100, 64000].
     """
 
-    machine_type: Optional[str] = "n1-standard-4"
+    machine_type: Optional[str] = "n1-standard-8"
     node_count: Optional[int] = 1
     accelerator_type: Optional[str] = None
     accelerator_count: Optional[int] = 0
@@ -64,7 +64,7 @@ class Cluster:
         python_version: Python version for the ray cluster (e.g. "3_10").
         ray_version: Ray version for the ray cluster (e.g. "2_4").
         head_node_type: The head node resource. Resources.node_count must be 1.
-            If not set, by default it is a CPU node with machine_type of n1-standard-4.
+            If not set, by default it is a CPU node with machine_type of n1-standard-8.
         worker_node_types: The list of Resources of the worker nodes. Should not
             duplicate the elements in the list.
         dashboard_address: For Ray Job API (JobSubmissionClient), with this
diff --git a/tests/unit/vertex_ray/test_constants.py b/tests/unit/vertex_ray/test_constants.py
@@ -102,7 +102,7 @@ class ClusterConstants:
     _TEST_RESOURCE_POOL_0 = ResourcePool(
         id="head-node",
         machine_spec=MachineSpec(
-            machine_type="n1-standard-4",
+            machine_type="n1-standard-8",
             accelerator_type="NVIDIA_TESLA_P100",
             accelerator_count=1,
         ),
@@ -156,7 +156,7 @@ class ClusterConstants:
     _TEST_RESOURCE_POOL_1 = ResourcePool(
         id="head-node",
         machine_spec=MachineSpec(
-            machine_type="n1-standard-4",
+            machine_type="n1-standard-8",
         ),
         disk_spec=DiskSpec(
             boot_disk_type="pd-ssd",
