new: expose some onnx session options #578
Conversation
📝 WalkthroughWalkthroughIntroduces controlled exposure and application of ONNX Runtime SessionOptions. OnnxModel gains EXPOSED_SESSION_OPTIONS, _select_exposed_session_options, and add_extra_session_options; _load_onnx_model accepts extra_session_options and applies validated options to ort.SessionOptions before creating an InferenceSession. Multiple model classes (text, image, sparse, rerank, late-interaction, multimodal) extract _extra_session_options from kwargs and thread them through load_onnx_model, embedding calls, and parallel worker params. Tests added to assert enable_cpu_mem_arena behavior. Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (5)
🧰 Additional context used🧬 Code graph analysis (3)tests/test_sparse_embeddings.py (3)
tests/test_text_cross_encoder.py (5)
tests/test_image_onnx_embeddings.py (1)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (9)
🔇 Additional comments (6)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (6)
fastembed/image/onnx_image_model.py (1)
127-135: Fix critical None-unpacking bug in parallel embedding.Line 133 unpacks
extra_session_optionswith**extra_session_options, which raisesTypeError: 'NoneType' object is not a mappingwhenextra_session_optionsisNone(the default).Apply this diff to fix the unpacking:
start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn" params = { "model_name": model_name, "cache_dir": cache_dir, "providers": providers, "local_files_only": local_files_only, "specific_model_path": specific_model_path, - **extra_session_options, + **(extra_session_options or {}), **kwargs, }fastembed/image/onnx_embedding.py (1)
174-186: Missing extra_session_options propagation to _embed_images.The
_extra_session_optionsare not passed to_embed_images, meaning parallel workers won't receive the session options. This breaks the feature for parallel embedding scenarios.Apply this diff to propagate the session options:
yield from self._embed_images( model_name=self.model_name, cache_dir=str(self.cache_dir), images=images, batch_size=batch_size, parallel=parallel, providers=self.providers, cuda=self.cuda, device_ids=self.device_ids, local_files_only=self._local_files_only, specific_model_path=self._specific_model_path, + extra_session_options=self._extra_session_options, **kwargs, )fastembed/text/onnx_text_model.py (1)
136-148: Critical: Fix unpacking of Optional dict.Unpacking
extra_session_optionsdirectly will fail type checking and cause a runtime error if it'sNone. The parameter isOptional[dict[str, Any]], so it must be guarded before unpacking.Apply this diff to fix the issue:
params = { "model_name": model_name, "cache_dir": cache_dir, "providers": providers, "local_files_only": local_files_only, "specific_model_path": specific_model_path, - **extra_session_options, + **(extra_session_options or {}), **kwargs, }fastembed/rerank/cross_encoder/onnx_text_model.py (1)
124-132: Critical: Fix unpacking of Optional dict.Unpacking
extra_session_optionsdirectly will fail type checking and cause a runtime error if it'sNone. The parameter isOptional[dict[str, Any]], so it must be guarded before unpacking.Apply this diff to fix the issue:
params = { "model_name": model_name, "cache_dir": cache_dir, "providers": providers, "local_files_only": local_files_only, "specific_model_path": specific_model_path, - **extra_session_options, + **(extra_session_options or {}), **kwargs, }fastembed/late_interaction_multimodal/onnx_multimodal_model.py (2)
150-158: Critical: Fix unpacking of Optional dict.Unpacking
extra_session_optionsdirectly will fail type checking and cause a runtime error if it'sNone. The parameter isOptional[dict[str, Any]], so it must be guarded before unpacking.Apply this diff to fix the issue:
params = { "model_name": model_name, "cache_dir": cache_dir, "providers": providers, "local_files_only": local_files_only, "specific_model_path": specific_model_path, - **extra_session_options, + **(extra_session_options or {}), **kwargs, }
219-227: Critical: Fix unpacking of Optional dict.Unpacking
extra_session_optionsdirectly will fail type checking and cause a runtime error if it'sNone. The parameter isOptional[dict[str, Any]], so it must be guarded before unpacking.Apply this diff to fix the issue:
params = { "model_name": model_name, "cache_dir": cache_dir, "providers": providers, "local_files_only": local_files_only, "specific_model_path": specific_model_path, - **extra_session_options, + **(extra_session_options or {}), **kwargs, }
🧹 Nitpick comments (1)
fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py (1)
88-89: Consider removing the redundant explicit parameter.The
extra_session_optionsparameter is explicitly declared but never used directly—it's only accessed viakwargsin_select_exposed_session_options(kwargs)on line 115. This creates redundancy and triggers the static analysis warning.Apply this diff to remove the redundant parameter:
device_id: Optional[int] = None, specific_model_path: Optional[str] = None, - extra_session_options: Optional[dict[str, Any]] = None, **kwargs: Any, ):Alternatively, if you prefer to keep explicit parameters for documentation, use it directly instead of extracting from kwargs:
self.providers = providers self.lazy_load = lazy_load - self._extra_session_options = self._select_exposed_session_options(kwargs) + if extra_session_options is None: + extra_session_options = {} + self._extra_session_options = extra_session_options
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
fastembed/common/onnx_model.py(4 hunks)fastembed/image/onnx_embedding.py(2 hunks)fastembed/image/onnx_image_model.py(4 hunks)fastembed/late_interaction/colbert.py(3 hunks)fastembed/late_interaction_multimodal/colpali.py(4 hunks)fastembed/late_interaction_multimodal/onnx_multimodal_model.py(6 hunks)fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py(4 hunks)fastembed/rerank/cross_encoder/onnx_text_model.py(4 hunks)fastembed/sparse/bm42.py(3 hunks)fastembed/sparse/minicoil.py(1 hunks)fastembed/sparse/splade_pp.py(3 hunks)fastembed/text/onnx_embedding.py(3 hunks)fastembed/text/onnx_text_model.py(4 hunks)
🧰 Additional context used
🧬 Code graph analysis (7)
fastembed/text/onnx_embedding.py (1)
fastembed/common/onnx_model.py (1)
_select_exposed_session_options(122-131)
fastembed/sparse/minicoil.py (1)
fastembed/common/onnx_model.py (1)
_select_exposed_session_options(122-131)
fastembed/late_interaction_multimodal/colpali.py (1)
fastembed/common/onnx_model.py (1)
_select_exposed_session_options(122-131)
fastembed/late_interaction/colbert.py (1)
fastembed/common/onnx_model.py (1)
_select_exposed_session_options(122-131)
fastembed/image/onnx_embedding.py (1)
fastembed/common/onnx_model.py (1)
_select_exposed_session_options(122-131)
fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py (1)
fastembed/common/onnx_model.py (1)
_select_exposed_session_options(122-131)
fastembed/sparse/bm42.py (1)
fastembed/common/onnx_model.py (1)
_select_exposed_session_options(122-131)
🪛 GitHub Actions: Tests
fastembed/image/onnx_image_model.py
[error] 127-127: TypeError: 'NoneType' object is not a mapping while constructing ONNX session options (extra_session_options is None).
fastembed/common/onnx_model.py
[error] 146-146: TypeError: 'NoneType' object is not iterable in add_extra_session_options (extra_session_options is None).
🪛 GitHub Actions: type-checkers
fastembed/text/onnx_text_model.py
[error] 146-146: Unpacked dict entry 5 has incompatible type "Optional[dict[str, Any]]"; expected "SupportsKeysAndGetItem[str, Any]" [dict-item]
fastembed/image/onnx_image_model.py
[error] 133-133: Unpacked dict entry 5 has incompatible type "Optional[dict[str, Any]]"; expected "SupportsKeysAndGetItem[str, Any]" [dict-item]
fastembed/common/onnx_model.py
[error] 105-105: Argument 2 to "add_extra_session_options" of "OnnxModel" has incompatible type "Optional[dict[str, Any]]"; expected "dict[str, Any]" [arg-type]
fastembed/rerank/cross_encoder/onnx_text_model.py
[error] 130-130: Unpacked dict entry 5 has incompatible type "Optional[dict[str, Any]]"; expected "SupportsKeysAndGetItem[str, Any]" [dict-item]
fastembed/late_interaction_multimodal/onnx_multimodal_model.py
[error] 156-156: Unpacked dict entry 5 has incompatible type "Optional[dict[str, Any]]"; expected "SupportsKeysAndGetItem[str, Any]" [dict-item]
[error] 225-225: Unpacked dict entry 5 has incompatible type "Optional[dict[str, Any]]"; expected "SupportsKeysAndGetItem[str, Any]" [dict-item]
🪛 Ruff (0.14.5)
fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py
88-88: Unused method argument: extra_session_options
(ARG002)
🔇 Additional comments (22)
fastembed/common/onnx_model.py (3)
27-28: LGTM! Session options exposure mechanism introduced.The
EXPOSED_SESSION_OPTIONStuple provides a whitelist of configurable ONNX Runtime session options, starting withenable_cpu_mem_arena. This is a sound approach to controlled exposure of session options.
121-131: LGTM! Helper method correctly filters exposed options.The
_select_exposed_session_optionsclassmethod correctly filters kwargs to extract only the exposed session options.
65-65: Type signature is correct here.The
extra_session_options: Optional[dict[str, Any]] = Noneparameter correctly acceptsNone. The issue is inadd_extra_session_optionswhich needs to handleNone(see related comment on lines 133-151).fastembed/sparse/splade_pp.py (2)
102-102: LGTM! Extra session options correctly initialized.The
_extra_session_optionsattribute is properly initialized via_select_exposed_session_options, which extracts exposed session options from kwargs.
137-137: LGTM! Extra session options correctly propagated.The
_extra_session_optionsare correctly threaded through both the ONNX model loading path (line 137) and the embedding path (line 173), ensuring session options are applied in both direct and parallel worker scenarios.Also applies to: 173-173
fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py (1)
115-115: LGTM! Extra session options correctly propagated.The
_extra_session_optionsare correctly initialized and threaded through both the ONNX model loading (line 155) and reranking (line 198) paths.Also applies to: 155-155, 198-198
fastembed/image/onnx_image_model.py (1)
58-59: LGTM! Extra session options correctly forwarded to parent.The
extra_session_optionsparameter is properly declared and forwarded to the parent class's_load_onnx_modelmethod.Also applies to: 67-67
fastembed/sparse/bm42.py (2)
106-106: LGTM! Extra session options correctly initialized.The
_extra_session_optionsattribute is properly initialized via_select_exposed_session_options.
150-150: LGTM! Extra session options correctly propagated.The
_extra_session_optionsare correctly threaded through both the ONNX model loading (line 150) and embedding (line 317) paths.Also applies to: 317-317
fastembed/image/onnx_embedding.py (1)
101-101: LGTM! Extra session options correctly initialized and passed to model loading.The
_extra_session_optionsare properly initialized and forwarded to_load_onnx_model.Also applies to: 138-138
fastembed/text/onnx_embedding.py (2)
236-236: LGTM! Extra session options correctly initialized.The
_extra_session_optionsattribute is properly initialized via_select_exposed_session_options.
294-294: LGTM! Extra session options correctly propagated.The
_extra_session_optionsare correctly threaded through both the embedding path (line 294) and the ONNX model loading path (line 331).Also applies to: 331-331
fastembed/late_interaction_multimodal/colpali.py (4)
83-83: LGTM!The initialization of
_extra_session_optionscorrectly filters kwargs using the exposed session options pattern.
121-130: LGTM!The session options are correctly propagated to the ONNX model loading path.
232-245: LGTM!The session options are correctly propagated through the text embedding path.
268-281: LGTM!The session options are correctly propagated through the image embedding path.
fastembed/late_interaction/colbert.py (3)
146-146: LGTM!The initialization of
_extra_session_optionscorrectly filters kwargs using the exposed session options pattern, consistent with other classes in this PR.
178-187: LGTM!The session options are correctly propagated to the ONNX model loading path.
229-242: LGTM!The session options are correctly propagated through the document embedding path.
fastembed/text/onnx_text_model.py (1)
49-68: LGTM!The
extra_session_optionsparameter is correctly added and propagated to the superclass.fastembed/rerank/cross_encoder/onnx_text_model.py (1)
28-47: LGTM!The
extra_session_optionsparameter is correctly added and propagated to the superclass.fastembed/late_interaction_multimodal/onnx_multimodal_model.py (1)
59-81: LGTM!The
extra_session_optionsparameter is correctly added and propagated to the superclass.
| self._extra_session_options = self._select_exposed_session_options(kwargs) | ||
|
|
There was a problem hiding this comment.
Extra session options initialized but not propagated.
The _extra_session_options are initialized here but never passed to _load_onnx_model (line 151-158) or _embed_documents (lines 211-227, 235-249). This means the session options won't be applied for MiniCOIL models.
Apply this diff to propagate the session options:
def load_onnx_model(self) -> None:
self._load_onnx_model(
model_dir=self._model_dir,
model_file=self.model_description.model_file,
threads=self.threads,
providers=self.providers,
cuda=self.cuda,
device_id=self.device_id,
+ extra_session_options=self._extra_session_options,
)Also add to the embed method (around line 211):
yield from self._embed_documents(
model_name=self.model_name,
cache_dir=str(self.cache_dir),
documents=documents,
batch_size=batch_size,
parallel=parallel,
providers=self.providers,
cuda=self.cuda,
device_ids=self.device_ids,
k=self.k,
b=self.b,
avg_len=self.avg_len,
is_query=False,
local_files_only=self._local_files_only,
specific_model_path=self._specific_model_path,
+ extra_session_options=self._extra_session_options,
**kwargs,
)And similarly for query_embed (around line 235).
Committable suggestion skipped: line range outside the PR's diff.
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
fastembed/common/onnx_model.py (1)
134-152: Fix the None-handling bug flagged in previous review.As noted in the previous review, this method needs to handle
Noneinput properly. The type signature should beOptional[dict[str, Any]]with an early return guard.Additionally, consider using
ValueErrorinstead ofassertfor the validation on lines 148-150, as assertions can be disabled with Python's-Oflag and may not provide clear error messages in production.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
fastembed/common/onnx_model.py(4 hunks)fastembed/image/onnx_image_model.py(4 hunks)fastembed/late_interaction_multimodal/onnx_multimodal_model.py(6 hunks)fastembed/rerank/cross_encoder/onnx_text_model.py(4 hunks)fastembed/text/onnx_text_model.py(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- fastembed/image/onnx_image_model.py
- fastembed/rerank/cross_encoder/onnx_text_model.py
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)
- GitHub Check: Python 3.9.x on windows-latest test
- GitHub Check: Python 3.13.x on windows-latest test
- GitHub Check: Python 3.13.x on macos-latest test
- GitHub Check: Python 3.9.x on macos-latest test
- GitHub Check: Python 3.10.x on ubuntu-latest test
🔇 Additional comments (5)
fastembed/common/onnx_model.py (2)
122-132: LGTM!The filtering logic correctly extracts only the exposed session options from the kwargs dictionary.
65-66: LGTM!The parameter is correctly typed as
Optional[dict[str, Any]]and the None-check before callingadd_extra_session_optionsappropriately guards against the method's current limitation.Also applies to: 105-106
fastembed/text/onnx_text_model.py (1)
49-67: LGTM!The parameter is correctly added and forwarded to the superclass, maintaining consistency with the parent class API.
fastembed/late_interaction_multimodal/onnx_multimodal_model.py (2)
59-77: LGTM!The parameter is correctly threaded through to the superclass.
115-128: LGTM!The parameter signatures are correctly typed and consistently applied to both text and image embedding methods.
Also applies to: 186-199
| if extra_session_options is not None: | ||
| params.update(extra_session_options) |
There was a problem hiding this comment.
Critical bug: session options are incorrectly flattened into params.
Same issue as in _embed_documents: the dictionary structure must be preserved when passing to workers.
Apply this diff:
- if extra_session_options is not None:
- params.update(extra_session_options)
+ if extra_session_options is not None:
+ params["extra_session_options"] = extra_session_options📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if extra_session_options is not None: | |
| params.update(extra_session_options) | |
| if extra_session_options is not None: | |
| params["extra_session_options"] = extra_session_options |
🤖 Prompt for AI Agents
In fastembed/late_interaction_multimodal/onnx_multimodal_model.py around lines
230-231, the code currently flattens extra_session_options into params with
params.update(extra_session_options); instead preserve the dictionary structure
by adding the options as a nested entry (e.g., params['extra_session_options']
or the same key used in _embed_documents) instead of updating params, so the
worker receives a single dict field containing the session options rather than
merging its keys into the top-level params.
| if extra_session_options is not None: | ||
| params.update(extra_session_options) |
There was a problem hiding this comment.
Critical bug: session options are incorrectly flattened into params.
Using params.update(extra_session_options) flattens the session options dictionary into top-level keys in params. This breaks the propagation to worker processes because:
extra_session_options = {"enable_cpu_mem_arena": True}becomesparams["enable_cpu_mem_arena"] = True- Workers receive
**params, so they getkwargs["enable_cpu_mem_arena"] = Trueinstead ofkwargs["extra_session_options"] - Workers cannot pass these correctly to
_load_onnx_model(extra_session_options=...)
Apply this diff to preserve the dictionary structure:
- if extra_session_options is not None:
- params.update(extra_session_options)
+ if extra_session_options is not None:
+ params["extra_session_options"] = extra_session_options📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if extra_session_options is not None: | |
| params.update(extra_session_options) | |
| if extra_session_options is not None: | |
| params["extra_session_options"] = extra_session_options |
🤖 Prompt for AI Agents
In fastembed/text/onnx_text_model.py around lines 149-150, the code flattens
extra_session_options into params which breaks worker propagation; instead of
params.update(extra_session_options) preserve the dictionary by assigning it to
a single key (e.g. params["extra_session_options"] = extra_session_options) so
that workers receive extra_session_options as a dict and can call
_load_onnx_model(extra_session_options=...) correctly; remove the update call
and add the single-key assignment, keeping params None-safe if needed.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py (2)
111-115: Session options extraction placement is fine; consider moving beforesuper().__init__Using
_select_exposed_session_options(kwargs)and storing the result inself._extra_session_optionsis a clean way to centralize ONNX session tuning. If_select_exposed_session_optionsmutateskwargs(e.g., viapop), you might want to call it beforesuper().__init__(...)so base classes never see these special keys, and to keep the pattern consistent with other ONNX-based models, if any. Not a blocker, just a small robustness/readability tweak to consider.
186-199: Guard against duplicateextra_session_optionsbetween attribute and**kwargsPassing
extra_session_options=self._extra_session_optionswhile also forwarding**kwargsmeans a caller who explicitly passesextra_session_optionsintorerank_pairs(...)will hit a “multiple values for keyword argument 'extra_session_options'” error.If you want to support per-call overrides, consider resolving it here, e.g.:
- yield from self._rerank_pairs( + extra_session_options = kwargs.pop( + "extra_session_options", self._extra_session_options + ) + yield from self._rerank_pairs( model_name=self.model_name, cache_dir=str(self.cache_dir), pairs=pairs, batch_size=batch_size, parallel=parallel, providers=self.providers, cuda=self.cuda, device_ids=self.device_ids, local_files_only=self._local_files_only, specific_model_path=self._specific_model_path, - extra_session_options=self._extra_session_options, + extra_session_options=extra_session_options, **kwargs, )If per-call overrides are not intended, it may still be worth documenting that
extra_session_optionsis constructor-only to avoid confusion.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
fastembed/image/onnx_embedding.py(3 hunks)fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py(3 hunks)fastembed/sparse/minicoil.py(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- fastembed/sparse/minicoil.py
- fastembed/image/onnx_embedding.py
🔇 Additional comments (1)
fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py (1)
146-155: ONNX session options plumbing into_load_onnx_modellooks correctForwarding
extra_session_options=self._extra_session_optionsinto_load_onnx_modelcleanly connects constructor-time configuration with actual session creation; this should give consistent behavior across eager and lazy loads. No issues spotted with this change.
|
|
||
|
|
||
| class OnnxModel(Generic[T]): | ||
| EXPOSED_SESSION_OPTIONS = ("enable_cpu_mem_arena",) |
There was a problem hiding this comment.
If we don't expect this to be expanded at runtime it could be an enum like
from enum import StrEnum, auto
class ExposedSessionOptions(StrEnum):
enable_cpu_mem_arena = auto()or similar, would make type hinting nicer and still allows in checks and iteration.
Nvm, StrEnum was added with 3.11.
There was a problem hiding this comment.
I guess it can be emulated with ExposedSessionOptions(str, Enum), but feels like a bit too much, users won't see this enum anyway
There was a problem hiding this comment.
yes, it's not worth it to emulate it that way, it was a low hanging fruit with StrEnum built in but definitely not worth reimplementing it (also just (str, Enum) won't give you auto iirc so it'd be even more work). I just misremembered when StrEnum was added.
There was a problem hiding this comment.
I guess we could type it as a tuple[Literal[...]] but again, not worth the effort for something that's only internal.
* new: expose some onnx session options * fix: fix extra session options is None case * fix: fix missing params * new: add tests
Uh oh!
There was an error while loading. Please reload this page.