chore: make final updates to include Ray v2.42.0 support on SDK, make 2.42.0 the default Ray version over 2.33.0, allow python 3.11 with v2.42.0

PiperOrigin-RevId: 740062827

Author: vertex-sdk-bot

.kokoro/presubmit/unit_ray_2-33.cfg

@@ -3,7 +3,7 @@
 # Run unit tests for Ray 2.33.0 on Python 3.10
 env_vars: {
     key: "NOX_SESSION"
-    value: "unit_ray(ray='2.33.0')"
+    value: "unit_ray-3.10(ray='2.33.0')"
 }
 
 # Run unit tests in parallel, splitting up by file

.kokoro/presubmit/unit_ray_2-42.cfg

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

.kokoro/presubmit/unit_ray_2-42_py311.cfg

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

.kokoro/presubmit/unit_ray_2-9.cfg

@@ -3,7 +3,7 @@
 # Run unit tests for Ray 2.9.3 on Python 3.10
 env_vars: {
     key: "NOX_SESSION"
-    value: "unit_ray(ray='2.9.3')"
+    value: "unit_ray-3.10(ray='2.9.3')"
 }
 
 # Run unit tests in parallel, splitting up by file

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

@@ -2,7 +2,7 @@
 
 # -*- coding: utf-8 -*-
 
-# Copyright 2022 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,7 +19,7 @@
 import sys
 
 from google.cloud.aiplatform.vertex_ray.bigquery_datasource import (
-    BigQueryDatasource,
+    _BigQueryDatasource,
 )
 from google.cloud.aiplatform.vertex_ray.client_builder import (
     VertexRayClientBuilder as ClientBuilder,
@@ -44,13 +44,13 @@
     get_job_submission_client_cluster_info,
 )
 
-if sys.version_info[1] != 10:
+if sys.version_info[1] not in (10, 11):
     print(
-        "[Ray on Vertex]: The client environment with Python version 3.10 is required."
+        "[Ray on Vertex]: The client environment with Python version 3.10 or 3.11 is required."
     )
 
 __all__ = (
-    "BigQueryDatasource",
+    "_BigQueryDatasource",
     "data",
     "ClientBuilder",
     "get_job_submission_client_cluster_info",

google/cloud/aiplatform/vertex_ray/__init__.py

@@ -2,7 +2,7 @@
 
 # -*- coding: utf-8 -*-
 
-# Copyright 2022 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -46,9 +46,9 @@
     get_job_submission_client_cluster_info,
 )
 
