Make Vertex AI evaluation optional to avoid litellm bloat - #69749
Make Vertex AI evaluation optional to avoid litellm bloat#69749Vamsi-klu wants to merge 16 commits into
Conversation
boring-cyborg
Bot
added
area:providers
kind:documentation
provider:google
|
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 |
8d5aab9 to
b54b2d6
Compare
|
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 |
|
@Vamsi-klu This PR has been converted to draft because it does not yet meet our Pull Request quality criteria. Issues found:
What to do next:
Converting a PR to draft is not a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates. If you have questions, feel free to ask on the Airflow Slack. Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you. |
1506a46 to
ce7370f
Compare
5820c69 to
9ee89ab
Compare
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.
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.
9ee89ab to
1141e4f
Compare
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.
Reworks the optional Vertex AI evaluation change on top of apache#71273 (2c9d91e), which had removed the google-cloud-aiplatform[evaluation] extra from the base deps and inlined the evaluation requirements directly instead, as a workaround for the litellm security range tracked in googleapis/python-aiplatform#7057. Rather than reverting that, the inlined block -- including the litellm python_version markers and the tracking comment -- moves verbatim into a new "evaluation" entry under [project.optional-dependencies], so base installs stop pulling litellm and scikit-learn while the 3.14 security pin is preserved for anyone installing the extra. providers/google/README.rst and docs/index.rst tables are regenerated to match, and uv.lock is regenerated with the pinned uv 0.11.29 (uv lock --check passes; note main's lock was already stale, so the diff also picks up a few unrelated requires-dist specifiers).
…-clean 22.3.0 shipped on 2026-08-08, so the breaking-change note moves into a new 23.0.0 section (major bump, next release) instead of sitting inside an already-released section. Collapses the two copies of the "install the evaluation extra" raise into a single _raise_if_evaluation_unavailable() helper, so the message and the install hint live in one place, and chains the original ImportError as the cause. ImportError (not ModuleNotFoundError) stays the caught type, so a partially installed evaluation stack gets the same hint rather than a raw traceback. EvalResult is only ever an annotation, so it is now imported under TYPE_CHECKING alone; the dead Any fallbacks are gone. RunEvaluationOperator has no evaluation-specific code of its own -- it calls GenerativeModelHook.run_evaluation -- so the single hook-side guard is what surfaces the exception from the operator too. Fixes the MyPy providers failure: assigning MagicMock over the imported MetricPromptTemplateExamples is "Cannot assign to a type" [misc]. The fallback is now built in one expression with a targeted ignore, and still engages only on ImportError, so environments with the real extra keep testing the real metric values.
|
Rebased onto current main, so this is no longer conflicting. The rework sits on top of #71273 rather than reverting it. The evaluation dependencies that PR inlined into the base deps, including the litellm version markers and the tracking comment, are moved verbatim into a new One thing to flag: |
The PR was blocked as DIRTY after main restored google-cloud-aiplatform[evaluation] on the base install. Keep evaluation optional so the provider does not pull litellm, and take main's 1.164.0 floor. Co-authored-by: Cursor <cursoragent@cursor.com>
CI images do not install the Google provider evaluation extra, so the module-level import guard is set before the happy-path test mocks apply and the suite fails on a missing sklearn install.
…on-optional-69323
CI images do not install the Google evaluation extra, so the happy-path evaluation tests fail on a missing sklearn import unless the optional extra guard is exercised for that exact ImportError.
Fork workflows were waiting on approval after a bot push.
…on-optional-69323 # Conflicts: # providers/google/docs/changelog.rst # uv.lock
4 participants
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
Important
🛠️ Maintainer triage note for @Vamsi-klu · by
@potiuk· 2026-08-13 12:55 UTCHelpful heads-up from the maintainers — please address before this PR can be reviewed:
Full list of what we check: Pull Request quality criteria.
The ball is in your court — you've been assigned to this PR. Fix the above, then mark it Ready for review.
Automated triage — may be imperfect; a maintainer takes the next look.