[V1][Frontend] Improve Shutdown And Logs

.buildkite/test-pipeline.yaml

@@ -197,6 +197,7 @@ steps:
   commands:
     # split the test to avoid interference
     - VLLM_USE_V1=1 pytest -v -s v1/core
+    - VLLM_USE_V1=1 pytest -v -s v1/shutdown
     - VLLM_USE_V1=1 pytest -v -s v1/engine
     - VLLM_USE_V1=1 pytest -v -s v1/sample
     - VLLM_USE_V1=1 pytest -v -s v1/worker

tests/v1/engine/test_async_llm.py

@@ -6,14 +6,14 @@
 
 import pytest
 
-from tests.v1.engine.utils import PLP_APC_UNSUPPORTED_MSG
 from vllm import SamplingParams
 from vllm.assets.image import ImageAsset
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.inputs import PromptType
 from vllm.platforms import current_platform
 from vllm.sampling_params import RequestOutputKind
 from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.engine.exceptions import EngineGenerateError
 
 if not current_platform.is_cuda():
     pytest.skip(reason="V1 currently only supported on CUDA.",
@@ -88,10 +88,11 @@ async def test_async_llm_refuses_prompt_logprobs_with_apc(
     apc_engine_args = AsyncEngineArgs(model="facebook/opt-125m",
                                       enable_prefix_caching=True,
                                       gpu_memory_utilization=0.8,
+                                      enforce_eager=True,
                                       disable_log_requests=True)
     engine = AsyncLLM.from_engine_args(apc_engine_args)
     try:
-        with pytest.raises(ValueError) as excinfo:
+        with pytest.raises(EngineGenerateError):
             # Issue a request with prompt logprobs enabled, which should fail
             await asyncio.create_task(
                 generate(engine,
@@ -100,8 +101,6 @@ async def test_async_llm_refuses_prompt_logprobs_with_apc(
                          output_kind,
                          10,
                          prompt_logprobs=5))
-        # Validate exception string is correct
-        assert str(excinfo.value) == PLP_APC_UNSUPPORTED_MSG
     finally:
         # Shut down engine
         engine.shutdown()

tests/v1/shutdown/test_forward_error.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Test that we handle an Error in model forward and shutdown."""
+
+import asyncio
+
+import pytest
+
+from tests.utils import wait_for_gpu_memory_to_clear
+from vllm import LLM, SamplingParams
+from vllm.distributed import get_tensor_model_parallel_rank
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.model_executor.models.llama import LlamaForCausalLM
+from vllm.utils import GiB_bytes, cuda_device_count_stateless
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.engine.exceptions import EngineDeadError
+
+
+def evil_forward(self, *args, **kwargs):
+    """Evil forward method that raises an exception after 10 calls."""
+    NUMBER_OF_GOOD_PASSES = 10
+
+    if not hasattr(self, "num_calls"):
+        self.num_calls = 0
+
+    if (self.num_calls == NUMBER_OF_GOOD_PASSES
+            and get_tensor_model_parallel_rank() == 0):
+        raise Exception("Simulated illegal memory access on Rank 0!")
+    self.num_calls += 1
+
+    return self.model(*args, **kwargs)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
+async def test_async_llm_model_error(monkeypatch, tensor_parallel_size):
+
+    if cuda_device_count_stateless() < tensor_parallel_size:
+        pytest.skip(reason="Not enough CUDA devices")
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+
+        # Monkeypatch an error in the model.
+        m.setattr(LlamaForCausalLM, "forward", evil_forward)
+
+        engine_args = AsyncEngineArgs(
+            model="meta-llama/Llama-3.2-1B",
+            enforce_eager=True,
+            tensor_parallel_size=tensor_parallel_size)
+        async_llm = AsyncLLM.from_engine_args(engine_args)
+
+        async def generate(request_id: str):
+            generator = async_llm.generate("Hello my name is",
+                                           request_id=request_id,
+                                           sampling_params=SamplingParams())
+            try:
+                async for _ in generator:
+                    pass
+            except Exception as e:
+                return e
+
+        NUM_REQS = 3
+        tasks = [generate(f"request-{idx}") for idx in range(NUM_REQS)]
+        outputs = await asyncio.gather(*tasks)
+
+        # Every request should get an EngineDeadError.
+        for output in outputs:
+            assert isinstance(output, EngineDeadError)
+
+        # AsyncLLM should be errored.
+        assert async_llm.errored
+
+        # We should not be able to make another request.
+        with pytest.raises(EngineDeadError):
+            async for _ in async_llm.generate(
+                    "Hello my name is",
+                    request_id="abc",
+                    sampling_params=SamplingParams()):
+                raise Exception("We should not get here.")
+
+        # Confirm all the processes are cleaned up.
+        wait_for_gpu_memory_to_clear(
+            devices=list(range(tensor_parallel_size)),
+            threshold_bytes=2 * GiB_bytes,
+            timeout_s=60,
+        )
+
+        # NOTE: shutdown is handled by the API Server if an exception
+        # occurs, so it is expected that we would need to call this.
+        async_llm.shutdown()
+
+
+@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
+def test_llm_model_error(monkeypatch, tensor_parallel_size):
+
+    if cuda_device_count_stateless() < tensor_parallel_size:
+        pytest.skip(reason="Not enough CUDA devices")
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+
+        # Monkeypatch an error in the model.
+        m.setattr(LlamaForCausalLM, "forward", evil_forward)
+
+        llm = LLM(model="meta-llama/Llama-3.2-1B",
+                  enforce_eager=True,
+                  tensor_parallel_size=tensor_parallel_size)
+
+        with pytest.raises(EngineDeadError):
+            llm.generate("Hello my name is Robert and I")
+
+    # Confirm all the processes are cleaned up.
+    wait_for_gpu_memory_to_clear(
+        devices=list(range(tensor_parallel_size)),
+        threshold_bytes=2 * GiB_bytes,
+        timeout_s=60,
+    )

tests/v1/shutdown/test_processor_error.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Test error handling in Processor. Should not impact other reqs."""
+
+import asyncio
+
+import pytest
+
+from vllm import SamplingParams
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.inputs.data import TokensPrompt
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.engine.exceptions import EngineGenerateError
+
+
+@pytest.mark.asyncio
+async def test_async_llm_processor_error(monkeypatch):
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+
+        engine_args = AsyncEngineArgs(model="meta-llama/Llama-3.2-1B",
+                                      enforce_eager=True)
+        async_llm = AsyncLLM.from_engine_args(engine_args)
+
+        async def generate(request_id: str):
+            # [] is not allowed and will raise a ValueError in Processor.
+            generator = async_llm.generate(TokensPrompt([]),
+                                           request_id=request_id,
+                                           sampling_params=SamplingParams())
+            try:
+                async for _ in generator:
+                    pass
+            except Exception as e:
+                return e
+
+        NUM_REQS = 3
+        tasks = [generate(f"request-{idx}") for idx in range(NUM_REQS)]
+        outputs = await asyncio.gather(*tasks)
+
+        # Every request should have get an EngineGenerateError.
+        for output in outputs:
+            with pytest.raises(EngineGenerateError):
+                raise output
+
+        # AsyncLLM should be errored.
+        assert not async_llm.errored
+
+        # This should be no problem.
+        EXPECTED_TOKENS = 5
+        outputs = []
+        async for out in async_llm.generate(
+                "Hello my name is",
+                request_id="abc",
+                sampling_params=SamplingParams(
+                    max_tokens=EXPECTED_TOKENS,
+                    output_kind=RequestOutputKind.DELTA)):
+            outputs.append(out)
+
+        generated_tokens = []
+        for out in outputs:
+            generated_tokens.extend(out.outputs[0].token_ids)
+        assert len(generated_tokens) == EXPECTED_TOKENS
+
+        async_llm.shutdown()

tests/v1/shutdown/test_startup_error.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Test that we handle a startup Error and shutdown."""
+
+import pytest
+
+from tests.utils import wait_for_gpu_memory_to_clear
+from vllm import LLM
+from vllm.distributed import get_tensor_model_parallel_rank
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.model_executor.models.llama import LlamaForCausalLM
+from vllm.utils import GiB_bytes, cuda_device_count_stateless
+from vllm.v1.engine.async_llm import AsyncLLM
+
+
+def evil_forward(self, *args, **kwargs):
+    """Evil forward method that raises an exception."""
+
+    if get_tensor_model_parallel_rank() == 0:
+        raise RuntimeError("Simulated Error during forward pass!")
+
+    return self.model(*args, **kwargs)
+
+
+def evil_load_weights(self, *args, **kwargs):
+    """Evil load_weights method that raises an exception."""
+
+    raise RuntimeError("Simulated OOM Error during weight loading!")
+
+
+MODELS = [
+    "meta-llama/Llama-3.2-1B",
+]
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
+def test_async_llm_forward_pass_error(monkeypatch, model,
+                                      tensor_parallel_size):
+    """Test failure during first forward pass"""
+
+    if cuda_device_count_stateless() < tensor_parallel_size:
+        pytest.skip(reason="Not enough CUDA devices")
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+
+        # Monkeypatch an error in the model.
+        m.setattr(LlamaForCausalLM, "load_weights", evil_load_weights)
+
+        engine_args = AsyncEngineArgs(
+            model=model,
+            enforce_eager=True,
+            tensor_parallel_size=tensor_parallel_size)
+
+        # Confirm we get an exception.
+        with pytest.raises(Exception,
+                           match="EngineCore initialization failed"):
+            _ = AsyncLLM.from_engine_args(engine_args)
+
+        # Confirm all the processes are cleaned up.
+        wait_for_gpu_memory_to_clear(
+            devices=list(range(tensor_parallel_size)),
+            threshold_bytes=2 * GiB_bytes,
+            timeout_s=60,
+        )
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
+def test_async_llm_weight_loading_failure(monkeypatch, model,
+                                          tensor_parallel_size):
+    """Test failure during first forward pass"""
+
+    if cuda_device_count_stateless() < tensor_parallel_size:
+        pytest.skip(reason="Not enough CUDA devices")
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+
+        # Monkeypatch an error in the model.
+        m.setattr(LlamaForCausalLM, "forward", evil_forward)
+
+        engine_args = AsyncEngineArgs(
+            model=model,
+            enforce_eager=True,
+            tensor_parallel_size=tensor_parallel_size)
+
+        # Confirm we get an exception.
+        with pytest.raises(Exception,
+                           match="EngineCore initialization failed"):
+            _ = AsyncLLM.from_engine_args(engine_args)
+
+        # Confirm all the processes are cleaned up.
+        wait_for_gpu_memory_to_clear(
+            devices=list(range(tensor_parallel_size)),
+            threshold_bytes=2 * GiB_bytes,
+            timeout_s=60,
+        )
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
+def test_llm_forward_pass_failure(monkeypatch, model, tensor_parallel_size):
+    """Test failure during first forward pass (after IPC setup)."""
+
+    if cuda_device_count_stateless() < tensor_parallel_size:
+        pytest.skip(reason="Not enough CUDA devices")
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+
+        # Simulate error during forward pass
+        m.setattr(LlamaForCausalLM, "forward", evil_forward)
+
+        with pytest.raises(Exception,
+                           match="EngineCore initialization failed"):
+            _ = LLM(model=model,
+                    enforce_eager=True,
+                    tensor_parallel_size=tensor_parallel_size)
+
+        wait_for_gpu_memory_to_clear(
+            devices=list(range(tensor_parallel_size)),
+            threshold_bytes=2 * GiB_bytes,
+            timeout_s=60,
+        )
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
+def test_llm_weight_loading_failure(monkeypatch, model, tensor_parallel_size):
+    """Test failure during weight loading (before IPC setup)."""
+
+    if cuda_device_count_stateless() < tensor_parallel_size:
+        pytest.skip(reason="Not enough CUDA devices")
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+
+        # Simulate error during weight loading
+        m.setattr(LlamaForCausalLM, "load_weights", evil_load_weights)
+
+        with pytest.raises(Exception,
+                           match="EngineCore initialization failed"):
+            _ = LLM(model=model,
+                    enforce_eager=True,
+                    tensor_parallel_size=tensor_parallel_size)
+
+        wait_for_gpu_memory_to_clear(
+            devices=list(range(tensor_parallel_size)),
+            threshold_bytes=2 * GiB_bytes,
+            timeout_s=60,
+        )