diff --git a/providers/google/docs/operators/cloud/vertex_ai.rst b/providers/google/docs/operators/cloud/vertex_ai.rst index 233bddcf266e5..be535a7641dc8 100644 --- a/providers/google/docs/operators/cloud/vertex_ai.rst +++ b/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 diff --git a/providers/google/pyproject.toml b/providers/google/pyproject.toml index afeae2a406955..b09621ca9d8cf 100644 --- a/providers/google/pyproject.toml +++ b/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", diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/ray.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/ray.py index 0cd222ec6d2fa..e06c18ea1c42a 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/ray.py +++ b/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, diff --git a/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_ray.py b/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_ray.py index a0e2e6c1124d1..52008d4020cdb 100644 --- a/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_ray.py +++ b/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]