feat(service): opt-in per-request model endpoint overrides (embed, VLM, LLM)

[jdye64]

Summary

Lets clients override the default model endpoints on a per-request basis when submitting jobs / asking questions in service mode, instead of being locked to the hardcoded values in the operator's retriever-service.yaml. Overrides are opt-in by the operator and travel over a dedicated, audited channel so the existing trust boundary stays intact. This PR also wires the reranker into service mode for the first time.

Covered models: embed NIM and caption/VLM NIM (ingest pipeline), and the LLM and reranker (/v1/answer).

SDK ergonomics (ingest)

Each fluent stage method carries the endpoint URL for its own model — there is no separate .endpoints(...) method:

create_ingestor(run_mode="service", base_url=...)
    .files("doc.pdf")
    .extract()
    .embed(embed_invoke_url="https://my-embed-nim/v1", embed_model_name="my-embed")
    .caption(endpoint_url="https://my-vlm/v1", model_name="my-vlm", prompt="Describe")
    .ingest()

• .embed(...) accepts embed_invoke_url (or embedding_endpoint), embed_model_name (or model_name), embed_model_provider_prefix, optional api_key.
• .caption(...) accepts endpoint_url, model_name, optional api_key. Supplying endpoint_url also unlocks the caption stage even when the cluster has no VLM configured.

Under the hood these are peeled off and routed into PipelineSpec.endpoint_overrides — a structured, audited field kept separate from the ordinary shape-params blocks so the server's denylist stays effective. A bare api_key with no accompanying endpoint override stays server-owned (dropped).

Query / answer path (POST /v1/answer)

• LLM override: llm_api_base / llm_model / llm_api_key.
• Reranking (new wiring): reranking was defined in the library/query paths but never wired into service mode. The answer path now, when a rerank endpoint is available, over-fetches top_k * refine_factor candidates, reranks them via the remote NIM, and keeps the best top_k before answer generation.
  • Behaviour config: rerank.enabled / rerank.refine_factor / rerank.max_length.
  • Endpoint config (server-owned): nim_endpoints.rerank_invoke_url / rerank_model_name / api_key.
  • Per-request override: rerank_invoke_url / rerank_model_name / rerank_api_key, plus a rerank: true|false toggle. Supplying a rerank override (or rerank: true) enables reranking for that request.

Server side

• pipeline_overrides.endpoint_overrides gates each model type (embed, caption, llm, rerank) plus an allowed_url_prefixes allowlist.
• EndpointOverridePolicy validates client overrides (check_embed/caption/llm/rerank); requests are rejected with 403 when the operator hasn't opted in or the URL prefix is disallowed, and 400 for a bare API key.
• Ingest overrides are merged over server defaults in pipeline_executor.py; the answer route builds per-request LiteLLMClient / rerank calls so transient credentials never poison cached clients.

Helm

Auto-wire nim_endpoints.rerank_invoke_url / rerank_model_name from the operator-managed llama-nemotron-rerank-vl-1b-v2 NIM (mirroring the caption auto-wiring), and render the rerank section — rerank.enabled flips true automatically when a rerank URL resolves. New values: serviceConfig.nimEndpoints.rerankInvokeUrl / rerankModelName and serviceConfig.rerank.{enabled,refineFactor,maxLength}. The pipeline_overrides.endpoint_overrides opt-in flags are set via the mounted retriever-service.yaml (not templated), consistent with the existing chart.

Tests

• tests/test_service_endpoint_overrides.py — schema, config→policy plumbing, policy accept/reject (URL-prefix allowlist, bare-key rejection), caption-stage unlock, worker merge, SDK routing via .embed()/.caption(), and the answer-path LLM and rerank overrides (server-default rerank, over-fetch, override applied/rejected/prefix-enforced, 400 when requested without an endpoint).
• tests/test_service_caption.py — .caption(endpoint_url=...)/model_name route into endpoint_overrides rather than being rejected.
• tests/test_helm_rerank_endpoint.py — source-level + helm template integration coverage of the rerank auto-wiring (integration cases skip cleanly when helm is absent).

All service/policy/config/rerank/helm test selections pass.

  

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds opt-in per-request model endpoint overrides for the embed NIM, caption/VLM NIM, LLM, and reranker in service mode, and wires the reranker into the /v1/answer path for the first time. Overrides travel through a dedicated EndpointOverrides field on PipelineSpec (kept separate from the existing *_params denylist surface) and are guarded by a new EndpointOverridePolicy that enforces operator opt-in flags and an optional URL-prefix allowlist.

• Ingest path: service_ingestor.embed() / .caption() peel endpoint fields into PipelineSpec.endpoint_overrides; pipeline_executor folds them onto server-owned base dicts before the _merge_server_owned trust-key pass.
• Answer path: ServiceAnswerRequest gains llm_api_base/llm_model/llm_api_key and rerank_invoke_url/rerank_model_name/rerank_api_key; the answer handler builds one-off LiteLLMClient instances for LLM overrides and runs over-fetch + rerank via asyncio.to_thread when a rerank endpoint is available.
• Helm: rerank NIM auto-wiring (mirroring the existing caption wiring) and new serviceConfig.rerank behaviour knobs are added to configmap.yaml and values.yaml.

Confidence Score: 5/5

Safe to merge; the two minor inconsistencies found do not affect security boundaries or data correctness.

The override gating is well-layered: the policy class enforces opt-in flags and URL-prefix allowlists, endpoint fields travel through a dedicated schema field rather than the existing denylist surface, and one-off LLM clients avoid poisoning the cached server default. The two findings are: a short-circuit pop in embed() that can leave an alias key in the shape-params dict when both aliases are supplied simultaneously (very unlikely in practice, caught by the server denylist at worst), and _spec_has_only_endpoint_overrides not guarding result_schema (allowing result_schema='compact' to slip past the mode='reject' gate alongside an endpoint override). Neither affects credentials, data integrity, or pipeline execution.

Files Needing Attention: nemo_retriever/src/nemo_retriever/service/service_ingestor.py (alias pop short-circuit) and nemo_retriever/src/nemo_retriever/common/policy.py (_spec_has_only_endpoint_overrides missing result_schema check)

Important Files Changed

Filename	Overview
nemo_retriever/src/nemo_retriever/common/policy.py	Adds EndpointOverridePolicy class and helpers; wires them into validate_pipeline_spec. _spec_has_only_endpoint_overrides is missing a result_schema guard.
nemo_retriever/src/nemo_retriever/service/service_ingestor.py	Adds _record_endpoint_overrides and rewrites embed()/caption(). The or short-circuit for alias keys may leave a second alias in params_dict.
nemo_retriever/src/nemo_retriever/service/routers/ingest.py	Adds LLM and rerank override fields to ServiceAnswerRequest; implements over-fetch + rerank on the answer path. Logic is correct.
nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py	Adds _apply_embed_endpoint_override / _apply_caption_endpoint_override; folds validated client overrides into server-owned base dicts correctly.
nemo_retriever/src/nemo_retriever/service/config.py	Adds RerankConfig, EndpointOverridesConfig, rerank_model_name; wires EndpointOverridePolicy into to_policy(). Well structured.
nemo_retriever/src/nemo_retriever/common/schemas/pipeline_spec.py	Adds EndpointOverrides Pydantic model; wires into PipelineSpec and updates is_empty() correctly.

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into cursor/per-requ..." | Re-trigger Greptile}