Adjust ray operators in google provider. Add ray package as a dependency. (#53182)

Co-authored-by: Anton Nitochkin <nitochkin@google.com>

Author: Crowiant

providers/google/docs/operators/cloud/vertex_ai.rst

@@ -756,6 +756,10 @@ Interacting with Ray on Vertex AI Cluster
 To create a Ray cluster you can use
 :class:`~airflow.providers.google.cloud.operators.vertex_ai.ray.CreateRayClusterOperator`.
 
+Please note that you need to specify python_version and ray_version in :class:`~airflow.providers.google.cloud.operators.vertex_ai.ray.CreateRayClusterOperator`.
+Currently supported versions of ray package in ray cluster are: 2.9.3, 2.33, 2.42.
+For more information you can check: https://github.com/googleapis/python-aiplatform/blob/main/setup.py#L101
+
 .. exampleinclude:: /../../google/tests/system/google/cloud/vertex_ai/example_vertex_ai_ray.py
     :language: python
     :dedent: 4

providers/google/pyproject.toml

@@ -77,10 +77,9 @@ dependencies = [
     # google-cloud-aiplatform doesn't install ray for python 3.12 (issue: https://github.com/googleapis/python-aiplatform/issues/5252).
     # Temporarily lock in ray 2.42.0 which is compatible with python 3.12 until linked issue is solved.
     # Remove the ray dependency as well as google-cloud-bigquery-storage once linked issue is fixed
-    "google-cloud-aiplatform[evaluation,ray]>=1.73.0;python_version < '3.12'",
-    "google-cloud-aiplatform[evaluation]>=1.73.0;python_version >= '3.12'",
-    "ray[default]>=2.42.0 ; python_version >= '3.12' and python_version < '3.13'",
-    "google-cloud-bigquery-storage>=2.31.0; python_version >= '3.12'",
+    "google-cloud-aiplatform[evaluation]>=1.73.0",
+    "ray[default]>=2.42.0 ; python_version < '3.13'",
+    "google-cloud-bigquery-storage>=2.31.0 ; python_version < '3.13'",
     "google-cloud-alloydb>=0.4.0",
     "google-cloud-automl>=2.12.0",
     "google-cloud-bigquery>=3.24.0",

providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/ray.py

@@ -21,7 +21,7 @@
 
 from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 from google.api_core.exceptions import NotFound
 from google.cloud.aiplatform.vertex_ray.util import resources
@@ -93,8 +93,10 @@ class CreateRayClusterOperator(RayBaseOperator):
     :param location: Required. The ID of the Google Cloud region that the service belongs to.
     :param head_node_type: The head node resource. Resources.node_count must be 1. If not set, default
         value of Resources() class will be used.
-    :param python_version: Python version for the ray cluster.
-    :param ray_version: Ray version for the ray cluster. Default is 2.33.0.
+    :param python_version: Required. Python version for the ray cluster.
+    :param ray_version: Required. Ray version for the ray cluster.
+        Currently only 3 version are available: 2.9.3, 2.33, 2.42. For more information please refer to
+        https://github.com/googleapis/python-aiplatform/blob/main/setup.py#L101
     :param 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 because Ray Cluster connection can be accessed through dashboard address.
@@ -136,9 +138,9 @@ class CreateRayClusterOperator(RayBaseOperator):
 
     def __init__(
         self,
+        python_version: str,
+        ray_version: Literal["2.9.3", "2.33", "2.42"],
         head_node_type: resources.Resources = resources.Resources(),
-        python_version: str = "3.10",
-        ray_version: str = "2.33",
         network: str | None = None,
         service_account: str | None = None,
         cluster_name: str | None = None,

providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_ray.py

@@ -43,6 +43,9 @@
 DAG_ID = "vertex_ai_ray_operations"
 LOCATION = "us-central1"
 WORKER_NODE_RESOURCES = resources.Resources(
+    node_count=1,
+)
+WORKER_NODE_RESOURCES_NEW = resources.Resources(
     node_count=2,
 )
 
@@ -58,6 +61,9 @@
         task_id="create_ray_cluster",
         project_id=PROJECT_ID,
         location=LOCATION,
+        worker_node_types=[WORKER_NODE_RESOURCES],
+        python_version="3.10",
+        ray_version="2.33",
     )
     # [END how_to_cloud_vertex_ai_create_ray_cluster_operator]
 
@@ -67,7 +73,7 @@
         project_id=PROJECT_ID,
         location=LOCATION,
         cluster_id="{{ task_instance.xcom_pull(task_ids='create_ray_cluster', key='cluster_id') }}",
-        worker_node_types=[WORKER_NODE_RESOURCES],
+        worker_node_types=[WORKER_NODE_RESOURCES_NEW],
     )
     # [END how_to_cloud_vertex_ai_update_ray_cluster_operator]