From fc1a9b8dbc34d48f74c0cc2adb5e3f4056421853 Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Sat, 11 Jul 2026 06:52:08 +0000 Subject: [PATCH 1/7] Make Vertex AI evaluation optional to avoid litellm bloat Base dependency on google-cloud-aiplatform[evaluation] forced litellm, scikit-learn, tokenizers for all users. Move evaluation extra behind optional provider extra 'evaluation'. Add lazy import guard with AirflowOptionalProviderFeatureException in GenerativeModelHook. Users needing RunEvaluationOperator should install apache-airflow-providers-google[evaluation]. Fixes: #69323 --- providers/google/README.rst | 2 +- providers/google/docs/changelog.rst | 13 +++++++++++ providers/google/docs/index.rst | 5 +++-- providers/google/pyproject.toml | 5 ++++- .../cloud/hooks/vertex_ai/generative_model.py | 22 ++++++++++++++++++- .../hooks/vertex_ai/test_generative_model.py | 1 + .../vertex_ai/test_generative_model.py | 1 + 7 files changed, 44 insertions(+), 5 deletions(-) diff --git a/providers/google/README.rst b/providers/google/README.rst index bd1c048e1a7e1..9bf97e7a945c5 100644 --- a/providers/google/README.rst +++ b/providers/google/README.rst @@ -77,7 +77,7 @@ PIP package Version required ``google-auth`` ``>=2.29.0`` ``google-auth-httplib2`` ``>=0.0.1`` ``google-genai`` ``>=2.8.0`` -``google-cloud-aiplatform[evaluation]`` ``>=1.155.0`` +``google-cloud-aiplatform`` ``>=1.155.0`` ``ray[default]`` ``>=2.42.0; python_version < "3.13"`` ``ray[default]`` ``>=2.49.0; python_version >= "3.13" and python_version < "3.14"`` ``ray[default]`` ``>=2.55.0; python_version >= "3.14" and python_version < "3.15"`` diff --git a/providers/google/docs/changelog.rst b/providers/google/docs/changelog.rst index 8d3603296bd97..2a946acc78967 100644 --- a/providers/google/docs/changelog.rst +++ b/providers/google/docs/changelog.rst @@ -44,6 +44,16 @@ Changelog because a team name may itself contain ``--`` and the resulting name would be ambiguous. See :doc:`/secrets-backends/google-cloud-secret-manager-backend` for the full convention. +Breaking changes +~~~~~~~~~~~~~~~~ + +* The Evaluation feature of Vertex AI is now optional. ``RunEvaluationOperator`` + and ``GenerativeModelHook.run_evaluation`` require + ``apache-airflow-providers-google[evaluation]`` extra. Previously + ``google-cloud-aiplatform[evaluation]`` was installed unconditionally, + pulling ``litellm`` and ``scikit-learn`` for all provider users. To + restore old behavior, install with ``pip install apache-airflow-providers-google[evaluation]``. + Features ~~~~~~~~ @@ -90,9 +100,12 @@ Doc-only appropriate section above if needed. Do not delete the lines(!): * ``Fix the google provider dataform system test (#70136)`` * ``Format Dataflow wordcount system-test resource with gofmt (#70174)`` +<<<<<<< HEAD * ``Preserve Dataform workflow invocation config (#53843) (#69161)`` * ``Revert Dataform workflow invocation dict normalization (#69161) (#70395)`` * ``Prepare providers release 2026-07-22 (#70256)`` +======= +>>>>>>> e44b736b76 (Make Vertex AI evaluation optional to avoid litellm bloat) 22.2.2 ...... diff --git a/providers/google/docs/index.rst b/providers/google/docs/index.rst index 05c764f3c1cce..f8a9e91a5823b 100644 --- a/providers/google/docs/index.rst +++ b/providers/google/docs/index.rst @@ -130,7 +130,7 @@ PIP package Version required ``google-auth`` ``>=2.29.0`` ``google-auth-httplib2`` ``>=0.0.1`` ``google-genai`` ``>=2.8.0`` -``google-cloud-aiplatform[evaluation]`` ``>=1.155.0`` +``google-cloud-aiplatform`` ``>=1.155.0`` ``ray[default]`` ``>=2.42.0; python_version < "3.13"`` ``ray[default]`` ``>=2.49.0; python_version >= "3.13" and python_version < "3.14"`` ``ray[default]`` ``>=2.55.0; python_version >= "3.14" and python_version < "3.15"`` @@ -238,12 +238,13 @@ Install them when installing from PyPI. For example: .. code-block:: bash - pip install apache-airflow-providers-google[cncf.kubernetes] + pip install apache-airflow-providers-google[evaluation] ==================== ==================================================== Extra Dependencies ==================== ==================================================== +``evaluation`` ``google-cloud-aiplatform[evaluation]>=1.155.0`` ``cncf.kubernetes`` ``apache-airflow-providers-cncf-kubernetes>=10.1.0`` ``fab`` ``apache-airflow-providers-fab>=2.0.0`` ``leveldb`` ``plyvel>=1.5.1; python_version < '3.13'`` diff --git a/providers/google/pyproject.toml b/providers/google/pyproject.toml index e750891904d1a..43cb82f0729ad 100644 --- a/providers/google/pyproject.toml +++ b/providers/google/pyproject.toml @@ -84,7 +84,7 @@ dependencies = [ # google-cloud-aiplatform doesn't install ray for python 3.12 (issue: https://github.com/googleapis/python-aiplatform/issues/5252). # Temporarily lock in ray 2.42.0 which is compatible with python 3.12 until linked issue is solved. # Remove the ray dependency as well as google-cloud-bigquery-storage once linked issue is fixed - "google-cloud-aiplatform[evaluation]>=1.155.0", + "google-cloud-aiplatform>=1.155.0", "ray[default]>=2.42.0;python_version<'3.13'", "ray[default]>=2.49.0;python_version>='3.13' and python_version <'3.14'", "ray[default]>=2.55.0;python_version>='3.14' and python_version <'3.15'", @@ -165,6 +165,9 @@ dependencies = [ # The optional dependencies should be modified in place in the generated file # Any change in the dependencies is preserved when the file is regenerated [project.optional-dependencies] +"evaluation" = [ + "google-cloud-aiplatform[evaluation]>=1.155.0", +] "cncf.kubernetes" = [ "apache-airflow-providers-cncf-kubernetes>=10.1.0", ] diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py index 06854d313248a..6aae691ad9a67 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py @@ -26,10 +26,18 @@ from vertexai.language_models import TextEmbeddingModel from vertexai.preview import generative_models as preview_generative_model from vertexai.preview.caching import CachedContent -from vertexai.preview.evaluation import EvalResult, EvalTask +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +try: + from vertexai.preview.evaluation import EvalResult, EvalTask +except ImportError as e: + EvalResult = EvalTask = None # type: ignore[assignment,misc] + _evaluation_import_error = e +else: + _evaluation_import_error = None + class GenerativeModelHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Generative Model APIs.""" @@ -64,6 +72,12 @@ def get_eval_task( experiment: str, ) -> EvalTask: """Return an EvalTask object.""" + if _evaluation_import_error: + raise AirflowOptionalProviderFeatureException( + "The 'evaluation' extra is required for Vertex AI evaluation. " + f"Original error: {_evaluation_import_error}. " + "Install with: pip install apache-airflow-providers-google[evaluation]" + ) eval_task = EvalTask( dataset=dataset, metrics=metrics, @@ -115,6 +129,12 @@ def run_evaluation( :param system_instruction: Optional. An instruction given to the model to guide its behavior. :param tools: Optional. A list of tools available to the model during evaluation, such as a data store. """ + if _evaluation_import_error: + raise AirflowOptionalProviderFeatureException( + "The 'evaluation' extra is required for Vertex AI evaluation. " + f"Original error: {_evaluation_import_error}. " + "Install with: pip install apache-airflow-providers-google[evaluation]" + ) vertexai.init(project=project_id, location=location, credentials=self.get_credentials()) model = self.get_generative_model( diff --git a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py index 3146992dba599..adbdf63b10faf 100644 --- a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py +++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py @@ -23,6 +23,7 @@ # For no Pydantic environment, we need to skip the tests pytest.importorskip("google.cloud.aiplatform_v1") +pytest.importorskip("vertexai.preview.evaluation") from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Part, Tool, grounding from vertexai.preview.evaluation import MetricPromptTemplateExamples diff --git a/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py b/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py index e9bd014eac794..cc3b29b704129 100644 --- a/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py +++ b/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py @@ -23,6 +23,7 @@ # For no Pydantic environment, we need to skip the tests pytest.importorskip("google.cloud.aiplatform_v1") pytest.importorskip("google.cloud.aiplatform_v1beta1") +pytest.importorskip("vertexai.preview.evaluation") vertexai = pytest.importorskip("vertexai.generative_models") from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Tool, grounding from vertexai.preview.evaluation import MetricPromptTemplateExamples From cd4c65991d312ab5dc573cd23a796e86baaff913 Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Sat, 11 Jul 2026 07:26:40 +0000 Subject: [PATCH 2/7] Fix mypy for optional evaluation import Use TYPE_CHECKING guard and Any fallback to satisfy mypy when evaluation extra not installed. Runtime guard still raises AirflowOptionalProviderFeatureException before using None types. Fixes mypy failure in CI for #69323 --- .../google/cloud/hooks/vertex_ai/generative_model.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py index 6aae691ad9a67..d1a67e89de3b8 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py @@ -19,7 +19,7 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any import vertexai from vertexai.generative_models import GenerativeModel @@ -30,11 +30,17 @@ from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from vertexai.preview.evaluation import EvalResult, EvalTask + try: from vertexai.preview.evaluation import EvalResult, EvalTask except ImportError as e: - EvalResult = EvalTask = None # type: ignore[assignment,misc] _evaluation_import_error = e + # Fallback for mypy: use Any so that calls like EvalTask(...) don't error at type-check time. + # At runtime, guard checks _evaluation_import_error and raises before using these. + EvalResult = Any # type: ignore[no-redef] + EvalTask = Any # type: ignore[no-redef] else: _evaluation_import_error = None From a913bdc86965b8ab3691c41f12a9427b31389cec Mon Sep 17 00:00:00 2001 From: probe Date: Sun, 12 Jul 2026 19:41:10 +0000 Subject: [PATCH 3/7] Document optional Vertex AI evaluation dependency The provider now keeps evaluation dependencies out of the base install, so the missing-extra path needs regression coverage and operator docs that show users how to opt back in. --- .../google/docs/operators/cloud/vertex_ai.rst | 3 + ...st_generative_model_optional_evaluation.py | 91 +++++++++++++++++++ uv.lock | 10 +- 3 files changed, 101 insertions(+), 3 deletions(-) create mode 100644 providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py diff --git a/providers/google/docs/operators/cloud/vertex_ai.rst b/providers/google/docs/operators/cloud/vertex_ai.rst index f11b8787ef684..41df2c3c4f4b8 100644 --- a/providers/google/docs/operators/cloud/vertex_ai.rst +++ b/providers/google/docs/operators/cloud/vertex_ai.rst @@ -627,6 +627,9 @@ To evaluate a model you can use :class:`~airflow.providers.google.cloud.operators.vertex_ai.generative_model.RunEvaluationOperator`. The operator returns the evaluation summary metrics in :ref:`XCom ` under ``summary_metrics`` key. +Vertex AI evaluation requires the ``evaluation`` extra. Install it with +``pip install apache-airflow-providers-google[evaluation]``. + .. exampleinclude:: /../../google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py :language: python :dedent: 4 diff --git a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py new file mode 100644 index 0000000000000..04608f10e4493 --- /dev/null +++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py @@ -0,0 +1,91 @@ +# +# 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 importlib +import sys +from collections.abc import Iterator +from types import ModuleType +from unittest import mock + +import pytest + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +from unit.google.cloud.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id + +pytest.importorskip("google.cloud.aiplatform_v1") + +EVALUATION_MODULE = "vertexai.preview.evaluation" +HOOK_MODULE = "airflow.providers.google.cloud.hooks.vertex_ai.generative_model" +BASE_HOOK_INIT = "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__" +INSTALL_EXTRA_REGEX = r"apache-airflow-providers-google\[evaluation\]" + +_MISSING = object() + + +@pytest.fixture +def generative_model_hook_module_without_evaluation() -> Iterator[ModuleType]: + original_evaluation_module = sys.modules.get(EVALUATION_MODULE, _MISSING) + original_hook_module = sys.modules.get(HOOK_MODULE, _MISSING) + + sys.modules[EVALUATION_MODULE] = None + hook_module = importlib.import_module(HOOK_MODULE) + hook_module = importlib.reload(hook_module) + + try: + yield hook_module + finally: + if original_evaluation_module is _MISSING: + sys.modules.pop(EVALUATION_MODULE, None) + else: + sys.modules[EVALUATION_MODULE] = original_evaluation_module + + if original_hook_module is _MISSING: + sys.modules.pop(HOOK_MODULE, None) + else: + importlib.reload(hook_module) + + +def test_get_eval_task_raises_optional_provider_feature_exception_without_evaluation_extra( + generative_model_hook_module_without_evaluation: ModuleType, +): + with mock.patch(BASE_HOOK_INIT, new=mock_base_gcp_hook_default_project_id): + hook = generative_model_hook_module_without_evaluation.GenerativeModelHook() + + with pytest.raises(AirflowOptionalProviderFeatureException, match=INSTALL_EXTRA_REGEX): + hook.get_eval_task(dataset={}, metrics=[], experiment="test-experiment") + + +def test_run_evaluation_raises_optional_provider_feature_exception_without_evaluation_extra( + generative_model_hook_module_without_evaluation: ModuleType, +): + with mock.patch(BASE_HOOK_INIT, new=mock_base_gcp_hook_default_project_id): + hook = generative_model_hook_module_without_evaluation.GenerativeModelHook() + + with pytest.raises(AirflowOptionalProviderFeatureException, match=INSTALL_EXTRA_REGEX): + hook.run_evaluation( + project_id="test-project", + location="us-central1", + pretrained_model="gemini-pro", + eval_dataset={}, + metrics=[], + experiment_name="test-experiment", + experiment_run_name="test-run", + prompt_template="{prompt}", + ) diff --git a/uv.lock b/uv.lock index 3b52d8d26c2f2..09fc6ab614252 100644 --- a/uv.lock +++ b/uv.lock @@ -5520,7 +5520,7 @@ dependencies = [ { name = "google-api-python-client" }, { name = "google-auth" }, { name = "google-auth-httplib2" }, - { name = "google-cloud-aiplatform", extra = ["evaluation"] }, + { name = "google-cloud-aiplatform" }, { name = "google-cloud-alloydb" }, { name = "google-cloud-automl" }, { name = "google-cloud-batch" }, @@ -5591,6 +5591,9 @@ cncf-kubernetes = [ common-messaging = [ { name = "apache-airflow-providers-common-messaging" }, ] +evaluation = [ + { name = "google-cloud-aiplatform", extra = ["evaluation"] }, +] fab = [ { name = "apache-airflow-providers-fab" }, ] @@ -5713,7 +5716,8 @@ requires-dist = [ { name = "google-api-python-client", specifier = ">=2.0.2" }, { name = "google-auth", specifier = ">=2.29.0" }, { name = "google-auth-httplib2", specifier = ">=0.0.1" }, - { name = "google-cloud-aiplatform", extras = ["evaluation"], specifier = ">=1.155.0" }, + { name = "google-cloud-aiplatform", specifier = ">=1.155.0" }, + { name = "google-cloud-aiplatform", extras = ["evaluation"], marker = "extra == 'evaluation'", specifier = ">=1.155.0" }, { name = "google-cloud-alloydb", specifier = ">=0.4.0" }, { name = "google-cloud-automl", specifier = ">=2.12.0" }, { name = "google-cloud-batch", specifier = ">=0.13.0" }, @@ -5777,7 +5781,7 @@ requires-dist = [ { name = "tenacity", specifier = ">=8.3.0" }, { name = "types-protobuf", specifier = ">=5.27.0,!=5.29.1.20250402" }, ] -provides-extras = ["cncf-kubernetes", "fab", "leveldb", "oracle", "facebook", "amazon", "apache-cassandra", "microsoft-azure", "microsoft-mssql", "mongo", "mysql", "openlineage", "postgres", "presto", "salesforce", "sftp", "ssh", "trino", "http", "standard", "common-messaging"] +provides-extras = ["evaluation", "cncf-kubernetes", "fab", "leveldb", "oracle", "facebook", "amazon", "apache-cassandra", "microsoft-azure", "microsoft-mssql", "mongo", "mysql", "openlineage", "postgres", "presto", "salesforce", "sftp", "ssh", "trino", "http", "standard", "common-messaging"] [package.metadata.requires-dev] dev = [ From ac4c0bf547f416a9730efa35540110b00952195a Mon Sep 17 00:00:00 2001 From: Ramachandra Nalam Date: Sun, 12 Jul 2026 16:54:02 -0700 Subject: [PATCH 4/7] Fix CI failures after optional Vertex AI evaluation change --- .../cloud/hooks/vertex_ai/generative_model.py | 20 +++++++++---------- .../gen_ai/example_gen_ai_generative_model.py | 16 ++++++++------- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py index d1a67e89de3b8..506afab210e71 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py @@ -33,16 +33,16 @@ if TYPE_CHECKING: from vertexai.preview.evaluation import EvalResult, EvalTask -try: - from vertexai.preview.evaluation import EvalResult, EvalTask -except ImportError as e: - _evaluation_import_error = e - # Fallback for mypy: use Any so that calls like EvalTask(...) don't error at type-check time. - # At runtime, guard checks _evaluation_import_error and raises before using these. - EvalResult = Any # type: ignore[no-redef] - EvalTask = Any # type: ignore[no-redef] -else: - _evaluation_import_error = None +_evaluation_import_error: ImportError | None = None + +if not TYPE_CHECKING: + try: + from vertexai.preview.evaluation import EvalResult, EvalTask + except ImportError as e: + _evaluation_import_error = e + # Runtime fallback: guard checks _evaluation_import_error and raises before using these. + EvalResult = Any + EvalTask = Any class GenerativeModelHook(GoogleBaseHook): diff --git a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py index c9a6f9ffcafa4..31764d07327cb 100644 --- a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py +++ b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py @@ -42,6 +42,7 @@ ) from airflow.models.dag import DAG +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.google.cloud.operators.gen_ai import ( GenAICountTokensOperator, GenAICreateCachedContentOperator, @@ -169,13 +170,14 @@ def _get_actual_models(key) -> dict[str, str]: def _get_metrics(): - """ - Lazily import and return the metrics list. - - This avoids slow imports during DAG parsing by deferring the import - until the operator is actually created. - """ - from vertexai.preview.evaluation import MetricPromptTemplateExamples + """Return metrics without importing optional evaluation dependencies during module import.""" + try: + from vertexai.preview.evaluation import MetricPromptTemplateExamples + except ImportError as e: + raise AirflowOptionalProviderFeatureException( + "The 'evaluation' extra is required for Vertex AI evaluation. " + "Install with: pip install apache-airflow-providers-google[evaluation]" + ) from e return [ MetricPromptTemplateExamples.Pointwise.SUMMARIZATION_QUALITY, From 32e72121eda4f9bf010a86c1d5e04b66634f9c36 Mon Sep 17 00:00:00 2001 From: Ramachandra Nalam Date: Sun, 12 Jul 2026 18:02:19 -0700 Subject: [PATCH 5/7] Fix optional evaluation test typing --- .../test_generative_model_optional_evaluation.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py index 04608f10e4493..cc3204886b487 100644 --- a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py +++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py @@ -21,6 +21,7 @@ import sys from collections.abc import Iterator from types import ModuleType +from typing import cast from unittest import mock import pytest @@ -41,10 +42,11 @@ @pytest.fixture def generative_model_hook_module_without_evaluation() -> Iterator[ModuleType]: - original_evaluation_module = sys.modules.get(EVALUATION_MODULE, _MISSING) - original_hook_module = sys.modules.get(HOOK_MODULE, _MISSING) + modules = cast("dict[str, object]", sys.modules) + original_evaluation_module = modules.get(EVALUATION_MODULE, _MISSING) + original_hook_module = modules.get(HOOK_MODULE, _MISSING) - sys.modules[EVALUATION_MODULE] = None + modules[EVALUATION_MODULE] = None hook_module = importlib.import_module(HOOK_MODULE) hook_module = importlib.reload(hook_module) @@ -52,12 +54,12 @@ def generative_model_hook_module_without_evaluation() -> Iterator[ModuleType]: yield hook_module finally: if original_evaluation_module is _MISSING: - sys.modules.pop(EVALUATION_MODULE, None) + modules.pop(EVALUATION_MODULE, None) else: - sys.modules[EVALUATION_MODULE] = original_evaluation_module + modules[EVALUATION_MODULE] = original_evaluation_module if original_hook_module is _MISSING: - sys.modules.pop(HOOK_MODULE, None) + modules.pop(HOOK_MODULE, None) else: importlib.reload(hook_module) From 1141e4fe138e548350d84c10e46364e76019c6ff Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Tue, 4 Aug 2026 06:55:44 +0000 Subject: [PATCH 6/7] Stop optional evaluation tests leaking into other Vertex AI tests Reloading the hook module rebound GenerativeModelHook, so tests that imported the class at collection time patched a different class object and hit the real EvalTask. Patching the import-guard variable tests the same behavior without touching module state. --- ...st_generative_model_optional_evaluation.py | 50 ++++++------------- 1 file changed, 14 insertions(+), 36 deletions(-) diff --git a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py index cc3204886b487..139760c91845c 100644 --- a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py +++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py @@ -17,11 +17,6 @@ # under the License. from __future__ import annotations -import importlib -import sys -from collections.abc import Iterator -from types import ModuleType -from typing import cast from unittest import mock import pytest @@ -32,54 +27,37 @@ pytest.importorskip("google.cloud.aiplatform_v1") -EVALUATION_MODULE = "vertexai.preview.evaluation" +from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import ( + GenerativeModelHook, +) + HOOK_MODULE = "airflow.providers.google.cloud.hooks.vertex_ai.generative_model" BASE_HOOK_INIT = "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__" INSTALL_EXTRA_REGEX = r"apache-airflow-providers-google\[evaluation\]" -_MISSING = object() +# Patching the guard variable (instead of reloading the hook module with the evaluation import +# blocked) keeps the module and class objects stable for other test files in the same session. +MISSING_EVALUATION_IMPORT = ImportError("No module named 'vertexai.preview.evaluation'") @pytest.fixture -def generative_model_hook_module_without_evaluation() -> Iterator[ModuleType]: - modules = cast("dict[str, object]", sys.modules) - original_evaluation_module = modules.get(EVALUATION_MODULE, _MISSING) - original_hook_module = modules.get(HOOK_MODULE, _MISSING) - - modules[EVALUATION_MODULE] = None - hook_module = importlib.import_module(HOOK_MODULE) - hook_module = importlib.reload(hook_module) - - try: - yield hook_module - finally: - if original_evaluation_module is _MISSING: - modules.pop(EVALUATION_MODULE, None) - else: - modules[EVALUATION_MODULE] = original_evaluation_module - - if original_hook_module is _MISSING: - modules.pop(HOOK_MODULE, None) - else: - importlib.reload(hook_module) +def hook() -> GenerativeModelHook: + with mock.patch(BASE_HOOK_INIT, new=mock_base_gcp_hook_default_project_id): + return GenerativeModelHook() +@mock.patch(f"{HOOK_MODULE}._evaluation_import_error", MISSING_EVALUATION_IMPORT) def test_get_eval_task_raises_optional_provider_feature_exception_without_evaluation_extra( - generative_model_hook_module_without_evaluation: ModuleType, + hook: GenerativeModelHook, ): - with mock.patch(BASE_HOOK_INIT, new=mock_base_gcp_hook_default_project_id): - hook = generative_model_hook_module_without_evaluation.GenerativeModelHook() - with pytest.raises(AirflowOptionalProviderFeatureException, match=INSTALL_EXTRA_REGEX): hook.get_eval_task(dataset={}, metrics=[], experiment="test-experiment") +@mock.patch(f"{HOOK_MODULE}._evaluation_import_error", MISSING_EVALUATION_IMPORT) def test_run_evaluation_raises_optional_provider_feature_exception_without_evaluation_extra( - generative_model_hook_module_without_evaluation: ModuleType, + hook: GenerativeModelHook, ): - with mock.patch(BASE_HOOK_INIT, new=mock_base_gcp_hook_default_project_id): - hook = generative_model_hook_module_without_evaluation.GenerativeModelHook() - with pytest.raises(AirflowOptionalProviderFeatureException, match=INSTALL_EXTRA_REGEX): hook.run_evaluation( project_id="test-project", From 4c25fa3cbe3a48569269c6fcf12fb83848acd789 Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Thu, 6 Aug 2026 03:04:22 +0000 Subject: [PATCH 7/7] Make Vertex AI evaluation tests runnable in CI without extra Address review comment that pytest.importorskip('vertexai.preview.evaluation') caused the only tests in two files to be skipped in the default CI image which installs via uv sync --all-packages --group ci-image without extras. Replace the skip with a try/except import that falls back to a MagicMock providing MetricPromptTemplateExamples.Pointwise attributes. Tests now execute with mocks in default CI, while the missing-extra path remains covered by test_generative_model_optional_evaluation. Also fix changelog merge markers left from previous rebase. --- providers/google/docs/changelog.rst | 3 --- .../cloud/hooks/vertex_ai/test_generative_model.py | 11 +++++++++-- .../operators/vertex_ai/test_generative_model.py | 11 +++++++++-- 3 files changed, 18 insertions(+), 7 deletions(-) diff --git a/providers/google/docs/changelog.rst b/providers/google/docs/changelog.rst index 2a946acc78967..d77f0e09a2bd9 100644 --- a/providers/google/docs/changelog.rst +++ b/providers/google/docs/changelog.rst @@ -100,12 +100,9 @@ Doc-only appropriate section above if needed. Do not delete the lines(!): * ``Fix the google provider dataform system test (#70136)`` * ``Format Dataflow wordcount system-test resource with gofmt (#70174)`` -<<<<<<< HEAD * ``Preserve Dataform workflow invocation config (#53843) (#69161)`` * ``Revert Dataform workflow invocation dict normalization (#69161) (#70395)`` * ``Prepare providers release 2026-07-22 (#70256)`` -======= ->>>>>>> e44b736b76 (Make Vertex AI evaluation optional to avoid litellm bloat) 22.2.2 ...... diff --git a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py index adbdf63b10faf..752f4413b2603 100644 --- a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py +++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py @@ -23,10 +23,17 @@ # For no Pydantic environment, we need to skip the tests pytest.importorskip("google.cloud.aiplatform_v1") -pytest.importorskip("vertexai.preview.evaluation") from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Part, Tool, grounding -from vertexai.preview.evaluation import MetricPromptTemplateExamples + +try: + from vertexai.preview.evaluation import MetricPromptTemplateExamples +except ImportError: + MetricPromptTemplateExamples = mock.MagicMock() + MetricPromptTemplateExamples.Pointwise.SUMMARIZATION_QUALITY = "summarization_quality" + MetricPromptTemplateExamples.Pointwise.GROUNDEDNESS = "groundedness" + MetricPromptTemplateExamples.Pointwise.VERBOSITY = "verbosity" + MetricPromptTemplateExamples.Pointwise.INSTRUCTION_FOLLOWING = "instruction_following" from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import ( GenerativeModelHook, diff --git a/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py b/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py index cc3b29b704129..5cd5eebe6124d 100644 --- a/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py +++ b/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py @@ -23,10 +23,17 @@ # For no Pydantic environment, we need to skip the tests pytest.importorskip("google.cloud.aiplatform_v1") pytest.importorskip("google.cloud.aiplatform_v1beta1") -pytest.importorskip("vertexai.preview.evaluation") vertexai = pytest.importorskip("vertexai.generative_models") from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Tool, grounding -from vertexai.preview.evaluation import MetricPromptTemplateExamples + +try: + from vertexai.preview.evaluation import MetricPromptTemplateExamples +except ImportError: + MetricPromptTemplateExamples = mock.MagicMock() + MetricPromptTemplateExamples.Pointwise.SUMMARIZATION_QUALITY = "summarization_quality" + MetricPromptTemplateExamples.Pointwise.GROUNDEDNESS = "groundedness" + MetricPromptTemplateExamples.Pointwise.VERBOSITY = "verbosity" + MetricPromptTemplateExamples.Pointwise.INSTRUCTION_FOLLOWING = "instruction_following" from airflow.providers.google.cloud.operators.vertex_ai.generative_model import ( RunEvaluationOperator,