From 2a6d1bbe90630d3106d97a2d6c3b54de77c8e70f Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Sat, 11 Jul 2026 06:52:08 +0000 Subject: [PATCH 1/3] 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 35bc9fe5f23ea..4e5b411787f66 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 6c0993664a4ef..d35f225a894b1 100644 --- a/providers/google/docs/changelog.rst +++ b/providers/google/docs/changelog.rst @@ -27,6 +27,19 @@ Changelog --------- +22.3.0 +...... + +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]``. + 22.2.2 ...... diff --git a/providers/google/docs/index.rst b/providers/google/docs/index.rst index e3bec173806d8..70b026aa208be 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 1698491c65ff8..56b0e5dbfcebf 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'", @@ -156,6 +156,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 7c1a1a9c192cd5b8442ec3901c1c265e1fc467a1 Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Sat, 11 Jul 2026 07:26:40 +0000 Subject: [PATCH 2/3] 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 b54b2d67ee30789c2a6e6213c4d0d4e01ef9b6fd Mon Sep 17 00:00:00 2001 From: probe Date: Sun, 12 Jul 2026 19:41:10 +0000 Subject: [PATCH 3/3] 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 2973c8da2681a..1a4b3f81c94e0 100644 --- a/uv.lock +++ b/uv.lock @@ -5477,7 +5477,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" }, @@ -5548,6 +5548,9 @@ cncf-kubernetes = [ common-messaging = [ { name = "apache-airflow-providers-common-messaging" }, ] +evaluation = [ + { name = "google-cloud-aiplatform", extra = ["evaluation"] }, +] fab = [ { name = "apache-airflow-providers-fab" }, ] @@ -5670,7 +5673,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" }, @@ -5734,7 +5738,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 = [