feat(text): bound vLLM embedding generation

benchmarking/scripts/embedding_generation_benchmark.py

@@ -78,6 +78,7 @@ def _create_embedding_stages(
     model_inference_batch_size: int,
     max_seq_length: int,
     embedding_pooling: str,
+    metadata_fields: list[str] | None = None,
     cache_dir: str | None = None,
 ) -> list:
     """Create the embedding stage(s) for the given model variation."""
@@ -113,6 +114,8 @@ def _create_embedding_stages(
                 model_identifier=model_identifier,
                 text_field="text",
                 embedding_field="embeddings",
+                metadata_fields=metadata_fields,
+                model_inference_batch_size=model_inference_batch_size,
                 pretokenize=model_variation == EmbeddingModelVariation.VLLM_TEXT_PRETOKENIZED,
                 vllm_init_kwargs=vllm_init_kwargs,
                 cache_dir=cache_dir,
@@ -134,6 +137,7 @@ def run_embedding_generation_benchmark(
     embedding_pooling: str,
     input_format: str = "parquet",
     cache_dir: str | None = None,
+    metadata_fields: list[str] | None = None,
     **kwargs: Any,  # noqa: ANN401, ARG001
 ) -> dict[str, Any]:
     """Run the embedding generation benchmark and collect comprehensive metrics."""
@@ -156,6 +160,9 @@ def run_embedding_generation_benchmark(
 
     run_start_time = time.perf_counter()
 
+    metadata_fields = list(dict.fromkeys(metadata_fields or []))
+    input_fields = list(dict.fromkeys(["text", *metadata_fields]))
+    output_fields = [*metadata_fields, "embeddings"]
     keep_ext = "jsonl" if input_format == "jsonl" else "parquet"
     input_files = load_dataset_files(input_path, dataset_size_gb, keep_extensions=keep_ext)
     executor_obj = setup_executor(executor)
@@ -166,15 +173,16 @@ def run_embedding_generation_benchmark(
         model_inference_batch_size=model_inference_batch_size,
         max_seq_length=max_seq_length,
         embedding_pooling=embedding_pooling,
+        metadata_fields=metadata_fields,
         cache_dir=cache_dir,
     )
 
     if input_format == "jsonl":
-        reader = JsonlReader(file_paths=input_files, files_per_partition=1, fields=["text"], _generate_ids=False)
-        writer = JsonlWriter(path=str(output_path), fields=["embeddings"])
+        reader = JsonlReader(file_paths=input_files, files_per_partition=1, fields=input_fields, _generate_ids=False)
+        writer = JsonlWriter(path=str(output_path), fields=output_fields)
     else:
-        reader = ParquetReader(file_paths=input_files, files_per_partition=1, fields=["text"], _generate_ids=False)
-        writer = ParquetWriter(path=str(output_path), fields=["embeddings"])
+        reader = ParquetReader(file_paths=input_files, files_per_partition=1, fields=input_fields, _generate_ids=False)
+        writer = ParquetWriter(path=str(output_path), fields=output_fields)
 
     pipeline = Pipeline(
         name="embedding_generation_pipeline",
@@ -240,6 +248,13 @@ def main() -> int:
         default=None,
         help="HuggingFace cache directory for model weights (uses default HF cache if not set)",
     )
+    parser.add_argument(
+        "--metadata-field",
+        dest="metadata_fields",
+        action="append",
+        default=None,
+        help="Input metadata field to preserve in output; may be repeated",
+    )
 
     args = parser.parse_args()
 

nemo_curator/stages/text/deduplication/semantic.py

@@ -278,6 +278,7 @@ def _run_embedding_generation(self, executor: BaseExecutor) -> list[Task]:
             model_identifier=self.model_identifier,
             text_field=self.text_field,
             embedding_field=self.embedding_field,
+            metadata_fields=list(dict.fromkeys([self.id_field, *(self.metadata_fields or [])])),
             max_chars=self.embedding_max_chars,
             pretokenize=self.embedding_pretokenize,
             vllm_init_kwargs=self.embedding_vllm_init_kwargs,

nemo_curator/stages/text/embedders/vllm.py

@@ -12,10 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import gc
 import time
-from typing import TYPE_CHECKING, Any
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING, Any, Literal
 
+import numpy as np
+import pyarrow as pa
 import torch
 from huggingface_hub import snapshot_download
 
@@ -30,16 +35,22 @@ class LLM:  # dummy for type hints
         pass
 
 
-from nemo_curator.backends.base import NodeInfo, WorkerMetadata
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.stages.resources import Resources
 from nemo_curator.stages.text.models.utils import format_name_with_suffix
 from nemo_curator.tasks import DocumentBatch
+from nemo_curator.utils.vllm_utils import create_vllm_llm_with_retry
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from concurrent.futures import Future
+
     from transformers import AutoTokenizer
 
+    from nemo_curator.backends.base import NodeInfo, WorkerMetadata
+
 _VLLM_INSTALL_HINT = "vLLM is required for VLLMEmbeddingModelStage. Install with: pip install nemo_curator[vllm]"
+_MAX_LIST_ARRAY_VALUES = np.iinfo(np.int32).max
 
 
 class VLLMEmbeddingModelStage(ProcessingStage[DocumentBatch, DocumentBatch]):
@@ -54,13 +65,28 @@ def __init__(  # noqa: PLR0913
         cache_dir: str | None = None,
         hf_token: str | None = None,
         verbose: bool = False,
+        *,
+        metadata_fields: list[str] | None = None,
+        model_inference_batch_size: int | None = 8192,
+        # Keep float32 when feeding semantic dedup; cuDF cannot read nested Float16 Parquet as numeric list values.
+        embedding_output_dtype: Literal["float16", "float32", "float64"] = "float32",
     ):
         self.model_identifier = model_identifier
         self.vllm_init_kwargs = vllm_init_kwargs or {}
 
         self.text_field = text_field
         self.pretokenize = pretokenize
         self.embedding_field = embedding_field
+        self.embedding_output_dtype = embedding_output_dtype
+        # Retained columns are opt-in so large source-text columns are not carried
+        # alongside embeddings unless a caller explicitly requests them.
+        self.metadata_fields = list(dict.fromkeys(metadata_fields or []))
+        if model_inference_batch_size is not None and model_inference_batch_size < 0:
+            msg = (
+                f"model_inference_batch_size must be a non-negative integer or None, got {model_inference_batch_size}"
+            )
+            raise ValueError(msg)
+        self.model_inference_batch_size = model_inference_batch_size
         self.max_chars = max_chars
 
         self.cache_dir = cache_dir
@@ -81,7 +107,10 @@ def inputs(self) -> tuple[list[str], list[str]]:
         return ["data"], [self.text_field]
 
     def outputs(self) -> tuple[list[str], list[str]]:
-        return ["data"], [self.text_field, self.embedding_field]
+        output_fields = list(self.metadata_fields)
+        if self.embedding_field not in output_fields:
+            output_fields.append(self.embedding_field)
+        return ["data"], output_fields
 
     def _initialize_vllm(self, local_files_only: bool) -> None:
         """Download (or locate) the model and initialize vLLM.
@@ -115,7 +144,7 @@ def _initialize_vllm(self, local_files_only: bool) -> None:
         if not self.verbose and "disable_log_stats" not in vllm_init_kwargs:
             vllm_init_kwargs["disable_log_stats"] = True
 
-        self.model = LLM(model=model_path, **vllm_init_kwargs)
+        self.model = create_vllm_llm_with_retry(model=model_path, **vllm_init_kwargs)
 
     def setup_on_node(self, node_info: NodeInfo | None = None, worker_metadata: WorkerMetadata | None = None) -> None:  # noqa: ARG002
         if not self.verbose:
@@ -149,43 +178,149 @@ def setup(self, worker_metadata: WorkerMetadata | None = None) -> None:  # noqa:
                 local_files_only=True,
             )
 
-    def process(self, batch: DocumentBatch) -> DocumentBatch:
-        df = batch.to_pandas()
+    def _prepare_input_chunk(
+        self, text_column: pa.ChunkedArray, offset: int, chunk_size: int
+    ) -> tuple[list[Any], float]:
+        input_data = text_column.slice(offset, chunk_size).to_pylist()
         if self.max_chars is not None:
-            df[self.text_field] = df[self.text_field].str.slice(0, self.max_chars)
-        input_data = df[self.text_field].tolist()
-        metrics = {}
+            input_data = [text[: self.max_chars] if text is not None else None for text in input_data]
+        if not self.pretokenize:
+            return input_data, 0.0
 
-        if self.pretokenize:
-            from vllm.inputs import TokensPrompt
+        from vllm.inputs import TokensPrompt
+
+        t0 = time.perf_counter()
+        tokenized_data = self.tokenizer.batch_encode_plus(
+            input_data,
+            truncation=True,
+            max_length=self.model.model_config.max_model_len,
+        )
+        prompts = [TokensPrompt(prompt_token_ids=ids) for ids in tokenized_data.input_ids]
+        return prompts, time.perf_counter() - t0
+
+    def _iter_prepared_chunks(self, text_column: pa.ChunkedArray, num_rows: int) -> Iterator[tuple[list[Any], float]]:
+        """Prepare one chunk ahead while the caller embeds the current chunk.
 
-            if self.tokenizer is None:
-                msg = (
-                    "Tokenizer is not initialized. Please call setup() before processing or set pretokenize to False."
+        One worker is intentional: this is a single-producer prefetch pipeline,
+        not parallel tokenization. While the caller embeds the current chunk on
+        the GPU, that worker prepares only the immediately following chunk on the
+        CPU. This overlaps CPU and GPU work without allowing an unbounded queue of
+        prepared chunks. ``Future.result`` blocks, so the generator never polls.
+        """
+        inference_batch_size = self.model_inference_batch_size or num_rows
+        chunk_specs = iter(
+            (offset, min(inference_batch_size, num_rows - offset))
+            for offset in range(0, num_rows, inference_batch_size)
+        )
+        pending_offset, pending_chunk_size = next(chunk_specs)
+
+        with ThreadPoolExecutor(max_workers=1, thread_name_prefix="vllm-tokenizer") as executor:
+            pending_input: Future[tuple[list[Any], float]] = executor.submit(
+                self._prepare_input_chunk,
+                text_column,
+                pending_offset,
+                pending_chunk_size,
+            )
+            for next_offset, next_chunk_size in chunk_specs:
+                input_data, tokenization_time = pending_input.result()
+                next_input = executor.submit(
+                    self._prepare_input_chunk,
+                    text_column,
+                    next_offset,
+                    next_chunk_size,
                 )
-                raise ValueError(msg)
 
-            t0 = time.perf_counter()
-            max_model_len = self.model.model_config.max_model_len
-            tokenized_data = self.tokenizer.batch_encode_plus(input_data, truncation=True, max_length=max_model_len)
-            input_data = [TokensPrompt(prompt_token_ids=ids) for ids in tokenized_data.input_ids]
-            metrics["tokenization_time"] = time.perf_counter() - t0
+                yield input_data, tokenization_time
+                pending_offset = next_offset
+                pending_chunk_size = next_chunk_size
+                pending_input = next_input
 
+            input_data, tokenization_time = pending_input.result()
+            yield input_data, tokenization_time
+
+    def _embed_chunk(self, input_data: list[Any]) -> tuple[np.ndarray, dict[str, float]]:
         t0 = time.perf_counter()
         vllm_output = self.model.embed(
             input_data,
             tokenization_kwargs={"truncate_prompt_tokens": -1},
             use_tqdm=self.verbose,
         )
-        metrics["vllm_embedding_time"] = time.perf_counter() - t0
+        elapsed = time.perf_counter() - t0
+        chunk_embedding_matrix = np.asarray(
+            [output.outputs.embedding for output in vllm_output],
+            dtype=self.embedding_output_dtype,
+        )
+        return chunk_embedding_matrix, {
+            "vllm_embedding_time": elapsed,
+            "input_tokens": sum(len(output.prompt_token_ids) for output in vllm_output),
+        }
+
+    def _select_output_table(self, input_table: pa.Table) -> pa.Table:
+        """Validate the input and select columns retained beside embeddings."""
+        if self.text_field not in input_table.column_names:
+            msg = f"Input batch is missing required text field {self.text_field!r}"
+            raise ValueError(msg)
+
+        missing_fields = [field for field in self.metadata_fields if field not in input_table.column_names]
+        if missing_fields:
+            msg = f"Input batch is missing metadata fields: {missing_fields}"
+            raise ValueError(msg)
+        return input_table.select(self.metadata_fields)
+
+    def _collect_embeddings(
+        self, text_column: pa.ChunkedArray, num_rows: int
+    ) -> tuple[pa.ChunkedArray, dict[str, float]]:
+        """Embed bounded chunks and assemble one ordered Arrow array."""
+        embedding_chunks: list[pa.Array] = []
+        tokenization_time = 0.0
+        vllm_embedding_time = 0.0
+        input_tokens = 0
+
+        for input_data, chunk_tokenization_time in self._iter_prepared_chunks(text_column, num_rows):
+            tokenization_time += chunk_tokenization_time
+            chunk_embedding_matrix, chunk_metrics = self._embed_chunk(input_data)
+            vllm_embedding_time += chunk_metrics["vllm_embedding_time"]
+            input_tokens += chunk_metrics["input_tokens"]
+            embedding_chunks.extend(self._to_arrow_embeddings(chunk_embedding_matrix).chunks)
+            del chunk_embedding_matrix
+
+        return pa.chunked_array(embedding_chunks), {
+            "tokenization_time": tokenization_time,
+            "vllm_embedding_time": vllm_embedding_time,
+            "input_tokens": input_tokens,
+        }
+
+    @staticmethod
+    def _to_arrow_embeddings(embedding_matrix: np.ndarray) -> pa.ChunkedArray:
+        """Convert a dense matrix to bounded Arrow list-array chunks."""
+        embedding_dim = embedding_matrix.shape[1]
+        rows_per_chunk = _MAX_LIST_ARRAY_VALUES // embedding_dim
+        value_type = pa.from_numpy_dtype(embedding_matrix.dtype)
+        chunks = []
+        for offset in range(0, embedding_matrix.shape[0], rows_per_chunk):
+            matrix_chunk = embedding_matrix[offset : offset + rows_per_chunk]
+            values = pa.array(matrix_chunk.reshape(-1), type=value_type, from_pandas=False)
+            offsets = pa.array(
+                np.arange(0, matrix_chunk.size + 1, embedding_dim, dtype=np.int32),
+            )
+            chunks.append(pa.ListArray.from_arrays(offsets, values))
+        return pa.chunked_array(chunks, type=pa.list_(value_type))
 
-        df[self.embedding_field] = [e.outputs.embedding for e in vllm_output]
+    def process(self, batch: DocumentBatch) -> DocumentBatch:
+        input_table = batch.to_pyarrow()
+        output_table = self._select_output_table(input_table)
+        embedding_array, metrics = self._collect_embeddings(input_table[self.text_field], input_table.num_rows)
+        if self.embedding_field in output_table.column_names:
+            embedding_index = output_table.column_names.index(self.embedding_field)
+            output_table = output_table.set_column(embedding_index, self.embedding_field, embedding_array)
+        else:
+            output_table = output_table.append_column(self.embedding_field, embedding_array)
 
         self._log_metrics(metrics)
 
         return DocumentBatch(
             dataset_name=batch.dataset_name,
-            data=df,
+            data=output_table,
             _metadata=batch._metadata,
             _stage_perf=batch._stage_perf,
         )

nemo_curator/utils/vllm_utils.py

@@ -161,12 +161,6 @@ def create_vllm_llm(  # noqa: PLR0913
         (e.g. ``gpu_memory_utilization``, ``max_num_batched_tokens``). Keys here
         override the explicit defaults above when they collide.
     """
-    import os
-    import random
-    import time
-
-    from vllm import LLM
-
     if limit_mm_per_prompt is None:
         limit_mm_per_prompt = {"image": 1}
 
@@ -180,6 +174,17 @@ def create_vllm_llm(  # noqa: PLR0913
         **extra_engine_kwargs,
     }
 
+    return create_vllm_llm_with_retry(max_port_retries=max_port_retries, **engine_kwargs)
+
+
+def create_vllm_llm_with_retry(*, max_port_retries: int = 3, **engine_kwargs: object) -> "vllm.LLM":  # noqa: F821,UP037
+    """Create a vLLM engine with port-collision retries and no added defaults."""
+    import os
+    import random
+    import time
+
+    from vllm import LLM
+
     for attempt in range(1, max_port_retries + 1):
         free_port = pick_free_port()
         os.environ["MASTER_PORT"] = str(free_port)

tests/stages/text/deduplication/test_semantic.py

@@ -104,6 +104,7 @@ def capture_pipeline_run(self, executor) -> list[object]:  # noqa: ANN001, ARG00
     workflow._run_embedding_generation(executor=object())
 
     assert captured_stages[0].file_extensions == expected_extensions
+    assert captured_stages[1].metadata_fields == [workflow.id_field]
 
 
 @pytest.mark.gpu