diff --git a/.gitignore b/.gitignore index 1256a84..e272dc2 100644 --- a/.gitignore +++ b/.gitignore @@ -49,3 +49,4 @@ dashboard/dist/ dashboard/.vite/ dashboard/tsconfig.tsbuildinfo dashboard/src/**/*.js +projects/proxy_readers/reports/assistant_axis_message_score/report_22bceb208eaf_fc421b6a/ diff --git a/pipelines_v2/engine/vllm/capture.py b/pipelines_v2/engine/vllm/capture.py index 8a1612a..36b9739 100644 --- a/pipelines_v2/engine/vllm/capture.py +++ b/pipelines_v2/engine/vllm/capture.py @@ -70,6 +70,10 @@ def run_vllm_capture( llm_kwargs["served_model_name"] = engine.canonical_model_name() if engine.max_model_len: llm_kwargs["max_model_len"] = int(engine.max_model_len) + if engine.max_num_batched_tokens is not None: + llm_kwargs["max_num_batched_tokens"] = int(engine.max_num_batched_tokens) + if engine.async_scheduling: + llm_kwargs["async_scheduling"] = True llm_kwargs.update(engine.extra_llm_kwargs()) reasoning_parser = (engine.reasoning_parser or "").strip() if spec.generation.capture_reasoning and not reasoning_parser and "qwen3" in str(engine.model_id).lower(): @@ -107,7 +111,21 @@ def run_vllm_capture( ) tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) - llm = LLM(**llm_kwargs) + _preflight_prompt_lengths( + tokenizer=tokenizer, + examples=examples, + max_model_len=engine.max_model_len, + add_generation_prompt=bool(engine.add_generation_prompt), + generation_max_tokens=spec.generation.max_tokens if wants_generation else None, + tools=spec.generation.chat_tools, + tool_choice=spec.generation.tool_choice, + enable_thinking=engine.enable_thinking, + chat_template_kwargs=dict(engine.extra.get("chat_template_kwargs", {})), + ) + try: + llm = LLM(**llm_kwargs) + except RuntimeError as exc: + raise RuntimeError(_engine_init_error_message(exc, llm_kwargs=llm_kwargs)) from exc reasoning_parser = _build_reasoning_parser( tokenizer=tokenizer, parser_name=reasoning_parser, @@ -123,6 +141,83 @@ def run_vllm_capture( ) +def _preflight_prompt_lengths( + *, + tokenizer: Any, + examples: list[Example], + max_model_len: int | None, + add_generation_prompt: bool, + generation_max_tokens: int | None, + tools: Any = None, + tool_choice: Any = None, + enable_thinking: bool | None = None, + chat_template_kwargs: dict[str, Any] | None = None, +) -> None: + """Fail fast (before model weights load) if any prompt exceeds the context window. + + A single over-long prompt otherwise surfaces minutes later as a per-request + vLLM error after the engine has spun up, which wastes a GPU cold start and + buries the offending example. + """ + + if not max_model_len: + return + reserved = max(1, int(generation_max_tokens or 0)) + limit = int(max_model_len) - reserved + offenders: list[tuple[str, int]] = [] + for example in examples: + token_ids = _tokenize_prompt( + tokenizer=tokenizer, + prompt=example.prompt, + add_generation_prompt=add_generation_prompt, + tools=tools, + tool_choice=tool_choice, + enable_thinking=enable_thinking, + chat_template_kwargs=chat_template_kwargs, + ) + if len(token_ids) > limit: + offenders.append((str(example.key), len(token_ids))) + if offenders: + preview = ", ".join(f"{key}={count} tokens" for key, count in offenders[:10]) + raise SpecValidationError( + f"{len(offenders)} example(s) exceed the usable context window " + f"({limit} tokens = max_model_len {max_model_len} minus {reserved} reserved for generation): " + f"{preview}" + + ("" if len(offenders) <= 10 else f" … and {len(offenders) - 10} more") + + ". Raise max_model_len, shorten the examples, or drop them from the dataset." + ) + + +def _engine_init_error_message(exc: BaseException, *, llm_kwargs: dict[str, Any]) -> str: + """Attach actionable engine config context to vLLM's opaque init failures. + + vLLM's EngineCore subprocess prints the root cause (commonly: KV cache does + not fit) to stdout and re-raises a generic RuntimeError, which is all that + survives into run catalogs. + """ + + summary_keys = ( + "model", + "max_model_len", + "tensor_parallel_size", + "gpu_memory_utilization", + "enforce_eager", + "max_num_seqs", + "max_num_batched_tokens", + "enable_chunked_prefill", + "enable_prefix_caching", + ) + config = ", ".join(f"{key}={llm_kwargs[key]!r}" for key in summary_keys if key in llm_kwargs) + return ( + f"{exc} | engine config: {config} | The root cause is printed by the EngineCore " + "process in the worker logs above this traceback. A common cause is the KV cache " + "not fitting at max_model_len after model weights and the profiled activation peak: " + "enable chunked prefill with a bounded max_num_batched_tokens (the activation peak " + "is profiled at max_model_len when chunked prefill is off), reduce max_model_len, " + "or use a GPU with more memory." + ) + + def run_vllm_capture_with_runtime( *, runtime: Any, diff --git a/pipelines_v2/engine/vllm/engine.py b/pipelines_v2/engine/vllm/engine.py index 88f8ff0..fcc2265 100644 --- a/pipelines_v2/engine/vllm/engine.py +++ b/pipelines_v2/engine/vllm/engine.py @@ -150,10 +150,17 @@ def runtime_spec(self) -> PythonRuntimeSpec: debug_value = str(os.getenv("XENON_ACTIVATION_PATCH_DEBUG", "") or "").strip() if debug_value: env["XENON_ACTIVATION_PATCH_DEBUG"] = debug_value + # vllm is pinned because residual capture relies on the + # `extract_hidden_states` speculative-decoding path, which has changed + # shape across minor releases (0.19.1 shipped a regression in the v1 + # engine's input validator). Unpinned installs mean any Modal image + # rebuild can silently change engine behavior; bump deliberately and + # re-verify capture before raising the pin. Keep in sync with + # yora/modal_base.py. return PythonRuntimeSpec( pip_packages=( "matplotlib", - "vllm", + "vllm==0.19.0", "torch", "transformers", "safetensors", diff --git a/pipelines_v2/operations/readouts/linear.py b/pipelines_v2/operations/readouts/linear.py index d52f8a8..9287849 100644 --- a/pipelines_v2/operations/readouts/linear.py +++ b/pipelines_v2/operations/readouts/linear.py @@ -7,6 +7,7 @@ from pipelines_v2.core.types import OperationSpec, RuntimeSecret from pipelines_v2.operations.common._shared import ( + analysis_runtime_spec, analysis_runtime_spec_for_refs, row_selector_from_dict, runtime_secrets_from_refs, @@ -80,6 +81,9 @@ class PersistedProbeImportSpec(OperationSpec): kind: ClassVar[str] = "persisted_probe_import" + def runtime_spec(self) -> Any | None: + return analysis_runtime_spec() + @classmethod def from_dict(cls, payload: dict[str, Any]) -> "PersistedProbeImportSpec": return cls( diff --git a/tests/pipelines_v2/engine/test_vllm_capture_helpers.py b/tests/pipelines_v2/engine/test_vllm_capture_helpers.py index e8e1fc5..9655278 100644 --- a/tests/pipelines_v2/engine/test_vllm_capture_helpers.py +++ b/tests/pipelines_v2/engine/test_vllm_capture_helpers.py @@ -212,3 +212,64 @@ def test_generation_rows_preserve_reasoning_and_structured_payloads_and_reject_m with pytest.raises(RuntimeError, match="output count does not match"): _generation_rows_from_outputs(examples, rows[:1]) + + +class _CountingTokenizer: + """Tokenizer stub: one token per character of the prompt string.""" + + def __call__(self, text: str, *, add_special_tokens: bool, return_offsets_mapping: bool) -> Any: + return types.SimpleNamespace( + input_ids=list(range(len(text))), + offset_mapping=[(i, i + 1) for i in range(len(text))], + ) + + +@pytest.mark.unit +@pytest.mark.vllm +def test_preflight_prompt_lengths_rejects_overlong_examples() -> None: + from pipelines_v2.core.types import SpecValidationError + from pipelines_v2.engine.vllm.capture import _preflight_prompt_lengths + + tokenizer = _CountingTokenizer() + examples = [ + Example(key="fits", prompt="ab"), + Example(key="too_long", prompt="abcdefgh"), + ] + + # limit = max_model_len - 1 reserved token = 7; "too_long" is 8 tokens. + with pytest.raises(SpecValidationError, match=r"too_long=8 tokens"): + _preflight_prompt_lengths( + tokenizer=tokenizer, + examples=examples, + max_model_len=8, + add_generation_prompt=False, + generation_max_tokens=None, + ) + + # Same prompts pass once the window covers them. + _preflight_prompt_lengths( + tokenizer=tokenizer, + examples=examples, + max_model_len=9, + add_generation_prompt=False, + generation_max_tokens=None, + ) + + # No max_model_len means no constraint to enforce. + _preflight_prompt_lengths( + tokenizer=tokenizer, + examples=examples, + max_model_len=None, + add_generation_prompt=False, + generation_max_tokens=None, + ) + + # Generation budget tightens the usable window. + with pytest.raises(SpecValidationError, match=r"2 example\(s\) exceed"): + _preflight_prompt_lengths( + tokenizer=tokenizer, + examples=examples, + max_model_len=8, + add_generation_prompt=False, + generation_max_tokens=7, + )