You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2025/03/31 11:47:55 UTC
(airflow) branch main updated: Remove deprecated for February and March 2025 in providers/google. Deprecation period for LifeScienceHook extended to july 08, 2025. (#48018)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b5ccbe50f5b Remove deprecated for February and March 2025 in providers/google. Deprecation period for LifeScienceHook extended to july 08, 2025. (#48018)
b5ccbe50f5b is described below
commit b5ccbe50f5bfac27b597af80105aad922a20543e
Author: Nitochkin <62...@users.noreply.github.com>
AuthorDate: Mon Mar 31 13:47:48 2025 +0200
Remove deprecated for February and March 2025 in providers/google. Deprecation period for LifeScienceHook extended to july 08, 2025. (#48018)
Co-authored-by: Anton Nitochkin <ni...@google.com>
---
.../tests/unit/always/test_project_structure.py | 14 -
providers/google/docs/changelog.rst | 62 +
providers/google/docs/operators/cloud/mlengine.rst | 171 +--
.../src/airflow/providers/google/ads/hooks/ads.py | 40 +-
.../providers/google/ads/transfers/ads_to_gcs.py | 18 +-
.../providers/google/cloud/hooks/cloud_build.py | 55 +-
.../providers/google/cloud/hooks/compute.py | 7 +-
.../providers/google/cloud/hooks/dataflow.py | 137 --
.../providers/google/cloud/hooks/life_sciences.py | 2 +-
.../google/cloud/hooks/vertex_ai/auto_ml.py | 5 +-
.../cloud/hooks/vertex_ai/batch_prediction_job.py | 2 +-
.../google/cloud/hooks/vertex_ai/custom_job.py | 274 +---
.../cloud/hooks/vertex_ai/endpoint_service.py | 3 +-
.../hooks/vertex_ai/hyperparameter_tuning_job.py | 2 +-
.../google/cloud/hooks/vertex_ai/model_service.py | 3 +-
.../providers/google/cloud/links/dataproc.py | 1 -
.../providers/google/cloud/operators/datafusion.py | 23 +-
.../providers/google/cloud/operators/dataproc.py | 141 --
.../providers/google/cloud/operators/mlengine.py | 1409 +-----------------
.../google/cloud/operators/vertex_ai/auto_ml.py | 12 -
.../google/cloud/operators/vertex_ai/custom_job.py | 23 +-
.../google/cloud/utils/mlengine_operator_utils.py | 273 ----
.../providers/google/common/hooks/base_google.py | 23 +-
providers/google/tests/deprecations_ignore.yml | 10 -
.../google/tests/unit/google/ads/hooks/test_ads.py | 31 -
.../unit/google/ads/transfers/test_ads_to_gcs.py | 36 +-
.../unit/google/cloud/hooks/test_cloud_build.py | 49 +-
.../tests/unit/google/cloud/hooks/test_dataflow.py | 514 +------
.../cloud/hooks/vertex_ai/test_custom_job.py | 204 ---
.../unit/google/cloud/operators/test_dataproc.py | 121 --
.../unit/google/cloud/operators/test_mlengine.py | 1506 +-------------------
.../unit/google/cloud/operators/test_vertex_ai.py | 12 +-
.../cloud/utils/test_mlengine_operator_utils.py | 305 ----
33 files changed, 121 insertions(+), 5367 deletions(-)
diff --git a/airflow-core/tests/unit/always/test_project_structure.py b/airflow-core/tests/unit/always/test_project_structure.py
index 95d87e2f2b4..88d9bb9e708 100644
--- a/airflow-core/tests/unit/always/test_project_structure.py
+++ b/airflow-core/tests/unit/always/test_project_structure.py
@@ -452,19 +452,7 @@ class TestGoogleProviderProjectStructure(ExampleCoverageTest, AssetsCoverageTest
"airflow.providers.google.cloud.operators.bigquery.BigQueryCreateExternalTableOperator",
"airflow.providers.google.cloud.operators.datapipeline.CreateDataPipelineOperator",
"airflow.providers.google.cloud.operators.datapipeline.RunDataPipelineOperator",
- "airflow.providers.google.cloud.operators.dataproc.DataprocScaleClusterOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineManageModelOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineManageVersionOperator",
"airflow.providers.google.cloud.operators.mlengine.MLEngineCreateModelOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineCreateVersionOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineDeleteModelOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineDeleteVersionOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineGetModelOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineListVersionsOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineSetDefaultVersionOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineStartBatchPredictionJobOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineStartTrainingJobOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineTrainingCancelJobOperator",
"airflow.providers.google.cloud.operators.vertex_ai.generative_model.TextGenerationModelPredictOperator",
"airflow.providers.google.marketing_platform.operators.GoogleDisplayVideo360CreateQueryOperator",
"airflow.providers.google.marketing_platform.operators.GoogleDisplayVideo360RunQueryOperator",
@@ -556,8 +544,6 @@ class TestGoogleProviderProjectStructure(ExampleCoverageTest, AssetsCoverageTest
"airflow.providers.google.cloud.operators.gcs.GCSDeleteBucketOperator",
"airflow.providers.google.cloud.operators.gcs.GCSDeleteObjectsOperator",
"airflow.providers.google.cloud.operators.kubernetes_engine.GKEDeleteClusterOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineDeleteModelOperator",
- "airflow.providers.google.cloud.operators.mlengine.MLEngineDeleteVersionOperator",
"airflow.providers.google.cloud.operators.pubsub.PubSubDeleteSubscriptionOperator",
"airflow.providers.google.cloud.operators.pubsub.PubSubDeleteTopicOperator",
"airflow.providers.google.cloud.operators.spanner.SpannerDeleteDatabaseInstanceOperator",
diff --git a/providers/google/docs/changelog.rst b/providers/google/docs/changelog.rst
index 684e3cfa136..53db811ce52 100644
--- a/providers/google/docs/changelog.rst
+++ b/providers/google/docs/changelog.rst
@@ -27,6 +27,68 @@
Changelog
---------
+15.0.0
+......
+
+Breaking changes
+~~~~~~~~~~~~~~~~
+
+.. warning::
+ Deprecated classes, parameters and features have been removed from the Google provider package.
+ The following breaking changes were introduced:
+
+ * Operators
+
+ * Removed ``MLEngineStartBatchPredictionJobOperator`` use ``CreateBatchPredictionJobOperator`` instead
+ * Removed ``MLEngineManageModelOperator`` use ``MLEngineCreateModelOperator`` or ``GetModelOperator`` from Vertex AI
+ instead
+ * Removed ``MLEngineDeleteModelOperator`` use ``DeleteModelOperator`` from Vertex AI instead
+ * Removed ``MLEngineManageVersionOperator`` use appropriate operator from Vertex AI instead,
+ e.g ``AddVersionAliasesOnModelOperator``
+ * Removed ``MLEngineCreateVersionOperator`` use ``parent_model`` parameter for VertexAI operators
+ * Removed ``MLEngineSetDefaultVersionOperator`` use ``SetDefaultVersionOnModelOperator`` from Vertex AI
+ instead
+ * Removed ``MLEngineListVersionsOperator`` use ``ListModelVersionsOperator`` from Vertex AI instead
+ * Removed ``MLEngineDeleteVersionOperator`` use ``DeleteModelVersionOperator`` from Vertex AI instead
+ * Removed ``MLEngineStartTrainingJobOperator`` use ``CreateCustomPythonPackageTrainingJobOperator`` from
+ Vertex AI instead
+ * Removed ``MLEngineTrainingCancelJobOperator`` use ``CancelCustomTrainingJobOperator`` from Vertex AI
+ instead
+ * Removed ``DataProcScaleClusterOperator`` use ``DataprocUpdateClusterOperator`` instead
+ * Removed ``page_size`` parameter from ``GoogleAdsToGcsOperator``
+ * Removed ``DeleteAutoMLTrainingJobOperator.training_pipeline`` method use ``training_pipeline_id`` parameter instead
+ * Removed ``DeleteCustomTrainingJobOperator.training_pipeline`` method use ``training_pipeline_id`` parameter instead
+ * Removed ``DeleteCustomTrainingJobOperator.custom_job`` method use ``custom_job_id`` parameter instead
+
+ * Links
+
+ * Removed ``DataFusionPipelineLinkHelper`` use ``airflow.providers.google.cloud.utils.helpers.resource_path_to_dict``
+ instead
+
+ * Hooks
+
+ * Removed ``page_size`` parameter from ``GoogleAdsHook.search``, ``GoogleAdsHook.search_proto_plus``,
+ ``GoogleAdsHook._search`` methods
+ * Removed ``CloudBuildHook.create_build`` method use ``CloudBuildHook.create_build_without_waiting_for_result`` method
+ * Removed ``DataflowHook.start_java_dataflow`` method please use ``airflow.providers.apache.beam.hooks.beam.start.start_java_pipeline``
+ method to start pipeline and ``providers.google.cloud.hooks.dataflow.DataflowHook.wait_for_done`` method
+ to wait for the required pipeline state instead
+ * Removed ``DataflowHook.start_python_dataflow`` method please use ``airflow.providers.apache.beam.hooks.beam.start.start_python_pipeline``
+ method to start pipeline and ``providers.google.cloud.hooks.dataflow.DataflowHook.wait_for_done`` method
+ to wait for the required pipeline state instead
+ * Removed ``CustomJobHook.cancel_pipeline_job`` method use ``PipelineJobHook.cancel_pipeline_job`` method instead
+ * Removed ``CustomJobHook.create_pipeline_job`` method use ``PipelineJobHook.create_pipeline_job`` method instead
+ * Removed ``CustomJobHook.get_pipeline_job`` method use ``PipelineJobHook.get_pipeline_job`` method instead
+ * Removed ``CustomJobHook.list_pipeline_jobs`` method use ``PipelineJobHook.list_pipeline_jobs`` method instead
+ * Removed ``CustomJobHook.delete_pipeline_job`` method use ``PipelineJobHook.delete_pipeline_job`` method instead
+ * Removed ``GoogleBaseHook.get_client_info`` method use ``airflow.providers.google.common.consts.CLIENT_INFO`` instead
+ * ``LifeSciencesHook`` life period extended to July 08, 2025
+
+ * Utils
+
+ * Removed ``create_evaluate_ops`` method
+
+* ``Remove deprecated for February and March 2025 in providers/google. #48018``
14.1.0
......
diff --git a/providers/google/docs/operators/cloud/mlengine.rst b/providers/google/docs/operators/cloud/mlengine.rst
index fe3a72c0b75..5c10c0258e6 100644
--- a/providers/google/docs/operators/cloud/mlengine.rst
+++ b/providers/google/docs/operators/cloud/mlengine.rst
@@ -36,25 +36,10 @@ Prerequisite tasks
.. include:: /operators/_partials/prerequisite_tasks.rst
-.. _howto/operator:MLEngineStartTrainingJobOperator:
-
Launching a Job
^^^^^^^^^^^^^^^
-To start a machine learning operation with AI Platform, you must launch a training job.
-This creates a virtual machine that can run code specified in the trainer file, which
-contains the main application code. A job can be initiated with the
-:class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineStartTrainingJobOperator`.
-
-.. warning::
- This operator is deprecated. Please, use
- :class:`~airflow.providers.google.cloud.operators.vertex_ai.custom_job.CreateCustomPythonPackageTrainingJobOperator`
- instead.
-
-.. exampleinclude:: /../../providers/google/tests/system/google/cloud/ml_engine/example_mlengine.py
- :language: python
- :dedent: 4
- :start-after: [START howto_operator_create_custom_python_training_job_v1]
- :end-before: [END howto_operator_create_custom_python_training_job_v1]
+This function is deprecated. All the functionality of legacy MLEngine and new features are available on
+the Vertex AI platform.
.. _howto/operator:MLEngineCreateModelOperator:
@@ -77,161 +62,41 @@ The ``model`` field should be defined with a dictionary containing the informati
:start-after: [START howto_operator_create_custom_python_training_job_v1]
:end-before: [END howto_operator_create_custom_python_training_job_v1]
-.. _howto/operator:MLEngineGetModelOperator:
-
Getting a model
^^^^^^^^^^^^^^^
-The :class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineGetModelOperator`
-can be used to obtain a model previously created. To obtain the correct model, ``model_name``
-must be defined in the operator.
-
-.. warning::
- This operator is deprecated. Please, use
- :class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.GetModelOperator`
- instead.
-
-.. exampleinclude:: /../../providers/google/tests/system/google/cloud/ml_engine/example_mlengine.py
- :language: python
- :dedent: 4
- :start-after: [START howto_operator_gcp_mlengine_get_model]
- :end-before: [END howto_operator_gcp_mlengine_get_model]
-
-You can use :ref:`Jinja templating <concepts:jinja-templating>` with the ``project_id`` and ``model``
-fields to dynamically determine their values. The result are saved to :ref:`XCom <concepts:xcom>`,
-allowing them to be used by other operators. In this case, the
-:class:`~airflow.providers.standard.operators.bash.BashOperator` is used to print the model information.
-
-.. exampleinclude:: /../../providers/google/tests/system/google/cloud/ml_engine/example_mlengine.py
- :language: python
- :dedent: 4
- :start-after: [START howto_operator_gcp_mlengine_print_model]
- :end-before: [END howto_operator_gcp_mlengine_print_model]
-
-.. _howto/operator:MLEngineCreateVersionOperator:
+This function is deprecated. All the functionality of legacy MLEngine and new features are available on
+the Vertex AI platform.
Creating model versions
^^^^^^^^^^^^^^^^^^^^^^^
-A model version is a subset of the model container where the code runs. A new version of the model can be created
-through the :class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineCreateVersionOperator`.
-The model must be specified by ``model_name``, and the ``version`` parameter should contain a dictionary of
-all the information about the version. Within the ``version`` parameter's dictionary, the ``name`` field is
-required.
-
-.. warning::
- This operator is deprecated. Please, use
- :class:`~airflow.providers.google.cloud.operators.vertex_ai.custom_job.CreateCustomPythonPackageTrainingJobOperator`
- instead. In this case, the new version of specific model could be created by specifying existing model id in
- ``parent_model`` parameter when running Training Job. This will ensure that new version of model will be trained except
- of creating new model.
-
-.. exampleinclude:: /../../providers/google/tests/system/google/cloud/ml_engine/example_mlengine.py
- :language: python
- :dedent: 4
- :start-after: [START howto_operator_create_custom_python_training_job_v1]
- :end-before: [END howto_operator_create_custom_python_training_job_v1]
-
-The :class:`~airflow.providers.google.cloud.operators.vertex_ai.custom_job.CreateCustomPythonPackageTrainingJobOperator`
-can also be used to create more versions with varying parameters.
-
-.. exampleinclude:: /../../providers/google/tests/system/google/cloud/ml_engine/example_mlengine.py
- :language: python
- :dedent: 4
- :start-after: [START howto_operator_create_custom_python_training_job_v2]
- :end-before: [END howto_operator_create_custom_python_training_job_v2]
-.. _howto/operator:MLEngineSetDefaultVersionOperator:
-.. _howto/operator:MLEngineListVersionsOperator:
+This function is deprecated. All the functionality of legacy MLEngine and new features are available on
+the Vertex AI platform.
+For model versioning please check:
+`Model versioning with Vertex AI
+<https://cloud.google.com/vertex-ai/docs/model-registry/versioning>`__
Managing model versions
^^^^^^^^^^^^^^^^^^^^^^^
-By default, the model code will run using the default model version. You can set the model version through the
-:class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineSetDefaultVersionOperator`
-by specifying the ``model_name`` and ``version_name`` parameters.
-
-.. warning::
- This operator is deprecated. Please, use
- :class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.SetDefaultVersionOnModelOperator`
- instead. The desired model version to be set as default could be passed with the model ID in ``model_id`` parameter
- in format ``projects/{project}/locations/{location}/models/{model_id}@{version_id}`` or
- ``projects/{project}/locations/{location}/models/{model_id}@{version_alias}``. By default, the first model version
- created will be marked as default.
-
-.. exampleinclude:: /../../providers/google/tests/system/google/cloud/ml_engine/example_mlengine.py
- :language: python
- :dedent: 4
- :start-after: [START howto_operator_gcp_mlengine_default_version]
- :end-before: [END howto_operator_gcp_mlengine_default_version]
-To list the model versions available, use the
-:class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineListVersionsOperator`
-while specifying the ``model_name`` parameter.
-
-.. warning::
- This operator is deprecated. Please, use
- :class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.ListModelVersionsOperator`
- instead. You can pass the name of the desired model in ``model_id`` parameter. If the model ID is passed
- with version aliases, the operator will output all the versions available for this model.
-
-.. exampleinclude:: /../../providers/google/tests/system/google/cloud/ml_engine/example_mlengine.py
- :language: python
- :dedent: 4
- :start-after: [START howto_operator_gcp_mlengine_list_versions]
- :end-before: [END howto_operator_gcp_mlengine_list_versions]
-
-.. _howto/operator:MLEngineStartBatchPredictionJobOperator:
+This function is deprecated. All the functionality of legacy MLEngine and new features are available on
+the Vertex AI platform.
+For model versioning please check:
+`Model versioning with Vertex AI
+<https://cloud.google.com/vertex-ai/docs/model-registry/versioning>`__
Making predictions
^^^^^^^^^^^^^^^^^^
-A Google Cloud AI Platform prediction job can be started with the
-:class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineStartBatchPredictionJobOperator`.
-For specifying the model origin, you need to provide either the ``model_name``, ``uri``, or ``model_name`` and
-``version_name``. If you do not provide the ``version_name``, the operator will use the default model version.
-
-.. warning::
- This operator is deprecated. Please, use
- :class:`~airflow.providers.google.cloud.operators.vertex_ai.batch_prediction_job.CreateBatchPredictionJobOperator`
- instead.
-.. exampleinclude:: /../../providers/google/tests/system/google/cloud/ml_engine/example_mlengine.py
- :language: python
- :dedent: 4
- :start-after: [START howto_operator_start_batch_prediction]
- :end-before: [END howto_operator_start_batch_prediction]
+This function is deprecated. All the functionality of legacy MLEngine and new features are available on
+the Vertex AI platform.
-.. _howto/operator:MLEngineDeleteVersionOperator:
-.. _howto/operator:MLEngineDeleteModelOperator:
Cleaning up
^^^^^^^^^^^
-A model version can be deleted with the
-:class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineDeleteVersionOperator` by
-the ``version_name`` and ``model_name`` parameters.
-.. warning::
- This operator is deprecated. Please, use
- :class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelVersionOperator`
- instead. The default version could not be deleted on the model.
-
-.. exampleinclude:: /../../providers/google/tests/system/google/cloud/ml_engine/example_mlengine.py
- :language: python
- :dedent: 4
- :start-after: [START howto_operator_gcp_mlengine_delete_version]
- :end-before: [END howto_operator_gcp_mlengine_delete_version]
-
-You can also delete a model with the
-:class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineDeleteModelOperator`
-by providing the ``model_name`` parameter.
-
-.. warning::
- This operator is deprecated. Please, use
- :class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelOperator`
- instead.
-
-.. exampleinclude:: /../../providers/google/tests/system/google/cloud/ml_engine/example_mlengine.py
- :language: python
- :dedent: 4
- :start-after: [START howto_operator_gcp_mlengine_delete_model]
- :end-before: [END howto_operator_gcp_mlengine_delete_model]
+This function is deprecated. All the functionality of legacy MLEngine and new features are available on
+the Vertex AI platform.
Evaluating a model
^^^^^^^^^^^^^^^^^^
diff --git a/providers/google/src/airflow/providers/google/ads/hooks/ads.py b/providers/google/src/airflow/providers/google/ads/hooks/ads.py
index 2e9f1873140..14c279132bd 100644
--- a/providers/google/src/airflow/providers/google/ads/hooks/ads.py
+++ b/providers/google/src/airflow/providers/google/ads/hooks/ads.py
@@ -19,7 +19,6 @@
from __future__ import annotations
-import warnings
from functools import cached_property
from tempfile import NamedTemporaryFile
from typing import IO, TYPE_CHECKING, Any, Literal
@@ -28,7 +27,7 @@ from google.ads.googleads.client import GoogleAdsClient
from google.ads.googleads.errors import GoogleAdsException
from google.auth.exceptions import GoogleAuthError
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
from airflow.hooks.base import BaseHook
from airflow.providers.google.common.hooks.base_google import get_field
@@ -116,9 +115,7 @@ class GoogleAdsHook(BaseHook):
self.google_ads_config: dict[str, Any] = {}
self.authentication_method: Literal["service_account", "developer_token"] = "service_account"
- def search(
- self, client_ids: list[str], query: str, page_size: int | None = None, **kwargs
- ) -> list[GoogleAdsRow]:
+ def search(self, client_ids: list[str], query: str, **kwargs) -> list[GoogleAdsRow]:
"""
Pull data from the Google Ads API.
@@ -134,18 +131,14 @@ class GoogleAdsHook(BaseHook):
:param client_ids: Google Ads client ID(s) to query the API for.
:param query: Google Ads Query Language query.
- :param page_size: Number of results to return per page. Max 10000 (for version 16 and 16.1)
- This parameter deprecated. After February 05, 2025, it will be removed.
:return: Google Ads API response, converted to Google Ads Row objects.
"""
- data_proto_plus = self._search(client_ids, query, page_size, **kwargs)
+ data_proto_plus = self._search(client_ids, query, **kwargs)
data_native_pb = [row._pb for row in data_proto_plus]
return data_native_pb
- def search_proto_plus(
- self, client_ids: list[str], query: str, page_size: int | None = None, **kwargs
- ) -> list[GoogleAdsRow]:
+ def search_proto_plus(self, client_ids: list[str], query: str, **kwargs) -> list[GoogleAdsRow]:
"""
Pull data from the Google Ads API.
@@ -154,11 +147,9 @@ class GoogleAdsHook(BaseHook):
:param client_ids: Google Ads client ID(s) to query the API for.
:param query: Google Ads Query Language query.
- :param page_size: Number of results to return per page. Max 10000 (for version 16 and 16.1)
- This parameter is deprecated. After February 05, 2025, it will be removed.
:return: Google Ads API response, converted to Google Ads Row objects
"""
- return self._search(client_ids, query, page_size, **kwargs)
+ return self._search(client_ids, query, **kwargs)
def list_accessible_customers(self) -> list[str]:
"""
@@ -269,37 +260,20 @@ class GoogleAdsHook(BaseHook):
self.google_ads_config["json_key_file_path"] = secrets_temp.name
- def _search(
- self, client_ids: list[str], query: str, page_size: int | None = None, **kwargs
- ) -> list[GoogleAdsRow]:
+ def _search(self, client_ids: list[str], query: str, **kwargs) -> list[GoogleAdsRow]:
"""
Pull data from the Google Ads API.
:param client_ids: Google Ads client ID(s) to query the API for.
:param query: Google Ads Query Language query.
- :param page_size: Number of results to return per page. Max 10000 (for version 16 and 16.1)
- This parameter is deprecated. After February 05, 2025, it will be removed.
:return: Google Ads API response, converted to Google Ads Row objects
"""
service = self._get_service
- extra_req_params = {}
- if self.api_version == "v16": # TODO: remove this after deprecation removal for page_size parameter
- extra_req_params["page_size"] = page_size or 10000
- else:
- if page_size:
- warnings.warn(
- "page_size parameter for the GoogleAdsHook.search and "
- "GoogleAdsHook.search_proto_plus method is deprecated and will be removed "
- "after February 05, 2025.",
- AirflowProviderDeprecationWarning,
- stacklevel=2,
- )
-
iterators = []
for client_id in client_ids:
- iterator = service.search(request={"customer_id": client_id, "query": query, **extra_req_params})
+ iterator = service.search(request={"customer_id": client_id, "query": query})
iterators.append(iterator)
self.log.info("Fetched Google Ads Iterators")
diff --git a/providers/google/src/airflow/providers/google/ads/transfers/ads_to_gcs.py b/providers/google/src/airflow/providers/google/ads/transfers/ads_to_gcs.py
index 51ed0662842..591eaf3b4f0 100644
--- a/providers/google/src/airflow/providers/google/ads/transfers/ads_to_gcs.py
+++ b/providers/google/src/airflow/providers/google/ads/transfers/ads_to_gcs.py
@@ -17,13 +17,11 @@
from __future__ import annotations
import csv
-import warnings
from collections.abc import Sequence
from operator import attrgetter
from tempfile import NamedTemporaryFile
from typing import TYPE_CHECKING
-from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.models import BaseOperator
from airflow.providers.google.ads.hooks.ads import GoogleAdsHook
from airflow.providers.google.cloud.hooks.gcs import GCSHook
@@ -54,8 +52,6 @@ class GoogleAdsToGcsOperator(BaseOperator):
:param obj: GCS path to save the object. Must be the full file path (ex. `path/to/file.txt`)
:param gcp_conn_id: Airflow Google Cloud connection ID
:param google_ads_conn_id: Airflow Google Ads connection ID
- :param page_size: The number of results per API page request. Max 10,000 (for version 16 and 16.1)
- This parameter deprecated. After March 01, 2025, it will be removed.
:param gzip: Option to compress local file or file data for upload
:param impersonation_chain: Optional service account to impersonate using short-term
credentials, or chained list of accounts required to get the access_token
@@ -87,7 +83,6 @@ class GoogleAdsToGcsOperator(BaseOperator):
obj: str,
gcp_conn_id: str = "google_cloud_default",
google_ads_conn_id: str = "google_ads_default",
- page_size: int | None = None,
gzip: bool = False,
impersonation_chain: str | Sequence[str] | None = None,
api_version: str | None = None,
@@ -101,8 +96,6 @@ class GoogleAdsToGcsOperator(BaseOperator):
self.obj = obj
self.gcp_conn_id = gcp_conn_id
self.google_ads_conn_id = google_ads_conn_id
- # TODO: remove this after deprecation removal for page_size parameter
- self.page_size = page_size or 10000 if api_version == "v16" else None
self.gzip = gzip
self.impersonation_chain = impersonation_chain
self.api_version = api_version
@@ -114,16 +107,7 @@ class GoogleAdsToGcsOperator(BaseOperator):
api_version=self.api_version,
)
- if self.api_version != "v16" and self.page_size:
- warnings.warn(
- "page_size parameter for the GoogleAdsToGcsOperator is deprecated and will be removed "
- "after March 01, 2025.",
- AirflowProviderDeprecationWarning,
- stacklevel=2,
- )
- rows = service.search(client_ids=self.client_ids, query=self.query)
- else:
- rows = service.search(client_ids=self.client_ids, query=self.query, page_size=self.page_size)
+ rows = service.search(client_ids=self.client_ids, query=self.query)
try:
getter = attrgetter(*self.attributes)
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/cloud_build.py b/providers/google/src/airflow/providers/google/cloud/hooks/cloud_build.py
index c2f4a41ec74..4c74fbd02bb 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/cloud_build.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/cloud_build.py
@@ -27,9 +27,8 @@ from google.api_core.exceptions import AlreadyExists
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
from google.cloud.devtools.cloudbuild_v1 import CloudBuildAsyncClient, CloudBuildClient, GetBuildRequest
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.deprecated import deprecated
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
@@ -178,58 +177,6 @@ class CloudBuildHook(GoogleBaseHook, OperationHelper):
return operation, id_
- @GoogleBaseHook.fallback_to_default_project_id
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="create_build_without_waiting_for_result",
- category=AirflowProviderDeprecationWarning,
- )
- def create_build(
- self,
- build: dict | Build,
- project_id: str = PROVIDE_PROJECT_ID,
- wait: bool = True,
- retry: Retry | _MethodDefault = DEFAULT,
- timeout: float | None = None,
- metadata: Sequence[tuple[str, str]] = (),
- ) -> Build:
- """
- Start a build with the specified configuration.
-
- :param build: The build resource to create. If a dict is provided, it must be of the same form
- as the protobuf message `google.cloud.devtools.cloudbuild_v1.types.Build`
- :param project_id: Optional, Google Cloud Project project_id where the function belongs.
- If set to None or missing, the default project_id from the GCP connection is used.
- :param wait: Optional, wait for operation to finish.
- :param retry: Optional, a retry object used to retry requests. If `None` is specified, requests
- will not be retried.
- :param timeout: Optional, the amount of time, in seconds, to wait for the request to complete.
- Note that if `retry` is specified, the timeout applies to each individual attempt.
- :param metadata: Optional, additional metadata that is provided to the method.
-
- """
- client = self.get_conn()
-
- self.log.info("Start creating build...")
-
- operation = client.create_build(
- request={"project_id": project_id, "build": build},
- retry=retry,
- timeout=timeout,
- metadata=metadata,
- )
-
- id_ = self._get_build_id_from_operation(operation)
-
- if not wait:
- return self.get_build(id_=id_, project_id=project_id)
-
- operation.result()
-
- self.log.info("Build has been created: %s.", id_)
-
- return self.get_build(id_=id_, project_id=project_id)
-
@GoogleBaseHook.fallback_to_default_project_id
def create_build_trigger(
self,
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/compute.py b/providers/google/src/airflow/providers/google/cloud/hooks/compute.py
index bb86fbfa255..1b97d061c98 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/compute.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/compute.py
@@ -29,6 +29,7 @@ from google.cloud.compute_v1.services.instances import InstancesClient
from googleapiclient.discovery import build
from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
if TYPE_CHECKING:
@@ -85,15 +86,15 @@ class ComputeEngineHook(GoogleBaseHook):
def get_compute_instance_template_client(self):
"""Return Compute Engine Instance Template Client."""
- return InstanceTemplatesClient(credentials=self.get_credentials(), client_info=self.client_info)
+ return InstanceTemplatesClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
def get_compute_instance_client(self):
"""Return Compute Engine Instance Client."""
- return InstancesClient(credentials=self.get_credentials(), client_info=self.client_info)
+ return InstancesClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
def get_compute_instance_group_managers_client(self):
"""Return Compute Engine Instance Group Managers Client."""
- return InstanceGroupManagersClient(credentials=self.get_credentials(), client_info=self.client_info)
+ return InstanceGroupManagersClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
@GoogleBaseHook.fallback_to_default_project_id
def insert_instance_template(
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py b/providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py
index e508d629937..a419d599179 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py
@@ -51,7 +51,6 @@ from googleapiclient.discovery import Resource, build
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.apache.beam.hooks.beam import BeamHook, BeamRunnerType, beam_options_to_args
-from airflow.providers.google.common.deprecated import deprecated
from airflow.providers.google.common.hooks.base_google import (
PROVIDE_PROJECT_ID,
GoogleBaseAsyncHook,
@@ -586,66 +585,6 @@ class DataflowHook(GoogleBaseHook):
http_authorized = self._authorize()
return build("datapipelines", "v1", http=http_authorized, cache_discovery=False)
- @_fallback_to_location_from_variables
- @_fallback_to_project_id_from_variables
- @GoogleBaseHook.fallback_to_default_project_id
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="airflow.providers.apache.beam.hooks.beam.start.start_java_pipeline, "
- "providers.google.cloud.hooks.dataflow.DataflowHook.wait_for_done",
- instructions="Please use airflow.providers.apache.beam.hooks.beam.start.start_java_pipeline "
- "to start pipeline and providers.google.cloud.hooks.dataflow.DataflowHook.wait_for_done method "
- "to wait for the required pipeline state instead.",
- category=AirflowProviderDeprecationWarning,
- )
- def start_java_dataflow(
- self,
- job_name: str,
- variables: dict,
- jar: str,
- project_id: str,
- job_class: str | None = None,
- append_job_name: bool = True,
- multiple_jobs: bool = False,
- on_new_job_id_callback: Callable[[str], None] | None = None,
- location: str = DEFAULT_DATAFLOW_LOCATION,
- ) -> None:
- """
- Start Dataflow java job.
-
- :param job_name: The name of the job.
- :param variables: Variables passed to the job.
- :param project_id: Optional, the Google Cloud project ID in which to start a job.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- :param jar: Name of the jar for the job
- :param job_class: Name of the java class for the job.
- :param append_job_name: True if unique suffix has to be appended to job name.
- :param multiple_jobs: True if to check for multiple job in dataflow
- :param on_new_job_id_callback: Callback called when the job ID is known.
- :param location: Job location.
- """
- name = self.build_dataflow_job_name(job_name, append_job_name)
-
- variables["jobName"] = name
- variables["region"] = location
- variables["project"] = project_id
-
- if "labels" in variables:
- variables["labels"] = json.dumps(variables["labels"], separators=(",", ":"))
-
- self.beam_hook.start_java_pipeline(
- variables=variables,
- jar=jar,
- job_class=job_class,
- process_line_callback=process_line_and_extract_dataflow_job_id_callback(on_new_job_id_callback),
- )
- self.wait_for_done(
- job_name=name,
- location=location,
- job_id=self.job_id,
- multiple_jobs=multiple_jobs,
- )
-
@_fallback_to_location_from_variables
@_fallback_to_project_id_from_variables
@GoogleBaseHook.fallback_to_default_project_id
@@ -1027,82 +966,6 @@ class DataflowHook(GoogleBaseHook):
"While reading job object after template execution error occurred. Job object has no id."
)
- @_fallback_to_location_from_variables
- @_fallback_to_project_id_from_variables
- @GoogleBaseHook.fallback_to_default_project_id
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="airflow.providers.apache.beam.hooks.beam.start.start_python_pipeline method, "
- "providers.google.cloud.hooks.dataflow.DataflowHook.wait_for_done",
- instructions="Please use airflow.providers.apache.beam.hooks.beam.start.start_python_pipeline method "
- "to start pipeline and providers.google.cloud.hooks.dataflow.DataflowHook.wait_for_done method "
- "to wait for the required pipeline state instead.",
- category=AirflowProviderDeprecationWarning,
- )
- def start_python_dataflow(
- self,
- job_name: str,
- variables: dict,
- dataflow: str,
- py_options: list[str],
- project_id: str,
- py_interpreter: str = "python3",
- py_requirements: list[str] | None = None,
- py_system_site_packages: bool = False,
- append_job_name: bool = True,
- on_new_job_id_callback: Callable[[str], None] | None = None,
- location: str = DEFAULT_DATAFLOW_LOCATION,
- ):
- """
- Start Dataflow job.
-
- :param job_name: The name of the job.
- :param variables: Variables passed to the job.
- :param dataflow: Name of the Dataflow process.
- :param py_options: Additional options.
- :param project_id: The ID of the GCP project that owns the job.
- If set to ``None`` or missing, the default project_id from the GCP connection is used.
- :param py_interpreter: Python version of the beam pipeline.
- If None, this defaults to the python3.
- To track python versions supported by beam and related
- issues check: https://issues.apache.org/jira/browse/BEAM-1251
- :param py_requirements: Additional python package(s) to install.
- If a value is passed to this parameter, a new virtual environment has been created with
- additional packages installed.
-
- You could also install the apache-beam package if it is not installed on your system or you want
- to use a different version.
- :param py_system_site_packages: Whether to include system_site_packages in your virtualenv.
- See virtualenv documentation for more information.
-
- This option is only relevant if the ``py_requirements`` parameter is not None.
- :param append_job_name: True if unique suffix has to be appended to job name.
- :param project_id: Optional, the Google Cloud project ID in which to start a job.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- :param on_new_job_id_callback: Callback called when the job ID is known.
- :param location: Job location.
- """
- name = self.build_dataflow_job_name(job_name, append_job_name)
- variables["job_name"] = name
- variables["region"] = location
- variables["project"] = project_id
-
- self.beam_hook.start_python_pipeline(
- variables=variables,
- py_file=dataflow,
- py_options=py_options,
- py_interpreter=py_interpreter,
- py_requirements=py_requirements,
- py_system_site_packages=py_system_site_packages,
- process_line_callback=process_line_and_extract_dataflow_job_id_callback(on_new_job_id_callback),
- )
-
- self.wait_for_done(
- job_name=name,
- location=location,
- job_id=self.job_id,
- )
-
@staticmethod
def build_dataflow_job_name(job_name: str, append_job_name: bool = True) -> str:
"""Build Dataflow job name."""
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/life_sciences.py b/providers/google/src/airflow/providers/google/cloud/hooks/life_sciences.py
index a3552db5c6b..4ccd05a8ceb 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/life_sciences.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/life_sciences.py
@@ -34,7 +34,7 @@ TIME_TO_SLEEP_IN_SECONDS = 5
@deprecated(
- planned_removal_date="March 01, 2025",
+ planned_removal_date="July 08, 2025",
use_instead="Google Cloud Batch Operators' hook",
reason="The Life Sciences API (beta) will be discontinued on July 8, 2025 "
"in favor of Google Cloud Batch.",
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py
index 900519acdcb..d0a3407ae30 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py
@@ -37,6 +37,7 @@ from google.cloud.aiplatform import (
from google.cloud.aiplatform_v1 import JobServiceClient, PipelineServiceClient
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.providers.google.common.consts import CLIENT_INFO
from airflow.providers.google.common.deprecated import deprecated
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
@@ -81,7 +82,7 @@ class AutoMLHook(GoogleBaseHook, OperationHelper):
client_options = ClientOptions()
return PipelineServiceClient(
- credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
)
def get_job_service_client(
@@ -95,7 +96,7 @@ class AutoMLHook(GoogleBaseHook, OperationHelper):
client_options = ClientOptions()
return JobServiceClient(
- credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
)
def get_auto_ml_tabular_training_job(
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py
index 544bcb9f11e..66a1e7ee0a6 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py
@@ -63,7 +63,7 @@ class BatchPredictionJobHook(GoogleBaseHook, OperationHelper):
client_options = ClientOptions()
return JobServiceClient(
- credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
)
@staticmethod
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py
index 89f917f1a3b..b2e1fc302ee 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py
@@ -42,9 +42,8 @@ from google.cloud.aiplatform_v1 import (
types,
)
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.deprecated import deprecated
from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
@@ -54,10 +53,9 @@ if TYPE_CHECKING:
from google.auth.credentials import Credentials
from google.cloud.aiplatform_v1.services.job_service.pagers import ListCustomJobsPager
from google.cloud.aiplatform_v1.services.pipeline_service.pagers import (
- ListPipelineJobsPager,
ListTrainingPipelinesPager,
)
- from google.cloud.aiplatform_v1.types import CustomJob, PipelineJob, TrainingPipeline
+ from google.cloud.aiplatform_v1.types import CustomJob, TrainingPipeline
class CustomJobHook(GoogleBaseHook, OperationHelper):
@@ -368,54 +366,6 @@ class CustomJobHook(GoogleBaseHook, OperationHelper):
)
return model, training_id, custom_job_id
- @GoogleBaseHook.fallback_to_default_project_id
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="PipelineJobHook.cancel_pipeline_job",
- category=AirflowProviderDeprecationWarning,
- )
- def cancel_pipeline_job(
- self,
- project_id: str,
- region: str,
- pipeline_job: str,
- retry: Retry | _MethodDefault = DEFAULT,
- timeout: float | None = None,
- metadata: Sequence[tuple[str, str]] = (),
- ) -> None:
- """
- Cancel a PipelineJob.
-
- Starts asynchronous cancellation on the PipelineJob. The server makes the best
- effort to cancel the pipeline, but success is not guaranteed. Clients can use
- [PipelineService.GetPipelineJob][google.cloud.aiplatform.v1.PipelineService.GetPipelineJob] or other
- methods to check whether the cancellation succeeded or whether the pipeline completed despite
- cancellation. On successful cancellation, the PipelineJob is not deleted; instead it becomes a
- pipeline with a [PipelineJob.error][google.cloud.aiplatform.v1.PipelineJob.error] value with a
- [google.rpc.Status.code][google.rpc.Status.code] of 1, corresponding to ``Code.CANCELLED``, and
- [PipelineJob.state][google.cloud.aiplatform.v1.PipelineJob.state] is set to ``CANCELLED``.
-
- This method is deprecated, please use `PipelineJobHook.cancel_pipeline_job` method.
-
- :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
- :param region: Required. The ID of the Google Cloud region that the service belongs to.
- :param pipeline_job: The name of the PipelineJob to cancel.
- :param retry: Designation of what errors, if any, should be retried.
- :param timeout: The timeout for this request.
- :param metadata: Strings which should be sent along with the request as metadata.
- """
- client = self.get_pipeline_service_client(region)
- name = client.pipeline_job_path(project_id, region, pipeline_job)
-
- client.cancel_pipeline_job(
- request={
- "name": name,
- },
- retry=retry,
- timeout=timeout,
- metadata=metadata,
- )
-
@GoogleBaseHook.fallback_to_default_project_id
def cancel_training_pipeline(
self,
@@ -498,53 +448,6 @@ class CustomJobHook(GoogleBaseHook, OperationHelper):
metadata=metadata,
)
- @GoogleBaseHook.fallback_to_default_project_id
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="PipelineJobHook.create_pipeline_job",
- category=AirflowProviderDeprecationWarning,
- )
- def create_pipeline_job(
- self,
- project_id: str,
- region: str,
- pipeline_job: PipelineJob,
- pipeline_job_id: str,
- retry: Retry | _MethodDefault = DEFAULT,
- timeout: float | None = None,
- metadata: Sequence[tuple[str, str]] = (),
- ) -> PipelineJob:
- """
- Create a PipelineJob. A PipelineJob will run immediately when created.
-
- This method is deprecated, please use `PipelineJobHook.create_pipeline_job` method.
-
- :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
- :param region: Required. The ID of the Google Cloud region that the service belongs to.
- :param pipeline_job: Required. The PipelineJob to create.
- :param pipeline_job_id: The ID to use for the PipelineJob, which will become the final component of
- the PipelineJob name. If not provided, an ID will be automatically generated.
-
- This value should be less than 128 characters, and valid characters are /[a-z][0-9]-/.
- :param retry: Designation of what errors, if any, should be retried.
- :param timeout: The timeout for this request.
- :param metadata: Strings which should be sent along with the request as metadata.
- """
- client = self.get_pipeline_service_client(region)
- parent = client.common_location_path(project_id, region)
-
- result = client.create_pipeline_job(
- request={
- "parent": parent,
- "pipeline_job": pipeline_job,
- "pipeline_job_id": pipeline_job_id,
- },
- retry=retry,
- timeout=timeout,
- metadata=metadata,
- )
- return result
-
@GoogleBaseHook.fallback_to_default_project_id
def create_training_pipeline(
self,
@@ -2970,46 +2873,6 @@ class CustomJobHook(GoogleBaseHook, OperationHelper):
)
return result
- @GoogleBaseHook.fallback_to_default_project_id
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="PipelineJobHook.get_pipeline_job",
- category=AirflowProviderDeprecationWarning,
- )
- def get_pipeline_job(
- self,
- project_id: str,
- region: str,
- pipeline_job: str,
- retry: Retry | _MethodDefault = DEFAULT,
- timeout: float | None = None,
- metadata: Sequence[tuple[str, str]] = (),
- ) -> PipelineJob:
- """
- Get a PipelineJob.
-
- This method is deprecated, please use `PipelineJobHook.get_pipeline_job` method.
-
- :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
- :param region: Required. The ID of the Google Cloud region that the service belongs to.
- :param pipeline_job: Required. The name of the PipelineJob resource.
- :param retry: Designation of what errors, if any, should be retried.
- :param timeout: The timeout for this request.
- :param metadata: Strings which should be sent along with the request as metadata.
- """
- client = self.get_pipeline_service_client(region)
- name = client.pipeline_job_path(project_id, region, pipeline_job)
-
- result = client.get_pipeline_job(
- request={
- "name": name,
- },
- retry=retry,
- timeout=timeout,
- metadata=metadata,
- )
- return result
-
@GoogleBaseHook.fallback_to_default_project_id
def get_training_pipeline(
self,
@@ -3076,101 +2939,6 @@ class CustomJobHook(GoogleBaseHook, OperationHelper):
)
return result
- @GoogleBaseHook.fallback_to_default_project_id
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="PipelineJobHook.list_pipeline_jobs",
- category=AirflowProviderDeprecationWarning,
- )
- def list_pipeline_jobs(
- self,
- project_id: str,
- region: str,
- page_size: int | None = None,
- page_token: str | None = None,
- filter: str | None = None,
- order_by: str | None = None,
- retry: Retry | _MethodDefault = DEFAULT,
- timeout: float | None = None,
- metadata: Sequence[tuple[str, str]] = (),
- ) -> ListPipelineJobsPager:
- """
- List PipelineJobs in a Location.
-
- This method is deprecated, please use `PipelineJobHook.list_pipeline_jobs` method.
-
- :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
- :param region: Required. The ID of the Google Cloud region that the service belongs to.
- :param filter: Optional. Lists the PipelineJobs that match the filter expression. The
- following fields are supported:
-
- - ``pipeline_name``: Supports ``=`` and ``!=`` comparisons.
- - ``display_name``: Supports ``=``, ``!=`` comparisons, and
- ``:`` wildcard.
- - ``pipeline_job_user_id``: Supports ``=``, ``!=``
- comparisons, and ``:`` wildcard. for example, can check
- if pipeline's display_name contains *step* by doing
- display_name:"*step*"
- - ``create_time``: Supports ``=``, ``!=``, ``<``, ``>``,
- ``<=``, and ``>=`` comparisons. Values must be in RFC
- 3339 format.
- - ``update_time``: Supports ``=``, ``!=``, ``<``, ``>``,
- ``<=``, and ``>=`` comparisons. Values must be in RFC
- 3339 format.
- - ``end_time``: Supports ``=``, ``!=``, ``<``, ``>``,
- ``<=``, and ``>=`` comparisons. Values must be in RFC
- 3339 format.
- - ``labels``: Supports key-value equality and key presence.
-
- Filter expressions can be combined together using logical
- operators (``AND`` & ``OR``). For example:
- ``pipeline_name="test" AND create_time>"2020-05-18T13:30:00Z"``.
-
- The syntax to define filter expression is based on
- https://google.aip.dev/160.
- :param page_size: Optional. The standard list page size.
- :param page_token: Optional. The standard list page token. Typically obtained via
- [ListPipelineJobsResponse.next_page_token][google.cloud.aiplatform.v1.ListPipelineJobsResponse.next_page_token]
- of the previous
- [PipelineService.ListPipelineJobs][google.cloud.aiplatform.v1.PipelineService.ListPipelineJobs]
- call.
- :param order_by: Optional. A comma-separated list of fields to order by. The default
- sort order is in ascending order. Use "desc" after a field
- name for descending. You can have multiple order_by fields
- provided e.g. "create_time desc, end_time", "end_time,
- start_time, update_time" For example, using "create_time
- desc, end_time" will order results by create time in
- descending order, and if there are multiple jobs having the
- same create time, order them by the end time in ascending
- order. if order_by is not specified, it will order by
- default order is create time in descending order. Supported
- fields:
-
- - ``create_time``
- - ``update_time``
- - ``end_time``
- - ``start_time``
- :param retry: Designation of what errors, if any, should be retried.
- :param timeout: The timeout for this request.
- :param metadata: Strings which should be sent along with the request as metadata.
- """
- client = self.get_pipeline_service_client(region)
- parent = client.common_location_path(project_id, region)
-
- result = client.list_pipeline_jobs(
- request={
- "parent": parent,
- "page_size": page_size,
- "page_token": page_token,
- "filter": filter,
- "order_by": order_by,
- },
- retry=retry,
- timeout=timeout,
- metadata=metadata,
- )
- return result
-
@GoogleBaseHook.fallback_to_default_project_id
def list_training_pipelines(
self,
@@ -3293,44 +3061,6 @@ class CustomJobHook(GoogleBaseHook, OperationHelper):
)
return result
- @GoogleBaseHook.fallback_to_default_project_id
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="PipelineJobHook.delete_pipeline_job",
- category=AirflowProviderDeprecationWarning,
- )
- def delete_pipeline_job(
- self,
- project_id: str,
- region: str,
- pipeline_job: str,
- retry: Retry | _MethodDefault = DEFAULT,
- timeout: float | None = None,
- metadata: Sequence[tuple[str, str]] = (),
- ) -> Operation:
- """
- Delete a PipelineJob.
-
- This method is deprecated, please use `PipelineJobHook.delete_pipeline_job` method.
-
- :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
- :param region: Required. The ID of the Google Cloud region that the service belongs to.
- :param pipeline_job: Required. The name of the PipelineJob resource to be deleted.
- :param retry: Designation of what errors, if any, should be retried.
- :param timeout: The timeout for this request.
- :param metadata: Strings which should be sent along with the request as metadata.
- """
- client = self.get_pipeline_service_client(region)
- name = client.pipeline_job_path(project_id, region, pipeline_job)
-
- result = client.delete_pipeline_job(
- request={"name": name},
- retry=retry,
- timeout=timeout,
- metadata=metadata,
- )
- return result
-
class CustomJobAsyncHook(GoogleBaseAsyncHook):
"""Async hook for Custom Job Service Client."""
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py
index 6d6d661686b..60b20d840ea 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py
@@ -26,6 +26,7 @@ from google.api_core.client_options import ClientOptions
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
from google.cloud.aiplatform_v1 import EndpointServiceClient
+from airflow.providers.google.common.consts import CLIENT_INFO
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
@@ -48,7 +49,7 @@ class EndpointServiceHook(GoogleBaseHook, OperationHelper):
client_options = ClientOptions()
return EndpointServiceClient(
- credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
)
@staticmethod
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py
index 4cb85f5037a..a48a686a2c1 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py
@@ -69,7 +69,7 @@ class HyperparameterTuningJobHook(GoogleBaseHook, OperationHelper):
client_options = ClientOptions()
return JobServiceClient(
- credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
)
def get_hyperparameter_tuning_job_object(
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py
index 5604d3314ef..fe9c8a6d011 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py
@@ -28,6 +28,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
from google.cloud.aiplatform_v1 import ModelServiceClient
from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
if TYPE_CHECKING:
@@ -53,7 +54,7 @@ class ModelServiceHook(GoogleBaseHook, OperationHelper):
client_options = ClientOptions()
return ModelServiceClient(
- credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
)
@staticmethod
diff --git a/providers/google/src/airflow/providers/google/cloud/links/dataproc.py b/providers/google/src/airflow/providers/google/cloud/links/dataproc.py
index 72dfef308e6..4fafd9da90f 100644
--- a/providers/google/src/airflow/providers/google/cloud/links/dataproc.py
+++ b/providers/google/src/airflow/providers/google/cloud/links/dataproc.py
@@ -126,7 +126,6 @@ class DataprocLink(BaseOperatorLink):
def __attrs_post_init__(self):
# This link is still used into the selected operators
- # - airflow.providers.google.cloud.operators.dataproc.DataprocScaleClusterOperator
# - airflow.providers.google.cloud.operators.dataproc.DataprocJobBaseOperator
# As soon as we remove reference to this link we might deprecate it by add warning message
# with `stacklevel=3` below in this method.
diff --git a/providers/google/src/airflow/providers/google/cloud/operators/datafusion.py b/providers/google/src/airflow/providers/google/cloud/operators/datafusion.py
index 02169748b39..6c6d42fb5cd 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/datafusion.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/datafusion.py
@@ -26,7 +26,7 @@ from google.api_core.retry import exponential_sleep_generator
from googleapiclient.errors import HttpError
from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
from airflow.providers.google.cloud.hooks.datafusion import SUCCESS_STATES, DataFusionHook, PipelineStates
from airflow.providers.google.cloud.links.datafusion import (
DataFusionInstanceLink,
@@ -37,33 +37,12 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseO
from airflow.providers.google.cloud.triggers.datafusion import DataFusionStartPipelineTrigger
from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
from airflow.providers.google.cloud.utils.helpers import resource_path_to_dict
-from airflow.providers.google.common.deprecated import deprecated
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
if TYPE_CHECKING:
from airflow.utils.context import Context
-class DataFusionPipelineLinkHelper:
- """
- Helper class for Pipeline links.
-
- .. warning::
- This class is deprecated. Consider using ``resource_path_to_dict()`` instead.
- """
-
- @staticmethod
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="airflow.providers.google.cloud.utils.helpers.resource_path_to_dict",
- category=AirflowProviderDeprecationWarning,
- )
- def get_project_id(instance):
- instance = instance["name"]
- project_id = next(x for x in instance.split("/") if x.startswith("airflow"))
- return project_id
-
-
class CloudDataFusionRestartInstanceOperator(GoogleCloudBaseOperator):
"""
Restart a single Data Fusion instance.
diff --git a/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py b/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py
index 6b52813a34b..64ba220f499 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py
@@ -44,7 +44,6 @@ from airflow.providers.google.cloud.hooks.dataproc import (
)
from airflow.providers.google.cloud.links.dataproc import (
DATAPROC_BATCH_LINK,
- DATAPROC_CLUSTER_LINK_DEPRECATED,
DATAPROC_JOB_LINK_DEPRECATED,
DataprocBatchesListLink,
DataprocBatchLink,
@@ -63,7 +62,6 @@ from airflow.providers.google.cloud.triggers.dataproc import (
DataprocSubmitTrigger,
)
from airflow.providers.google.cloud.utils.dataproc import DataprocOperationType
-from airflow.providers.google.common.deprecated import deprecated
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
from airflow.utils import timezone
@@ -919,145 +917,6 @@ class DataprocCreateClusterOperator(GoogleCloudBaseOperator):
return event["cluster"]
-# TODO: Remove one day
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="DataprocUpdateClusterOperator",
- category=AirflowProviderDeprecationWarning,
-)
-class DataprocScaleClusterOperator(GoogleCloudBaseOperator):
- """
- Scale, up or down, a cluster on Google Cloud Dataproc.
-
- The operator will wait until the cluster is re-scaled.
-
- Example usage:
-
- .. code-block:: python
-
- t1 = DataprocClusterScaleOperator(
- task_id="dataproc_scale",
- project_id="my-project",
- cluster_name="cluster-1",
- num_workers=10,
- num_preemptible_workers=10,
- graceful_decommission_timeout="1h",
- )
-
- .. seealso::
- For more detail on about scaling clusters have a look at the reference:
- https://cloud.google.com/dataproc/docs/concepts/configuring-clusters/scaling-clusters
-
- :param cluster_name: The name of the cluster to scale. (templated)
- :param project_id: The ID of the google cloud project in which
- the cluster runs. (templated)
- :param region: The region for the dataproc cluster. (templated)
- :param num_workers: The new number of workers
- :param num_preemptible_workers: The new number of preemptible workers
- :param graceful_decommission_timeout: Timeout for graceful YARN decommissioning.
- Maximum value is 1d
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
- :param 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: Sequence[str] = ("cluster_name", "project_id", "region", "impersonation_chain")
-
- operator_extra_links = (DataprocLink(),)
-
- def __init__(
- self,
- *,
- cluster_name: str,
- project_id: str = PROVIDE_PROJECT_ID,
- region: str = "global",
- num_workers: int = 2,
- num_preemptible_workers: int = 0,
- graceful_decommission_timeout: str | None = None,
- gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.region = region
- self.cluster_name = cluster_name
- self.num_workers = num_workers
- self.num_preemptible_workers = num_preemptible_workers
- self.graceful_decommission_timeout = graceful_decommission_timeout
- self.gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
-
- def _build_scale_cluster_data(self) -> dict:
- scale_data = {
- "config": {
- "worker_config": {"num_instances": self.num_workers},
- "secondary_worker_config": {"num_instances": self.num_preemptible_workers},
- }
- }
- return scale_data
-
- @property
- def _graceful_decommission_timeout_object(self) -> dict[str, int] | None:
- if not self.graceful_decommission_timeout:
- return None
-
- timeout = None
- match = re.fullmatch(r"(\d+)([smdh])", self.graceful_decommission_timeout)
- if match:
- val = int(match.group(1))
- unit = match.group(2)
- if unit == "s":
- timeout = val
- elif unit == "m":
- timeout = int(timedelta(minutes=val).total_seconds())
- elif unit == "h":
- timeout = int(timedelta(hours=val).total_seconds())
- elif unit == "d":
- timeout = int(timedelta(days=val).total_seconds())
-
- if not timeout:
- raise AirflowException(
- "DataprocClusterScaleOperator "
- " should be expressed in day, hours, minutes or seconds. "
- " i.e. 1d, 4h, 10m, 30s"
- )
-
- return {"seconds": timeout}
-
- def execute(self, context: Context) -> None:
- """Scale, up or down, a cluster on Google Cloud Dataproc."""
- self.log.info("Scaling cluster: %s", self.cluster_name)
-
- scaling_cluster_data = self._build_scale_cluster_data()
- update_mask = ["config.worker_config.num_instances", "config.secondary_worker_config.num_instances"]
-
- hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
- # Save data required to display extra link no matter what the cluster status will be
- DataprocLink.persist(
- context=context,
- task_instance=self,
- url=DATAPROC_CLUSTER_LINK_DEPRECATED,
- resource=self.cluster_name,
- )
- operation = hook.update_cluster(
- project_id=self.project_id,
- region=self.region,
- cluster_name=self.cluster_name,
- cluster=scaling_cluster_data,
- graceful_decommission_timeout=self._graceful_decommission_timeout_object,
- update_mask={"paths": update_mask},
- )
- operation.result()
- self.log.info("Cluster scaling finished")
-
-
class DataprocDeleteClusterOperator(GoogleCloudBaseOperator):
"""
Delete a cluster in a project.
diff --git a/providers/google/src/airflow/providers/google/cloud/operators/mlengine.py b/providers/google/src/airflow/providers/google/cloud/operators/mlengine.py
index 630c8b78ce2..cf432491395 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/mlengine.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/mlengine.py
@@ -20,25 +20,13 @@
from __future__ import annotations
import logging
-import re
-import time
from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
-from googleapiclient.errors import HttpError
-
-from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.google.cloud.hooks.mlengine import MLEngineHook
-from airflow.providers.google.cloud.links.mlengine import (
- MLEngineJobDetailsLink,
- MLEngineJobSListLink,
- MLEngineModelLink,
- MLEngineModelsListLink,
- MLEngineModelVersionDetailsLink,
-)
+from airflow.providers.google.cloud.links.mlengine import MLEngineModelLink
from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
-from airflow.providers.google.cloud.triggers.mlengine import MLEngineStartTrainingJobTrigger
from airflow.providers.google.common.deprecated import deprecated
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
@@ -49,326 +37,6 @@ if TYPE_CHECKING:
log = logging.getLogger(__name__)
-def _normalize_mlengine_job_id(job_id: str) -> str:
- """
- Replace invalid MLEngine job_id characters with '_'.
-
- This also adds a leading 'z' in case job_id starts with an invalid
- character.
-
- :param job_id: A job_id str that may have invalid characters.
- :return: A valid job_id representation.
- """
- # Add a prefix when a job_id starts with a digit or a template
- match = re.search(r"\d|\{{2}", job_id)
- if match and match.start() == 0:
- job = f"z_{job_id}"
- else:
- job = job_id
-
- # Clean up 'bad' characters except templates
- tracker = 0
- cleansed_job_id = ""
- for match in re.finditer(r"\{{2}.+?\}{2}", job):
- cleansed_job_id += re.sub(r"[^0-9a-zA-Z]+", "_", job[tracker : match.start()])
- cleansed_job_id += job[match.start() : match.end()]
- tracker = match.end()
-
- # Clean up last substring or the full string if no templates
- cleansed_job_id += re.sub(r"[^0-9a-zA-Z]+", "_", job[tracker:])
-
- return cleansed_job_id
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="CreateBatchPredictionJobOperator",
- reason="All the functionality of legacy MLEngine and new features are available on the Vertex AI.",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineStartBatchPredictionJobOperator(GoogleCloudBaseOperator):
- """
- Start a Google Cloud ML Engine prediction job.
-
- .. warning::
- This operator is deprecated. Please use
- :class:`airflow.providers.google.cloud.operators.vertex_ai.batch_prediction.CreateBatchPredictionJobOperator`
- instead.
-
- .. seealso::
- For more information on how to use this operator, take a look at the guide:
- :ref:`howto/operator:MLEngineStartBatchPredictionJobOperator`
-
- NOTE: For model origin, users should consider exactly one from the
- three options below:
-
- 1. Populate ``uri`` field only, which should be a GCS location that
- points to a tensorflow savedModel directory.
- 2. Populate ``model_name`` field only, which refers to an existing
- model, and the default version of the model will be used.
- 3. Populate both ``model_name`` and ``version_name`` fields, which
- refers to a specific version of a specific model.
-
- In options 2 and 3, both model and version name should contain the
- minimal identifier. For instance, call::
-
- MLEngineStartBatchPredictionJobOperator(
- ...,
- model_name='my_model',
- version_name='my_version',
- ...)
-
- if the desired model version is
- ``projects/my_project/models/my_model/versions/my_version``.
-
- See https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs
- for further documentation on the parameters.
-
- :param job_id: A unique id for the prediction job on Google Cloud
- ML Engine. (templated)
- :param data_format: The format of the input data.
- It will default to 'DATA_FORMAT_UNSPECIFIED' if is not provided
- or is not one of ["TEXT", "TF_RECORD", "TF_RECORD_GZIP"].
- :param input_paths: A list of GCS paths of input data for batch
- prediction. Accepting wildcard operator ``*``, but only at the end. (templated)
- :param output_path: The GCS path where the prediction results are
- written to. (templated)
- :param region: The Google Compute Engine region to run the
- prediction job in. (templated)
- :param model_name: The Google Cloud ML Engine model to use for prediction.
- If version_name is not provided, the default version of this
- model will be used.
- Should not be None if version_name is provided.
- Should be None if uri is provided. (templated)
- :param version_name: The Google Cloud ML Engine model version to use for
- prediction.
- Should be None if uri is provided. (templated)
- :param uri: The GCS path of the saved model to use for prediction.
- Should be None if model_name is provided.
- It should be a GCS path pointing to a tensorflow SavedModel. (templated)
- :param max_worker_count: The maximum number of workers to be used
- for parallel processing. Defaults to 10 if not specified. Should be a
- string representing the worker count ("10" instead of 10, "50" instead
- of 50, etc.)
- :param runtime_version: The Google Cloud ML Engine runtime version to use
- for batch prediction.
- :param signature_name: The name of the signature defined in the SavedModel
- to use for this job.
- :param project_id: The Google Cloud project name where the prediction job is submitted.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- (templated)
- :param gcp_conn_id: The connection ID used for connection to Google
- Cloud Platform.
- :param labels: a dictionary containing labels for the job; passed to BigQuery
- :param 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).
-
- :raises: ``ValueError``: if a unique model/version origin cannot be
- determined.
- """
-
- template_fields: Sequence[str] = (
- "project_id",
- "job_id",
- "region",
- "input_paths",
- "output_path",
- "model_name",
- "version_name",
- "uri",
- "impersonation_chain",
- )
-
- def __init__(
- self,
- *,
- job_id: str,
- region: str,
- data_format: str,
- input_paths: list[str],
- output_path: str,
- model_name: str | None = None,
- version_name: str | None = None,
- uri: str | None = None,
- max_worker_count: int | None = None,
- runtime_version: str | None = None,
- signature_name: str | None = None,
- project_id: str | None = None,
- gcp_conn_id: str = "google_cloud_default",
- labels: dict[str, str] | None = None,
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
-
- self.project_id = project_id
- self.job_id = job_id
- self.region = region
- self._data_format = data_format
- self.input_paths = input_paths
- self.output_path = output_path
- self.model_name = model_name
- self.version_name = version_name
- self.uri = uri
- self._max_worker_count = max_worker_count
- self._runtime_version = runtime_version
- self._signature_name = signature_name
- self._gcp_conn_id = gcp_conn_id
- self._labels = labels
- self.impersonation_chain = impersonation_chain
-
- def execute(self, context: Context):
- if not self.project_id:
- raise AirflowException("Google Cloud project id is required.")
- if not self.job_id:
- raise AirflowException("An unique job id is required for Google MLEngine prediction job.")
-
- if self.uri:
- if self.model_name or self.version_name:
- raise AirflowException(
- "Ambiguous model origin: Both uri and model/version name are provided."
- )
-
- if self.version_name and not self.model_name:
- raise AirflowException(
- "Missing model: Batch prediction expects a model name when a version name is provided."
- )
-
- if not (self.uri or self.model_name):
- raise AirflowException(
- "Missing model origin: Batch prediction expects a model, "
- "a model & version combination, or a URI to a savedModel."
- )
- job_id = _normalize_mlengine_job_id(self.job_id)
- prediction_request: dict[str, Any] = {
- "jobId": job_id,
- "predictionInput": {
- "dataFormat": self._data_format,
- "inputPaths": self.input_paths,
- "outputPath": self.output_path,
- "region": self.region,
- },
- }
- if self._labels:
- prediction_request["labels"] = self._labels
-
- if self.uri:
- prediction_request["predictionInput"]["uri"] = self.uri
- elif self.model_name:
- origin_name = f"projects/{self.project_id}/models/{self.model_name}"
- if not self.version_name:
- prediction_request["predictionInput"]["modelName"] = origin_name
- else:
- prediction_request["predictionInput"]["versionName"] = (
- origin_name + f"/versions/{self.version_name}"
- )
-
- if self._max_worker_count:
- prediction_request["predictionInput"]["maxWorkerCount"] = self._max_worker_count
-
- if self._runtime_version:
- prediction_request["predictionInput"]["runtimeVersion"] = self._runtime_version
-
- if self._signature_name:
- prediction_request["predictionInput"]["signatureName"] = self._signature_name
-
- hook = MLEngineHook(gcp_conn_id=self._gcp_conn_id, impersonation_chain=self.impersonation_chain)
-
- # Helper method to check if the existing job's prediction input is the
- # same as the request we get here.
- def check_existing_job(existing_job):
- return existing_job.get("predictionInput") == prediction_request["predictionInput"]
-
- finished_prediction_job = hook.create_job(
- project_id=self.project_id, job=prediction_request, use_existing_job_fn=check_existing_job
- )
-
- if finished_prediction_job["state"] != "SUCCEEDED":
- self.log.error("MLEngine batch prediction job failed: %s", finished_prediction_job)
- raise RuntimeError(finished_prediction_job["errorMessage"])
-
- return finished_prediction_job["predictionOutput"]
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="MLEngineCreateModelOperator, MLEngineGetModelOperator",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineManageModelOperator(GoogleCloudBaseOperator):
- """
- Operator for managing a Google Cloud ML Engine model.
-
- .. warning::
- This operator is deprecated. Consider using operators for specific operations:
- MLEngineCreateModelOperator, MLEngineGetModelOperator.
-
- :param model: A dictionary containing the information about the model.
- If the `operation` is `create`, then the `model` parameter should
- contain all the information about this model such as `name`.
-
- If the `operation` is `get`, the `model` parameter
- should contain the `name` of the model.
- :param operation: The operation to perform. Available operations are:
-
- * ``create``: Creates a new model as provided by the `model` parameter.
- * ``get``: Gets a particular model where the name is specified in `model`.
- :param project_id: The Google Cloud project name to which MLEngine model belongs.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- (templated)
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :param 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: Sequence[str] = (
- "project_id",
- "model",
- "impersonation_chain",
- )
-
- def __init__(
- self,
- *,
- model: dict,
- operation: str = "create",
- project_id: str = PROVIDE_PROJECT_ID,
- gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.model = model
- self._operation = operation
- self._gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
-
- def execute(self, context: Context):
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- )
- if self._operation == "create":
- return hook.create_model(project_id=self.project_id, model=self.model)
- elif self._operation == "get":
- return hook.get_model(project_id=self.project_id, model_name=self.model["name"])
- else:
- raise ValueError(f"Unknown operation: {self._operation}")
-
-
@deprecated(
planned_removal_date="November 01, 2025",
use_instead="appropriate VertexAI operator",
@@ -442,1074 +110,3 @@ class MLEngineCreateModelOperator(GoogleCloudBaseOperator):
)
return hook.create_model(project_id=self.project_id, model=self.model)
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="GetModelOperator",
- reason="All the functionality of legacy MLEngine and new features are available on the Vertex AI.",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineGetModelOperator(GoogleCloudBaseOperator):
- """
- Gets a particular model.
-
- .. warning::
- This operator is deprecated. Please use
- :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.GetModelOperator` instead.
-
- .. seealso::
- For more information on how to use this operator, take a look at the guide:
- :ref:`howto/operator:MLEngineGetModelOperator`
-
- The name of model should be specified in `model_name`.
-
- :param model_name: The name of the model.
- :param project_id: The Google Cloud project name to which MLEngine model belongs.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- (templated)
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :param 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: Sequence[str] = (
- "project_id",
- "model_name",
- "impersonation_chain",
- )
- operator_extra_links = (MLEngineModelLink(),)
-
- def __init__(
- self,
- *,
- model_name: str,
- project_id: str = PROVIDE_PROJECT_ID,
- gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.model_name = model_name
- self._gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
-
- def execute(self, context: Context):
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- )
- project_id = self.project_id or hook.project_id
- if project_id:
- MLEngineModelLink.persist(
- context=context,
- task_instance=self,
- project_id=project_id,
- model_id=self.model_name,
- )
-
- return hook.get_model(project_id=self.project_id, model_name=self.model_name)
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="DeleteModelOperator",
- reason="All the functionality of legacy MLEngine and new features are available on the Vertex AI.",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineDeleteModelOperator(GoogleCloudBaseOperator):
- """
- Deletes a model.
-
- .. warning::
- This operator is deprecated. Please use
- :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelOperator` instead.
-
- .. seealso::
- For more information on how to use this operator, take a look at the guide:
- :ref:`howto/operator:MLEngineDeleteModelOperator`
-
- The model should be provided by the `model_name` parameter.
-
- :param model_name: The name of the model.
- :param delete_contents: (Optional) Whether to force the deletion even if the models is not empty.
- Will delete all version (if any) in the dataset if set to True.
- The default value is False.
- :param project_id: The Google Cloud project name to which MLEngine model belongs.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- (templated)
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :param 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: Sequence[str] = (
- "project_id",
- "model_name",
- "impersonation_chain",
- )
- operator_extra_links = (MLEngineModelsListLink(),)
-
- def __init__(
- self,
- *,
- model_name: str,
- delete_contents: bool = False,
- project_id: str = PROVIDE_PROJECT_ID,
- gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.model_name = model_name
- self._delete_contents = delete_contents
- self._gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
-
- def execute(self, context: Context):
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- )
-
- project_id = self.project_id or hook.project_id
- if project_id:
- MLEngineModelsListLink.persist(
- context=context,
- task_instance=self,
- project_id=project_id,
- )
-
- return hook.delete_model(
- project_id=self.project_id, model_name=self.model_name, delete_contents=self._delete_contents
- )
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="MLEngineCreateVersion, MLEngineSetDefaultVersion, MLEngineListVersions, "
- "MLEngineDeleteVersion",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineManageVersionOperator(GoogleCloudBaseOperator):
- """
- Operator for managing a Google Cloud ML Engine version.
-
- .. warning::
- This operator is deprecated. Consider using operators for specific operations:
- MLEngineCreateVersionOperator, MLEngineSetDefaultVersionOperator,
- MLEngineListVersionsOperator, MLEngineDeleteVersionOperator.
-
- :param model_name: The name of the Google Cloud ML Engine model that the version
- belongs to. (templated)
- :param version_name: A name to use for the version being operated upon.
- If not None and the `version` argument is None or does not have a value for
- the `name` key, then this will be populated in the payload for the
- `name` key. (templated)
- :param version: A dictionary containing the information about the version.
- If the `operation` is `create`, `version` should contain all the
- information about this version such as name, and deploymentUrl.
- If the `operation` is `get` or `delete`, the `version` parameter
- should contain the `name` of the version.
- If it is None, the only `operation` possible would be `list`. (templated)
- :param operation: The operation to perform. Available operations are:
-
- * ``create``: Creates a new version in the model specified by `model_name`,
- in which case the `version` parameter should contain all the
- information to create that version
- (e.g. `name`, `deploymentUrl`).
-
- * ``set_defaults``: Sets a version in the model specified by `model_name` to be the default.
- The name of the version should be specified in the `version`
- parameter.
-
- * ``list``: Lists all available versions of the model specified
- by `model_name`.
-
- * ``delete``: Deletes the version specified in `version` parameter from the
- model specified by `model_name`).
- The name of the version should be specified in the `version`
- parameter.
- :param project_id: The Google Cloud project name to which MLEngine model belongs.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- (templated)
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :param 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: Sequence[str] = (
- "project_id",
- "model_name",
- "version_name",
- "version",
- "impersonation_chain",
- )
-
- def __init__(
- self,
- *,
- model_name: str,
- version_name: str | None = None,
- version: dict | None = None,
- operation: str = "create",
- project_id: str = PROVIDE_PROJECT_ID,
- gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.model_name = model_name
- self.version_name = version_name
- self.version = version or {}
- self._operation = operation
- self._gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
-
- def execute(self, context: Context):
- if "name" not in self.version:
- self.version["name"] = self.version_name
-
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- )
-
- if self._operation == "create":
- if not self.version:
- raise ValueError(f"version attribute of {self.__class__.__name__} could not be empty")
- return hook.create_version(
- project_id=self.project_id, model_name=self.model_name, version_spec=self.version
- )
- elif self._operation == "set_default":
- return hook.set_default_version(
- project_id=self.project_id, model_name=self.model_name, version_name=self.version["name"]
- )
- elif self._operation == "list":
- return hook.list_versions(project_id=self.project_id, model_name=self.model_name)
- elif self._operation == "delete":
- return hook.delete_version(
- project_id=self.project_id, model_name=self.model_name, version_name=self.version["name"]
- )
- else:
- raise ValueError(f"Unknown operation: {self._operation}")
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="parent_model parameter for VertexAI operators",
- reason="All the functionality of legacy MLEngine and new features are available on the Vertex AI.",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineCreateVersionOperator(GoogleCloudBaseOperator):
- """
- Creates a new version in the model.
-
- .. warning::
- This operator is deprecated. Please use parent_model parameter of VertexAI
- operators instead.
-
- .. seealso::
- For more information on how to use this operator, take a look at the guide:
- :ref:`howto/operator:MLEngineCreateVersionOperator`
-
- Model should be specified by `model_name`, in which case the `version` parameter should contain all the
- information to create that version
-
- :param model_name: The name of the Google Cloud ML Engine model that the version belongs to. (templated)
- :param version: A dictionary containing the information about the version. (templated)
- :param project_id: The Google Cloud project name to which MLEngine model belongs.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- (templated)
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :param 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: Sequence[str] = (
- "project_id",
- "model_name",
- "version",
- "impersonation_chain",
- )
- operator_extra_links = (MLEngineModelVersionDetailsLink(),)
-
- def __init__(
- self,
- *,
- model_name: str,
- version: dict,
- project_id: str = PROVIDE_PROJECT_ID,
- gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.model_name = model_name
- self.version = version
- self._gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
-
- def _validate_inputs(self):
- if not self.model_name:
- raise AirflowException("The model_name parameter could not be empty.")
-
- if not self.version:
- raise AirflowException("The version parameter could not be empty.")
-
- def execute(self, context: Context):
- self._validate_inputs()
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- )
-
- project_id = self.project_id or hook.project_id
- if project_id:
- MLEngineModelVersionDetailsLink.persist(
- context=context,
- task_instance=self,
- project_id=project_id,
- model_id=self.model_name,
- version_id=self.version["name"],
- )
-
- return hook.create_version(
- project_id=self.project_id, model_name=self.model_name, version_spec=self.version
- )
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="SetDefaultVersionOnModelOperator",
- reason="All the functionality of legacy MLEngine and new features are available on the Vertex AI.",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineSetDefaultVersionOperator(GoogleCloudBaseOperator):
- """
- Sets a version in the model.
-
- .. warning::
- This operator is deprecated. Please use
- :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.SetDefaultVersionOnModelOperator`
- instead.
-
- .. seealso::
- For more information on how to use this operator, take a look at the guide:
- :ref:`howto/operator:MLEngineSetDefaultVersionOperator`
-
- The model should be specified by `model_name` to be the default. The name of the version should be
- specified in the `version_name` parameter.
-
- :param model_name: The name of the Google Cloud ML Engine model that the version belongs to. (templated)
- :param version_name: A name to use for the version being operated upon. (templated)
- :param project_id: The Google Cloud project name to which MLEngine model belongs.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- (templated)
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :param 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: Sequence[str] = (
- "project_id",
- "model_name",
- "version_name",
- "impersonation_chain",
- )
- operator_extra_links = (MLEngineModelVersionDetailsLink(),)
-
- def __init__(
- self,
- *,
- model_name: str,
- version_name: str,
- project_id: str = PROVIDE_PROJECT_ID,
- gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.model_name = model_name
- self.version_name = version_name
- self._gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
-
- def _validate_inputs(self):
- if not self.model_name:
- raise AirflowException("The model_name parameter could not be empty.")
-
- if not self.version_name:
- raise AirflowException("The version_name parameter could not be empty.")
-
- def execute(self, context: Context):
- self._validate_inputs()
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- )
-
- project_id = self.project_id or hook.project_id
- if project_id:
- MLEngineModelVersionDetailsLink.persist(
- context=context,
- task_instance=self,
- project_id=project_id,
- model_id=self.model_name,
- version_id=self.version_name,
- )
-
- return hook.set_default_version(
- project_id=self.project_id, model_name=self.model_name, version_name=self.version_name
- )
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="istModelVersionsOperator",
- reason="All the functionality of legacy MLEngine and new features are available on the Vertex AI.",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineListVersionsOperator(GoogleCloudBaseOperator):
- """
- Lists all available versions of the model.
-
- .. warning::
- This operator is deprecated. Please use
- :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.ListModelVersionsOperator`
- instead.
-
- .. seealso::
- For more information on how to use this operator, take a look at the guide:
- :ref:`howto/operator:MLEngineListVersionsOperator`
-
- The model should be specified by `model_name`.
-
- :param model_name: The name of the Google Cloud ML Engine model that the version
- belongs to. (templated)
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :param project_id: The Google Cloud project name to which MLEngine model belongs.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- (templated)
- :param 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: Sequence[str] = (
- "project_id",
- "model_name",
- "impersonation_chain",
- )
- operator_extra_links = (MLEngineModelLink(),)
-
- def __init__(
- self,
- *,
- model_name: str,
- project_id: str = PROVIDE_PROJECT_ID,
- gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.model_name = model_name
- self._gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
-
- def _validate_inputs(self):
- if not self.model_name:
- raise AirflowException("The model_name parameter could not be empty.")
-
- def execute(self, context: Context):
- self._validate_inputs()
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- )
-
- project_id = self.project_id or hook.project_id
- if project_id:
- MLEngineModelLink.persist(
- context=context,
- task_instance=self,
- project_id=project_id,
- model_id=self.model_name,
- )
-
- return hook.list_versions(
- project_id=self.project_id,
- model_name=self.model_name,
- )
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="DeleteModelVersionOperator",
- reason="All the functionality of legacy MLEngine and new features are available on the Vertex AI.",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineDeleteVersionOperator(GoogleCloudBaseOperator):
- """
- Deletes the version from the model.
-
- .. warning::
- This operator is deprecated. Please use
- :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelVersionOperator`
- instead.
-
- .. seealso::
- For more information on how to use this operator, take a look at the guide:
- :ref:`howto/operator:MLEngineDeleteVersionOperator`
-
- The name of the version should be specified in `version_name` parameter from the model specified
- by `model_name`.
-
- :param model_name: The name of the Google Cloud ML Engine model that the version
- belongs to. (templated)
- :param version_name: A name to use for the version being operated upon. (templated)
- :param project_id: The Google Cloud project name to which MLEngine
- model belongs.
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :param 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: Sequence[str] = (
- "project_id",
- "model_name",
- "version_name",
- "impersonation_chain",
- )
- operator_extra_links = (MLEngineModelLink(),)
-
- def __init__(
- self,
- *,
- model_name: str,
- version_name: str,
- project_id: str = PROVIDE_PROJECT_ID,
- gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.model_name = model_name
- self.version_name = version_name
- self._gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
-
- def _validate_inputs(self):
- if not self.model_name:
- raise AirflowException("The model_name parameter could not be empty.")
-
- if not self.version_name:
- raise AirflowException("The version_name parameter could not be empty.")
-
- def execute(self, context: Context):
- self._validate_inputs()
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- )
-
- project_id = self.project_id or hook.project_id
- if project_id:
- MLEngineModelLink.persist(
- context=context,
- task_instance=self,
- project_id=project_id,
- model_id=self.model_name,
- )
-
- return hook.delete_version(
- project_id=self.project_id, model_name=self.model_name, version_name=self.version_name
- )
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="CreateCustomPythonPackageTrainingJobOperator",
- reason="All the functionality of legacy MLEngine and new features are available on the Vertex AI.",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineStartTrainingJobOperator(GoogleCloudBaseOperator):
- """
- Operator for launching a MLEngine training job.
-
- .. warning::
- This operator is deprecated. Please use
- :class:`airflow.providers.google.cloud.operators.vertex_ai.custom_job.CreateCustomPythonPackageTrainingJobOperator`
- instead.
-
- .. seealso::
- For more information on how to use this operator, take a look at the guide:
- :ref:`howto/operator:MLEngineStartTrainingJobOperator`
-
- For more information about used parameters, check:
- https://cloud.google.com/sdk/gcloud/reference/ml-engine/jobs/submit/training
-
- :param job_id: A unique templated id for the submitted Google MLEngine
- training job. (templated)
- :param region: The Google Compute Engine region to run the MLEngine training
- job in (templated).
- :param package_uris: A list of Python package locations for the training
- job, which should include the main training program and any additional
- dependencies. This is mutually exclusive with a custom image specified
- via master_config. (templated)
- :param training_python_module: The name of the Python module to run within
- the training job after installing the packages. This is mutually
- exclusive with a custom image specified via master_config. (templated)
- :param training_args: A list of command-line arguments to pass to the
- training program. (templated)
- :param scale_tier: Resource tier for MLEngine training job. (templated)
- :param master_type: The type of virtual machine to use for the master
- worker. It must be set whenever scale_tier is CUSTOM. (templated)
- :param master_config: The configuration for the master worker. If this is
- provided, master_type must be set as well. If a custom image is
- specified, this is mutually exclusive with package_uris and
- training_python_module. (templated)
- :param runtime_version: The Google Cloud ML runtime version to use for
- training. (templated)
- :param python_version: The version of Python used in training. (templated)
- :param job_dir: A Google Cloud Storage path in which to store training
- outputs and other data needed for training. (templated)
- :param service_account: Optional service account to use when running the training application.
- (templated)
- The specified service account must have the `iam.serviceAccounts.actAs` role. The
- Google-managed Cloud ML Engine service account must have the `iam.serviceAccountAdmin` role
- for the specified service account.
- If set to None or missing, the Google-managed Cloud ML Engine service account will be used.
- :param project_id: The Google Cloud project name within which MLEngine training job should run.
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :param mode: Can be one of 'DRY_RUN'/'CLOUD'. In 'DRY_RUN' mode, no real
- training job will be launched, but the MLEngine training job request
- will be printed out. In 'CLOUD' mode, a real MLEngine training job
- creation request will be issued.
- :param labels: a dictionary containing labels for the job; passed to BigQuery
- :param hyperparameters: Optional HyperparameterSpec dictionary for hyperparameter tuning.
- For further reference, check:
- https://cloud.google.com/ai-platform/training/docs/reference/rest/v1/projects.jobs#HyperparameterSpec
- :param 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).
- :param cancel_on_kill: Flag which indicates whether cancel the hook's job or not, when on_kill is called
- :param deferrable: Run operator in the deferrable mode
- """
-
- template_fields: Sequence[str] = (
- "project_id",
- "job_id",
- "region",
- "package_uris",
- "training_python_module",
- "training_args",
- "scale_tier",
- "master_type",
- "master_config",
- "runtime_version",
- "python_version",
- "job_dir",
- "service_account",
- "hyperparameters",
- "impersonation_chain",
- )
- operator_extra_links = (MLEngineJobDetailsLink(),)
-
- def __init__(
- self,
- *,
- job_id: str,
- region: str,
- project_id: str,
- package_uris: list[str] | None = None,
- training_python_module: str | None = None,
- training_args: list[str] | None = None,
- scale_tier: str | None = None,
- master_type: str | None = None,
- master_config: dict | None = None,
- runtime_version: str | None = None,
- python_version: str | None = None,
- job_dir: str | None = None,
- service_account: str | None = None,
- gcp_conn_id: str = "google_cloud_default",
- mode: str = "PRODUCTION",
- labels: dict[str, str] | None = None,
- impersonation_chain: str | Sequence[str] | None = None,
- hyperparameters: dict | None = None,
- deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
- cancel_on_kill: bool = True,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.job_id = job_id
- self.region = region
- self.package_uris = package_uris
- self.training_python_module = training_python_module
- self.training_args = training_args
- self.scale_tier = scale_tier
- self.master_type = master_type
- self.master_config = master_config
- self.runtime_version = runtime_version
- self.python_version = python_version
- self.job_dir = job_dir
- self.service_account = service_account
- self._gcp_conn_id = gcp_conn_id
- self._mode = mode
- self._labels = labels
- self.hyperparameters = hyperparameters
- self.impersonation_chain = impersonation_chain
- self.deferrable = deferrable
- self.cancel_on_kill = cancel_on_kill
-
- def _handle_job_error(self, finished_training_job) -> None:
- if finished_training_job["state"] != "SUCCEEDED":
- self.log.error("MLEngine training job failed: %s", finished_training_job)
- raise RuntimeError(finished_training_job["errorMessage"])
-
- def execute(self, context: Context):
- custom = self.scale_tier is not None and self.scale_tier.upper() == "CUSTOM"
- custom_image = (
- custom and self.master_config is not None and self.master_config.get("imageUri", None) is not None
- )
-
- if not self.project_id:
- raise AirflowException("Google Cloud project id is required.")
- if not self.job_id:
- raise AirflowException("An unique job id is required for Google MLEngine training job.")
- if not self.region:
- raise AirflowException("Google Compute Engine region is required.")
- if custom and not self.master_type:
- raise AirflowException("master_type must be set when scale_tier is CUSTOM")
- if self.master_config and not self.master_type:
- raise AirflowException("master_type must be set when master_config is provided")
- if not (self.package_uris and self.training_python_module) and not custom_image:
- raise AirflowException(
- "Either a Python package with a Python module or a custom Docker image should be provided."
- )
- if (self.package_uris or self.training_python_module) and custom_image:
- raise AirflowException(
- "Either a Python package with a Python module or "
- "a custom Docker image should be provided but not both."
- )
- job_id = _normalize_mlengine_job_id(self.job_id)
- self.job_id = job_id
- training_request: dict[str, Any] = {
- "jobId": self.job_id,
- "trainingInput": {
- "scaleTier": self.scale_tier,
- "region": self.region,
- },
- }
- if self.package_uris:
- training_request["trainingInput"]["packageUris"] = self.package_uris
-
- if self.training_python_module:
- training_request["trainingInput"]["pythonModule"] = self.training_python_module
-
- if self.training_args:
- training_request["trainingInput"]["args"] = self.training_args
-
- if self.master_type:
- training_request["trainingInput"]["masterType"] = self.master_type
-
- if self.master_config:
- training_request["trainingInput"]["masterConfig"] = self.master_config
-
- if self.runtime_version:
- training_request["trainingInput"]["runtimeVersion"] = self.runtime_version
-
- if self.python_version:
- training_request["trainingInput"]["pythonVersion"] = self.python_version
-
- if self.job_dir:
- training_request["trainingInput"]["jobDir"] = self.job_dir
-
- if self.service_account:
- training_request["trainingInput"]["serviceAccount"] = self.service_account
-
- if self.hyperparameters:
- training_request["trainingInput"]["hyperparameters"] = self.hyperparameters
-
- if self._labels:
- training_request["labels"] = self._labels
-
- if self._mode == "DRY_RUN":
- self.log.info("In dry_run mode.")
- self.log.info("MLEngine Training job request is: %s", training_request)
- return
-
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- )
- self.hook = hook
-
- try:
- self.log.info("Executing: %s'", training_request)
- self.job_id = self.hook.create_job_without_waiting_result(
- project_id=self.project_id,
- body=training_request,
- )
- except HttpError as e:
- if e.resp.status == 409:
- # If the job already exists retrieve it
- self.hook.get_job(project_id=self.project_id, job_id=self.job_id)
- if self.project_id:
- MLEngineJobDetailsLink.persist(
- context=context,
- task_instance=self,
- project_id=self.project_id,
- job_id=self.job_id,
- )
- self.log.error(
- "Failed to create new job with given name since it already exists. "
- "The existing one will be used."
- )
- else:
- raise e
-
- context["ti"].xcom_push(key="job_id", value=self.job_id)
- if self.deferrable:
- self.defer(
- timeout=self.execution_timeout,
- trigger=MLEngineStartTrainingJobTrigger(
- conn_id=self._gcp_conn_id,
- job_id=self.job_id,
- project_id=self.project_id,
- region=self.region,
- runtime_version=self.runtime_version,
- python_version=self.python_version,
- job_dir=self.job_dir,
- package_uris=self.package_uris,
- training_python_module=self.training_python_module,
- training_args=self.training_args,
- labels=self._labels,
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- ),
- method_name="execute_complete",
- )
- else:
- finished_training_job = self._wait_for_job_done(self.project_id, self.job_id)
- self._handle_job_error(finished_training_job)
- gcp_metadata = {
- "job_id": self.job_id,
- "project_id": self.project_id,
- }
- context["task_instance"].xcom_push("gcp_metadata", gcp_metadata)
-
- project_id = self.project_id or hook.project_id
- if project_id:
- MLEngineJobDetailsLink.persist(
- context=context,
- task_instance=self,
- project_id=project_id,
- job_id=job_id,
- )
-
- def _wait_for_job_done(self, project_id: str, job_id: str, interval: int = 30):
- """
- Wait for the Job to reach a terminal state.
-
- This method will periodically check the job state until the job reach
- a terminal state.
-
- :param project_id: The project in which the Job is located. If set to None or missing, the default
- project_id from the Google Cloud connection is used. (templated)
- :param job_id: A unique id for the Google MLEngine job. (templated)
- :param interval: Time expressed in seconds after which the job status is checked again. (templated)
- :raises: googleapiclient.errors.HttpError
- """
- self.log.info("Waiting for job. job_id=%s", job_id)
-
- if interval <= 0:
- raise ValueError("Interval must be > 0")
- while True:
- job = self.hook.get_job(project_id, job_id)
- if job["state"] in ["SUCCEEDED", "FAILED", "CANCELLED"]:
- return job
- time.sleep(interval)
-
- def execute_complete(self, context: Context, event: dict[str, Any]):
- """
- Act as a callback for when the trigger fires - returns immediately.
-
- Relies on trigger to throw an exception, otherwise it assumes execution was successful.
- """
- if event["status"] == "error":
- raise AirflowException(event["message"])
- self.log.info(
- "%s completed with response %s ",
- self.task_id,
- event["message"],
- )
- if self.project_id:
- MLEngineJobDetailsLink.persist(
- context=context,
- task_instance=self,
- project_id=self.project_id,
- job_id=self.job_id,
- )
-
- def on_kill(self) -> None:
- if self.job_id and self.cancel_on_kill:
- self.hook.cancel_job(job_id=self.job_id, project_id=self.project_id) # type: ignore[union-attr]
- else:
- self.log.info("Skipping to cancel job: %s:%s.%s", self.project_id, self.job_id)
-
-
-@deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="CancelCustomTrainingJobOperator",
- reason="All the functionality of legacy MLEngine and new features are available on the Vertex AI.",
- category=AirflowProviderDeprecationWarning,
-)
-class MLEngineTrainingCancelJobOperator(GoogleCloudBaseOperator):
- """
- Operator for cleaning up failed MLEngine training job.
-
- .. warning::
- This operator is deprecated. Please use
- :class:`airflow.providers.google.cloud.operators.vertex_ai.custom_job.CancelCustomTrainingJobOperator`
- instead.
-
- :param job_id: A unique templated id for the submitted Google MLEngine
- training job. (templated)
- :param project_id: The Google Cloud project name within which MLEngine training job should run.
- If set to None or missing, the default project_id from the Google Cloud connection is used.
- (templated)
- :param gcp_conn_id: The connection ID to use when fetching connection info.
- :param 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: Sequence[str] = (
- "project_id",
- "job_id",
- "impersonation_chain",
- )
- operator_extra_links = (MLEngineJobSListLink(),)
-
- def __init__(
- self,
- *,
- job_id: str,
- project_id: str = PROVIDE_PROJECT_ID,
- gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.project_id = project_id
- self.job_id = job_id
- self._gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
-
- @property
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="project_id",
- category=AirflowProviderDeprecationWarning,
- )
- def _project_id(self):
- """Alias for ``project_id``, used for compatibility (deprecated)."""
- return self.project_id
-
- @property
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="job_id",
- category=AirflowProviderDeprecationWarning,
- )
- def _job_id(self):
- """Alias for ``job_id``, used for compatibility (deprecated)."""
- return self.job_id
-
- @property
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="impersonation_chain",
- category=AirflowProviderDeprecationWarning,
- )
- def _impersonation_chain(self):
- """Alias for ``impersonation_chain``, used for compatibility (deprecated)."""
- return self.impersonation_chain
-
- def execute(self, context: Context):
- if not self.project_id:
- raise AirflowException("Google Cloud project id is required.")
- hook = MLEngineHook(
- gcp_conn_id=self._gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
- )
-
- project_id = self.project_id or hook.project_id
- if project_id:
- MLEngineJobSListLink.persist(
- context=context,
- task_instance=self,
- project_id=project_id,
- )
-
- hook.cancel_job(project_id=self.project_id, job_id=_normalize_mlengine_job_id(self.job_id))
diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py
index d92b2cef585..363b04d5d9b 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py
@@ -29,7 +29,6 @@ from google.cloud.aiplatform import datasets
from google.cloud.aiplatform.models import Model
from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
-from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.google.cloud.hooks.vertex_ai.auto_ml import AutoMLHook
from airflow.providers.google.cloud.links.vertex_ai import (
VertexAIModelLink,
@@ -37,7 +36,6 @@ from airflow.providers.google.cloud.links.vertex_ai import (
VertexAITrainingPipelinesLink,
)
from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
-from airflow.providers.google.common.deprecated import deprecated
if TYPE_CHECKING:
from google.api_core.retry import Retry
@@ -575,16 +573,6 @@ class DeleteAutoMLTrainingJobOperator(GoogleCloudBaseOperator):
self.gcp_conn_id = gcp_conn_id
self.impersonation_chain = impersonation_chain
- @property
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="training_pipeline_id",
- category=AirflowProviderDeprecationWarning,
- )
- def training_pipeline(self):
- """Alias for ``training_pipeline_id``, used for compatibility (deprecated)."""
- return self.training_pipeline_id
-
def execute(self, context: Context):
hook = AutoMLHook(
gcp_conn_id=self.gcp_conn_id,
diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py
index fa901b4cae6..4fdac797dca 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py
@@ -30,7 +30,7 @@ from google.cloud.aiplatform_v1.types.dataset import Dataset
from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
from airflow.providers.google.cloud.hooks.vertex_ai.custom_job import CustomJobHook
from airflow.providers.google.cloud.links.vertex_ai import (
VertexAIModelLink,
@@ -43,7 +43,6 @@ from airflow.providers.google.cloud.triggers.vertex_ai import (
CustomPythonPackageTrainingJobTrigger,
CustomTrainingJobTrigger,
)
-from airflow.providers.google.common.deprecated import deprecated
if TYPE_CHECKING:
from google.api_core.retry import Retry
@@ -1634,26 +1633,6 @@ class DeleteCustomTrainingJobOperator(GoogleCloudBaseOperator):
self.gcp_conn_id = gcp_conn_id
self.impersonation_chain = impersonation_chain
- @property
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="training_pipeline_id",
- category=AirflowProviderDeprecationWarning,
- )
- def training_pipeline(self):
- """Alias for ``training_pipeline_id``, used for compatibility (deprecated)."""
- return self.training_pipeline_id
-
- @property
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="custom_job_id",
- category=AirflowProviderDeprecationWarning,
- )
- def custom_job(self):
- """Alias for ``custom_job_id``, used for compatibility (deprecated)."""
- return self.custom_job_id
-
def execute(self, context: Context):
hook = CustomJobHook(
gcp_conn_id=self.gcp_conn_id,
diff --git a/providers/google/src/airflow/providers/google/cloud/utils/mlengine_operator_utils.py b/providers/google/src/airflow/providers/google/cloud/utils/mlengine_operator_utils.py
deleted file mode 100644
index 412b699421a..00000000000
--- a/providers/google/src/airflow/providers/google/cloud/utils/mlengine_operator_utils.py
+++ /dev/null
@@ -1,273 +0,0 @@
-# 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 helper functions for MLEngine operators."""
-
-from __future__ import annotations
-
-import base64
-import json
-import os
-import re
-from collections.abc import Iterable
-from typing import TYPE_CHECKING, Callable, TypeVar
-from urllib.parse import urlsplit
-
-import dill
-
-from airflow.exceptions import AirflowException
-from airflow.providers.apache.beam.hooks.beam import BeamRunnerType
-from airflow.providers.apache.beam.operators.beam import BeamRunPythonPipelineOperator
-from airflow.providers.common.compat.standard.operators import PythonOperator
-from airflow.providers.google.cloud.hooks.gcs import GCSHook
-from airflow.providers.google.cloud.operators.mlengine import MLEngineStartBatchPredictionJobOperator
-
-if TYPE_CHECKING:
- from airflow import DAG
-
-T = TypeVar("T", bound=Callable)
-
-
-def create_evaluate_ops(
- task_prefix: str,
- data_format: str,
- input_paths: list[str],
- prediction_path: str,
- metric_fn_and_keys: tuple[T, Iterable[str]],
- validate_fn: T,
- batch_prediction_job_id: str | None = None,
- region: str | None = None,
- project_id: str | None = None,
- dataflow_options: dict | None = None,
- model_uri: str | None = None,
- model_name: str | None = None,
- version_name: str | None = None,
- dag: DAG | None = None,
- py_interpreter="python3",
-) -> tuple[MLEngineStartBatchPredictionJobOperator, BeamRunPythonPipelineOperator, PythonOperator]:
- r"""
- Create Operators needed for model evaluation and returns.
-
- This function is deprecated. All the functionality of legacy MLEngine and new features are available
- on the Vertex AI platform.
-
- To create and view Model Evaluation, please check the documentation:
- https://cloud.google.com/vertex-ai/docs/evaluation/using-model-evaluation#create_an_evaluation.
-
- It gets prediction over inputs via Cloud ML Engine BatchPrediction API by
- calling MLEngineBatchPredictionOperator, then summarize and validate
- the result via Cloud Dataflow using DataFlowPythonOperator.
-
- For details and pricing about Batch prediction, please refer to the website
- https://cloud.google.com/ml-engine/docs/how-tos/batch-predict
- and for Cloud Dataflow, https://cloud.google.com/dataflow/docs/
-
- It returns three chained operators for prediction, summary, and validation,
- named as ``<prefix>-prediction``, ``<prefix>-summary``, and ``<prefix>-validation``,
- respectively.
- (``<prefix>`` should contain only alphanumeric characters or hyphen.)
-
- The upstream and downstream can be set accordingly like:
-
- .. code-block:: python
-
- pred, _, val = create_evaluate_ops(...)
- pred.set_upstream(upstream_op)
- ...
- downstream_op.set_upstream(val)
-
- Callers will provide two python callables, metric_fn and validate_fn, in
- order to customize the evaluation behavior as they wish.
-
- - metric_fn receives a dictionary per instance derived from json in the
- batch prediction result. The keys might vary depending on the model.
- It should return a tuple of metrics.
- - validation_fn receives a dictionary of the averaged metrics that metric_fn
- generated over all instances.
- The key/value of the dictionary matches to what's given by
- metric_fn_and_keys arg.
- The dictionary contains an additional metric, 'count' to represent the
- total number of instances received for evaluation.
- The function would raise an exception to mark the task as failed, in a
- case the validation result is not okay to proceed (i.e. to set the trained
- version as default).
-
- Typical examples are like this:
-
- .. code-block:: python
-
- def get_metric_fn_and_keys():
- import math # imports should be outside of the metric_fn below.
-
- def error_and_squared_error(inst):
- label = float(inst["input_label"])
- classes = float(inst["classes"]) # 0 or 1
- err = abs(classes - label)
- squared_err = math.pow(classes - label, 2)
- return (err, squared_err) # returns a tuple.
-
- return error_and_squared_error, ["err", "mse"] # key order must match.
-
-
- def validate_err_and_count(summary):
- if summary["err"] > 0.2:
- raise ValueError("Too high err>0.2; summary=%s" % summary)
- if summary["mse"] > 0.05:
- raise ValueError("Too high mse>0.05; summary=%s" % summary)
- if summary["count"] < 1000:
- raise ValueError("Too few instances<1000; summary=%s" % summary)
- return summary
-
- For the details on the other BatchPrediction-related arguments (project_id,
- job_id, region, data_format, input_paths, prediction_path, model_uri),
- please refer to MLEngineBatchPredictionOperator too.
-
- :param task_prefix: a prefix for the tasks. Only alphanumeric characters and
- hyphen are allowed (no underscores), since this will be used as dataflow
- job name, which doesn't allow other characters.
-
- :param data_format: either of 'TEXT', 'TF_RECORD', 'TF_RECORD_GZIP'
-
- :param input_paths: a list of input paths to be sent to BatchPrediction.
-
- :param prediction_path: GCS path to put the prediction results in.
-
- :param metric_fn_and_keys: a tuple of metric_fn and metric_keys:
-
- - metric_fn is a function that accepts a dictionary (for an instance),
- and returns a tuple of metric(s) that it calculates.
-
- - metric_keys is a list of strings to denote the key of each metric.
-
- :param validate_fn: a function to validate whether the averaged metric(s) is
- good enough to push the model.
-
- :param batch_prediction_job_id: the id to use for the Cloud ML Batch
- prediction job. Passed directly to the MLEngineBatchPredictionOperator as
- the job_id argument.
-
- :param project_id: the Google Cloud project id in which to execute
- Cloud ML Batch Prediction and Dataflow jobs. If None, then the `dag`'s
- `default_args['project_id']` will be used.
-
- :param region: the Google Cloud region in which to execute Cloud ML
- Batch Prediction and Dataflow jobs. If None, then the `dag`'s
- `default_args['region']` will be used.
-
- :param dataflow_options: options to run Dataflow jobs. If None, then the
- `dag`'s `default_args['dataflow_default_options']` will be used.
-
- :param model_uri: GCS path of the model exported by Tensorflow using
- ``tensorflow.estimator.export_savedmodel()``. It cannot be used with
- model_name or version_name below. See MLEngineBatchPredictionOperator for
- more detail.
-
- :param model_name: Used to indicate a model to use for prediction. Can be
- used in combination with version_name, but cannot be used together with
- model_uri. See MLEngineBatchPredictionOperator for more detail. If None,
- then the `dag`'s `default_args['model_name']` will be used.
-
- :param version_name: Used to indicate a model version to use for prediction,
- in combination with model_name. Cannot be used together with model_uri.
- See MLEngineBatchPredictionOperator for more detail. If None, then the
- `dag`'s `default_args['version_name']` will be used.
-
- :param dag: The `DAG` to use for all Operators.
-
- :param py_interpreter: Python version of the beam pipeline.
- If None, this defaults to the python3.
- To track python versions supported by beam and related
- issues check: https://issues.apache.org/jira/browse/BEAM-1251
-
- :returns: a tuple of three operators, (prediction, summary, validation)
- PythonOperator)
- """
- batch_prediction_job_id = batch_prediction_job_id or ""
- dataflow_options = dataflow_options or {}
- region = region or ""
-
- # Verify that task_prefix doesn't have any special characters except hyphen
- # '-', which is the only allowed non-alphanumeric character by Dataflow.
- if not re.fullmatch(r"[a-zA-Z][-A-Za-z0-9]*", task_prefix):
- raise AirflowException(
- "Malformed task_id for DataFlowPythonOperator (only alphanumeric "
- "and hyphens are allowed but got: " + task_prefix
- )
-
- metric_fn, metric_keys = metric_fn_and_keys
- if not callable(metric_fn):
- raise AirflowException("`metric_fn` param must be callable.")
- if not callable(validate_fn):
- raise AirflowException("`validate_fn` param must be callable.")
-
- if dag is not None and dag.default_args is not None:
- default_args = dag.default_args
- project_id = project_id or default_args.get("project_id")
- region = region or default_args["region"]
- model_name = model_name or default_args.get("model_name")
- version_name = version_name or default_args.get("version_name")
- dataflow_options = dataflow_options or default_args.get("dataflow_default_options")
-
- evaluate_prediction = MLEngineStartBatchPredictionJobOperator(
- task_id=(task_prefix + "-prediction"),
- project_id=project_id,
- job_id=batch_prediction_job_id,
- region=region,
- data_format=data_format,
- input_paths=input_paths,
- output_path=prediction_path,
- uri=model_uri,
- model_name=model_name,
- version_name=version_name,
- dag=dag,
- )
-
- metric_fn_encoded = base64.b64encode(dill.dumps(metric_fn, recurse=True)).decode()
- evaluate_summary = BeamRunPythonPipelineOperator(
- task_id=(task_prefix + "-summary"),
- runner=BeamRunnerType.DataflowRunner,
- py_file=os.path.join(os.path.dirname(__file__), "mlengine_prediction_summary.py"),
- default_pipeline_options=dataflow_options,
- pipeline_options={
- "prediction_path": prediction_path,
- "metric_fn_encoded": metric_fn_encoded,
- "metric_keys": ",".join(metric_keys),
- },
- py_interpreter=py_interpreter,
- py_requirements=["apache-beam[gcp]>=2.46.0"],
- dag=dag,
- )
- evaluate_summary.set_upstream(evaluate_prediction)
-
- def apply_validate_fn(*args, templates_dict, **kwargs):
- prediction_path = templates_dict["prediction_path"]
- scheme, bucket, obj, _, _ = urlsplit(prediction_path)
- if scheme != "gs" or not bucket or not obj:
- raise ValueError(f"Wrong format prediction_path: {prediction_path}")
- summary = os.path.join(obj.strip("/"), "prediction.summary.json")
- gcs_hook = GCSHook()
- summary = json.loads(gcs_hook.download(bucket, summary).decode("utf-8"))
- return validate_fn(summary)
-
- evaluate_validation = PythonOperator(
- task_id=(task_prefix + "-validation"),
- python_callable=apply_validate_fn,
- templates_dict={"prediction_path": prediction_path},
- dag=dag,
- )
- evaluate_validation.set_upstream(evaluate_summary)
-
- return evaluate_prediction, evaluate_summary, evaluate_validation
diff --git a/providers/google/src/airflow/providers/google/common/hooks/base_google.py b/providers/google/src/airflow/providers/google/common/hooks/base_google.py
index 21641d468c6..b12ebdb81e7 100644
--- a/providers/google/src/airflow/providers/google/common/hooks/base_google.py
+++ b/providers/google/src/airflow/providers/google/common/hooks/base_google.py
@@ -49,20 +49,17 @@ from googleapiclient.http import MediaIoBaseDownload, build_http, set_user_agent
from requests import Session
from airflow import version
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
from airflow.hooks.base import BaseHook
from airflow.providers.google.cloud.utils.credentials_provider import (
_get_scopes,
_get_target_principal_and_delegates,
get_credentials_and_project_id,
)
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.deprecated import deprecated
from airflow.utils.process_utils import patch_environ
if TYPE_CHECKING:
from aiohttp import ClientSession
- from google.api_core.gapic_v1.client_info import ClientInfo
from google.auth.credentials import Credentials
log = logging.getLogger(__name__)
@@ -442,24 +439,6 @@ class GoogleBaseHook(BaseHook):
f"Please check the connection configuration."
)
- @property
- @deprecated(
- planned_removal_date="March 01, 2025",
- use_instead="airflow.providers.google.common.consts.CLIENT_INFO",
- category=AirflowProviderDeprecationWarning,
- )
- def client_info(self) -> ClientInfo:
- """
- Return client information used to generate a user-agent for API calls.
-
- It allows for better errors tracking.
-
- This object is only used by the google-cloud-* libraries that are built specifically for
- the Google Cloud. It is not supported by The Google APIs Python Client that use Discovery
- based APIs.
- """
- return CLIENT_INFO
-
@property
def scopes(self) -> Sequence[str]:
"""
diff --git a/providers/google/tests/deprecations_ignore.yml b/providers/google/tests/deprecations_ignore.yml
index 56db413e96c..81ab445000b 100644
--- a/providers/google/tests/deprecations_ignore.yml
+++ b/providers/google/tests/deprecations_ignore.yml
@@ -68,16 +68,6 @@
- providers/google/tests/unit/google/cloud/hooks/test_life_sciences.py::TestLifeSciencesHookWithPassedProjectId::test_waiting_operation
- providers/google/tests/unit/google/cloud/hooks/test_life_sciences.py::TestLifeSciencesHookWithoutProjectId::test_life_science_client_creation
- providers/google/tests/unit/google/cloud/hooks/test_life_sciences.py::TestLifeSciencesHookWithoutProjectId::test_run_pipeline
-- providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithDefaultProjectIdHook::test_cancel_pipeline_job
-- providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithDefaultProjectIdHook::test_create_pipeline_job
-- providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithDefaultProjectIdHook::test_delete_pipeline_job
-- providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithDefaultProjectIdHook::test_get_pipeline_job
-- providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithDefaultProjectIdHook::test_list_pipeline_jobs
-- providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithoutDefaultProjectIdHook::test_cancel_pipeline_job
-- providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithoutDefaultProjectIdHook::test_create_pipeline_job
-- providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithoutDefaultProjectIdHook::test_delete_pipeline_job
-- providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithoutDefaultProjectIdHook::test_get_pipeline_job
-- providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithoutDefaultProjectIdHook::test_list_pipeline_jobs
- providers/google/tests/unit/google/cloud/operators/test_dataproc.py::TestDataprocClusterScaleOperator::test_execute
- providers/google/tests/unit/google/cloud/operators/test_dataproc.py::test_create_cluster_operator_extra_links
- providers/google/tests/unit/google/cloud/operators/test_dataproc.py::test_scale_cluster_operator_extra_links
diff --git a/providers/google/tests/unit/google/ads/hooks/test_ads.py b/providers/google/tests/unit/google/ads/hooks/test_ads.py
index 813a276da4e..5e32a1440a9 100644
--- a/providers/google/tests/unit/google/ads/hooks/test_ads.py
+++ b/providers/google/tests/unit/google/ads/hooks/test_ads.py
@@ -52,17 +52,6 @@ def mock_hook(request):
yield hook
-@pytest.fixture(
- params=[EXTRAS_DEVELOPER_TOKEN, EXTRAS_SERVICE_ACCOUNT], ids=["developer_token", "service_account"]
-)
-def mock_hook_v16(request):
- # TODO: remove this after deprecation removal for page_size parameter
- with mock.patch("airflow.hooks.base.BaseHook.get_connection") as conn:
- hook = GoogleAdsHook(api_version="v16")
- conn.return_value.extra_dejson = request.param
- yield hook
-
-
@pytest.fixture(
params=[
{"input": EXTRAS_DEVELOPER_TOKEN, "expected_result": "developer_token"},
@@ -112,26 +101,6 @@ class TestGoogleAdsHook:
assert kwargs["request"]["query"] == query
assert "page_size" not in kwargs["request"]
- # TODO: remove this after deprecation removal for page_size parameter
- @mock.patch("airflow.providers.google.ads.hooks.ads.GoogleAdsClient")
- def test_search_v16(self, mock_client, mock_hook_v16):
- service = mock_client.load_from_dict.return_value.get_service.return_value
- mock_client.load_from_dict.return_value.get_type.side_effect = [PropertyMock(), PropertyMock()]
- client_ids = ["1", "2"]
- rows = ["row1", "row2"]
- service.search.side_effects = rows
-
- # Here we mock _extract_rows to assert calls and
- # avoid additional __iter__ calls
- mock_hook_v16._extract_rows = list
- query = "QUERY"
- mock_hook_v16.search(client_ids=client_ids, query=query, page_size=2)
- for i, client_id in enumerate(client_ids):
- name, args, kwargs = service.search.mock_calls[i]
- assert kwargs["request"]["customer_id"] == client_id
- assert kwargs["request"]["query"] == query
- assert kwargs["request"]["page_size"] == 2
-
def test_extract_rows(self, mock_hook):
iterators = [[1, 2, 3], [4, 5, 6]]
assert mock_hook._extract_rows(iterators) == sum(iterators, [])
diff --git a/providers/google/tests/unit/google/ads/transfers/test_ads_to_gcs.py b/providers/google/tests/unit/google/ads/transfers/test_ads_to_gcs.py
index a5ba0d66dae..a5d83ccb308 100644
--- a/providers/google/tests/unit/google/ads/transfers/test_ads_to_gcs.py
+++ b/providers/google/tests/unit/google/ads/transfers/test_ads_to_gcs.py
@@ -55,41 +55,7 @@ class TestGoogleAdsToGcsOperator:
google_ads_conn_id=google_ads_conn_id,
api_version=api_version,
)
- mock_ads_hook.return_value.search.assert_called_once_with(
- client_ids=CLIENT_IDS, query=QUERY, page_size=None
- )
- mock_gcs_hook.assert_called_once_with(
- gcp_conn_id=gcp_conn_id,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
- mock_gcs_hook.return_value.upload.assert_called_once_with(
- bucket_name=BUCKET, object_name=GCS_OBJ_PATH, filename=mock.ANY, gzip=False
- )
-
- @mock.patch("airflow.providers.google.ads.transfers.ads_to_gcs.GoogleAdsHook")
- @mock.patch("airflow.providers.google.ads.transfers.ads_to_gcs.GCSHook")
- def test_execute_v16(self, mock_gcs_hook, mock_ads_hook):
- op = GoogleAdsToGcsOperator(
- gcp_conn_id=gcp_conn_id,
- google_ads_conn_id=google_ads_conn_id,
- client_ids=CLIENT_IDS,
- query=QUERY,
- attributes=FIELDS_TO_EXTRACT,
- obj=GCS_OBJ_PATH,
- bucket=BUCKET,
- task_id="run_operator",
- impersonation_chain=IMPERSONATION_CHAIN,
- api_version="v16",
- )
- op.execute({})
- mock_ads_hook.assert_called_once_with(
- gcp_conn_id=gcp_conn_id,
- google_ads_conn_id=google_ads_conn_id,
- api_version="v16",
- )
- mock_ads_hook.return_value.search.assert_called_once_with(
- client_ids=CLIENT_IDS, query=QUERY, page_size=10000
- )
+ mock_ads_hook.return_value.search.assert_called_once_with(client_ids=CLIENT_IDS, query=QUERY)
mock_gcs_hook.assert_called_once_with(
gcp_conn_id=gcp_conn_id,
impersonation_chain=IMPERSONATION_CHAIN,
diff --git a/providers/google/tests/unit/google/cloud/hooks/test_cloud_build.py b/providers/google/tests/unit/google/cloud/hooks/test_cloud_build.py
index 1a704bb2c47..455bb143f9a 100644
--- a/providers/google/tests/unit/google/cloud/hooks/test_cloud_build.py
+++ b/providers/google/tests/unit/google/cloud/hooks/test_cloud_build.py
@@ -29,7 +29,7 @@ import pytest
from google.api_core.gapic_v1.method import DEFAULT
from google.cloud.devtools.cloudbuild_v1 import CloudBuildAsyncClient, GetBuildRequest
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
from airflow.providers.google.cloud.hooks.cloud_build import CloudBuildAsyncHook, CloudBuildHook
from airflow.providers.google.common.consts import CLIENT_INFO
@@ -88,53 +88,6 @@ class TestCloudBuildHook:
request={"project_id": PROJECT_ID, "id": BUILD_ID}, retry=DEFAULT, timeout=None, metadata=()
)
- @mock.patch(
- "airflow.providers.google.cloud.hooks.cloud_build.CloudBuildHook._get_build_id_from_operation"
- )
- @mock.patch("airflow.providers.google.cloud.hooks.cloud_build.TIME_TO_SLEEP_IN_SECONDS")
- @mock.patch("airflow.providers.google.cloud.hooks.cloud_build.CloudBuildHook.get_conn")
- def test_create_build_with_wait(self, get_conn, wait_time, mock_get_id_from_operation):
- get_conn.return_value.run_build_trigger.return_value = mock.MagicMock()
- mock_get_id_from_operation.return_value = BUILD_ID
-
- wait_time.return_value = 0
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.hook.create_build(build=BUILD, project_id=PROJECT_ID)
-
- get_conn.return_value.create_build.assert_called_once_with(
- request={"project_id": PROJECT_ID, "build": BUILD}, retry=DEFAULT, timeout=None, metadata=()
- )
-
- get_conn.return_value.create_build.return_value.result.assert_called_once_with()
-
- get_conn.return_value.get_build.assert_called_once_with(
- request={"project_id": PROJECT_ID, "id": BUILD_ID}, retry=DEFAULT, timeout=None, metadata=()
- )
-
- @mock.patch(
- "airflow.providers.google.cloud.hooks.cloud_build.CloudBuildHook._get_build_id_from_operation"
- )
- @mock.patch("airflow.providers.google.cloud.hooks.cloud_build.CloudBuildHook.get_conn")
- def test_create_build_without_wait(self, get_conn, mock_get_id_from_operation):
- get_conn.return_value.run_build_trigger.return_value = mock.MagicMock()
- mock_get_id_from_operation.return_value = BUILD_ID
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.hook.create_build(build=BUILD, project_id=PROJECT_ID, wait=False)
-
- mock_operation = get_conn.return_value.create_build
-
- mock_operation.assert_called_once_with(
- request={"project_id": PROJECT_ID, "build": BUILD}, retry=DEFAULT, timeout=None, metadata=()
- )
-
- get_conn.return_value.get_build.assert_called_once_with(
- request={"project_id": PROJECT_ID, "id": BUILD_ID}, retry=DEFAULT, timeout=None, metadata=()
- )
-
- mock_get_id_from_operation.assert_called_once_with(mock_operation())
-
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_build.CloudBuildHook._get_build_id_from_operation"
)
diff --git a/providers/google/tests/unit/google/cloud/hooks/test_dataflow.py b/providers/google/tests/unit/google/cloud/hooks/test_dataflow.py
index 22e68565338..a673517cf71 100644
--- a/providers/google/tests/unit/google/cloud/hooks/test_dataflow.py
+++ b/providers/google/tests/unit/google/cloud/hooks/test_dataflow.py
@@ -23,7 +23,6 @@ import re
import shlex
import subprocess
from asyncio import Future
-from typing import Any
from unittest import mock
from unittest.mock import MagicMock
from uuid import UUID
@@ -38,8 +37,8 @@ from google.cloud.dataflow_v1beta3 import (
)
from google.cloud.dataflow_v1beta3.types import JobMessageImportance
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.providers.apache.beam.hooks.beam import BeamHook, run_beam_command
+from airflow.exceptions import AirflowException
+from airflow.providers.apache.beam.hooks.beam import run_beam_command
from airflow.providers.google.cloud.hooks.dataflow import (
DEFAULT_DATAFLOW_LOCATION,
AsyncDataflowHook,
@@ -51,10 +50,6 @@ from airflow.providers.google.cloud.hooks.dataflow import (
process_line_and_extract_dataflow_job_id_callback,
)
-DEFAULT_RUNNER = "DirectRunner"
-BEAM_STRING = "airflow.providers.apache.beam.hooks.beam.{}"
-
-TASK_ID = "test-dataflow-operator"
JOB_NAME = "test-dataflow-pipeline"
MOCK_UUID = UUID("cf4a56d2-8101-4217-b027-2af6216feb48")
MOCK_UUID_PREFIX = str(MOCK_UUID)[:8]
@@ -64,17 +59,6 @@ PARAMETERS = {
"inputFile": "gs://dataflow-samples/shakespeare/kinglear.txt",
"output": "gs://test/output/my_output",
}
-TEST_ENVIRONMENT: dict[str, str] = {}
-PY_FILE = "apache_beam.examples.wordcount"
-JAR_FILE = "unitest.jar"
-JOB_CLASS = "com.example.UnitTest"
-PY_OPTIONS = ["-m"]
-DATAFLOW_VARIABLES_PY = {"project": "test", "staging_location": "gs://test/staging", "labels": {"foo": "bar"}}
-DATAFLOW_VARIABLES_JAVA = {
- "project": "test",
- "stagingLocation": "gs://test/staging",
- "labels": {"foo": "bar"},
-}
RUNTIME_ENV = {
"additionalExperiments": ["exp_flag1", "exp_flag2"],
"additionalUserLabels": {"name": "wrench", "mass": "1.3kg", "count": "3"},
@@ -94,13 +78,11 @@ RUNTIME_ENV = {
"zone": "us-central1-f",
"machineType": "n1-standard-1",
}
-BASE_STRING = "airflow.providers.google.common.hooks.base_google.{}"
DATAFLOW_STRING = "airflow.providers.google.cloud.hooks.dataflow.{}"
TEST_PROJECT = "test-project"
TEST_JOB_ID = "test-job-id"
TEST_JOBS_FILTER = ListJobsRequest.Filter.ACTIVE
TEST_LOCATION = "custom-location"
-DEFAULT_PY_INTERPRETER = "python3"
TEST_FLEX_PARAMETERS = {
"containerSpecGcsPath": "gs://test-bucket/test-file",
"jobName": "test-job-name",
@@ -223,498 +205,6 @@ class TestDataflowHook:
)
assert mock_build.return_value == result
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_python_dataflow(self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid):
- mock_beam_start_python_pipeline = self.dataflow_hook.beam_hook.start_python_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- py_requirements = ["pandas", "numpy"]
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_python_dataflow(
- job_name=JOB_NAME,
- variables=DATAFLOW_VARIABLES_PY,
- dataflow=PY_FILE,
- py_options=PY_OPTIONS,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_requirements=py_requirements,
- on_new_job_id_callback=on_new_job_id_callback,
- )
-
- expected_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
- expected_variables["job_name"] = job_name
- expected_variables["region"] = DEFAULT_DATAFLOW_LOCATION
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_python_pipeline.assert_called_once_with(
- variables=expected_variables,
- py_file=PY_FILE,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_options=PY_OPTIONS,
- py_requirements=py_requirements,
- py_system_site_packages=False,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=DEFAULT_DATAFLOW_LOCATION
- )
-
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_python_dataflow_with_custom_region_as_variable(
- self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid
- ):
- mock_beam_start_python_pipeline = self.dataflow_hook.beam_hook.start_python_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- py_requirements = ["pandas", "numpy"]
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- passed_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
- passed_variables["region"] = TEST_LOCATION
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_python_dataflow(
- job_name=JOB_NAME,
- variables=passed_variables,
- dataflow=PY_FILE,
- py_options=PY_OPTIONS,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_requirements=py_requirements,
- on_new_job_id_callback=on_new_job_id_callback,
- )
-
- expected_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
- expected_variables["job_name"] = job_name
- expected_variables["region"] = TEST_LOCATION
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_python_pipeline.assert_called_once_with(
- variables=expected_variables,
- py_file=PY_FILE,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_options=PY_OPTIONS,
- py_requirements=py_requirements,
- py_system_site_packages=False,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=TEST_LOCATION
- )
-
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_python_dataflow_with_custom_region_as_parameter(
- self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid
- ):
- mock_beam_start_python_pipeline = self.dataflow_hook.beam_hook.start_python_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- py_requirements = ["pandas", "numpy"]
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- passed_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_python_dataflow(
- job_name=JOB_NAME,
- variables=passed_variables,
- dataflow=PY_FILE,
- py_options=PY_OPTIONS,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_requirements=py_requirements,
- on_new_job_id_callback=on_new_job_id_callback,
- location=TEST_LOCATION,
- )
-
- expected_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
- expected_variables["job_name"] = job_name
- expected_variables["region"] = TEST_LOCATION
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_python_pipeline.assert_called_once_with(
- variables=expected_variables,
- py_file=PY_FILE,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_options=PY_OPTIONS,
- py_requirements=py_requirements,
- py_system_site_packages=False,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=TEST_LOCATION
- )
-
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_python_dataflow_with_no_custom_region_or_region(
- self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid
- ):
- mock_beam_start_python_pipeline = self.dataflow_hook.beam_hook.start_python_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- py_requirements = ["pandas", "numpy"]
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- passed_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_python_dataflow(
- job_name=JOB_NAME,
- variables=passed_variables,
- dataflow=PY_FILE,
- py_options=PY_OPTIONS,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_requirements=py_requirements,
- on_new_job_id_callback=on_new_job_id_callback,
- )
-
- expected_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
- expected_variables["job_name"] = job_name
- expected_variables["region"] = DEFAULT_DATAFLOW_LOCATION
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_python_pipeline.assert_called_once_with(
- variables=expected_variables,
- py_file=PY_FILE,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_options=PY_OPTIONS,
- py_requirements=py_requirements,
- py_system_site_packages=False,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=DEFAULT_DATAFLOW_LOCATION
- )
-
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_python_dataflow_with_multiple_extra_packages(
- self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid
- ):
- mock_beam_start_python_pipeline = self.dataflow_hook.beam_hook.start_python_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- py_requirements = ["pandas", "numpy"]
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- passed_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
- passed_variables["extra-package"] = ["a.whl", "b.whl"]
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_python_dataflow(
- job_name=JOB_NAME,
- variables=passed_variables,
- dataflow=PY_FILE,
- py_options=PY_OPTIONS,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_requirements=py_requirements,
- on_new_job_id_callback=on_new_job_id_callback,
- )
-
- expected_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
- expected_variables["job_name"] = job_name
- expected_variables["region"] = DEFAULT_DATAFLOW_LOCATION
- expected_variables["extra-package"] = ["a.whl", "b.whl"]
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_python_pipeline.assert_called_once_with(
- variables=expected_variables,
- py_file=PY_FILE,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_options=PY_OPTIONS,
- py_requirements=py_requirements,
- py_system_site_packages=False,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=DEFAULT_DATAFLOW_LOCATION
- )
-
- @pytest.mark.parametrize("py_interpreter", ["python3", "python2", "python3.6"])
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_python_dataflow_with_custom_interpreter(
- self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid, py_interpreter
- ):
- mock_beam_start_python_pipeline = self.dataflow_hook.beam_hook.start_python_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_python_dataflow(
- job_name=JOB_NAME,
- variables=DATAFLOW_VARIABLES_PY,
- dataflow=PY_FILE,
- py_options=PY_OPTIONS,
- py_interpreter=py_interpreter,
- py_requirements=None,
- on_new_job_id_callback=on_new_job_id_callback,
- )
-
- expected_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
- expected_variables["job_name"] = job_name
- expected_variables["region"] = DEFAULT_DATAFLOW_LOCATION
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_python_pipeline.assert_called_once_with(
- variables=expected_variables,
- py_file=PY_FILE,
- py_interpreter=py_interpreter,
- py_options=PY_OPTIONS,
- py_requirements=None,
- py_system_site_packages=False,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=DEFAULT_DATAFLOW_LOCATION
- )
-
- @pytest.mark.parametrize(
- "current_py_requirements, current_py_system_site_packages",
- [
- (["foo-bar"], False),
- (["foo-bar"], True),
- ([], True),
- ],
- )
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_python_dataflow_with_non_empty_py_requirements_and_without_system_packages(
- self,
- mock_callback_on_job_id,
- mock_dataflow_wait_for_done,
- mock_uuid,
- current_py_requirements,
- current_py_system_site_packages,
- ):
- mock_beam_start_python_pipeline = self.dataflow_hook.beam_hook.start_python_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_python_dataflow(
- job_name=JOB_NAME,
- variables=DATAFLOW_VARIABLES_PY,
- dataflow=PY_FILE,
- py_options=PY_OPTIONS,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_requirements=current_py_requirements,
- py_system_site_packages=current_py_system_site_packages,
- on_new_job_id_callback=on_new_job_id_callback,
- )
-
- expected_variables = copy.deepcopy(DATAFLOW_VARIABLES_PY)
- expected_variables["job_name"] = job_name
- expected_variables["region"] = DEFAULT_DATAFLOW_LOCATION
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_python_pipeline.assert_called_once_with(
- variables=expected_variables,
- py_file=PY_FILE,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_options=PY_OPTIONS,
- py_requirements=current_py_requirements,
- py_system_site_packages=current_py_system_site_packages,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=DEFAULT_DATAFLOW_LOCATION
- )
-
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- def test_start_python_dataflow_with_empty_py_requirements_and_without_system_packages(
- self, mock_dataflow_wait_for_done, mock_uuid
- ):
- self.dataflow_hook.beam_hook = BeamHook(runner="DataflowRunner")
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- with pytest.raises(AirflowException, match=r"Invalid method invocation\."):
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_python_dataflow(
- job_name=JOB_NAME,
- variables=DATAFLOW_VARIABLES_PY,
- dataflow=PY_FILE,
- py_options=PY_OPTIONS,
- py_interpreter=DEFAULT_PY_INTERPRETER,
- py_requirements=[],
- on_new_job_id_callback=on_new_job_id_callback,
- )
-
- mock_dataflow_wait_for_done.assert_not_called()
-
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_java_dataflow(self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid):
- mock_beam_start_java_pipeline = self.dataflow_hook.beam_hook.start_java_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_java_dataflow(
- job_name=JOB_NAME,
- variables=DATAFLOW_VARIABLES_JAVA,
- jar=JAR_FILE,
- job_class=JOB_CLASS,
- on_new_job_id_callback=on_new_job_id_callback,
- )
-
- expected_variables = copy.deepcopy(DATAFLOW_VARIABLES_JAVA)
- expected_variables["jobName"] = job_name
- expected_variables["region"] = DEFAULT_DATAFLOW_LOCATION
- expected_variables["labels"] = '{"foo":"bar"}'
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_java_pipeline.assert_called_once_with(
- variables=expected_variables,
- jar=JAR_FILE,
- job_class=JOB_CLASS,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=DEFAULT_DATAFLOW_LOCATION, multiple_jobs=False
- )
-
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_java_dataflow_with_multiple_values_in_variables(
- self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid
- ):
- mock_beam_start_java_pipeline = self.dataflow_hook.beam_hook.start_java_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- passed_variables: dict[str, Any] = copy.deepcopy(DATAFLOW_VARIABLES_JAVA)
- passed_variables["mock-option"] = ["a.whl", "b.whl"]
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_java_dataflow(
- job_name=JOB_NAME,
- variables=passed_variables,
- jar=JAR_FILE,
- job_class=JOB_CLASS,
- on_new_job_id_callback=on_new_job_id_callback,
- )
-
- expected_variables = copy.deepcopy(passed_variables)
- expected_variables["jobName"] = job_name
- expected_variables["region"] = DEFAULT_DATAFLOW_LOCATION
- expected_variables["labels"] = '{"foo":"bar"}'
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_java_pipeline.assert_called_once_with(
- variables=expected_variables,
- jar=JAR_FILE,
- job_class=JOB_CLASS,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=DEFAULT_DATAFLOW_LOCATION, multiple_jobs=False
- )
-
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_java_dataflow_with_custom_region_as_variable(
- self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid
- ):
- mock_beam_start_java_pipeline = self.dataflow_hook.beam_hook.start_java_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- passed_variables: dict[str, Any] = copy.deepcopy(DATAFLOW_VARIABLES_JAVA)
- passed_variables["region"] = TEST_LOCATION
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_java_dataflow(
- job_name=JOB_NAME,
- variables=passed_variables,
- jar=JAR_FILE,
- job_class=JOB_CLASS,
- on_new_job_id_callback=on_new_job_id_callback,
- )
-
- expected_variables = copy.deepcopy(DATAFLOW_VARIABLES_JAVA)
- expected_variables["jobName"] = job_name
- expected_variables["region"] = TEST_LOCATION
- expected_variables["labels"] = '{"foo":"bar"}'
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_java_pipeline.assert_called_once_with(
- variables=expected_variables,
- jar=JAR_FILE,
- job_class=JOB_CLASS,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=TEST_LOCATION, multiple_jobs=False
- )
-
- @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
- @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
- @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
- def test_start_java_dataflow_with_custom_region_as_parameter(
- self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid
- ):
- mock_beam_start_java_pipeline = self.dataflow_hook.beam_hook.start_java_pipeline
- mock_uuid.return_value = MOCK_UUID
- on_new_job_id_callback = MagicMock()
- job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- self.dataflow_hook.start_java_dataflow(
- job_name=JOB_NAME,
- variables=DATAFLOW_VARIABLES_JAVA,
- jar=JAR_FILE,
- job_class=JOB_CLASS,
- on_new_job_id_callback=on_new_job_id_callback,
- location=TEST_LOCATION,
- )
-
- expected_variables = copy.deepcopy(DATAFLOW_VARIABLES_JAVA)
- expected_variables["jobName"] = job_name
- expected_variables["region"] = TEST_LOCATION
- expected_variables["labels"] = '{"foo":"bar"}'
-
- mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
- mock_beam_start_java_pipeline.assert_called_once_with(
- variables=expected_variables,
- jar=JAR_FILE,
- job_class=JOB_CLASS,
- process_line_callback=mock_callback_on_job_id.return_value,
- )
-
- mock_dataflow_wait_for_done.assert_called_once_with(
- job_id=mock.ANY, job_name=job_name, location=TEST_LOCATION, multiple_jobs=False
- )
-
@pytest.mark.parametrize(
"expected_result, job_name, append_job_name",
[
diff --git a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py
index 5ce42779ee1..cf4678c42e2 100644
--- a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py
+++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py
@@ -105,26 +105,6 @@ class TestCustomJobWithDefaultProjectIdHook:
):
self.hook = CustomJobHook(gcp_conn_id=TEST_GCP_CONN_ID)
- @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
- def test_cancel_pipeline_job(self, mock_client) -> None:
- self.hook.cancel_pipeline_job(
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- pipeline_job=TEST_PIPELINE_JOB_ID,
- )
- mock_client.assert_called_once_with(TEST_REGION)
- mock_client.return_value.cancel_pipeline_job.assert_called_once_with(
- request=dict(
- name=mock_client.return_value.pipeline_job_path.return_value,
- ),
- metadata=(),
- retry=DEFAULT,
- timeout=None,
- )
- mock_client.return_value.pipeline_job_path.assert_called_once_with(
- TEST_PROJECT_ID, TEST_REGION, TEST_PIPELINE_JOB_ID
- )
-
@mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
def test_cancel_training_pipeline(self, mock_client) -> None:
self.hook.cancel_training_pipeline(
@@ -145,27 +125,6 @@ class TestCustomJobWithDefaultProjectIdHook:
TEST_PROJECT_ID, TEST_REGION, TEST_TRAINING_PIPELINE_NAME
)
- @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
- def test_create_pipeline_job(self, mock_client) -> None:
- self.hook.create_pipeline_job(
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- pipeline_job=TEST_PIPELINE_JOB,
- pipeline_job_id=TEST_PIPELINE_JOB_ID,
- )
- mock_client.assert_called_once_with(TEST_REGION)
- mock_client.return_value.create_pipeline_job.assert_called_once_with(
- request=dict(
- parent=mock_client.return_value.common_location_path.return_value,
- pipeline_job=TEST_PIPELINE_JOB,
- pipeline_job_id=TEST_PIPELINE_JOB_ID,
- ),
- metadata=(),
- retry=DEFAULT,
- timeout=None,
- )
- mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID, TEST_REGION)
-
@mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
def test_create_training_pipeline(self, mock_client) -> None:
self.hook.create_training_pipeline(
@@ -185,26 +144,6 @@ class TestCustomJobWithDefaultProjectIdHook:
)
mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID, TEST_REGION)
- @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
- def test_delete_pipeline_job(self, mock_client) -> None:
- self.hook.delete_pipeline_job(
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- pipeline_job=TEST_PIPELINE_JOB_ID,
- )
- mock_client.assert_called_once_with(TEST_REGION)
- mock_client.return_value.delete_pipeline_job.assert_called_once_with(
- request=dict(
- name=mock_client.return_value.pipeline_job_path.return_value,
- ),
- metadata=(),
- retry=DEFAULT,
- timeout=None,
- )
- mock_client.return_value.pipeline_job_path.assert_called_once_with(
- TEST_PROJECT_ID, TEST_REGION, TEST_PIPELINE_JOB_ID
- )
-
@mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
def test_delete_training_pipeline(self, mock_client) -> None:
self.hook.delete_training_pipeline(
@@ -225,26 +164,6 @@ class TestCustomJobWithDefaultProjectIdHook:
TEST_PROJECT_ID, TEST_REGION, TEST_TRAINING_PIPELINE_NAME
)
- @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
- def test_get_pipeline_job(self, mock_client) -> None:
- self.hook.get_pipeline_job(
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- pipeline_job=TEST_PIPELINE_JOB_ID,
- )
- mock_client.assert_called_once_with(TEST_REGION)
- mock_client.return_value.get_pipeline_job.assert_called_once_with(
- request=dict(
- name=mock_client.return_value.pipeline_job_path.return_value,
- ),
- metadata=(),
- retry=DEFAULT,
- timeout=None,
- )
- mock_client.return_value.pipeline_job_path.assert_called_once_with(
- TEST_PROJECT_ID, TEST_REGION, TEST_PIPELINE_JOB_ID
- )
-
@mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
def test_get_training_pipeline(self, mock_client) -> None:
self.hook.get_training_pipeline(
@@ -265,27 +184,6 @@ class TestCustomJobWithDefaultProjectIdHook:
TEST_PROJECT_ID, TEST_REGION, TEST_TRAINING_PIPELINE_NAME
)
- @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
- def test_list_pipeline_jobs(self, mock_client) -> None:
- self.hook.list_pipeline_jobs(
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- )
- mock_client.assert_called_once_with(TEST_REGION)
- mock_client.return_value.list_pipeline_jobs.assert_called_once_with(
- request=dict(
- parent=mock_client.return_value.common_location_path.return_value,
- page_size=None,
- page_token=None,
- filter=None,
- order_by=None,
- ),
- metadata=(),
- retry=DEFAULT,
- timeout=None,
- )
- mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID, TEST_REGION)
-
@mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
def test_list_training_pipelines(self, mock_client) -> None:
self.hook.list_training_pipelines(
@@ -315,26 +213,6 @@ class TestCustomJobWithoutDefaultProjectIdHook:
):
self.hook = CustomJobHook(gcp_conn_id=TEST_GCP_CONN_ID)
- @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
- def test_cancel_pipeline_job(self, mock_client) -> None:
- self.hook.cancel_pipeline_job(
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- pipeline_job=TEST_PIPELINE_JOB_ID,
- )
- mock_client.assert_called_once_with(TEST_REGION)
- mock_client.return_value.cancel_pipeline_job.assert_called_once_with(
- request=dict(
- name=mock_client.return_value.pipeline_job_path.return_value,
- ),
- metadata=(),
- retry=DEFAULT,
- timeout=None,
- )
- mock_client.return_value.pipeline_job_path.assert_called_once_with(
- TEST_PROJECT_ID, TEST_REGION, TEST_PIPELINE_JOB_ID
- )
-
@mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
def test_cancel_training_pipeline(self, mock_client) -> None:
self.hook.cancel_training_pipeline(
@@ -355,27 +233,6 @@ class TestCustomJobWithoutDefaultProjectIdHook:
TEST_PROJECT_ID, TEST_REGION, TEST_TRAINING_PIPELINE_NAME
)
- @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
- def test_create_pipeline_job(self, mock_client) -> None:
- self.hook.create_pipeline_job(
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- pipeline_job=TEST_PIPELINE_JOB,
- pipeline_job_id=TEST_PIPELINE_JOB_ID,
- )
- mock_client.assert_called_once_with(TEST_REGION)
- mock_client.return_value.create_pipeline_job.assert_called_once_with(
- request=dict(
- parent=mock_client.return_value.common_location_path.return_value,
- pipeline_job=TEST_PIPELINE_JOB,
- pipeline_job_id=TEST_PIPELINE_JOB_ID,
- ),
- metadata=(),
- retry=DEFAULT,
- timeout=None,
- )
- mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID, TEST_REGION)
-
@mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
def test_create_training_pipeline(self, mock_client) -> None:
self.hook.create_training_pipeline(
@@ -395,26 +252,6 @@ class TestCustomJobWithoutDefaultProjectIdHook:
)
mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID, TEST_REGION)
- @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
- def test_delete_pipeline_job(self, mock_client) -> None:
- self.hook.delete_pipeline_job(
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- pipeline_job=TEST_PIPELINE_JOB_ID,
- )
- mock_client.assert_called_once_with(TEST_REGION)
- mock_client.return_value.delete_pipeline_job.assert_called_once_with(
- request=dict(
- name=mock_client.return_value.pipeline_job_path.return_value,
- ),
- metadata=(),
- retry=DEFAULT,
- timeout=None,
- )
- mock_client.return_value.pipeline_job_path.assert_called_once_with(
- TEST_PROJECT_ID, TEST_REGION, TEST_PIPELINE_JOB_ID
- )
-
@mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
def test_delete_training_pipeline(self, mock_client) -> None:
self.hook.delete_training_pipeline(
@@ -435,26 +272,6 @@ class TestCustomJobWithoutDefaultProjectIdHook:
TEST_PROJECT_ID, TEST_REGION, TEST_TRAINING_PIPELINE_NAME
)
- @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
- def test_get_pipeline_job(self, mock_client) -> None:
- self.hook.get_pipeline_job(
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- pipeline_job=TEST_PIPELINE_JOB_ID,
- )
- mock_client.assert_called_once_with(TEST_REGION)
- mock_client.return_value.get_pipeline_job.assert_called_once_with(
- request=dict(
- name=mock_client.return_value.pipeline_job_path.return_value,
- ),
- metadata=(),
- retry=DEFAULT,
- timeout=None,
- )
- mock_client.return_value.pipeline_job_path.assert_called_once_with(
- TEST_PROJECT_ID, TEST_REGION, TEST_PIPELINE_JOB_ID
- )
-
@mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
def test_get_training_pipeline(self, mock_client) -> None:
self.hook.get_training_pipeline(
@@ -475,27 +292,6 @@ class TestCustomJobWithoutDefaultProjectIdHook:
TEST_PROJECT_ID, TEST_REGION, TEST_TRAINING_PIPELINE_NAME
)
- @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
- def test_list_pipeline_jobs(self, mock_client) -> None:
- self.hook.list_pipeline_jobs(
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- )
- mock_client.assert_called_once_with(TEST_REGION)
- mock_client.return_value.list_pipeline_jobs.assert_called_once_with(
- request=dict(
- parent=mock_client.return_value.common_location_path.return_value,
- page_size=None,
- page_token=None,
- filter=None,
- order_by=None,
- ),
- metadata=(),
- retry=DEFAULT,
- timeout=None,
- )
- mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID, TEST_REGION)
-
@mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_pipeline_service_client"))
def test_list_training_pipelines(self, mock_client) -> None:
self.hook.list_training_pipelines(
diff --git a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py b/providers/google/tests/unit/google/cloud/operators/test_dataproc.py
index b9b804d0a54..35072befd72 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_dataproc.py
@@ -54,7 +54,6 @@ from airflow.providers.google.cloud.operators.dataproc import (
DataprocInstantiateInlineWorkflowTemplateOperator,
DataprocInstantiateWorkflowTemplateOperator,
DataprocListBatchesOperator,
- DataprocScaleClusterOperator,
DataprocStartClusterOperator,
DataprocStopClusterOperator,
DataprocSubmitJobOperator,
@@ -1154,126 +1153,6 @@ def test_create_cluster_operator_extra_links(
)
-class TestDataprocClusterScaleOperator(DataprocClusterTestBase):
- @classmethod
- def setup_class(cls):
- super().setup_class()
- if AIRFLOW_V_3_0_PLUS:
- cls.extra_links_expected_calls_base = [
- call.ti.xcom_push(key="conf", value=DATAPROC_CLUSTER_CONF_EXPECTED)
- ]
- else:
- cls.extra_links_expected_calls_base = [
- call.ti.xcom_push(key="conf", value=DATAPROC_CLUSTER_CONF_EXPECTED, execution_date=None)
- ]
-
- def test_deprecation_warning(self):
- with pytest.warns(AirflowProviderDeprecationWarning) as warnings:
- DataprocScaleClusterOperator(task_id=TASK_ID, cluster_name=CLUSTER_NAME, project_id=GCP_PROJECT)
- assert_warning("DataprocUpdateClusterOperator", warnings)
-
- @mock.patch(DATAPROC_PATH.format("DataprocHook"))
- def test_execute(self, mock_hook):
- self.extra_links_manager_mock.attach_mock(mock_hook, "hook")
- mock_hook.return_value.update_cluster.result.return_value = None
- cluster_update = {
- "config": {"worker_config": {"num_instances": 3}, "secondary_worker_config": {"num_instances": 4}}
- }
- update_cluster_args = {
- "project_id": GCP_PROJECT,
- "region": GCP_REGION,
- "cluster_name": CLUSTER_NAME,
- "cluster": cluster_update,
- "graceful_decommission_timeout": {"seconds": 600},
- "update_mask": UPDATE_MASK,
- }
- expected_calls = [
- *self.extra_links_expected_calls_base,
- call.hook().update_cluster(**update_cluster_args),
- ]
-
- op = DataprocScaleClusterOperator(
- task_id=TASK_ID,
- cluster_name=CLUSTER_NAME,
- project_id=GCP_PROJECT,
- region=GCP_REGION,
- num_workers=3,
- num_preemptible_workers=4,
- graceful_decommission_timeout="10m",
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- )
- op.execute(context=self.mock_context)
- mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN)
- mock_hook.return_value.update_cluster.assert_called_once_with(**update_cluster_args)
-
- # Test whether xcom push occurs before cluster is updated
- self.extra_links_manager_mock.assert_has_calls(expected_calls, any_order=False)
-
- if AIRFLOW_V_3_0_PLUS:
- self.mock_ti.xcom_push.assert_called_once_with(
- key="conf",
- value=DATAPROC_CLUSTER_CONF_EXPECTED,
- )
- else:
- self.mock_ti.xcom_push.assert_called_once_with(
- key="conf",
- value=DATAPROC_CLUSTER_CONF_EXPECTED,
- execution_date=None,
- )
-
-
-@pytest.mark.db_test
-@pytest.mark.need_serialized_dag
-def test_scale_cluster_operator_extra_links(
- dag_maker, create_task_instance_of_operator, mock_supervisor_comms
-):
- ti = create_task_instance_of_operator(
- DataprocScaleClusterOperator,
- dag_id=TEST_DAG_ID,
- task_id=TASK_ID,
- cluster_name=CLUSTER_NAME,
- project_id=GCP_PROJECT,
- region=GCP_REGION,
- num_workers=3,
- num_preemptible_workers=2,
- graceful_decommission_timeout="2m",
- gcp_conn_id=GCP_CONN_ID,
- )
-
- serialized_dag = dag_maker.get_serialized_data()
-
- # Assert operator links for serialized DAG
- deserialized_dag = SerializedDAG.deserialize_dag(serialized_dag["dag"])
- operator_extra_link = deserialized_dag.tasks[0].operator_extra_links[0]
- assert operator_extra_link.name == "Dataproc resource"
-
- if AIRFLOW_V_3_0_PLUS:
- mock_supervisor_comms.get_message.return_value = XComResult(
- key="key",
- value="",
- )
- # Assert operator link is empty when no XCom push occurred
- assert ti.task.operator_extra_links[0].get_link(operator=ti.task, ti_key=ti.key) == ""
-
- ti.xcom_push(
- key="conf",
- value=DATAPROC_CLUSTER_CONF_EXPECTED,
- )
-
- if AIRFLOW_V_3_0_PLUS:
- mock_supervisor_comms.get_message.return_value = XComResult(
- key="key",
- value=DATAPROC_CLUSTER_CONF_EXPECTED,
- )
-
- # Assert operator links after execution
- assert (
- ti.task.operator_extra_links[0].get_link(operator=ti.task, ti_key=ti.key)
- == DATAPROC_CLUSTER_LINK_EXPECTED
- )
-
-
class TestDataprocClusterDeleteOperator:
@mock.patch(DATAPROC_PATH.format("DataprocHook"))
def test_execute(self, mock_hook):
diff --git a/providers/google/tests/unit/google/cloud/operators/test_mlengine.py b/providers/google/tests/unit/google/cloud/operators/test_mlengine.py
index 2f7dee7328a..5ca6cf8cea7 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_mlengine.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_mlengine.py
@@ -16,481 +16,23 @@
# under the License.
from __future__ import annotations
-import copy
-from unittest import mock
-from unittest.mock import ANY, MagicMock, patch
+from unittest.mock import MagicMock, patch
-import httplib2
import pytest
-from googleapiclient.errors import HttpError
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, TaskDeferred
-from airflow.models.dag import DAG
-from airflow.models.dagrun import DagRun
-from airflow.models.taskinstance import TaskInstance
-from airflow.providers.google.cloud.operators.mlengine import ( # AIPlatformConsoleLink,
- MLEngineCreateModelOperator,
- MLEngineCreateVersionOperator,
- MLEngineDeleteModelOperator,
- MLEngineDeleteVersionOperator,
- MLEngineGetModelOperator,
- MLEngineListVersionsOperator,
- MLEngineManageModelOperator,
- MLEngineManageVersionOperator,
- MLEngineSetDefaultVersionOperator,
- MLEngineStartBatchPredictionJobOperator,
- MLEngineStartTrainingJobOperator,
- MLEngineTrainingCancelJobOperator,
-)
-from airflow.providers.google.cloud.triggers.mlengine import MLEngineStartTrainingJobTrigger
-from airflow.utils import timezone
-from airflow.utils.timezone import datetime
-from airflow.utils.types import DagRunType
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.google.cloud.operators.mlengine import MLEngineCreateModelOperator
-DEFAULT_DATE = timezone.datetime(2017, 6, 6)
-
-TEST_DAG_ID = "test-mlengine-operators"
TEST_PROJECT_ID = "test-project-id"
TEST_MODEL_NAME = "test-model-name"
-TEST_VERSION_NAME = "test-version"
TEST_GCP_CONN_ID = "test-gcp-conn-id"
TEST_IMPERSONATION_CHAIN = ["ACCOUNT_1", "ACCOUNT_2", "ACCOUNT_3"]
TEST_MODEL = {
"name": TEST_MODEL_NAME,
}
-TEST_VERSION = {
- "name": "v1",
- "deploymentUri": "gs://some-bucket/jobs/test_training/model.pb",
- "runtimeVersion": "1.6",
-}
MLENGINE_AI_PATH = "airflow.providers.google.cloud.operators.mlengine.{}"
-class TestMLEngineStartBatchPredictionJobOperator:
- INPUT_MISSING_ORIGIN = {
- "dataFormat": "TEXT",
- "inputPaths": ["gs://legal-bucket/fake-input-path/*"],
- "outputPath": "gs://legal-bucket/fake-output-path",
- "region": "us-east1",
- }
- SUCCESS_MESSAGE_MISSING_INPUT = {
- "jobId": "test_prediction",
- "labels": {"some": "labels"},
- "predictionOutput": {
- "outputPath": "gs://fake-output-path",
- "predictionCount": 5000,
- "errorCount": 0,
- "nodeHours": 2.78,
- },
- "state": "SUCCEEDED",
- }
- BATCH_PREDICTION_DEFAULT_ARGS = {
- "project_id": "test-project",
- "job_id": "test_prediction",
- "labels": {"some": "labels"},
- "region": "us-east1",
- "data_format": "TEXT",
- "input_paths": ["gs://legal-bucket-dash-Capital/legal-input-path/*"],
- "output_path": "gs://12_legal_bucket_underscore_number/legal-output-path",
- "task_id": "test-prediction",
- }
-
- def setup_method(self):
- self.dag = DAG(
- "test_dag",
- default_args={
- "owner": "airflow",
- "start_date": DEFAULT_DATE,
- "end_date": DEFAULT_DATE,
- },
- schedule="@daily",
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success_with_model(self, mock_hook):
- input_with_model = self.INPUT_MISSING_ORIGIN.copy()
- input_with_model["modelName"] = "projects/test-project/models/test_model"
- success_message = self.SUCCESS_MESSAGE_MISSING_INPUT.copy()
- success_message["predictionInput"] = input_with_model
-
- hook_instance = mock_hook.return_value
- hook_instance.get_job.side_effect = HttpError(
- resp=httplib2.Response({"status": 404}), content=b"some bytes"
- )
- hook_instance.create_job.return_value = success_message
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- prediction_task = MLEngineStartBatchPredictionJobOperator(
- job_id="test_prediction",
- project_id="test-project",
- region=input_with_model["region"],
- data_format=input_with_model["dataFormat"],
- input_paths=input_with_model["inputPaths"],
- output_path=input_with_model["outputPath"],
- model_name=input_with_model["modelName"].split("/")[-1],
- labels={"some": "labels"},
- dag=self.dag,
- task_id="test-prediction",
- )
- prediction_output = prediction_task.execute(None)
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- hook_instance.create_job.assert_called_once_with(
- project_id="test-project",
- job={
- "jobId": "test_prediction",
- "labels": {"some": "labels"},
- "predictionInput": input_with_model,
- },
- use_existing_job_fn=ANY,
- )
- assert success_message["predictionOutput"] == prediction_output
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success_with_version(self, mock_hook):
- input_with_version = self.INPUT_MISSING_ORIGIN.copy()
- input_with_version["versionName"] = "projects/test-project/models/test_model/versions/test_version"
- success_message = self.SUCCESS_MESSAGE_MISSING_INPUT.copy()
- success_message["predictionInput"] = input_with_version
-
- hook_instance = mock_hook.return_value
- hook_instance.get_job.side_effect = HttpError(
- resp=httplib2.Response({"status": 404}), content=b"some bytes"
- )
- hook_instance.create_job.return_value = success_message
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- prediction_task = MLEngineStartBatchPredictionJobOperator(
- job_id="test_prediction",
- project_id="test-project",
- region=input_with_version["region"],
- data_format=input_with_version["dataFormat"],
- input_paths=input_with_version["inputPaths"],
- output_path=input_with_version["outputPath"],
- model_name=input_with_version["versionName"].split("/")[-3],
- version_name=input_with_version["versionName"].split("/")[-1],
- dag=self.dag,
- task_id="test-prediction",
- )
- prediction_output = prediction_task.execute(None)
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- hook_instance.create_job.assert_called_once_with(
- project_id="test-project",
- job={"jobId": "test_prediction", "predictionInput": input_with_version},
- use_existing_job_fn=ANY,
- )
- assert success_message["predictionOutput"] == prediction_output
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success_with_uri(self, mock_hook):
- input_with_uri = self.INPUT_MISSING_ORIGIN.copy()
- input_with_uri["uri"] = "gs://my_bucket/my_models/savedModel"
- success_message = self.SUCCESS_MESSAGE_MISSING_INPUT.copy()
- success_message["predictionInput"] = input_with_uri
-
- hook_instance = mock_hook.return_value
- hook_instance.get_job.side_effect = HttpError(
- resp=httplib2.Response({"status": 404}), content=b"some bytes"
- )
- hook_instance.create_job.return_value = success_message
- with pytest.warns(AirflowProviderDeprecationWarning):
- prediction_task = MLEngineStartBatchPredictionJobOperator(
- job_id="test_prediction",
- project_id="test-project",
- region=input_with_uri["region"],
- data_format=input_with_uri["dataFormat"],
- input_paths=input_with_uri["inputPaths"],
- output_path=input_with_uri["outputPath"],
- uri=input_with_uri["uri"],
- dag=self.dag,
- task_id="test-prediction",
- )
- prediction_output = prediction_task.execute(None)
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- hook_instance.create_job.assert_called_once_with(
- project_id="test-project",
- job={"jobId": "test_prediction", "predictionInput": input_with_uri},
- use_existing_job_fn=ANY,
- )
- assert success_message["predictionOutput"] == prediction_output
-
- def test_invalid_model_origin(self):
- # Test that both uri and model is given
- task_args = self.BATCH_PREDICTION_DEFAULT_ARGS.copy()
- task_args["uri"] = "gs://fake-uri/saved_model"
- task_args["model_name"] = "fake_model"
- with pytest.raises(AirflowException) as ctx, pytest.warns(AirflowProviderDeprecationWarning):
- MLEngineStartBatchPredictionJobOperator(**task_args).execute(None)
- assert str(ctx.value) == "Ambiguous model origin: Both uri and model/version name are provided."
-
- # Test that both uri and model/version is given
- task_args = self.BATCH_PREDICTION_DEFAULT_ARGS.copy()
- task_args["uri"] = "gs://fake-uri/saved_model"
- task_args["model_name"] = "fake_model"
- task_args["version_name"] = "fake_version"
- with pytest.raises(AirflowException) as ctx, pytest.warns(AirflowProviderDeprecationWarning):
- MLEngineStartBatchPredictionJobOperator(**task_args).execute(None)
- assert str(ctx.value) == "Ambiguous model origin: Both uri and model/version name are provided."
-
- # Test that a version is given without a model
- task_args = self.BATCH_PREDICTION_DEFAULT_ARGS.copy()
- task_args["version_name"] = "bare_version"
- with pytest.raises(AirflowException) as ctx, pytest.warns(AirflowProviderDeprecationWarning):
- MLEngineStartBatchPredictionJobOperator(**task_args).execute(None)
- assert (
- str(ctx.value) == "Missing model: Batch prediction expects a model "
- "name when a version name is provided."
- )
-
- # Test that none of uri, model, model/version is given
- task_args = self.BATCH_PREDICTION_DEFAULT_ARGS.copy()
- with pytest.raises(AirflowException) as ctx, pytest.warns(AirflowProviderDeprecationWarning):
- MLEngineStartBatchPredictionJobOperator(**task_args).execute(None)
- assert (
- str(ctx.value) == "Missing model origin: Batch prediction expects a "
- "model, a model & version combination, or a URI to a savedModel."
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_http_error(self, mock_hook):
- http_error_code = 403
- input_with_model = self.INPUT_MISSING_ORIGIN.copy()
- input_with_model["modelName"] = "projects/experimental/models/test_model"
-
- hook_instance = mock_hook.return_value
- hook_instance.create_job.side_effect = HttpError(
- resp=httplib2.Response({"status": http_error_code}), content=b"Forbidden"
- )
- with pytest.warns(AirflowProviderDeprecationWarning):
- prediction_task = MLEngineStartBatchPredictionJobOperator(
- job_id="test_prediction",
- project_id="test-project",
- region=input_with_model["region"],
- data_format=input_with_model["dataFormat"],
- input_paths=input_with_model["inputPaths"],
- output_path=input_with_model["outputPath"],
- model_name=input_with_model["modelName"].split("/")[-1],
- dag=self.dag,
- task_id="test-prediction",
- )
- with pytest.raises(HttpError):
- prediction_task.execute(None)
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_failed_job_error(self, mock_hook):
- hook_instance = mock_hook.return_value
- hook_instance.create_job.return_value = {"state": "FAILED", "errorMessage": "A failure message"}
- task_args = self.BATCH_PREDICTION_DEFAULT_ARGS.copy()
- task_args["uri"] = "a uri"
-
- with pytest.raises(RuntimeError) as ctx, pytest.warns(AirflowProviderDeprecationWarning):
- MLEngineStartBatchPredictionJobOperator(**task_args).execute(None)
-
- assert str(ctx.value) == "A failure message"
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineStartBatchPredictionJobOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- job_id="{{ 'job_id' }}",
- region="{{ 'region' }}",
- input_paths="{{ 'input_paths' }}",
- output_path="{{ 'output_path' }}",
- model_name="{{ 'model_name' }}",
- version_name="{{ 'version_name' }}",
- uri="{{ 'uri' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- data_format="data_format",
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineStartBatchPredictionJobOperator = ti.task
- assert task.project_id == "project_id"
- assert task.job_id == "job_id"
- assert task.region == "region"
- assert task.input_paths == "input_paths"
- assert task.output_path == "output_path"
- assert task.model_name == "model_name"
- assert task.version_name == "version_name"
- assert task.uri == "uri"
- assert task.impersonation_chain == "impersonation_chain"
-
-
-class TestMLEngineTrainingCancelJobOperator:
- TRAINING_DEFAULT_ARGS = {
- "project_id": "test-project",
- "job_id": "test_training",
- "task_id": "test-training",
- }
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success_cancel_training_job(self, mock_hook):
- success_response = {}
- hook_instance = mock_hook.return_value
- hook_instance.cancel_job.return_value = success_response
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- cancel_training_op = MLEngineTrainingCancelJobOperator(**self.TRAINING_DEFAULT_ARGS)
- cancel_training_op.execute(context=MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- # Make sure only 'cancel_job' is invoked on hook instance
- assert len(hook_instance.mock_calls) == 1
- hook_instance.cancel_job.assert_called_once_with(
- project_id=self.TRAINING_DEFAULT_ARGS["project_id"], job_id=self.TRAINING_DEFAULT_ARGS["job_id"]
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_http_error(self, mock_hook):
- http_error_code = 403
- hook_instance = mock_hook.return_value
- hook_instance.cancel_job.side_effect = HttpError(
- resp=httplib2.Response({"status": http_error_code}), content=b"Forbidden"
- )
- with pytest.warns(AirflowProviderDeprecationWarning):
- cancel_training_op = MLEngineTrainingCancelJobOperator(**self.TRAINING_DEFAULT_ARGS)
- with pytest.raises(HttpError) as ctx:
- cancel_training_op.execute(context=MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- # Make sure only 'cancel_job' is invoked on hook instance
- assert len(hook_instance.mock_calls) == 1
- hook_instance.cancel_job.assert_called_once_with(
- project_id=self.TRAINING_DEFAULT_ARGS["project_id"], job_id=self.TRAINING_DEFAULT_ARGS["job_id"]
- )
- assert http_error_code == ctx.value.resp.status
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineTrainingCancelJobOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- job_id="{{ 'job_id' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineTrainingCancelJobOperator = ti.task
- assert task.project_id == "project_id"
- assert task.job_id == "job_id"
- assert task.impersonation_chain == "impersonation_chain"
-
-
-class TestMLEngineModelOperator:
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success_create_model(self, mock_hook):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineManageModelOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model=TEST_MODEL,
- operation="create",
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
-
- task.execute(context=MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
- mock_hook.return_value.create_model.assert_called_once_with(
- project_id=TEST_PROJECT_ID, model=TEST_MODEL
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success_get_model(self, mock_hook):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineManageModelOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model=TEST_MODEL,
- operation="get",
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
-
- result = task.execute(context=MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
- mock_hook.return_value.get_model.assert_called_once_with(
- project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME
- )
- assert mock_hook.return_value.get_model.return_value == result
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_fail(self, mock_hook):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineManageModelOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model=TEST_MODEL,
- operation="invalid",
- gcp_conn_id=TEST_GCP_CONN_ID,
- )
- with pytest.raises(ValueError):
- task.execute(None)
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineManageModelOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- model="{{ 'model' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineManageModelOperator = ti.task
- assert task.project_id == "project_id"
- assert task.model == "model"
- assert task.impersonation_chain == "impersonation_chain"
-
-
class TestMLEngineCreateModelOperator:
@patch(MLENGINE_AI_PATH.format("MLEngineHook"))
def test_success_create_model(self, mock_hook):
@@ -533,1045 +75,3 @@ class TestMLEngineCreateModelOperator:
assert task.project_id == "project_id"
assert task.model == "model"
assert task.impersonation_chain == "impersonation_chain"
-
-
-class TestMLEngineGetModelOperator:
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success_get_model(self, mock_hook):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineGetModelOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=TEST_MODEL_NAME,
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
-
- result = task.execute(context=MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
- mock_hook.return_value.get_model.assert_called_once_with(
- project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME
- )
- assert mock_hook.return_value.get_model.return_value == result
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineGetModelOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- model_name="{{ 'model_name' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineGetModelOperator = ti.task
- assert task.project_id == "project_id"
- assert task.model_name == "model_name"
- assert task.impersonation_chain == "impersonation_chain"
-
-
-class TestMLEngineDeleteModelOperator:
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success_delete_model(self, mock_hook):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineDeleteModelOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=TEST_MODEL_NAME,
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- delete_contents=True,
- )
-
- task.execute(context=MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
- mock_hook.return_value.delete_model.assert_called_once_with(
- project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME, delete_contents=True
- )
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineDeleteModelOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- model_name="{{ 'model_name' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineDeleteModelOperator = ti.task
- assert task.project_id == "project_id"
- assert task.model_name == "model_name"
- assert task.impersonation_chain == "impersonation_chain"
-
-
-class TestMLEngineVersionOperator:
- VERSION_DEFAULT_ARGS = {
- "project_id": "test-project",
- "model_name": "test-model",
- "task_id": "test-version",
- }
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success_create_version(self, mock_hook):
- success_response = {"name": "some-name", "done": True}
- hook_instance = mock_hook.return_value
- hook_instance.create_version.return_value = success_response
- with pytest.warns(AirflowProviderDeprecationWarning):
- training_op = MLEngineManageVersionOperator(version=TEST_VERSION, **self.VERSION_DEFAULT_ARGS)
- training_op.execute(None)
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- # Make sure only 'create_version' is invoked on hook instance
- assert len(hook_instance.mock_calls) == 1
- hook_instance.create_version.assert_called_once_with(
- project_id="test-project", model_name="test-model", version_spec=TEST_VERSION
- )
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineManageVersionOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- model_name="{{ 'model_name' }}",
- version="{{ 'version' }}",
- version_name="{{ 'version_name' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineManageVersionOperator = ti.task
- assert task.project_id == "project_id"
- assert task.model_name == "model_name"
- assert task.version == "version"
- assert task.version_name == "version_name"
- assert task.impersonation_chain == "impersonation_chain"
-
-
-class TestMLEngineCreateVersion:
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success(self, mock_hook):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineCreateVersionOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=TEST_MODEL_NAME,
- version=TEST_VERSION,
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
-
- task.execute(context=MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
- mock_hook.return_value.create_version.assert_called_once_with(
- project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME, version_spec=TEST_VERSION
- )
-
- def test_missing_model_name(self):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineCreateVersionOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=None,
- version=TEST_VERSION,
- gcp_conn_id=TEST_GCP_CONN_ID,
- )
- with pytest.raises(AirflowException):
- task.execute(context=MagicMock())
-
- def test_missing_version(self):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineCreateVersionOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=TEST_MODEL_NAME,
- version=None,
- gcp_conn_id=TEST_GCP_CONN_ID,
- )
- with pytest.raises(AirflowException):
- task.execute(context=MagicMock())
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineCreateVersionOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- model_name="{{ 'model_name' }}",
- version="{{ 'version' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineCreateVersionOperator = ti.task
- assert task.project_id == "project_id"
- assert task.model_name == "model_name"
- assert task.version == "version"
- assert task.impersonation_chain == "impersonation_chain"
-
-
-class TestMLEngineSetDefaultVersion:
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success(self, mock_hook):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineSetDefaultVersionOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=TEST_MODEL_NAME,
- version_name=TEST_VERSION_NAME,
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
-
- task.execute(context=MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
- mock_hook.return_value.set_default_version.assert_called_once_with(
- project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME, version_name=TEST_VERSION_NAME
- )
-
- def test_missing_model_name(self):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineSetDefaultVersionOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=None,
- version_name=TEST_VERSION_NAME,
- gcp_conn_id=TEST_GCP_CONN_ID,
- )
- with pytest.raises(AirflowException):
- task.execute(context=MagicMock())
-
- def test_missing_version_name(self):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineSetDefaultVersionOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=TEST_MODEL_NAME,
- version_name=None,
- gcp_conn_id=TEST_GCP_CONN_ID,
- )
- with pytest.raises(AirflowException):
- task.execute(context=MagicMock())
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineSetDefaultVersionOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- model_name="{{ 'model_name' }}",
- version_name="{{ 'version_name' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineSetDefaultVersionOperator = ti.task
- assert task.project_id == "project_id"
- assert task.model_name == "model_name"
- assert task.version_name == "version_name"
- assert task.impersonation_chain == "impersonation_chain"
-
-
-class TestMLEngineListVersions:
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success(self, mock_hook):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineListVersionsOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=TEST_MODEL_NAME,
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
-
- task.execute(context=MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
- mock_hook.return_value.list_versions.assert_called_once_with(
- project_id=TEST_PROJECT_ID,
- model_name=TEST_MODEL_NAME,
- )
-
- def test_missing_model_name(self):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineListVersionsOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=None,
- gcp_conn_id=TEST_GCP_CONN_ID,
- )
- with pytest.raises(AirflowException):
- task.execute(context=MagicMock())
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineListVersionsOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- model_name="{{ 'model_name' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineListVersionsOperator = ti.task
- assert task.project_id == "project_id"
- assert task.model_name == "model_name"
- assert task.impersonation_chain == "impersonation_chain"
-
-
-class TestMLEngineDeleteVersion:
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_success(self, mock_hook):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineDeleteVersionOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=TEST_MODEL_NAME,
- version_name=TEST_VERSION_NAME,
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
-
- task.execute(context=MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- )
- mock_hook.return_value.delete_version.assert_called_once_with(
- project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME, version_name=TEST_VERSION_NAME
- )
-
- def test_missing_version_name(self):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineDeleteVersionOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=TEST_MODEL_NAME,
- version_name=None,
- gcp_conn_id=TEST_GCP_CONN_ID,
- )
- with pytest.raises(AirflowException):
- task.execute(context=MagicMock())
-
- def test_missing_model_name(self):
- with pytest.warns(AirflowProviderDeprecationWarning):
- task = MLEngineDeleteVersionOperator(
- task_id="task-id",
- project_id=TEST_PROJECT_ID,
- model_name=None,
- version_name=TEST_VERSION_NAME,
- gcp_conn_id=TEST_GCP_CONN_ID,
- )
- with pytest.raises(AirflowException):
- task.execute(context=MagicMock())
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineDeleteVersionOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- model_name="{{ 'model_name' }}",
- version_name="{{ 'version_name' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineDeleteVersionOperator = ti.task
- assert task.project_id == "project_id"
- assert task.model_name == "model_name"
- assert task.version_name == "version_name"
- assert task.impersonation_chain == "impersonation_chain"
-
-
-class TestMLEngineStartTrainingJobOperator:
- TRAINING_DEFAULT_ARGS = {
- "project_id": "test-project",
- "job_id": "test_training",
- "package_uris": ["gs://some-bucket/package1"],
- "training_python_module": "trainer",
- "training_args": "--some_arg='aaa'",
- "region": "us-east1",
- "scale_tier": "STANDARD_1",
- "labels": {"some": "labels"},
- "task_id": "test-training",
- }
- TRAINING_INPUT = {
- "jobId": "test_training",
- "labels": {"some": "labels"},
- "trainingInput": {
- "scaleTier": "STANDARD_1",
- "packageUris": ["gs://some-bucket/package1"],
- "pythonModule": "trainer",
- "args": "--some_arg='aaa'",
- "region": "us-east1",
- },
- }
-
- @patch(MLENGINE_AI_PATH.format("MLEngineStartTrainingJobOperator._wait_for_job_done"))
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_create_training_job_should_execute_successfully(self, mock_hook, mock_wait_for_job):
- mock_hook.return_value.create_job_without_waiting_result.return_value = "test_training"
- mock_wait_for_job.return_value = {"state": "SUCCEEDED"}
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- training_op = MLEngineStartTrainingJobOperator(deferrable=False, **self.TRAINING_DEFAULT_ARGS)
- training_op.execute(MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- mock_hook.return_value.create_job_without_waiting_result.assert_called_once_with(
- project_id="test-project", body=self.TRAINING_INPUT
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineStartTrainingJobOperator._wait_for_job_done"))
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_create_training_job_with_master_config_should_execute_successfully(
- self, mock_hook, mock_wait_for_job
- ):
- custom_training_default_args: dict = copy.deepcopy(self.TRAINING_DEFAULT_ARGS)
- custom_training_default_args["scale_tier"] = "CUSTOM"
-
- training_input = copy.deepcopy(self.TRAINING_INPUT)
- training_input["trainingInput"]["runtimeVersion"] = "1.6"
- training_input["trainingInput"]["pythonVersion"] = "3.5"
- training_input["trainingInput"]["jobDir"] = "gs://some-bucket/jobs/test_training"
- training_input["trainingInput"]["scaleTier"] = "CUSTOM"
- training_input["trainingInput"]["masterType"] = "n1-standard-4"
- training_input["trainingInput"]["masterConfig"] = {
- "acceleratorConfig": {"count": "1", "type": "NVIDIA_TESLA_P4"},
- }
-
- success_response = training_input.copy()
- mock_wait_for_job.return_value = {"state": "SUCCEEDED"}
- mock_hook.return_value.create_job_without_waiting_result.return_value = success_response
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- training_op = MLEngineStartTrainingJobOperator(
- runtime_version="1.6",
- python_version="3.5",
- job_dir="gs://some-bucket/jobs/test_training",
- master_type="n1-standard-4",
- master_config={
- "acceleratorConfig": {"count": "1", "type": "NVIDIA_TESLA_P4"},
- },
- deferrable=False,
- **custom_training_default_args,
- )
- training_op.execute(MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- mock_hook.return_value.create_job_without_waiting_result.assert_called_once_with(
- project_id="test-project", body=training_input
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineStartTrainingJobOperator._wait_for_job_done"))
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_create_training_job_with_master_image_should_execute_successfully(
- self, mock_hook, mock_wait_for_job
- ):
- arguments = {
- "project_id": "test-project",
- "job_id": "test_training",
- "region": "europe-west1",
- "scale_tier": "CUSTOM",
- "master_type": "n1-standard-8",
- "master_config": {
- "imageUri": "eu.gcr.io/test-project/test-image:test-version",
- },
- "task_id": "test-training",
- "start_date": DEFAULT_DATE,
- "deferrable": False,
- }
- request = {
- "jobId": "test_training",
- "trainingInput": {
- "region": "europe-west1",
- "scaleTier": "CUSTOM",
- "masterType": "n1-standard-8",
- "masterConfig": {
- "imageUri": "eu.gcr.io/test-project/test-image:test-version",
- },
- },
- }
-
- response = request.copy()
- mock_wait_for_job.return_value = {"state": "SUCCEEDED"}
- mock_hook.return_value.create_job_without_waiting_result.return_value = response
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- training_op = MLEngineStartTrainingJobOperator(**arguments)
- training_op.execute(MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- mock_hook.return_value.create_job_without_waiting_result.assert_called_once_with(
- project_id="test-project",
- body=request,
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineStartTrainingJobOperator._wait_for_job_done"))
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_create_training_job_with_optional_args_should_execute_successfully(
- self, mock_hook, mock_wait_for_job
- ):
- training_input = copy.deepcopy(self.TRAINING_INPUT)
- training_input["trainingInput"]["runtimeVersion"] = "1.6"
- training_input["trainingInput"]["pythonVersion"] = "3.5"
- training_input["trainingInput"]["jobDir"] = "gs://some-bucket/jobs/test_training"
- training_input["trainingInput"]["serviceAccount"] = "test@serviceaccount.com"
-
- hyperparams = {
- "goal": "MAXIMIZE",
- "hyperparameterMetricTag": "metric1",
- "maxTrials": 30,
- "maxParallelTrials": 1,
- "enableTrialEarlyStopping": True,
- "params": [],
- }
-
- hyperparams["params"].append(
- {
- "parameterName": "hidden1",
- "type": "INTEGER",
- "minValue": 40,
- "maxValue": 400,
- "scaleType": "UNIT_LINEAR_SCALE",
- }
- )
-
- hyperparams["params"].append(
- {"parameterName": "numRnnCells", "type": "DISCRETE", "discreteValues": [1, 2, 3, 4]}
- )
-
- hyperparams["params"].append(
- {
- "parameterName": "rnnCellType",
- "type": "CATEGORICAL",
- "categoricalValues": [
- "BasicLSTMCell",
- "BasicRNNCell",
- "GRUCell",
- "LSTMCell",
- "LayerNormBasicLSTMCell",
- ],
- }
- )
-
- training_input["trainingInput"]["hyperparameters"] = hyperparams
-
- success_response = self.TRAINING_INPUT.copy()
- mock_wait_for_job.return_value = {"state": "SUCCEEDED"}
- mock_hook.return_value.create_job_without_waiting_result.return_value = success_response
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- training_op = MLEngineStartTrainingJobOperator(
- runtime_version="1.6",
- python_version="3.5",
- job_dir="gs://some-bucket/jobs/test_training",
- service_account="test@serviceaccount.com",
- **self.TRAINING_DEFAULT_ARGS,
- hyperparameters=hyperparams,
- deferrable=False,
- )
- training_op.execute(MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- mock_hook.return_value.create_job_without_waiting_result.assert_called_once_with(
- project_id="test-project", body=training_input
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineStartTrainingJobOperator._wait_for_job_done"))
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_create_training_job_when_http_error_409_should_execute_successfully(
- self, mock_hook, mock_wait_for_job
- ):
- mock_hook.return_value.create_job_without_waiting_result.return_value = HttpError(
- resp=httplib2.Response({"status": "409"}), content=b"content"
- )
- mock_hook.return_value.get_job.return_value = {"job_id": "test_training"}
- mock_wait_for_job.return_value = {"state": "SUCCEEDED"}
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- training_op = MLEngineStartTrainingJobOperator(**self.TRAINING_DEFAULT_ARGS)
- training_op.execute(MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_create_training_job_should_throw_exception_when_http_error_403(self, mock_hook):
- mock_hook.return_value.create_job_without_waiting_result.side_effect = HttpError(
- resp=httplib2.Response({"status": "403"}), content=b"content"
- )
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- training_op = MLEngineStartTrainingJobOperator(**self.TRAINING_DEFAULT_ARGS)
- with pytest.raises(HttpError):
- training_op.execute(MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
-
- @patch(MLENGINE_AI_PATH.format("MLEngineStartTrainingJobOperator._wait_for_job_done"))
- @patch(MLENGINE_AI_PATH.format("MLEngineHook"))
- def test_create_training_job_should_throw_exception_when_job_failed(self, mock_hook, mock_wait_for_job):
- failure_response = self.TRAINING_INPUT.copy()
-
- mock_wait_for_job.return_value = {"state": "FAILED", "errorMessage": "A failure message"}
- mock_hook.return_value.create_job_without_waiting_result.return_value = failure_response
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- training_op = MLEngineStartTrainingJobOperator(**self.TRAINING_DEFAULT_ARGS)
- with pytest.raises(RuntimeError) as ctx:
- training_op.execute(MagicMock())
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- mock_hook.return_value.create_job_without_waiting_result.assert_called_once_with(
- project_id="test-project", body=self.TRAINING_INPUT
- )
- assert str(ctx.value) == "A failure message"
-
- @pytest.mark.db_test
- def test_templating(self, create_task_instance_of_operator, session):
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- MLEngineStartTrainingJobOperator,
- # Templated fields
- project_id="{{ 'project_id' }}",
- job_id="{{ 'job_id' }}",
- region="{{ 'region' }}",
- package_uris="{{ 'package_uris' }}",
- training_python_module="{{ 'training_python_module' }}",
- training_args="{{ 'training_args' }}",
- scale_tier="{{ 'scale_tier' }}",
- master_type="{{ 'master_type' }}",
- master_config="{{ 'master_config' }}",
- runtime_version="{{ 'runtime_version' }}",
- python_version="{{ 'python_version' }}",
- job_dir="{{ 'job_dir' }}",
- service_account="{{ 'service_account' }}",
- hyperparameters="{{ 'hyperparameters' }}",
- impersonation_chain="{{ 'impersonation_chain' }}",
- # Other parameters
- dag_id="test_template_body_templating_dag",
- task_id="test_template_body_templating_task",
- )
- session.add(ti)
- session.commit()
- ti.render_templates()
- task: MLEngineStartTrainingJobOperator = ti.task
- assert task.project_id == "project_id"
- assert task.job_id == "job_id"
- assert task.region == "region"
- assert task.package_uris == "package_uris"
- assert task.training_python_module == "training_python_module"
- assert task.training_args == "training_args"
- assert task.scale_tier == "scale_tier"
- assert task.master_type == "master_type"
- assert task.master_config == "master_config"
- assert task.runtime_version == "runtime_version"
- assert task.python_version == "python_version"
- assert task.job_dir == "job_dir"
- assert task.service_account == "service_account"
- assert task.hyperparameters == "hyperparameters"
- assert task.impersonation_chain == "impersonation_chain"
-
-
-TEST_TASK_ID = "training"
-TEST_JOB_ID = "1234"
-TEST_GCP_PROJECT_ID = "test-project"
-TEST_REGION = "us-central1"
-TEST_RUNTIME_VERSION = "1.15"
-TEST_PYTHON_VERSION = "3.9"
-TEST_JOB_DIR = "gs://example_mlengine_bucket/job-dir"
-TEST_PACKAGE_URIS = ["gs://system-tests-resources/example_gcp_mlengine/trainer-0.1.tar.gz"]
-TEST_TRAINING_PYTHON_MODULE = "trainer.task"
-TEST_TRAINING_ARGS: list[str] = []
-TEST_LABELS = {"job_type": "training", "***-version": "v2-5-0-dev0"}
-
-
-@pytest.mark.db_test
-@patch(MLENGINE_AI_PATH.format("MLEngineHook"))
-def test_async_create_training_job_should_execute_successfully(
- mock_hook, create_task_instance_of_operator, session
-):
- """
- Asserts that a task is deferred and a MLEngineStartTrainingJobTrigger will be fired
- when the MLEngineStartTrainingJobOperator is executed in deferrable mode when deferrable=True.
- """
- mock_hook.return_value.create_job_without_waiting_result.return_value = "test_training"
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- task_id=TEST_TASK_ID,
- project_id=TEST_GCP_PROJECT_ID,
- region=TEST_REGION,
- job_id=TEST_JOB_ID,
- runtime_version=TEST_RUNTIME_VERSION,
- python_version=TEST_PYTHON_VERSION,
- job_dir=TEST_JOB_DIR,
- package_uris=TEST_PACKAGE_URIS,
- training_python_module=TEST_TRAINING_PYTHON_MODULE,
- training_args=TEST_TRAINING_ARGS,
- labels=TEST_LABELS,
- deferrable=True,
- operator_class=MLEngineStartTrainingJobOperator,
- dag_id="test_async_create_training",
- )
-
- with pytest.raises(TaskDeferred) as exc:
- ti.task.execute({"ti": ti})
-
- assert isinstance(exc.value.trigger, MLEngineStartTrainingJobTrigger), (
- "Trigger is not a MLEngineStartTrainingJobTrigger"
- )
-
-
-def test_async_create_training_job_should_throw_exception():
- """Tests that an AirflowException is raised in case of error event"""
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- op = MLEngineStartTrainingJobOperator(
- task_id=TEST_TASK_ID,
- project_id=TEST_GCP_PROJECT_ID,
- region=TEST_REGION,
- job_id=TEST_JOB_ID,
- runtime_version=TEST_RUNTIME_VERSION,
- python_version=TEST_PYTHON_VERSION,
- job_dir=TEST_JOB_DIR,
- package_uris=TEST_PACKAGE_URIS,
- training_python_module=TEST_TRAINING_PYTHON_MODULE,
- training_args=TEST_TRAINING_ARGS,
- labels=TEST_LABELS,
- deferrable=True,
- )
-
- with pytest.raises(AirflowException):
- op.execute_complete(context=None, event={"status": "error", "message": "test failure message"})
-
-
-def create_context(task):
- dag = DAG(dag_id="dag", schedule=None)
- logical_date = datetime(2022, 1, 1, 0, 0, 0)
- dag_run = DagRun(
- dag_id=dag.dag_id,
- logical_date=logical_date,
- run_id=DagRun.generate_run_id(DagRunType.MANUAL, logical_date),
- )
- task_instance = TaskInstance(task=task)
- task_instance.dag_run = dag_run
- task_instance.dag_id = dag.dag_id
- task_instance.xcom_push = MagicMock()
- return {
- "dag": dag,
- "run_id": dag_run.run_id,
- "task": task,
- "ti": task_instance,
- "task_instance": task_instance,
- "logical_date": logical_date,
- }
-
-
-@pytest.mark.db_test
-def test_async_create_training_job_logging_should_execute_successfully(
- create_task_instance_of_operator, session
-):
- """Asserts that logging occurs as expected"""
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- operator_class=MLEngineStartTrainingJobOperator,
- task_id=TEST_TASK_ID,
- project_id=TEST_GCP_PROJECT_ID,
- region=TEST_REGION,
- job_id=TEST_JOB_ID,
- runtime_version=TEST_RUNTIME_VERSION,
- python_version=TEST_PYTHON_VERSION,
- job_dir=TEST_JOB_DIR,
- package_uris=TEST_PACKAGE_URIS,
- training_python_module=TEST_TRAINING_PYTHON_MODULE,
- training_args=TEST_TRAINING_ARGS,
- labels=TEST_LABELS,
- deferrable=True,
- dag_id="test_async_create_training_job_logging",
- )
-
- with mock.patch.object(ti.task.log, "info") as mock_log_info:
- ti.task.execute_complete(
- context={"ti": ti},
- event={"status": "success", "message": "Job completed", "job_id": TEST_TASK_ID},
- )
-
- mock_log_info.assert_called_with("%s completed with response %s ", TEST_TASK_ID, "Job completed")
-
-
-@pytest.mark.db_test
-@patch(MLENGINE_AI_PATH.format("MLEngineHook"))
-def test_async_create_training_job_with_conflict_should_execute_successfully(
- mock_hook, create_task_instance_of_operator, session
-):
- """Test that the operator defers correctly when a conflict occurs."""
-
- mock_hook.return_value.create_job_without_waiting_result.side_effect = HttpError(
- resp=httplib2.Response({"status": "409"}), content=b"some bytes"
- )
- mock_hook.return_value.get_job.return_value = {"job_id": "test_training"}
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- operator_class=MLEngineStartTrainingJobOperator,
- task_id=TEST_TASK_ID,
- project_id=TEST_GCP_PROJECT_ID,
- region=TEST_REGION,
- job_id=TEST_JOB_ID,
- runtime_version=TEST_RUNTIME_VERSION,
- python_version=TEST_PYTHON_VERSION,
- job_dir=TEST_JOB_DIR,
- package_uris=TEST_PACKAGE_URIS,
- training_python_module=TEST_TRAINING_PYTHON_MODULE,
- training_args=TEST_TRAINING_ARGS,
- labels=TEST_LABELS,
- deferrable=True,
- dag_id="test_async_create_training_job_with_conflict",
- )
-
- with pytest.raises(TaskDeferred):
- ti.task.execute({"ti": ti})
-
- mock_hook.assert_called_once_with(
- gcp_conn_id="google_cloud_default",
- impersonation_chain=None,
- )
- mock_hook.return_value.create_job_without_waiting_result.assert_called_once()
-
-
-@pytest.mark.db_test
-def test_async_create_training_job_should_throw_exception_if_job_id_none(
- create_task_instance_of_operator, session
-):
- """Test that the operator throws an exception if the job_id is None."""
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- operator_class=MLEngineStartTrainingJobOperator,
- task_id=TEST_TASK_ID,
- project_id=TEST_GCP_PROJECT_ID,
- region=TEST_REGION,
- job_id=None,
- runtime_version=TEST_RUNTIME_VERSION,
- python_version=TEST_PYTHON_VERSION,
- job_dir=TEST_JOB_DIR,
- package_uris=TEST_PACKAGE_URIS,
- training_python_module=TEST_TRAINING_PYTHON_MODULE,
- training_args=TEST_TRAINING_ARGS,
- labels=TEST_LABELS,
- deferrable=True,
- dag_id="test_async_create_training_job",
- )
-
- with pytest.raises(
- AirflowException, match=r"An unique job id is required for Google MLEngine training job."
- ):
- ti.task.execute({"ti": ti})
-
-
-@pytest.mark.db_test
-def test_async_create_training_job_should_throw_exception_if_project_id_none(
- create_task_instance_of_operator, session
-):
- """Test that the operator throws an exception if the project_id is None."""
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- operator_class=MLEngineStartTrainingJobOperator,
- task_id=TEST_TASK_ID,
- project_id=None,
- region=TEST_REGION,
- job_id=TEST_JOB_ID,
- runtime_version=TEST_RUNTIME_VERSION,
- python_version=TEST_PYTHON_VERSION,
- job_dir=TEST_JOB_DIR,
- package_uris=TEST_PACKAGE_URIS,
- training_python_module=TEST_TRAINING_PYTHON_MODULE,
- training_args=TEST_TRAINING_ARGS,
- labels=TEST_LABELS,
- deferrable=True,
- dag_id="test_async_create_training_job_exception",
- )
-
- with pytest.raises(AirflowException, match=r"Google Cloud project id is required."):
- ti.task.execute({"ti": ti})
-
-
-@pytest.mark.db_test
-def test_async_create_training_job_should_throw_exception_if_custom_none(
- create_task_instance_of_operator, session
-):
- """Test that the operator throws an exception if master_type is None while master_config is provided."""
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- operator_class=MLEngineStartTrainingJobOperator,
- task_id=TEST_TASK_ID,
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- job_id=TEST_JOB_ID,
- runtime_version=TEST_RUNTIME_VERSION,
- python_version=TEST_PYTHON_VERSION,
- job_dir=TEST_JOB_DIR,
- package_uris=TEST_PACKAGE_URIS,
- training_python_module=TEST_TRAINING_PYTHON_MODULE,
- training_args=TEST_TRAINING_ARGS,
- labels=TEST_LABELS,
- master_config={"config": "config"},
- master_type=None,
- deferrable=True,
- dag_id="test_async_create_training_job_custom",
- )
-
- with pytest.raises(AirflowException, match=r"master_type must be set when master_config is provided"):
- ti.task.execute({"ti": ti})
-
-
-@pytest.mark.db_test
-def test_async_create_training_job_should_throw_exception_if_package_none(
- create_task_instance_of_operator, session
-):
- """Test that the operator throws an exception if both package_uris and training_python_module are None."""
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- operator_class=MLEngineStartTrainingJobOperator,
- task_id=TEST_TASK_ID,
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- job_id=TEST_JOB_ID,
- runtime_version=TEST_RUNTIME_VERSION,
- python_version=TEST_PYTHON_VERSION,
- job_dir=TEST_JOB_DIR,
- package_uris=None,
- training_python_module=None,
- training_args=TEST_TRAINING_ARGS,
- labels=TEST_LABELS,
- deferrable=True,
- dag_id="test_async_create_training_job_package",
- )
-
- with pytest.raises(
- AirflowException,
- match=r"Either a Python package with a Python module or a custom "
- r"Docker image should be provided.",
- ):
- ti.task.execute({"ti": ti})
-
-
-@pytest.mark.db_test
-def test_async_create_training_job_should_throw_exception_if_uris_none(
- create_task_instance_of_operator, session
-):
- """Test that the operator throws an exception if package_uris is None."""
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- ti = create_task_instance_of_operator(
- operator_class=MLEngineStartTrainingJobOperator,
- task_id=TEST_TASK_ID,
- project_id=TEST_PROJECT_ID,
- region=TEST_REGION,
- job_id=TEST_JOB_ID,
- runtime_version=TEST_RUNTIME_VERSION,
- python_version=TEST_PYTHON_VERSION,
- job_dir=TEST_JOB_DIR,
- package_uris=None,
- training_python_module=TEST_TRAINING_PYTHON_MODULE,
- training_args=TEST_TRAINING_ARGS,
- labels=TEST_LABELS,
- master_config={"config": "config"},
- master_type="type",
- deferrable=True,
- dag_id="test_async_create_training_job_with_uri",
- )
-
- with pytest.raises(
- AirflowException,
- match=r"Either a Python package with a Python module or a custom "
- r"Docker image should be provided.",
- ):
- ti.task.execute({"ti": ti})
diff --git a/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py b/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py
index 15852be62ab..e0468ff2b47 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py
@@ -1210,12 +1210,8 @@ class TestVertexAIDeleteCustomTrainingJobOperator:
assert task.region == "region"
assert task.project_id == "project_id"
assert task.impersonation_chain == "impersonation-chain"
-
- # Deprecated aliases
- with pytest.warns(AirflowProviderDeprecationWarning):
- assert task.training_pipeline == "training-pipeline-id"
- with pytest.warns(AirflowProviderDeprecationWarning):
- assert task.custom_job == "custom_job_id"
+ assert task.training_pipeline_id == "training-pipeline-id"
+ assert task.custom_job_id == "custom_job_id"
class TestVertexAIListCustomTrainingJobOperator:
@@ -1935,9 +1931,7 @@ class TestVertexAIDeleteAutoMLTrainingJobOperator:
assert task.region == "region"
assert task.project_id == "project-id"
assert task.impersonation_chain == "impersonation-chain"
-
- with pytest.warns(AirflowProviderDeprecationWarning):
- assert task.training_pipeline == "training-pipeline-id"
+ assert task.training_pipeline_id == "training-pipeline-id"
class TestVertexAIListAutoMLTrainingJobOperator:
diff --git a/providers/google/tests/unit/google/cloud/utils/test_mlengine_operator_utils.py b/providers/google/tests/unit/google/cloud/utils/test_mlengine_operator_utils.py
deleted file mode 100644
index 4384eb7fea8..00000000000
--- a/providers/google/tests/unit/google/cloud/utils/test_mlengine_operator_utils.py
+++ /dev/null
@@ -1,305 +0,0 @@
-# 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.
-from __future__ import annotations
-
-import base64
-import json
-from datetime import datetime
-from unittest import mock
-
-import dill
-import pytest
-
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.models import DAG
-from airflow.providers.apache.beam.operators.beam import BeamRunPythonPipelineOperator
-from airflow.providers.google.cloud.hooks.gcs import GCSHook
-from airflow.providers.google.cloud.utils.mlengine_operator_utils import create_evaluate_ops
-
-from tests_common.test_utils.compat import PythonOperator
-
-TASK_PREFIX = "test-task-prefix"
-TASK_PREFIX_PREDICTION = TASK_PREFIX + "-prediction"
-TASK_PREFIX_SUMMARY = TASK_PREFIX + "-summary"
-TASK_PREFIX_VALIDATION = TASK_PREFIX + "-validation"
-DATA_FORMAT = "TEXT"
-INPUT_PATHS = [
- "gs://path/to/input/file.json",
- "gs://path/to/input/file2.json",
- "gs://path/to/input/file3.json",
-]
-PREDICTION_PATH = "gs://path/to/output/predictions.json"
-BATCH_PREDICTION_JOB_ID = "test-batch-prediction-job-id"
-PROJECT_ID = "test-project-id"
-REGION = "test-region"
-DATAFLOW_OPTIONS = {
- "project": "my-gcp-project",
- "zone": "us-central1-f",
- "stagingLocation": "gs://bucket/tmp/dataflow/staging/",
-}
-MODEL_URI = "gs://path/to/model/model"
-MODEL_NAME = "test-model-name"
-VERSION_NAME = "test-version-name"
-DAG_DEFAULT_ARGS = {
- "project_id": PROJECT_ID,
- "region": REGION,
- "model_name": MODEL_NAME,
- "version_name": VERSION_NAME,
- "dataflow_default_options": DATAFLOW_OPTIONS,
-}
-TEST_DAG = DAG(dag_id="test-dag-id", start_date=datetime(2000, 1, 1), default_args=DAG_DEFAULT_ARGS)
-
-
-def get_metric_fn_and_keys():
- import math
-
- def error_and_squared_error(inst):
- label = float(inst["input_label"])
- classes = float(inst["classes"])
- err = abs(classes - label)
- squared_err = math.pow(classes - label, 2)
- return err, squared_err
-
- return error_and_squared_error, ["err", "mse"]
-
-
-METRIC_FN, METRIC_KEYS = get_metric_fn_and_keys()
-METRIC_KEYS_EXPECTED = ",".join(METRIC_KEYS)
-
-
-def validate_err_and_count(summary):
- if summary["err"] > 0.2:
- raise ValueError(f"Too high err>0.2; summary={summary}")
- if summary["mse"] > 0.05:
- raise ValueError(f"Too high mse>0.05; summary={summary}")
- if summary["count"] < 1000:
- raise ValueError(f"Too few instances<1000; summary={summary}")
- return summary
-
-
-class TestMlengineOperatorUtils:
- @mock.patch.object(PythonOperator, "set_upstream")
- @mock.patch.object(BeamRunPythonPipelineOperator, "set_upstream")
- def test_create_evaluate_ops(self, mock_beam_pipeline, mock_python):
- with pytest.warns(AirflowProviderDeprecationWarning):
- result = create_evaluate_ops(
- task_prefix=TASK_PREFIX,
- data_format=DATA_FORMAT,
- input_paths=INPUT_PATHS,
- prediction_path=PREDICTION_PATH,
- metric_fn_and_keys=get_metric_fn_and_keys(),
- validate_fn=validate_err_and_count,
- batch_prediction_job_id=BATCH_PREDICTION_JOB_ID,
- project_id=PROJECT_ID,
- region=REGION,
- dataflow_options=DATAFLOW_OPTIONS,
- model_uri=MODEL_URI,
- )
-
- evaluate_prediction, evaluate_summary, evaluate_validation = result
-
- mock_beam_pipeline.assert_called_once_with(evaluate_prediction)
- mock_python.assert_called_once_with(evaluate_summary)
-
- # importing apache_beam elsewhere modifies the metrics. In order to avoid metrics being modified
- # by apache_beam import happening after importing this test, we retrieve the metrics here rather than
- # at the top of the file.
- METRIC_FN_ENCODED = base64.b64encode(dill.dumps(METRIC_FN, recurse=True)).decode()
-
- assert evaluate_prediction.task_id == TASK_PREFIX_PREDICTION
- assert evaluate_prediction.project_id == PROJECT_ID
- assert evaluate_prediction.job_id == BATCH_PREDICTION_JOB_ID
- assert evaluate_prediction.region == REGION
- assert evaluate_prediction._data_format == DATA_FORMAT
- assert evaluate_prediction.input_paths == INPUT_PATHS
- assert evaluate_prediction.output_path == PREDICTION_PATH
- assert evaluate_prediction.uri == MODEL_URI
-
- assert evaluate_summary.task_id == TASK_PREFIX_SUMMARY
- assert evaluate_summary.default_pipeline_options == DATAFLOW_OPTIONS
- assert evaluate_summary.pipeline_options["prediction_path"] == PREDICTION_PATH
- assert evaluate_summary.pipeline_options["metric_fn_encoded"] == METRIC_FN_ENCODED
- assert evaluate_summary.pipeline_options["metric_keys"] == METRIC_KEYS_EXPECTED
-
- assert evaluate_validation.task_id == TASK_PREFIX_VALIDATION
- assert evaluate_validation.templates_dict["prediction_path"] == PREDICTION_PATH
-
- @mock.patch.object(PythonOperator, "set_upstream")
- @mock.patch.object(BeamRunPythonPipelineOperator, "set_upstream")
- def test_create_evaluate_ops_model_and_version_name(self, mock_beam_pipeline, mock_python):
- with pytest.warns(AirflowProviderDeprecationWarning):
- result = create_evaluate_ops(
- task_prefix=TASK_PREFIX,
- data_format=DATA_FORMAT,
- input_paths=INPUT_PATHS,
- prediction_path=PREDICTION_PATH,
- metric_fn_and_keys=get_metric_fn_and_keys(),
- validate_fn=validate_err_and_count,
- batch_prediction_job_id=BATCH_PREDICTION_JOB_ID,
- project_id=PROJECT_ID,
- region=REGION,
- dataflow_options=DATAFLOW_OPTIONS,
- model_name=MODEL_NAME,
- version_name=VERSION_NAME,
- )
-
- evaluate_prediction, evaluate_summary, evaluate_validation = result
-
- mock_beam_pipeline.assert_called_once_with(evaluate_prediction)
- mock_python.assert_called_once_with(evaluate_summary)
-
- # importing apache_beam elsewhere modifies the metrics. In order to avoid metrics being modified
- # by apache_beam import happening after importing this test, we retrieve the metrics here rather than
- # at the top of the file.
- METRIC_FN_ENCODED = base64.b64encode(dill.dumps(METRIC_FN, recurse=True)).decode()
-
- assert evaluate_prediction.task_id == TASK_PREFIX_PREDICTION
- assert evaluate_prediction.project_id == PROJECT_ID
- assert evaluate_prediction.job_id == BATCH_PREDICTION_JOB_ID
- assert evaluate_prediction.region == REGION
- assert evaluate_prediction._data_format == DATA_FORMAT
- assert evaluate_prediction.input_paths == INPUT_PATHS
- assert evaluate_prediction.output_path == PREDICTION_PATH
- assert evaluate_prediction.model_name == MODEL_NAME
- assert evaluate_prediction.version_name == VERSION_NAME
-
- assert evaluate_summary.task_id == TASK_PREFIX_SUMMARY
- assert evaluate_summary.default_pipeline_options == DATAFLOW_OPTIONS
- assert evaluate_summary.pipeline_options["prediction_path"] == PREDICTION_PATH
- assert evaluate_summary.pipeline_options["metric_fn_encoded"] == METRIC_FN_ENCODED
- assert evaluate_summary.pipeline_options["metric_keys"] == METRIC_KEYS_EXPECTED
-
- assert evaluate_validation.task_id == TASK_PREFIX_VALIDATION
- assert evaluate_validation.templates_dict["prediction_path"] == PREDICTION_PATH
-
- @mock.patch.object(PythonOperator, "set_upstream")
- @mock.patch.object(BeamRunPythonPipelineOperator, "set_upstream")
- def test_create_evaluate_ops_dag(self, mock_dataflow, mock_python):
- with pytest.warns(AirflowProviderDeprecationWarning):
- result = create_evaluate_ops(
- task_prefix=TASK_PREFIX,
- data_format=DATA_FORMAT,
- input_paths=INPUT_PATHS,
- prediction_path=PREDICTION_PATH,
- metric_fn_and_keys=get_metric_fn_and_keys(),
- validate_fn=validate_err_and_count,
- batch_prediction_job_id=BATCH_PREDICTION_JOB_ID,
- dag=TEST_DAG,
- )
-
- evaluate_prediction, evaluate_summary, evaluate_validation = result
-
- mock_dataflow.assert_called_once_with(evaluate_prediction)
- mock_python.assert_called_once_with(evaluate_summary)
-
- # importing apache_beam elsewhere modifies the metrics. In order to avoid metrics being modified
- # by apache_beam import happening after importing this test, we retrieve the metrics here rather than
- # at the top of the file.
- METRIC_FN_ENCODED = base64.b64encode(dill.dumps(METRIC_FN, recurse=True)).decode()
-
- assert evaluate_prediction.task_id == TASK_PREFIX_PREDICTION
- assert evaluate_prediction.project_id == PROJECT_ID
- assert evaluate_prediction.job_id == BATCH_PREDICTION_JOB_ID
- assert evaluate_prediction.region == REGION
- assert evaluate_prediction._data_format == DATA_FORMAT
- assert evaluate_prediction.input_paths == INPUT_PATHS
- assert evaluate_prediction.output_path == PREDICTION_PATH
- assert evaluate_prediction.model_name == MODEL_NAME
- assert evaluate_prediction.version_name == VERSION_NAME
-
- assert evaluate_summary.task_id == TASK_PREFIX_SUMMARY
- assert evaluate_summary.default_pipeline_options == DATAFLOW_OPTIONS
- assert evaluate_summary.pipeline_options["prediction_path"] == PREDICTION_PATH
- assert evaluate_summary.pipeline_options["metric_fn_encoded"] == METRIC_FN_ENCODED
- assert evaluate_summary.pipeline_options["metric_keys"] == METRIC_KEYS_EXPECTED
-
- assert evaluate_validation.task_id == TASK_PREFIX_VALIDATION
- assert evaluate_validation.templates_dict["prediction_path"] == PREDICTION_PATH
-
- @pytest.mark.db_test
- @mock.patch.object(GCSHook, "download")
- @mock.patch.object(PythonOperator, "set_upstream")
- @mock.patch.object(BeamRunPythonPipelineOperator, "set_upstream")
- def test_apply_validate_fn(self, mock_beam_pipeline, mock_python, mock_download):
- with pytest.warns(AirflowProviderDeprecationWarning):
- result = create_evaluate_ops(
- task_prefix=TASK_PREFIX,
- data_format=DATA_FORMAT,
- input_paths=INPUT_PATHS,
- prediction_path=PREDICTION_PATH,
- metric_fn_and_keys=get_metric_fn_and_keys(),
- validate_fn=validate_err_and_count,
- batch_prediction_job_id=BATCH_PREDICTION_JOB_ID,
- project_id=PROJECT_ID,
- region=REGION,
- dataflow_options=DATAFLOW_OPTIONS,
- model_uri=MODEL_URI,
- )
-
- _, _, evaluate_validation = result
-
- mock_download.return_value = json.dumps({"err": 0.3, "mse": 0.04, "count": 1100}).encode("utf-8")
- templates_dict = {"prediction_path": PREDICTION_PATH}
- with pytest.raises(ValueError) as ctx:
- evaluate_validation.python_callable(templates_dict=templates_dict)
-
- assert str(ctx.value) == "Too high err>0.2; summary={'err': 0.3, 'mse': 0.04, 'count': 1100}"
- mock_download.assert_called_once_with("path", "to/output/predictions.json/prediction.summary.json")
-
- invalid_prediction_paths = ["://path/to/output/predictions.json", "gs://", ""]
-
- for path in invalid_prediction_paths:
- templates_dict = {"prediction_path": path}
- with pytest.raises(ValueError) as ctx:
- evaluate_validation.python_callable(templates_dict=templates_dict)
- assert str(ctx.value)[:29] == "Wrong format prediction_path:"
-
- def test_invalid_task_prefix(self):
- invalid_task_prefix_values = ["test-task-prefix&", "~test-task-prefix", "test-task(-prefix"]
-
- for invalid_task_prefix_value in invalid_task_prefix_values:
- with pytest.raises(AirflowException):
- create_evaluate_ops(
- task_prefix=invalid_task_prefix_value,
- data_format=DATA_FORMAT,
- input_paths=INPUT_PATHS,
- prediction_path=PREDICTION_PATH,
- metric_fn_and_keys=get_metric_fn_and_keys(),
- validate_fn=validate_err_and_count,
- )
-
- def test_non_callable_metric_fn(self):
- with pytest.raises(AirflowException):
- create_evaluate_ops(
- task_prefix=TASK_PREFIX,
- data_format=DATA_FORMAT,
- input_paths=INPUT_PATHS,
- prediction_path=PREDICTION_PATH,
- metric_fn_and_keys=("error_and_squared_error", ["err", "mse"]),
- validate_fn=validate_err_and_count,
- )
-
- def test_non_callable_validate_fn(self):
- with pytest.raises(AirflowException):
- create_evaluate_ops(
- task_prefix=TASK_PREFIX,
- data_format=DATA_FORMAT,
- input_paths=INPUT_PATHS,
- prediction_path=PREDICTION_PATH,
- metric_fn_and_keys=get_metric_fn_and_keys(),
- validate_fn="validate_err_and_count",
- )