[Router] vectorstore: thread context into Embedder interface#2535
[Router] vectorstore: thread context into Embedder interface#2535Peterren wants to merge 5 commits into
Conversation
✅ Deploy Preview for vllm-semantic-router ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
👥 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 |
Xunzhuo
left a comment
There was a problem hiding this comment.
The A2 direction is sound: the incremental diff threads the lifecycle/request context through every vector-store Embedder call site, and it accurately documents that Candle FFI is only pre-call cancellable. I am requesting changes for two correctness/validation issues below. The new test currently fails in both test-and-build and the pre-commit fast gate, and cancellation is recorded as an embedding failure. Also, please remove “backward-compatible” from the description (this is a source-incompatible Go interface change) and satisfy check-linked-issue with a dedicated closing issue link; Part of #2474 does not populate GitHub closing references, while closing the umbrella issue from this partial slice would be premature. Keep this PR stacked on #2517 or rebase it after #2517 lands.
11678e1 to
d2a073a
Compare
|
@Xunzhuo addressed both review points and rebased on latest
Still stacked on #2517 — the A2-only change is the final commit, |
|
Tick the box to add this pull request to the merge queue (same as
|
The ingestion pipeline processed each job with a hardcoded context.Background() and shut down via an unbounded wg.Wait(). If an embedder or backend wedged inside a stage, Stop() blocked forever and the process could not be shut down. Introduce a pipeline lifecycle root context created in Start() and cancelled on shutdown. processJob now derives its work from that context and checks cancellation between stages (read, chunk, per-chunk embed, insert), so a cancelled lifecycle aborts in-flight jobs at the next checkpoint instead of running to completion. Stop() becomes Stop(ctx) error: it stops accepting jobs, fails queued jobs, and drains in-flight work bounded by ctx. If ctx elapses before the drain finishes, the root context is cancelled and Stop returns ctx.Err() without blocking further. A nil ctx uses a default 30s bound. The runtime Shutdown path now bounds pipeline drain with a 30s deadline so a wedged backend cannot hang process shutdown. Add lifecycle and integration tests covering: bounded Stop returning within its deadline when a job is wedged, Stop idempotency, graceful drain with no leaked worker goroutines, and post-shutdown searchability. Part of vllm-project#2474 Signed-off-by: Yincheng Ren <27608815+Peterren@users.noreply.github.com>
Address review: the shutdown drain bound was a hardcoded 30s. It is not a per-file ingest budget but a graceful-shutdown grace window (stop accepting jobs, wait for in-flight jobs to drain, then cancel), so it must be sized per deployment — above typical single-job latency but below the platform's shutdown grace period (e.g. Kubernetes terminationGracePeriodSeconds). Bubble it up as vector_store.ingestion_drain_timeout_seconds (default 30), sourced by VectorStoreRuntime; keep a bounded default fallback so Stop is never unbounded when constructed without config. Add the key to the reference config, Helm values, and the RAG e2e profile. Also check the error return of Stop(ctx) at all test call sites (errcheck), which the pre-commit go-lint hook flagged after Stop's signature changed. Part of vllm-project#2474 Signed-off-by: Yincheng Ren <27608815+Peterren@users.noreply.github.com>
processJob exceeded the agent lint funlen limit (41 > 40 statements). Extract the per-chunk embedding loop (Step 5) into embedChunks, which preserves the same per-chunk ctx cancellation checks and error handling and returns ok=false to signal the caller to stop. No behavior change. Signed-off-by: Yincheng Ren <27608815+Peterren@users.noreply.github.com>
Change Embedder.Embed(text) -> Embed(ctx, text) so embedding work derives from the caller's context. CandleEmbedder checks ctx.Err() before the FFI call; the ingestion pipeline passes the job's lifecycle-derived context, and the two query-side callers (apiserver search, extproc RAG filter) pass their request context. This makes a cooperative embedder interruptible: when the pipeline lifecycle is cancelled, an in-flight embed unwinds on its own instead of running to completion, closing the follow-up called out in the bounded Stop(ctx) work. The underlying candle_binding calls remain synchronous; interrupting a call already in flight is a further follow-up. Adds a ctx-aware embedder test asserting a parked embed unwinds via lifecycle cancellation with no leaked worker goroutines and the job recorded as failed. Part of vllm-project#2474 Signed-off-by: Yincheng Ren <27608815+Peterren@users.noreply.github.com>
…baseline race Address review on the embedder-context PR: - embedChunks: when a cooperative embedder returns context.Canceled or context.DeadlineExceeded (lifecycle root cancelled during shutdown), record the job with the cancelled code instead of embedding_error, so status/metrics do not misclassify an intentional termination as a backend fault. - pipeline_embedder_ctx_test: capture the goroutine baseline before the fixture starts the worker (the worker was previously counted in the baseline, making the <= baseline+1 leak assertion trivially true while the worker was still unwinding). Read the failed status via Eventually (failJob runs concurrently as the worker unwinds) and assert LastError.Code == "cancelled". Signed-off-by: Yincheng Ren <27608815+Peterren@users.noreply.github.com>
06a88f7 to
4a40006
Compare
Purpose
context.Contextthrough the vector-storeEmbedderinterface:Embed(text)becomesEmbed(ctx, text).Stop(ctx)work. This PR closes it for cooperative embedders.Router(withE2E-style tests inside thevectorstorepackage).What changed
Embedder.Embed(text)→Embed(ctx context.Context, text). The interface doc notes that a cancelled lifecycle/request aborts embedding at the next checkpoint.CandleEmbedder.Embedchecksctx.Err()before the FFI call. The underlyingcandle_bindingcalls are synchronous and not themselves cancellable; interrupting a call already in flight remains a deliberate follow-up.Embed(the per-chunkctx.Err()checkpoint from A1 is unchanged), so cancelling the root context now propagates into an in-flight embed for any embedder that honors it.route_vectorstore_search.gousesr.Context(), and the extproc RAG filter uses its trace context.This is a standalone, source-incompatible Go interface change:
Embedder.Embedgains actxparameter, so every implementation and call site must be updated (all in-tree ones are, in this PR). It does not alter routing, ranking, or ingestion semantics; it only makes cancellation reach the embedding stage.Test Plan
New test (
pipeline_embedder_ctx_test.go): a ctx-aware embedder that blocks insideEmbeduntil either released or its context is cancelled. With a job parked inside the embed, a boundedStop(ctx)elapses and cancels the lifecycle root; the test asserts the parked embed then unwinds on its own (no explicit release), no worker goroutines are leaked, and the interrupted job is recorded asfailedwithLastError.Code == "cancelled"(a context cancellation is classified as an intentional termination, not anembedding_error). The goroutine baseline is captured before the fixture starts the worker, so the leak assertion is a real barrier rather than a tautology, and the failed-status read is anEventuallybecausefailJobruns as the worker unwinds. A ctx-ignoring embedder would stay parked — that contrast is the behavior this PR adds.Test Result
go test -race ./pkg/vectorstore ./pkg/routerruntimepasses; race detector clean.pkg/apiserverandpkg/extproctests also pass with the updated signature.gofmtandgolangci-lint(repo config) are clean on the changed files.candle_bindingembed and backend insert calls natively cancellable so a job wedged inside a single stage can be interrupted before the stage returns; streaming/bounded batches (B2) build on this ctx-aware interface.Semantic Router PR Checklist
[Router]git commit -sCloses #2603
Part of the umbrella #2474 (stacked on #2602 / #2517).