Skip to content

Make Vertex AI evaluation optional to avoid litellm bloat - #69749

Open
Vamsi-klu wants to merge 16 commits into
apache:mainfrom
Vamsi-klu:fix/google-evaluation-optional-69323
Open

Make Vertex AI evaluation optional to avoid litellm bloat#69749
Vamsi-klu wants to merge 16 commits into
apache:mainfrom
Vamsi-klu:fix/google-evaluation-optional-69323

Conversation

@Vamsi-klu

@Vamsi-kluVamsi-klu commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

What is the change?

The google provider no longer installs google-cloud-aiplatform[evaluation] by default. The base dependency drops to google-cloud-aiplatform>=1.155.0, and a new evaluation provider extra pulls the evaluation stack for the users who actually need it.

Why did I do it?

Fixes: #69323

The base install dragged in evaluation-only dependencies such as litellm, scikit-learn, huggingface-hub, and tokenizers. The only code that needs vertexai.preview.evaluation is GenerativeModelHook.get_eval_task(), GenerativeModelHook.run_evaluation(), and RunEvaluationOperator.

How did I do it?

I split the dependency in pyproject.toml, added an import guard that raises AirflowOptionalProviderFeatureException with an install instruction when evaluation APIs are used without the extra, regenerated uv.lock and the provider dependency docs, documented the extra next to RunEvaluationOperator, and added a changelog note for the breaking dependency split.

What's the impact?

Base google provider installs get lighter. Users of RunEvaluationOperator or the evaluation hook methods must install apache-airflow-providers-google[evaluation]. Everyone else sees no runtime behavior change.

What's the test plan?

test_generative_model_optional_evaluation.py covers the missing-extra failure path, including a missing-sklearn case, and runs without the extra installed. The existing hook and operator evaluation tests run with --extra evaluation. Both passed locally with the commands recorded in the commit history, and provider CI runs them.


Was generative AI tooling used to co-author this PR?
  • Yes — Codex (GPT-5)

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 UTC

Helpful heads-up from the maintainers — please address before this PR can be reviewed:

  • Provider tests. See docs.

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.

@boring-cyborgboring-cyborgBot added area:providers kind:documentation provider:google Google (including GCP) related issues labels Jul 11, 2026
@Vamsi-klu

Copy link
Copy Markdown
ContributorAuthor

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

@Vamsi-klu
Vamsi-kluforce-pushed the fix/google-evaluation-optional-69323 branch from 8d5aab9 to b54b2d6CompareJuly 12, 2026 19:56
@Vamsi-klu
Vamsi-klu marked this pull request as ready for review July 13, 2026 03:36
@Vamsi-klu
Vamsi-klu requested a review from shahar1 as a code ownerJuly 13, 2026 03:36
@github-actions

github-actionsBot commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

uv.lock on main just moved via commit fd0f7e2 and this PR currently conflicts.

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-lease

Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future uv.lock bumps.

@potiuk
potiuk marked this pull request as draft July 20, 2026 11:54
@potiuk

Copy link
Copy Markdown
Member

@Vamsi-klu This PR has been converted to draft because it does not yet meet our Pull Request quality criteria.

Issues found:

  • Merge conflicts: This PR has merge conflicts with main and cannot be merged as-is. Rebase locally and resolve them: git fetch upstream main && git rebase upstream/main, resolve the conflicts, then force-push.

Note: Your branch is 147 commits behind main. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.

What to do next:

  • Fix each issue listed above.
  • Make sure static checks pass locally (prek run --from-ref main --stage pre-commit).
  • Mark the PR as "Ready for review" when you're done.

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.

@Vamsi-klu
Vamsi-klu marked this pull request as ready for review July 20, 2026 12:03
@Vamsi-klu
Vamsi-kluforce-pushed the fix/google-evaluation-optional-69323 branch from 1506a46 to ce7370fCompareJuly 20, 2026 14:07
@Vamsi-klu
Vamsi-kluforce-pushed the fix/google-evaluation-optional-69323 branch 2 times, most recently from 5820c69 to 9ee89abCompareAugust 1, 2026 04:36
Vamsi-kluand others added 6 commits August 4, 2026 06:37
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.
@Vamsi-klu
Vamsi-kluforce-pushed the fix/google-evaluation-optional-69323 branch from 9ee89ab to 1141e4fCompareAugust 4, 2026 06:55

@aaron-y-chenaaron-y-chen left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR!

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.
@Vamsi-klu

Copy link
Copy Markdown
ContributorAuthor

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 evaluation entry under [project.optional-dependencies]. The base keeps google-cloud-aiplatform at main's floor and the old [evaluation] extra form is not reintroduced, so the security pinning from #71273 is preserved while litellm and scikit-learn stay out of a default provider install.

One thing to flag: uv.lock on main is already stale, uv lock --check fails on files this PR does not touch, so the regenerated lock carries a few hunks unrelated to this change.

@potiuk
potiuk marked this pull request as draft August 13, 2026 12:57
Vamsi-kluand others added 2 commits August 16, 2026 20:01
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.
cursoragentand others added 4 commits August 19, 2026 01:13
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.
@Vamsi-klu
Vamsi-klu marked this pull request as ready for review August 23, 2026 23:42
…on-optional-69323
# Conflicts:
#	providers/google/docs/changelog.rst
#	uv.lock
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:providerskind:documentationprovider:googleGoogle (including GCP) related issues

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Google provider base-depends on google-cloud-aiplatform[evaluation], force-installing litellm/scikit-learn for every user

4 participants

@Vamsi-klu@potiuk@aaron-y-chen@cursoragent