Skip to content

Make Vertex AI evaluation optional to avoid litellm bloat#69749

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

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

Conversation

@Vamsi-klu

Copy link
Copy Markdown
Contributor

Problem

apache-airflow-providers-google base depends on google-cloud-aiplatform[evaluation] which pulls litellm, scikit-learn, ruamel-yaml, tqdm, huggingface-hub, tokenizers. Users of GCS/BigQuery get ML eval stack they never use. litellm has frequent proxy-server CVEs (auth bypass, RCE) triaged by every user.

Root cause: providers/google/pyproject.toml:86
"google-cloud-aiplatform[evaluation]>=1.155.0"

Code needing evaluation: vertexai.preview.evaluation.EvalResult, EvalTask used only in GenerativeModelHook.get_eval_task() and run_evaluation() and RunEvaluationOperator.

What I did

Part A: pyproject.toml

  • Change base from "google-cloud-aiplatform[evaluation]>=1.155.0" to "google-cloud-aiplatform>=1.155.0"
  • Add optional extra "evaluation" = ["google-cloud-aiplatform[evaluation]>=1.155.0"]

Part B: Lazy import guard in hook

  • Try/except ImportError for EvalResult/EvalTask, store _evaluation_import_error
  • In get_eval_task() and run_evaluation(), raise AirflowOptionalProviderFeatureException with install instruction if import failed
  • Operator file doesn't import EvalTask at top-level, hook guard enough

Part C: Tests use pytest.importorskip("vertexai.preview.evaluation") to skip gracefully when extra not installed

Part D: Changelog breaking change note directly edited per provider rules

Generated files via prek update-providers-dependencies --all-files: README, docs/index.rst, provider info

Impact

  • Base install no longer pulls litellm/scikit-learn (security + image size + dep hell improvement)
  • Users needing RunEvaluationOperator must install apache-airflow-providers-google[evaluation]
  • Breaking change for Eval users, mitigated via changelog and clear error message
  • pandas remains base dep, ray remains

Testing

  • python -c "from airflow.providers.google.cloud.hooks.gcs import GCSHook" works without evaluation extra
  • Hook.get_eval_task() raises AirflowOptionalProviderFeatureException with helpful msg when mocked missing
  • uv run --project providers-google --extra evaluation pytest ... -k evaluation → 1 passed
  • ruff format/check passed
  • Note: breeze ci-image rebuild needed

Fixes: #69323

Was generative AI tooling used to co-author this PR?
  • Yes — Muse Spark 1.1
    Generated-by: Muse Spark 1.1

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
@boring-cyborg boring-cyborg Bot added area:providers kind:documentation provider:google Google (including GCP) related issues labels Jul 11, 2026
@Vamsi-klu

Copy link
Copy Markdown
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

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:providers kind:documentation provider:google Google (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

1 participant