[Bugfix] Add compilation timeout for JSON schema and grammar in xgrammar

tests/v1/structured_output/test_grammar_compilation_timeout.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for grammar and JSON schema compilation timeout guard.
+
+Verifies that deeply-nested schemas or complex grammars that would cause
+exponential state-space explosion are rejected with a timeout rather than
+hanging worker threads indefinitely (#54003).
+"""
+
+import json
+import time
+from unittest.mock import patch
+
+import pytest
+
+from vllm.v1.structured_output.utils import compile_grammar_with_timeout
+
+
+class TestCompileGrammarWithTimeout:
+    """Unit tests for the compile_grammar_with_timeout utility."""
+
+    def test_normal_json_schema_compiles_successfully(self):
+        schema = '{"type": "object", "properties": {"name": {"type": "string"}}}'
+        result = compile_grammar_with_timeout(
+            lambda s: "compiled",
+            schema,
+            grammar_spec=schema,
+            grammar_type="JSON schema",
+        )
+        assert result == "compiled"
+
+    def test_normal_grammar_compiles_successfully(self):
+        grammar = 'root ::= "hello" | "world"'
+        result = compile_grammar_with_timeout(
+            lambda g: "compiled",
+            grammar,
+            grammar_spec=grammar,
+            grammar_type="Grammar",
+        )
+        assert result == "compiled"
+
+    def test_timeout_raises_value_error(self):
+        def slow_compile(spec: str):
+            time.sleep(10)
+            return "never"
+
+        with (
+            patch("vllm.envs.VLLM_GRAMMAR_COMPILATION_TIMEOUT_S", 1),
+            pytest.raises(ValueError, match="JSON schema compilation timed out"),
+        ):
+            compile_grammar_with_timeout(
+                slow_compile,
+                '{"type": "slow"}',
+                grammar_spec='{"type": "slow"}',
+                grammar_type="JSON schema",
+            )
+
+    def test_timeout_disabled_when_zero(self):
+        result = None
+        with patch("vllm.envs.VLLM_GRAMMAR_COMPILATION_TIMEOUT_S", 0):
+            result = compile_grammar_with_timeout(
+                lambda s: "no_timeout",
+                '{"type": "test"}',
+                grammar_spec='{"type": "test"}',
+                grammar_type="JSON schema",
+            )
+        assert result == "no_timeout"
+
+    def test_compilation_error_propagates(self):
+        def failing_compile(spec: str):
+            raise RuntimeError("compilation failed")
+
+        with pytest.raises(RuntimeError, match="compilation failed"):
+            compile_grammar_with_timeout(
+                failing_compile,
+                "bad",
+                grammar_spec="bad",
+                grammar_type="JSON schema",
+            )
+
+    def test_spec_included_in_error_message(self):
+        def slow_compile(spec: str):
+            time.sleep(10)
+            return "never"
+
+        schema = '{"oneOf": [{"type": "string"}]}'
+        with (
+            patch("vllm.envs.VLLM_GRAMMAR_COMPILATION_TIMEOUT_S", 1),
+            pytest.raises(ValueError, match=r'Spec: \{"oneOf"'),
+        ):
+            compile_grammar_with_timeout(
+                slow_compile,
+                schema,
+                grammar_spec=schema,
+                grammar_type="JSON schema",
+            )
+
+
+def generate_nested_schema(depth: int) -> dict:
+    """Generate a nested oneOf schema that causes exponential branching."""
+    schema: dict = {"type": "string"}
+    for _ in range(depth):
+        schema = {
+            "oneOf": [
+                {"type": "object", "properties": {"x": schema}, "required": ["x"]},
+                {"type": "object", "properties": {"y": schema}, "required": ["y"]},
+            ]
+        }
+    return schema
+
+
+def test_real_xgrammar_compilation_timeout():
+    """Real E2E test verifying that actual xgrammar compilation aborts on timeout."""
+    import xgrammar as xgr
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
+    tok_info = xgr.TokenizerInfo.from_huggingface(tokenizer)
+    compiler = xgr.GrammarCompiler(tok_info)
+
+    # 1. Normal shallow schema compiles quickly (<0.5s)
+    shallow_schema = json.dumps(generate_nested_schema(4))
+    t0 = time.time()
+    result = compile_grammar_with_timeout(
+        compiler.compile_json_schema,
+        shallow_schema,
+        grammar_spec=shallow_schema,
+        grammar_type="JSON schema",
+        timeout=5,
+    )
+    assert result is not None
+    assert (time.time() - t0) < 1.0
+
+    # 2. Deep schema taking >3s to compile must abort at 1s timeout
+    deep_schema = json.dumps(generate_nested_schema(13))
+    t0 = time.time()
+    with pytest.raises(ValueError, match="JSON schema compilation timed out after 1s"):
+        compile_grammar_with_timeout(
+            compiler.compile_json_schema,
+            deep_schema,
+            grammar_spec=deep_schema,
+            grammar_type="JSON schema",
+            timeout=1,
+        )
+    elapsed = time.time() - t0
+    assert 0.9 <= elapsed <= 2.5, f"Expected ~1s timeout, took {elapsed:.2f}s"

vllm/envs.py

@@ -220,6 +220,7 @@
     VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024
     VLLM_XGRAMMAR_CACHE_MB: int = 0
     VLLM_REGEX_COMPILATION_TIMEOUT_S: int = 5
+    VLLM_GRAMMAR_COMPILATION_TIMEOUT_S: int = 10
     VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
     VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
     VLLM_DISABLE_REQUEST_ID_RANDOMIZATION: bool = False
@@ -1634,6 +1635,13 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_REGEX_COMPILATION_TIMEOUT_S": lambda: int(
         os.getenv("VLLM_REGEX_COMPILATION_TIMEOUT_S", "5")
     ),
+    # Maximum time in seconds allowed for JSON schema, CFG grammar, and structural
+    # tag compilation in structured output backends (xgrammar). Prevents unbounded
+    # worker stalls on deeply-nested or pathological schemas.
+    # Set to 0 to disable the timeout (not recommended in production).
+    "VLLM_GRAMMAR_COMPILATION_TIMEOUT_S": lambda: int(
+        os.getenv("VLLM_GRAMMAR_COMPILATION_TIMEOUT_S", "10")
+    ),
     # Control the threshold for msgspec to use 'zero copy' for
     # serialization/deserialization of tensors. Tensors below
     # this limit will be encoded into the msgpack buffer, and

vllm/v1/structured_output/backend_xgrammar.py

@@ -20,6 +20,7 @@
 )
 from vllm.v1.structured_output.utils import (
     choice_as_grammar,
+    compile_grammar_with_timeout,
     compile_regex_with_timeout,
     convert_lark_to_ebnf,
     grammar_is_likely_lark,
@@ -83,15 +84,28 @@ def compile_grammar(
         stop_token_ids: set[int] | None = None,
     ) -> StructuredOutputGrammar:
         if request_type == StructuredOutputOptions.JSON:
-            ctx = self.compiler.compile_json_schema(
-                grammar_spec, any_whitespace=not self.disable_any_whitespace
+            ctx = compile_grammar_with_timeout(
+                self.compiler.compile_json_schema,
+                grammar_spec,
+                grammar_spec=grammar_spec,
+                grammar_type="JSON schema",
+                any_whitespace=not self.disable_any_whitespace,
             )
         elif request_type == StructuredOutputOptions.JSON_OBJECT:
-            ctx = self.compiler.compile_json_schema(
-                '{"type": "object"}', any_whitespace=not self.disable_any_whitespace
+            ctx = compile_grammar_with_timeout(
+                self.compiler.compile_json_schema,
+                '{"type": "object"}',
+                grammar_spec='{"type": "object"}',
+                grammar_type="JSON object schema",
+                any_whitespace=not self.disable_any_whitespace,
             )
         elif request_type == StructuredOutputOptions.GRAMMAR:
-            ctx = self.compiler.compile_grammar(grammar_spec)
+            ctx = compile_grammar_with_timeout(
+                self.compiler.compile_grammar,
+                grammar_spec,
+                grammar_spec=grammar_spec,
+                grammar_type="Grammar",
+            )
         elif request_type == StructuredOutputOptions.REGEX:
             ctx = compile_regex_with_timeout(
                 self.compiler.compile_regex,
@@ -109,9 +123,20 @@ def compile_grammar(
                     )
                     for s in s_tag["structures"]
                 ]
-                ctx = self.compiler.compile_structural_tag(tags, s_tag["triggers"])
+                ctx = compile_grammar_with_timeout(
+                    self.compiler.compile_structural_tag,
+                    tags,
+                    s_tag["triggers"],
+                    grammar_spec=grammar_spec,
+                    grammar_type="Structural tag",
+                )
             else:
-                ctx = self.compiler.compile_structural_tag(grammar_spec)
+                ctx = compile_grammar_with_timeout(
+                    self.compiler.compile_structural_tag,
+                    grammar_spec,
+                    grammar_spec=grammar_spec,
+                    grammar_type="Structural tag",
+                )
         else:
             logger.error(
                 "Validation should have already occurred. Please file an issue."
@@ -329,7 +354,12 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None:
     if so_params.choice:
         choice_grammar = choice_as_grammar(so_params.choice)
         try:
-            xgr.Grammar.from_ebnf(choice_grammar)
+            compile_grammar_with_timeout(
+                xgr.Grammar.from_ebnf,
+                choice_grammar,
+                grammar_spec=choice_grammar,
+                grammar_type="Choice grammar",
+            )
         except Exception as err:
             raise VLLMValidationError(
                 f"Failed to transform choices into a grammar: {err}"
@@ -353,7 +383,12 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None:
             )
 
         try:
-            xgr.Grammar.from_json_schema(schema)
+            compile_grammar_with_timeout(
+                xgr.Grammar.from_json_schema,
+                schema,
+                grammar_spec=str(schema),
+                grammar_type="JSON schema",
+            )
         except Exception as err:
             raise VLLMValidationError(
                 f"Failed to transform json schema into a grammar: {err}"
@@ -373,7 +408,12 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None:
         # Test parsing EBNF grammar, possibly already converted from Lark
         try:
             # parse the grammar, but we aren't compiling it.
-            xgr.Grammar.from_ebnf(so_params.grammar)
+            compile_grammar_with_timeout(
+                xgr.Grammar.from_ebnf,
+                so_params.grammar,
+                grammar_spec=so_params.grammar,
+                grammar_type="Grammar",
+            )
         except Exception as e:
             raise VLLMValidationError("Invalid grammar specification.") from e
         return
@@ -392,8 +432,19 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None:
                     )
                     for s in s_tag["structures"]
                 ]
-                xgr.Grammar.from_structural_tag(tags, s_tag["triggers"])
+                compile_grammar_with_timeout(
+                    xgr.Grammar.from_structural_tag,
+                    tags,
+                    s_tag["triggers"],
+                    grammar_spec=so_params.structural_tag,
+                    grammar_type="Structural tag",
+                )
             else:
-                xgr.Grammar.from_structural_tag(so_params.structural_tag)
+                compile_grammar_with_timeout(
+                    xgr.Grammar.from_structural_tag,
+                    so_params.structural_tag,
+                    grammar_spec=so_params.structural_tag,
+                    grammar_type="Structural tag",
+                )
         except Exception as e:
             raise VLLMValidationError("Invalid structural tag specification.") from e