# # Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. # This file is a part of the vllm-ascend project. # Adapted from vllm/tests/entrypoints/llm/test_guided_generate.py # Copyright 2023 The vLLM team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import json import os from unittest.mock import patch import jsonschema import pytest import regex as re from vllm.outputs import RequestOutput from vllm.sampling_params import SamplingParams, StructuredOutputsParams from tests.e2e.conftest import ModelName from vllm_ascend.utils import vllm_version_is os.environ["VLLM_BATCH_INVARIANT"] = "1" MODEL_NAME = ModelName.QWEN3_06B GuidedDecodingBackend = ["xgrammar", "guidance", "outlines"] REGEX_COMPILATION_TIMEOUT_ENV = {"VLLM_REGEX_COMPILATION_TIMEOUT_S": "30"} @pytest.fixture(params=[False, True], ids=["v1", "v2"]) def model_runner_env(request): use_v2_model_runner = request.param if use_v2_model_runner and vllm_version_is("0.23.0"): pytest.skip("No need to support v2 model runner for vLLM tag version.") with patch.dict(os.environ, {"VLLM_USE_V2_MODEL_RUNNER": "1" if use_v2_model_runner else "0"}): yield @pytest.fixture(scope="module") def sample_regex(): return ( r"((25[0-5]|(2[0-4]|1\d|[1-9]|)\d)\.){3}" r"(25[0-5]|(2[0-4]|1\d|[1-9]|)\d)" ) @pytest.fixture(scope="module") def sample_json_schema(): return { "type": "object", "properties": { "name": {"type": "string"}, "age": {"type": "integer"}, "skills": {"type": "array", "items": {"type": "string", "maxLength": 10}, "minItems": 3}, "work_history": { "type": "array", "items": { "type": "object", "properties": { "company": {"type": "string"}, "duration": {"type": "number"}, "position": {"type": "string"}, }, "required": ["company", "position"], }, }, }, "required": ["name", "age", "skills", "work_history"], } @pytest.mark.timeout(1000) @pytest.mark.model( model_name=MODEL_NAME, compilation_config={"cudagraph_capture_sizes": [1, 2, 4, 8]}, extra_kwargs={"seed": 0, "structured_outputs_config": {"backend": "xgrammar"}}, ) def test_guided_json_completion_xgrammar(sample_json_schema, request): sampling_params = SamplingParams( temperature=1.0, max_tokens=500, structured_outputs=StructuredOutputsParams(json=sample_json_schema) ) if not vllm_version_is("0.23.0"): model_marker = request.node.get_closest_marker("model") model_marker.kwargs["env_vars"] = REGEX_COMPILATION_TIMEOUT_ENV with patch.dict(os.environ, REGEX_COMPILATION_TIMEOUT_ENV, clear=False): vllm_runner = request.getfixturevalue("vllm_runner") prompts = [f"Give an example JSON for an employee profile that fits this schema: {sample_json_schema}"] * 2 inputs = vllm_runner.get_inputs(prompts) outputs = vllm_runner.model.generate(inputs, sampling_params=sampling_params) assert outputs is not None for output in outputs: assert output is not None assert isinstance(output, RequestOutput) prompt = output.prompt generated_text = output.outputs[0].text assert generated_text is not None print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") output_json = json.loads(generated_text) jsonschema.validate(instance=output_json, schema=sample_json_schema) else: vllm_runner = request.getfixturevalue("vllm_runner") prompts = [f"Give an example JSON for an employee profile that fits this schema: {sample_json_schema}"] * 2 inputs = vllm_runner.get_inputs(prompts) outputs = vllm_runner.model.generate(inputs, sampling_params=sampling_params) assert outputs is not None for output in outputs: assert output is not None assert isinstance(output, RequestOutput) prompt = output.prompt generated_text = output.outputs[0].text assert generated_text is not None print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") output_json = json.loads(generated_text) jsonschema.validate(instance=output_json, schema=sample_json_schema) @pytest.mark.timeout(1000) @pytest.mark.model( model_name=MODEL_NAME, compilation_config={"cudagraph_capture_sizes": [1, 2, 4, 8]}, extra_kwargs={"seed": 0, "structured_outputs_config": {"backend": "xgrammar"}}, ) def test_guided_regex_xgrammar(sample_regex, vllm_runner): sampling_params = SamplingParams( temperature=0.8, top_p=0.95, structured_outputs=StructuredOutputsParams(regex=sample_regex) ) prompts = [f"Give an example IPv4 address with this regex: {sample_regex}"] * 2 inputs = vllm_runner.get_inputs(prompts) outputs = vllm_runner.model.generate(inputs, sampling_params=sampling_params) assert outputs is not None for output in outputs: assert output is not None assert isinstance(output, RequestOutput) prompt = output.prompt generated_text = output.outputs[0].text print(generated_text) assert generated_text is not None assert re.fullmatch(".*", generated_text) is not None print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") @pytest.mark.timeout(1000) @pytest.mark.model( model_name=MODEL_NAME, compilation_config={"cudagraph_capture_sizes": [1, 2, 4, 8]}, extra_kwargs={"seed": 0, "structured_outputs_config": {"backend": "guidance"}}, ) def test_guided_json_completion_guidance(sample_json_schema, vllm_runner): sampling_params = SamplingParams( temperature=1.0, max_tokens=500, structured_outputs=StructuredOutputsParams(json=sample_json_schema) ) prompts = [f"Give an example JSON for an employee profile that fits this schema: {sample_json_schema}"] * 2 inputs = vllm_runner.get_inputs(prompts) outputs = vllm_runner.model.generate(inputs, sampling_params=sampling_params) assert outputs is not None for output in outputs: assert output is not None assert isinstance(output, RequestOutput) prompt = output.prompt generated_text = output.outputs[0].text assert generated_text is not None print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") output_json = json.loads(generated_text) jsonschema.validate(instance=output_json, schema=sample_json_schema) @pytest.mark.timeout(1000) @pytest.mark.model( model_name=MODEL_NAME, compilation_config={"cudagraph_capture_sizes": [1, 2, 4, 8]}, extra_kwargs={"seed": 0, "structured_outputs_config": {"backend": "guidance"}}, ) def test_guided_regex_guidance(sample_regex, vllm_runner): sampling_params = SamplingParams( temperature=0.8, top_p=0.95, structured_outputs=StructuredOutputsParams(regex=sample_regex) ) prompts = [f"Give an example IPv4 address with this regex: {sample_regex}"] * 2 inputs = vllm_runner.get_inputs(prompts) outputs = vllm_runner.model.generate(inputs, sampling_params=sampling_params) assert outputs is not None for output in outputs: assert output is not None assert isinstance(output, RequestOutput) prompt = output.prompt generated_text = output.outputs[0].text print(generated_text) assert generated_text is not None assert re.fullmatch(".*", generated_text) is not None print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") @pytest.mark.timeout(1000) @pytest.mark.model( model_name=MODEL_NAME, compilation_config={"cudagraph_capture_sizes": [1, 2, 4, 8]}, extra_kwargs={"seed": 0, "structured_outputs_config": {"backend": "auto"}}, ) def test_guided_auto_rejects_mixed_structured_output_backends(vllm_runner): xgrammar_schema = { "type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"], } guidance_schema = { "type": "object", "properties": {"count": {"type": "integer", "multipleOf": 2}}, "required": ["count"], } xgrammar_params = SamplingParams( temperature=0.0, max_tokens=32, structured_outputs=StructuredOutputsParams(json=xgrammar_schema), ) prompts = [f"Give an example JSON that fits this schema: {xgrammar_schema}"] inputs = vllm_runner.get_inputs(prompts) outputs = vllm_runner.model.generate(inputs, sampling_params=xgrammar_params) assert outputs is not None assert outputs[0] is not None guidance_params = SamplingParams( temperature=0.0, max_tokens=32, structured_outputs=StructuredOutputsParams(json=guidance_schema), ) prompts = [f"Give an example JSON that fits this schema: {guidance_schema}"] inputs = vllm_runner.get_inputs(prompts) with pytest.raises(ValueError, match="already using 'xgrammar'.*'guidance'"): vllm_runner.model.generate(inputs, sampling_params=guidance_params) @pytest.mark.timeout(1000) @pytest.mark.model( model_name=MODEL_NAME, compilation_config={"cudagraph_capture_sizes": [1, 2, 4, 8]}, extra_kwargs={"seed": 0, "structured_outputs_config": {"backend": "outlines"}}, ) def test_guided_json_completion_outlines(sample_json_schema, request): sampling_params = SamplingParams( temperature=1.0, max_tokens=500, structured_outputs=StructuredOutputsParams(json=sample_json_schema) ) if not vllm_version_is("0.23.0"): model_marker = request.node.get_closest_marker("model") model_marker.kwargs["env_vars"] = REGEX_COMPILATION_TIMEOUT_ENV with patch.dict(os.environ, REGEX_COMPILATION_TIMEOUT_ENV, clear=False): vllm_runner = request.getfixturevalue("vllm_runner") prompts = [f"Give an example JSON for an employee profile that fits this schema: {sample_json_schema}"] * 2 inputs = vllm_runner.get_inputs(prompts) outputs = vllm_runner.model.generate(inputs, sampling_params=sampling_params) assert outputs is not None for output in outputs: assert output is not None assert isinstance(output, RequestOutput) prompt = output.prompt generated_text = output.outputs[0].text assert generated_text is not None print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") output_json = json.loads(generated_text) jsonschema.validate(instance=output_json, schema=sample_json_schema) else: vllm_runner = request.getfixturevalue("vllm_runner") prompts = [f"Give an example JSON for an employee profile that fits this schema: {sample_json_schema}"] * 2 inputs = vllm_runner.get_inputs(prompts) outputs = vllm_runner.model.generate(inputs, sampling_params=sampling_params) assert outputs is not None for output in outputs: assert output is not None assert isinstance(output, RequestOutput) prompt = output.prompt generated_text = output.outputs[0].text assert generated_text is not None print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") output_json = json.loads(generated_text) jsonschema.validate(instance=output_json, schema=sample_json_schema)