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

[CI/Test] improve robustness of test by replacing del with context manager (vllm_runner) (#5357)
2024-06-08 01:59:20 -07:00
parent 9fb900f90c
commit 8ea5e44a43
28 changed files with 431 additions and 470 deletions
--- a/tests/basic_correctness/test_preemption.py
+++ b/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