| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from vllm import SamplingParams |
| 5 | |
| 6 | |
| 7 | def test_gpu_memory_utilization(vllm_runner): |
| 8 | prompts = [ |
| 9 | "Hello, my name is", |
| 10 | "The president of the United States is", |
| 11 | "The capital of France is", |
| 12 | "The future of AI is", |
| 13 | ] |
| 14 | sampling_params = SamplingParams(temperature=0.8, top_p=0.95) |
| 15 | |
| 16 | # makes sure gpu_memory_utilization is per-instance limit, |
| 17 | # not a global limit |
| 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}") |