diff --git a/benchmarking/scripts/nemotron_parse_pdf_benchmark.py b/benchmarking/scripts/nemotron_parse_pdf_benchmark.py index 223cf61127..718b5632c4 100644 --- a/benchmarking/scripts/nemotron_parse_pdf_benchmark.py +++ b/benchmarking/scripts/nemotron_parse_pdf_benchmark.py @@ -28,7 +28,7 @@ from typing import Any from loguru import logger -from utils import collect_parquet_output_metrics, setup_executor, write_benchmark_results +from utils import setup_executor, write_benchmark_results REPO_ROOT = Path(__file__).parent.parent.parent sys.path.insert(0, str(REPO_ROOT / "tutorials" / "interleaved" / "nemotron_parse_pdf")) @@ -38,6 +38,26 @@ create_nemotron_parse_pdf_pipeline, ) +from nemo_curator.tasks.utils import TaskPerfUtils # noqa: E402 + + +def _safe_div(numerator: float, denominator: float) -> float: + return numerator / denominator if denominator else 0.0 + + +def _compute_pdf_parse_metrics(output_tasks: list, run_time_taken: float) -> dict[str, float]: + """Compute benchmark-level throughput metrics from additive task stats.""" + task_metrics = TaskPerfUtils.aggregate_task_metrics(output_tasks, prefix="task") + metric_prefix = "task_nemotron_parse_inference_custom" + + num_valid_pages = task_metrics.get(f"{metric_prefix}.num_valid_pages_sum", 0.0) + total_output_tokens = task_metrics.get(f"{metric_prefix}.total_output_tokens_sum", 0.0) + + return { + "throughput_pages_per_sec": _safe_div(num_valid_pages, run_time_taken), + "throughput_output_tokens_per_sec": _safe_div(total_output_tokens, run_time_taken), + } + def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any]: """Run the Nemotron-Parse PDF benchmark and collect metrics.""" @@ -72,27 +92,13 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any] unique_samples.update(task.data.column("sample_id").to_pylist()) num_pdfs_processed = len(unique_samples) - parquet_metrics = collect_parquet_output_metrics(output_dir) - - stage_perf: dict[str, list[float]] = {} - for task in output_tasks: - for perf in task._stage_perf: - stage_perf.setdefault(perf.stage_name, []).append(perf.process_time) - - stage_summary = {} - for stage_name, times in stage_perf.items(): - stage_summary[stage_name] = { - "count": len(times), - "total_s": sum(times), - "mean_s": sum(times) / len(times) if times else 0, - "min_s": min(times) if times else 0, - "max_s": max(times) if times else 0, - } + pdf_parse_metrics = _compute_pdf_parse_metrics(output_tasks, run_time_taken) logger.success(f"Benchmark completed in {run_time_taken:.2f}s") logger.success(f"Processed {num_pdfs_processed} PDFs") + logger.success(f"Page throughput: {pdf_parse_metrics['throughput_pages_per_sec']:.2f} pages/s") logger.success( - f"Output: {parquet_metrics.get('num_rows', 0)} rows in {parquet_metrics.get('num_output_files', 0)} files" + f"Output token throughput: {pdf_parse_metrics['throughput_output_tokens_per_sec']:.2f} tokens/s" ) success = True @@ -102,8 +108,7 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any] logger.debug(f"Full traceback:\n{error_traceback}") run_time_taken = time.perf_counter() - run_start_time num_pdfs_processed = 0 - parquet_metrics = {} - stage_summary = {} + pdf_parse_metrics = {} return { "params": { @@ -128,8 +133,7 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any] "num_pdfs_processed": num_pdfs_processed, "num_output_tasks": len(output_tasks), "throughput_pdfs_per_sec": num_pdfs_processed / run_time_taken if run_time_taken > 0 else 0, - **parquet_metrics, - "stage_performance": stage_summary, + **pdf_parse_metrics, }, "tasks": output_tasks, } diff --git a/nemo_curator/stages/interleaved/pdf/nemotron_parse/inference.py b/nemo_curator/stages/interleaved/pdf/nemotron_parse/inference.py index 2eb130530c..9ee70cb19c 100644 --- a/nemo_curator/stages/interleaved/pdf/nemotron_parse/inference.py +++ b/nemo_curator/stages/interleaved/pdf/nemotron_parse/inference.py @@ -18,7 +18,9 @@ import contextlib import io +import time from dataclasses import dataclass, field +from typing import Any import pyarrow as pa import torch @@ -69,6 +71,9 @@ class NemotronParseInferenceStage(ProcessingStage[InterleavedBatch, InterleavedB Pages per GPU forward pass (HF backend only). max_num_seqs Maximum concurrent sequences (vLLM backend only). + engine_kwargs + Extra keyword arguments forwarded to the vLLM engine (e.g. + ``gpu_memory_utilization``, ``max_num_batched_tokens``). vLLM backend only. """ model_path: str = DEFAULT_MODEL_PATH @@ -78,6 +83,7 @@ class NemotronParseInferenceStage(ProcessingStage[InterleavedBatch, InterleavedB inference_batch_size: int = 4 max_num_seqs: int = 64 enforce_eager: bool = False + engine_kwargs: dict[str, Any] | None = None name: str = "nemotron_parse_inference" resources: Resources = field(default_factory=lambda: Resources(cpus=4.0, gpus=1.0)) @@ -134,11 +140,12 @@ def _setup_vllm(self) -> None: from nemo_curator.utils.vllm_utils import create_vllm_llm, resolve_local_model_path resolved_path = resolve_local_model_path(self.model_path) - self._llm = create_vllm_llm( - resolved_path, - max_num_seqs=self.max_num_seqs, - enforce_eager=self.enforce_eager, - ) + engine_kwargs = { + "max_num_seqs": self.max_num_seqs, + "enforce_eager": self.enforce_eager, + **(self.engine_kwargs or {}), + } + self._llm = create_vllm_llm(resolved_path, **engine_kwargs) self._sampling_params = SamplingParams( temperature=0, top_k=1, @@ -188,23 +195,80 @@ def _reset_vllm(self) -> None: torch.cuda.empty_cache() self._setup_vllm() - def _infer_vllm(self, images: list[Image.Image]) -> list[str]: + @staticmethod + def _vllm_metrics_from_outputs( # noqa: PLR0913 + outputs: list[Any], + *, + inference_time_s: float, + num_input_pages: int, + num_valid_pages: int, + num_skipped_pages: int, + vllm_retries: int = 0, + ) -> dict[str, float]: + """Build additive per-task vLLM metrics for TaskPerfUtils aggregation.""" + total_prompt_tokens = 0 + total_output_tokens = 0 + total_output_chars = 0 + num_length_truncated = 0 + num_empty_outputs = 0 + + for req_out in outputs: + prompt_ids = getattr(req_out, "prompt_token_ids", None) + if prompt_ids is not None: + total_prompt_tokens += len(prompt_ids) + + if not req_out.outputs: + num_empty_outputs += 1 + continue + + completion = req_out.outputs[0] + token_ids = getattr(completion, "token_ids", None) + if token_ids is not None: + total_output_tokens += len(token_ids) + + text = getattr(completion, "text", "") or "" + total_output_chars += len(text) + if not text.strip(): + num_empty_outputs += 1 + + if getattr(completion, "finish_reason", None) == "length": + num_length_truncated += 1 + + return { + "vllm_inference_time": inference_time_s, + "num_input_pages": float(num_input_pages), + "num_valid_pages": float(num_valid_pages), + "num_skipped_pages": float(num_skipped_pages), + "total_prompt_tokens": float(total_prompt_tokens), + "total_output_tokens": float(total_output_tokens), + "total_output_chars": float(total_output_chars), + "num_output_length_truncated": float(num_length_truncated), + "num_empty_outputs": float(num_empty_outputs), + "vllm_retries": float(vllm_retries), + } + + def _infer_vllm(self, images: list[Image.Image]) -> tuple[list[str], list[Any], int]: if not images: - return [] + return [], [], 0 prompts = [{"prompt": self.task_prompt, "multi_modal_data": {"image": img}} for img in images] max_retries = 3 + vllm_retries = 0 for attempt in range(1, max_retries + 1): try: outputs = self._llm.generate(prompts, self._sampling_params) - return [output.outputs[0].text for output in outputs] except Exception as e: logger.warning(f"[vLLM] Inference failed (attempt {attempt}/{max_retries}): {e}") if attempt < max_retries: + vllm_retries += 1 self._reset_vllm() else: raise - return [] + else: + texts = [output.outputs[0].text if output.outputs else "" for output in outputs] + return texts, outputs, vllm_retries + msg = "unreachable" + raise RuntimeError(msg) def _infer_hf(self, images: list[Image.Image]) -> list[str]: all_outputs: list[str] = [] @@ -233,19 +297,42 @@ def _infer_hf_single_fallback(self, images: list[Image.Image]) -> list[str]: def process(self, task: InterleavedBatch) -> InterleavedBatch | None: task_df = task.to_pandas() images = [] + image_t0 = time.perf_counter() for idx, b in enumerate(task_df["binary_content"]): try: images.append(Image.open(io.BytesIO(b))) except Exception as e: # noqa: BLE001 logger.warning(f"Skipping page {idx} in {task.task_id}: {e}") images.append(None) - + self._log_metrics({"image_load_time": time.perf_counter() - image_t0}) valid_mask = [img is not None for img in images] valid_images = [img for img in images if img is not None] if not valid_images: return None - valid_outputs = self._infer_vllm(valid_images) if self.backend == "vllm" else self._infer_hf(valid_images) + if self.backend == "vllm": + t0 = time.perf_counter() + valid_outputs, raw_outputs, vllm_retries = self._infer_vllm(valid_images) + inference_time_s = time.perf_counter() - t0 + self._log_metrics( + self._vllm_metrics_from_outputs( + raw_outputs, + inference_time_s=inference_time_s, + num_input_pages=len(images), + num_valid_pages=len(valid_images), + num_skipped_pages=len(images) - len(valid_images), + vllm_retries=vllm_retries, + ) + ) + else: + valid_outputs = self._infer_hf(valid_images) + self._log_metrics( + { + "num_input_pages": float(len(images)), + "num_valid_pages": float(len(valid_images)), + "num_skipped_pages": float(len(images) - len(valid_images)), + } + ) all_outputs = [] valid_iter = iter(valid_outputs) diff --git a/nemo_curator/stages/interleaved/pdf/nemotron_parse/partitioning.py b/nemo_curator/stages/interleaved/pdf/nemotron_parse/partitioning.py index 9af5a7cf57..c0c04538e8 100644 --- a/nemo_curator/stages/interleaved/pdf/nemotron_parse/partitioning.py +++ b/nemo_curator/stages/interleaved/pdf/nemotron_parse/partitioning.py @@ -22,6 +22,7 @@ from loguru import logger +from nemo_curator.backends.utils import RayStageSpecKeys from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources from nemo_curator.tasks import EmptyTask, FileGroupTask @@ -80,6 +81,9 @@ def inputs(self) -> tuple[list[str], list[str]]: def outputs(self) -> tuple[list[str], list[str]]: return [], [] + def ray_stage_spec(self) -> dict[str, Any]: + return {RayStageSpecKeys.IS_FANOUT_STAGE: True} + def xenna_stage_spec(self) -> dict[str, Any]: return {"num_workers_per_node": 1} diff --git a/nemo_curator/utils/vllm_utils.py b/nemo_curator/utils/vllm_utils.py index 2ae52dc9e5..4f1c96d907 100644 --- a/nemo_curator/utils/vllm_utils.py +++ b/nemo_curator/utils/vllm_utils.py @@ -30,6 +30,46 @@ from loguru import logger +# Errors that should not be retried. Add entries here when vLLM exposes +# other fatal startup errors through generic EngineCore wrapper messages. +_NON_RETRYABLE_MARKERS = ( + "out of memory", + "cudaerrormemoryallocation", + "device-side assert", + "invalid config", + "invalid model config", +) + +# Substrings that indicate a (usually transient) vLLM engine-startup failure +# caused by a MASTER_PORT collision. In vLLM v1 the bind happens in the +# EngineCore subprocess, so the parent often only sees the wrapped messages +# below rather than the raw "EADDRINUSE". +_ENGINE_STARTUP_FAILURE_MARKERS = ( + "eaddrinuse", + "address already in use", + "engine core initialization failed", + "enginecore failed to start", +) + + +def _exception_chain_text(exc: BaseException) -> str: + parts: list[str] = [] + seen: set[int] = set() + current: BaseException | None = exc + while current is not None and id(current) not in seen: + seen.add(id(current)) + parts.append(str(current)) + current = current.__cause__ or current.__context__ + return " ".join(parts).lower() + + +def _is_engine_startup_failure(exc: BaseException) -> bool: + """Return True if ``exc`` looks like a retryable vLLM engine-startup failure.""" + message = _exception_chain_text(exc) + if any(marker in message for marker in _NON_RETRYABLE_MARKERS): + return False + return any(marker in message for marker in _ENGINE_STARTUP_FAILURE_MARKERS) + def pick_free_port() -> int: """Return a free TCP port on the local machine.""" @@ -49,6 +89,7 @@ def create_vllm_llm( # noqa: PLR0913 trust_remote_code: bool = True, limit_mm_per_prompt: dict | None = None, max_port_retries: int = 3, + **extra_engine_kwargs: object, ) -> "vllm.LLM": # noqa: F821,UP037 """Create a :class:`vllm.LLM` instance with automatic port-collision retry. @@ -74,8 +115,13 @@ def create_vllm_llm( # noqa: PLR0913 Defaults to ``{"image": 1}`` when ``None``. max_port_retries: Number of port-pick attempts before re-raising the error. + extra_engine_kwargs: + Additional keyword arguments forwarded verbatim to :class:`vllm.LLM` + (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 @@ -83,26 +129,40 @@ def create_vllm_llm( # noqa: PLR0913 if limit_mm_per_prompt is None: limit_mm_per_prompt = {"image": 1} + engine_kwargs: dict = { + "model": model_path, + "max_num_seqs": max_num_seqs, + "limit_mm_per_prompt": limit_mm_per_prompt, + "dtype": dtype, + "trust_remote_code": trust_remote_code, + "enforce_eager": enforce_eager, + **extra_engine_kwargs, + } + for attempt in range(1, max_port_retries + 1): free_port = pick_free_port() os.environ["MASTER_PORT"] = str(free_port) try: - return LLM( - model=model_path, - max_num_seqs=max_num_seqs, - limit_mm_per_prompt=limit_mm_per_prompt, - dtype=dtype, - trust_remote_code=trust_remote_code, - enforce_eager=enforce_eager, - ) + return LLM(**engine_kwargs) except RuntimeError as e: - if "EADDRINUSE" in str(e) or "address already in use" in str(e): - logger.warning(f"[vLLM] Port {free_port} collision on attempt {attempt}, retrying...") - time.sleep(2) - if attempt == max_port_retries: - raise - else: + # MASTER_PORT collisions happen when many engines start concurrently + # on one node (e.g. 8 single-GPU replicas under xenna). In vLLM v1 the + # distributed-store bind runs inside the EngineCore subprocess, so the + # original "EADDRINUSE" text does not reach this parent process: the + # caller only sees a generic "Engine core initialization failed". We + # therefore retry on both the direct bind error and that wrapped + # startup failure, picking a fresh port and adding jitter so racing + # workers de-stagger instead of re-colliding on the same retry. + if not _is_engine_startup_failure(e): raise + if attempt == max_port_retries: + logger.error(f"[vLLM] Engine startup failed after {max_port_retries} attempt(s): {e}") + raise + logger.warning( + f"[vLLM] Engine startup failed on attempt {attempt}/{max_port_retries} " + f"(port {free_port}, likely a MASTER_PORT collision), retrying: {e}" + ) + time.sleep(2 + random.uniform(0, 3)) # noqa: S311 - jitter only, not security-sensitive msg = "unreachable" raise RuntimeError(msg) # pragma: no cover diff --git a/tests/stages/interleaved/pdf/nemotron_parse/test_stages.py b/tests/stages/interleaved/pdf/nemotron_parse/test_stages.py index bad16441e3..bce9fc3e43 100644 --- a/tests/stages/interleaved/pdf/nemotron_parse/test_stages.py +++ b/tests/stages/interleaved/pdf/nemotron_parse/test_stages.py @@ -20,7 +20,9 @@ import io import json import zipfile +from types import SimpleNamespace from typing import TYPE_CHECKING +from unittest.mock import patch import pytest @@ -28,6 +30,7 @@ from pathlib import Path from PIL import Image +from nemo_curator.backends.utils import RayStageSpecKeys from nemo_curator.stages.interleaved.pdf.nemotron_parse.partitioning import PDFPartitioningStage from nemo_curator.stages.interleaved.pdf.nemotron_parse.postprocess import NemotronParsePostprocessStage from nemo_curator.stages.interleaved.pdf.nemotron_parse.preprocess import PDFPreprocessStage @@ -39,6 +42,14 @@ def _empty_task() -> EmptyTask: class TestPDFPartitioningStage: + def test_ray_stage_spec_has_fanout(self, tmp_path: Path): + manifest = tmp_path / "manifest.jsonl" + manifest.write_text(json.dumps({"file_name": "a.pdf"}) + "\n") + + stage = PDFPartitioningStage(manifest_path=str(manifest)) + + assert stage.ray_stage_spec()[RayStageSpecKeys.IS_FANOUT_STAGE] is True + def test_simple_manifest(self, tmp_path: Path): manifest = tmp_path / "manifest.jsonl" manifest.write_text( @@ -352,3 +363,160 @@ def test_no_valid_output_returns_none(self): assert result is not None out_df = result.to_pandas() assert out_df.iloc[0]["modality"] == "metadata" + + +class TestNemotronParseInferenceStageMetrics: + def test_vllm_metrics_from_outputs(self) -> None: + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import NemotronParseInferenceStage + + outputs = [ + SimpleNamespace( + prompt_token_ids=[1, 2, 3], + outputs=[SimpleNamespace(text="hello", token_ids=[4, 5, 6], finish_reason="stop")], + ), + SimpleNamespace( + prompt_token_ids=[7, 8], + outputs=[SimpleNamespace(text="", token_ids=[], finish_reason="length")], + ), + ] + + metrics = NemotronParseInferenceStage._vllm_metrics_from_outputs( + outputs, + inference_time_s=1.5, + num_input_pages=3, + num_valid_pages=2, + num_skipped_pages=1, + vllm_retries=1, + ) + + assert metrics["vllm_inference_time"] == 1.5 + assert metrics["num_input_pages"] == 3.0 + assert metrics["num_valid_pages"] == 2.0 + assert metrics["num_skipped_pages"] == 1.0 + assert metrics["total_prompt_tokens"] == 5.0 + assert metrics["total_output_tokens"] == 3.0 + assert metrics["total_output_chars"] == 5.0 + assert metrics["num_output_length_truncated"] == 1.0 + assert metrics["num_empty_outputs"] == 1.0 + assert metrics["vllm_retries"] == 1.0 + assert "avg_output_tokens_per_page" not in metrics + assert "avg_output_chars_per_page" not in metrics + + def test_process_logs_vllm_metrics(self) -> None: + import pandas as pd + import pyarrow as pa + + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import NemotronParseInferenceStage + from nemo_curator.tasks import InterleavedBatch + + img = Image.new("RGB", (10, 10), color="white") + buf = io.BytesIO() + img.save(buf, format="PNG") + + task_df = pd.DataFrame( + [ + { + "sample_id": "s1", + "position": 0, + "modality": "page_image", + "content_type": "image/png", + "text_content": None, + "binary_content": buf.getvalue(), + "source_ref": None, + } + ] + ) + task = InterleavedBatch(dataset_name="test", data=pa.Table.from_pandas(task_df)) + + stage = NemotronParseInferenceStage(backend="vllm") + stage._proc_size = (100, 100) + + raw_outputs = [ + SimpleNamespace( + prompt_token_ids=[1, 2], + outputs=[SimpleNamespace(text="parsed", token_ids=[3, 4, 5], finish_reason="stop")], + ) + ] + + with patch.object(stage, "_infer_vllm", return_value=(["parsed"], raw_outputs, 0)): + stage.process(task) + + assert hasattr(stage, "_custom_metrics") + assert stage._custom_metrics["num_valid_pages"] == 1.0 + assert stage._custom_metrics["total_output_tokens"] == 3.0 + assert stage._custom_metrics["total_output_chars"] == 6.0 + assert "vllm_inference_time" in stage._custom_metrics + + def test_setup_vllm_engine_kwargs_override_stage_defaults(self, monkeypatch: pytest.MonkeyPatch) -> None: + import sys + import types + + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import NemotronParseInferenceStage + from nemo_curator.utils import vllm_utils + + fake_vllm = types.ModuleType("vllm") + + class FakeSamplingParams: + def __init__(self, **kwargs): + self.kwargs = kwargs + + fake_vllm.SamplingParams = FakeSamplingParams + monkeypatch.setitem(sys.modules, "vllm", fake_vllm) + + captured_kwargs: dict = {} + + def fake_create_vllm_llm(model_path: str, **kwargs) -> object: + captured_kwargs["model_path"] = model_path + captured_kwargs.update(kwargs) + return object() + + fake_processor = SimpleNamespace(image_processor=SimpleNamespace(final_size=(100, 100))) + monkeypatch.setattr(vllm_utils, "resolve_local_model_path", lambda _path: "/models/nemotron") + monkeypatch.setattr(vllm_utils, "create_vllm_llm", fake_create_vllm_llm) + monkeypatch.setattr("transformers.AutoProcessor.from_pretrained", lambda *_args, **_kwargs: fake_processor) + + stage = NemotronParseInferenceStage( + backend="vllm", + max_num_seqs=64, + enforce_eager=False, + engine_kwargs={ + "max_num_seqs": 8, + "enforce_eager": True, + "gpu_memory_utilization": 0.9, + }, + ) + + stage._setup_vllm() + + assert captured_kwargs["model_path"] == "/models/nemotron" + assert captured_kwargs["max_num_seqs"] == 8 + assert captured_kwargs["enforce_eager"] is True + assert captured_kwargs["gpu_memory_utilization"] == 0.9 + assert stage._proc_size == (100, 100) + + def test_infer_vllm_empty_outputs_produces_empty_string(self) -> None: + """RequestOutput with no completions should yield '' rather than IndexError.""" + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import NemotronParseInferenceStage + + stage = NemotronParseInferenceStage(backend="vllm") + stage._sampling_params = SimpleNamespace() + # Simulate a RequestOutput where the model returned no completions. + empty_req_output = SimpleNamespace(prompt_token_ids=[1, 2], outputs=[]) + stage._llm = SimpleNamespace(generate=lambda _p, _s: [empty_req_output]) + + texts, raw, retries = stage._infer_vllm([Image.new("RGB", (10, 10))]) + + assert texts == [""] + assert raw == [empty_req_output] + assert retries == 0 + + def test_infer_vllm_unreachable_loop_path_raises(self) -> None: + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import NemotronParseInferenceStage + + stage = NemotronParseInferenceStage(backend="vllm") + stage._sampling_params = SimpleNamespace() + stage._llm = SimpleNamespace(generate=lambda _p, _s: []) + image = Image.new("RGB", (10, 10)) + + with patch("builtins.range", return_value=()), pytest.raises(RuntimeError, match="unreachable"): + stage._infer_vllm([image]) diff --git a/tests/utils/test_vllm_utils.py b/tests/utils/test_vllm_utils.py index 6b79dd180d..97cfec5366 100644 --- a/tests/utils/test_vllm_utils.py +++ b/tests/utils/test_vllm_utils.py @@ -114,6 +114,70 @@ def __init__(self, **_kw): with pytest.raises(RuntimeError, match="address already in use"): _vllm_utils.create_vllm_llm("fake/model", max_port_retries=2) + def test_wrapped_engine_startup_failure_retries(self, monkeypatch: pytest.MonkeyPatch): + """In vLLM v1 the bind error is wrapped; the generic startup message must still retry.""" + call_count = 0 + # This is what the parent process actually sees in v1 (EADDRINUSE is buried + # in the EngineCore subprocess and never reaches here). + err_msg = "Engine core initialization failed. See root cause above. Failed core proc(s): {}" + + class FakeLLM: + def __init__(self, **_kw): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise RuntimeError(err_msg) + + self._inject_fake_vllm(monkeypatch, FakeLLM) + monkeypatch.setattr(_vllm_utils, "pick_free_port", lambda: 12345) + monkeypatch.setattr("time.sleep", lambda _: None) + + result = _vllm_utils.create_vllm_llm("fake/model", max_port_retries=5) + assert isinstance(result, FakeLLM) + assert call_count == 3 + + def test_wrapped_non_retryable_startup_failure_raises_without_retry(self, monkeypatch: pytest.MonkeyPatch): + """Fatal EngineCore failures should fail fast when the exception chain exposes the cause.""" + call_count = 0 + + class FakeLLM: + def __init__(self, **_kw): + nonlocal call_count + call_count += 1 + outer_msg = "Engine core initialization failed" + cause_msg = "CUDA out of memory" + raise RuntimeError(outer_msg) from RuntimeError(cause_msg) + + self._inject_fake_vllm(monkeypatch, FakeLLM) + monkeypatch.setattr(_vllm_utils, "pick_free_port", lambda: 12345) + monkeypatch.setattr("time.sleep", lambda _: None) + + with pytest.raises(RuntimeError, match="Engine core initialization failed"): + _vllm_utils.create_vllm_llm("fake/model", max_port_retries=3) + + assert call_count == 1 + + def test_extra_engine_kwargs_forwarded(self, monkeypatch: pytest.MonkeyPatch): + """Extra engine kwargs (e.g. gpu_memory_utilization) reach vllm.LLM.""" + captured_kwargs: dict = {} + + class FakeLLM: + def __init__(self, **kw): + captured_kwargs.update(kw) + + self._inject_fake_vllm(monkeypatch, FakeLLM) + monkeypatch.setattr(_vllm_utils, "pick_free_port", lambda: 12345) + + _vllm_utils.create_vllm_llm( + "fake/model", + max_num_seqs=128, + gpu_memory_utilization=0.95, + max_num_batched_tokens=16384, + ) + assert captured_kwargs.get("max_num_seqs") == 128 + assert captured_kwargs.get("gpu_memory_utilization") == 0.95 + assert captured_kwargs.get("max_num_batched_tokens") == 16384 + def test_non_eaddrinuse_raises_immediately(self, monkeypatch: pytest.MonkeyPatch): """A non-port-collision RuntimeError should propagate without retry.""" call_count = 0