Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
potiuk merged 1 commit into from
Jul 11, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions providers/google/docs/operators/cloud/vertex_ai.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +759 to +761
Copy link
Contributor

@eladkal eladkal Sep 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't like that our docs reference to internal details of other packages - especially when we are referencing main branch where any commit can make this link not accurate. doesn't the library serve a user facing doc that we can reference instead?

Copy link
Contributor

@VladaZakharova VladaZakharova Sep 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As I understand, this is specifically for users to see the limitations of the package they are trying to use, because without this it will be hard to determine what is wrong if the user will use the operator with unsupported versions of Ray.
For the link to the comment I agree, maybe better to refer to it in another way, more flexible one

Copy link
Contributor

@eladkal eladkal Sep 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The package has doc it looks like they just don't document which ray version they support. Might worth asking them to do so?
The sentence of

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

Should be something like:
To read more about compatible versions of Ray check the Google docs in <Link to upstream doc>

This avoids the need to update our docs for every change upstream library makes.


.. exampleinclude:: /../../google/tests/system/google/cloud/vertex_ai/example_vertex_ai_ray.py
:language: python
:dedent: 4
Expand Down
7 changes: 3 additions & 4 deletions providers/google/pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
12 changes: 7 additions & 5 deletions providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/ray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 7 additions & 1 deletion providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_ray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand All @@ -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]

Expand All @@ -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]

Expand Down
Loading