fix(embedding): re-normalize per-item index after batch split

[m199369309]

Problem

Under concurrency, the /v1/embeddings endpoint intermittently returns responses whose data[].index is out of range: a 2-input request can come back with index = [1, 2] instead of [0, 1]. The HTTP response is otherwise correct (200, correct count and dimensions, vectors still in input order) — only the index labels are wrong. This violates the OpenAI-compatible contract that data[].index must cover [0, n), and clients that rely on index to map vectors back to their inputs will silently mismatch them.

Root cause

EmbeddingModel.create_embedding.batch (xoscar batching) works in four steps:

1. Group concurrent calls by kwargs and concatenate their inputs into one list (recording per-call offsets).
2. Call _create_embedding once on the whole concatenated list.
3. Engines number EmbeddingData.index with enumerate over the entire batch (global 0..N-1), e.g. sentence_transformers/core.py.
4. Split the result back per call: data[offset:offset+n].

Step 4 slices the correct vectors for each caller, but every EmbeddingData.index still holds the global batch value — the splitter never re-normalizes index to a per-call [0, n). So any call whose inputs land at offset > 0 leaks the global index ([1, 2] instead of [0, 1]).

The defect is in the shared EmbeddingModel base, so it affects every batched embedding engine (sentence_transformers, flag, vllm, llama_cpp). It is intermittent: it only triggers when ≥2 concurrent requests land inside the batch window (XINFERENCE_BATCH_INTERVAL, default 3 ms; allow_batch=True by default) and one sits at a non-zero offset.

Fix

After slicing, rebuild each EmbeddingData with a local index of 0..n-1, preserving object and embedding verbatim. Minimal, localized change — no effect on grouping, sorting, the single-request (non-batch) path, vectors, count, dimensions, or usage. Works for both dense (List[float]) and sparse (Dict[str, float]) embeddings.

 for (offset, n), idx in zip(offsets, indices):
     data = embedding_list["data"][offset : offset + n]
+    # Re-normalize the per-item index to [0, n) for this caller.
+    data = [
+        EmbeddingData(
+            index=j, object=item["object"], embedding=item["embedding"]
+        )
+        for j, item in enumerate(data)
+    ]
     result = Embedding(...)

Test

Adds xinference/model/embedding/tests/test_embedding_batch.py: a stub EmbeddingModel whose _create_embedding reproduces the real engines' global-enumerate behaviour, with no model download. It calls create_embedding.batch(...) directly to deterministically place one call at a non-zero offset:

• test_create_embedding_batch_normalizes_index_per_request — a multi-input call at offset 1 must return [0, 1] (not [1, 2]), and the vectors must be unchanged.
• test_create_embedding_batch_single_input_at_nonzero_offset — a single-input call batched behind another must still return [0].

Both fail on unfixed code ([1, 2] != [0, 1]) and pass after the fix; runs in ~0.01 s.

Compatibility

Backward compatible: only the index label is corrected. No changes to request/response schemas, model registration, CLI flags, or the non-batch path.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request addresses an issue where concurrent batched embedding requests leak global batch indices to individual callers instead of returning normalized indices starting from 0. It resolves this by re-normalizing the per-item index to [0, n) for each caller and adds corresponding unit tests. The review feedback suggests optimizing this process by updating the indices in-place rather than reconstructing the EmbeddingData objects, which also eliminates the need to import EmbeddingData.

[qinxuye]

LGTM