chore: make final updates to include Ray v2.47.1 support on SDK, make 2.47.1 the default Ray version.

PiperOrigin-RevId: 785583340

Author: vertex-sdk-bot

.kokoro/presubmit/unit_ray_2-47_py311.cfg

@@ -0,0 +1,13 @@
+# Format: //devtools/kokoro/config/proto/build.proto
+
+# Run unit tests for Ray 2.47.1 on Python 3.11
+env_vars: {
+    key: "NOX_SESSION"
+    value: "unit_ray-3.11(ray='2.47.1')"
+}
+
+# Run unit tests in parallel, splitting up by file
+env_vars: {
+    key: "PYTEST_ADDOPTS"
+    value: "-n=auto --dist=loadscope"
+}

google/cloud/aiplatform/vertex_ray/cluster_init.py

@@ -53,8 +53,8 @@
 
 def create_ray_cluster(
     head_node_type: Optional[resources.Resources] = resources.Resources(),
-    python_version: Optional[str] = "3.10",
-    ray_version: Optional[str] = "2.42",
+    python_version: Optional[str] = None,
+    ray_version: Optional[str] = "2.47",
     network: Optional[str] = None,
     service_account: Optional[str] = None,
     cluster_name: Optional[str] = None,
@@ -78,15 +78,15 @@ def create_ray_cluster(
         node_count=1,
         accelerator_type="NVIDIA_TESLA_T4",
         accelerator_count=1,
-        custom_image="us-docker.pkg.dev/my-project/ray-cpu-image.2.33:latest",  # Optional
+        custom_image="us-docker.pkg.dev/my-project/ray-2-47-cpu-py3.11:latest",  # Optional
     )
 
     worker_node_types = [Resources(
         machine_type="n1-standard-8",
         node_count=2,
         accelerator_type="NVIDIA_TESLA_T4",
         accelerator_count=1,
-        custom_image="us-docker.pkg.dev/my-project/ray-gpu-image.2.33:latest",  # Optional
+        custom_image="us-docker.pkg.dev/my-project/ray-2-47-gpu-py3.11:latest",  # Optional
     )]
 
     cluster_resource_name = vertex_ray.create_ray_cluster(
@@ -95,7 +95,7 @@ def create_ray_cluster(
         service_account="my-service-account@my-project-number.iam.gserviceaccount.com",  # Optional
         cluster_name="my-cluster-name",  # Optional
         worker_node_types=worker_node_types,
-        ray_version="2.33",
+        ray_version="2.47",
     )
 
     After a ray cluster is set up, you can call
@@ -109,7 +109,7 @@ def create_ray_cluster(
         head_node_type: The head node resource. Resources.node_count must be 1.
             If not set, default value of Resources() class will be used.
         python_version: Python version for the ray cluster.
-        ray_version: Ray version for the ray cluster. Default is 2.42.0.
+        ray_version: Ray version for the ray cluster. Default is 2.47.1.
         network: Virtual private cloud (VPC) network. For Ray Client, VPC
             peering is required to connect to the Ray Cluster managed in the
             Vertex API service. For Ray Job API, VPC network is not required
@@ -162,7 +162,7 @@ def create_ray_cluster(
     local_ray_verion = _validation_utils.get_local_ray_version()
     if ray_version != local_ray_verion:
         if custom_images is None and head_node_type.custom_image is None:
-            install_ray_version = "2.42.0"
+            install_ray_version = "2.47.1"
             logging.info(
                 "[Ray on Vertex]: Local runtime has Ray version %s"
                 ", but the requested cluster runtime has %s. Please "

google/cloud/aiplatform/vertex_ray/util/_validation_utils.py

@@ -40,7 +40,7 @@
 )
 _V2_9_WARNING_MESSAGE = (
     "In March 2025, using Ray version = 2.9 will result in an error. "
-    "Please use Ray version = 2.33.0, 2.42.0, or 2.47.1 (default) instead."
+    "Please use Ray version = 2.33.0, 2.42.0 or 2.47.1 (default) instead."
 )
 
 
@@ -107,6 +107,11 @@ def get_image_uri(ray_version, python_version, enable_cuda):
                 list(SUPPORTED_RAY_VERSIONS.values())[1],
             )
         )
+    if python_version is None:
+        for pv, ray_versions in SUPPORTED_RAY_VERSIONS_FROM_PYTHON_VERSIONS.items():
+            if ray_version in ray_versions:
+                python_version = pv
+                break
     if python_version not in SUPPORTED_RAY_VERSIONS_FROM_PYTHON_VERSIONS:
         raise ValueError(
             "[Ray on Vertex AI]: The supported Python versions are 3.10 or 3.11."

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

@@ -70,21 +70,23 @@ class NodeImages:
     """Custom images for a ray cluster.
 
     We currently support Ray v2.9, v2.33, v2.42 and python v3.10.
-    We also support python v3.11 for Ray v2.42.
+    We also support python v3.11 for Ray v2.42 and v2.47.
     The custom images must be extended from the following base images:
     "{region}-docker.pkg.dev/vertex-ai/training/ray-cpu.2-9.py310:latest",
     "{region}-docker.pkg.dev/vertex-ai/training/ray-gpu.2-9.py310:latest",
     "{region}-docker.pkg.dev/vertex-ai/training/ray-cpu.2-33.py310:latest",
     "{region}-docker.pkg.dev/vertex-ai/training/ray-gpu.2-33.py310:latest",
     "{region}-docker.pkg.dev/vertex-ai/training/ray-cpu.2-42.py310:latest",
     "{region}-docker.pkg.dev/vertex-ai/training/ray-gpu.2-42.py310:latest",
-    "{region}-docker.pkg.dev/vertex-ai/training/ray-cpu.2-42.py311:latest", or
-    "{region}-docker.pkg.dev/vertex-ai/training/ray-gpu.2-42.py311:latest". In
+    "{region}-docker.pkg.dev/vertex-ai/training/ray-cpu.2-42.py311:latest",
+    "{region}-docker.pkg.dev/vertex-ai/training/ray-gpu.2-42.py311:latest".
+    "{region}-docker.pkg.dev/vertex-ai/training/ray-cpu.2-47.py311:latest", or
+    "{region}-docker.pkg.dev/vertex-ai/training/ray-gpu.2-47.py311:latest". In
     order to use custom images, need to specify both head and worker images.
 
     Attributes:
-        head: image for head node (eg. us-docker.pkg.dev/my-project/ray-cpu.2-33.py310-tf:latest).
-        worker: image for all worker nodes (eg. us-docker.pkg.dev/my-project/ray-gpu.2-33.py310-tf:latest).
+        head: image for head node (eg. us-docker.pkg.dev/my-project/ray-cpu.2-47.py311:latest).
+        worker: image for all worker nodes (eg. us-docker.pkg.dev/my-project/ray-gpu.2-47.py311:latest).
     """
 
     head: str = None
@@ -151,8 +153,8 @@ class Cluster:
         service_account: Service account to be used for running Ray programs on
             the cluster.
         state: Describes the cluster state (defined in PersistentResource.State).
-        python_version: Python version for the ray cluster (e.g. "3.10").
-        ray_version: Ray version for the ray cluster (e.g. "2.33").
+        python_version: Python version for the ray cluster (e.g. "3.11").
+        ray_version: Ray version for the ray cluster (e.g. "2.47").
         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-8.
         worker_node_types: The list of Resources of the worker nodes. Should not

noxfile.py

@@ -57,7 +57,7 @@
 UNIT_TEST_LLAMA_INDEX_PYTHON_VERSIONS = ["3.9", "3.10", "3.11", "3.12"]
 PYTHON_TO_RAY_VERSIONS = {
     "3.10": ["2.9.3", "2.33.0", "2.42.0"],
-    "3.11": ["2.42.0"],
+    "3.11": ["2.42.0", "2.47.1"],
 }
 UNIT_TEST_STANDARD_DEPENDENCIES = [
     "mock",
@@ -261,7 +261,7 @@ def unit_genai_minimal_dependencies(session):
 
 
 @nox.session(python=["3.10", "3.11"])
-@nox.parametrize("ray", ["2.9.3", "2.33.0", "2.42.0"])
+@nox.parametrize("ray", ["2.9.3", "2.33.0", "2.42.0", "2.47.1"])
 def unit_ray(session, ray):
     # check if ray version is supported
     if ray not in PYTHON_TO_RAY_VERSIONS[session.python]:

testing/constraints-ray-2.47.1.txt

@@ -0,0 +1,13 @@
+ray==2.47.1
+# Below constraints are inherited from constraints-3.10.txt
+google-api-core
+proto-plus==1.22.3
+protobuf
+mock==4.0.2
+google-cloud-storage==2.2.1 # Increased for kfp 2.0 compatibility
+packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
+grpcio-testing==1.34.0
+mlflow==1.30.1 # Pinned to speed up installation
+pytest-xdist==3.3.1 # Pinned to unbreak unit tests
+IPython # Added to test supernova rich html buttons
+

tests/unit/vertex_ray/test_cluster_init.py

@@ -290,7 +290,7 @@ def cluster_eq(returned_cluster, expected_cluster):
     assert returned_cluster.state == expected_cluster.state
 
 
-@pytest.mark.parametrize("ray_version", ["2.9", "2.33", "2.42"])
+@pytest.mark.parametrize("ray_version", ["2.9", "2.33", "2.42", "2.47"])
 @pytest.mark.usefixtures("google_auth_mock", "get_project_number_mock")
 class TestClusterManagement:
     def setup_method(self, ray_version):
@@ -322,8 +322,10 @@ def test_create_ray_cluster_1_pool_gpu_success(
             head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_9
         elif ray_version == "2.33":
             head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_33
-        else:
+        elif ray_version == "2.42":
             head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_42
+        else:
+            head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_47
         test_persistent_resource.resource_runtime_spec.ray_spec.resource_pool_images[
             "head-node"
         ] = head_node_image
@@ -395,8 +397,10 @@ def test_create_ray_cluster_1_pool_gpu_with_labels_success(
             head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_9
         elif ray_version == "2.33":
             head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_33
-        else:
+        elif ray_version == "2.42":
             head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_42
+        else:
+            head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_47
 
         test_persistent_resource.resource_runtime_spec.ray_spec.resource_pool_images[
             "head-node"
@@ -461,9 +465,12 @@ def test_create_ray_cluster_2_pools_success(
         elif ray_version == "2.33":
             head_node_image = tc.ClusterConstants.TEST_CPU_IMAGE_2_33
             worker_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_33
-        else:
+        elif ray_version == "2.42":
             head_node_image = tc.ClusterConstants.TEST_CPU_IMAGE_2_42
             worker_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_42
+        else:
+            head_node_image = tc.ClusterConstants.TEST_CPU_IMAGE_2_47
+            worker_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_47
 
         test_persistent_resource.resource_runtime_spec.ray_spec.resource_pool_images[
             "head-node"
@@ -539,8 +546,10 @@ def test_create_ray_cluster_byosa_success(
             head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_9
         elif ray_version == "2.33":
             head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_33
-        else:
+        elif ray_version == "2.42":
             head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_42
+        else:
+            head_node_image = tc.ClusterConstants.TEST_GPU_IMAGE_2_47
 
         test_persistent_resource.resource_runtime_spec.ray_spec.resource_pool_images[
             "head-node"

tests/unit/vertex_ray/test_constants.py

@@ -129,6 +129,12 @@ class ClusterConstants:
     TEST_GPU_IMAGE_2_42 = (
         "us-docker.pkg.dev/vertex-ai/training/ray-gpu.2-42.py310:latest"
     )
+    TEST_CPU_IMAGE_2_47 = (
+        "us-docker.pkg.dev/vertex-ai/training/ray-cpu.2-47.py311:latest"
+    )
+    TEST_GPU_IMAGE_2_47 = (
+        "us-docker.pkg.dev/vertex-ai/training/ray-gpu.2-47.py311:latest"
+    )
     TEST_CUSTOM_IMAGE = "us-docker.pkg.dev/my-project/ray-custom-image.2.9:latest"
     TEST_PSC_NETWORK_ATTACHMENT = "my-network-attachment"
     # RUNNING Persistent Cluster w/o Ray

tests/unit/vertex_ray/test_vertex_ray_client.py

@@ -29,7 +29,6 @@
     python_version="MOCK_PYTHON_VERSION",
     ray_version="MOCK_RAY_VERSION",
     ray_commit="MOCK_RAY_COMMIT",
-    protocol_version="MOCK_PROTOCOL_VERSION",
     _num_clients=1,
     _context_to_restore=None,
 )