| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test that we handle an Error in model forward and shutdown.""" |
| 4 | |
| 5 | import asyncio |
| 6 | import inspect |
| 7 | |
| 8 | import pytest |
| 9 | |
| 10 | from tests.utils import wait_for_gpu_memory_to_clear |
| 11 | from tests.v1.shutdown.utils import ( |
| 12 | SHUTDOWN_TEST_THRESHOLD_BYTES, |
| 13 | SHUTDOWN_TEST_TIMEOUT_SEC, |
| 14 | ) |
| 15 | from vllm import LLM, AsyncEngineArgs, SamplingParams |
| 16 | from vllm.distributed import get_tensor_model_parallel_rank |
| 17 | from vllm.model_executor.models.llama import LlamaForCausalLM |
| 18 | from vllm.platforms import current_platform |
| 19 | from vllm.v1.engine.async_llm import AsyncLLM |
| 20 | from vllm.v1.engine.exceptions import EngineDeadError |
| 21 | |
| 22 | MODELS = ["hmellor/tiny-random-LlamaForCausalLM"] |
| 23 | |
| 24 | |
| 25 | def evil_forward(self, *args, **kwargs): |
| 26 | """Evil forward method that raise an exception after 10 calls.""" |
| 27 | NUMBER_OF_GOOD_PASSES = 10 |
| 28 | |
| 29 | if not hasattr(self, "num_calls"): |
| 30 | self.num_calls = 0 |
| 31 | |
| 32 | if ( |
| 33 | self.num_calls == NUMBER_OF_GOOD_PASSES |
| 34 | and get_tensor_model_parallel_rank() == 0 |
| 35 | ): |
| 36 | raise Exception("Simulated illegal memory access on Rank 0!") |
| 37 | self.num_calls += 1 |
| 38 | |
| 39 | kwargs.setdefault("intermediate_tensors", None) # required for MRV2 |
| 40 | return self.model(*args, **kwargs) |
| 41 | |
| 42 | |
| 43 | @pytest.fixture |
| 44 | def rocm_evil_forward(rocm_sitecustomize_factory): |
| 45 | lines = [ |
| 46 | "from vllm.distributed import get_tensor_model_parallel_rank", |
| 47 | "from vllm.model_executor.models.llama import LlamaForCausalLM", |
| 48 | inspect.getsource(evil_forward), |
| 49 | f"LlamaForCausalLM.forward = {evil_forward.__name__}", |
| 50 | ] |
| 51 | rocm_sitecustomize_factory(lines) |
| 52 | |
| 53 | |
| 54 | @pytest.mark.asyncio |
| 55 | @pytest.mark.parametrize("tensor_parallel_size", [2, 1]) |
| 56 | @pytest.mark.parametrize("model", MODELS) |
| 57 | async def test_async_llm_model_error( |
| 58 | monkeypatch, rocm_evil_forward, tensor_parallel_size: int, model: str |
| 59 | ) -> None: |
| 60 | """Test that AsyncLLM propagates a forward pass error and frees memory. |
| 61 | |
| 62 | AsyncLLM always uses an MP client. |
| 63 | """ |
| 64 | if current_platform.device_count() < tensor_parallel_size: |
| 65 | pytest.skip(reason="Not enough CUDA devices") |
| 66 | |
| 67 | # Monkeypatch an error in the model. |
| 68 | monkeypatch.setattr(LlamaForCausalLM, "forward", evil_forward) |
| 69 | |
| 70 | engine_args = AsyncEngineArgs( |
| 71 | model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size |
| 72 | ) |
| 73 | async_llm = AsyncLLM.from_engine_args(engine_args) |
| 74 | |
| 75 | async def generate(request_id: str): |
| 76 | generator = async_llm.generate( |
| 77 | "Hello my name is", request_id=request_id, sampling_params=SamplingParams() |
| 78 | ) |
| 79 | try: |
| 80 | async for _ in generator: |
| 81 | pass |
| 82 | except Exception as e: |
| 83 | return e |
| 84 | |
| 85 | NUM_REQS = 3 |
| 86 | tasks = [generate(f"request-{idx}") for idx in range(NUM_REQS)] |
| 87 | outputs = await asyncio.gather(*tasks) |
| 88 | |
| 89 | # Every request should get an EngineDeadError. |
| 90 | for output in outputs: |
| 91 | assert isinstance(output, EngineDeadError) |
| 92 | |
| 93 | # AsyncLLM should be errored. |
| 94 | assert async_llm.errored |
| 95 | |
| 96 | # We should not be able to make another request. |
| 97 | with pytest.raises(EngineDeadError): |
| 98 | async for _ in async_llm.generate( |
| 99 | "Hello my name is", request_id="abc", sampling_params=SamplingParams() |
| 100 | ): |
| 101 | raise Exception("We should not get here.") |
| 102 | |
| 103 | # Confirm all the processes are cleaned up. |
| 104 | wait_for_gpu_memory_to_clear( |
| 105 | devices=list(range(tensor_parallel_size)), |
| 106 | threshold_bytes=2 * 2**30, |
| 107 | timeout_s=60, |
| 108 | ) |
| 109 | |
| 110 | # NOTE: shutdown is handled by the API Server if an exception |
| 111 | # occurs, so it is expected that we would need to call this. |
| 112 | async_llm.shutdown() |
| 113 | |
| 114 | |
| 115 | @pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC) |
| 116 | @pytest.mark.parametrize("enable_multiprocessing", [True]) |
| 117 | @pytest.mark.parametrize("tensor_parallel_size", [2, 1]) |
| 118 | @pytest.mark.parametrize("model", MODELS) |
| 119 | def test_llm_model_error( |
| 120 | monkeypatch, |
| 121 | rocm_evil_forward, |
| 122 | tensor_parallel_size: int, |
| 123 | enable_multiprocessing: bool, |
| 124 | model: str, |
| 125 | ) -> None: |
| 126 | """Test that LLM propagates a forward pass error and frees memory. |
| 127 | TODO(andy) - LLM without multiprocessing; LLM with multiprocessing |
| 128 | and >1 rank |
| 129 | """ |
| 130 | if current_platform.device_count() < tensor_parallel_size: |
| 131 | pytest.skip(reason="Not enough CUDA devices") |
| 132 | |
| 133 | with monkeypatch.context() as m: |
| 134 | MP_VALUE = "1" if enable_multiprocessing else "0" |
| 135 | m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", MP_VALUE) |
| 136 | |
| 137 | # Monkeypatch an error in the model. |
| 138 | m.setattr(LlamaForCausalLM, "forward", evil_forward) |
| 139 | |
| 140 | llm = LLM( |
| 141 | model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size |
| 142 | ) |
| 143 | |
| 144 | with pytest.raises(EngineDeadError if enable_multiprocessing else Exception): |
| 145 | llm.generate("Hello my name is Robert and I") |
| 146 | |
| 147 | # Confirm all the processes are cleaned up. |
| 148 | wait_for_gpu_memory_to_clear( |
| 149 | devices=list(range(tensor_parallel_size)), |
| 150 | threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES, |
| 151 | ) |