Make Vertex AI evaluation optional to avoid litellm bloat#69749
Open
Vamsi-klu wants to merge 6 commits into
Open
Make Vertex AI evaluation optional to avoid litellm bloat#69749Vamsi-klu wants to merge 6 commits into
Vamsi-klu wants to merge 6 commits into
Conversation
Contributor
Author
|
Reviewers: @shahar1 (google provider CODEOWNER). Breaking change: evaluation extra now optional to avoid litellm bloat. Changelog entry added, provider docs regenerated via update-providers-dependencies. Fixes #69323 Drafted-by: Muse Spark 1.1; reviewed by @Vamsi-klu before posting |
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: apache#69323
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 apache#69323
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.
8d5aab9 to
b54b2d6
Compare
Contributor
|
Quickest fix: git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-leaseAutomated nudge — ignore if you're not ready to rebase. This comment is updated in place on future |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
apache-airflow-providers-googlepulledgoogle-cloud-aiplatform[evaluation]into the base provider install. That extra brings in evaluation-only dependencies such aslitellm,scikit-learn,huggingface-hub, andtokenizersfor users who may only need unrelated Google services.The only code path that needs
vertexai.preview.evaluationis Vertex AI model evaluation throughGenerativeModelHook.get_eval_task(),GenerativeModelHook.run_evaluation(), andRunEvaluationOperator.What Changed
google-cloud-aiplatform>=1.155.0.evaluationprovider extra that installsgoogle-cloud-aiplatform[evaluation]>=1.155.0.AirflowOptionalProviderFeatureExceptionwith an install instruction when evaluation APIs are used without the extra.uv.lock.evaluationextra in the Vertex AI operator docs.Why The Follow-Up Commit Exists
The feedback was correct: the original PR proved the opt-in path but did not directly test the failure mode when the extra is missing, did not document the extra next to
RunEvaluationOperator, and left the lockfile inconsistent with the dependency split.Impact
RunEvaluationOperatoror direct Vertex AI evaluation hook methods must installapache-airflow-providers-google[evaluation].Fixes: #69323
Testing
UV_CACHE_DIR=/tmp/uv-cache-69749 uv run --python 3.11 ruff format providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.pyUV_CACHE_DIR=/tmp/uv-cache-69749 uv run --python 3.11 ruff check --fix providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.pyAIRFLOW_HOME=/tmp/airflow-home-69749 UV_CACHE_DIR=/tmp/uv-cache-69749 uv run --python 3.11 --project providers/google pytest providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py --with-db-init -xvs2 passed, 1 warningAIRFLOW_HOME=/tmp/airflow-home-69749 UV_CACHE_DIR=/tmp/uv-cache-69749 uv run --python 3.11 --project providers/google --extra evaluation pytest providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py -xvs2 passed, 1 warningWas generative AI tooling used to co-author this PR?
Generated-by: Codex (GPT-5) following the guidelines
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting