(lib) Add support for vLLM-backed text embedder

[charlesbluca]

Description

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• If adjusting docker-compose.yaml environment variables have you ensured those are mimicked in the Helm values.yaml file.

[greptile-apps]

Greptile Summary

This PR adds vLLM-backed text embedding support to the NeMo Retriever library, introducing LlamaNemotronEmbed1BV2Embedder (now vLLM-backed, replacing the former HF class), LlamaNemotronEmbedVL1BV2VLLMEmbedder, LlamaNemotronEmbed1BV2HFEmbedder (new explicit HF class for query embedding), and the embed_with_vllm_llm / embed_multimodal_with_vllm_llm helpers. The factory create_local_embedder is restructured around a backend parameter, Retriever now uses create_local_query_embedder with embed_queries(), and the [local] extra gains vllm==0.17.0, flashinfer-cubin==0.6.4, and flashinfer-python==0.6.4 dependencies.

Several previously-flagged P1 concerns from earlier review rounds (missing _llm dataclass field in LlamaNemotronEmbedVL1BV2VLLMEmbedder, bool(\"False\") silently enabling eager mode in gpu_operator.py and processor.py, CUDA_VISIBLE_DEVICES global mutation) remain unresolved in the current diff and should be addressed before merging.

Confidence Score: 4/5

Safe to merge with caution — core embedding logic is correct, but several previously-flagged P1 issues remain unresolved.

The double-prefix bug is fixed, aarch64 wheel sources are present, version pins are consistent (0.17.0), and PoolerConfig exception handling is now specific to ImportError. However, three previously-flagged P1 concerns are still unaddressed: _llm is not declared as a dataclass field in LlamaNemotronEmbedVL1BV2VLLMEmbedder (can hide construction errors), bool('False') silently enables enforce_eager in gpu_operator.py and processor.py, and CUDA_VISIBLE_DEVICES is still mutated globally in LlamaNemotronEmbedVL1BV2VLLMEmbedder.post_init.

nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_vl_1b_v2_embedder.py (missing _llm field + CUDA_VISIBLE_DEVICES), nemo_retriever/src/nemo_retriever/text_embed/gpu_operator.py and processor.py (bool('False') enforce_eager)

Important Files Changed

Filename	Overview
nemo_retriever/src/nemo_retriever/text_embed/vllm.py	New helpers create_vllm_llm, embed_with_vllm_llm, embed_multimodal_with_vllm_llm; PoolerConfig outer catch now correctly limits to ImportError; well-structured with proper batching logic.
nemo_retriever/src/nemo_retriever/model/init.py	Factory now routes via backend=auto/hf/vllm; non-VL create_local_embedder always returns vLLM embedder with no auto-fallback on non-Linux platforms.
nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_embedder.py	Class repurposed as vLLM-backed embedder; eager GPU load in post_init and missing unload() violate model-lazy-loading/gpu-memory-lifecycle rules; device deprecation warning is correctly implemented.
nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_vl_1b_v2_embedder.py	New LlamaNemotronEmbedVL1BV2VLLMEmbedder added; _llm not declared as dataclass field (could cause AttributeError on failed construction); CUDA_VISIBLE_DEVICES global mutation still present.
nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_hf_embedder.py	New HF-only embedder for query-side use; loads model eagerly in post_init; embed() signature lacks prefix kwarg (inconsistent with vLLM embedder's embed()).
nemo_retriever/src/nemo_retriever/text_embed/processor.py	Double-prefix bug fixed by passing prefix='' to embed(); VL/non-VL routing correct; bool('False') for enforce_eager still silently enables eager mode when YAML injects string values.
nemo_retriever/src/nemo_retriever/text_embed/gpu_operator.py	Correctly routes HF vs vLLM ingest backends; bool('False') for enforce_eager still present (string from YAML evaluates truthy), disabling CUDA graphs silently.
nemo_retriever/src/nemo_retriever/text_embed/runtime.py	local_batch_size bug fixed (now uses effective_batch_size); GPU cache clear on exception added; skip_prefix heuristic works correctly for all new model types.
nemo_retriever/src/nemo_retriever/retriever.py	_get_local_embedder now uses create_local_query_embedder with backend selection; embed_queries() replaces manual 'query: ' prefixing; cache keyed by (resolved_name, backend) tuple.
nemo_retriever/pyproject.toml	vllm bumped to 0.17.0 with both x86_64 and aarch64 CUDA 13 wheels in uv.sources; flashinfer-cubin and flashinfer-python added without uv.sources wheel pinning.
nemo_retriever/tests/test_vllm_embed.py	New test file covering embed_with_vllm_llm, embed_multimodal_with_vllm_llm, create_vllm_llm, and VLLMEmbedder image/text-image paths; good coverage of edge cases.
nemo_retriever/tests/test_create_local_embedder.py	New deprecation test for device kwarg; gpu_memory_utilization kwarg forwarding assertion added; tests look correct.

Class Diagram

%%{init: {'theme': 'neutral'}}%%
classDiagram
    class LlamaNemotronEmbed1BV2Embedder {
        +model_id: Optional[str]
        +device: Optional[str] deprecated
        +gpu_memory_utilization: float
        +enforce_eager: bool
        -_llm: Any
        +embed(texts, batch_size, prefix) Tensor
        +embed_queries(texts, batch_size) Tensor
    }
    class LlamaNemotronEmbed1BV2HFEmbedder {
        +device: Optional[str]
        +model_id: Optional[str]
        -_tokenizer: Any
        -_model: Any
        +embed(texts, batch_size) Tensor
        +embed_queries(texts, batch_size) Tensor
    }
    class LlamaNemotronEmbedVL1BV2Embedder {
        +device: Optional[str]
        +model_id: Optional[str]
        -_model: Any
        +embed(texts, batch_size) Tensor
        +embed_queries(texts, batch_size) Tensor
        +embed_images(images_b64, batch_size) Tensor
        +embed_text_image(texts, images_b64, batch_size) Tensor
    }
    class LlamaNemotronEmbedVL1BV2VLLMEmbedder {
        +model_id: Optional[str]
        +gpu_memory_utilization: float
        +enforce_eager: bool
        +embed(texts, batch_size) Tensor
        +embed_queries(texts, batch_size) Tensor
        +embed_images(images_b64, batch_size) Tensor
        +embed_text_image(texts, images_b64, batch_size) Tensor
    }
    class create_local_embedder {
        <<factory>>
        VL+auto/vllm --> LlamaNemotronEmbedVL1BV2VLLMEmbedder
        VL+hf --> LlamaNemotronEmbedVL1BV2Embedder
        nonVL --> LlamaNemotronEmbed1BV2Embedder
    }
    class create_local_query_embedder {
        <<factory>>
        nonVL+hf --> LlamaNemotronEmbed1BV2HFEmbedder
        otherwise --> create_local_embedder
    }
    create_local_embedder --> LlamaNemotronEmbed1BV2Embedder : non-VL
    create_local_embedder --> LlamaNemotronEmbedVL1BV2VLLMEmbedder : VL + auto/vllm
    create_local_embedder --> LlamaNemotronEmbedVL1BV2Embedder : VL + hf
    create_local_query_embedder --> LlamaNemotronEmbed1BV2HFEmbedder : non-VL + hf
    create_local_query_embedder --> create_local_embedder : otherwise

Loading

Prompt To Fix All With AI

This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_hf_embedder.py
Line: 116-128

Comment:
**Inconsistent `embed()` interface with vLLM counterpart**

`LlamaNemotronEmbed1BV2HFEmbedder.embed()` does not accept a `prefix` kwarg, while `LlamaNemotronEmbed1BV2Embedder.embed()` does (`prefix: str = "passage: "`). Code that switches backends via `create_local_query_embedder(..., backend="hf")` and then calls `embedder.embed(..., prefix="")` would get a `TypeError`. The `public-api-contract` rule requires interchangeable backends to share the same call signature. Consider adding `prefix: str = "passage: "` here (even if it's effectively ignored since the HF path self-prefixes), so callers don't need to know which backend is active.

**Rule Used:** Every public class and function in nemo_retriever ... ([source](https://app.greptile.com/review/custom-context?memory=public-api-contract))

How can I resolve this? If you propose a fix, please make it concise.

_{Reviews (27): Last reviewed commit: "fix(embed): add missing --embed-local-in..." | Re-trigger Greptile}

[greptile-apps]

Non-Linux environments silently broken by factory-forcing vLLM

create_local_embedder now always returns LlamaNemotronEmbed1BV2VLLMEmbedder for non-VL models. Because vllm is gated to sys_platform == 'linux' in pyproject.toml, calling this factory on macOS or Windows immediately raises RuntimeError: vLLM is not installed — with no fallback to the perfectly functional HF path (LlamaNemotronEmbed1BV2Embedder). Users doing development, CI, or CPU inference on non-Linux hosts now have no path through the public factory.

Consider keeping a use_vllm: bool parameter (defaulting to True on Linux, False elsewhere) and falling back to the HF class when vLLM is unavailable:

def create_local_embedder(
    model_name: str | None = None,
    *,
    use_vllm: bool | None = None,  # None → auto-detect
    ...
) -> Any:
    if use_vllm is None:
        try:
            import vllm  # noqa: F401
            use_vllm = True
        except ImportError:
            use_vllm = False

    if not is_vl_embed_model(model_name):
        if use_vllm:
            from ...llama_nemotron_embed_1b_v2_embedder import LlamaNemotronEmbed1BV2VLLMEmbedder
            return LlamaNemotronEmbed1BV2VLLMEmbedder(...)
        else:
            from ...llama_nemotron_embed_1b_v2_embedder import LlamaNemotronEmbed1BV2Embedder
            return LlamaNemotronEmbed1BV2Embedder(device=device, hf_cache_dir=hf_cache_dir, model_id=model_id)

Prompt To Fix With AI

This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/model/__init__.py
Line: 71-83

Comment:
**Non-Linux environments silently broken by factory-forcing vLLM**

`create_local_embedder` now _always_ returns `LlamaNemotronEmbed1BV2VLLMEmbedder` for non-VL models. Because `vllm` is gated to `sys_platform == 'linux'` in `pyproject.toml`, calling this factory on macOS or Windows immediately raises `RuntimeError: vLLM is not installed` — with no fallback to the perfectly functional HF path (`LlamaNemotronEmbed1BV2Embedder`). Users doing development, CI, or CPU inference on non-Linux hosts now have no path through the public factory.

Consider keeping a `use_vllm: bool` parameter (defaulting to `True` on Linux, `False` elsewhere) and falling back to the HF class when vLLM is unavailable:

```python
def create_local_embedder(
    model_name: str | None = None,
    *,
    use_vllm: bool | None = None,  # None → auto-detect
    ...
) -> Any:
    if use_vllm is None:
        try:
            import vllm  # noqa: F401
            use_vllm = True
        except ImportError:
            use_vllm = False

    if not is_vl_embed_model(model_name):
        if use_vllm:
            from ...llama_nemotron_embed_1b_v2_embedder import LlamaNemotronEmbed1BV2VLLMEmbedder
            return LlamaNemotronEmbed1BV2VLLMEmbedder(...)
        else:
            from ...llama_nemotron_embed_1b_v2_embedder import LlamaNemotronEmbed1BV2Embedder
            return LlamaNemotronEmbed1BV2Embedder(device=device, hf_cache_dir=hf_cache_dir, model_id=model_id)
```

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

bool("False") evaluates to True, silently enabling eager mode

task_config values can arrive as strings when the dict is populated from YAML or environment variables. bool("False") returns True in Python because any non-empty string is truthy, so a user setting enforce_eager: false in YAML (parsed as the string "False") would silently get enforce_eager=True, disabling CUDA graphs and degrading throughput with no error.

Suggested change

	enforce_eager=bool(task_config.get("enforce_eager", False)),
	enforce_eager=task_config.get("enforce_eager", False) is True or str(task_config.get("enforce_eager", "false")).lower() in ("1", "true", "yes"),

Or use a small helper:

def _to_bool(v, default=False):
    if isinstance(v, bool): return v
    if v is None: return default
    return str(v).strip().lower() in ("1", "true", "yes")

Prompt To Fix With AI

This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/text_embed/processor.py
Line: 108

Comment:
**`bool("False")` evaluates to `True`, silently enabling eager mode**

`task_config` values can arrive as strings when the dict is populated from YAML or environment variables. `bool("False")` returns `True` in Python because any non-empty string is truthy, so a user setting `enforce_eager: false` in YAML (parsed as the string `"False"`) would silently get `enforce_eager=True`, disabling CUDA graphs and degrading throughput with no error.

```suggestion
            enforce_eager=task_config.get("enforce_eager", False) is True or str(task_config.get("enforce_eager", "false")).lower() in ("1", "true", "yes"),
```

Or use a small helper:

```python
def _to_bool(v, default=False):
    if isinstance(v, bool): return v
    if v is None: return default
    return str(v).strip().lower() in ("1", "true", "yes")
```

How can I resolve this? If you propose a fix, please make it concise.

[jperez999]

Might need to coordinate with @edknv and the VL reranker PR

[greptile-apps]

Eager model load in __post_init__ violates model-lazy-loading rule

LlamaNemotronEmbed1BV2HFEmbedder.__post_init__ calls AutoTokenizer.from_pretrained and AutoModel.from_pretrained unconditionally at construction time, including a .to(dev) call that places weights on GPU immediately. The model-lazy-loading rule requires models to load on first use, not at construction. Any code that constructs this class (e.g. create_local_query_embedder(..., backend="hf")) will trigger a full HF download/load and GPU transfer before any embed call is made.

Defer loading to the first _embed_local call and guard with an _ensure_loaded() helper, or document (in the class docstring) that GPU allocation occurs at construction time as an acknowledged design trade-off.

Rule Used: Models must be loaded lazily (on first use or expl... (source)

Prompt To Fix With AI

This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_hf_embedder.py
Line: 39-66

Comment:
**Eager model load in `__post_init__` violates `model-lazy-loading` rule**

`LlamaNemotronEmbed1BV2HFEmbedder.__post_init__` calls `AutoTokenizer.from_pretrained` and `AutoModel.from_pretrained` unconditionally at construction time, including a `.to(dev)` call that places weights on GPU immediately. The `model-lazy-loading` rule requires models to load on first use, not at construction. Any code that constructs this class (e.g. `create_local_query_embedder(..., backend="hf")`) will trigger a full HF download/load and GPU transfer before any embed call is made.

Defer loading to the first `_embed_local` call and guard with an `_ensure_loaded()` helper, or document (in the class docstring) that GPU allocation occurs at construction time as an acknowledged design trade-off.

**Rule Used:** Models must be loaded lazily (on first use or expl... ([source](https://app.greptile.com/review/custom-context?memory=model-lazy-loading))

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

_llm not declared as a dataclass field

LlamaNemotronEmbedVL1BV2VLLMEmbedder assigns self._llm in __post_init__ but never declares it as a typed dataclass field. Compare with LlamaNemotronEmbed1BV2Embedder which has _llm: Any = field(default=None, init=False, repr=False). Without the field declaration there is no None sentinel before __post_init__ runs, so if anything in lines 188–215 raises before the assignment on line 210, a subsequent call to embed() / embed_queries() will raise an opaque AttributeError: 'LlamaNemotronEmbedVL1BV2VLLMEmbedder' object has no attribute '_llm' rather than propagating the original construction error.

Suggested change

	model_id: Optional[str] = None
	device: Optional[str] = None
	hf_cache_dir: Optional[str] = None
	gpu_memory_utilization: float = 0.45
	enforce_eager: bool = False
	model_id: Optional[str] = None
	device: Optional[str] = None
	hf_cache_dir: Optional[str] = None
	gpu_memory_utilization: float = 0.45
	enforce_eager: bool = False
	_llm: Any = field(default=None, init=False, repr=False)

Also requires adding field to the from dataclasses import dataclass, field import.

Prompt To Fix With AI

This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_vl_1b_v2_embedder.py
Line: 176-180

Comment:
**`_llm` not declared as a dataclass field**

`LlamaNemotronEmbedVL1BV2VLLMEmbedder` assigns `self._llm` in `__post_init__` but never declares it as a typed dataclass field. Compare with `LlamaNemotronEmbed1BV2Embedder` which has `_llm: Any = field(default=None, init=False, repr=False)`. Without the field declaration there is no `None` sentinel before `__post_init__` runs, so if anything in lines 188–215 raises before the assignment on line 210, a subsequent call to `embed()` / `embed_queries()` will raise an opaque `AttributeError: 'LlamaNemotronEmbedVL1BV2VLLMEmbedder' object has no attribute '_llm'` rather than propagating the original construction error.

```suggestion
    model_id: Optional[str] = None
    device: Optional[str] = None
    hf_cache_dir: Optional[str] = None
    gpu_memory_utilization: float = 0.45
    enforce_eager: bool = False

    _llm: Any = field(default=None, init=False, repr=False)
```

Also requires adding `field` to the `from dataclasses import dataclass, field` import.

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

bool("False") silently enables eager mode

self._kwargs.get("enforce_eager", False) can return the string "False" when kwargs are populated from YAML or environment variables. bool("False") evaluates to True in Python because any non-empty string is truthy, so a user setting enforce_eager: false in config silently gets enforce_eager=True, disabling CUDA graphs and degrading throughput. The same pattern appears in processor.py line 108.

Or introduce a small shared helper — e.g. _to_bool(v, default=False) — and call it in both places.

Prompt To Fix With AI

This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/text_embed/gpu_operator.py
Line: 44-46

Comment:
**`bool("False")` silently enables eager mode**

`self._kwargs.get("enforce_eager", False)` can return the string `"False"` when kwargs are populated from YAML or environment variables. `bool("False")` evaluates to `True` in Python because any non-empty string is truthy, so a user setting `enforce_eager: false` in config silently gets `enforce_eager=True`, disabling CUDA graphs and degrading throughput. The same pattern appears in `processor.py` line 108.

Or introduce a small shared helper — e.g. `_to_bool(v, default=False)` — and call it in both places.

How can I resolve this? If you propose a fix, please make it concise.