[Router] vectorstore: bound ingestion memory with streaming batches - #2666
[Router] vectorstore: bound ingestion memory with streaming batches#2666Peterren wants to merge 1 commit into
Conversation
Previously processJob embedded every chunk of a file and then inserted them all in a single call, holding the whole file's chunk text and embedding vectors in memory at once. For large files with high-dimensional embeddings this is an unbounded-by-count per-job memory spike, multiplied across workers. Embed and store chunks in fixed-size windows instead (default 64, configurable via vector_store.ingestion_batch_size). This bounds peak per-job memory to roughly O(batch_size x embedding_dimension) regardless of file size. ctx is checked before each batch so a cancelled lifecycle aborts promptly between batches. Chunk IDs remain globally correct across batch boundaries. Batching changes failure semantics: earlier batches are already inserted when a later one fails. To preserve the prior all-or-nothing invariant (a failed file leaves nothing searchable), a partial failure now best-effort removes the already-inserted chunks via DeleteByFileID, detached from the job context so cleanup still runs during shutdown but bounded so a wedged backend cannot block the worker. Full transactional reconciliation of ingestion remains tracked separately under the vector-store hardening issue. Adds tests for batch-boundary correctness across several chunk/batch sizes, the configurable default, partial-failure cleanup, and the first-batch-failure no-op cleanup path. go test -race is clean for pkg/vectorstore, pkg/routerruntime, and pkg/config. Signed-off-by: Yincheng Ren <27608815+Peterren@users.noreply.github.com>
✅ Deploy Preview for vllm-semantic-router ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
github-actions
Bot
assigned
AayushSaini101,
drivebyer,
FAUST-BENCHOU,
ramkrishs and
shraderdm
👥 vLLM Semantic Team NotificationThe following members have been identified for the changed files in this PR and have been automatically assigned when their GitHub accounts are assignable in this repository: 📁
|
✅ Supply Chain Security Report — All Clear
Scanned at |
| } | ||
| cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), batchCleanupTimeout) | ||
| defer cancel() | ||
| _ = p.backend.DeleteByFileID(cleanupCtx, job.VectorStoreID, job.FileID) |
There was a problem hiding this comment.
+1 on keeping the all or nothing invariant, nice touch. One thing I've been bitten by before with compensation deletes keyed on file ID: AttachFile doesn't dedupe, so a retry of the same file can run while the old attachment is still completed. If the retry fails here, this delete wipes the previous attachment's chunks too, and it still shows completed with nothing searchable behind it. Tracking the chunk IDs this job inserted and deleting only those would avoid that.
| // whose ingestion failed partway through. The delete is detached from the | ||
| // job context (which may be cancelled during shutdown) but bounded by | ||
| // batchCleanupTimeout so a wedged backend cannot block the worker. Failure to | ||
| // clean up is logged implicitly by leaving the file marked failed; it is not |
There was a problem hiding this comment.
Just thinking out loud here: If DeleteByFileID itself fails or times out, the stated invariant (failed file leaves nothing searchable) is broken with zero operator signal. Can we log a warning with the store and file IDs (and ideally a metric) when cleanup fails?
|
Nice work bounding the embed/insert memory! A few things worth a look:
|
8 participants
Part of #2474 (vector-store ingestion hardening) — implements the "bounded streaming batches" acceptance item. Not a full close of the umbrella issue, whose remaining slices (status retention, saga reconciliation, parser streaming) are tracked separately.
Purpose
InsertChunkscall. That holds the whole file's chunk text and embedding vectors in memory simultaneously — an unbounded-by-count per-job memory spike, multiplied across ingestion workers. This PR embeds and stores chunks in fixed-size windows instead (default 64, configurable viavector_store.ingestion_batch_size), bounding peak per-job memory to roughlyO(batch_size × embedding_dimension)regardless of file size.Router(primary), with config-surface updates toDashboardandE2E/Helm values.Behavior notes
ctxis checked before each batch, so a cancelled lifecycle aborts promptly between batches. Chunk IDs remain globally correct across batch boundaries (indexing is unchanged).DeleteByFileID, detached from the (possibly cancelled) job context so cleanup still runs during shutdown, but bounded so a wedged backend cannot block the worker.Test Plan
From
src/semantic-router:New unit tests in
pkg/vectorstore/pipeline_batch_test.go:InsertChunkscalls equalsceil(N/batch), and no batch exceeds the configured size.Test Result
go test -raceis clean forpkg/vectorstore,pkg/routerruntime, andpkg/config;golangci-lintreports 0 issues;gofmt/go vetclean; the router binary builds with themilvustag.rag-hybrid-searchE2E profile is not run locally (requires a cluster); it exercises this path in CI.Semantic Router PR Checklist
[Router],[CLI],[Dashboard],[Operator],[Fleet-Sim],[Bindings],[Training],[E2E],[Docs], or[CI/Build]git commit -s