Skip to content

[BUG](ef): Cache sentence models per device - #7882

Open
WitBlue6 wants to merge 1 commit into
chroma-core:mainfrom
WitBlue6:fix/sentence-transformer-device-cache
Open

WitBlue6 wants to merge 1 commit into
chroma-core:mainfrom
WitBlue6:fix/sentence-transformer-device-cache

Conversation

@WitBlue6

Copy link
Copy Markdown

Description of changes

  • Improvements & Bug fixes
    • Include the configured device in the SentenceTransformer model cache key. Previously creating the same model on CPU and then on CUDA reused the CPU instance: the second constructor never reached SentenceTransformer, despite its configuration saying CUDA. The reverse order similarly makes a CPU instance use an existing GPU model.
    • Preserve reuse for the same model/device and isolation for different model names.
  • New functionality
    • None.

Test plan

  • Python targeted tests pass locally.
  • pytest -q chromadb/test/ef/test_sentence_transformer_cache.py chromadb/test/utils/test_embedding_function_config_validation.py chromadb/test/utils/test_embedding_function_schemas.py --disable-warnings: 145 passed, 7 skipped (optional provider dependencies unavailable).
  • New device-isolation case fails before the fix; the two reuse/model-isolation controls pass before and after. Tests fake only the SentenceTransformer constructor, with no model downloads or GPU requirement.
  • Black 23.3 check, flake8 with repository hook arguments, and strict mypy 1.10 with repository hook arguments: passed on both changed files.
  • JS/Rust suites not run; no JS/Rust code changed.

Migration plan

No migration or public configuration change. The process-local cache now distinguishes the device already supplied by callers.

Observability plan

No new instrumentation is required for this local cache correction.

Documentation Changes

Existing device argument documentation remains accurate. No public API changes.

Found by local source inspection and regression testing, not a production incident. AI assistance was used to prepare and review this narrow fix and execute validation. Related open PRs #1915 (default device selection) and #6904 (parallel embedding) were checked; neither corrects model cache keys.

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@github-actions

Copy link
Copy Markdown

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

  • Can you think of any use case in which the code does not behave as intended? Have they been tested?
  • Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
  • If appropriate, are there adequate property based tests?
  • If appropriate, are there adequate unit tests?
  • Should any logging, debugging, tracing information be added or removed?
  • Are error messages user-friendly?
  • Have all documentation changes needed been made?
  • Have all non-obvious changes been commented?

System Compatibility

  • Are there any potential impacts on other parts of the system or backward compatibility?
  • Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

  • Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

This branch has not been deployed

No deployments
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.

1 participant