Skip to content

[Router] vectorstore: thread context into Embedder interface#2535

Open
Peterren wants to merge 5 commits into
vllm-project:mainfrom
Peterren:feat/vectorstore-embedder-ctx
Open

[Router] vectorstore: thread context into Embedder interface#2535
Peterren wants to merge 5 commits into
vllm-project:mainfrom
Peterren:feat/vectorstore-embedder-ctx

Conversation

@Peterren

@Peterren Peterren commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

Stacked on #2517 (A1). This PR targets main, so until #2517 merges the diff also shows the A1 commits (lifecycle root context + bounded Stop(ctx)). Please review #2517 first; the A2-only change is the final commit, refactor(vectorstore): thread context into Embedder interface.

Purpose

  • Thread a context.Context through the vector-store Embedder interface: Embed(text) becomes Embed(ctx, text).
  • Today the ingestion pipeline derives a cancellable lifecycle context (added in [Router] vectorstore: add lifecycle root context and bounded Stop(ctx) #2517), but the embedding stage still takes no context, so a job wedged inside an embed cannot be interrupted — that gap was explicitly called out as the follow-up in the bounded Stop(ctx) work. This PR closes it for cooperative embedders.
  • Module: Router (with E2E-style tests inside the vectorstore package).

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.Embed checks ctx.Err() before the FFI call. The underlying candle_binding calls are synchronous and not themselves cancellable; interrupting a call already in flight remains a deliberate follow-up.
  • The pipeline passes the job's lifecycle-derived context into Embed (the per-chunk ctx.Err() checkpoint from A1 is unchanged), so cancelling the root context now propagates into an in-flight embed for any embedder that honors it.
  • The two query-side callers pass their existing request context: route_vectorstore_search.go uses r.Context(), and the extproc RAG filter uses its trace context.
  • Test embedders and mocks updated to the new signature.

This is a standalone, source-incompatible Go interface change: Embedder.Embed gains a ctx parameter, 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

cd src/semantic-router && go test -race ./pkg/vectorstore ./pkg/routerruntime

New test (pipeline_embedder_ctx_test.go): a ctx-aware embedder that blocks inside Embed until either released or its context is cancelled. With a job parked inside the embed, a bounded Stop(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 as failed with LastError.Code == "cancelled" (a context cancellation is classified as an intentional termination, not an embedding_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 an Eventually because failJob runs 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/routerruntime passes; race detector clean. pkg/apiserver and pkg/extproc tests also pass with the updated signature.
  • gofmt and golangci-lint (repo config) are clean on the changed files.
  • Follow-up (not in this PR, tracked under vectorstore: make ingestion cancellable, bounded, and transactional #2474): make the candle_binding embed 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
  • PR title uses module-aligned prefix [Router]
  • Commits in this PR are signed off with git commit -s
  • The Purpose, Test Plan, and Test Result sections reflect the actual scope, commands, and blockers for this change

Closes #2603

Part of the umbrella #2474 (stacked on #2602 / #2517).

@Peterren
Peterren requested review from Xunzhuo and rootfs as code owners July 14, 2026 20:01
@netlify

netlify Bot commented Jul 14, 2026

Copy link
Copy Markdown

Deploy Preview for vllm-semantic-router ready!

Name Link
🔨 Latest commit 4a40006
🔍 Latest deploy log https://app.netlify.com/projects/vllm-semantic-router/deploys/6a5e3fae74d7c9000885d094
😎 Deploy Preview https://deploy-preview-2535--vllm-semantic-router.netlify.app
📱 Preview on mobile
Toggle QR Code...

QR Code

Use your smartphone camera to open QR code link.
🤖 Make changes Run an agent on this branch

To edit notification comments on pull requests, go to your Netlify project configuration.

@github-actions

github-actions Bot commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

👥 vLLM Semantic Team Notification

The 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:

📁 config

Owners: @FAUST-BENCHOU, @shraderdm, @drivebyer, @ramkrishs, @WUKUNTAI-0211, @AayushSaini101, @siloteemu
Files changed:

  • config/config.yaml

📁 deploy

Owners: @FAUST-BENCHOU, @shraderdm, @drivebyer, @ramkrishs, @WUKUNTAI-0211, @AayushSaini101, @siloteemu
Files changed:

  • deploy/helm/semantic-router/values.yaml

📁 e2e

Owners: @FAUST-BENCHOU, @shraderdm, @drivebyer, @ramkrishs, @WUKUNTAI-0211, @AayushSaini101, @siloteemu
Files changed:

  • e2e/profiles/rag-hybrid-search/values.yaml

📁 src/semantic-router

Owners: @FAUST-BENCHOU, @shraderdm, @drivebyer, @ramkrishs, @WUKUNTAI-0211, @AayushSaini101, @siloteemu
Files changed:

  • src/semantic-router/pkg/apiserver/route_vectorstore_search.go
  • src/semantic-router/pkg/apiserver/runtime_globals_test.go
  • src/semantic-router/pkg/config/vectorstore.go
  • src/semantic-router/pkg/extproc/req_filter_rag_vectorstore.go
  • src/semantic-router/pkg/extproc/router_runtime_services_test.go
  • src/semantic-router/pkg/routerruntime/vectorstore_runtime.go
  • src/semantic-router/pkg/vectorstore/candle_embedder.go
  • src/semantic-router/pkg/vectorstore/integration_test.go
  • src/semantic-router/pkg/vectorstore/pipeline.go
  • src/semantic-router/pkg/vectorstore/pipeline_embedder_ctx_test.go
  • src/semantic-router/pkg/vectorstore/pipeline_lifecycle_test.go
  • src/semantic-router/pkg/vectorstore/pipeline_test.go

vLLM Semantic Router

🎉 Thanks for your contributions!

This comment was automatically generated based on the OWNER files in the repository.

@github-actions

github-actions Bot commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

✅ Supply Chain Security Report — All Clear

Scanner Status Findings
AST Codebase Scan (Py, Go, JS/TS, Rust) 29 finding(s) — MEDIUM: 22 · LOW: 7
AST PR Diff Scan No issues detected
Regex Fallback Scan No issues detected

Scanned at 2026-07-20T15:36:38.623Z · View full workflow logs

@github-actions github-actions Bot deleted a comment from codecov-commenter Jul 14, 2026

@Xunzhuo Xunzhuo left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/semantic-router/pkg/vectorstore/pipeline_embedder_ctx_test.go
Comment thread src/semantic-router/pkg/vectorstore/pipeline.go
@Peterren
Peterren force-pushed the feat/vectorstore-embedder-ctx branch from 11678e1 to d2a073a Compare July 19, 2026 03:05
@Peterren

Copy link
Copy Markdown
Contributor Author

@Xunzhuo addressed both review points and rebased on latest main:

  1. Test race fixed (d2a073ab): the goroutine baseline is now captured in BeforeEach before the fixture starts the worker, so NumGoroutine() <= baseline+1 is a real leak barrier rather than trivially true while the worker unwinds. The failed-status read is now an Eventually (since failJob runs concurrently as the worker unwinds), so it no longer races and observes in_progress. The cleanup release stays only as a timeout safety net.
  2. Cancellation classification fixed: embedChunks now checks errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) and records the job with the existing cancelled code instead of embedding_error. The new test asserts status.LastError.Code == "cancelled".
  3. Description: removed "backward-compatible" — this is a source-incompatible Go interface change (Embed gains a ctx param; all in-tree implementations/callers are updated here).
  4. check-linked-issue now passes — added Closes #2603, a dedicated slice issue for the embedding-stage cancellation, rather than Part of #2474 (which doesn't populate closing references) or prematurely closing the umbrella.

Still stacked on #2517 — the A2-only change is the final commit, vectorstore: classify embed cancellation as cancelled; fix leak-test baseline race sitting on top of refactor(vectorstore): thread context into Embedder interface. Happy to rebase onto main once #2517 lands.

@github-actions github-actions Bot deleted a comment from codecov-commenter Jul 19, 2026
@github-actions github-actions Bot deleted a comment from codecov-commenter Jul 19, 2026
@mergify

mergify Bot commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

Tick the box to add this pull request to the merge queue (same as @mergifyio queue).

  • Queue this pull request

@github-actions github-actions Bot deleted a comment from codecov-commenter Jul 20, 2026
Peterren added 5 commits July 20, 2026 11:32
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

vectorstore: thread context into the Embedder interface (embedding-stage cancellation)

8 participants