feat(text): bound vLLM embedding generation

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 concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Any
 
+import numpy as np
+import pyarrow as pa
 import torch
 from huggingface_hub import snapshot_download
 
@@ -30,15 +35,20 @@ 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]"
 
 
@@ -54,13 +64,23 @@ 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,
     ):
         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
+        # 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 +101,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 +138,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 +172,155 @@ 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
 
-            if self.tokenizer is None:
-                msg = (
-                    "Tokenizer is not initialized. Please call setup() before processing or set pretokenize to False."
+        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[int, int, list[Any], float]]:
+        """Prepare one chunk ahead while the caller embeds the current chunk.
+
+        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 pending_offset, pending_chunk_size, 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 pending_offset, pending_chunk_size, 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=np.float32,
+        )
+        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[np.ndarray, dict[str, float]]:
+        """Embed bounded chunks and assemble one ordered float32 matrix."""
+        embedding_matrix: np.ndarray | None = None
+        tokenization_time = 0.0
+        vllm_embedding_time = 0.0
+        input_tokens = 0
+
+        for offset, chunk_size, 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"]
+            if embedding_matrix is None:
+                embedding_matrix = np.empty(
+                    (num_rows, chunk_embedding_matrix.shape[1]),
+                    dtype=np.float32,
+                )
+            embedding_matrix[offset : offset + chunk_size] = chunk_embedding_matrix
+            del chunk_embedding_matrix
+
+        return embedding_matrix, {
+            "tokenization_time": tokenization_time,
+            "vllm_embedding_time": vllm_embedding_time,
+            "input_tokens": input_tokens,
+        }
+
+    @staticmethod
+    def _to_arrow_embeddings(embedding_matrix: np.ndarray) -> pa.ListArray:
+        """Convert a dense float32 matrix to an Arrow list array."""
+        embedding_values = pa.array(embedding_matrix.reshape(-1), type=pa.float32(), from_pandas=False)
+        embedding_offsets = pa.array(
+            np.arange(
+                0,
+                (embedding_matrix.shape[0] + 1) * embedding_matrix.shape[1],
+                embedding_matrix.shape[1],
+                dtype=np.int64,
+            )
+        )
+        return pa.ListArray.from_arrays(embedding_offsets, embedding_values)
 
-        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_matrix, metrics = self._collect_embeddings(input_table[self.text_field], input_table.num_rows)
+        embedding_array = self._to_arrow_embeddings(embedding_matrix)
+        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