airflow.providers.google.cloud.operators.vertex_ai.ray

This module contains Google Vertex AI Ray operators.

Classes

RayBaseOperator

Base class for Ray operators.

CreateRayClusterOperator

Create a Ray cluster on the Vertex AI.

ListRayClustersOperator

List Ray clusters under the currently authenticated project.

GetRayClusterOperator

Get Ray cluster.

UpdateRayClusterOperator

Update Ray cluster (currently support resizing node counts for worker nodes).

DeleteRayClusterOperator

Delete Ray cluster.

Module Contents

class airflow.providers.google.cloud.operators.vertex_ai.ray.RayBaseOperator(project_id, location, gcp_conn_id='google_cloud_default', impersonation_chain=None, *args, **kwargs)[source]

Bases: airflow.providers.google.cloud.operators.cloud_base.GoogleCloudBaseOperator

Base class for Ray operators.

Parameters:

  • project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.

  • location (str) – Required. The ID of the Google Cloud region that the service belongs to.

  • gcp_conn_id (str) – The connection ID to use connecting to Google Cloud.

  • impersonation_chain (str | collections.abc.Sequence[str] | None) – Optional service account to impersonate using short-term credentials, or chained list of accounts required to get the access_token of the last account in the list, which will be impersonated in the request. If set as a string, the account must grant the originating account the Service Account Token Creator IAM role. If set as a sequence, the identities from the list must grant Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account (templated).

template_fields: collections.abc.Sequence[str] = ('location', 'gcp_conn_id', 'project_id', 'impersonation_chain')[source]
location[source]
project_id[source]
gcp_conn_id = 'google_cloud_default'[source]
impersonation_chain = None[source]
property hook: airflow.providers.google.cloud.hooks.vertex_ai.ray.RayHook[source]
class airflow.providers.google.cloud.operators.vertex_ai.ray.CreateRayClusterOperator(python_version, ray_version, head_node_type=None, network=None, service_account=None, cluster_name=None, worker_node_types=None, custom_images=None, enable_metrics_collection=True, enable_logging=True, psc_interface_config=None, reserved_ip_ranges=None, labels=None, *args, **kwargs)[source]

Bases: RayBaseOperator

Create a Ray cluster on the Vertex AI.

Parameters:

  • project_id – Required. The ID of the Google Cloud project that the service belongs to.

  • location – Required. The ID of the Google Cloud region that the service belongs to.

  • head_node_type (google.cloud.aiplatform.vertex_ray.util.resources.Resources | None) – The head node resource. Resources.node_count must be 1. If not set, default value of Resources() class will be used.

  • python_version (str) – Required. Python version for the ray cluster.

  • ray_version (Literal['2.9.3', '2.33', '2.42']) – 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

  • network (str | None) – 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.

  • service_account (str | None) – Service account to be used for running Ray programs on the cluster.

  • cluster_name (str | None) – This value may be up to 63 characters, and valid characters are [a-z0-9_-]. The first character cannot be a number or hyphen.

  • worker_node_types (list[google.cloud.aiplatform.vertex_ray.util.resources.Resources] | None) – The list of Resources of the worker nodes. The same Resources object should not appear multiple times in the list.

  • custom_images (google.cloud.aiplatform.vertex_ray.util.resources.NodeImages | None) – The NodeImages which specifies head node and worker nodes images. All the workers will share the same image. If each Resource has a specific custom image, use Resources.custom_image for head/worker_node_type(s). Note that configuring Resources.custom_image will override custom_images here. Allowlist only.

  • enable_metrics_collection (bool) – Enable Ray metrics collection for visualization.

  • enable_logging (bool) – Enable exporting Ray logs to Cloud Logging.

  • psc_interface_config (google.cloud.aiplatform.vertex_ray.util.resources.PscIConfig | None) – PSC-I config.

  • reserved_ip_ranges (list[str] | None) – A list of names for the reserved IP ranges under the VPC network that can be used for this cluster. If set, we will deploy the cluster within the provided IP ranges. Otherwise, the cluster is deployed to any IP ranges under the provided VPC network. Example: [“vertex-ai-ip-range”].

  • labels (dict[str, str] | None) – The labels with user-defined metadata to organize Ray cluster. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.

  • gcp_conn_id – The connection ID to use connecting to Google Cloud.

  • impersonation_chain – Optional service account to impersonate using short-term credentials, or chained list of accounts required to get the access_token of the last account in the list, which will be impersonated in the request. If set as a string, the account must grant the originating account the Service Account Token Creator IAM role. If set as a sequence, the identities from the list must grant Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account (templated).

