[CI/Test] improve robustness of test (vllm_runner) (#5357)

[CI/Test] improve robustness of test by replacing del with context manager (vllm_runner) (#5357)

tests/basic_correctness/test_basic_correctness.py

@@ -46,12 +46,11 @@ def test_models(
     with hf_runner(model, dtype=dtype) as hf_model:
         hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
 
-    vllm_model = vllm_runner(model,
-                             dtype=dtype,
-                             enforce_eager=enforce_eager,
-                             gpu_memory_utilization=0.7)
-    vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
-    del vllm_model
+    with vllm_runner(model,
+                     dtype=dtype,
+                     enforce_eager=enforce_eager,
+                     gpu_memory_utilization=0.7) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
 
     for i in range(len(example_prompts)):
         hf_output_ids, hf_output_str = hf_outputs[i]

tests/basic_correctness/test_chunked_prefill.py

@@ -43,17 +43,16 @@ def test_models(
     with hf_runner(model, dtype=dtype) as hf_model:
         hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
 
-    vllm_model = vllm_runner(
-        model,
-        dtype=dtype,
-        max_num_batched_tokens=max_num_batched_tokens,
-        enable_chunked_prefill=enable_chunked_prefill,
-        tensor_parallel_size=tensor_parallel_size,
-        enforce_eager=enforce_eager,
-        max_num_seqs=max_num_seqs,
-    )
-    vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
-    del vllm_model
+    with vllm_runner(
+            model,
+            dtype=dtype,
+            max_num_batched_tokens=max_num_batched_tokens,
+            enable_chunked_prefill=enable_chunked_prefill,
+            tensor_parallel_size=tensor_parallel_size,
+            enforce_eager=enforce_eager,
+            max_num_seqs=max_num_seqs,
+    ) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
 
     for i in range(len(example_prompts)):
         hf_output_ids, hf_output_str = hf_outputs[i]

tests/basic_correctness/test_preemption.py

@@ -46,17 +46,16 @@ def test_chunked_prefill_recompute(
     with hf_runner(model, dtype=dtype) as hf_model:
         hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
 
-    vllm_model = vllm_runner(
-        model,
-        dtype=dtype,
-        max_num_batched_tokens=max_num_batched_tokens,
-        enable_chunked_prefill=enable_chunked_prefill,
-        max_num_seqs=max_num_seqs,
-    )
-    vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
-    assert (vllm_model.model.llm_engine.scheduler.artificial_preempt_cnt <
-            ARTIFICIAL_PREEMPTION_MAX_CNT)
-    del vllm_model
+    with vllm_runner(
+            model,
+            dtype=dtype,
+            max_num_batched_tokens=max_num_batched_tokens,
+            enable_chunked_prefill=enable_chunked_prefill,
+            max_num_seqs=max_num_seqs,
+    ) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
+        assert (vllm_model.model.llm_engine.scheduler.artificial_preempt_cnt <
+                ARTIFICIAL_PREEMPTION_MAX_CNT)
 
     for i in range(len(example_prompts)):
         hf_output_ids, hf_output_str = hf_outputs[i]
@@ -84,17 +83,16 @@ def test_preemption(
     with hf_runner(model, dtype=dtype) as hf_model:
         hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
 
-    vllm_model = vllm_runner(
-        model,
-        dtype=dtype,
-        disable_log_stats=False,
-    )
-    vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
-    assert (vllm_model.model.llm_engine.scheduler.artificial_preempt_cnt <
-            ARTIFICIAL_PREEMPTION_MAX_CNT)
-    total_preemption = (
-        vllm_model.model.llm_engine.scheduler.num_cumulative_preemption)
-    del vllm_model
+    with vllm_runner(
+            model,
+            dtype=dtype,
+            disable_log_stats=False,
+    ) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
+        assert (vllm_model.model.llm_engine.scheduler.artificial_preempt_cnt <
+                ARTIFICIAL_PREEMPTION_MAX_CNT)
+        total_preemption = (
+            vllm_model.model.llm_engine.scheduler.num_cumulative_preemption)
 
     for i in range(len(example_prompts)):
         hf_output_ids, hf_output_str = hf_outputs[i]
@@ -139,19 +137,18 @@ def test_swap(
         hf_outputs = hf_model.generate_beam_search(example_prompts, beam_width,
                                                    max_tokens)
 
-    vllm_model = vllm_runner(
-        model,
-        dtype=dtype,
-        swap_space=10,
-        disable_log_stats=False,
-    )
-    vllm_outputs = vllm_model.generate_beam_search(example_prompts, beam_width,
-                                                   max_tokens)
-    assert (vllm_model.model.llm_engine.scheduler.artificial_preempt_cnt <
-            ARTIFICIAL_PREEMPTION_MAX_CNT)
-    total_preemption = (
-        vllm_model.model.llm_engine.scheduler.num_cumulative_preemption)
-    del vllm_model
+    with vllm_runner(
+            model,
+            dtype=dtype,
+            swap_space=10,
+            disable_log_stats=False,
+    ) as vllm_model:
+        vllm_outputs = vllm_model.generate_beam_search(example_prompts,
+                                                       beam_width, max_tokens)
+        assert (vllm_model.model.llm_engine.scheduler.artificial_preempt_cnt <
+                ARTIFICIAL_PREEMPTION_MAX_CNT)
+        total_preemption = (
+            vllm_model.model.llm_engine.scheduler.num_cumulative_preemption)
 
     for i in range(len(example_prompts)):
         hf_output_ids, _ = hf_outputs[i]
@@ -196,28 +193,28 @@ def test_swap_infeasible(
     decode_blocks = max_tokens // BLOCK_SIZE
     example_prompts = example_prompts[:1]
 
-    vllm_model = vllm_runner(
-        model,
-        dtype=dtype,
-        swap_space=10,
-        block_size=BLOCK_SIZE,
-        # Since beam search have more than 1 sequence, prefill + decode blocks
-        # are not enough to finish.
-        num_gpu_blocks_override=prefill_blocks + decode_blocks,
-        max_model_len=(prefill_blocks + decode_blocks) * BLOCK_SIZE,
-    )
-    sampling_params = SamplingParams(n=beam_width,
-                                     use_beam_search=True,
-                                     temperature=0.0,
-                                     max_tokens=max_tokens,
-                                     ignore_eos=True)
-    req_outputs = vllm_model.model.generate(
-        example_prompts,
-        sampling_params=sampling_params,
-    )
-    assert (vllm_model.model.llm_engine.scheduler.artificial_preempt_cnt <
-            ARTIFICIAL_PREEMPTION_MAX_CNT)
-    del vllm_model
+    with vllm_runner(
+            model,
+            dtype=dtype,
+            swap_space=10,
+            block_size=BLOCK_SIZE,
+            # Since beam search have more than 1 sequence, prefill +
+            # decode blocks are not enough to finish.
+            num_gpu_blocks_override=prefill_blocks + decode_blocks,
+            max_model_len=(prefill_blocks + decode_blocks) * BLOCK_SIZE,
+    ) as vllm_model:
+        sampling_params = SamplingParams(n=beam_width,
+                                         use_beam_search=True,
+                                         temperature=0.0,
+                                         max_tokens=max_tokens,
+                                         ignore_eos=True)
+        req_outputs = vllm_model.model.generate(
+            example_prompts,
+            sampling_params=sampling_params,
+        )
+        assert (vllm_model.model.llm_engine.scheduler.artificial_preempt_cnt <
+                ARTIFICIAL_PREEMPTION_MAX_CNT)
+
     # Verify the request is ignored and not hang.
     assert req_outputs[0].outputs[0].finish_reason == "length"
 
@@ -236,25 +233,26 @@ def test_preemption_infeasible(
     BLOCK_SIZE = 16
     prefill_blocks = 2
     decode_blocks = max_tokens // BLOCK_SIZE
-    vllm_model = vllm_runner(
-        model,
-        dtype=dtype,
-        block_size=BLOCK_SIZE,
-        # Not enough gpu blocks to complete a single sequence.
-        # preemption should happen, and the sequence should be
-        # ignored instead of hanging forever.
-        num_gpu_blocks_override=prefill_blocks + decode_blocks // 2,
-        max_model_len=((prefill_blocks + decode_blocks // 2) * BLOCK_SIZE),
-    )
-    sampling_params = SamplingParams(max_tokens=max_tokens, ignore_eos=True)
-    req_outputs = vllm_model.model.generate(
-        example_prompts,
-        sampling_params=sampling_params,
-    )
+    with vllm_runner(
+            model,
+            dtype=dtype,
+            block_size=BLOCK_SIZE,
+            # Not enough gpu blocks to complete a single sequence.
+            # preemption should happen, and the sequence should be
+            # ignored instead of hanging forever.
+            num_gpu_blocks_override=prefill_blocks + decode_blocks // 2,
+            max_model_len=((prefill_blocks + decode_blocks // 2) * BLOCK_SIZE),
+    ) as vllm_model:
+        sampling_params = SamplingParams(max_tokens=max_tokens,
+                                         ignore_eos=True)
+        req_outputs = vllm_model.model.generate(
+            example_prompts,
+            sampling_params=sampling_params,
+        )
+
+        assert (vllm_model.model.llm_engine.scheduler.artificial_preempt_cnt <
+                ARTIFICIAL_PREEMPTION_MAX_CNT)
 
-    assert (vllm_model.model.llm_engine.scheduler.artificial_preempt_cnt <
-            ARTIFICIAL_PREEMPTION_MAX_CNT)
-    del vllm_model
     # Verify the request is ignored and not hang.
     for req_output in req_outputs:
         outputs = req_output.outputs

tests/conftest.py

@@ -493,7 +493,10 @@ class VllmRunner:
             outputs.append(embedding)
         return outputs
 
-    def __del__(self):
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
         del self.model
         cleanup()
 

tests/distributed/test_basic_distributed_correctness.py

@@ -45,14 +45,13 @@ def test_models(
     with hf_runner(model, dtype=dtype) as hf_model:
         hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
 
-    vllm_model = vllm_runner(
-        model,
-        dtype=dtype,
-        tensor_parallel_size=2,
-        enforce_eager=enforce_eager,
-        distributed_executor_backend=distributed_executor_backend)
-    vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
-    del vllm_model
+    with vllm_runner(model,
+                     dtype=dtype,
+                     tensor_parallel_size=2,
+                     enforce_eager=enforce_eager,
+                     distributed_executor_backend=distributed_executor_backend
+                     ) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
 
     for i in range(len(example_prompts)):
         hf_output_ids, hf_output_str = hf_outputs[i]

tests/distributed/test_chunked_prefill_distributed.py

@@ -48,17 +48,16 @@ def test_models(
     with hf_runner(model, dtype=dtype) as hf_model:
         hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
 
-    vllm_model = vllm_runner(
-        model,
-        dtype=dtype,
-        tensor_parallel_size=2,
-        max_num_seqs=max_num_seqs,
-        enable_chunked_prefill=enable_chunked_prefill,
-        max_num_batched_tokens=max_num_batched_tokens,
-        distributed_executor_backend=distributed_executor_backend,
-    )
-    vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
-    del vllm_model
+    with vllm_runner(
+            model,
+            dtype=dtype,
+            tensor_parallel_size=2,
+            max_num_seqs=max_num_seqs,
+            enable_chunked_prefill=enable_chunked_prefill,
+            max_num_batched_tokens=max_num_batched_tokens,
+            distributed_executor_backend=distributed_executor_backend,
+    ) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
 
     for i in range(len(example_prompts)):
         hf_output_ids, hf_output_str = hf_outputs[i]

tests/engine/test_stop_reason.py

@@ -19,9 +19,8 @@ MAX_TOKENS = 1024
 
 @pytest.fixture
 def vllm_model(vllm_runner):
-    vllm_model = vllm_runner(MODEL)
-    yield vllm_model
-    del vllm_model
+    with vllm_runner(MODEL) as vllm_model:
+        yield vllm_model
 
 
 def test_stop_reason(vllm_model, example_prompts):

tests/engine/test_stop_strings.py

@@ -10,7 +10,8 @@ MAX_TOKENS = 200
 
 @pytest.fixture(scope="session")
 def vllm_model(vllm_runner):
-    return vllm_runner(MODEL)
+    with vllm_runner(MODEL) as vllm_model:
+        yield vllm_model
 
 
 @pytest.mark.skip_global_cleanup

tests/metrics/test_metrics.py

@@ -23,23 +23,25 @@ def test_metric_counter_prompt_tokens(
     dtype: str,
     max_tokens: int,
 ) -> None:
-    vllm_model = vllm_runner(model,
-                             dtype=dtype,
-                             disable_log_stats=False,
-                             gpu_memory_utilization=0.4)
-    tokenizer = vllm_model.model.get_tokenizer()
-    prompt_token_counts = [len(tokenizer.encode(p)) for p in example_prompts]
-    # This test needs at least 2 prompts in a batch of different lengths to
-    # verify their token count is correct despite padding.
-    assert len(example_prompts) > 1, "at least 2 prompts are required"
-    assert prompt_token_counts[0] != prompt_token_counts[1], (
-        "prompts of different lengths are required")
-    vllm_prompt_token_count = sum(prompt_token_counts)
-
-    _ = vllm_model.generate_greedy(example_prompts, max_tokens)
-    stat_logger = vllm_model.model.llm_engine.stat_logger
-    metric_count = stat_logger.metrics.counter_prompt_tokens.labels(
-        **stat_logger.labels)._value.get()
+    with vllm_runner(model,
+                     dtype=dtype,
+                     disable_log_stats=False,
+                     gpu_memory_utilization=0.4) as vllm_model:
+        tokenizer = vllm_model.model.get_tokenizer()
+        prompt_token_counts = [
+            len(tokenizer.encode(p)) for p in example_prompts
+        ]
+        # This test needs at least 2 prompts in a batch of different lengths to
+        # verify their token count is correct despite padding.
+        assert len(example_prompts) > 1, "at least 2 prompts are required"
+        assert prompt_token_counts[0] != prompt_token_counts[1], (
+            "prompts of different lengths are required")
+        vllm_prompt_token_count = sum(prompt_token_counts)
+
+        _ = vllm_model.generate_greedy(example_prompts, max_tokens)
+        stat_logger = vllm_model.model.llm_engine.stat_logger
+        metric_count = stat_logger.metrics.counter_prompt_tokens.labels(
+            **stat_logger.labels)._value.get()
 
     assert vllm_prompt_token_count == metric_count, (
         f"prompt token count: {vllm_prompt_token_count!r}\n"
@@ -56,22 +58,22 @@ def test_metric_counter_generation_tokens(
     dtype: str,
     max_tokens: int,
 ) -> None:
-    vllm_model = vllm_runner(model,
-                             dtype=dtype,
-                             disable_log_stats=False,
-                             gpu_memory_utilization=0.4)
-    vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
-    tokenizer = vllm_model.model.get_tokenizer()
-    stat_logger = vllm_model.model.llm_engine.stat_logger
-    metric_count = stat_logger.metrics.counter_generation_tokens.labels(
-        **stat_logger.labels)._value.get()
-    vllm_generation_count = 0
-    for i in range(len(example_prompts)):
-        vllm_output_ids, vllm_output_str = vllm_outputs[i]
-        prompt_ids = tokenizer.encode(example_prompts[i])
-        # vllm_output_ids contains both prompt tokens and generation tokens.
-        # We're interested only in the count of the generation tokens.
-        vllm_generation_count += len(vllm_output_ids) - len(prompt_ids)
+    with vllm_runner(model,
+                     dtype=dtype,
+                     disable_log_stats=False,
+                     gpu_memory_utilization=0.4) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
+        tokenizer = vllm_model.model.get_tokenizer()
+        stat_logger = vllm_model.model.llm_engine.stat_logger
+        metric_count = stat_logger.metrics.counter_generation_tokens.labels(
+            **stat_logger.labels)._value.get()
+        vllm_generation_count = 0
+        for i in range(len(example_prompts)):
+            vllm_output_ids, vllm_output_str = vllm_outputs[i]
+            prompt_ids = tokenizer.encode(example_prompts[i])
+            # vllm_output_ids contains both prompt tokens and generation tokens.
+            # We're interested only in the count of the generation tokens.
+            vllm_generation_count += len(vllm_output_ids) - len(prompt_ids)
 
     assert vllm_generation_count == metric_count, (
         f"generation token count: {vllm_generation_count!r}\n"
@@ -85,15 +87,13 @@ def test_metric_counter_generation_tokens(
     [None, [], ["ModelName0"], ["ModelName0", "ModelName1", "ModelName2"]])
 def test_metric_set_tag_model_name(vllm_runner, model: str, dtype: str,
                                    served_model_name: List[str]) -> None:
-    vllm_model = vllm_runner(model,
-                             dtype=dtype,
-                             disable_log_stats=False,
-                             gpu_memory_utilization=0.3,
-                             served_model_name=served_model_name)
-    stat_logger = vllm_model.model.llm_engine.stat_logger
-    metrics_tag_content = stat_logger.labels["model_name"]
-
-    del vllm_model
+    with vllm_runner(model,
+                     dtype=dtype,
+                     disable_log_stats=False,
+                     gpu_memory_utilization=0.3,
+                     served_model_name=served_model_name) as vllm_model:
+        stat_logger = vllm_model.model.llm_engine.stat_logger
+        metrics_tag_content = stat_logger.labels["model_name"]
 
     if served_model_name is None or served_model_name == []:
         assert metrics_tag_content == model, (

tests/models/test_aqlm.py

@@ -82,10 +82,9 @@ def test_models(
     num_logprobs: int,
 ) -> None:
 
-    vllm_model = vllm_runner(model, dtype=dtype)
-    vllm_outputs = vllm_model.generate_greedy_logprobs(example_prompts,
-                                                       max_tokens,
-                                                       num_logprobs)
+    with vllm_runner(model, dtype=dtype) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy_logprobs(
+            example_prompts, max_tokens, num_logprobs)
 
     # loop through the prompts to compare against the ground truth generations
     for prompt_idx in range(len(example_prompts)):

tests/models/test_big_models.py

@@ -37,9 +37,8 @@ def test_models(
     with hf_runner(model, dtype=dtype) as hf_model:
         hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
 
-    vllm_model = vllm_runner(model, dtype=dtype)
-    vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
-    del vllm_model
+    with vllm_runner(model, dtype=dtype) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
 
     for i in range(len(example_prompts)):
         hf_output_ids, hf_output_str = hf_outputs[i]
@@ -57,9 +56,8 @@ def test_model_print(
     model: str,
     dtype: str,
 ) -> None:
-    vllm_model = vllm_runner(model, dtype=dtype)
-    # This test is for verifying whether the model's extra_repr
-    # can be printed correctly.
-    print(vllm_model.model.llm_engine.model_executor.driver_worker.
-          model_runner.model)
-    del vllm_model
+    with vllm_runner(model, dtype=dtype) as vllm_model:
+        # This test is for verifying whether the model's extra_repr
+        # can be printed correctly.
+        print(vllm_model.model.llm_engine.model_executor.driver_worker.
+              model_runner.model)

tests/models/test_embedding.py

@@ -31,9 +31,8 @@ def test_models(
     with hf_runner(model, dtype=dtype, is_embedding_model=True) as hf_model:
         hf_outputs = hf_model.encode(example_prompts)
 
-    vllm_model = vllm_runner(model, dtype=dtype)
-    vllm_outputs = vllm_model.encode(example_prompts)
-    del vllm_model
+    with vllm_runner(model, dtype=dtype) as vllm_model:
+        vllm_outputs = vllm_model.encode(example_prompts)
 
     similarities = compare_embeddings(hf_outputs, vllm_outputs)
     all_similarities = torch.stack(similarities)

tests/models/test_gptq_marlin.py

@@ -70,32 +70,29 @@ def test_models(
     model_name, revision = model
 
     # Run marlin.
-    gptq_marlin_model = vllm_runner(model_name=model_name,
-                                    revision=revision,
-                                    dtype=dtype,
-                                    quantization="marlin",
-                                    max_model_len=MAX_MODEL_LEN,
-                                    tensor_parallel_size=1)
-
-    gptq_marlin_outputs = gptq_marlin_model.generate_greedy_logprobs(
-        example_prompts[:-1], max_tokens, num_logprobs)
-    del gptq_marlin_model
+    with vllm_runner(model_name=model_name,
+                     revision=revision,
+                     dtype=dtype,
+                     quantization="marlin",
+                     max_model_len=MAX_MODEL_LEN,
+                     tensor_parallel_size=1) as gptq_marlin_model:
+
+        gptq_marlin_outputs = gptq_marlin_model.generate_greedy_logprobs(
+            example_prompts[:-1], max_tokens, num_logprobs)
     _ROPE_DICT.clear()  # clear rope cache to avoid rope dtype error
 
     # Run gptq.
     # The naive gptq kernel doesn't support bf16 yet.
     # Here we always compare fp16/bf16 gpt marlin kernel
     # to fp16 gptq kernel.
-    gptq_model = vllm_runner(model_name=model_name,
-                             revision=revision,
-                             dtype="half",
-                             quantization="gptq",
-                             max_model_len=MAX_MODEL_LEN,
-                             tensor_parallel_size=1)
-    gptq_outputs = gptq_model.generate_greedy_logprobs(example_prompts[:-1],
-                                                       max_tokens,
-                                                       num_logprobs)
-    del gptq_model
+    with vllm_runner(model_name=model_name,
+                     revision=revision,
+                     dtype="half",
+                     quantization="gptq",
+                     max_model_len=MAX_MODEL_LEN,
+                     tensor_parallel_size=1) as gptq_model:
+        gptq_outputs = gptq_model.generate_greedy_logprobs(
+            example_prompts[:-1], max_tokens, num_logprobs)
 
     check_logprobs_close(
         outputs_0_lst=gptq_outputs,

tests/models/test_gptq_marlin_24.py

@@ -61,20 +61,16 @@ def test_models(
     max_tokens: int,
     num_logprobs: int,
 ) -> None:
-    marlin_24_model = vllm_runner(model_pair.model_marlin,
-                                  dtype=dtype,
-                                  quantization="gptq_marlin_24")
-    marlin_24_outputs = marlin_24_model.generate_greedy_logprobs(
-        example_prompts, max_tokens, num_logprobs)
-    del marlin_24_model
+    with vllm_runner(model_pair.model_marlin,
+                     dtype=dtype,
+                     quantization="gptq_marlin_24") as marlin_24_model:
+        marlin_24_outputs = marlin_24_model.generate_greedy_logprobs(
+            example_prompts, max_tokens, num_logprobs)
 
-    gptq_model = vllm_runner(model_pair.model_gptq,
-                             dtype=dtype,
-                             quantization="gptq")
-    gptq_outputs = gptq_model.generate_greedy_logprobs(example_prompts,
-                                                       max_tokens,
-                                                       num_logprobs)
-    del gptq_model
+    with vllm_runner(model_pair.model_gptq, dtype=dtype,
+                     quantization="gptq") as gptq_model:
+        gptq_outputs = gptq_model.generate_greedy_logprobs(
+            example_prompts, max_tokens, num_logprobs)
 
     check_logprobs_close(
         outputs_0_lst=gptq_outputs,

tests/models/test_llava.py

@@ -94,14 +94,13 @@ def test_models(hf_runner, vllm_runner, hf_images, vllm_images,
         for p in HF_IMAGE_PROMPTS
     ]
 
-    vllm_model = vllm_runner(model_id,
-                             dtype=dtype,
-                             enforce_eager=True,
-                             **vlm_config.as_cli_args_dict())
-    vllm_outputs = vllm_model.generate_greedy(vllm_image_prompts,
-                                              max_tokens,
-                                              images=vllm_images)
-    del vllm_model
+    with vllm_runner(model_id,
+                     dtype=dtype,
+                     enforce_eager=True,
+                     **vlm_config.as_cli_args_dict()) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy(vllm_image_prompts,
+                                                  max_tokens,
+                                                  images=vllm_images)
 
     for i in range(len(HF_IMAGE_PROMPTS)):
         hf_output_ids, hf_output_str = hf_outputs[i]

tests/models/test_marlin.py

@@ -59,20 +59,16 @@ def test_models(
     max_tokens: int,
     num_logprobs: int,
 ) -> None:
-    marlin_model = vllm_runner(model_pair.model_marlin,
-                               dtype=dtype,
-                               quantization="marlin")
-    marlin_outputs = marlin_model.generate_greedy_logprobs(
-        example_prompts, max_tokens, num_logprobs)
-    del marlin_model
-
-    gptq_model = vllm_runner(model_pair.model_gptq,
-                             dtype=dtype,
-                             quantization="gptq")
-    gptq_outputs = gptq_model.generate_greedy_logprobs(example_prompts,
-                                                       max_tokens,
-                                                       num_logprobs)
-    del gptq_model
+    with vllm_runner(model_pair.model_marlin,
+                     dtype=dtype,
+                     quantization="marlin") as marlin_model:
+        marlin_outputs = marlin_model.generate_greedy_logprobs(
+            example_prompts, max_tokens, num_logprobs)
+
+    with vllm_runner(model_pair.model_gptq, dtype=dtype,
+                     quantization="gptq") as gptq_model:
+        gptq_outputs = gptq_model.generate_greedy_logprobs(
+            example_prompts, max_tokens, num_logprobs)
 
     check_logprobs_close(
         outputs_0_lst=gptq_outputs,

tests/models/test_mistral.py

@@ -30,11 +30,9 @@ def test_models(
         hf_outputs = hf_model.generate_greedy_logprobs_limit(
             example_prompts, max_tokens, num_logprobs)
 
-    vllm_model = vllm_runner(model, dtype=dtype)
-    vllm_outputs = vllm_model.generate_greedy_logprobs(example_prompts,
-                                                       max_tokens,
-                                                       num_logprobs)
-    del vllm_model
+    with vllm_runner(model, dtype=dtype) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy_logprobs(
+            example_prompts, max_tokens, num_logprobs)
     check_logprobs_close(
         outputs_0_lst=hf_outputs,
         outputs_1_lst=vllm_outputs,

tests/models/test_models.py

@@ -37,9 +37,8 @@ def test_models(
     with hf_runner(model, dtype=dtype) as hf_model:
         hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
 
-    vllm_model = vllm_runner(model, dtype=dtype)
-    vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
-    del vllm_model
+    with vllm_runner(model, dtype=dtype) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
 
     for i in range(len(example_prompts)):
         hf_output_ids, hf_output_str = hf_outputs[i]
@@ -57,9 +56,8 @@ def test_model_print(
     model: str,
     dtype: str,
 ) -> None:
-    vllm_model = vllm_runner(model, dtype=dtype)
-    # This test is for verifying whether the model's extra_repr
-    # can be printed correctly.
-    print(vllm_model.model.llm_engine.model_executor.driver_worker.
-          model_runner.model)
-    del vllm_model
+    with vllm_runner(model, dtype=dtype) as vllm_model:
+        # This test is for verifying whether the model's extra_repr
+        # can be printed correctly.
+        print(vllm_model.model.llm_engine.model_executor.driver_worker.
+              model_runner.model)

tests/quantization/test_bitsandbytes.py

@@ -16,65 +16,65 @@ capability = capability[0] * 10 + capability[1]
     capability < QUANTIZATION_METHODS['bitsandbytes'].get_min_capability(),
     reason='bitsandbytes is not supported on this GPU type.')
 def test_load_bnb_model(vllm_runner) -> None:
-    llm = vllm_runner('huggyllama/llama-7b',
-                      quantization='bitsandbytes',
-                      load_format='bitsandbytes',
-                      enforce_eager=True)
-
-    model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model
-
-    # check the weights in MLP & SelfAttention are quantized to torch.uint8
-    qweight = model.model.layers[0].mlp.gate_up_proj.qweight
-    assert qweight.dtype == torch.uint8, (
-        f'Expected gate_up_proj dtype torch.uint8 but got {qweight.dtype}')
-
-    qweight = model.model.layers[0].mlp.down_proj.qweight
-    assert qweight.dtype == torch.uint8, (
-        f'Expected down_proj dtype torch.uint8 but got {qweight.dtype}')
-
-    qweight = model.model.layers[0].self_attn.o_proj.qweight
-    assert qweight.dtype == torch.uint8, (
-        f'Expected o_proj dtype torch.uint8 but got {qweight.dtype}')
-
-    qweight = model.model.layers[0].self_attn.qkv_proj.qweight
-    assert qweight.dtype == torch.uint8, (
-        f'Expected qkv_proj dtype torch.uint8 but got {qweight.dtype}')
-
-    # some weights should not be quantized
-    weight = model.lm_head.weight
-    assert weight.dtype != torch.uint8, (
-        'lm_head weight dtype should not be torch.uint8')
-
-    weight = model.model.embed_tokens.weight
-    assert weight.dtype != torch.uint8, (
-        'embed_tokens weight dtype should not be torch.uint8')
-
-    weight = model.model.layers[0].input_layernorm.weight
-    assert weight.dtype != torch.uint8, (
-        'input_layernorm weight dtype should not be torch.uint8')
-
-    weight = model.model.layers[0].post_attention_layernorm.weight
-    assert weight.dtype != torch.uint8, (
-        'input_layernorm weight dtype should not be torch.uint8')
-
-    # check the output of the model is expected
-    sampling_params = SamplingParams(temperature=0.0,
-                                     logprobs=1,
-                                     prompt_logprobs=1,
-                                     max_tokens=8)
-
-    prompts = ['That which does not kill us', 'To be or not to be,']
-    expected_outputs = [
-        'That which does not kill us makes us stronger.',
-        'To be or not to be, that is the question.'
-    ]
-    outputs = llm.generate(prompts, sampling_params=sampling_params)
-
-    assert len(outputs) == len(prompts)
-
-    for index in range(len(outputs)):
-        # compare the first line of the output
-        actual_output = outputs[index][1][0].split('\n', 1)[0]
-        expected_output = expected_outputs[index].split('\n', 1)[0]
-        assert actual_output == expected_output, (
-            f'Expected: {expected_output}, but got: {actual_output}')
+    with vllm_runner('huggyllama/llama-7b',
+                     quantization='bitsandbytes',
+                     load_format='bitsandbytes',
+                     enforce_eager=True) as llm:
+
+        model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model  # noqa: E501
+
+        # check the weights in MLP & SelfAttention are quantized to torch.uint8
+        qweight = model.model.layers[0].mlp.gate_up_proj.qweight
+        assert qweight.dtype == torch.uint8, (
+            f'Expected gate_up_proj dtype torch.uint8 but got {qweight.dtype}')
+
+        qweight = model.model.layers[0].mlp.down_proj.qweight
+        assert qweight.dtype == torch.uint8, (
+            f'Expected down_proj dtype torch.uint8 but got {qweight.dtype}')
+
+        qweight = model.model.layers[0].self_attn.o_proj.qweight
+        assert qweight.dtype == torch.uint8, (
+            f'Expected o_proj dtype torch.uint8 but got {qweight.dtype}')
+
+        qweight = model.model.layers[0].self_attn.qkv_proj.qweight
+        assert qweight.dtype == torch.uint8, (
+            f'Expected qkv_proj dtype torch.uint8 but got {qweight.dtype}')
+
+        # some weights should not be quantized
+        weight = model.lm_head.weight
+        assert weight.dtype != torch.uint8, (
+            'lm_head weight dtype should not be torch.uint8')
+
+        weight = model.model.embed_tokens.weight
+        assert weight.dtype != torch.uint8, (
+            'embed_tokens weight dtype should not be torch.uint8')
+
+        weight = model.model.layers[0].input_layernorm.weight
+        assert weight.dtype != torch.uint8, (
+            'input_layernorm weight dtype should not be torch.uint8')
+
+        weight = model.model.layers[0].post_attention_layernorm.weight
+        assert weight.dtype != torch.uint8, (
+            'input_layernorm weight dtype should not be torch.uint8')
+
+        # check the output of the model is expected
+        sampling_params = SamplingParams(temperature=0.0,
+                                         logprobs=1,
+                                         prompt_logprobs=1,
+                                         max_tokens=8)
+
+        prompts = ['That which does not kill us', 'To be or not to be,']
+        expected_outputs = [
+            'That which does not kill us makes us stronger.',
+            'To be or not to be, that is the question.'
+        ]
+        outputs = llm.generate(prompts, sampling_params=sampling_params)
+
+        assert len(outputs) == len(prompts)
+
+        for index in range(len(outputs)):
+            # compare the first line of the output
+            actual_output = outputs[index][1][0].split('\n', 1)[0]
+            expected_output = expected_outputs[index].split('\n', 1)[0]
+            assert actual_output == expected_output, (
+                f'Expected: {expected_output}, but got: {actual_output}')

tests/quantization/test_compressed_tensors.py

@@ -12,42 +12,45 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
 
 def test_compressed_tensors_w8a8_static_setup(vllm_runner):
     model_path = "nm-testing/tinyllama-one-shot-static-quant-test-compressed"
-    llm = vllm_runner(model_path, quantization="sparseml", enforce_eager=True)
-    model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model
-    layer = model.model.layers[0]
+    with vllm_runner(model_path, quantization="sparseml",
+                     enforce_eager=True) as llm:
+        model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model  # noqa: E501
+        layer = model.model.layers[0]
 
-    qkv_proj = layer.self_attn.qkv_proj
-    o_proj = layer.self_attn.o_proj
-    gate_up_proj = layer.mlp.gate_up_proj
-    down_proj = layer.mlp.down_proj
+        qkv_proj = layer.self_attn.qkv_proj
+        o_proj = layer.self_attn.o_proj
+        gate_up_proj = layer.mlp.gate_up_proj
+        down_proj = layer.mlp.down_proj
 
-    assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
-    assert isinstance(o_proj.quant_method, CompressedTensorsLinearMethod)
-    assert isinstance(gate_up_proj.quant_method, CompressedTensorsLinearMethod)
-    assert isinstance(down_proj.quant_method, CompressedTensorsLinearMethod)
+        assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
+        assert isinstance(o_proj.quant_method, CompressedTensorsLinearMethod)
+        assert isinstance(gate_up_proj.quant_method,
+                          CompressedTensorsLinearMethod)
+        assert isinstance(down_proj.quant_method,
+                          CompressedTensorsLinearMethod)
 
-    assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8StaticTensor)
+        assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8StaticTensor)
 
-    assert qkv_proj.weight.dtype is torch.int8
-    assert o_proj.weight.dtype is torch.int8
-    assert gate_up_proj.weight.dtype is torch.int8
+        assert qkv_proj.weight.dtype is torch.int8
+        assert o_proj.weight.dtype is torch.int8
+        assert gate_up_proj.weight.dtype is torch.int8
 
-    assert qkv_proj.weight_scale.shard_splitter is not None
-    assert qkv_proj.weight_scale.logical_widths is not None
-    assert qkv_proj.input_scale.dtype is torch.float32
+        assert qkv_proj.weight_scale.shard_splitter is not None
+        assert qkv_proj.weight_scale.logical_widths is not None
+        assert qkv_proj.input_scale.dtype is torch.float32
 
 
 def test_compressed_tensors_w8a8_dynanmic_per_token(vllm_runner):
     model_path = "nm-testing/tinyllama-one-shot-dynamic-test"
-    llm = vllm_runner(model_path,
-                      quantization="sparseml",
-                      enforce_eager=True,
-                      dtype=torch.float16)
-    model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model
-    layer = model.model.layers[0]
-
-    qkv_proj = layer.self_attn.qkv_proj
-
-    assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
-    assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8DynamicToken)
-    assert qkv_proj.weight.dtype is torch.int8
+    with vllm_runner(model_path,
+                     quantization="sparseml",
+                     enforce_eager=True,
+                     dtype=torch.float16) as llm:
+        model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model  # noqa: E501
+        layer = model.model.layers[0]
+
+        qkv_proj = layer.self_attn.qkv_proj
+
+        assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
+        assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8DynamicToken)
+        assert qkv_proj.weight.dtype is torch.int8