feat: bundle google-cloud-aiplatform for Vertex partner models#2910
feat: bundle google-cloud-aiplatform for Vertex partner models#2910abguymon wants to merge 5 commits into
Conversation
abguymon
force-pushed
the
adam/install-aiplatform-for-vertex-ai
branch
from
6bd2f19 to
8c21f2d
Compare
|
[Automatic Post]: I have assigned @neubig as a reviewer based on the repository MAINTAINERS file. Thanks in advance for the help! |
|
@OpenHands audit the packages in this PR and make sure that they're appropriate (and don't, e.g. pose supply chain risks). Then consider the increase in total package size based on this. We may want to consider accepting this PR as-is, or we may want to make this an optional dependency. In the latter case, propose an implementation that would allow us to do that. |
|
I'm on it! neubig can track my progress at all-hands.dev |
|
I audited the dependency changes in this PR. Packages added / changedNew packages in
Existing dependency changes pulled in by the resolver:
Supply-chain / package auditI don't see obvious supply-chain red flags here:
I also ran a known-vulnerability check against the exact versions added/updated by this PR: uvx --with pip-audit pip-audit -r /tmp/pr2910-reqs.txtThat came back clean. So from a provenance / vuln perspective, this looks acceptable to me. The real tradeoff is footprint, not package trust. Size impactI measured this two ways. Python install footprint (exact wheel versions in this PR):
PyInstaller binary footprint:
So this is a fairly meaningful size increase for a provider-specific feature. RecommendationMy take:
If we make it optional, here's the implementation I'd suggestThis can follow the existing
Given the numbers above, my preference would be:
That keeps the base Python install much leaner while preserving a straightforward path for users who do need Vertex partner models. AI note: this comment was generated by an AI assistant (OpenHands) on behalf of the user. |
|
[Automatic Post]: It has been a while since there was any activity on this PR. @abguymon, are you still working on it? If so, please go ahead, if not then please request review, close it, or request that someone else follow up. |
|
[Automatic Post]: This PR seems to be currently waiting for review. @neubig, could you please take a look when you have a chance? |
Adds google-cloud-aiplatform as a required dep so LiteLLM's
vertex_ai_partner_models handler can route to Vertex AI Model Garden
MaaS endpoints (MiniMax, Qwen, Kimi, etc.). Without it, requests like
vertex_ai/minimaxai/minimax-m2-maas, vertex_ai/qwen/qwen3-coder-*-maas,
vertex_ai/moonshotai/kimi-k2-thinking-maas fail with:
litellm.BadRequestError: Vertex_aiException BadRequestError -
vertexai import failed please run
`pip install -U "google-cloud-aiplatform>=1.38"`.
Got error: No module named 'vertexai'
Changes:
- openhands-sdk/pyproject.toml: add "google-cloud-aiplatform>=1.38"
- uv.lock: regenerated (+7 transitive packages)
- agent-server.spec: collect vertexai/google.cloud.aiplatform/google.api_core
submodules as hidden imports (the vertexai import is deferred inside
litellm function bodies, so PyInstaller static analysis misses it),
pin google.rpc.status_pb2 (dynamic gRPC proto stub), and copy package
metadata.
Supersedes OpenHands#2373 (closed without merging) and lands the spec-side
changes needed for the bundled binary to actually work at runtime.
Addresses review feedback on OpenHands#2910. Mirrors the existing `boto3` extra pattern so the base `openhands-sdk` install stays lean (the Vertex SDK adds ~80 MiB to site-packages and ~62 MiB to the PyInstaller binary). - pyproject.toml: move `google-cloud-aiplatform` from base deps into `[project.optional-dependencies]` as `vertex`. - llm.py / vertex_preflight.py: when the inferred provider is `vertex_ai` and `vertexai` is not importable, raise a friendly `LLMBadRequestError` with an install hint (`pip install "openhands-sdk[vertex]"`) instead of letting LiteLLM crash with a low-level `ModuleNotFoundError`. - agent-server.spec: skip the Vertex `collect_all` block (and the `google.rpc.status_pb2` pin) when `vertexai` is not installed, so local PyInstaller builds without the extra still succeed. - docker/Dockerfile: add `--extra vertex` to both `uv sync` calls so the published agent-server binary continues to bundle the Vertex SDK by default. - tests: cover the preflight raise/no-op paths. - uv.lock: regenerated.
abguymon
force-pushed
the
adam/install-aiplatform-for-vertex-ai
branch
from
8c21f2d to
8120d23
Compare
|
Thanks for the audit @neubig — agreed on the size tradeoff. I've pushed a follow-up commit that makes Vertex an optional extra along the lines you suggested:
Re-requesting your review. |
neubig
left a comment
neubig
left a comment
There was a problem hiding this comment.
Thanks, overall this looks very solid but I think we'd like to not bake this in to the default image, as we'd like to keep it lean. I think we could consider making a separate docker file with more of the optional dependencies backed in though.
Per follow-up review on OpenHands#2910: don't bundle the Vertex SDK in the default agent-server image — it's a ~62 MiB / ~72% size hit for provider-specific support that most users won't need. - Dockerfile: drop `--extra vertex` from both `uv sync` invocations, so the published image stays as lean as it was on `main`. - agent-server.spec: rewrite the header comment to make the new policy explicit (off by default; opt in with `uv sync ... --extra vertex` before running pyinstaller). The `if _VERTEX_AVAILABLE:` guard already added in the previous commit is what makes the default build succeed without the extra installed; keep it as the supported opt-in path for downstream builds (e.g. fork-specific images) that do want Vertex bundled.
Closes the gap left by the previous commit: with the default Docker
build now lean, there was no way to opt back into a Vertex-bundled
binary without forking the Dockerfile. Adds `ENABLE_VERTEX` (default
0) so downstream builds can do:
docker build --build-arg ENABLE_VERTEX=1 ...
ARG declared globally and re-imported into both the `builder` and
`binary-builder` stages, with a small inline shell conditional that
appends `--extra vertex` to each `uv sync` call. The PyInstaller
spec already auto-detects an installed `vertexai` and bundles
accordingly, so no further wiring is needed.
Spec header comment updated to point at this build arg as the
canonical opt-in path.
|
Thanks @neubig — pushed two follow-up commits (4ba82047, 1490482d):
If you'd actually prefer a separate |
|
[Automatic Post]: It has been a while since there was any activity on this PR. @abguymon, are you still working on it? If so, please go ahead, if not then please request review, close it, or request that someone else follow up. |
|
[Automatic Post]: It has been a while since there was any activity on this PR. @abguymon, are you still working on it? If so, please go ahead, if not then please request review, close it, or request that someone else follow up. This comment was created by an AI agent (OpenHands) on behalf of the user. |
|
[Automatic Post]: It has been a while since there was any activity on this PR. @abguymon, are you still working on it? If so, please go ahead, if not then please request review, close it, or request that someone else follow up. This comment was generated by an AI agent (OpenHands) on behalf of the user. |
|
[Automatic Post]: It has been a while since there was any activity on this PR. @abguymon, are you still working on it? If so, please go ahead, if not then please request review, close it, or request that someone else follow up. This comment was created by an AI agent (OpenHands) on behalf of the user. |
# Conflicts: # openhands-agent-server/openhands/agent_server/agent-server.spec # uv.lock
|
@neubig I believe this PR is ready to merge |
|
[Automatic Post]: It has been a while since there was any activity on this PR. @abguymon, are you still working on it? If so, please go ahead, if not then please request review, close it, or request that someone else follow up. This comment was created by an AI agent (OpenHands) on behalf of the user. |
1 similar comment
|
[Automatic Post]: It has been a while since there was any activity on this PR. @abguymon, are you still working on it? If so, please go ahead, if not then please request review, close it, or request that someone else follow up. This comment was created by an AI agent (OpenHands) on behalf of the user. |
|
✅ Review complete. This review was performed through OpenHands Cloud Automation. You can log in and view the conversation here. |
all-hands-bot
left a comment
all-hands-bot
left a comment
There was a problem hiding this comment.
Code Review: feat: bundle google-cloud-aiplatform for Vertex partner models
🟡 Acceptable - Works but could be cleaner.
Overview
This PR adds an optional openhands-sdk[vertex] extra that bundles google-cloud-aiplatform for Vertex AI partner model support (MiniMax, Qwen, Kimi MaaS). The implementation is thoughtful: it uses conditional bundling in PyInstaller, adds a preflight check for missing SDK, and keeps the default image lean via --build-arg ENABLE_VERTEX=1.
[CRITICAL ISSUES]
- None identified.
[IMPROVEMENT OPPORTUNITIES]
-
[agent-server.spec, Line 38] Special Case: The loop
for _pkg in _vertex_pkgswith repeated unpacking (_d, _b, _h = collect_all(_pkg)) followed by+=concatenation is harder to read. Consider usingitertools.chainor flattening with a list comprehension. -
[agent-server.spec, Lines 28-32] Comment Volume: The 23-line block comment explaining Vertex AI bundling is excessive for an internal build file. This could be a 3-4 line docstring or a comment inline with the conditional check. The explanation of PEP-420 namespace packages and
collect_allvscollect_submodulesbelongs in a commit message or internal docs, not inline.
[TESTING GAPS]
- The tests in
test_vertex_preflight.pyare clean and cover the key paths well. No issues.
[DEPENDENCY CHANGES]
uv.lock⚠️ Dependency Downgrade:pypdfdowngraded from6.10.2to6.9.2. Was this intentional? Check if this introduces any regressions.uv.lock⚠️ Dependency Upgrade:google-authupgraded from2.41.1to2.49.2. This is a minor bump but worth verifying compatibility.uv.lock⚠️ Dependency Change:rsapackage removed. This appears to be related to thegoogle-authupgrade (they replaced thersadependency withcryptographyin newer versions). Confirm this is expected.
[RISK ASSESSMENT]
- [Overall PR]
⚠️ Risk Assessment: 🟡 MEDIUM
The core implementation is sound. The medium rating reflects the dependency changes inuv.lock(pypdf downgrade, rsa removal) which should be verified as intentional. The newgoogle-cloud-aiplatformdependency is substantial but opt-in, so existing users are unaffected.
VERDICT:
✅ Worth merging - Core logic is sound, minor improvements suggested
KEY INSIGHT:
The opt-in design (both for SDK users via pip extra and for Docker builds via ENABLE_VERTEX=1) is the right approach for a large optional dependency; the preflight check provides a good user experience.
Improve this review? If any feedback above seems incorrect or irrelevant to this repository, you can teach the reviewer to do better:
- Add a
.agents/skills/custom-codereview-guide.mdfile to your branch (or edit it if one already exists) with the/codereviewtrigger and the context the reviewer is missing. See the customization docs for the required frontmatter format.- Re-request a review - the reviewer reads guidelines from the PR branch, so your changes take effect immediately.
- When your PR is merged, the guideline file goes through normal code review by repository maintainers.
Resolve with AI? Install the iterate skill in your agent and run
/iterateto automatically drive this PR through CI, review, and QA until it's merge-ready.Was this review helpful? React with 👍 or 👎 to give feedback.
This review was generated by an AI agent (OpenHands) on behalf of the user through OpenHands Automation. View conversation
| import importlib.util as _vertex_importlib_util | ||
|
|
||
| _VERTEX_AVAILABLE = _vertex_importlib_util.find_spec("vertexai") is not None | ||
|
|
There was a problem hiding this comment.
🟡 Suggestion: The loop with repeated unpacking and += concatenation could be clearer. Consider using itertools.chain to flatten the collected lists.
| return importlib.util.find_spec("vertexai") is not None | ||
|
|
||
|
|
||
| def assert_vertex_sdk_available(provider: str | None) -> None: |
There was a problem hiding this comment.
🟡 Suggestion: Consider adding a type hint for the provider parameter: def assert_vertex_sdk_available(provider: str | None) -> None:. This matches the actual usage.
|
I could not push to the original fork branch ( That branch merges current |
3 participants
Summary
Adds
google-cloud-aiplatformas a required dep ofopenhands-sdkso LiteLLM'svertex_ai_partner_modelshandler can route to Vertex AI Model Garden MaaS endpoints — MiniMax, Qwen, Kimi, etc.Without this dep, requests like
vertex_ai/minimaxai/minimax-m2-maas,vertex_ai/qwen/qwen3-coder-*-maas,vertex_ai/moonshotai/kimi-k2-thinking-maasfail at the first completion call with:What changed
openhands-sdk/pyproject.toml: add"google-cloud-aiplatform>=1.38"uv.lock: regenerated (+7 transitive packages)openhands-agent-server/openhands/agent_server/agent-server.spec: usecollect_all(...)for the Vertex SDK and itsgoogle.cloud.*namespace siblings (vertexai,google.cloud.aiplatform,aiplatform_v1,aiplatform_v1beta1,bigquery,storage,resourcemanager,google.api_core,google.auth,google.rpc,google.genai,proto,grpc_status). Also pingoogle.rpc.status_pb2explicitly (gRPC proto stub imported dynamically) and copygoogle-cloud-aiplatformmetadata.Why
collect_alland not justcollect_submodulesI first tried the same minimal change as #2373 (just
pyproject.toml), then addedcollect_submodules("vertexai")etc. to the spec. Neither worked at runtime — the bundled binary still failed withModuleNotFoundError: No module named 'vertexai', then withNo module named 'google.rpc.status_pb2'after I hand-injectedvertexai.Root cause:
google.cloud.*is a PEP-420 implicit namespace package shared across many distributions. PyInstaller'scollect_submodules("google.cloud.aiplatform")doesn't traverse PEP-420 namespaces correctly — the package'sdist-infoends up in the bundle (viacopy_metadata), but the actual module files do not.collect_all(...)walks the installed package directory directly and pulls in modules + datas + binaries, which is what the bundle needs.I verified by inspecting the extracted
_MEI*/directory after running the freshly-built binary:_MEI*/vertexai/__init__.py✅_MEI*/google/cloud/aiplatform/✅_MEI*/google/rpc/status_pb2.py✅Background
This supersedes #2373 (closed without merging earlier this year). That PR only touched
pyproject.toml, which would have shipped the dep into the build env but not into the PyInstaller bundle.Checklist
Test plan
-pythonimage withvertexaiactually present in the bundle (not just the dist-info)vertex_ai/minimaxai/minimax-m2-maas(or another partner MaaS model) in a deployment with Vertex credentials