[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/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)
+    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
+        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}')
+        # 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].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.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}')
+        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')
+        # 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.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].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')
+        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)
+        # 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)
+        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)
+        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}')
+        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]
+    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
+        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
+        assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
+        assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8DynamicToken)
+        assert qkv_proj.weight.dtype is torch.int8

tests/quantization/test_fp8.py

@@ -16,9 +16,9 @@ capability = capability[0] * 10 + capability[1]
     capability < QUANTIZATION_METHODS["fp8"].get_min_capability(),
     reason="FP8 is not supported on this GPU type.")
 def test_load_fp16_model(vllm_runner) -> None:
-    llm = vllm_runner("facebook/opt-125m", quantization="fp8")
+    with vllm_runner("facebook/opt-125m", quantization="fp8") as llm:
 
-    model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model
-    fc1 = model.model.decoder.layers[0].fc1
-    assert isinstance(fc1.quant_method, Fp8LinearMethod)
-    assert fc1.weight.dtype == torch.float8_e4m3fn
+        model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model  # noqa: E501
+        fc1 = model.model.decoder.layers[0].fc1
+        assert isinstance(fc1.quant_method, Fp8LinearMethod)
+        assert fc1.weight.dtype == torch.float8_e4m3fn