feat(text): bound vLLM embedding generation

[praateekmahajan]

Description

Main Change - Add model_inference_batch_size support for vLLM too. Otherwise it generates a gigantic tokenized column for all texts which can result in CPU OOM. Now we only tokenize what we feed into the LLM. The output returned is a pyarrow table (earlier it was a pd.DataFrame with a embedding column which was list[float] i.e. fp64. Now we do fp32.

Also related to OOM: we don't now persist text column always

Other changes

1. Use create_vllm_llm_with_retry inside setup
2. Add metadata_fields for embedding (removed text as the default column that gets preserved)

Linear: NMCUR-240

Validation

• 27 focused embedder, utility, and workflow-configuration tests passed
• 12 semantic-dedup workflow tests passed across Xenna and Ray Data
• Pre-commit passed

Benchmark comparison

Compared nightly baseline (ae713b5e) with the PR run (878cf463). Both entries succeeded and processed the same 5,807,132 documents. Time values are seconds; delta is (PR - nightly) / nightly, so negative time deltas are improvements.

Executor	Runtime metric	Nightly	PR	Delta
Ray Data	time_taken_s	901.925	847.189	-6.07%
Ray Data	throughput_docs_per_sec	6,438.60	6,854.59	+6.46%
Ray Data	vLLM process_time_mean	2.962	2.850	-3.79%
Ray Data	custom tokenization_time_mean	0.388	0.543	+39.78%
Ray Data	custom vllm_embedding_time_mean	2.417	2.461	+1.81%
Xenna	time_taken_s	768.562	743.618	-3.25%
Xenna	throughput_docs_per_sec	7,555.84	7,809.30	+3.35%
Xenna	vLLM process_time_mean	3.031	2.910	-4.00%
Xenna	custom tokenization_time_mean	0.413	0.560	+35.44%
Xenna	custom vllm_embedding_time_mean	2.454	2.510	+2.31%

Overall end-to-end runtime improved by 6.1% on Ray Data and 3.2% on Xenna, while embedding-stage process_time_mean improved by about 4% for both executors. The custom tokenization mean increased, but its overlap with vLLM inference keeps the aggregate stage time lower. This is a single-run comparison on different EOS hosts (eos0288 before and eos0183 after), so modest differences should be treated as directional.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[sarahyurick]

LGTM, left one minor minor comment.

[sarahyurick]

Does this one matter at all?

[praateekmahajan]

That was a really good catch 🙏

Before since it was pandas, we didn't have the 2bn list-elements limit, since we moved to pyarrow now, we would've had that limit (iow if embedding_dim=1024, our document batch couldn't have had more than 2mn rows; so i then ended up moving it to ChunkedArray)

This issue isn't really an issue for SemDedup since we're anyway constrained that a single file MUST contain less than 2bn/EMBEDDING_DIM number of rows, however I can imagine there might be some embedding use case out there where we do want to support this

[VibhuJawa]

LGTM