apache-airflow-providers-google 15.1.0rc1__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

airflow/providers/google/cloud/hooks/vertex_ai/ray.py

@@ -0,0 +1,223 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""This module contains a Google Cloud Vertex AI hook."""
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import MutableMapping
+from typing import Any
+
+import vertex_ray
+from google.cloud import aiplatform
+from google.cloud.aiplatform.vertex_ray.util import resources
+from google.cloud.aiplatform_v1 import (
+    PersistentResourceServiceClient,
+)
+from proto.marshal.collections.repeated import Repeated
+
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class RayHook(GoogleBaseHook):
+    """Hook for Google Cloud Vertex AI Ray APIs."""
+
+    def extract_cluster_id(self, cluster_path) -> str:
+        """Extract cluster_id from cluster_path."""
+        cluster_id = PersistentResourceServiceClient.parse_persistent_resource_path(cluster_path)[
+            "persistent_resource"
+        ]
+        return cluster_id
+
+    def serialize_cluster_obj(self, cluster_obj: resources.Cluster) -> dict:
+        """Serialize Cluster dataclass to dict."""
+
+        def __encode_value(value: Any) -> Any:
+            if isinstance(value, (list, Repeated)):
+                return [__encode_value(nested_value) for nested_value in value]
+            if not isinstance(value, dict) and isinstance(value, MutableMapping):
+                return {key: __encode_value(nested_value) for key, nested_value in dict(value).items()}
+            if dataclasses.is_dataclass(value):
+                return dataclasses.asdict(value)
+            return value
+
+        return {
+            field.name: __encode_value(getattr(cluster_obj, field.name))
+            for field in dataclasses.fields(cluster_obj)
+        }
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def create_ray_cluster(
+        self,
+        project_id: str,
+        location: str,
+        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,
+        worker_node_types: list[resources.Resources] | None = None,
+        custom_images: resources.NodeImages | None = None,
+        enable_metrics_collection: bool = True,
+        enable_logging: bool = True,
+        psc_interface_config: resources.PscIConfig | None = None,
+        reserved_ip_ranges: list[str] | None = None,
+        labels: dict[str, str] | None = None,
+    ) -> str:
+        """
+        Create a Ray cluster on the Vertex AI.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location 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 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.
+        :param service_account: Service account to be used for running Ray programs on the cluster.
+        :param cluster_name: This value may be up to 63 characters, and valid characters are `[a-z0-9_-]`.
+            The first character cannot be a number or hyphen.
+        :param worker_node_types: The list of Resources of the worker nodes. The same Resources object should
+            not appear multiple times in the list.
+        :param custom_images: 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.
+        :param enable_metrics_collection: Enable Ray metrics collection for visualization.
+        :param enable_logging: Enable exporting Ray logs to Cloud Logging.
+        :param psc_interface_config: PSC-I config.
+        :param reserved_ip_ranges: 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"].
+        :param labels: 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.
+        """
+        aiplatform.init(project=project_id, location=location, credentials=self.get_credentials())
+        cluster_path = vertex_ray.create_ray_cluster(
+            head_node_type=head_node_type,
+            python_version=python_version,
+            ray_version=ray_version,
+            network=network,
+            service_account=service_account,
+            cluster_name=cluster_name,
+            worker_node_types=worker_node_types,
+            custom_images=custom_images,
+            enable_metrics_collection=enable_metrics_collection,
+            enable_logging=enable_logging,
+            psc_interface_config=psc_interface_config,
+            reserved_ip_ranges=reserved_ip_ranges,
+            labels=labels,
+        )
+        return cluster_path
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def list_ray_clusters(
+        self,
+        project_id: str,
+        location: str,
+    ) -> list[resources.Cluster]:
+        """
+        List Ray clusters under the currently authenticated project.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        """
+        aiplatform.init(project=project_id, location=location, credentials=self.get_credentials())
+        ray_clusters = vertex_ray.list_ray_clusters()
+        return ray_clusters
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def get_ray_cluster(
+        self,
+        project_id: str,
+        location: str,
+        cluster_id: str,
+    ) -> resources.Cluster:
+        """
+        Get Ray cluster.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param cluster_id: Cluster resource ID.
+        """
+        aiplatform.init(project=project_id, location=location, credentials=self.get_credentials())
+        ray_cluster_name = PersistentResourceServiceClient.persistent_resource_path(
+            project=project_id,
+            location=location,
+            persistent_resource=cluster_id,
+        )
+        ray_cluster = vertex_ray.get_ray_cluster(
+            cluster_resource_name=ray_cluster_name,
+        )
+        return ray_cluster
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def update_ray_cluster(
+        self,
+        project_id: str,
+        location: str,
+        cluster_id: str,
+        worker_node_types: list[resources.Resources],
+    ) -> str:
+        """
+        Update Ray cluster (currently support resizing node counts for worker nodes).
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param cluster_id: Cluster resource ID.
+        :param worker_node_types: The list of Resources of the resized worker nodes. The same Resources
+            object should not appear multiple times in the list.
+        """
+        aiplatform.init(project=project_id, location=location, credentials=self.get_credentials())
+        ray_cluster_name = PersistentResourceServiceClient.persistent_resource_path(
+            project=project_id,
+            location=location,
+            persistent_resource=cluster_id,
+        )
+        updated_ray_cluster_name = vertex_ray.update_ray_cluster(
+            cluster_resource_name=ray_cluster_name, worker_node_types=worker_node_types
+        )
+        return updated_ray_cluster_name
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def delete_ray_cluster(
+        self,
+        project_id: str,
+        location: str,
+        cluster_id: str,
+    ) -> None:
+        """
+        Delete Ray cluster.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param cluster_id: Cluster resource ID.
+        """
+        aiplatform.init(project=project_id, location=location, credentials=self.get_credentials())
+        ray_cluster_name = PersistentResourceServiceClient.persistent_resource_path(
+            project=project_id,
+            location=location,
+            persistent_resource=cluster_id,
+        )
+        vertex_ray.delete_ray_cluster(cluster_resource_name=ray_cluster_name)

airflow/providers/google/cloud/hooks/vision.py

@@ -19,10 +19,10 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from copy import deepcopy
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.vision_v1 import (

airflow/providers/google/cloud/hooks/workflows.py

@@ -65,7 +65,7 @@ class WorkflowsHook(GoogleBaseHook):
         Create a new workflow.
 
         If a workflow with the specified name already exists in the
-        specified project and location, the long running operation will
+        specified project and location, the long-running operation will
         return [ALREADY_EXISTS][google.rpc.Code.ALREADY_EXISTS] error.
 
         :param workflow: Required. Workflow to be created.

airflow/providers/google/cloud/links/alloy_db.py

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.utils.context import Context
-
 ALLOY_DB_BASE_LINK = "/alloydb"
 ALLOY_DB_CLUSTER_LINK = (
     ALLOY_DB_BASE_LINK + "/locations/{location_id}/clusters/{cluster_id}?project={project_id}"
@@ -44,20 +38,6 @@ class AlloyDBClusterLink(BaseGoogleLink):
     key = "alloy_db_cluster"
     format_str = ALLOY_DB_CLUSTER_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        location_id: str,
-        cluster_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=AlloyDBClusterLink.key,
-            value={"location_id": location_id, "cluster_id": cluster_id, "project_id": project_id},
-        )
-
 
 class AlloyDBUsersLink(BaseGoogleLink):
     """Helper class for constructing AlloyDB users Link."""
@@ -66,20 +46,6 @@ class AlloyDBUsersLink(BaseGoogleLink):
     key = "alloy_db_users"
     format_str = ALLOY_DB_USERS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        location_id: str,
-        cluster_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=AlloyDBUsersLink.key,
-            value={"location_id": location_id, "cluster_id": cluster_id, "project_id": project_id},
-        )
-
 
 class AlloyDBBackupsLink(BaseGoogleLink):
     """Helper class for constructing AlloyDB backups Link."""
@@ -87,15 +53,3 @@ class AlloyDBBackupsLink(BaseGoogleLink):
     name = "AlloyDB Backups"
     key = "alloy_db_backups"
     format_str = ALLOY_DB_BACKUPS_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=AlloyDBBackupsLink.key,
-            value={"project_id": project_id},
-        )

airflow/providers/google/cloud/links/base.py

@@ -18,19 +18,15 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, ClassVar
+from urllib.parse import urlparse
 
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.providers.common.compat.sdk import BaseOperatorLink, BaseSensorOperator, XCom
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator
 
 if TYPE_CHECKING:
-    from airflow.models import BaseOperator
     from airflow.models.taskinstancekey import TaskInstanceKey
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseOperatorLink
-    from airflow.sdk.execution_time.xcom import XCom
-else:
-    from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
-    from airflow.models.xcom import XCom  # type: ignore[no-redef]
+    from airflow.providers.common.compat.sdk import Context
+    from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 
 BASE_LINK = "https://console.cloud.google.com"
 
@@ -39,6 +35,12 @@ class BaseGoogleLink(BaseOperatorLink):
     """
     Base class for all Google links.
 
