From b96bcd0b474fe469c42423f6a4e7de8ba16baa01 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 13 Aug 2026 12:24:02 -0500 Subject: [PATCH] [ROCm][CI] Solidify entrypoint LLM lifecycle (#51280) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas --- .../llm/offline_mode/test_offline_mode.py | 26 +- tests/entrypoints/llm/test_chat.py | 57 +- tests/entrypoints/llm/test_collective_rpc.py | 12 +- tests/entrypoints/llm/test_generate.py | 62 +- tests/entrypoints/llm/test_gpu_utilization.py | 37 +- .../entrypoints/llm/test_prompt_validation.py | 39 +- .../llm/test_struct_output_generate.py | 1234 +++++++++-------- tests/entrypoints/multimodal/conftest.py | 73 - tests/entrypoints/multimodal/llm/test_chat.py | 13 +- .../llm/test_mm_cache_external_injection.py | 156 +-- .../multimodal/llm/test_mm_cache_stats.py | 45 +- .../multimodal/llm/test_mm_embeds_only.py | 15 +- .../entrypoints/pooling/basic/test_encode.py | 22 +- .../pooling/basic/test_tiling_engine.py | 22 +- .../pooling/classify/test_offline.py | 22 +- .../entrypoints/pooling/embed/test_offline.py | 23 +- .../reward/test_token_reward_offline.py | 22 +- .../scoring/test_bi_encoder_offline.py | 23 +- .../scoring/test_cross_encoder_offline.py | 22 +- .../scoring/test_late_interaction_offline.py | 24 +- .../test_late_interaction_offline_vision.py | 24 +- .../pooling/token_classify/test_offline.py | 22 +- .../pooling/token_embed/test_offline.py | 23 +- .../test_weight_transfer_llm.py | 87 +- 24 files changed, 1002 insertions(+), 1103 deletions(-) diff --git a/tests/entrypoints/llm/offline_mode/test_offline_mode.py b/tests/entrypoints/llm/offline_mode/test_offline_mode.py index 0708597079..555e8b9b7d 100644 --- a/tests/entrypoints/llm/offline_mode/test_offline_mode.py +++ b/tests/entrypoints/llm/offline_mode/test_offline_mode.py @@ -9,9 +9,6 @@ import pytest import regex as re import urllib3 -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory - MODEL_CONFIGS = [ { "model": "facebook/opt-125m", @@ -55,19 +52,26 @@ MODEL_CONFIGS = [ ] +def _create_runner(vllm_runner, model_config): + runner_config = model_config.copy() + model = runner_config.pop("model") + tokenizer_name = runner_config.pop("tokenizer", None) + return vllm_runner(model, tokenizer_name=tokenizer_name, **runner_config) + + @pytest.fixture(scope="module") -def cache_models(): +def cache_models(vllm_runner): # Cache model files first for model_config in MODEL_CONFIGS: - LLM(**model_config) - cleanup_dist_env_and_memory() + with _create_runner(vllm_runner, model_config): + pass yield @pytest.mark.skip_global_cleanup @pytest.mark.usefixtures("cache_models") -def test_offline_mode(monkeypatch: pytest.MonkeyPatch): +def test_offline_mode(monkeypatch: pytest.MonkeyPatch, vllm_runner): # Set HF to offline mode and ensure we can still construct an LLM with monkeypatch.context() as m: try: @@ -93,7 +97,8 @@ def test_offline_mode(monkeypatch: pytest.MonkeyPatch): _re_import_modules() # Cached model files should be used in offline mode for model_config in MODEL_CONFIGS: - LLM(**model_config) + with _create_runner(vllm_runner, model_config): + pass finally: # Reset the environment after the test # NB: Assuming tests are run in online mode @@ -136,7 +141,7 @@ def _re_import_modules(): @pytest.mark.skip_global_cleanup @pytest.mark.usefixtures("cache_models") -def test_model_from_huggingface_offline(monkeypatch: pytest.MonkeyPatch): +def test_model_from_huggingface_offline(monkeypatch: pytest.MonkeyPatch, vllm_runner): # Set HF to offline mode and ensure we can still construct an LLM with monkeypatch.context() as m: try: @@ -159,7 +164,8 @@ def test_model_from_huggingface_offline(monkeypatch: pytest.MonkeyPatch): # Need to re-import huggingface_hub # and friends to set up offline mode _re_import_modules() - LLM(model="facebook/opt-125m") + with vllm_runner("facebook/opt-125m"): + pass finally: # Reset the environment after the test # NB: Assuming tests are run in online mode diff --git a/tests/entrypoints/llm/test_chat.py b/tests/entrypoints/llm/test_chat.py index cbc57b80da..6ee8d34ce1 100644 --- a/tests/entrypoints/llm/test_chat.py +++ b/tests/entrypoints/llm/test_chat.py @@ -4,46 +4,36 @@ import weakref import pytest -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams @pytest.fixture(scope="function") -def text_llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True, seed=0) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() +def text_llm(vllm_runner): + with vllm_runner( + "meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True, seed=0 + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.fixture(scope="function") -def llm_for_failure_test(): +def llm_for_failure_test(vllm_runner): """ Fixture for testing issue #26081. Uses a small max_model_len to easily trigger length errors. """ - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model="meta-llama/Llama-3.2-1B-Instruct", + with vllm_runner( + "meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True, seed=0, max_model_len=128, disable_log_stats=True, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) def test_chat(text_llm): @@ -99,21 +89,16 @@ def test_llm_chat_tokenization_no_double_bos(text_llm): @pytest.fixture(scope="function") -def thinking_llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model="Qwen/Qwen3-0.6B", +def thinking_llm(vllm_runner): + with vllm_runner( + "Qwen/Qwen3-0.6B", max_model_len=4096, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.parametrize("enable_thinking", [True, False]) diff --git a/tests/entrypoints/llm/test_collective_rpc.py b/tests/entrypoints/llm/test_collective_rpc.py index d664558893..36f1436e32 100644 --- a/tests/entrypoints/llm/test_collective_rpc.py +++ b/tests/entrypoints/llm/test_collective_rpc.py @@ -4,15 +4,13 @@ import pytest import torch -from vllm import LLM - from ...utils import create_new_process_for_each_test @pytest.mark.parametrize("tp_size", [1, 2]) @pytest.mark.parametrize("backend", ["mp", "ray"]) @create_new_process_for_each_test() -def test_collective_rpc(tp_size, backend, monkeypatch): +def test_collective_rpc(tp_size, backend, monkeypatch, vllm_runner): if torch.accelerator.device_count() < tp_size: pytest.skip(f"Not enough GPUs for tensor parallelism {tp_size}") if tp_size == 1 and backend == "ray": @@ -26,11 +24,11 @@ def test_collective_rpc(tp_size, backend, monkeypatch): return self.rank monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") - llm = LLM( - model="hmellor/tiny-random-LlamaForCausalLM", + with vllm_runner( + "hmellor/tiny-random-LlamaForCausalLM", enforce_eager=True, load_format="dummy", tensor_parallel_size=tp_size, distributed_executor_backend=backend, - ) - assert llm.collective_rpc(echo_rank) == list(range(tp_size)) + ) as runner: + assert runner.llm.collective_rpc(echo_rank) == list(range(tp_size)) diff --git a/tests/entrypoints/llm/test_generate.py b/tests/entrypoints/llm/test_generate.py index 82f38adfb7..e79e62b361 100644 --- a/tests/entrypoints/llm/test_generate.py +++ b/tests/entrypoints/llm/test_generate.py @@ -6,7 +6,6 @@ import weakref import pytest from vllm import LLM, SamplingParams -from vllm.distributed import cleanup_dist_env_and_memory MODEL_NAME = "distilbert/distilgpt2" @@ -26,22 +25,17 @@ TOKEN_IDS = [ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, max_num_batched_tokens=4096, tensor_parallel_size=1, gpu_memory_utilization=0.10, enforce_eager=True, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup @@ -97,34 +91,34 @@ def test_single_prompt_priority(llm: LLM): assert len(outputs) == 1 -def test_max_model_len(): +def test_max_model_len(vllm_runner): max_model_len = 20 - llm = LLM( - model=MODEL_NAME, + with vllm_runner( + MODEL_NAME, max_model_len=max_model_len, gpu_memory_utilization=0.10, enforce_eager=True, # reduce test time - ) - sampling_params = SamplingParams(max_tokens=max_model_len + 10) - outputs = llm.generate(PROMPTS, sampling_params) - for output in outputs: - num_total_tokens = len(output.prompt_token_ids) + len( - output.outputs[0].token_ids - ) - # Total tokens must not exceed max_model_len. - # It can be less if generation finishes due to other reasons (e.g., EOS) - # before reaching the absolute model length limit. - assert num_total_tokens <= max_model_len + ) as runner: + sampling_params = SamplingParams(max_tokens=max_model_len + 10) + outputs = runner.llm.generate(PROMPTS, sampling_params) + for output in outputs: + num_total_tokens = len(output.prompt_token_ids) + len( + output.outputs[0].token_ids + ) + # Total tokens must not exceed max_model_len. + # It can be less if generation finishes due to other reasons (e.g., EOS) + # before reaching the absolute model length limit. + assert num_total_tokens <= max_model_len -def test_log_stats(): - llm = LLM( - model=MODEL_NAME, +def test_log_stats(vllm_runner): + with vllm_runner( + MODEL_NAME, disable_log_stats=False, gpu_memory_utilization=0.10, enforce_eager=True, # reduce test time - ) - outputs = llm.generate(PROMPTS, sampling_params=None) + ) as runner: + outputs = runner.llm.generate(PROMPTS, sampling_params=None) - # disable_log_stats is False, every output should have metrics - assert all(output.metrics is not None for output in outputs) + # disable_log_stats is False, every output should have metrics + assert all(output.metrics is not None for output in outputs) diff --git a/tests/entrypoints/llm/test_gpu_utilization.py b/tests/entrypoints/llm/test_gpu_utilization.py index 896091533a..fc24a73dd4 100644 --- a/tests/entrypoints/llm/test_gpu_utilization.py +++ b/tests/entrypoints/llm/test_gpu_utilization.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm import LLM, SamplingParams +from vllm import SamplingParams -def test_gpu_memory_utilization(): +def test_gpu_memory_utilization(vllm_runner): prompts = [ "Hello, my name is", "The president of the United States is", @@ -15,13 +15,26 @@ def test_gpu_memory_utilization(): # makes sure gpu_memory_utilization is per-instance limit, # not a global limit - llms = [ - LLM(model="facebook/opt-125m", gpu_memory_utilization=0.3, enforce_eager=True) - for i in range(3) - ] - for llm in llms: - outputs = llm.generate(prompts, sampling_params) - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + with ( + vllm_runner( + "facebook/opt-125m", + gpu_memory_utilization=0.3, + enforce_eager=True, + ) as runner_0, + vllm_runner( + "facebook/opt-125m", + gpu_memory_utilization=0.3, + enforce_eager=True, + ) as runner_1, + vllm_runner( + "facebook/opt-125m", + gpu_memory_utilization=0.3, + enforce_eager=True, + ) as runner_2, + ): + for runner in (runner_0, runner_1, runner_2): + outputs = runner.llm.generate(prompts, sampling_params) + for output in outputs: + prompt = output.prompt + generated_text = output.outputs[0].text + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") diff --git a/tests/entrypoints/llm/test_prompt_validation.py b/tests/entrypoints/llm/test_prompt_validation.py index 8dd55c6b10..6000c03da9 100644 --- a/tests/entrypoints/llm/test_prompt_validation.py +++ b/tests/entrypoints/llm/test_prompt_validation.py @@ -4,30 +4,35 @@ import pytest import torch -from vllm import LLM from vllm.exceptions import VLLMValidationError -def test_empty_prompt(): - llm = LLM(model="openai-community/gpt2", enforce_eager=True) - with pytest.raises(VLLMValidationError, match="decoder prompt cannot be empty"): - llm.generate([""]) +def test_empty_prompt(vllm_runner): + with ( + vllm_runner("openai-community/gpt2", enforce_eager=True) as runner, + pytest.raises(VLLMValidationError, match="decoder prompt cannot be empty"), + ): + runner.llm.generate([""]) -def test_out_of_vocab_token(): - llm = LLM(model="openai-community/gpt2", enforce_eager=True) - with pytest.raises(VLLMValidationError, match="out of vocabulary"): - llm.generate({"prompt_token_ids": [999999]}) +def test_out_of_vocab_token(vllm_runner): + with ( + vllm_runner("openai-community/gpt2", enforce_eager=True) as runner, + pytest.raises(VLLMValidationError, match="out of vocabulary"), + ): + runner.llm.generate({"prompt_token_ids": [999999]}) -def test_require_mm_embeds(): - llm = LLM( - model="llava-hf/llava-1.5-7b-hf", - enforce_eager=True, - enable_mm_embeds=False, - ) - with pytest.raises(ValueError, match="--enable-mm-embeds"): - llm.generate( +def test_require_mm_embeds(vllm_runner): + with ( + vllm_runner( + "llava-hf/llava-1.5-7b-hf", + enforce_eager=True, + enable_mm_embeds=False, + ) as runner, + pytest.raises(ValueError, match="--enable-mm-embeds"), + ): + runner.llm.generate( { "prompt": "", "multi_modal_data": {"image": torch.empty(1, 1, 1)}, diff --git a/tests/entrypoints/llm/test_struct_output_generate.py b/tests/entrypoints/llm/test_struct_output_generate.py index 219ee7cd38..7b90aeacda 100644 --- a/tests/entrypoints/llm/test_struct_output_generate.py +++ b/tests/entrypoints/llm/test_struct_output_generate.py @@ -9,13 +9,10 @@ from typing import Any import jsonschema import pytest import regex as re -import torch from pydantic import BaseModel from tests.reasoning.utils import run_reasoning_extraction from vllm.config import StructuredOutputsConfig -from vllm.distributed import cleanup_dist_env_and_memory -from vllm.entrypoints.llm import LLM from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager @@ -212,11 +209,11 @@ class CarDescription(BaseModel): PARAMS_MODELS_BACKENDS_TOKENIZER_MODE, ) def test_structured_output( - request: pytest.FixtureRequest, backend: str, tokenizer_mode: str, model_name: str, speculative_config: dict[str, Any], + vllm_runner, ): sample_json_schema = SAMPLE_JSON_SCHEMA unsupported_json_schema = UNSUPPORTED_JSON_SCHEMA @@ -229,8 +226,8 @@ def test_structured_output( # Use a single LLM instance for several scenarios to # speed up the test suite. - llm = LLM( - model=model_name, + with vllm_runner( + model_name, enforce_eager=True, max_model_len=1024, structured_outputs_config=dict( @@ -242,65 +239,357 @@ def test_structured_output( config_format="auto" if not model_name.startswith("mistralai/") else "hf", speculative_config=speculative_config, **platform_args, - ) - request.addfinalizer(llm.llm_engine.engine_core.shutdown) - - # - # Test 1: Generate JSON output based on a provided schema - # - sampling_params = SamplingParams( - temperature=1.0, - max_tokens=4096, - structured_outputs=StructuredOutputsParams(json=sample_json_schema), - ) - - prompt = ( - "Give an example JSON for an employee profile that fits this " - "schema. Make the response as short as possible. Schema: " - f"{sample_json_schema}" - ) - outputs = llm.generate( - [prompt] * 2, - sampling_params=sampling_params, - use_tqdm=True, - ) - - 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 - if backend != "lm-format-enforcer": - assert "\n" not in generated_text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - try: - output_json = json.loads(generated_text) - except json.JSONDecodeError as e: - pytest.fail( - f"Invalid JSON from backend={backend}: {generated_text!r}\n" - f"Schema: {sample_json_schema}\nError: {e}" - ) - jsonschema.validate(instance=output_json, schema=sample_json_schema) - - # - # Test 2: Generate JSON object without a schema - # - if backend != "outlines": + ) as runner: + # + # Test 1: Generate JSON output based on a provided schema + # sampling_params = SamplingParams( temperature=1.0, max_tokens=4096, - n=2, - structured_outputs=StructuredOutputsParams(json_object=True), + structured_outputs=StructuredOutputsParams(json=sample_json_schema), ) - outputs = llm.generate( - prompts=( - "Generate a JSON object with curly braces for a person with " - "name and age fields for John Smith who is 31 years old. " + prompt = ( + "Give an example JSON for an employee profile that fits this " + "schema. Make the response as short as possible. Schema: " + f"{sample_json_schema}" + ) + outputs = runner.llm.generate( + [prompt] * 2, + sampling_params=sampling_params, + use_tqdm=True, + ) + + 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 + if backend != "lm-format-enforcer": + assert "\n" not in generated_text + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + try: + output_json = json.loads(generated_text) + except json.JSONDecodeError as e: + pytest.fail( + f"Invalid JSON from backend={backend}: {generated_text!r}\n" + f"Schema: {sample_json_schema}\nError: {e}" + ) + jsonschema.validate(instance=output_json, schema=sample_json_schema) + + # + # Test 2: Generate JSON object without a schema + # + if backend != "outlines": + sampling_params = SamplingParams( + temperature=1.0, + max_tokens=4096, + n=2, + structured_outputs=StructuredOutputsParams(json_object=True), + ) + + outputs = runner.llm.generate( + prompts=( + "Generate a JSON object with curly braces for a person with " + "name and age fields for John Smith who is 31 years old. " + "Make the response as short as possible." + ), + sampling_params=sampling_params, + use_tqdm=True, + ) + + assert outputs is not None + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + + for i in range(2): + generated_text = output.outputs[i].text + print(generated_text) + assert generated_text is not None + + # Parse to verify it is a valid JSON object + parsed_json = json.loads(generated_text) + assert isinstance(parsed_json, dict) + + # + # Test 3: test a jsonschema incompatible with xgrammar + # + sampling_params = SamplingParams( + temperature=1.0, + max_tokens=4096, + structured_outputs=StructuredOutputsParams(json=unsupported_json_schema), + ) + if backend.startswith("xgrammar"): + with pytest.raises( + ValueError, + match="The provided JSON schema contains features " + "not supported by xgrammar.", + ): + prompt = ( + f"Give an example JSON for an employee profile that " + f"fits this schema: {unsupported_json_schema}. " + f"Make the response as short as possible." + ) + runner.llm.generate( + [prompt] * 2, + sampling_params=sampling_params, + use_tqdm=True, + ) + else: + prompt = ( + f"Give an example JSON object for a grade that " + f"fits this schema: {unsupported_json_schema}. " + f"Make the response as short as possible." + ) + outputs = runner.llm.generate( + prompt, + sampling_params=sampling_params, + use_tqdm=True, + ) + assert outputs is not None + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + generated_text = output.outputs[0].text + assert generated_text is not None + print(generated_text) + + # Parse to verify it is valid JSON + parsed_json = json.loads(generated_text) + assert isinstance(parsed_json, dict) + + if backend not in ["outlines", "lm-format-enforcer"]: + # + # Test 4: Generate SQL statement using EBNF grammar + # + sampling_params = SamplingParams( + temperature=0.8, + top_p=0.95, + max_tokens=1000, + structured_outputs=StructuredOutputsParams(grammar=sample_sql_ebnf), + ) + outputs = runner.llm.generate( + ( + "Generate a sql statement that selects col_1 from " + "table_1 where it is equal to 1. Make the response as short as " + "possible." + ), + sampling_params=sampling_params, + use_tqdm=True, + ) + + 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 + + # remove spaces for comparison b/c we removed them in the grammar + ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace( + " ", "" + ) + + assert generated_text.strip() == ground_truth + + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + + # + # Test 5: Generate SQL statement using Lark grammar + # + sampling_params = SamplingParams( + temperature=0.8, + top_p=0.95, + max_tokens=1000, + structured_outputs=StructuredOutputsParams(grammar=sample_sql_lark), + ) + outputs = runner.llm.generate( + ( + "Generate a sql statement that selects col_1 from " + "table_1 where it is equal to 1. Make the response as short as " + "possible." + ), + sampling_params=sampling_params, + use_tqdm=True, + ) + + 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 + + # use Lark to parse the output, and make sure it's a valid parse tree + from lark import Lark + + parser = Lark(sample_sql_lark) + parser.parse(generated_text) + + # remove spaces for comparison b/c we removed them in the grammar + ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace( + " ", "" + ) + + assert generated_text.strip() == ground_truth + + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + + # + # Test 6: Test invalid grammar input + # + sampling_params = SamplingParams( + temperature=0.8, + top_p=0.95, + max_tokens=1000, + structured_outputs=StructuredOutputsParams(grammar="not a grammar"), + ) + with pytest.raises(ValueError, match="Failed to convert the grammar "): + runner.llm.generate( + ( + "Generate a sql statement that selects col_1 from " + "table_1 where it is equal to 1. Make the response as short " + "as possible." + ), + sampling_params=sampling_params, + use_tqdm=True, + ) + + # + # Test 7: Generate text based on a regex pattern + # + sampling_params = SamplingParams( + temperature=0.8, + top_p=0.95, + structured_outputs=StructuredOutputsParams(regex=sample_regex), + ) + + prompt = ( + f"Give an example IPv4 address with this regex: {sample_regex}. " + f"Make the response as short as possible." + ) + outputs = runner.llm.generate( + [prompt] * 2, + sampling_params=sampling_params, + use_tqdm=True, + ) + + 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(sample_regex, generated_text) is not None + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + + # + # Test 8: Generate text based on a choices + # + sampling_params = SamplingParams( + temperature=0.8, + top_p=0.95, + structured_outputs=StructuredOutputsParams( + choice=sample_structured_outputs_choices + ), + ) + + outputs = runner.llm.generate( + ( + "The best language for type-safe systems programming is " + "(Make the response as short as possible.) " + ), + sampling_params=sampling_params, + use_tqdm=True, + ) + 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 generated_text in sample_structured_outputs_choices + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + + # + # Test 9: Generate structured output using a Pydantic model with an enum + # + json_schema = CarDescription.model_json_schema() + sampling_params = SamplingParams( + temperature=1.0, + max_tokens=1000, + structured_outputs=StructuredOutputsParams(json=json_schema), + ) + + outputs = runner.llm.generate( + ( + "Generate a JSON with the brand, model and car_type of the most " + "iconic car from the 90's. Make the response as short as " + "possible." + ), + sampling_params=sampling_params, + use_tqdm=True, + ) + + 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}") + try: + output_json = json.loads(generated_text) + except json.JSONDecodeError as e: + pytest.fail( + f"Invalid JSON from backend={backend}: {generated_text!r}\n" + f"Schema: {json_schema}\nError: {e}" + ) + jsonschema.validate(instance=output_json, schema=json_schema) + + # + # Test 10: Generate structured with minLength and maxLength + # + min_length = 50 + max_length = 50 + json_schema = { + "type": "object", + "properties": { + "description": { + "type": "string", + "maxLength": max_length, + "minLength": min_length, + } + }, + "required": ["description"], + "additionalProperties": False, + } + + sampling_params = SamplingParams( + temperature=1.0, + max_tokens=4096, + structured_outputs=StructuredOutputsParams(json=json_schema), + ) + + outputs = runner.llm.generate( + ( + "Generate a description of a frog using 50 characters. " "Make the response as short as possible." ), sampling_params=sampling_params, @@ -308,87 +597,7 @@ def test_structured_output( ) assert outputs is not None - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - for i in range(2): - generated_text = output.outputs[i].text - print(generated_text) - assert generated_text is not None - - # Parse to verify it is a valid JSON object - parsed_json = json.loads(generated_text) - assert isinstance(parsed_json, dict) - - # - # Test 3: test a jsonschema incompatible with xgrammar - # - sampling_params = SamplingParams( - temperature=1.0, - max_tokens=4096, - structured_outputs=StructuredOutputsParams(json=unsupported_json_schema), - ) - if backend.startswith("xgrammar"): - with pytest.raises( - ValueError, - match="The provided JSON schema contains features " - "not supported by xgrammar.", - ): - prompt = ( - f"Give an example JSON for an employee profile that " - f"fits this schema: {unsupported_json_schema}. " - f"Make the response as short as possible." - ) - llm.generate( - [prompt] * 2, - sampling_params=sampling_params, - use_tqdm=True, - ) - else: - prompt = ( - f"Give an example JSON object for a grade that " - f"fits this schema: {unsupported_json_schema}. " - f"Make the response as short as possible." - ) - outputs = llm.generate( - prompt, - sampling_params=sampling_params, - use_tqdm=True, - ) - assert outputs is not None - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - generated_text = output.outputs[0].text - assert generated_text is not None - print(generated_text) - - # Parse to verify it is valid JSON - parsed_json = json.loads(generated_text) - assert isinstance(parsed_json, dict) - - if backend not in ["outlines", "lm-format-enforcer"]: - # - # Test 4: Generate SQL statement using EBNF grammar - # - sampling_params = SamplingParams( - temperature=0.8, - top_p=0.95, - max_tokens=1000, - structured_outputs=StructuredOutputsParams(grammar=sample_sql_ebnf), - ) - outputs = llm.generate( - ( - "Generate a sql statement that selects col_1 from " - "table_1 where it is equal to 1. Make the response as short as " - "possible." - ), - sampling_params=sampling_params, - use_tqdm=True, - ) - - assert outputs is not None for output in outputs: assert output is not None assert isinstance(output, RequestOutput) @@ -396,255 +605,45 @@ def test_structured_output( generated_text = output.outputs[0].text assert generated_text is not None - - # remove spaces for comparison b/c we removed them in the grammar - ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(" ", "") - - assert generated_text.strip() == ground_truth - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + try: + output_json = json.loads(generated_text) + except json.JSONDecodeError as e: + pytest.fail( + f"Invalid JSON from backend={backend}: {generated_text!r}\n" + f"Schema: {json_schema}\nError: {e}" + ) + jsonschema.validate(instance=output_json, schema=json_schema) - # - # Test 5: Generate SQL statement using Lark grammar - # - sampling_params = SamplingParams( - temperature=0.8, - top_p=0.95, - max_tokens=1000, - structured_outputs=StructuredOutputsParams(grammar=sample_sql_lark), - ) - outputs = llm.generate( - ( - "Generate a sql statement that selects col_1 from " - "table_1 where it is equal to 1. Make the response as short as " - "possible." - ), - sampling_params=sampling_params, - use_tqdm=True, - ) - - 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 - - # use Lark to parse the output, and make sure it's a valid parse tree - from lark import Lark - - parser = Lark(sample_sql_lark) - parser.parse(generated_text) - - # remove spaces for comparison b/c we removed them in the grammar - ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(" ", "") - - assert generated_text.strip() == ground_truth - - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - # - # Test 6: Test invalid grammar input - # - sampling_params = SamplingParams( - temperature=0.8, - top_p=0.95, - max_tokens=1000, - structured_outputs=StructuredOutputsParams(grammar="not a grammar"), - ) - with pytest.raises(ValueError, match="Failed to convert the grammar "): - llm.generate( - ( - "Generate a sql statement that selects col_1 from " - "table_1 where it is equal to 1. Make the response as short " - "as possible." - ), - sampling_params=sampling_params, - use_tqdm=True, - ) - - # - # Test 7: Generate text based on a regex pattern - # - sampling_params = SamplingParams( - temperature=0.8, - top_p=0.95, - structured_outputs=StructuredOutputsParams(regex=sample_regex), - ) - - prompt = ( - f"Give an example IPv4 address with this regex: {sample_regex}. " - f"Make the response as short as possible." - ) - outputs = llm.generate( - [prompt] * 2, - sampling_params=sampling_params, - use_tqdm=True, - ) - - 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(sample_regex, generated_text) is not None - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - # - # Test 8: Generate text based on a choices - # - sampling_params = SamplingParams( - temperature=0.8, - top_p=0.95, - structured_outputs=StructuredOutputsParams( - choice=sample_structured_outputs_choices - ), - ) - - outputs = llm.generate( - ( - "The best language for type-safe systems programming is " - "(Make the response as short as possible.) " - ), - sampling_params=sampling_params, - use_tqdm=True, - ) - 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 generated_text in sample_structured_outputs_choices - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - # - # Test 9: Generate structured output using a Pydantic model with an enum - # - json_schema = CarDescription.model_json_schema() - sampling_params = SamplingParams( - temperature=1.0, - max_tokens=1000, - structured_outputs=StructuredOutputsParams(json=json_schema), - ) - - outputs = llm.generate( - ( - "Generate a JSON with the brand, model and car_type of the most " - "iconic car from the 90's. Make the response as short as " - "possible." - ), - sampling_params=sampling_params, - use_tqdm=True, - ) - - 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}") - try: - output_json = json.loads(generated_text) - except json.JSONDecodeError as e: - pytest.fail( - f"Invalid JSON from backend={backend}: {generated_text!r}\n" - f"Schema: {json_schema}\nError: {e}" - ) - jsonschema.validate(instance=output_json, schema=json_schema) - - # - # Test 10: Generate structured with minLength and maxLength - # - min_length = 50 - max_length = 50 - json_schema = { - "type": "object", - "properties": { - "description": { - "type": "string", - "maxLength": max_length, - "minLength": min_length, + if backend not in ["outlines", "lm-format-enforcer"]: + # + # Test 11: Generate structured output using structural_tag format + # + structural_tag_config = { + "type": "structural_tag", + "structures": [ + { + "begin": "", + "schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "additionalProperties": False, + }, + "end": "", + } + ], + "triggers": ["", - "schema": { - "type": "object", - "properties": {"city": {"type": "string"}}, - "additionalProperties": False, - }, - "end": "", - } - ], - "triggers": ["(.*?)" - matches = re.findall(function_call_pattern, generated_text) + # Search for function call pattern in the response + function_call_pattern = r"(.*?)" + matches = re.findall(function_call_pattern, generated_text) - if not matches: - print( - f"Warning: No function calls found in response: {generated_text!r}" - ) - continue + if not matches: + print( + f"Warning: No function calls found in response: {generated_text!r}" + ) + continue - # Take the first function call if multiple are found - json_str = matches[0] - try: - json_content = json.loads(json_str) - assert "city" in json_content - assert isinstance(json_content["city"], str) - print(f"Found valid function call: {generated_text!r}") - except (json.JSONDecodeError, AssertionError) as e: - pytest.fail( - f"Invalid function call format: {generated_text!r}\nError: {str(e)}" - ) + # Take the first function call if multiple are found + json_str = matches[0] + try: + json_content = json.loads(json_str) + assert "city" in json_content + assert isinstance(json_content["city"], str) + print(f"Found valid function call: {generated_text!r}") + except (json.JSONDecodeError, AssertionError) as e: + pytest.fail( + f"Invalid function call format: {generated_text!r}\nError: {str(e)}" + ) @pytest.mark.parametrize( @@ -738,14 +739,15 @@ def test_structured_output_with_reasoning_matrices( model_name: str, speculative_config: dict[str, Any] | None, async_scheduling: bool, + vllm_runner, ): if current_platform.is_tpu() and speculative_config: pytest.skip("TPU does not support speculative decoding") # Use a single LLM instance for several scenarios to # speed up the test suite. - llm = LLM( - model=model_name, + with vllm_runner( + model_name, # Don't use eager execution on TPUs because we want to test for no # recompilation at runtime enforce_eager=bool(not current_platform.is_tpu()), @@ -759,271 +761,273 @@ def test_structured_output_with_reasoning_matrices( tokenizer_mode=tokenizer_mode, speculative_config=speculative_config, async_scheduling=async_scheduling, - ) - tokenizer = llm.get_tokenizer() - reasoner = ReasoningParserManager.get_reasoning_parser(reasoning_parser)( - tokenizer=tokenizer - ) + ) as runner: + tokenizer = runner.llm.get_tokenizer() + reasoner = ReasoningParserManager.get_reasoning_parser(reasoning_parser)( + tokenizer=tokenizer + ) - reasoning_prompt = "Solve the following math problem step-by-step, then provide the final answer as JSON object with a single key 'result'. Make sure to correct your reasoning if there are any issue should it arise.\nProblem: What is 5 * 8 + 2?" # noqa: E501 - reasoning_schema = { - "type": "object", - "properties": {"result": {"type": "integer"}}, - "required": ["result"], - "additionalProperties": False, - } - if "Qwen3" in model_name: - reasoning_prompt += "\n" + reasoning_prompt = "Solve the following math problem step-by-step, then provide the final answer as JSON object with a single key 'result'. Make sure to correct your reasoning if there are any issue should it arise.\nProblem: What is 5 * 8 + 2?" # noqa: E501 + reasoning_schema = { + "type": "object", + "properties": {"result": {"type": "integer"}}, + "required": ["result"], + "additionalProperties": False, + } + if "Qwen3" in model_name: + reasoning_prompt += "\n" - sampling_params = SamplingParams( - temperature=0.1, - max_tokens=8192, - structured_outputs=StructuredOutputsParams(json=reasoning_schema), - ) - outputs = llm.generate( - [reasoning_prompt], - sampling_params=sampling_params, - use_tqdm=True, - ) + sampling_params = SamplingParams( + temperature=0.1, + max_tokens=8192, + structured_outputs=StructuredOutputsParams(json=reasoning_schema), + ) + outputs = runner.llm.generate( + [reasoning_prompt], + sampling_params=sampling_params, + use_tqdm=True, + ) - assert outputs is not None - output = outputs[0] - assert output is not None and isinstance(output, RequestOutput) - prompt = output.prompt - generated_text = output.outputs[0].text - reasoning, content = run_reasoning_extraction(reasoner, [generated_text]) - print(f"Prompt: {prompt!r}\nReasoning: {reasoning!r}\nContent: {content!r}") + assert outputs is not None + output = outputs[0] + assert output is not None and isinstance(output, RequestOutput) + prompt = output.prompt + generated_text = output.outputs[0].text + reasoning, content = run_reasoning_extraction(reasoner, [generated_text]) + print(f"Prompt: {prompt!r}\nReasoning: {reasoning!r}\nContent: {content!r}") - if "Qwen3" in model_name: - assert content is not None + if "Qwen3" in model_name: + assert content is not None - assert reasoning is not None + assert reasoning is not None - if content is not None: - output_json = json.loads(content) - jsonschema.validate(instance=output_json, schema=reasoning_schema) + if content is not None: + output_json = json.loads(content) + jsonschema.validate(instance=output_json, schema=reasoning_schema) @pytest.mark.parametrize("model_name, tokenizer_mode", PARAMS_MODELS_TOKENIZER_MODE) def test_structured_output_auto_mode( model_name: str, tokenizer_mode: str, + vllm_runner, ): unsupported_json_schema = UNSUPPORTED_JSON_SCHEMA - llm = LLM( - model=model_name, + with vllm_runner( + model_name, max_model_len=1024, structured_outputs_config=dict(backend="auto"), tokenizer_mode=tokenizer_mode, load_format="auto", config_format="auto", - ) + ) as runner: + sampling_params = SamplingParams( + temperature=1.0, + max_tokens=1000, + structured_outputs=StructuredOutputsParams(json=unsupported_json_schema), + ) - sampling_params = SamplingParams( - temperature=1.0, - max_tokens=1000, - structured_outputs=StructuredOutputsParams(json=unsupported_json_schema), - ) + prompts = ( + "Give an example JSON object for a grade " + "that fits this schema: " + f"{unsupported_json_schema}. Make the response as short as possible." + ) + # This would fail with the default of "xgrammar", but in "auto" + # we will handle fallback automatically. + outputs = runner.llm.generate( + prompts, sampling_params=sampling_params, use_tqdm=True + ) + # Make sure `auto` backend handling doesn't mess up sampling_params + # and that we can reuse it without error. + outputs.extend( + runner.llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True) + ) - prompts = ( - "Give an example JSON object for a grade " - "that fits this schema: " - f"{unsupported_json_schema}. Make the response as short as possible." - ) - # This would fail with the default of "xgrammar", but in "auto" - # we will handle fallback automatically. - outputs = llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True) - # Make sure `auto` backend handling doesn't mess up sampling_params - # and that we can reuse it without error. - outputs.extend( - llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True) - ) + assert outputs is not None + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + generated_text = output.outputs[0].text + assert generated_text is not None + print(generated_text) - assert outputs is not None - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - generated_text = output.outputs[0].text - assert generated_text is not None - print(generated_text) - - # Parse to verify it is valid JSON - parsed_json = json.loads(generated_text) - assert isinstance(parsed_json, dict) + # Parse to verify it is valid JSON + parsed_json = json.loads(generated_text) + assert isinstance(parsed_json, dict) -def test_guidance_no_additional_properties(): - llm = LLM( - model="Qwen/Qwen2.5-1.5B-Instruct", +def test_guidance_no_additional_properties(vllm_runner): + with vllm_runner( + "Qwen/Qwen2.5-1.5B-Instruct", max_model_len=1024, structured_outputs_config=dict( backend="guidance", disable_any_whitespace=True, disable_additional_properties=True, ), - ) + ) as runner: + schema = { + "type": "object", + "properties": { + "a1": {"type": "string"}, + "a2": {"type": "string"}, + "a3": {"type": "string"}, + }, + "required": ["a1", "a2", "a3"], + } - schema = { - "type": "object", - "properties": { - "a1": {"type": "string"}, - "a2": {"type": "string"}, - "a3": {"type": "string"}, - }, - "required": ["a1", "a2", "a3"], - } - - prompt = ( - "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a " - "helpful assistant.<|im_end|>\n<|im_start|>user\nPlease generate a " - "large JSON object with key-value pairs a1=b1, a2=b2, ..., a20=b20. " - "Make the response as short as possible." - "<|im_end|>\n<|im_start|>assistant\n" - ) - - def generate_with_backend(backend): - structured_outputs_params = StructuredOutputsParams( - json=schema, - backend=backend, - disable_any_whitespace=True, - disable_additional_properties=True, - ) - sampling_params = SamplingParams( - temperature=0, max_tokens=256, structured_outputs=structured_outputs_params + prompt = ( + "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a " + "helpful assistant.<|im_end|>\n<|im_start|>user\nPlease generate a " + "large JSON object with key-value pairs a1=b1, a2=b2, ..., a20=b20. " + "Make the response as short as possible." + "<|im_end|>\n<|im_start|>assistant\n" ) - outputs = llm.generate(prompt, sampling_params=sampling_params) - assert outputs is not None - generated_text = outputs[0].outputs[0].text - assert generated_text is not None - parsed_json = json.loads(generated_text) - assert isinstance(parsed_json, dict) - jsonschema.validate(instance=parsed_json, schema=schema) - return parsed_json + def generate_with_backend(backend): + structured_outputs_params = StructuredOutputsParams( + json=schema, + backend=backend, + disable_any_whitespace=True, + disable_additional_properties=True, + ) + sampling_params = SamplingParams( + temperature=0, + max_tokens=256, + structured_outputs=structured_outputs_params, + ) - generated = generate_with_backend("guidance") - assert "a1" in generated - assert "a2" in generated - assert "a3" in generated - assert "a4" not in generated - assert "a5" not in generated - assert "a6" not in generated + outputs = runner.llm.generate(prompt, sampling_params=sampling_params) + assert outputs is not None + generated_text = outputs[0].outputs[0].text + assert generated_text is not None + parsed_json = json.loads(generated_text) + assert isinstance(parsed_json, dict) + jsonschema.validate(instance=parsed_json, schema=schema) + return parsed_json + + generated = generate_with_backend("guidance") + assert "a1" in generated + assert "a2" in generated + assert "a3" in generated + assert "a4" not in generated + assert "a5" not in generated + assert "a6" not in generated @pytest.mark.parametrize("backend", ["guidance", "xgrammar", "outlines"]) def test_structured_output_batched_with_non_structured_outputs_requests( backend: str, + vllm_runner, ): sample_json_schema = SAMPLE_JSON_SCHEMA # Don't use eager execution on TPUs because we want to test for no # recompilation at runtime enforce_eager = bool(not current_platform.is_tpu()) - llm = LLM( - model="meta-llama/Meta-Llama-3.1-8B-Instruct", + with vllm_runner( + "meta-llama/Meta-Llama-3.1-8B-Instruct", enforce_eager=enforce_eager, max_model_len=1024, structured_outputs_config=StructuredOutputsConfig( backend=backend, disable_any_whitespace=backend in {"xgrammar", "guidance"}, ), - ) + ) as runner: + structured_outputs_prompt = ( + "Give an example JSON for an employee profile that fits this " + "schema. Make the response as short as possible. Schema: " + f"{sample_json_schema}" + ) - structured_outputs_prompt = ( - "Give an example JSON for an employee profile that fits this " - "schema. Make the response as short as possible. Schema: " - f"{sample_json_schema}" - ) + non_structured_outputs_prompt = "The diameter of the Earth in kilometers is " - non_structured_outputs_prompt = "The diameter of the Earth in kilometers is " + prompts = [structured_outputs_prompt, non_structured_outputs_prompt] + sampling_params = [ + SamplingParams( + temperature=1.0, + max_tokens=400, + structured_outputs=StructuredOutputsParams(json=sample_json_schema), + ), + # No max tokens, temp=0 to assert on contents + SamplingParams( + seed=42, + temperature=0, + top_p=1.0, + ), + ] - prompts = [structured_outputs_prompt, non_structured_outputs_prompt] - sampling_params = [ - SamplingParams( - temperature=1.0, - max_tokens=400, - structured_outputs=StructuredOutputsParams(json=sample_json_schema), - ), - # No max tokens, temp=0 to assert on contents - SamplingParams( - seed=42, - temperature=0, - top_p=1.0, - ), - ] + outputs = runner.llm.generate( + prompts=prompts, sampling_params=sampling_params, use_tqdm=True + ) - outputs = llm.generate( - prompts=prompts, sampling_params=sampling_params, use_tqdm=True - ) + assert outputs is not None - assert outputs is not None + for index, output in enumerate(outputs): + assert output is not None + assert isinstance(output, RequestOutput) + prompt = output.prompt - # Free memory as soon as possible as failed assertions - # will short circuit and not free up memory - del llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + generated_text = output.outputs[0].text + assert generated_text is not None + print(f"Prompt:\n{prompt!r}\nGenerated text:\n{generated_text!r}") - for index, output in enumerate(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:\n{prompt!r}\nGenerated text:\n{generated_text!r}") - - if index == 0: - # First prompt is structured outputs, expect valid JSON - assert "\n" not in generated_text - output_json = json.loads(generated_text) - jsonschema.validate(instance=output_json, schema=sample_json_schema) - else: - # Second prompt is not structured outputs, expect valid output - # Cannot assert on exact output, but we can expect it to be factual - assert "12,742" in generated_text - - # non-structured outputs requests should not return a valid JSON here - with pytest.raises(ValueError): + if index == 0: + # First prompt is structured outputs, expect valid JSON + assert "\n" not in generated_text output_json = json.loads(generated_text) + jsonschema.validate(instance=output_json, schema=sample_json_schema) + else: + # Second prompt is not structured outputs, expect valid output + # Cannot assert on exact output, but we can expect it to be factual + assert "12,742" in generated_text + + # non-structured outputs requests should not return a valid JSON here + with pytest.raises(ValueError): + output_json = json.loads(generated_text) @pytest.mark.parametrize("backend", ["xgrammar"]) -def test_structured_output_with_structural_tag(backend: str): - llm = LLM( - model="Qwen/Qwen2.5-1.5B-Instruct", +def test_structured_output_with_structural_tag(backend: str, vllm_runner): + with vllm_runner( + "Qwen/Qwen2.5-1.5B-Instruct", structured_outputs_config=StructuredOutputsConfig(backend=backend), - ) + ) as runner: + structural_tag_config = { + "type": "structural_tag", + "format": { + "type": "triggered_tags", + "tags": [ + { + "begin": "hello_flag", + "content": {"type": "any_text"}, + "end": "hello", + } + ], + "triggers": ["hello"], + "stop_after_first": False, + }, + } - structural_tag_config = { - "type": "structural_tag", - "format": { - "type": "triggered_tags", - "tags": [ - {"begin": "hello_flag", "content": {"type": "any_text"}, "end": "hello"} - ], - "triggers": ["hello"], - "stop_after_first": False, - }, - } - - sampling_params = SamplingParams( - temperature=0.0, - max_tokens=500, - structured_outputs=StructuredOutputsParams( - structural_tag=json.dumps(structural_tag_config) - ), - ) - - prompt = "Hello and repeat hello 10 times, do not say anything else. Only say hello hello hello, now start" - outputs = llm.generate(prompt, sampling_params=sampling_params, use_tqdm=True) - 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 - assert "hello_flag" in generated_text, ( - f"Expected 'hello_flag' to be in generated text, but got: {generated_text}" + sampling_params = SamplingParams( + temperature=0.0, + max_tokens=500, + structured_outputs=StructuredOutputsParams( + structural_tag=json.dumps(structural_tag_config) + ), ) + + prompt = "Hello and repeat hello 10 times, do not say anything else. Only say hello hello hello, now start" + outputs = runner.llm.generate( + prompt, sampling_params=sampling_params, use_tqdm=True + ) + 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 + assert "hello_flag" in generated_text, ( + f"Expected 'hello_flag' to be in generated text, but got: {generated_text}" + ) diff --git a/tests/entrypoints/multimodal/conftest.py b/tests/entrypoints/multimodal/conftest.py index 8003f1bf7d..f27c042d21 100644 --- a/tests/entrypoints/multimodal/conftest.py +++ b/tests/entrypoints/multimodal/conftest.py @@ -1,11 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable, Iterator -from contextlib import contextmanager -from typing import Any - -import pytest - # Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA) TEST_IMAGE_ASSETS = [ "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg" @@ -13,70 +7,3 @@ TEST_IMAGE_ASSETS = [ "1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png", "RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png", ] - - -def _shutdown_llm(llm: Any, gpu_memory_utilization: float) -> None: - from vllm.distributed import cleanup_dist_env_and_memory - from vllm.platforms import current_platform - - try: - shutdown_timeout = 60.0 if current_platform.is_rocm() else None - llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout) - except Exception: - pass - - del llm - - try: - import torch - - torch._dynamo.reset() - except Exception: - pass - - cleanup_dist_env_and_memory() - - if current_platform.is_rocm(): - from tests.utils import wait_for_rocm_memory_to_settle - - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) - - -@contextmanager -def managed_llm(*args: Any, **kwargs: Any) -> Iterator[Any]: - from vllm import LLM - - llm = LLM(*args, **kwargs) - gpu_memory_utilization = ( - llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization - ) - try: - yield llm - finally: - _shutdown_llm(llm, gpu_memory_utilization) - - -def _make_managed_llm_factory() -> Iterator[Callable[..., Any]]: - from vllm import LLM - - llms: list[tuple[Any, float]] = [] - - def make_llm(*args: Any, **kwargs: Any) -> Any: - llm = LLM(*args, **kwargs) - gpu_memory_utilization = ( - llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization - ) - llms.append((llm, gpu_memory_utilization)) - return llm - - try: - yield make_llm - finally: - while llms: - llm, gpu_memory_utilization = llms.pop() - _shutdown_llm(llm, gpu_memory_utilization) - - -@pytest.fixture -def multimodal_llm_factory() -> Iterator[Callable[..., Any]]: - yield from _make_managed_llm_factory() diff --git a/tests/entrypoints/multimodal/llm/test_chat.py b/tests/entrypoints/multimodal/llm/test_chat.py index 4de1f5cb80..1e19900ce8 100644 --- a/tests/entrypoints/multimodal/llm/test_chat.py +++ b/tests/entrypoints/multimodal/llm/test_chat.py @@ -1,21 +1,26 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import weakref + import pytest from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS @pytest.fixture(scope="function") -def vision_llm(multimodal_llm_factory): - return multimodal_llm_factory( - model="microsoft/Phi-3.5-vision-instruct", +def vision_llm(vllm_runner): + with vllm_runner( + "microsoft/Phi-3.5-vision-instruct", max_model_len=4096, max_num_seqs=5, enforce_eager=True, trust_remote_code=True, limit_mm_per_prompt={"image": 2}, seed=0, - ) + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.parametrize( diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py index 076a381f6c..5c378de35d 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py +++ b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py @@ -69,7 +69,7 @@ def test_inject_into_mm_cache( image_urls, mm_processor_cache_type, caplog_vllm, - multimodal_llm_factory, + vllm_runner, ): """Test that inject_into_mm_cache() injects pre-processed mm_kwargs into the processor cache and MM cache hit metrics are updated correctly. @@ -79,117 +79,115 @@ def test_inject_into_mm_cache( 2. Extract cached kwargs, call inject_into_mm_cache with a new hash, then generate with a pre-rendered input -> verifies injection works """ - llm = multimodal_llm_factory( - model="llava-hf/llava-1.5-7b-hf", + with vllm_runner( + "llava-hf/llava-1.5-7b-hf", max_model_len=4096, max_num_seqs=5, enforce_eager=True, disable_log_stats=False, limit_mm_per_prompt={"image": 2}, mm_processor_cache_type=mm_processor_cache_type, - ) + ) as runner: + # Step 1: Normal requests to populate the cache + runner.llm.chat(_make_messages(image_urls[0])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (1, 0) - # Step 1: Normal requests to populate the cache - llm.chat(_make_messages(image_urls[0])) - assert _get_mm_cache_stats(llm.get_metrics()) == (1, 0) + runner.llm.chat(_make_messages(image_urls[0])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (2, 1) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(50.0) - llm.chat(_make_messages(image_urls[0])) - assert _get_mm_cache_stats(llm.get_metrics()) == (2, 1) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(50.0) + # Step 2: Use a second image to get valid expanded tokens and + # placeholder positions via the renderer. + runner.llm.chat(_make_messages(image_urls[1])) + queries_before = _get_mm_cache_stats(runner.llm.get_metrics())[0] # 3 - # Step 2: Use a second image to get valid expanded tokens and - # placeholder positions via the renderer. - llm.chat(_make_messages(image_urls[1])) - queries_before = _get_mm_cache_stats(llm.get_metrics())[0] # 3 + renderer = runner.llm.llm_engine.renderer + cache = renderer.mm_processor_cache + assert cache is not None, "Processor cache should be enabled" - renderer = llm.llm_engine.renderer - cache = renderer.mm_processor_cache - assert cache is not None, "Processor cache should be enabled" + _, eng_prompts = renderer.render_chat( + [_make_messages(image_urls[1])], + ChatParams(), + ) + eng_input = eng_prompts[0] - _, eng_prompts = renderer.render_chat( - [_make_messages(image_urls[1])], - ChatParams(), - ) - eng_input = eng_prompts[0] + # Inject pre-processed mm_kwargs with a NEW hash via public API + new_mm_hash = "deadbeef" * 8 + mm_hashes = {"image": [new_mm_hash]} + mm_kwargs = eng_input["mm_kwargs"] - # Inject pre-processed mm_kwargs with a NEW hash via public API - new_mm_hash = "deadbeef" * 8 - mm_hashes = {"image": [new_mm_hash]} - mm_kwargs = eng_input["mm_kwargs"] + runner.llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs) - llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs) + # Build pre-rendered input (no externally_processed flag needed) + pre_rendered_input = { + "type": "multimodal", + "prompt_token_ids": eng_input["prompt_token_ids"], + "mm_kwargs": mm_kwargs, + "mm_hashes": mm_hashes, + "mm_placeholders": eng_input["mm_placeholders"], + } - # Build pre-rendered input (no externally_processed flag needed) - pre_rendered_input = { - "type": "multimodal", - "prompt_token_ids": eng_input["prompt_token_ids"], - "mm_kwargs": mm_kwargs, - "mm_hashes": mm_hashes, - "mm_placeholders": eng_input["mm_placeholders"], - } + runner.llm.generate( + pre_rendered_input, + sampling_params=SamplingParams(max_tokens=1), + ) - llm.generate( - pre_rendered_input, - sampling_params=SamplingParams(max_tokens=1), - ) - - # Verify cache was queried and injection happened - queries_after = _get_mm_cache_stats(llm.get_metrics())[0] - assert queries_after > queries_before, ( - "Cache should have been queried for the injected item" - ) - mm_rate = _get_mm_cache_log(llm, caplog_vllm) - assert mm_rate >= 0.0, "MM cache hit rate should be reported" + # Verify cache was queried and injection happened + queries_after = _get_mm_cache_stats(runner.llm.get_metrics())[0] + assert queries_after > queries_before, ( + "Cache should have been queried for the injected item" + ) + mm_rate = _get_mm_cache_log(runner.llm, caplog_vllm) + assert mm_rate >= 0.0, "MM cache hit rate should be reported" @pytest.mark.parametrize("image_urls", [TEST_IMAGE_ASSETS[:1]], indirect=True) def test_inject_into_mm_cache_without_cache( num_gpus_available, image_urls, - multimodal_llm_factory, + vllm_runner, ): """Test that inject_into_mm_cache works gracefully when processor cache is disabled (mm_processor_cache_gb=0). Should not crash. """ - llm = multimodal_llm_factory( - model="llava-hf/llava-1.5-7b-hf", + with vllm_runner( + "llava-hf/llava-1.5-7b-hf", max_model_len=4096, max_num_seqs=5, enforce_eager=True, disable_log_stats=False, limit_mm_per_prompt={"image": 2}, mm_processor_cache_gb=0, - ) + ) as runner: + # Run a normal chat request first to warm up the model. + runner.llm.chat(_make_messages(image_urls[0])) - # Run a normal chat request first to warm up the model. - llm.chat(_make_messages(image_urls[0])) + # Use the renderer to get a proper EngineInput with expanded tokens + renderer = runner.llm.llm_engine.renderer + _, eng_prompts = renderer.render_chat( + [_make_messages(image_urls[0])], + ChatParams(), + ) + eng_input = eng_prompts[0] - # Use the renderer to get a proper EngineInput with expanded tokens - renderer = llm.llm_engine.renderer - _, eng_prompts = renderer.render_chat( - [_make_messages(image_urls[0])], - ChatParams(), - ) - eng_input = eng_prompts[0] + mm_hashes = {"image": ["abcd1234" * 8]} + mm_kwargs = eng_input["mm_kwargs"] - mm_hashes = {"image": ["abcd1234" * 8]} - mm_kwargs = eng_input["mm_kwargs"] + # inject_into_mm_cache should not crash even without cache + runner.llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs) - # inject_into_mm_cache should not crash even without cache - llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs) + # Build and generate with pre-rendered input + pre_rendered_input = { + "type": "multimodal", + "prompt_token_ids": eng_input["prompt_token_ids"], + "mm_kwargs": mm_kwargs, + "mm_hashes": mm_hashes, + "mm_placeholders": eng_input["mm_placeholders"], + } - # Build and generate with pre-rendered input - pre_rendered_input = { - "type": "multimodal", - "prompt_token_ids": eng_input["prompt_token_ids"], - "mm_kwargs": mm_kwargs, - "mm_hashes": mm_hashes, - "mm_placeholders": eng_input["mm_placeholders"], - } - - result = llm.generate( - pre_rendered_input, - sampling_params=SamplingParams(max_tokens=1), - ) - assert len(result) == 1, "Should produce one output" - assert len(result[0].outputs) >= 1, "Should have at least one output sequence" + result = runner.llm.generate( + pre_rendered_input, + sampling_params=SamplingParams(max_tokens=1), + ) + assert len(result) == 1, "Should produce one output" + assert len(result[0].outputs) >= 1, "Should have at least one output sequence" diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py b/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py index dbea37f64e..ac86788f69 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py +++ b/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py @@ -61,38 +61,37 @@ def test_mm_cache_stats( image_urls, mm_processor_cache_type, caplog_vllm, - multimodal_llm_factory, + vllm_runner, ): - llm = multimodal_llm_factory( - model="llava-hf/llava-1.5-7b-hf", + with vllm_runner( + "llava-hf/llava-1.5-7b-hf", max_model_len=4096, max_num_seqs=5, enforce_eager=True, mm_processor_cache_type=mm_processor_cache_type, disable_log_stats=False, limit_mm_per_prompt={"image": 2}, - ) + ) as runner: + runner.llm.chat(_make_messages(image_urls[0])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (1, 0) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0) - llm.chat(_make_messages(image_urls[0])) - assert _get_mm_cache_stats(llm.get_metrics()) == (1, 0) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0) + runner.llm.chat(_make_messages(image_urls[1])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (2, 0) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0) - llm.chat(_make_messages(image_urls[1])) - assert _get_mm_cache_stats(llm.get_metrics()) == (2, 0) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0) + runner.llm.chat(_make_messages(image_urls[0])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (3, 1) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(33.3) - llm.chat(_make_messages(image_urls[0])) - assert _get_mm_cache_stats(llm.get_metrics()) == (3, 1) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(33.3) + # NOTE: This only resets hit rate stats in CachingMetrics + # The raw queries and hits counts remain unaffected + runner.llm.reset_mm_cache() - # NOTE: This only resets hit rate stats in CachingMetrics - # The raw queries and hits counts remain unaffected - llm.reset_mm_cache() + runner.llm.chat(_make_messages(image_urls[0])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (4, 1) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0) - llm.chat(_make_messages(image_urls[0])) - assert _get_mm_cache_stats(llm.get_metrics()) == (4, 1) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0) - - llm.chat(_make_messages(image_urls[1])) - assert _get_mm_cache_stats(llm.get_metrics()) == (5, 1) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0) + runner.llm.chat(_make_messages(image_urls[1])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (5, 1) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0) diff --git a/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py b/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py index 0ea1807124..f63e9463bc 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py +++ b/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import weakref + import pytest -from tests.entrypoints.multimodal.conftest import managed_llm from vllm import LLM, SamplingParams from vllm.assets.image import ImageAsset from vllm.exceptions import VLLMValidationError @@ -14,17 +15,19 @@ TEXT_ONLY_PROMPT = "USER: What is 2 + 2?\nASSISTANT:" @pytest.fixture(scope="module") -def llm(): +def llm(vllm_runner): """LLM with enable_mm_embeds=True and all modality limits zeroed out.""" - with managed_llm( - model=MODEL, + with vllm_runner( + MODEL, max_model_len=2048, enforce_eager=True, gpu_memory_utilization=0.8, enable_mm_embeds=True, limit_mm_per_prompt={"image": 0}, - ) as llm: - yield llm + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/basic/test_encode.py b/tests/entrypoints/pooling/basic/test_encode.py index 61222eb780..b385fc8697 100644 --- a/tests/entrypoints/pooling/basic/test_encode.py +++ b/tests/entrypoints/pooling/basic/test_encode.py @@ -6,7 +6,6 @@ import weakref import pytest from vllm import LLM, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory from vllm.exceptions import VLLMValidationError MODEL_NAME = "intfloat/multilingual-e5-small" @@ -29,23 +28,20 @@ TOKEN_IDS = [ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/basic/test_tiling_engine.py b/tests/entrypoints/pooling/basic/test_tiling_engine.py index 329f91debb..88c9b434b3 100644 --- a/tests/entrypoints/pooling/basic/test_tiling_engine.py +++ b/tests/entrypoints/pooling/basic/test_tiling_engine.py @@ -6,28 +6,26 @@ from unittest import mock import pytest -from vllm import LLM, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory +from vllm import PoolingParams MODEL_NAME = "intfloat/multilingual-e5-small" @pytest.fixture(scope="module") -def llm(): - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_seqs=2, # small to trigger tiling tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/classify/test_offline.py b/tests/entrypoints/pooling/classify/test_offline.py index 2f6a9f1db4..9db7b48642 100644 --- a/tests/entrypoints/pooling/classify/test_offline.py +++ b/tests/entrypoints/pooling/classify/test_offline.py @@ -7,7 +7,6 @@ import torch from tests.models.utils import softmax from vllm import LLM, ClassificationRequestOutput, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory from vllm.tasks import PoolingTask MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach" @@ -18,23 +17,20 @@ num_labels = 2 @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/embed/test_offline.py b/tests/entrypoints/pooling/embed/test_offline.py index 6dd2955ead..641f8b714e 100644 --- a/tests/entrypoints/pooling/embed/test_offline.py +++ b/tests/entrypoints/pooling/embed/test_offline.py @@ -7,7 +7,6 @@ import torch import torch.nn.functional as F from vllm import LLM, EmbeddingRequestOutput, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory from vllm.tasks import PoolingTask MODEL_NAME = "intfloat/multilingual-e5-small" @@ -18,23 +17,21 @@ embedding_size = 384 @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - assert embedding_size == llm.model_config.embedding_size - - yield weakref.proxy(llm) - - del llm - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + assert embedding_size == runner.llm.model_config.embedding_size + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/reward/test_token_reward_offline.py b/tests/entrypoints/pooling/reward/test_token_reward_offline.py index 50a4b54682..135861e700 100644 --- a/tests/entrypoints/pooling/reward/test_token_reward_offline.py +++ b/tests/entrypoints/pooling/reward/test_token_reward_offline.py @@ -8,7 +8,6 @@ import torch from tests.models.utils import softmax from vllm import LLM, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory MODEL_NAME = "internlm/internlm2-1_8b-reward" @@ -16,24 +15,21 @@ prompts = ["The chef prepared a delicious meal."] @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, trust_remote_code=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/scoring/test_bi_encoder_offline.py b/tests/entrypoints/pooling/scoring/test_bi_encoder_offline.py index e68fb61a99..5f0ba9c60e 100644 --- a/tests/entrypoints/pooling/scoring/test_bi_encoder_offline.py +++ b/tests/entrypoints/pooling/scoring/test_bi_encoder_offline.py @@ -6,8 +6,6 @@ import weakref import pytest from tests.entrypoints.pooling.scoring.util import EncoderScoringHfRunner -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory MODEL_NAME = "intfloat/multilingual-e5-small" PROMPT = "The chef prepared a delicious meal." @@ -27,23 +25,20 @@ DTYPE = "half" @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.fixture(scope="module") diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py index 3ad044987a..c9bc291e58 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py @@ -8,7 +8,6 @@ import torch from tests.models.utils import softmax from vllm import LLM, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls" PROMPT = "The chef prepared a delicious meal." @@ -24,23 +23,20 @@ TEXTS_2 = [ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.fixture(scope="module") diff --git a/tests/entrypoints/pooling/scoring/test_late_interaction_offline.py b/tests/entrypoints/pooling/scoring/test_late_interaction_offline.py index 9a4db3d3f8..b39aa7abb9 100644 --- a/tests/entrypoints/pooling/scoring/test_late_interaction_offline.py +++ b/tests/entrypoints/pooling/scoring/test_late_interaction_offline.py @@ -5,9 +5,6 @@ import weakref import pytest -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory - from .util import ColBERTScoringHfRunner MODEL_NAME = "answerdotai/answerai-colbert-small-v1" @@ -30,23 +27,20 @@ DTYPE = "half" @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.fixture(scope="module") diff --git a/tests/entrypoints/pooling/scoring/test_late_interaction_offline_vision.py b/tests/entrypoints/pooling/scoring/test_late_interaction_offline_vision.py index 4812bd2ab3..ca7843f8cf 100644 --- a/tests/entrypoints/pooling/scoring/test_late_interaction_offline_vision.py +++ b/tests/entrypoints/pooling/scoring/test_late_interaction_offline_vision.py @@ -5,32 +5,26 @@ import weakref import pytest -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory - from .util import make_base64_image, make_image_mm_param MODEL_NAME = "vidore/colpali-v1.3-hf" @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/token_classify/test_offline.py b/tests/entrypoints/pooling/token_classify/test_offline.py index 3f59b37517..9715cd63b1 100644 --- a/tests/entrypoints/pooling/token_classify/test_offline.py +++ b/tests/entrypoints/pooling/token_classify/test_offline.py @@ -6,7 +6,6 @@ import pytest from vllm import LLM, PoolingRequestOutput from vllm.config import PoolerConfig -from vllm.distributed import cleanup_dist_env_and_memory from vllm.tasks import PoolingTask MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach" @@ -17,24 +16,21 @@ num_labels = 2 @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, pooler_config=PoolerConfig(task="token_classify"), max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/token_embed/test_offline.py b/tests/entrypoints/pooling/token_embed/test_offline.py index 7bc7be4866..e0bda65422 100644 --- a/tests/entrypoints/pooling/token_embed/test_offline.py +++ b/tests/entrypoints/pooling/token_embed/test_offline.py @@ -6,7 +6,6 @@ import pytest from vllm import LLM, PoolingRequestOutput from vllm.config import PoolerConfig -from vllm.distributed import cleanup_dist_env_and_memory from vllm.tasks import PoolingTask MODEL_NAME = "intfloat/multilingual-e5-small" @@ -17,24 +16,22 @@ embedding_size = 384 @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, pooler_config=PoolerConfig(task="token_embed"), max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - assert embedding_size == llm.model_config.embedding_size - - yield weakref.proxy(llm) - - del llm - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + assert embedding_size == runner.llm.model_config.embedding_size + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py index 68119a5683..e5c3616d6e 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py @@ -8,13 +8,13 @@ actual NCCL communication. """ import os +import weakref from dataclasses import dataclass from unittest.mock import patch import pytest import torch -from vllm import LLM from vllm.config import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( WeightTransferEngine, @@ -102,25 +102,24 @@ def mock_create_engine(config, vllm_config, device, model): @create_new_process_for_each_test() -def test_get_world_size_tp1(): +def test_get_world_size_tp1(vllm_runner): """Test world_size is correctly configured for TP=1.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") - llm = LLM( - model=MODEL_NAME, + with vllm_runner( + MODEL_NAME, enforce_eager=True, load_format="dummy", tensor_parallel_size=1, weight_transfer_config=WeightTransferConfig(backend="nccl"), - ) - - world_size = llm.llm_engine.vllm_config.parallel_config.world_size - assert world_size == 1 + ) as runner: + world_size = runner.llm.llm_engine.vllm_config.parallel_config.world_size + assert world_size == 1 @create_new_process_for_each_test() -def test_init_weight_transfer_engine_calls_engine(): +def test_init_weight_transfer_engine_calls_engine(vllm_runner): """Test that init_weight_transfer_engine calls the engine's init_transfer_engine method.""" if torch.accelerator.device_count() < 1: @@ -131,17 +130,20 @@ def test_init_weight_transfer_engine_calls_engine(): # Enable insecure serialization to allow pickling functions for collective_rpc os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" - with patch( - "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", - mock_create_engine, - ): - llm = LLM( - model=MODEL_NAME, + with ( + patch( + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", + mock_create_engine, + ), + vllm_runner( + MODEL_NAME, enforce_eager=True, load_format="dummy", tensor_parallel_size=1, weight_transfer_config=WeightTransferConfig(backend="nccl"), - ) + ) as runner, + ): + llm = weakref.proxy(runner.llm) # Verify engine was created def check_engine_exists(self): @@ -170,7 +172,7 @@ def test_init_weight_transfer_engine_calls_engine(): @create_new_process_for_each_test() -def test_update_weights_calls_engine(): +def test_update_weights_calls_engine(vllm_runner): """Test that update_weights calls the engine's receive_weights method.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") @@ -180,17 +182,20 @@ def test_update_weights_calls_engine(): # Enable insecure serialization to allow pickling functions for collective_rpc os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" - with patch( - "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", - mock_create_engine, - ): - llm = LLM( - model=MODEL_NAME, + with ( + patch( + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", + mock_create_engine, + ), + vllm_runner( + MODEL_NAME, enforce_eager=True, load_format="dummy", tensor_parallel_size=1, weight_transfer_config=WeightTransferConfig(backend="nccl"), - ) + ) as runner, + ): + llm = weakref.proxy(runner.llm) # First init the weight transfer llm.init_weight_transfer_engine( @@ -233,7 +238,7 @@ def test_update_weights_calls_engine(): @create_new_process_for_each_test() -def test_full_weight_transfer_flow(): +def test_full_weight_transfer_flow(vllm_runner): """Test the complete weight transfer flow: init -> start -> update -> finish.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") @@ -243,17 +248,20 @@ def test_full_weight_transfer_flow(): # Enable insecure serialization to allow pickling functions for collective_rpc os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" - with patch( - "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", - mock_create_engine, - ): - llm = LLM( - model=MODEL_NAME, + with ( + patch( + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", + mock_create_engine, + ), + vllm_runner( + MODEL_NAME, enforce_eager=True, load_format="dummy", tensor_parallel_size=1, weight_transfer_config=WeightTransferConfig(backend="nccl"), - ) + ) as runner, + ): + llm = weakref.proxy(runner.llm) assert llm.get_weight_version() == "default" @@ -309,20 +317,19 @@ def test_full_weight_transfer_flow(): @create_new_process_for_each_test() -def test_weight_transfer_config_backend(): +def test_weight_transfer_config_backend(vllm_runner): """Test that WeightTransferConfig backend is properly configured.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") # Test with nccl backend - llm = LLM( - model=MODEL_NAME, + with vllm_runner( + MODEL_NAME, enforce_eager=True, load_format="dummy", tensor_parallel_size=1, weight_transfer_config=WeightTransferConfig(backend="nccl"), - ) - - config = llm.llm_engine.vllm_config.weight_transfer_config - assert config is not None - assert config.backend == "nccl" + ) as runner: + config = runner.llm.llm_engine.vllm_config.weight_transfer_config + assert config is not None + assert config.backend == "nccl"