Repository navigation
feat(vertex): Google Cloud Vertex AI Vector Search engine (pure KNN) - #144
Merged
Merged
Conversation
New cloud-only engine `vertex` (like Turbopuffer — no local server), talking to the Vertex AI REST API v1 via reqwest::blocking: - configure: create a STREAM_UPDATE tree-AH Index + public IndexEndpoint + deploy (LRO polling with a configurable long timeout). Reuse an already-deployed index via VERTEX_INDEX / VERTEX_INDEX_ENDPOINT / VERTEX_DEPLOYED_INDEX_ID to skip the tens-of-minutes deploy. - upload: concurrent upsertDatapoints (streaming index), each worker holding and refreshing its own bearer token so a long upload can't die on a 401. - search: findNeighbors against the public endpoint, N persistent workers each with its own blocking client, timing only the RPC + reply parse, routed through the canonical compute_search_stats. - delete: undeploy + delete endpoint + delete index (skipped for a reused index, gated on the same three-var condition as configure). Pure vector KNN only — no metadata filters, no mixed workload. Auth via VERTEX_ACCESS_TOKEN or `gcloud auth print-access-token`, re-fetched per phase. Registered in both engine registries; adds vertex-single-node.json + README row and note. 8 unit tests cover every request-body builder and reply parser (the REST shapes were cross-checked against the Vertex AI v1 API in adversarial review; the engine is cloud-only so there is no live integration test). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Docker Build ValidationPlatforms tested:
Tests performed:
|
…IZE, gated test, runbook
- CI Check was red because a newer stable clippy (0.1.97) flags a `question_mark`
pattern in weaviate.rs's `where_json_to_grpc_filters` (pre-existing, unrelated
to Vertex; blocks every PR). Rewrite the trailing `else { return None }` with
`?` — behavior identical, weaviate tests unchanged.
- vertex: add `VERTEX_SHARD_SIZE` (e.g. SHARD_SIZE_SMALL). Vertex ties valid
deploy machine types to the index shard size (default SHARD_SIZE_MEDIUM needs
e2-standard-16+); exposing it lets callers pick a compatible machine. Unit
test covers the omitted/set cases.
- Add gated `integration_vertex` test: drives the binary end-to-end against
`random-100` and asserts a recall floor, self-skipping unless VERTEX_PROJECT
is set (CI stays green). Plus a README runbook for benchmarking against Vertex.
Verified against the CI toolchain (rustc 1.97.1 / clippy 0.1.97):
`cargo clippy --all-targets -D warnings` and `cargo fmt --check` both clean;
466 unit tests pass.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Docker Build ValidationPlatforms tested:
Tests performed:
|
Bring Vertex to parity with the pure-KNN reference engine (Dragonfly) on the two
optional Engine telemetry hooks, both best-effort and outside every timed window:
- get_memory_usage: Vertex is managed (no client-visible memory), so report the
index's post-upload stats (shardsCount, vectorsCount) from a describe.
- server_metadata: reproducibility snapshot — deployment config (region, machine
type, shard size, distance measure, tree-AH params) plus, once the index
exists, its resource names and a live describe (config + indexStats +
deployment state). Called before upload (static config only) and after search.
Validated the index-describe shape against the live API: `indexStats` is a
top-level field (`{shardsCount, vectorsCount}`) and `metadata.config` carries
`shardSize` — matching what these hooks read.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Docker Build ValidationPlatforms tested:
Tests performed:
|
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
1 participant
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.
Adds a new Vertex AI Vector Search (Google Cloud) engine — pure vector KNN, per the "start simple" ask.
What
Cloud-only engine
vertex(no local server, like Turbopuffer), over the Vertex AI REST API v1 (reqwest::blocking):Index+ publicIndexEndpoint+deployIndex(LRO polling)upsertDatapoints(streaming index)findNeighborsagainst the public endpoint, N persistent workersundeployIndex+ delete endpoint + delete indexTimed search window wraps only the RPC + reply parse (matches Turbopuffer/Milvus), routed through the canonical
compute_search_stats.Config & auth
VERTEX_PROJECT. Auth:VERTEX_ACCESS_TOKEN, elsegcloud auth print-access-token(re-fetched per phase; upload workers refresh their own token so long uploads can't 401).VERTEX_REGION(defaultus-central1),VERTEX_MACHINE_TYPE,VERTEX_DEPLOY_TIMEOUT_SECS, and tree-AH knobsVERTEX_APPROX_NEIGHBORS/VERTEX_LEAF_EMBEDDING_COUNT/VERTEX_LEAF_SEARCH_PERCENT.VERTEX_INDEX+VERTEX_INDEX_ENDPOINT+VERTEX_DEPLOYED_INDEX_IDto reuse an already-deployed index (cleanup then leaves it in place).search_params.fraction_leaf_nodes_to_search_override(0..1).Pure KNN only — no metadata filters, no mixed workload.
Testing
Cloud-only (needs a GCP project + a ~tens-of-minutes deploy), so no live integration test — same posture as Turbopuffer. 8 unit tests pin every pure request-body builder and reply parser. An adversarial review cross-checked all REST field names, nesting, the regional-vs-public-endpoint URL split, LRO handling, and search measurement fidelity against the Vertex AI v1 API (no shape bugs found), and surfaced 4 operational issues — all fixed here: per-worker token refresh in upload, real (honest) upload parallelism, the reuse-guard consistency in
delete(), and a README note onbatch_sizefor wide datasets.Registered in both engine registries; adds
vertex-single-node.json+ README row/note.cargo clippy --all-targets -D warnings+ fmt clean;--describe engineslistsvertex-default.🤖 Generated with Claude Code