+    When you inherit this class in a Link class;
+      - You can call the persist method to push data to the XCom to use it later in the get_link method.
+      - If you have an operator which inherit the GoogleCloudBaseOperator or BaseSensorOperator
+        You can define extra_links_params method in the operator to pass the operator properties
+        to the get_link method.
+
     :meta private:
     """
 
@@ -46,15 +48,77 @@ class BaseGoogleLink(BaseOperatorLink):
     key: ClassVar[str]
     format_str: ClassVar[str]
 
+    @property
+    def xcom_key(self) -> str:
+        # NOTE: in Airflow 3 we need to have xcom_key property in the Link class.
+        # Since we have the key property already, this is just a proxy property method to use same
+        # key as in Airflow 2.
+        return self.key
+
+    @classmethod
+    def persist(cls, context: Context, **value):
+        """
+        Push arguments to the XCom to use later for link formatting at the `get_link` method.
+
+        Note: for Airflow 2 we need to call this function with context variable only
+        where we have the extra_links_params property method defined
+        """
+        params = {}
+        # TODO: remove after Airflow v2 support dropped
+        if not AIRFLOW_V_3_0_PLUS:
+            common_params = getattr(context["task"], "extra_links_params", None)
+            if common_params:
+                params.update(common_params)
+
+        context["ti"].xcom_push(
+            key=cls.key,
+            value={
+                **params,
+                **value,
+            },
+        )
+
+    def get_config(self, operator, ti_key):
+        conf = {}
+        conf.update(getattr(operator, "extra_links_params", {}))
+        conf.update(XCom.get_value(key=self.key, ti_key=ti_key) or {})
+
+        # if the config did not define, return None to stop URL formatting
+        if not conf:
+            return None
+
+        # Add a default value for the 'namespace' parameter for backward compatibility.
+        # This is for datafusion
+        conf.setdefault("namespace", "default")
+        return conf
+
     def get_link(
         self,
         operator: BaseOperator,
         *,
         ti_key: TaskInstanceKey,
     ) -> str:
-        conf = XCom.get_value(key=self.key, ti_key=ti_key)
+        if TYPE_CHECKING:
+            assert isinstance(operator, (GoogleCloudBaseOperator, BaseSensorOperator))
+
+        # In cases when worker passes execution to trigger, the value that is put to XCom
+        # already contains link to the object in string format. In this case we don't want to execute
+        # get_config() again. Instead we can leave this value without any changes
+        link_value = XCom.get_value(key=self.key, ti_key=ti_key)
+        if link_value and isinstance(link_value, str):
+            if urlparse(link_value).scheme in ("http", "https"):
+                return link_value
+
+        conf = self.get_config(operator, ti_key)
         if not conf:
             return ""
-        if self.format_str.startswith("http"):
-            return self.format_str.format(**conf)
-        return BASE_LINK + self.format_str.format(**conf)
+        return self._format_link(**conf)
+
+    def _format_link(self, **kwargs):
+        try:
+            formatted_str = self.format_str.format(**kwargs)
+            if formatted_str.startswith("http"):
+                return formatted_str
+            return BASE_LINK + formatted_str
+        except KeyError:
+            return ""

airflow/providers/google/cloud/links/bigquery.py

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.utils.context import Context
-
 BIGQUERY_BASE_LINK = "/bigquery"
 BIGQUERY_DATASET_LINK = (
     BIGQUERY_BASE_LINK + "?referrer=search&project={project_id}&d={dataset_id}&p={project_id}&page=dataset"
@@ -47,19 +41,6 @@ class BigQueryDatasetLink(BaseGoogleLink):
     key = "bigquery_dataset"
     format_str = BIGQUERY_DATASET_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        dataset_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryDatasetLink.key,
-            value={"dataset_id": dataset_id, "project_id": project_id},
-        )
-
 
 class BigQueryTableLink(BaseGoogleLink):
     """Helper class for constructing BigQuery Table Link."""
@@ -68,20 +49,6 @@ class BigQueryTableLink(BaseGoogleLink):
     key = "bigquery_table"
     format_str = BIGQUERY_TABLE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str,
-        table_id: str,
-        dataset_id: str | None = None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryTableLink.key,
-            value={"dataset_id": dataset_id, "project_id": project_id, "table_id": table_id},
-        )
-
 
 class BigQueryJobDetailLink(BaseGoogleLink):
     """Helper class for constructing BigQuery Job Detail Link."""
@@ -89,17 +56,3 @@ class BigQueryJobDetailLink(BaseGoogleLink):
     name = "BigQuery Job Detail"
     key = "bigquery_job_detail"
     format_str = BIGQUERY_JOB_DETAIL_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str,
-        location: str,
-        job_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryJobDetailLink.key,
-            value={"project_id": project_id, "location": location, "job_id": job_id},
-        )

airflow/providers/google/cloud/links/bigquery_dts.py

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.utils.context import Context
-
 BIGQUERY_BASE_LINK = "/bigquery/transfers"
 BIGQUERY_DTS_LINK = BIGQUERY_BASE_LINK + "/locations/{region}/configs/{config_id}/runs?project={project_id}"
 
@@ -37,17 +31,3 @@ class BigQueryDataTransferConfigLink(BaseGoogleLink):
     name = "BigQuery Data Transfer Config"
     key = "bigquery_dts_config"
     format_str = BIGQUERY_DTS_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        region: str,
-        config_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryDataTransferConfigLink.key,
-            value={"project_id": project_id, "region": region, "config_id": config_id},
-        )

airflow/providers/google/cloud/links/bigtable.py

@@ -16,13 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.utils.context import Context
-
 BIGTABLE_BASE_LINK = "/bigtable"
 BIGTABLE_INSTANCE_LINK = BIGTABLE_BASE_LINK + "/instances/{instance_id}/overview?project={project_id}"
 BIGTABLE_CLUSTER_LINK = (
@@ -38,20 +33,6 @@ class BigtableInstanceLink(BaseGoogleLink):
     key = "instance_key"
     format_str = BIGTABLE_INSTANCE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=BigtableInstanceLink.key,
-            value={
-                "instance_id": task_instance.instance_id,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class BigtableClusterLink(BaseGoogleLink):
     """Helper class for constructing Bigtable Cluster link."""
@@ -60,21 +41,6 @@ class BigtableClusterLink(BaseGoogleLink):
     key = "cluster_key"
     format_str = BIGTABLE_CLUSTER_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=BigtableClusterLink.key,
-            value={
-                "instance_id": task_instance.instance_id,
-                "cluster_id": task_instance.cluster_id,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class BigtableTablesLink(BaseGoogleLink):
     """Helper class for constructing Bigtable Tables link."""
@@ -82,17 +48,3 @@ class BigtableTablesLink(BaseGoogleLink):
     name = "Bigtable Tables"
     key = "tables_key"
     format_str = BIGTABLE_TABLES_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=BigtableTablesLink.key,
-            value={
-                "instance_id": task_instance.instance_id,
-                "project_id": task_instance.project_id,
-            },
-        )