A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Add unit tests for batched guided and non-guided requests (#23389)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
Aug 22, 2025 at 10:31 UTC
b6d7d34fc62947eadf9adcfbc0264da388cb830c
1 file changed
+82
tests/v1/entrypoints/llm/test_struct_output_generate.py
+82
@@ -11,9 +11,11 @@ from typing import TYPE_CHECKING, Any
11
import jsonschema
12
import pytest
13
import regex as re
14
+import torch
15
from pydantic import BaseModel
16
17
from tests.reasoning.utils import run_reasoning_extraction
18
+from vllm.distributed import cleanup_dist_env_and_memory
19
from vllm.entrypoints.llm import LLM
20
from vllm.outputs import RequestOutput
21
from vllm.platforms import current_platform
@@ -727,3 +729,83 @@ def test_guidance_no_additional_properties(monkeypatch: pytest.MonkeyPatch):
729
assert "a4" not in generated
730
assert "a5" not in generated
731
assert "a6" not in generated
732
+
733
+
734
+@pytest.mark.parametrize("guided_decoding_backend",
735
+ ["guidance", "xgrammar", "outlines"])
736
+def test_structured_output_batched_with_non_guided_requests(
737
+ monkeypatch: pytest.MonkeyPatch,
738
+ sample_json_schema: dict[str, Any],
739
+ guided_decoding_backend: str,
740
+):
741
+ monkeypatch.setenv("VLLM_USE_V1", "1")
742
+
743
+ # Don't use eager execution on TPUs because we want to test for no
744
+ # recompilation at runtime
745
+ enforce_eager = bool(not current_platform.is_tpu())
746
+
747
+ llm = LLM(
748
+ model="meta-llama/Meta-Llama-3.1-8B-Instruct",
749
+ enforce_eager=enforce_eager,
750
+ max_model_len=1024,
751
+ guided_decoding_backend=guided_decoding_backend,
752
+ guided_decoding_disable_any_whitespace=(guided_decoding_backend
753
+ in {"xgrammar", "guidance"}),
754
+ )
755
+
756
+ guided_prompt = (
757
+ "Give an example JSON for an employee profile that fits this "
758
+ "schema. Make the response as short as possible. Schema: "
759
+ f"{sample_json_schema}")
760
+
761
+ non_guided_prompt = "The diameter of the Earth in kilometers is "
762
+
763
+ prompts = [guided_prompt, non_guided_prompt]
764
+ sampling_params = [
765
+ SamplingParams(
766
+ temperature=1.0,
767
+ max_tokens=400,
768
+ guided_decoding=GuidedDecodingParams(json=sample_json_schema)),
769
+ # No max tokens, temp=0 to assert on contents
770
+ SamplingParams(
771
+ seed=42,
772
+ temperature=0,
773
+ top_p=1.0,
774
+ ),
775
+ ]
776
+
777
+ outputs = llm.generate(prompts=prompts,
778
+ sampling_params=sampling_params,
779
+ use_tqdm=True)
780
+
781
+ assert outputs is not None
782
+
783
+ # Free memory as soon as possible as failed assertions
784
+ # will short circuit and not free up memory
785
+ del llm
786
+ torch.cuda.empty_cache()
787
+ cleanup_dist_env_and_memory()
788
+
789
+ for index, output in enumerate(outputs):
790
+ assert output is not None
791
+ assert isinstance(output, RequestOutput)
792
+ prompt = output.prompt
793
+
794
+ generated_text = output.outputs[0].text
795
+ assert generated_text is not None
796
+ print(f"Prompt:\n{prompt!r}\nGenerated text:\n{generated_text!r}")
797
+
798
+ if index == 0:
799
+ # First prompt is guided, expect valid JSON
800
+ assert "\n" not in generated_text
801
+ output_json = json.loads(generated_text)
802
+ jsonschema.validate(instance=output_json,
803
+ schema=sample_json_schema)
804
+ else:
805
+ # Second prompt is not guided, expect valid output
806
+ # Cannot assert on exact output, but we can expect it to be factual
807
+ assert "12,742" in generated_text
808
+
809
+ # non-guided requests should not return a valid JSON here
810
+ with pytest.raises(ValueError):
811
+ output_json = json.loads(generated_text)