@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b96bcd0b47

[ROCm][CI] Solidify entrypoint LLM lifecycle (#51280)

Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed Aug 13, 2026 at 12:24 UTC b96bcd0b474fe469c42423f6a4e7de8ba16baa01

24 files changed +974 -1075

tests/entrypoints/llm/offline_mode/test_offline_mode.py

+16 -10

@@ -9,9 +9,6 @@ import pytest

9 import regex as re

10 import urllib3

12 -from vllm import LLM

13 -from vllm.distributed import cleanup_dist_env_and_memory

14 -

12 MODEL_CONFIGS = [

13 {

14 "model": "facebook/opt-125m",

@@ -55,19 +52,26 @@ MODEL_CONFIGS = [

52 ]

55 +def _create_runner(vllm_runner, model_config):

56 + runner_config = model_config.copy()

57 + model = runner_config.pop("model")

58 + tokenizer_name = runner_config.pop("tokenizer", None)

59 + return vllm_runner(model, tokenizer_name=tokenizer_name, **runner_config)

60 +

61 +

62 @pytest.fixture(scope="module")

59 -def cache_models():

63 +def cache_models(vllm_runner):

64 # Cache model files first

65 for model_config in MODEL_CONFIGS:

62 - LLM(**model_config)

63 - cleanup_dist_env_and_memory()

66 + with _create_runner(vllm_runner, model_config):

67 + pass

69 yield

72 @pytest.mark.skip_global_cleanup

73 @pytest.mark.usefixtures("cache_models")

70 -def test_offline_mode(monkeypatch: pytest.MonkeyPatch):

74 +def test_offline_mode(monkeypatch: pytest.MonkeyPatch, vllm_runner):

75 # Set HF to offline mode and ensure we can still construct an LLM

76 with monkeypatch.context() as m:

77 try:

@@ -93,7 +97,8 @@ def test_offline_mode(monkeypatch: pytest.MonkeyPatch):

97 _re_import_modules()

98 # Cached model files should be used in offline mode

99 for model_config in MODEL_CONFIGS:

96 - LLM(**model_config)

100 + with _create_runner(vllm_runner, model_config):

101 + pass

102 finally:

103 # Reset the environment after the test

104 # NB: Assuming tests are run in online mode

@@ -136,7 +141,7 @@ def _re_import_modules():

141

142 @pytest.mark.skip_global_cleanup

143 @pytest.mark.usefixtures("cache_models")

139 -def test_model_from_huggingface_offline(monkeypatch: pytest.MonkeyPatch):

144 +def test_model_from_huggingface_offline(monkeypatch: pytest.MonkeyPatch, vllm_runner):

145 # Set HF to offline mode and ensure we can still construct an LLM

146 with monkeypatch.context() as m:

147 try:

@@ -159,7 +164,8 @@ def test_model_from_huggingface_offline(monkeypatch: pytest.MonkeyPatch):

164 # Need to re-import huggingface_hub

165 # and friends to set up offline mode

166 _re_import_modules()

162 - LLM(model="facebook/opt-125m")

167 + with vllm_runner("facebook/opt-125m"):

168 + pass

169 finally:

170 # Reset the environment after the test

171 # NB: Assuming tests are run in online mode

tests/entrypoints/llm/test_chat.py

+21 -36

@@ -4,46 +4,36 @@ import weakref

5 import pytest

7 -from vllm import LLM

8 -from vllm.distributed import cleanup_dist_env_and_memory

7 from vllm.exceptions import VLLMValidationError

8 from vllm.sampling_params import SamplingParams

11 @pytest.fixture(scope="function")

14 -def text_llm():

15 - # pytest caches the fixture so we use weakref.proxy to

16 - # enable garbage collection

17 - llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True, seed=0)

18 -

19 - yield weakref.proxy(llm)

20 -

21 - del llm

22 -

23 - cleanup_dist_env_and_memory()

12 +def text_llm(vllm_runner):

13 + with vllm_runner(

14 + "meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True, seed=0

15 + ) as runner:

16 + # pytest caches yielded fixtures until after teardown, so use a proxy to

17 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

18 + yield weakref.proxy(runner.llm)

21 @pytest.fixture(scope="function")

27 -def llm_for_failure_test():

22 +def llm_for_failure_test(vllm_runner):

23 """

24 Fixture for testing issue #26081.

25 Uses a small max_model_len to easily trigger length errors.

26 """

32 - # pytest caches the fixture so we use weakref.proxy to

33 - # enable garbage collection

34 - llm = LLM(

35 - model="meta-llama/Llama-3.2-1B-Instruct",

27 + with vllm_runner(

28 + "meta-llama/Llama-3.2-1B-Instruct",

29 enforce_eager=True,

30 seed=0,

31 max_model_len=128,

32 disable_log_stats=True,

40 - )

41 -

42 - yield weakref.proxy(llm)

43 -

44 - del llm

45 -

46 - cleanup_dist_env_and_memory()

33 + ) as runner:

34 + # pytest caches yielded fixtures until after teardown, so use a proxy to

35 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

36 + yield weakref.proxy(runner.llm)

39 def test_chat(text_llm):

@@ -99,21 +89,16 @@ def test_llm_chat_tokenization_no_double_bos(text_llm):

91 @pytest.fixture(scope="function")

102 -def thinking_llm():

103 - # pytest caches the fixture so we use weakref.proxy to

104 - # enable garbage collection

105 - llm = LLM(

106 - model="Qwen/Qwen3-0.6B",

92 +def thinking_llm(vllm_runner):

93 + with vllm_runner(

94 + "Qwen/Qwen3-0.6B",

95 max_model_len=4096,

96 enforce_eager=True,

97 seed=0,

110 - )

111 -

112 - yield weakref.proxy(llm)

113 -

114 - del llm

115 -

116 - cleanup_dist_env_and_memory()

98 + ) as runner:

99 + # pytest caches yielded fixtures until after teardown, so use a proxy to

100 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

101 + yield weakref.proxy(runner.llm)

102

103

104 @pytest.mark.parametrize("enable_thinking", [True, False])

tests/entrypoints/llm/test_collective_rpc.py

+5 -7

@@ -4,15 +4,13 @@

4 import pytest

5 import torch

7 -from vllm import LLM

8 -

7 from ...utils import create_new_process_for_each_test

10 @pytest.mark.parametrize("tp_size", [1, 2])

11 @pytest.mark.parametrize("backend", ["mp", "ray"])

12 @create_new_process_for_each_test()

15 -def test_collective_rpc(tp_size, backend, monkeypatch):

13 +def test_collective_rpc(tp_size, backend, monkeypatch, vllm_runner):

14 if torch.accelerator.device_count() < tp_size:

15 pytest.skip(f"Not enough GPUs for tensor parallelism {tp_size}")

16 if tp_size == 1 and backend == "ray":

@@ -26,11 +24,11 @@ def test_collective_rpc(tp_size, backend, monkeypatch):

24 return self.rank

26 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

29 - llm = LLM(

30 - model="hmellor/tiny-random-LlamaForCausalLM",

27 + with vllm_runner(

28 + "hmellor/tiny-random-LlamaForCausalLM",

29 enforce_eager=True,

30 load_format="dummy",

31 tensor_parallel_size=tp_size,

32 distributed_executor_backend=backend,

35 - )

36 - assert llm.collective_rpc(echo_rank) == list(range(tp_size))

33 + ) as runner:

34 + assert runner.llm.collective_rpc(echo_rank) == list(range(tp_size))

tests/entrypoints/llm/test_generate.py

+30 -36

@@ -6,7 +6,6 @@ import weakref

6 import pytest

8 from vllm import LLM, SamplingParams

9 -from vllm.distributed import cleanup_dist_env_and_memory

10 MODEL_NAME = "distilbert/distilgpt2"

@@ -26,22 +25,17 @@ TOKEN_IDS = [

27 @pytest.fixture(scope="module")

29 -def llm():

30 - # pytest caches the fixture so we use weakref.proxy to

31 - # enable garbage collection

32 - llm = LLM(

33 - model=MODEL_NAME,

28 +def llm(vllm_runner):

29 + with vllm_runner(

30 + MODEL_NAME,

31 max_num_batched_tokens=4096,

32 tensor_parallel_size=1,

33 gpu_memory_utilization=0.10,

34 enforce_eager=True,

38 - )

39 -

40 - yield weakref.proxy(llm)

41 -

42 - del llm

43 -

44 - cleanup_dist_env_and_memory()

35 + ) as runner:

36 + # pytest caches yielded fixtures until after teardown, so use a proxy to

37 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

38 + yield weakref.proxy(runner.llm)

41 @pytest.mark.skip_global_cleanup

@@ -97,34 +91,34 @@ def test_single_prompt_priority(llm: LLM):

91 assert len(outputs) == 1

100 -def test_max_model_len():

94 +def test_max_model_len(vllm_runner):

95 max_model_len = 20

102 - llm = LLM(

103 - model=MODEL_NAME,

96 + with vllm_runner(

97 + MODEL_NAME,

98 max_model_len=max_model_len,

99 gpu_memory_utilization=0.10,

100 enforce_eager=True, # reduce test time

107 - )

108 - sampling_params = SamplingParams(max_tokens=max_model_len + 10)

109 - outputs = llm.generate(PROMPTS, sampling_params)

110 - for output in outputs:

111 - num_total_tokens = len(output.prompt_token_ids) + len(

112 - output.outputs[0].token_ids

113 - )

114 - # Total tokens must not exceed max_model_len.

115 - # It can be less if generation finishes due to other reasons (e.g., EOS)

116 - # before reaching the absolute model length limit.

117 - assert num_total_tokens <= max_model_len

118 -

119 -

120 -def test_log_stats():

121 - llm = LLM(

122 - model=MODEL_NAME,

101 + ) as runner:

102 + sampling_params = SamplingParams(max_tokens=max_model_len + 10)

103 + outputs = runner.llm.generate(PROMPTS, sampling_params)

104 + for output in outputs:

105 + num_total_tokens = len(output.prompt_token_ids) + len(

106 + output.outputs[0].token_ids

107 + )

108 + # Total tokens must not exceed max_model_len.

109 + # It can be less if generation finishes due to other reasons (e.g., EOS)

110 + # before reaching the absolute model length limit.

111 + assert num_total_tokens <= max_model_len

112 +

113 +

114 +def test_log_stats(vllm_runner):

115 + with vllm_runner(

116 + MODEL_NAME,

117 disable_log_stats=False,

118 gpu_memory_utilization=0.10,

119 enforce_eager=True, # reduce test time

126 - )

127 - outputs = llm.generate(PROMPTS, sampling_params=None)

120 + ) as runner:

121 + outputs = runner.llm.generate(PROMPTS, sampling_params=None)

122

129 - # disable_log_stats is False, every output should have metrics

130 - assert all(output.metrics is not None for output in outputs)

123 + # disable_log_stats is False, every output should have metrics

124 + assert all(output.metrics is not None for output in outputs)

tests/entrypoints/llm/test_gpu_utilization.py

+25 -12

@@ -1,10 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from vllm import LLM, SamplingParams

4 +from vllm import SamplingParams

7 -def test_gpu_memory_utilization():

7 +def test_gpu_memory_utilization(vllm_runner):

8 prompts = [

9 "Hello, my name is",

10 "The president of the United States is",

@@ -15,13 +15,26 @@ def test_gpu_memory_utilization():

16 # makes sure gpu_memory_utilization is per-instance limit,

17 # not a global limit

18 - llms = [

19 - LLM(model="facebook/opt-125m", gpu_memory_utilization=0.3, enforce_eager=True)

20 - for i in range(3)

21 - ]

22 - for llm in llms:

23 - outputs = llm.generate(prompts, sampling_params)

24 - for output in outputs:

25 - prompt = output.prompt

26 - generated_text = output.outputs[0].text

27 - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

18 + with (

19 + vllm_runner(

20 + "facebook/opt-125m",

21 + gpu_memory_utilization=0.3,

22 + enforce_eager=True,

23 + ) as runner_0,

24 + vllm_runner(

25 + "facebook/opt-125m",

26 + gpu_memory_utilization=0.3,

27 + enforce_eager=True,

28 + ) as runner_1,

29 + vllm_runner(

30 + "facebook/opt-125m",

31 + gpu_memory_utilization=0.3,

32 + enforce_eager=True,

33 + ) as runner_2,

34 + ):

35 + for runner in (runner_0, runner_1, runner_2):

36 + outputs = runner.llm.generate(prompts, sampling_params)

37 + for output in outputs:

38 + prompt = output.prompt

39 + generated_text = output.outputs[0].text

40 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

tests/entrypoints/llm/test_prompt_validation.py

+26 -21

@@ -4,30 +4,35 @@

4 import pytest

5 import torch

7 -from vllm import LLM

7 from vllm.exceptions import VLLMValidationError

11 -def test_empty_prompt():

12 - llm = LLM(model="openai-community/gpt2", enforce_eager=True)

13 - with pytest.raises(VLLMValidationError, match="decoder prompt cannot be empty"):

14 - llm.generate([""])

15 -

16 -

17 -def test_out_of_vocab_token():

18 - llm = LLM(model="openai-community/gpt2", enforce_eager=True)

19 - with pytest.raises(VLLMValidationError, match="out of vocabulary"):

20 - llm.generate({"prompt_token_ids": [999999]})

21 -

22 -

23 -def test_require_mm_embeds():

24 - llm = LLM(

25 - model="llava-hf/llava-1.5-7b-hf",

26 - enforce_eager=True,

27 - enable_mm_embeds=False,

28 - )

29 - with pytest.raises(ValueError, match="--enable-mm-embeds"):

30 - llm.generate(

10 +def test_empty_prompt(vllm_runner):

11 + with (

12 + vllm_runner("openai-community/gpt2", enforce_eager=True) as runner,

13 + pytest.raises(VLLMValidationError, match="decoder prompt cannot be empty"),

14 + ):

15 + runner.llm.generate([""])

16 +

17 +

18 +def test_out_of_vocab_token(vllm_runner):

19 + with (

20 + vllm_runner("openai-community/gpt2", enforce_eager=True) as runner,

21 + pytest.raises(VLLMValidationError, match="out of vocabulary"),

22 + ):

23 + runner.llm.generate({"prompt_token_ids": [999999]})

24 +

25 +

26 +def test_require_mm_embeds(vllm_runner):

27 + with (

28 + vllm_runner(

29 + "llava-hf/llava-1.5-7b-hf",

30 + enforce_eager=True,

31 + enable_mm_embeds=False,

32 + ) as runner,

33 + pytest.raises(ValueError, match="--enable-mm-embeds"),

34 + ):

35 + runner.llm.generate(

36 {

37 "prompt": "<image>",

38 "multi_modal_data": {"image": torch.empty(1, 1, 1)},

tests/entrypoints/llm/test_struct_output_generate.py

+566 -562

@@ -9,13 +9,10 @@ from typing import Any

9 import jsonschema

10 import pytest

11 import regex as re

12 -import torch

12 from pydantic import BaseModel

14 from tests.reasoning.utils import run_reasoning_extraction

15 from vllm.config import StructuredOutputsConfig

17 -from vllm.distributed import cleanup_dist_env_and_memory

18 -from vllm.entrypoints.llm import LLM

16 from vllm.outputs import RequestOutput

17 from vllm.platforms import current_platform

18 from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager

@@ -212,11 +209,11 @@ class CarDescription(BaseModel):

209 PARAMS_MODELS_BACKENDS_TOKENIZER_MODE,

210 )

211 def test_structured_output(

215 - request: pytest.FixtureRequest,

212 backend: str,

213 tokenizer_mode: str,

214 model_name: str,

215 speculative_config: dict[str, Any],

216 + vllm_runner,

217 ):

218 sample_json_schema = SAMPLE_JSON_SCHEMA

219 unsupported_json_schema = UNSUPPORTED_JSON_SCHEMA

@@ -229,8 +226,8 @@ def test_structured_output(

226

227 # Use a single LLM instance for several scenarios to

228 # speed up the test suite.

232 - llm = LLM(

233 - model=model_name,

229 + with vllm_runner(

230 + model_name,

231 enforce_eager=True,

232 max_model_len=1024,

233 structured_outputs_config=dict(

@@ -242,181 +239,305 @@ def test_structured_output(

239 config_format="auto" if not model_name.startswith("mistralai/") else "hf",

240 speculative_config=speculative_config,

241 **platform_args,

245 - )

246 - request.addfinalizer(llm.llm_engine.engine_core.shutdown)

247 -

248 - #

249 - # Test 1: Generate JSON output based on a provided schema

250 - #

251 - sampling_params = SamplingParams(

252 - temperature=1.0,

253 - max_tokens=4096,

254 - structured_outputs=StructuredOutputsParams(json=sample_json_schema),

255 - )

256 -

257 - prompt = (

258 - "Give an example JSON for an employee profile that fits this "

259 - "schema. Make the response as short as possible. Schema: "

260 - f"{sample_json_schema}"

261 - )

262 - outputs = llm.generate(

263 - [prompt] * 2,

264 - sampling_params=sampling_params,

265 - use_tqdm=True,

266 - )

267 -

268 - assert outputs is not None

269 -

270 - for output in outputs:

271 - assert output is not None

272 - assert isinstance(output, RequestOutput)

273 - prompt = output.prompt

274 -

275 - generated_text = output.outputs[0].text

276 - assert generated_text is not None

277 - if backend != "lm-format-enforcer":

278 - assert "\n" not in generated_text

279 - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

280 - try:

281 - output_json = json.loads(generated_text)

282 - except json.JSONDecodeError as e:

283 - pytest.fail(

284 - f"Invalid JSON from backend={backend}: {generated_text!r}\n"

285 - f"Schema: {sample_json_schema}\nError: {e}"

286 - )

287 - jsonschema.validate(instance=output_json, schema=sample_json_schema)

288 -

289 - #

290 - # Test 2: Generate JSON object without a schema

291 - #

292 - if backend != "outlines":

242 + ) as runner:

243 + #

244 + # Test 1: Generate JSON output based on a provided schema

245 + #

246 sampling_params = SamplingParams(

247 temperature=1.0,

248 max_tokens=4096,

296 - n=2,

297 - structured_outputs=StructuredOutputsParams(json_object=True),

249 + structured_outputs=StructuredOutputsParams(json=sample_json_schema),

250 )

251

300 - outputs = llm.generate(

301 - prompts=(

302 - "Generate a JSON object with curly braces for a person with "

303 - "name and age fields for John Smith who is 31 years old. "

304 - "Make the response as short as possible."

305 - ),

252 + prompt = (

253 + "Give an example JSON for an employee profile that fits this "

254 + "schema. Make the response as short as possible. Schema: "

255 + f"{sample_json_schema}"

256 + )

257 + outputs = runner.llm.generate(

258 + [prompt] * 2,

259 sampling_params=sampling_params,

260 use_tqdm=True,

261 )

262

263 assert outputs is not None

264 +

265 for output in outputs:

266 assert output is not None

267 assert isinstance(output, RequestOutput)

268 + prompt = output.prompt

269

315 - for i in range(2):

316 - generated_text = output.outputs[i].text

317 - print(generated_text)

318 - assert generated_text is not None

270 + generated_text = output.outputs[0].text

271 + assert generated_text is not None

272 + if backend != "lm-format-enforcer":

273 + assert "\n" not in generated_text

274 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

275 + try:

276 + output_json = json.loads(generated_text)

277 + except json.JSONDecodeError as e:

278 + pytest.fail(

279 + f"Invalid JSON from backend={backend}: {generated_text!r}\n"

280 + f"Schema: {sample_json_schema}\nError: {e}"

281 + )

282 + jsonschema.validate(instance=output_json, schema=sample_json_schema)

283

320 - # Parse to verify it is a valid JSON object

321 - parsed_json = json.loads(generated_text)

322 - assert isinstance(parsed_json, dict)

284 + #

285 + # Test 2: Generate JSON object without a schema

286 + #

287 + if backend != "outlines":

288 + sampling_params = SamplingParams(

289 + temperature=1.0,

290 + max_tokens=4096,

291 + n=2,

292 + structured_outputs=StructuredOutputsParams(json_object=True),

293 + )

294 +

295 + outputs = runner.llm.generate(

296 + prompts=(

297 + "Generate a JSON object with curly braces for a person with "

298 + "name and age fields for John Smith who is 31 years old. "

299 + "Make the response as short as possible."

300 + ),

301 + sampling_params=sampling_params,

302 + use_tqdm=True,

303 + )

304 +

305 + assert outputs is not None

306 + for output in outputs:

307 + assert output is not None

308 + assert isinstance(output, RequestOutput)

309 +

310 + for i in range(2):

311 + generated_text = output.outputs[i].text

312 + print(generated_text)

313 + assert generated_text is not None

314

324 - #

325 - # Test 3: test a jsonschema incompatible with xgrammar

326 - #

327 - sampling_params = SamplingParams(

328 - temperature=1.0,

329 - max_tokens=4096,

330 - structured_outputs=StructuredOutputsParams(json=unsupported_json_schema),

331 - )

332 - if backend.startswith("xgrammar"):

333 - with pytest.raises(

334 - ValueError,

335 - match="The provided JSON schema contains features "

336 - "not supported by xgrammar.",

337 - ):

315 + # Parse to verify it is a valid JSON object

316 + parsed_json = json.loads(generated_text)

317 + assert isinstance(parsed_json, dict)

318 +

319 + #

320 + # Test 3: test a jsonschema incompatible with xgrammar

321 + #

322 + sampling_params = SamplingParams(

323 + temperature=1.0,

324 + max_tokens=4096,

325 + structured_outputs=StructuredOutputsParams(json=unsupported_json_schema),

326 + )

327 + if backend.startswith("xgrammar"):

328 + with pytest.raises(

329 + ValueError,

330 + match="The provided JSON schema contains features "

331 + "not supported by xgrammar.",

332 + ):

333 + prompt = (

334 + f"Give an example JSON for an employee profile that "

335 + f"fits this schema: {unsupported_json_schema}. "

336 + f"Make the response as short as possible."

337 + )

338 + runner.llm.generate(

339 + [prompt] * 2,

340 + sampling_params=sampling_params,

341 + use_tqdm=True,

342 + )

343 + else:

344 prompt = (

339 - f"Give an example JSON for an employee profile that "

345 + f"Give an example JSON object for a grade that "

346 f"fits this schema: {unsupported_json_schema}. "

347 f"Make the response as short as possible."

348 )

343 - llm.generate(

344 - [prompt] * 2,

349 + outputs = runner.llm.generate(

350 + prompt,

351 + sampling_params=sampling_params,

352 + use_tqdm=True,

353 + )

354 + assert outputs is not None

355 + for output in outputs:

356 + assert output is not None

357 + assert isinstance(output, RequestOutput)

358 + generated_text = output.outputs[0].text

359 + assert generated_text is not None

360 + print(generated_text)

361 +

362 + # Parse to verify it is valid JSON

363 + parsed_json = json.loads(generated_text)

364 + assert isinstance(parsed_json, dict)

365 +

366 + if backend not in ["outlines", "lm-format-enforcer"]:

367 + #

368 + # Test 4: Generate SQL statement using EBNF grammar

369 + #

370 + sampling_params = SamplingParams(

371 + temperature=0.8,

372 + top_p=0.95,

373 + max_tokens=1000,

374 + structured_outputs=StructuredOutputsParams(grammar=sample_sql_ebnf),

375 + )

376 + outputs = runner.llm.generate(

377 + (

378 + "Generate a sql statement that selects col_1 from "

379 + "table_1 where it is equal to 1. Make the response as short as "

380 + "possible."

381 + ),

382 + sampling_params=sampling_params,

383 + use_tqdm=True,

384 + )

385 +

386 + assert outputs is not None

387 + for output in outputs:

388 + assert output is not None

389 + assert isinstance(output, RequestOutput)

390 + prompt = output.prompt

391 +

392 + generated_text = output.outputs[0].text

393 + assert generated_text is not None

394 +

395 + # remove spaces for comparison b/c we removed them in the grammar

396 + ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(

397 + " ", ""

398 + )

399 +

400 + assert generated_text.strip() == ground_truth

401 +

402 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

403 +

404 + #

405 + # Test 5: Generate SQL statement using Lark grammar

406 + #

407 + sampling_params = SamplingParams(

408 + temperature=0.8,

409 + top_p=0.95,

410 + max_tokens=1000,

411 + structured_outputs=StructuredOutputsParams(grammar=sample_sql_lark),

412 + )

413 + outputs = runner.llm.generate(

414 + (

415 + "Generate a sql statement that selects col_1 from "

416 + "table_1 where it is equal to 1. Make the response as short as "

417 + "possible."

418 + ),

419 sampling_params=sampling_params,

420 use_tqdm=True,

421 )

348 - else:

422 +

423 + assert outputs is not None

424 + for output in outputs:

425 + assert output is not None

426 + assert isinstance(output, RequestOutput)

427 + prompt = output.prompt

428 +

429 + generated_text = output.outputs[0].text

430 + assert generated_text is not None

431 +

432 + # use Lark to parse the output, and make sure it's a valid parse tree

433 + from lark import Lark

434 +

435 + parser = Lark(sample_sql_lark)

436 + parser.parse(generated_text)

437 +

438 + # remove spaces for comparison b/c we removed them in the grammar

439 + ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(

440 + " ", ""

441 + )

442 +

443 + assert generated_text.strip() == ground_truth

444 +

445 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

446 +

447 + #

448 + # Test 6: Test invalid grammar input

449 + #

450 + sampling_params = SamplingParams(

451 + temperature=0.8,

452 + top_p=0.95,

453 + max_tokens=1000,

454 + structured_outputs=StructuredOutputsParams(grammar="not a grammar"),

455 + )

456 + with pytest.raises(ValueError, match="Failed to convert the grammar "):

457 + runner.llm.generate(

458 + (

459 + "Generate a sql statement that selects col_1 from "

460 + "table_1 where it is equal to 1. Make the response as short "

461 + "as possible."

462 + ),

463 + sampling_params=sampling_params,

464 + use_tqdm=True,

465 + )

466 +

467 + #

468 + # Test 7: Generate text based on a regex pattern

469 + #

470 + sampling_params = SamplingParams(

471 + temperature=0.8,

472 + top_p=0.95,

473 + structured_outputs=StructuredOutputsParams(regex=sample_regex),

474 + )

475 +

476 prompt = (

350 - f"Give an example JSON object for a grade that "

351 - f"fits this schema: {unsupported_json_schema}. "

477 + f"Give an example IPv4 address with this regex: {sample_regex}. "

478 f"Make the response as short as possible."

479 )

354 - outputs = llm.generate(

355 - prompt,

480 + outputs = runner.llm.generate(

481 + [prompt] * 2,

482 sampling_params=sampling_params,

483 use_tqdm=True,

484 )

485 +

486 assert outputs is not None

487 for output in outputs:

488 assert output is not None

489 assert isinstance(output, RequestOutput)

490 + prompt = output.prompt

491 generated_text = output.outputs[0].text

364 - assert generated_text is not None

492 print(generated_text)

493 + assert generated_text is not None

494 + assert re.fullmatch(sample_regex, generated_text) is not None

495 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

496

367 - # Parse to verify it is valid JSON

368 - parsed_json = json.loads(generated_text)

369 - assert isinstance(parsed_json, dict)

370 -

371 - if backend not in ["outlines", "lm-format-enforcer"]:

497 #

373 - # Test 4: Generate SQL statement using EBNF grammar

498 + # Test 8: Generate text based on a choices

499 #

500 sampling_params = SamplingParams(

501 temperature=0.8,

502 top_p=0.95,

378 - max_tokens=1000,

379 - structured_outputs=StructuredOutputsParams(grammar=sample_sql_ebnf),

503 + structured_outputs=StructuredOutputsParams(

504 + choice=sample_structured_outputs_choices

505 + ),

506 )

381 - outputs = llm.generate(

507 +

508 + outputs = runner.llm.generate(

509 (

383 - "Generate a sql statement that selects col_1 from "

384 - "table_1 where it is equal to 1. Make the response as short as "

385 - "possible."

510 + "The best language for type-safe systems programming is "

511 + "(Make the response as short as possible.) "

512 ),

513 sampling_params=sampling_params,

514 use_tqdm=True,

515 )

390 -

516 assert outputs is not None

517 for output in outputs:

518 assert output is not None

519 assert isinstance(output, RequestOutput)

520 prompt = output.prompt

396 -

521 generated_text = output.outputs[0].text

522 + print(generated_text)

523 assert generated_text is not None

399 -

400 - # remove spaces for comparison b/c we removed them in the grammar

401 - ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(" ", "")

402 -

403 - assert generated_text.strip() == ground_truth

404 -

524 + assert generated_text in sample_structured_outputs_choices

525 print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

526

527 #

408 - # Test 5: Generate SQL statement using Lark grammar

528 + # Test 9: Generate structured output using a Pydantic model with an enum

529 #

530 + json_schema = CarDescription.model_json_schema()

531 sampling_params = SamplingParams(

411 - temperature=0.8,

412 - top_p=0.95,

532 + temperature=1.0,

533 max_tokens=1000,

414 - structured_outputs=StructuredOutputsParams(grammar=sample_sql_lark),

534 + structured_outputs=StructuredOutputsParams(json=json_schema),

535 )

416 - outputs = llm.generate(

536 +

537 + outputs = runner.llm.generate(

538 (

418 - "Generate a sql statement that selects col_1 from "

419 - "table_1 where it is equal to 1. Make the response as short as "

539 + "Generate a JSON with the brand, model and car_type of the most "

540 + "iconic car from the 90's. Make the response as short as "

541 "possible."

542 ),

543 sampling_params=sampling_params,

@@ -424,6 +545,7 @@ def test_structured_output(

545 )

546

547 assert outputs is not None

548 +

549 for output in outputs:

550 assert output is not None

551 assert isinstance(output, RequestOutput)

@@ -431,220 +553,97 @@ def test_structured_output(

553

554 generated_text = output.outputs[0].text

555 assert generated_text is not None

434 -

435 - # use Lark to parse the output, and make sure it's a valid parse tree

436 - from lark import Lark

437 -

438 - parser = Lark(sample_sql_lark)

439 - parser.parse(generated_text)

440 -

441 - # remove spaces for comparison b/c we removed them in the grammar

442 - ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(" ", "")

443 -

444 - assert generated_text.strip() == ground_truth

445 -

556 print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

557 + try:

558 + output_json = json.loads(generated_text)

559 + except json.JSONDecodeError as e:

560 + pytest.fail(

561 + f"Invalid JSON from backend={backend}: {generated_text!r}\n"

562 + f"Schema: {json_schema}\nError: {e}"

563 + )

564 + jsonschema.validate(instance=output_json, schema=json_schema)

565

566 #

449 - # Test 6: Test invalid grammar input

567 + # Test 10: Generate structured with minLength and maxLength

568 #

569 + min_length = 50

570 + max_length = 50

571 + json_schema = {

572 + "type": "object",

573 + "properties": {

574 + "description": {

575 + "type": "string",

576 + "maxLength": max_length,

577 + "minLength": min_length,

578 + }

579 + },

580 + "required": ["description"],

581 + "additionalProperties": False,

582 + }

583 +

584 sampling_params = SamplingParams(

452 - temperature=0.8,

453 - top_p=0.95,

454 - max_tokens=1000,

455 - structured_outputs=StructuredOutputsParams(grammar="not a grammar"),

585 + temperature=1.0,

586 + max_tokens=4096,

587 + structured_outputs=StructuredOutputsParams(json=json_schema),

588 )

457 - with pytest.raises(ValueError, match="Failed to convert the grammar "):

458 - llm.generate(

459 - (

460 - "Generate a sql statement that selects col_1 from "

461 - "table_1 where it is equal to 1. Make the response as short "

462 - "as possible."

463 - ),

464 - sampling_params=sampling_params,

465 - use_tqdm=True,

466 - )

467 -

468 - #

469 - # Test 7: Generate text based on a regex pattern

470 - #

471 - sampling_params = SamplingParams(

472 - temperature=0.8,

473 - top_p=0.95,

474 - structured_outputs=StructuredOutputsParams(regex=sample_regex),

475 - )

476 -

477 - prompt = (

478 - f"Give an example IPv4 address with this regex: {sample_regex}. "

479 - f"Make the response as short as possible."

480 - )

481 - outputs = llm.generate(

482 - [prompt] * 2,

483 - sampling_params=sampling_params,

484 - use_tqdm=True,

485 - )

486 -

487 - assert outputs is not None

488 - for output in outputs:

489 - assert output is not None

490 - assert isinstance(output, RequestOutput)

491 - prompt = output.prompt

492 - generated_text = output.outputs[0].text

493 - print(generated_text)

494 - assert generated_text is not None

495 - assert re.fullmatch(sample_regex, generated_text) is not None

496 - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

497 -

498 - #

499 - # Test 8: Generate text based on a choices

500 - #

501 - sampling_params = SamplingParams(

502 - temperature=0.8,

503 - top_p=0.95,

504 - structured_outputs=StructuredOutputsParams(

505 - choice=sample_structured_outputs_choices

506 - ),

507 - )

589

509 - outputs = llm.generate(

510 - (

511 - "The best language for type-safe systems programming is "

512 - "(Make the response as short as possible.) "

513 - ),

514 - sampling_params=sampling_params,

515 - use_tqdm=True,

516 - )

517 - assert outputs is not None

518 - for output in outputs:

519 - assert output is not None

520 - assert isinstance(output, RequestOutput)

521 - prompt = output.prompt

522 - generated_text = output.outputs[0].text

523 - print(generated_text)

524 - assert generated_text is not None

525 - assert generated_text in sample_structured_outputs_choices

526 - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

527 -

528 - #

529 - # Test 9: Generate structured output using a Pydantic model with an enum

530 - #

531 - json_schema = CarDescription.model_json_schema()

532 - sampling_params = SamplingParams(

533 - temperature=1.0,

534 - max_tokens=1000,

535 - structured_outputs=StructuredOutputsParams(json=json_schema),

536 - )

537 -

538 - outputs = llm.generate(

539 - (

540 - "Generate a JSON with the brand, model and car_type of the most "

541 - "iconic car from the 90's. Make the response as short as "

542 - "possible."

543 - ),

544 - sampling_params=sampling_params,

545 - use_tqdm=True,

546 - )

590 + outputs = runner.llm.generate(

591 + (

592 + "Generate a description of a frog using 50 characters. "

593 + "Make the response as short as possible."

594 + ),

595 + sampling_params=sampling_params,

596 + use_tqdm=True,

597 + )

598

548 - assert outputs is not None

599 + assert outputs is not None

600

550 - for output in outputs:

551 - assert output is not None

552 - assert isinstance(output, RequestOutput)

553 - prompt = output.prompt

601 + for output in outputs:

602 + assert output is not None

603 + assert isinstance(output, RequestOutput)

604 + prompt = output.prompt

605

555 - generated_text = output.outputs[0].text

556 - assert generated_text is not None

557 - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

558 - try:

559 - output_json = json.loads(generated_text)

560 - except json.JSONDecodeError as e:

561 - pytest.fail(

562 - f"Invalid JSON from backend={backend}: {generated_text!r}\n"

563 - f"Schema: {json_schema}\nError: {e}"

564 - )

565 - jsonschema.validate(instance=output_json, schema=json_schema)

566 -

567 - #

568 - # Test 10: Generate structured with minLength and maxLength

569 - #

570 - min_length = 50

571 - max_length = 50

572 - json_schema = {

573 - "type": "object",

574 - "properties": {

575 - "description": {

576 - "type": "string",

577 - "maxLength": max_length,

578 - "minLength": min_length,

606 + generated_text = output.outputs[0].text

607 + assert generated_text is not None

608 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

609 + try:

610 + output_json = json.loads(generated_text)

611 + except json.JSONDecodeError as e:

612 + pytest.fail(

613 + f"Invalid JSON from backend={backend}: {generated_text!r}\n"

614 + f"Schema: {json_schema}\nError: {e}"

615 + )

616 + jsonschema.validate(instance=output_json, schema=json_schema)

617 +

618 + if backend not in ["outlines", "lm-format-enforcer"]:

619 + #

620 + # Test 11: Generate structured output using structural_tag format

621 + #

622 + structural_tag_config = {

623 + "type": "structural_tag",

624 + "structures": [

625 + {

626 + "begin": "<function=get_weather>",

627 + "schema": {

628 + "type": "object",

629 + "properties": {"city": {"type": "string"}},

630 + "additionalProperties": False,

631 + },

632 + "end": "</function>",

633 + }

634 + ],

635 + "triggers": ["<function="],

636 }

580 - },

581 - "required": ["description"],

582 - "additionalProperties": False,

583 - }

584 -

585 - sampling_params = SamplingParams(

586 - temperature=1.0,

587 - max_tokens=4096,

588 - structured_outputs=StructuredOutputsParams(json=json_schema),

589 - )

590 -

591 - outputs = llm.generate(

592 - (

593 - "Generate a description of a frog using 50 characters. "

594 - "Make the response as short as possible."

595 - ),

596 - sampling_params=sampling_params,

597 - use_tqdm=True,

598 - )

599 -

600 - assert outputs is not None

637

602 - for output in outputs:

603 - assert output is not None

604 - assert isinstance(output, RequestOutput)

605 - prompt = output.prompt

606 -

607 - generated_text = output.outputs[0].text

608 - assert generated_text is not None

609 - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

610 - try:

611 - output_json = json.loads(generated_text)

612 - except json.JSONDecodeError as e:

613 - pytest.fail(

614 - f"Invalid JSON from backend={backend}: {generated_text!r}\n"

615 - f"Schema: {json_schema}\nError: {e}"

638 + sampling_params = SamplingParams(

639 + temperature=0.0,

640 + max_tokens=4096,

641 + structured_outputs=StructuredOutputsParams(

642 + structural_tag=json.dumps(structural_tag_config)

643 + ),

644 )

617 - jsonschema.validate(instance=output_json, schema=json_schema)

618 -

619 - if backend not in ["outlines", "lm-format-enforcer"]:

620 - #

621 - # Test 11: Generate structured output using structural_tag format

622 - #

623 - structural_tag_config = {

624 - "type": "structural_tag",

625 - "structures": [

626 - {

627 - "begin": "<function=get_weather>",

628 - "schema": {

629 - "type": "object",

630 - "properties": {"city": {"type": "string"}},

631 - "additionalProperties": False,

632 - },

633 - "end": "</function>",

634 - }

635 - ],

636 - "triggers": ["<function="],

637 - }

638 -

639 - sampling_params = SamplingParams(

640 - temperature=0.0,

641 - max_tokens=4096,

642 - structured_outputs=StructuredOutputsParams(

643 - structural_tag=json.dumps(structural_tag_config)

644 - ),

645 - )

645

647 - prompt = """

646 + prompt = """

647 You have access to the following function to retrieve the weather in a city:

648

649 {

@@ -683,37 +682,39 @@ Given the previous instructions, what is the weather in New York City? \

682 Make the response as short as possible.

683 """

684

686 - # Change this once other backends support structural_tag

687 - outputs = llm.generate(prompt, sampling_params=sampling_params, use_tqdm=True)

688 - assert outputs is not None

689 -

690 - for output in outputs:

691 - assert output is not None

692 - assert isinstance(output, RequestOutput)

693 - generated_text = output.outputs[0].text

694 - assert generated_text is not None

695 -

696 - # Search for function call pattern in the response

697 - function_call_pattern = r"<function=get_weather>(.*?)</function>"

698 - matches = re.findall(function_call_pattern, generated_text)

685 + # Change this once other backends support structural_tag

686 + outputs = runner.llm.generate(

687 + prompt, sampling_params=sampling_params, use_tqdm=True

688 + )

689 + assert outputs is not None

690

700 - if not matches:

701 - print(

702 - f"Warning: No function calls found in response: {generated_text!r}"

703 - )

704 - continue

691 + for output in outputs:

692 + assert output is not None

693 + assert isinstance(output, RequestOutput)

694 + generated_text = output.outputs[0].text

695 + assert generated_text is not None

696

706 - # Take the first function call if multiple are found

707 - json_str = matches[0]

708 - try:

709 - json_content = json.loads(json_str)

710 - assert "city" in json_content

711 - assert isinstance(json_content["city"], str)

712 - print(f"Found valid function call: {generated_text!r}")

713 - except (json.JSONDecodeError, AssertionError) as e:

714 - pytest.fail(

715 - f"Invalid function call format: {generated_text!r}\nError: {str(e)}"

716 - )

697 + # Search for function call pattern in the response

698 + function_call_pattern = r"<function=get_weather>(.*?)</function>"

699 + matches = re.findall(function_call_pattern, generated_text)

700 +

701 + if not matches:

702 + print(

703 + f"Warning: No function calls found in response: {generated_text!r}"

704 + )

705 + continue

706 +

707 + # Take the first function call if multiple are found

708 + json_str = matches[0]

709 + try:

710 + json_content = json.loads(json_str)

711 + assert "city" in json_content

712 + assert isinstance(json_content["city"], str)

713 + print(f"Found valid function call: {generated_text!r}")

714 + except (json.JSONDecodeError, AssertionError) as e:

715 + pytest.fail(

716 + f"Invalid function call format: {generated_text!r}\nError: {str(e)}"

717 + )

718

719

720 @pytest.mark.parametrize(

@@ -738,14 +739,15 @@ def test_structured_output_with_reasoning_matrices(

739 model_name: str,

740 speculative_config: dict[str, Any] | None,

741 async_scheduling: bool,

742 + vllm_runner,

743 ):

744 if current_platform.is_tpu() and speculative_config:

745 pytest.skip("TPU does not support speculative decoding")

746

747 # Use a single LLM instance for several scenarios to

748 # speed up the test suite.

747 - llm = LLM(

748 - model=model_name,

749 + with vllm_runner(

750 + model_name,

751 # Don't use eager execution on TPUs because we want to test for no

752 # recompilation at runtime

753 enforce_eager=bool(not current_platform.is_tpu()),

@@ -759,271 +761,273 @@ def test_structured_output_with_reasoning_matrices(

761 tokenizer_mode=tokenizer_mode,

762 speculative_config=speculative_config,

763 async_scheduling=async_scheduling,

762 - )

763 - tokenizer = llm.get_tokenizer()

764 - reasoner = ReasoningParserManager.get_reasoning_parser(reasoning_parser)(

765 - tokenizer=tokenizer

766 - )

767 -

768 - 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

769 - reasoning_schema = {

770 - "type": "object",

771 - "properties": {"result": {"type": "integer"}},

772 - "required": ["result"],

773 - "additionalProperties": False,

774 - }

775 - if "Qwen3" in model_name:

776 - reasoning_prompt += "<think>\n"

777 -

778 - sampling_params = SamplingParams(

779 - temperature=0.1,

780 - max_tokens=8192,

781 - structured_outputs=StructuredOutputsParams(json=reasoning_schema),

782 - )

783 - outputs = llm.generate(

784 - [reasoning_prompt],

785 - sampling_params=sampling_params,

786 - use_tqdm=True,

787 - )

788 -

789 - assert outputs is not None

790 - output = outputs[0]

791 - assert output is not None and isinstance(output, RequestOutput)

792 - prompt = output.prompt

793 - generated_text = output.outputs[0].text

794 - reasoning, content = run_reasoning_extraction(reasoner, [generated_text])

795 - print(f"Prompt: {prompt!r}\nReasoning: {reasoning!r}\nContent: {content!r}")

796 -

797 - if "Qwen3" in model_name:

798 - assert content is not None

799 -

800 - assert reasoning is not None

801 -

802 - if content is not None:

803 - output_json = json.loads(content)

804 - jsonschema.validate(instance=output_json, schema=reasoning_schema)

764 + ) as runner:

765 + tokenizer = runner.llm.get_tokenizer()

766 + reasoner = ReasoningParserManager.get_reasoning_parser(reasoning_parser)(

767 + tokenizer=tokenizer

768 + )

769 +

770 + 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

771 + reasoning_schema = {

772 + "type": "object",

773 + "properties": {"result": {"type": "integer"}},

774 + "required": ["result"],

775 + "additionalProperties": False,

776 + }

777 + if "Qwen3" in model_name:

778 + reasoning_prompt += "<think>\n"

779 +

780 + sampling_params = SamplingParams(

781 + temperature=0.1,

782 + max_tokens=8192,

783 + structured_outputs=StructuredOutputsParams(json=reasoning_schema),

784 + )

785 + outputs = runner.llm.generate(

786 + [reasoning_prompt],

787 + sampling_params=sampling_params,

788 + use_tqdm=True,

789 + )

790 +

791 + assert outputs is not None

792 + output = outputs[0]

793 + assert output is not None and isinstance(output, RequestOutput)

794 + prompt = output.prompt

795 + generated_text = output.outputs[0].text

796 + reasoning, content = run_reasoning_extraction(reasoner, [generated_text])

797 + print(f"Prompt: {prompt!r}\nReasoning: {reasoning!r}\nContent: {content!r}")

798 +

799 + if "Qwen3" in model_name:

800 + assert content is not None

801 +

802 + assert reasoning is not None

803 +

804 + if content is not None:

805 + output_json = json.loads(content)

806 + jsonschema.validate(instance=output_json, schema=reasoning_schema)

807

808

809 @pytest.mark.parametrize("model_name, tokenizer_mode", PARAMS_MODELS_TOKENIZER_MODE)

810 def test_structured_output_auto_mode(

811 model_name: str,

812 tokenizer_mode: str,

813 + vllm_runner,

814 ):

815 unsupported_json_schema = UNSUPPORTED_JSON_SCHEMA

813 - llm = LLM(

814 - model=model_name,

816 + with vllm_runner(

817 + model_name,

818 max_model_len=1024,

819 structured_outputs_config=dict(backend="auto"),

820 tokenizer_mode=tokenizer_mode,

821 load_format="auto",

822 config_format="auto",

820 - )

821 -

822 - sampling_params = SamplingParams(

823 - temperature=1.0,

824 - max_tokens=1000,

825 - structured_outputs=StructuredOutputsParams(json=unsupported_json_schema),

826 - )

827 -

828 - prompts = (

829 - "Give an example JSON object for a grade "

830 - "that fits this schema: "

831 - f"{unsupported_json_schema}. Make the response as short as possible."

832 - )

833 - # This would fail with the default of "xgrammar", but in "auto"

834 - # we will handle fallback automatically.

835 - outputs = llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True)

836 - # Make sure `auto` backend handling doesn't mess up sampling_params

837 - # and that we can reuse it without error.

838 - outputs.extend(

839 - llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True)

840 - )

841 -

842 - assert outputs is not None

843 - for output in outputs:

844 - assert output is not None

845 - assert isinstance(output, RequestOutput)

846 - generated_text = output.outputs[0].text

847 - assert generated_text is not None

848 - print(generated_text)

823 + ) as runner:

824 + sampling_params = SamplingParams(

825 + temperature=1.0,

826 + max_tokens=1000,

827 + structured_outputs=StructuredOutputsParams(json=unsupported_json_schema),

828 + )

829 +

830 + prompts = (

831 + "Give an example JSON object for a grade "

832 + "that fits this schema: "

833 + f"{unsupported_json_schema}. Make the response as short as possible."

834 + )

835 + # This would fail with the default of "xgrammar", but in "auto"

836 + # we will handle fallback automatically.

837 + outputs = runner.llm.generate(

838 + prompts, sampling_params=sampling_params, use_tqdm=True

839 + )

840 + # Make sure `auto` backend handling doesn't mess up sampling_params

841 + # and that we can reuse it without error.

842 + outputs.extend(

843 + runner.llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True)

844 + )

845 +

846 + assert outputs is not None

847 + for output in outputs:

848 + assert output is not None

849 + assert isinstance(output, RequestOutput)

850 + generated_text = output.outputs[0].text

851 + assert generated_text is not None

852 + print(generated_text)

853

850 - # Parse to verify it is valid JSON

851 - parsed_json = json.loads(generated_text)

852 - assert isinstance(parsed_json, dict)

854 + # Parse to verify it is valid JSON

855 + parsed_json = json.loads(generated_text)

856 + assert isinstance(parsed_json, dict)

857

858

855 -def test_guidance_no_additional_properties():

856 - llm = LLM(

857 - model="Qwen/Qwen2.5-1.5B-Instruct",

859 +def test_guidance_no_additional_properties(vllm_runner):

860 + with vllm_runner(

861 + "Qwen/Qwen2.5-1.5B-Instruct",

862 max_model_len=1024,

863 structured_outputs_config=dict(

864 backend="guidance",

865 disable_any_whitespace=True,

866 disable_additional_properties=True,

867 ),

864 - )

865 -

866 - schema = {

867 - "type": "object",

868 - "properties": {

869 - "a1": {"type": "string"},

870 - "a2": {"type": "string"},

871 - "a3": {"type": "string"},

872 - },

873 - "required": ["a1", "a2", "a3"],

874 - }

868 + ) as runner:

869 + schema = {

870 + "type": "object",

871 + "properties": {

872 + "a1": {"type": "string"},

873 + "a2": {"type": "string"},

874 + "a3": {"type": "string"},

875 + },

876 + "required": ["a1", "a2", "a3"],

877 + }

878

876 - prompt = (

877 - "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a "

878 - "helpful assistant.<|im_end|>\n<|im_start|>user\nPlease generate a "

879 - "large JSON object with key-value pairs a1=b1, a2=b2, ..., a20=b20. "

880 - "Make the response as short as possible."

881 - "<|im_end|>\n<|im_start|>assistant\n"

882 - )

883 -

884 - def generate_with_backend(backend):

885 - structured_outputs_params = StructuredOutputsParams(

886 - json=schema,

887 - backend=backend,

888 - disable_any_whitespace=True,

889 - disable_additional_properties=True,

890 - )

891 - sampling_params = SamplingParams(

892 - temperature=0, max_tokens=256, structured_outputs=structured_outputs_params

879 + prompt = (

880 + "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a "

881 + "helpful assistant.<|im_end|>\n<|im_start|>user\nPlease generate a "

882 + "large JSON object with key-value pairs a1=b1, a2=b2, ..., a20=b20. "

883 + "Make the response as short as possible."

884 + "<|im_end|>\n<|im_start|>assistant\n"

885 )

886

895 - outputs = llm.generate(prompt, sampling_params=sampling_params)

896 - assert outputs is not None

897 - generated_text = outputs[0].outputs[0].text

898 - assert generated_text is not None

899 - parsed_json = json.loads(generated_text)

900 - assert isinstance(parsed_json, dict)

901 - jsonschema.validate(instance=parsed_json, schema=schema)

902 - return parsed_json

903 -

904 - generated = generate_with_backend("guidance")

905 - assert "a1" in generated

906 - assert "a2" in generated

907 - assert "a3" in generated

908 - assert "a4" not in generated

909 - assert "a5" not in generated

910 - assert "a6" not in generated

887 + def generate_with_backend(backend):

888 + structured_outputs_params = StructuredOutputsParams(

889 + json=schema,

890 + backend=backend,

891 + disable_any_whitespace=True,

892 + disable_additional_properties=True,

893 + )

894 + sampling_params = SamplingParams(

895 + temperature=0,

896 + max_tokens=256,

897 + structured_outputs=structured_outputs_params,

898 + )

899 +

900 + outputs = runner.llm.generate(prompt, sampling_params=sampling_params)

901 + assert outputs is not None

902 + generated_text = outputs[0].outputs[0].text

903 + assert generated_text is not None

904 + parsed_json = json.loads(generated_text)

905 + assert isinstance(parsed_json, dict)

906 + jsonschema.validate(instance=parsed_json, schema=schema)

907 + return parsed_json

908 +

909 + generated = generate_with_backend("guidance")

910 + assert "a1" in generated

911 + assert "a2" in generated

912 + assert "a3" in generated

913 + assert "a4" not in generated

914 + assert "a5" not in generated

915 + assert "a6" not in generated

916

917

918 @pytest.mark.parametrize("backend", ["guidance", "xgrammar", "outlines"])

919 def test_structured_output_batched_with_non_structured_outputs_requests(

920 backend: str,

921 + vllm_runner,

922 ):

923 sample_json_schema = SAMPLE_JSON_SCHEMA

924 # Don't use eager execution on TPUs because we want to test for no

925 # recompilation at runtime

926 enforce_eager = bool(not current_platform.is_tpu())

927

922 - llm = LLM(

923 - model="meta-llama/Meta-Llama-3.1-8B-Instruct",

928 + with vllm_runner(

929 + "meta-llama/Meta-Llama-3.1-8B-Instruct",

930 enforce_eager=enforce_eager,

931 max_model_len=1024,

932 structured_outputs_config=StructuredOutputsConfig(

933 backend=backend,

934 disable_any_whitespace=backend in {"xgrammar", "guidance"},

935 ),

930 - )

931 -

932 - structured_outputs_prompt = (

933 - "Give an example JSON for an employee profile that fits this "

934 - "schema. Make the response as short as possible. Schema: "

935 - f"{sample_json_schema}"

936 - )

937 -

938 - non_structured_outputs_prompt = "The diameter of the Earth in kilometers is "

939 -

940 - prompts = [structured_outputs_prompt, non_structured_outputs_prompt]

941 - sampling_params = [

942 - SamplingParams(

943 - temperature=1.0,

944 - max_tokens=400,

945 - structured_outputs=StructuredOutputsParams(json=sample_json_schema),

946 - ),

947 - # No max tokens, temp=0 to assert on contents

948 - SamplingParams(

949 - seed=42,

950 - temperature=0,

951 - top_p=1.0,

952 - ),

953 - ]

936 + ) as runner:

937 + structured_outputs_prompt = (

938 + "Give an example JSON for an employee profile that fits this "

939 + "schema. Make the response as short as possible. Schema: "

940 + f"{sample_json_schema}"

941 + )

942

955 - outputs = llm.generate(

956 - prompts=prompts, sampling_params=sampling_params, use_tqdm=True

957 - )

943 + non_structured_outputs_prompt = "The diameter of the Earth in kilometers is "

944

959 - assert outputs is not None

945 + prompts = [structured_outputs_prompt, non_structured_outputs_prompt]

946 + sampling_params = [

947 + SamplingParams(

948 + temperature=1.0,

949 + max_tokens=400,

950 + structured_outputs=StructuredOutputsParams(json=sample_json_schema),

951 + ),

952 + # No max tokens, temp=0 to assert on contents

953 + SamplingParams(

954 + seed=42,

955 + temperature=0,

956 + top_p=1.0,

957 + ),

958 + ]

959

961 - # Free memory as soon as possible as failed assertions

962 - # will short circuit and not free up memory

963 - del llm

964 - torch.accelerator.empty_cache()

965 - cleanup_dist_env_and_memory()

960 + outputs = runner.llm.generate(

961 + prompts=prompts, sampling_params=sampling_params, use_tqdm=True

962 + )

963

967 - for index, output in enumerate(outputs):

968 - assert output is not None

969 - assert isinstance(output, RequestOutput)

970 - prompt = output.prompt

964 + assert outputs is not None

965

972 - generated_text = output.outputs[0].text

973 - assert generated_text is not None

974 - print(f"Prompt:\n{prompt!r}\nGenerated text:\n{generated_text!r}")

966 + for index, output in enumerate(outputs):

967 + assert output is not None

968 + assert isinstance(output, RequestOutput)

969 + prompt = output.prompt

970

976 - if index == 0:

977 - # First prompt is structured outputs, expect valid JSON

978 - assert "\n" not in generated_text

979 - output_json = json.loads(generated_text)

980 - jsonschema.validate(instance=output_json, schema=sample_json_schema)

981 - else:

982 - # Second prompt is not structured outputs, expect valid output

983 - # Cannot assert on exact output, but we can expect it to be factual

984 - assert "12,742" in generated_text

971 + generated_text = output.outputs[0].text

972 + assert generated_text is not None

973 + print(f"Prompt:\n{prompt!r}\nGenerated text:\n{generated_text!r}")

974

986 - # non-structured outputs requests should not return a valid JSON here

987 - with pytest.raises(ValueError):

975 + if index == 0:

976 + # First prompt is structured outputs, expect valid JSON

977 + assert "\n" not in generated_text

978 output_json = json.loads(generated_text)

979 + jsonschema.validate(instance=output_json, schema=sample_json_schema)

980 + else:

981 + # Second prompt is not structured outputs, expect valid output

982 + # Cannot assert on exact output, but we can expect it to be factual

983 + assert "12,742" in generated_text

984 +

985 + # non-structured outputs requests should not return a valid JSON here

986 + with pytest.raises(ValueError):

987 + output_json = json.loads(generated_text)

988

989

990 @pytest.mark.parametrize("backend", ["xgrammar"])

992 -def test_structured_output_with_structural_tag(backend: str):

993 - llm = LLM(

994 - model="Qwen/Qwen2.5-1.5B-Instruct",

991 +def test_structured_output_with_structural_tag(backend: str, vllm_runner):

992 + with vllm_runner(

993 + "Qwen/Qwen2.5-1.5B-Instruct",

994 structured_outputs_config=StructuredOutputsConfig(backend=backend),

996 - )

997 -

998 - structural_tag_config = {

999 - "type": "structural_tag",

1000 - "format": {

1001 - "type": "triggered_tags",

1002 - "tags": [

1003 - {"begin": "hello_flag", "content": {"type": "any_text"}, "end": "hello"}

1004 - ],

1005 - "triggers": ["hello"],

1006 - "stop_after_first": False,

1007 - },

1008 - }

995 + ) as runner:

996 + structural_tag_config = {

997 + "type": "structural_tag",

998 + "format": {

999 + "type": "triggered_tags",

1000 + "tags": [

1001 + {

1002 + "begin": "hello_flag",

1003 + "content": {"type": "any_text"},

1004 + "end": "hello",

1005 + }

1006 + ],

1007 + "triggers": ["hello"],

1008 + "stop_after_first": False,

1009 + },

1010 + }

1011

1010 - sampling_params = SamplingParams(

1011 - temperature=0.0,

1012 - max_tokens=500,

1013 - structured_outputs=StructuredOutputsParams(

1014 - structural_tag=json.dumps(structural_tag_config)

1015 - ),

1016 - )

1017 -

1018 - prompt = "Hello and repeat hello 10 times, do not say anything else. Only say hello hello hello, now start"

1019 - outputs = llm.generate(prompt, sampling_params=sampling_params, use_tqdm=True)

1020 - assert outputs is not None

1021 - for output in outputs:

1022 - assert output is not None

1023 - assert isinstance(output, RequestOutput)

1024 - prompt = output.prompt

1025 - generated_text = output.outputs[0].text

1026 - assert generated_text is not None

1027 - assert "hello_flag" in generated_text, (

1028 - f"Expected 'hello_flag' to be in generated text, but got: {generated_text}"

1012 + sampling_params = SamplingParams(

1013 + temperature=0.0,

1014 + max_tokens=500,

1015 + structured_outputs=StructuredOutputsParams(

1016 + structural_tag=json.dumps(structural_tag_config)

1017 + ),

1018 )

1019 +

1020 + prompt = "Hello and repeat hello 10 times, do not say anything else. Only say hello hello hello, now start"

1021 + outputs = runner.llm.generate(

1022 + prompt, sampling_params=sampling_params, use_tqdm=True

1023 + )

1024 + assert outputs is not None

1025 + for output in outputs:

1026 + assert output is not None

1027 + assert isinstance(output, RequestOutput)

1028 + prompt = output.prompt

1029 + generated_text = output.outputs[0].text

1030 + assert generated_text is not None

1031 + assert "hello_flag" in generated_text, (

1032 + f"Expected 'hello_flag' to be in generated text, but got: {generated_text}"

1033 + )

tests/entrypoints/multimodal/conftest.py

-73

@@ -1,11 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from collections.abc import Callable, Iterator

4 -from contextlib import contextmanager

5 -from typing import Any

6 -

7 -import pytest

8 -

3 # Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)

4 TEST_IMAGE_ASSETS = [

5 "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 = [

7 "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",

8 "RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",

9 ]

16 -

17 -

18 -def _shutdown_llm(llm: Any, gpu_memory_utilization: float) -> None:

19 - from vllm.distributed import cleanup_dist_env_and_memory

20 - from vllm.platforms import current_platform

21 -

22 - try:

23 - shutdown_timeout = 60.0 if current_platform.is_rocm() else None

24 - llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout)

25 - except Exception:

26 - pass

27 -

28 - del llm

29 -

30 - try:

31 - import torch

32 -

33 - torch._dynamo.reset()

34 - except Exception:

35 - pass

36 -

37 - cleanup_dist_env_and_memory()

38 -

39 - if current_platform.is_rocm():

40 - from tests.utils import wait_for_rocm_memory_to_settle

41 -

42 - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

43 -

44 -

45 -@contextmanager

46 -def managed_llm(*args: Any, **kwargs: Any) -> Iterator[Any]:

47 - from vllm import LLM

48 -

49 - llm = LLM(*args, **kwargs)

50 - gpu_memory_utilization = (

51 - llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization

52 - )

53 - try:

54 - yield llm

55 - finally:

56 - _shutdown_llm(llm, gpu_memory_utilization)

57 -

58 -

59 -def _make_managed_llm_factory() -> Iterator[Callable[..., Any]]:

60 - from vllm import LLM

61 -

62 - llms: list[tuple[Any, float]] = []

63 -

64 - def make_llm(*args: Any, **kwargs: Any) -> Any:

65 - llm = LLM(*args, **kwargs)

66 - gpu_memory_utilization = (

67 - llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization

68 - )

69 - llms.append((llm, gpu_memory_utilization))

70 - return llm

71 -

72 - try:

73 - yield make_llm

74 - finally:

75 - while llms:

76 - llm, gpu_memory_utilization = llms.pop()

77 - _shutdown_llm(llm, gpu_memory_utilization)

78 -

79 -

80 -@pytest.fixture

81 -def multimodal_llm_factory() -> Iterator[Callable[..., Any]]:

82 - yield from _make_managed_llm_factory()

tests/entrypoints/multimodal/llm/test_chat.py

+9 -4

@@ -1,21 +1,26 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import weakref

4 +

5 import pytest

7 from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS

10 @pytest.fixture(scope="function")

9 -def vision_llm(multimodal_llm_factory):

10 - return multimodal_llm_factory(

11 - model="microsoft/Phi-3.5-vision-instruct",

11 +def vision_llm(vllm_runner):

12 + with vllm_runner(

13 + "microsoft/Phi-3.5-vision-instruct",

14 max_model_len=4096,

15 max_num_seqs=5,

16 enforce_eager=True,

17 trust_remote_code=True,

18 limit_mm_per_prompt={"image": 2},

19 seed=0,

18 - )

20 + ) as runner:

21 + # pytest caches yielded fixtures until after teardown, so use a proxy to

22 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

23 + yield weakref.proxy(runner.llm)

26 @pytest.mark.parametrize(

tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py

+91 -93

@@ -69,7 +69,7 @@ def test_inject_into_mm_cache(

69 image_urls,

70 mm_processor_cache_type,

71 caplog_vllm,

72 - multimodal_llm_factory,

72 + vllm_runner,

73 ):

74 """Test that inject_into_mm_cache() injects pre-processed mm_kwargs into

75 the processor cache and MM cache hit metrics are updated correctly.

@@ -79,117 +79,115 @@ def test_inject_into_mm_cache(

79 2. Extract cached kwargs, call inject_into_mm_cache with a new hash,

80 then generate with a pre-rendered input -> verifies injection works

81 """

82 - llm = multimodal_llm_factory(

83 - model="llava-hf/llava-1.5-7b-hf",

82 + with vllm_runner(

83 + "llava-hf/llava-1.5-7b-hf",

84 max_model_len=4096,

85 max_num_seqs=5,

86 enforce_eager=True,

87 disable_log_stats=False,

88 limit_mm_per_prompt={"image": 2},

89 mm_processor_cache_type=mm_processor_cache_type,

90 - )

91 -

92 - # Step 1: Normal requests to populate the cache

93 - llm.chat(_make_messages(image_urls[0]))

94 - assert _get_mm_cache_stats(llm.get_metrics()) == (1, 0)

95 -

96 - llm.chat(_make_messages(image_urls[0]))

97 - assert _get_mm_cache_stats(llm.get_metrics()) == (2, 1)

98 - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(50.0)

99 -

100 - # Step 2: Use a second image to get valid expanded tokens and

101 - # placeholder positions via the renderer.

102 - llm.chat(_make_messages(image_urls[1]))

103 - queries_before = _get_mm_cache_stats(llm.get_metrics())[0] # 3

104 -

105 - renderer = llm.llm_engine.renderer

106 - cache = renderer.mm_processor_cache

107 - assert cache is not None, "Processor cache should be enabled"

108 -

109 - _, eng_prompts = renderer.render_chat(

110 - [_make_messages(image_urls[1])],

111 - ChatParams(),

112 - )

113 - eng_input = eng_prompts[0]

114 -

115 - # Inject pre-processed mm_kwargs with a NEW hash via public API

116 - new_mm_hash = "deadbeef" * 8

117 - mm_hashes = {"image": [new_mm_hash]}

118 - mm_kwargs = eng_input["mm_kwargs"]

119 -

120 - llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs)

121 -

122 - # Build pre-rendered input (no externally_processed flag needed)

123 - pre_rendered_input = {

124 - "type": "multimodal",

125 - "prompt_token_ids": eng_input["prompt_token_ids"],

126 - "mm_kwargs": mm_kwargs,

127 - "mm_hashes": mm_hashes,

128 - "mm_placeholders": eng_input["mm_placeholders"],

129 - }

130 -

131 - llm.generate(

132 - pre_rendered_input,

133 - sampling_params=SamplingParams(max_tokens=1),

134 - )

135 -

136 - # Verify cache was queried and injection happened

137 - queries_after = _get_mm_cache_stats(llm.get_metrics())[0]

138 - assert queries_after > queries_before, (

139 - "Cache should have been queried for the injected item"

140 - )

141 - mm_rate = _get_mm_cache_log(llm, caplog_vllm)

142 - assert mm_rate >= 0.0, "MM cache hit rate should be reported"

90 + ) as runner:

91 + # Step 1: Normal requests to populate the cache

92 + runner.llm.chat(_make_messages(image_urls[0]))

93 + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (1, 0)

94 +

95 + runner.llm.chat(_make_messages(image_urls[0]))

96 + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (2, 1)

97 + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(50.0)

98 +

99 + # Step 2: Use a second image to get valid expanded tokens and

100 + # placeholder positions via the renderer.

101 + runner.llm.chat(_make_messages(image_urls[1]))

102 + queries_before = _get_mm_cache_stats(runner.llm.get_metrics())[0] # 3

103 +

104 + renderer = runner.llm.llm_engine.renderer

105 + cache = renderer.mm_processor_cache

106 + assert cache is not None, "Processor cache should be enabled"

107 +

108 + _, eng_prompts = renderer.render_chat(

109 + [_make_messages(image_urls[1])],

110 + ChatParams(),

111 + )

112 + eng_input = eng_prompts[0]

113 +

114 + # Inject pre-processed mm_kwargs with a NEW hash via public API

115 + new_mm_hash = "deadbeef" * 8

116 + mm_hashes = {"image": [new_mm_hash]}

117 + mm_kwargs = eng_input["mm_kwargs"]

118 +

119 + runner.llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs)

120 +

121 + # Build pre-rendered input (no externally_processed flag needed)

122 + pre_rendered_input = {

123 + "type": "multimodal",

124 + "prompt_token_ids": eng_input["prompt_token_ids"],

125 + "mm_kwargs": mm_kwargs,

126 + "mm_hashes": mm_hashes,

127 + "mm_placeholders": eng_input["mm_placeholders"],

128 + }

129 +

130 + runner.llm.generate(

131 + pre_rendered_input,

132 + sampling_params=SamplingParams(max_tokens=1),

133 + )

134 +

135 + # Verify cache was queried and injection happened

136 + queries_after = _get_mm_cache_stats(runner.llm.get_metrics())[0]

137 + assert queries_after > queries_before, (

138 + "Cache should have been queried for the injected item"

139 + )

140 + mm_rate = _get_mm_cache_log(runner.llm, caplog_vllm)

141 + assert mm_rate >= 0.0, "MM cache hit rate should be reported"

142

143

144 @pytest.mark.parametrize("image_urls", [TEST_IMAGE_ASSETS[:1]], indirect=True)

145 def test_inject_into_mm_cache_without_cache(

146 num_gpus_available,

147 image_urls,

149 - multimodal_llm_factory,

148 + vllm_runner,

149 ):

150 """Test that inject_into_mm_cache works gracefully when processor cache

151 is disabled (mm_processor_cache_gb=0). Should not crash.

152 """

154 - llm = multimodal_llm_factory(

155 - model="llava-hf/llava-1.5-7b-hf",

153 + with vllm_runner(

154 + "llava-hf/llava-1.5-7b-hf",

155 max_model_len=4096,

156 max_num_seqs=5,

157 enforce_eager=True,

158 disable_log_stats=False,

159 limit_mm_per_prompt={"image": 2},

160 mm_processor_cache_gb=0,

162 - )

163 -

164 - # Run a normal chat request first to warm up the model.

165 - llm.chat(_make_messages(image_urls[0]))

166 -

167 - # Use the renderer to get a proper EngineInput with expanded tokens

168 - renderer = llm.llm_engine.renderer

169 - _, eng_prompts = renderer.render_chat(

170 - [_make_messages(image_urls[0])],

171 - ChatParams(),

172 - )

173 - eng_input = eng_prompts[0]

174 -

175 - mm_hashes = {"image": ["abcd1234" * 8]}

176 - mm_kwargs = eng_input["mm_kwargs"]

177 -

178 - # inject_into_mm_cache should not crash even without cache

179 - llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs)

180 -

181 - # Build and generate with pre-rendered input

182 - pre_rendered_input = {

183 - "type": "multimodal",

184 - "prompt_token_ids": eng_input["prompt_token_ids"],

185 - "mm_kwargs": mm_kwargs,

186 - "mm_hashes": mm_hashes,

187 - "mm_placeholders": eng_input["mm_placeholders"],

188 - }

189 -

190 - result = llm.generate(

191 - pre_rendered_input,

192 - sampling_params=SamplingParams(max_tokens=1),

193 - )

194 - assert len(result) == 1, "Should produce one output"

195 - assert len(result[0].outputs) >= 1, "Should have at least one output sequence"

161 + ) as runner:

162 + # Run a normal chat request first to warm up the model.

163 + runner.llm.chat(_make_messages(image_urls[0]))

164 +

165 + # Use the renderer to get a proper EngineInput with expanded tokens

166 + renderer = runner.llm.llm_engine.renderer

167 + _, eng_prompts = renderer.render_chat(

168 + [_make_messages(image_urls[0])],

169 + ChatParams(),

170 + )

171 + eng_input = eng_prompts[0]

172 +

173 + mm_hashes = {"image": ["abcd1234" * 8]}

174 + mm_kwargs = eng_input["mm_kwargs"]

175 +

176 + # inject_into_mm_cache should not crash even without cache

177 + runner.llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs)

178 +

179 + # Build and generate with pre-rendered input

180 + pre_rendered_input = {

181 + "type": "multimodal",

182 + "prompt_token_ids": eng_input["prompt_token_ids"],

183 + "mm_kwargs": mm_kwargs,

184 + "mm_hashes": mm_hashes,

185 + "mm_placeholders": eng_input["mm_placeholders"],

186 + }

187 +

188 + result = runner.llm.generate(

189 + pre_rendered_input,

190 + sampling_params=SamplingParams(max_tokens=1),

191 + )

192 + assert len(result) == 1, "Should produce one output"

193 + assert len(result[0].outputs) >= 1, "Should have at least one output sequence"

tests/entrypoints/multimodal/llm/test_mm_cache_stats.py

+27 -28

@@ -61,38 +61,37 @@ def test_mm_cache_stats(

61 image_urls,

62 mm_processor_cache_type,

63 caplog_vllm,

64 - multimodal_llm_factory,

64 + vllm_runner,

65 ):

66 - llm = multimodal_llm_factory(

67 - model="llava-hf/llava-1.5-7b-hf",

66 + with vllm_runner(

67 + "llava-hf/llava-1.5-7b-hf",

68 max_model_len=4096,

69 max_num_seqs=5,

70 enforce_eager=True,

71 mm_processor_cache_type=mm_processor_cache_type,

72 disable_log_stats=False,

73 limit_mm_per_prompt={"image": 2},

74 - )

75 -

76 - llm.chat(_make_messages(image_urls[0]))

77 - assert _get_mm_cache_stats(llm.get_metrics()) == (1, 0)

78 - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0)

79 -

80 - llm.chat(_make_messages(image_urls[1]))

81 - assert _get_mm_cache_stats(llm.get_metrics()) == (2, 0)

82 - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0)

83 -

84 - llm.chat(_make_messages(image_urls[0]))

85 - assert _get_mm_cache_stats(llm.get_metrics()) == (3, 1)

86 - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(33.3)

87 -

88 - # NOTE: This only resets hit rate stats in CachingMetrics

89 - # The raw queries and hits counts remain unaffected

90 - llm.reset_mm_cache()

91 -

92 - llm.chat(_make_messages(image_urls[0]))

93 - assert _get_mm_cache_stats(llm.get_metrics()) == (4, 1)

94 - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0)

95 -

96 - llm.chat(_make_messages(image_urls[1]))

97 - assert _get_mm_cache_stats(llm.get_metrics()) == (5, 1)

98 - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0)

74 + ) as runner:

75 + runner.llm.chat(_make_messages(image_urls[0]))

76 + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (1, 0)

77 + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0)

78 +

79 + runner.llm.chat(_make_messages(image_urls[1]))

80 + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (2, 0)

81 + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0)

82 +

83 + runner.llm.chat(_make_messages(image_urls[0]))

84 + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (3, 1)

85 + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(33.3)

86 +

87 + # NOTE: This only resets hit rate stats in CachingMetrics

88 + # The raw queries and hits counts remain unaffected

89 + runner.llm.reset_mm_cache()

90 +

91 + runner.llm.chat(_make_messages(image_urls[0]))

92 + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (4, 1)

93 + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0)

94 +

95 + runner.llm.chat(_make_messages(image_urls[1]))

96 + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (5, 1)

97 + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0)

tests/entrypoints/multimodal/llm/test_mm_embeds_only.py

+9 -6

@@ -1,9 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import weakref

5 +

6 import pytest

6 -from tests.entrypoints.multimodal.conftest import managed_llm

8 from vllm import LLM, SamplingParams

9 from vllm.assets.image import ImageAsset

10 from vllm.exceptions import VLLMValidationError

@@ -14,17 +15,19 @@ TEXT_ONLY_PROMPT = "USER: What is 2 + 2?\nASSISTANT:"

17 @pytest.fixture(scope="module")

17 -def llm():

18 +def llm(vllm_runner):

19 """LLM with enable_mm_embeds=True and all modality limits zeroed out."""

19 - with managed_llm(

20 - model=MODEL,

20 + with vllm_runner(

21 + MODEL,

22 max_model_len=2048,

23 enforce_eager=True,

24 gpu_memory_utilization=0.8,

25 enable_mm_embeds=True,

26 limit_mm_per_prompt={"image": 0},

26 - ) as llm:

27 - yield llm

27 + ) as runner:

28 + # pytest caches yielded fixtures until after teardown, so use a proxy to

29 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

30 + yield weakref.proxy(runner.llm)

33 @pytest.mark.skip_global_cleanup

tests/entrypoints/pooling/basic/test_encode.py

+9 -13

@@ -6,7 +6,6 @@ import weakref

6 import pytest

8 from vllm import LLM, PoolingParams

9 -from vllm.distributed import cleanup_dist_env_and_memory

9 from vllm.exceptions import VLLMValidationError

11 MODEL_NAME = "intfloat/multilingual-e5-small"

@@ -29,23 +28,20 @@ TOKEN_IDS = [

30 @pytest.fixture(scope="module")

32 -def llm():

33 - # pytest caches the fixture so we use weakref.proxy to

34 - # enable garbage collection

35 - llm = LLM(

36 - model=MODEL_NAME,

31 +def llm(vllm_runner):

32 + with vllm_runner(

33 + MODEL_NAME,

34 + max_model_len=None,

35 max_num_batched_tokens=32768,

36 tensor_parallel_size=1,

37 gpu_memory_utilization=0.75,

38 enforce_eager=True,

39 seed=0,

42 - )

43 -

44 - yield weakref.proxy(llm)

45 -

46 - del llm

47 -

48 - cleanup_dist_env_and_memory()

40 + enable_chunked_prefill=None,

41 + ) as runner:

42 + # pytest caches yielded fixtures until after teardown, so use a proxy to

43 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

44 + yield weakref.proxy(runner.llm)

47 @pytest.mark.skip_global_cleanup

tests/entrypoints/pooling/basic/test_tiling_engine.py

+10 -12

@@ -6,28 +6,26 @@ from unittest import mock

7 import pytest

9 -from vllm import LLM, PoolingParams

10 -from vllm.distributed import cleanup_dist_env_and_memory

9 +from vllm import PoolingParams

11 MODEL_NAME = "intfloat/multilingual-e5-small"

14 @pytest.fixture(scope="module")

16 -def llm():

17 - llm = LLM(

18 - model=MODEL_NAME,

15 +def llm(vllm_runner):

16 + with vllm_runner(

17 + MODEL_NAME,

18 + max_model_len=None,

19 max_num_seqs=2, # small to trigger tiling

20 tensor_parallel_size=1,

21 gpu_memory_utilization=0.75,

22 enforce_eager=True,

23 seed=0,

24 - )

25 -

26 - yield weakref.proxy(llm)

27 -

28 - del llm

29 -

30 - cleanup_dist_env_and_memory()

24 + enable_chunked_prefill=None,

25 + ) as runner:

26 + # pytest caches yielded fixtures until after teardown, so use a proxy to

27 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

28 + yield weakref.proxy(runner.llm)

31 @pytest.mark.skip_global_cleanup

tests/entrypoints/pooling/classify/test_offline.py

+9 -13

@@ -7,7 +7,6 @@ import torch

8 from tests.models.utils import softmax

9 from vllm import LLM, ClassificationRequestOutput, PoolingParams

10 -from vllm.distributed import cleanup_dist_env_and_memory

10 from vllm.tasks import PoolingTask

12 MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"

@@ -18,23 +17,20 @@ num_labels = 2

19 @pytest.fixture(scope="module")

21 -def llm():

22 - # pytest caches the fixture so we use weakref.proxy to

23 - # enable garbage collection

24 - llm = LLM(

25 - model=MODEL_NAME,

20 +def llm(vllm_runner):

21 + with vllm_runner(

22 + MODEL_NAME,

23 + max_model_len=None,

24 max_num_batched_tokens=32768,

25 tensor_parallel_size=1,

26 gpu_memory_utilization=0.75,

27 enforce_eager=True,

28 seed=0,

31 - )

32 -

33 - yield weakref.proxy(llm)

34 -

35 - del llm

36 -

37 - cleanup_dist_env_and_memory()

29 + enable_chunked_prefill=None,

30 + ) as runner:

31 + # pytest caches yielded fixtures until after teardown, so use a proxy to

32 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

33 + yield weakref.proxy(runner.llm)

36 @pytest.mark.skip_global_cleanup

tests/entrypoints/pooling/embed/test_offline.py

+10 -13

@@ -7,7 +7,6 @@ import torch

7 import torch.nn.functional as F

9 from vllm import LLM, EmbeddingRequestOutput, PoolingParams

10 -from vllm.distributed import cleanup_dist_env_and_memory

10 from vllm.tasks import PoolingTask

12 MODEL_NAME = "intfloat/multilingual-e5-small"

@@ -18,23 +17,21 @@ embedding_size = 384

19 @pytest.fixture(scope="module")

21 -def llm():

22 - # pytest caches the fixture so we use weakref.proxy to

23 - # enable garbage collection

24 - llm = LLM(

25 - model=MODEL_NAME,

20 +def llm(vllm_runner):

21 + with vllm_runner(

22 + MODEL_NAME,

23 + max_model_len=None,

24 max_num_batched_tokens=32768,

25 tensor_parallel_size=1,

26 gpu_memory_utilization=0.75,

27 enforce_eager=True,

28 seed=0,

31 - )

32 - assert embedding_size == llm.model_config.embedding_size

33 -

34 - yield weakref.proxy(llm)

35 -

36 - del llm

37 - cleanup_dist_env_and_memory()

29 + enable_chunked_prefill=None,

30 + ) as runner:

31 + assert embedding_size == runner.llm.model_config.embedding_size

32 + # pytest caches yielded fixtures until after teardown, so use a proxy to

33 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

34 + yield weakref.proxy(runner.llm)

37 @pytest.mark.skip_global_cleanup

tests/entrypoints/pooling/reward/test_token_reward_offline.py

+9 -13

@@ -8,7 +8,6 @@ import torch

9 from tests.models.utils import softmax

10 from vllm import LLM, PoolingParams

11 -from vllm.distributed import cleanup_dist_env_and_memory

12 MODEL_NAME = "internlm/internlm2-1_8b-reward"

@@ -16,24 +15,21 @@ prompts = ["The chef prepared a delicious meal."]

17 @pytest.fixture(scope="module")

19 -def llm():

20 - # pytest caches the fixture so we use weakref.proxy to

21 - # enable garbage collection

22 - llm = LLM(

23 - model=MODEL_NAME,

18 +def llm(vllm_runner):

19 + with vllm_runner(

20 + MODEL_NAME,

21 + max_model_len=None,

22 max_num_batched_tokens=32768,

23 tensor_parallel_size=1,

24 gpu_memory_utilization=0.75,

25 enforce_eager=True,

26 trust_remote_code=True,

27 seed=0,

30 - )

31 -

32 - yield weakref.proxy(llm)

33 -

34 - del llm

35 -

36 - cleanup_dist_env_and_memory()

28 + enable_chunked_prefill=None,

29 + ) as runner:

30 + # pytest caches yielded fixtures until after teardown, so use a proxy to

31 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

32 + yield weakref.proxy(runner.llm)

35 @pytest.mark.skip_global_cleanup

tests/entrypoints/pooling/scoring/test_bi_encoder_offline.py

+9 -14

@@ -6,8 +6,6 @@ import weakref

6 import pytest

8 from tests.entrypoints.pooling.scoring.util import EncoderScoringHfRunner

9 -from vllm import LLM

10 -from vllm.distributed import cleanup_dist_env_and_memory

10 MODEL_NAME = "intfloat/multilingual-e5-small"

11 PROMPT = "The chef prepared a delicious meal."

@@ -27,23 +25,20 @@ DTYPE = "half"

27 @pytest.fixture(scope="module")

30 -def llm():

31 - # pytest caches the fixture so we use weakref.proxy to

32 - # enable garbage collection

33 - llm = LLM(

34 - model=MODEL_NAME,

28 +def llm(vllm_runner):

29 + with vllm_runner(

30 + MODEL_NAME,

31 + max_model_len=None,

32 max_num_batched_tokens=32768,

33 tensor_parallel_size=1,

34 gpu_memory_utilization=0.75,

35 enforce_eager=True,

36 seed=0,

40 - )

41 -

42 - yield weakref.proxy(llm)

43 -

44 - del llm

45 -

46 - cleanup_dist_env_and_memory()

37 + enable_chunked_prefill=None,

38 + ) as runner:

39 + # pytest caches yielded fixtures until after teardown, so use a proxy to

40 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

41 + yield weakref.proxy(runner.llm)

44 @pytest.fixture(scope="module")

tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py

+9 -13

@@ -8,7 +8,6 @@ import torch

9 from tests.models.utils import softmax

10 from vllm import LLM, PoolingParams

11 -from vllm.distributed import cleanup_dist_env_and_memory

12 MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"

13 PROMPT = "The chef prepared a delicious meal."

@@ -24,23 +23,20 @@ TEXTS_2 = [

25 @pytest.fixture(scope="module")

27 -def llm():

28 - # pytest caches the fixture so we use weakref.proxy to

29 - # enable garbage collection

30 - llm = LLM(

31 - model=MODEL_NAME,

26 +def llm(vllm_runner):

27 + with vllm_runner(

28 + MODEL_NAME,

29 + max_model_len=None,

30 max_num_batched_tokens=32768,

31 tensor_parallel_size=1,

32 gpu_memory_utilization=0.75,

33 enforce_eager=True,

34 seed=0,

37 - )

38 -

39 - yield weakref.proxy(llm)

40 -

41 - del llm

42 -

43 - cleanup_dist_env_and_memory()

35 + enable_chunked_prefill=None,

36 + ) as runner:

37 + # pytest caches yielded fixtures until after teardown, so use a proxy to

38 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

39 + yield weakref.proxy(runner.llm)

42 @pytest.fixture(scope="module")

tests/entrypoints/pooling/scoring/test_late_interaction_offline.py

+9 -15

@@ -5,9 +5,6 @@ import weakref

6 import pytest

8 -from vllm import LLM

9 -from vllm.distributed import cleanup_dist_env_and_memory

10 -

8 from .util import ColBERTScoringHfRunner

10 MODEL_NAME = "answerdotai/answerai-colbert-small-v1"

@@ -30,23 +27,20 @@ DTYPE = "half"

29 @pytest.fixture(scope="module")

33 -def llm():

34 - # pytest caches the fixture so we use weakref.proxy to

35 - # enable garbage collection

36 - llm = LLM(

37 - model=MODEL_NAME,

30 +def llm(vllm_runner):

31 + with vllm_runner(

32 + MODEL_NAME,

33 + max_model_len=None,

34 max_num_batched_tokens=32768,

35 tensor_parallel_size=1,

36 gpu_memory_utilization=0.75,

37 enforce_eager=True,

38 seed=0,

43 - )

44 -

45 - yield weakref.proxy(llm)

46 -

47 - del llm

48 -

49 - cleanup_dist_env_and_memory()

39 + enable_chunked_prefill=None,

40 + ) as runner:

41 + # pytest caches yielded fixtures until after teardown, so use a proxy to

42 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

43 + yield weakref.proxy(runner.llm)

46 @pytest.fixture(scope="module")

tests/entrypoints/pooling/scoring/test_late_interaction_offline_vision.py

+9 -15

@@ -5,32 +5,26 @@ import weakref

6 import pytest

8 -from vllm import LLM

9 -from vllm.distributed import cleanup_dist_env_and_memory

10 -

8 from .util import make_base64_image, make_image_mm_param

10 MODEL_NAME = "vidore/colpali-v1.3-hf"

13 @pytest.fixture(scope="module")

17 -def llm():

18 - # pytest caches the fixture so we use weakref.proxy to

19 - # enable garbage collection

20 - llm = LLM(

21 - model=MODEL_NAME,

14 +def llm(vllm_runner):

15 + with vllm_runner(

16 + MODEL_NAME,

17 + max_model_len=None,

18 max_num_batched_tokens=32768,

19 tensor_parallel_size=1,

20 gpu_memory_utilization=0.75,

21 enforce_eager=True,

22 seed=0,

27 - )

28 -

29 - yield weakref.proxy(llm)

30 -

31 - del llm

32 -

33 - cleanup_dist_env_and_memory()

23 + enable_chunked_prefill=None,

24 + ) as runner:

25 + # pytest caches yielded fixtures until after teardown, so use a proxy to

26 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

27 + yield weakref.proxy(runner.llm)

30 @pytest.mark.skip_global_cleanup

tests/entrypoints/pooling/token_classify/test_offline.py

+9 -13

@@ -6,7 +6,6 @@ import pytest

7 from vllm import LLM, PoolingRequestOutput

8 from vllm.config import PoolerConfig

9 -from vllm.distributed import cleanup_dist_env_and_memory

9 from vllm.tasks import PoolingTask

11 MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"

@@ -17,24 +16,21 @@ num_labels = 2

18 @pytest.fixture(scope="module")

20 -def llm():

21 - # pytest caches the fixture so we use weakref.proxy to

22 - # enable garbage collection

23 - llm = LLM(

24 - model=MODEL_NAME,

19 +def llm(vllm_runner):

20 + with vllm_runner(

21 + MODEL_NAME,

22 + max_model_len=None,

23 pooler_config=PoolerConfig(task="token_classify"),

24 max_num_batched_tokens=32768,

25 tensor_parallel_size=1,

26 gpu_memory_utilization=0.75,

27 enforce_eager=True,

28 seed=0,

31 - )

32 -

33 - yield weakref.proxy(llm)

34 -

35 - del llm

36 -

37 - cleanup_dist_env_and_memory()

29 + enable_chunked_prefill=None,

30 + ) as runner:

31 + # pytest caches yielded fixtures until after teardown, so use a proxy to

32 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

33 + yield weakref.proxy(runner.llm)

36 @pytest.mark.skip_global_cleanup

tests/entrypoints/pooling/token_embed/test_offline.py

+10 -13

@@ -6,7 +6,6 @@ import pytest

7 from vllm import LLM, PoolingRequestOutput

8 from vllm.config import PoolerConfig

9 -from vllm.distributed import cleanup_dist_env_and_memory

9 from vllm.tasks import PoolingTask

11 MODEL_NAME = "intfloat/multilingual-e5-small"

@@ -17,24 +16,22 @@ embedding_size = 384

18 @pytest.fixture(scope="module")

20 -def llm():

21 - # pytest caches the fixture so we use weakref.proxy to

22 - # enable garbage collection

23 - llm = LLM(

24 - model=MODEL_NAME,

19 +def llm(vllm_runner):

20 + with vllm_runner(

21 + MODEL_NAME,

22 + max_model_len=None,

23 pooler_config=PoolerConfig(task="token_embed"),

24 max_num_batched_tokens=32768,

25 tensor_parallel_size=1,

26 gpu_memory_utilization=0.75,

27 enforce_eager=True,

28 seed=0,

31 - )

32 - assert embedding_size == llm.model_config.embedding_size

33 -

34 - yield weakref.proxy(llm)

35 -

36 - del llm

37 - cleanup_dist_env_and_memory()

29 + enable_chunked_prefill=None,

30 + ) as runner:

31 + assert embedding_size == runner.llm.model_config.embedding_size

32 + # pytest caches yielded fixtures until after teardown, so use a proxy to

33 + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.

34 + yield weakref.proxy(runner.llm)

37 @pytest.mark.skip_global_cleanup

tests/entrypoints/weight_transfer/test_weight_transfer_llm.py

+47 -40

@@ -8,13 +8,13 @@ actual NCCL communication.

8 """

10 import os

11 +import weakref

12 from dataclasses import dataclass

13 from unittest.mock import patch

15 import pytest

16 import torch

17 -from vllm import LLM

18 from vllm.config import WeightTransferConfig

19 from vllm.distributed.weight_transfer.base import (

20 WeightTransferEngine,

@@ -102,25 +102,24 @@ def mock_create_engine(config, vllm_config, device, model):

102

103

104 @create_new_process_for_each_test()

105 -def test_get_world_size_tp1():

105 +def test_get_world_size_tp1(vllm_runner):

106 """Test world_size is correctly configured for TP=1."""

107 if torch.accelerator.device_count() < 1:

108 pytest.skip("Need at least 1 GPU for this test")

109

110 - llm = LLM(

111 - model=MODEL_NAME,

110 + with vllm_runner(

111 + MODEL_NAME,

112 enforce_eager=True,

113 load_format="dummy",

114 tensor_parallel_size=1,

115 weight_transfer_config=WeightTransferConfig(backend="nccl"),

116 - )

117 -

118 - world_size = llm.llm_engine.vllm_config.parallel_config.world_size

119 - assert world_size == 1

116 + ) as runner:

117 + world_size = runner.llm.llm_engine.vllm_config.parallel_config.world_size

118 + assert world_size == 1

119

120

121 @create_new_process_for_each_test()

123 -def test_init_weight_transfer_engine_calls_engine():

122 +def test_init_weight_transfer_engine_calls_engine(vllm_runner):

123 """Test that init_weight_transfer_engine calls the engine's

124 init_transfer_engine method."""

125 if torch.accelerator.device_count() < 1:

@@ -131,17 +130,20 @@ def test_init_weight_transfer_engine_calls_engine():

130 # Enable insecure serialization to allow pickling functions for collective_rpc

131 os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"

132

134 - with patch(

135 - "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",

136 - mock_create_engine,

137 - ):

138 - llm = LLM(

139 - model=MODEL_NAME,

133 + with (

134 + patch(

135 + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",

136 + mock_create_engine,

137 + ),

138 + vllm_runner(

139 + MODEL_NAME,

140 enforce_eager=True,

141 load_format="dummy",

142 tensor_parallel_size=1,

143 weight_transfer_config=WeightTransferConfig(backend="nccl"),

144 - )

144 + ) as runner,

145 + ):

146 + llm = weakref.proxy(runner.llm)

147

148 # Verify engine was created

149 def check_engine_exists(self):

@@ -170,7 +172,7 @@ def test_init_weight_transfer_engine_calls_engine():

172

173

174 @create_new_process_for_each_test()

173 -def test_update_weights_calls_engine():

175 +def test_update_weights_calls_engine(vllm_runner):

176 """Test that update_weights calls the engine's receive_weights method."""

177 if torch.accelerator.device_count() < 1:

178 pytest.skip("Need at least 1 GPU for this test")

@@ -180,17 +182,20 @@ def test_update_weights_calls_engine():

182 # Enable insecure serialization to allow pickling functions for collective_rpc

183 os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"

184

183 - with patch(

184 - "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",

185 - mock_create_engine,

186 - ):

187 - llm = LLM(

188 - model=MODEL_NAME,

185 + with (

186 + patch(

187 + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",

188 + mock_create_engine,

189 + ),

190 + vllm_runner(

191 + MODEL_NAME,

192 enforce_eager=True,

193 load_format="dummy",

194 tensor_parallel_size=1,

195 weight_transfer_config=WeightTransferConfig(backend="nccl"),

193 - )

196 + ) as runner,

197 + ):

198 + llm = weakref.proxy(runner.llm)

199

200 # First init the weight transfer

201 llm.init_weight_transfer_engine(

@@ -233,7 +238,7 @@ def test_update_weights_calls_engine():

238

239

240 @create_new_process_for_each_test()

236 -def test_full_weight_transfer_flow():

241 +def test_full_weight_transfer_flow(vllm_runner):

242 """Test the complete weight transfer flow: init -> start -> update -> finish."""

243 if torch.accelerator.device_count() < 1:

244 pytest.skip("Need at least 1 GPU for this test")

@@ -243,17 +248,20 @@ def test_full_weight_transfer_flow():

248 # Enable insecure serialization to allow pickling functions for collective_rpc

249 os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"

250

246 - with patch(

247 - "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",

248 - mock_create_engine,

249 - ):

250 - llm = LLM(

251 - model=MODEL_NAME,

251 + with (

252 + patch(

253 + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",

254 + mock_create_engine,

255 + ),

256 + vllm_runner(

257 + MODEL_NAME,

258 enforce_eager=True,

259 load_format="dummy",

260 tensor_parallel_size=1,

261 weight_transfer_config=WeightTransferConfig(backend="nccl"),

256 - )

262 + ) as runner,

263 + ):

264 + llm = weakref.proxy(runner.llm)

265

266 assert llm.get_weight_version() == "default"

267

@@ -309,20 +317,19 @@ def test_full_weight_transfer_flow():

317

318

319 @create_new_process_for_each_test()

312 -def test_weight_transfer_config_backend():

320 +def test_weight_transfer_config_backend(vllm_runner):

321 """Test that WeightTransferConfig backend is properly configured."""

322 if torch.accelerator.device_count() < 1:

323 pytest.skip("Need at least 1 GPU for this test")

324

325 # Test with nccl backend

318 - llm = LLM(

319 - model=MODEL_NAME,

326 + with vllm_runner(

327 + MODEL_NAME,

328 enforce_eager=True,

329 load_format="dummy",

330 tensor_parallel_size=1,

331 weight_transfer_config=WeightTransferConfig(backend="nccl"),

324 - )

325 -

326 - config = llm.llm_engine.vllm_config.weight_transfer_config

327 - assert config is not None

328 - assert config.backend == "nccl"

332 + ) as runner:

333 + config = runner.llm.llm_engine.vllm_config.weight_transfer_config

334 + assert config is not None

335 + assert config.backend == "nccl"