[ENH]: add parallel embedding support to SentenceTransformerEmbeddingFunction

[Ashish-Abraham]

Description of changes

Summarize the changes made by this PR.

This PR adds support for data-parallel embedding computation in the SentenceTransformerEmbeddingFunction. Fixes #6903 . This brings it to parity with other local embedding functions (like FastEmbed) and allows users to fully utilize multi-core CPUs and multi-GPU setups.

• Improvements & Bug fixes
  • N/A
• New functionality
  • batch_size support: Users can now specify a custom batch_size for local encoding. It defaults to None, which falls back to the sentence-transformers internal default (32).
  • Multi-process pool support: Introduced the multiprocess_devices parameter. When provided (e.g., ["cpu", "cpu"] or ["cuda:0", "cuda:1"]), a persistent multi-process pool is initialized and used for all embedding calls.
  • Automatic Cleanup: Implemented __del__ to ensure the multi-process pool is properly stopped when the embedding function instance is garbage collected.
  • Schema Integration: Updated the sentence_transformer JSON schema to include the new parameters as nullable fields, ensuring configuration persistence and validation work correctly.

Test plan

How are these changes tested?

• Tests pass locally with pytest for python
  • Created chromadb/test/ef/test_sentence_transformer_ef.py with 7 new unit tests covering initialization, execution, resource cleanup, and configuration roundtrips.
  • Updated chromadb/test/utils/test_embedding_function_schemas.py to verify that the updated schema correctly validates the new fields (including null values).

Migration plan

Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?

No migrations are required. The changes are fully backwards compatible:

• Existing collections without batch_size or multiprocess_devices in their config will default to None, maintaining the previous behavior (single-process, library-default batch size).
• The JSON schema updates include null in the permitted types for the new fields to handle existing deployments.

Observability plan

What is the plan to instrument and monitor this change?

Standard python logging is used. No additional observability instrumentation is required for this local compute change.

Documentation Changes

Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs section]?

• For the [docs section], a small update to the Embedding Functions page would be beneficial to highlight the parallelization capabilities.

[github-actions]

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)

[propel-code-bot]

Add parallel and batched encoding support to SentenceTransformerEmbeddingFunction

This PR introduces new functionality to SentenceTransformerEmbeddingFunction by adding configurable batch_size and persistent multi-process encoding via multiprocess_devices. The implementation initializes a shared sentence-transformers process pool when devices are provided, passes pool/batch options into encode, and includes cleanup logic in __del__ to stop the pool.

It also updates config persistence and validation paths so these new options round-trip correctly through get_config/build_from_config and JSON schema validation. Test coverage is expanded with a new focused test file for initialization, encode argument behavior, pool usage/cleanup, and config roundtrip, plus a schema test validating nullable fields.

This summary was automatically generated by @propel-code-bot

[propel-code-bot]

Review found no issues; changes appear well-implemented, tested, and low risk.

Status: No Issues Found | Risk: Low

Review Details

📁 4 files reviewed | 💬 0 comments