-if sys.version_info[1] != 10:
+if sys.version_info[1] not in (10, 11):
     print(
-        "[Ray on Vertex]: The client environment with Python version 3.10 is required."
+        "[Ray on Vertex]: The client environment with Python version 3.10 or 3.11 is required."
     )
 
 __all__ = (

google/cloud/aiplatform/vertex_ray/cluster_init.py

@@ -54,7 +54,7 @@
 def create_ray_cluster(
     head_node_type: Optional[resources.Resources] = resources.Resources(),
     python_version: Optional[str] = "3.10",
-    ray_version: Optional[str] = "2.33",
+    ray_version: Optional[str] = "2.42",
     network: Optional[str] = None,
     service_account: Optional[str] = None,
     cluster_name: Optional[str] = None,
@@ -76,17 +76,17 @@ def create_ray_cluster(
     head_node_type = Resources(
         machine_type="n1-standard-8",
         node_count=1,
-        accelerator_type="NVIDIA_TESLA_K80",
+        accelerator_type="NVIDIA_TESLA_T4",
         accelerator_count=1,
-        custom_image="us-docker.pkg.dev/my-project/ray-cpu-image.2.9:latest",  # Optional
+        custom_image="us-docker.pkg.dev/my-project/ray-cpu-image.2.33:latest",  # Optional
     )
 
     worker_node_types = [Resources(
         machine_type="n1-standard-8",
         node_count=2,
-        accelerator_type="NVIDIA_TESLA_K80",
+        accelerator_type="NVIDIA_TESLA_T4",
         accelerator_count=1,
-        custom_image="us-docker.pkg.dev/my-project/ray-gpu-image.2.9:latest",  # Optional
+        custom_image="us-docker.pkg.dev/my-project/ray-gpu-image.2.33:latest",  # Optional
     )]
 
     cluster_resource_name = vertex_ray.create_ray_cluster(
@@ -95,7 +95,7 @@ def create_ray_cluster(
         service_account="[email protected]",  # Optional
         cluster_name="my-cluster-name",  # Optional
         worker_node_types=worker_node_types,
-        ray_version="2.9",
+        ray_version="2.33",
     )
 
     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.33.0.
+        ray_version: Ray version for the ray cluster. Default is 2.42.0.
         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.33.0"
+            install_ray_version = "2.42.0"
             logging.info(
                 "[Ray on Vertex]: Local runtime has Ray version %s"
                 ", but the requested cluster runtime has %s. Please "

google/cloud/aiplatform/vertex_ray/predict/sklearn/register.py

@@ -45,12 +45,15 @@
     if TYPE_CHECKING:
         import sklearn
 
-except ModuleNotFoundError as mnfe:
-    raise ModuleNotFoundError("Sklearn isn't installed.") from mnfe
+except ImportError as ie:
+    if ray.__version__ < "2.42.0":
+        raise ModuleNotFoundError("Sklearn isn't installed.") from ie
+    else:
+        sklearn = None
 
 
 def register_sklearn(
-    checkpoint: ray_sklearn.SklearnCheckpoint,
+    checkpoint: "ray_sklearn.SklearnCheckpoint",
     artifact_uri: Optional[str] = None,
     display_name: Optional[str] = None,
     **kwargs,
@@ -120,7 +123,7 @@ def register_sklearn(
 
 
 def _get_estimator_from(
-    checkpoint: ray_sklearn.SklearnCheckpoint,
+    checkpoint: "ray_sklearn.SklearnCheckpoint",
 ) -> "sklearn.base.BaseEstimator":
     """Converts a SklearnCheckpoint to sklearn estimator.
 

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

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-# Copyright 2023 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -28,7 +28,12 @@
 SUPPORTED_RAY_VERSIONS = immutabledict(
     {"2.9": "2.9.3", "2.33": "2.33.0", "2.42": "2.42.0"}
 )
-SUPPORTED_PY_VERSION = ["3.10"]
+SUPPORTED_RAY_VERSIONS_FROM_PYTHON_VERSIONS = immutabledict(
+    {
+        "3.10": ("2.9", "2.33", "2.42"),
+        "3.11": ("2.42"),
+    }
+)
 _V2_4_WARNING_MESSAGE = (
     "After google-cloud-aiplatform>1.53.0, using Ray version = 2.4 will result"
     " in an error. Please use Ray version = 2.33.0 or 2.42.0 (default) instead."
@@ -102,18 +107,30 @@ def get_image_uri(ray_version, python_version, enable_cuda):
                 list(SUPPORTED_RAY_VERSIONS.values())[1],
             )
         )
-    if python_version not in SUPPORTED_PY_VERSION:
-        raise ValueError("[Ray on Vertex AI]: The supported Python version is 3.10.")
+    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."
+        )
+
+    if ray_version not in SUPPORTED_RAY_VERSIONS_FROM_PYTHON_VERSIONS[python_version]:
+        raise ValueError(
+            "[Ray on Vertex AI]: The supported Ray version(s) for Python version %s: %s."
+            % (
+                python_version,
+                SUPPORTED_RAY_VERSIONS_FROM_PYTHON_VERSIONS[python_version],
+            )
+        )
 
     location = initializer.global_config.location
     region = location.split("-")[0]
     if region not in _AVAILABLE_REGIONS:
         region = _DEFAULT_REGION
     ray_version = ray_version.replace(".", "-")
+    python_version = python_version.replace(".", "")
     if enable_cuda:
-        return f"{region}-docker.pkg.dev/vertex-ai/training/ray-gpu.{ray_version}.py310:latest"
+        return f"{region}-docker.pkg.dev/vertex-ai/training/ray-gpu.{ray_version}.py{python_version}:latest"
     else:
-        return f"{region}-docker.pkg.dev/vertex-ai/training/ray-cpu.{ray_version}.py310:latest"
+        return f"{region}-docker.pkg.dev/vertex-ai/training/ray-cpu.{ray_version}.py{python_version}:latest"
 
 
 def get_versions_from_image_uri(image_uri):
@@ -122,7 +139,10 @@ def get_versions_from_image_uri(image_uri):
     image_label = image_uri.split("/")[-1].split(":")[0]
     py_version = image_label[-3] + "." + image_label[-2:]
     ray_version = image_label.split(".")[1].replace("-", ".")
-    if ray_version in SUPPORTED_RAY_VERSIONS and py_version in SUPPORTED_PY_VERSION:
+    if (
+        py_version in SUPPORTED_RAY_VERSIONS_FROM_PYTHON_VERSIONS
+        and ray_version in SUPPORTED_RAY_VERSIONS_FROM_PYTHON_VERSIONS[py_version]
+    ):
         return py_version, ray_version
     else:
         # May not parse custom image and get the versions correctly

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

@@ -69,17 +69,22 @@ class Resources:
 class NodeImages:
     """Custom images for a ray cluster.
 
-    We currently support Ray v2.9 and v2.33 and python v3.10.
+    We currently support Ray v2.9, v2.33, v2.42 and python v3.10.
+    We also support python v3.11 for Ray v2.42.
     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", or
-    "{region}-docker.pkg.dev/vertex-ai/training/ray-gpu.2-33.py310:latest". In
+    "{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
     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-9.py310-tf:latest).
-        worker: image for all worker nodes (eg. us-docker.pkg.dev/my-project/ray-gpu.2-9.py310-tf:latest).
+        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: str = None
@@ -147,7 +152,7 @@ class Cluster:
             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.9").
+        ray_version: Ray version for the ray cluster (e.g. "2.33").
         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

@@ -55,6 +55,10 @@
 UNIT_TEST_LANGCHAIN_PYTHON_VERSIONS = ["3.9", "3.10", "3.11", "3.12"]
 UNIT_TEST_AG2_PYTHON_VERSIONS = ["3.9", "3.10", "3.11", "3.12"]
 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"],
+}
 UNIT_TEST_STANDARD_DEPENDENCIES = [
     "mock",
     "asyncmock",
@@ -255,9 +259,13 @@ def unit_genai_minimal_dependencies(session):
     )
 
 
-@nox.session(python="3.10")
-@nox.parametrize("ray", ["2.9.3", "2.33.0"])
+@nox.session(python=["3.10", "3.11"])
+@nox.parametrize("ray", ["2.9.3", "2.33.0", "2.42.0"])
 def unit_ray(session, ray):
+    # check if ray version is supported
+    if ray not in PYTHON_TO_RAY_VERSIONS[session.python]:
+        session.skip(f"Ray version {ray} is not supported for python {session.python}")
+
     # Install all test dependencies, then install this package in-place.
 
     constraints_path = str(CURRENT_DIRECTORY / "testing" / f"constraints-ray-{ray}.txt")
@@ -271,7 +279,7 @@ def unit_ray(session, ray):
     session.run(
         "py.test",
         "--quiet",
-        f"--junitxml=unit_ray_{ray}_sponge_log.xml",
+        f"--junitxml=unit_ray_{ray}_py_{session.python}_sponge_log.xml",
         "--cov=google",
         "--cov-append",
         "--cov-config=.coveragerc",

testing/constraints-ray-2.42.0.txt

@@ -0,0 +1,13 @@
+ray==2.42.0
+# 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
+