template_fields: collections.abc.Sequence[str][source]
head_node_type[source]
python_version[source]
ray_version[source]
network = None[source]
service_account = None[source]
cluster_name = None[source]
worker_node_types = None[source]
custom_images = None[source]
enable_metrics_collection = True[source]
enable_logging = True[source]
psc_interface_config = None[source]
reserved_ip_ranges = None[source]
labels = None[source]
execute(context)[source]

Derive when creating an operator.

The main method to execute the task. Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

class airflow.providers.google.cloud.operators.vertex_ai.ray.ListRayClustersOperator(project_id, location, gcp_conn_id='google_cloud_default', impersonation_chain=None, *args, **kwargs)[source]

Bases: RayBaseOperator

List Ray clusters under the currently authenticated project.

Parameters:

  • project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.

  • location (str) – Required. The ID of the Google Cloud region that the service belongs to.

  • gcp_conn_id (str) – The connection ID to use connecting to Google Cloud.

  • impersonation_chain (str | collections.abc.Sequence[str] | None) – Optional service account to impersonate using short-term credentials, or chained list of accounts required to get the access_token of the last account in the list, which will be impersonated in the request. If set as a string, the account must grant the originating account the Service Account Token Creator IAM role. If set as a sequence, the identities from the list must grant Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account (templated).

execute(context)[source]

Derive when creating an operator.

The main method to execute the task. Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

class airflow.providers.google.cloud.operators.vertex_ai.ray.GetRayClusterOperator(cluster_id, *args, **kwargs)[source]

Bases: RayBaseOperator

Get Ray cluster.

Parameters:

  • project_id – Required. The ID of the Google Cloud project that the service belongs to.

  • location – Required. The ID of the Google Cloud region that the service belongs to.

  • cluster_id (str) – Cluster resource ID.

  • gcp_conn_id – The connection ID to use connecting to Google Cloud.

  • impersonation_chain – Optional service account to impersonate using short-term credentials, or chained list of accounts required to get the access_token of the last account in the list, which will be impersonated in the request. If set as a string, the account must grant the originating account the Service Account Token Creator IAM role. If set as a sequence, the identities from the list must grant Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account (templated).

template_fields: collections.abc.Sequence[str][source]
cluster_id[source]
execute(context)[source]

Derive when creating an operator.

The main method to execute the task. Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

class airflow.providers.google.cloud.operators.vertex_ai.ray.UpdateRayClusterOperator(cluster_id, worker_node_types, *args, **kwargs)[source]

Bases: RayBaseOperator

Update Ray cluster (currently support resizing node counts for worker nodes).

Parameters:

  • project_id – Required. The ID of the Google Cloud project that the service belongs to.

  • location – Required. The ID of the Google Cloud region that the service belongs to.

  • cluster_id (str) – Cluster resource ID.

  • worker_node_types (list[google.cloud.aiplatform.vertex_ray.util.resources.Resources]) – The list of Resources of the resized worker nodes. The same Resources object should not appear multiple times in the list.

  • gcp_conn_id – The connection ID to use connecting to Google Cloud.

  • impersonation_chain – Optional service account to impersonate using short-term credentials, or chained list of accounts required to get the access_token of the last account in the list, which will be impersonated in the request. If set as a string, the account must grant the originating account the Service Account Token Creator IAM role. If set as a sequence, the identities from the list must grant Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account (templated).

template_fields: collections.abc.Sequence[str][source]
cluster_id[source]
worker_node_types[source]
execute(context)[source]

Derive when creating an operator.

The main method to execute the task. Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

class airflow.providers.google.cloud.operators.vertex_ai.ray.DeleteRayClusterOperator(cluster_id, *args, **kwargs)[source]

Bases: RayBaseOperator

Delete Ray cluster.

Parameters:

  • project_id – Required. The ID of the Google Cloud project that the service belongs to.

  • location – Required. The ID of the Google Cloud region that the service belongs to.

  • cluster_id (str) – Cluster resource ID.

  • gcp_conn_id – The connection ID to use connecting to Google Cloud.

  • impersonation_chain – Optional service account to impersonate using short-term credentials, or chained list of accounts required to get the access_token of the last account in the list, which will be impersonated in the request. If set as a string, the account must grant the originating account the Service Account Token Creator IAM role. If set as a sequence, the identities from the list must grant Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account (templated).

template_fields: collections.abc.Sequence[str][source]
cluster_id[source]
execute(context)[source]

Derive when creating an operator.

The main method to execute the task. Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

Was this entry helpful?