Skip to content

Embeddings limiting max batch size - #4596

Open
przepeck wants to merge 2 commits into
mainfrom
przepeck/max_embeddings_batch_size
Open

przepeck wants to merge 2 commits into
mainfrom
przepeck/max_embeddings_batch_size

Conversation

@przepeck

Copy link
Copy Markdown
Collaborator

🛠 Summary

CVS-194659
Setting max batch size for embeddings endpoint to 1024
Additionally for adding boundary for maximum tensor size to 1GB

🧪 Checklist

  • Unit tests added.
  • The documentation updated.
  • Change follows security best practices.
    ``

Copilot AI 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.

Copilot review overview

🟡 Changes recommended

The tensor-size guard does not currently bound pre-check allocations, and documentation and regression tests are still needed.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1 High severity · 2 Low severity

Open (3)
What changed in this PR

Adds a 1024-document batch limit and 1 GiB input-tensor cap for embeddings requests.

Changes:

  • Validates string and pre-tokenized batch sizes.
  • Rejects requests exceeding the projected tensor-size limit.
File Summary Findings
src/​embeddings/​embeddings_calculator_ov.cc Implements embeddings batch and tensor-size validation. Critical (3 votes): Validation occurs after tokenization and may not prevent oversized allocations. Nits (3 votes each): Update API documentation and add boundary/regression tests.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment on lines +214 to +218
size_t totalInputBytes = tokens.input_ids.get_byte_size() + tokens.attention_mask.get_byte_size();
if (embeddings_session->getNumberOfModelInputs() == 3) {
totalInputBytes += tokens.input_ids.get_byte_size();
}
auto byteSizeCheckStatus = this->isInputByteSizeOk(totalInputBytes);
Comment on lines +191 to +194
auto batchSizeCheckStatus = this->isBatchSizeOk(receivedBatchSize);
if (!batchSizeCheckStatus.ok()) {
return batchSizeCheckStatus;
}
Comment on lines +455 to +456
const size_t EmbeddingsCalculatorOV::MAX_BATCH_SIZE{1024};
const size_t EmbeddingsCalculatorOV::MAX_INPUT_TENSORS_BYTE_SIZE{1024ull * 1024 * 1024}; // 1 GB
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.

2 participants