565 lines
17 KiB
Python
565 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Integration tests for shutdown behavior, timeout, and signal handling."""
|
|
|
|
import asyncio
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
|
|
import httpx
|
|
import openai
|
|
import psutil
|
|
import pytest
|
|
|
|
from tests.utils import RemoteOpenAIServer
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.network_utils import get_open_port
|
|
|
|
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
|
|
|
|
# GPU initialization might take take longer
|
|
_IS_ROCM = current_platform.is_rocm()
|
|
_SERVER_STARTUP_TIMEOUT = 120
|
|
_PROCESS_EXIT_TIMEOUT = 15
|
|
_SHUTDOWN_DETECTION_TIMEOUT = 10
|
|
_CHILD_CLEANUP_TIMEOUT = 10
|
|
|
|
|
|
def _get_child_pids(parent_pid: int) -> list[int]:
|
|
try:
|
|
parent = psutil.Process(parent_pid)
|
|
return [c.pid for c in parent.children(recursive=True)]
|
|
except psutil.NoSuchProcess:
|
|
return []
|
|
|
|
|
|
async def _assert_children_cleaned_up(
|
|
child_pids: list[int],
|
|
timeout: float = _CHILD_CLEANUP_TIMEOUT,
|
|
):
|
|
"""Wait for child processes to exit and fail if any remain."""
|
|
if not child_pids:
|
|
return
|
|
|
|
deadline = time.time() + timeout
|
|
while time.time() < deadline:
|
|
still_alive = []
|
|
for pid in child_pids:
|
|
try:
|
|
p = psutil.Process(pid)
|
|
if p.is_running() and p.status() != psutil.STATUS_ZOMBIE:
|
|
still_alive.append(pid)
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
if not still_alive:
|
|
return
|
|
await asyncio.sleep(0.5)
|
|
|
|
pytest.fail(
|
|
f"Child processes {still_alive} still alive after {timeout}s. "
|
|
f"Process cleanup may not be working correctly."
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class ShutdownState:
|
|
got_503: bool = False
|
|
got_500: bool = False
|
|
requests_after_sigterm: int = 0
|
|
aborted_requests: int = 0
|
|
connection_errors: int = 0
|
|
stop_requesting: bool = False
|
|
errors: list[str] = field(default_factory=list)
|
|
|
|
|
|
async def _concurrent_request_loop(
|
|
client: openai.AsyncOpenAI,
|
|
state: ShutdownState,
|
|
sigterm_sent: asyncio.Event | None = None,
|
|
concurrency: int = 10,
|
|
):
|
|
"""Run multiple concurrent requests to keep the server busy."""
|
|
|
|
async def single_request():
|
|
while not state.stop_requesting:
|
|
try:
|
|
response = await client.completions.create(
|
|
model=MODEL_NAME,
|
|
prompt="Write a story: ",
|
|
max_tokens=200,
|
|
)
|
|
if sigterm_sent is not None and sigterm_sent.is_set():
|
|
state.requests_after_sigterm += 1
|
|
# Check if any choice has finish_reason='abort'
|
|
if any(choice.finish_reason == "abort" for choice in response.choices):
|
|
state.aborted_requests += 1
|
|
except openai.APIStatusError as e:
|
|
if e.status_code == 503:
|
|
state.got_503 = True
|
|
elif e.status_code == 500:
|
|
state.got_500 = True
|
|
else:
|
|
state.errors.append(f"API error: {e}")
|
|
except (openai.APIConnectionError, httpx.RemoteProtocolError):
|
|
state.connection_errors += 1
|
|
if sigterm_sent is not None and sigterm_sent.is_set():
|
|
break
|
|
except Exception as e:
|
|
state.errors.append(f"Unexpected error: {e}")
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
|
|
tasks = [asyncio.create_task(single_request()) for _ in range(concurrency)]
|
|
try:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
finally:
|
|
for t in tasks:
|
|
if not t.done():
|
|
t.cancel()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_on_engine_failure():
|
|
"""Verify that API returns connection error when server process is killed.
|
|
|
|
Starts a vLLM server, kills it to simulate a crash, then verifies that
|
|
subsequent API calls fail appropriately.
|
|
"""
|
|
|
|
port = get_open_port()
|
|
|
|
proc = subprocess.Popen(
|
|
[
|
|
# dtype, max-len etc set so that this can run in CI
|
|
sys.executable,
|
|
"-m",
|
|
"vllm.entrypoints.openai.api_server",
|
|
"--model",
|
|
MODEL_NAME,
|
|
"--dtype",
|
|
"bfloat16",
|
|
"--max-model-len",
|
|
"128",
|
|
"--enforce-eager",
|
|
"--port",
|
|
str(port),
|
|
"--gpu-memory-utilization",
|
|
"0.05",
|
|
"--max-num-seqs",
|
|
"2",
|
|
"--disable-frontend-multiprocessing",
|
|
],
|
|
# ROCm: Disable stdout/stderr pipe capture. Subprocess hangs when
|
|
# stdout/stderr pipes are enabled during ROCm GPU initialization.
|
|
stdout=None if _IS_ROCM else subprocess.PIPE,
|
|
stderr=None if _IS_ROCM else subprocess.PIPE,
|
|
text=None if _IS_ROCM else True,
|
|
preexec_fn=lambda: signal.signal(signal.SIGINT, signal.SIG_IGN),
|
|
)
|
|
|
|
# Wait for server startup
|
|
start_time = time.time()
|
|
client = openai.AsyncOpenAI(
|
|
base_url=f"http://localhost:{port}/v1",
|
|
api_key="dummy",
|
|
max_retries=0,
|
|
timeout=10,
|
|
)
|
|
|
|
# Poll until server is ready
|
|
while time.time() - start_time < _SERVER_STARTUP_TIMEOUT:
|
|
try:
|
|
await client.completions.create(
|
|
model=MODEL_NAME, prompt="Hello", max_tokens=1
|
|
)
|
|
break
|
|
except Exception:
|
|
time.sleep(0.5)
|
|
if proc.poll() is not None:
|
|
if _IS_ROCM:
|
|
pytest.fail(f"Server died during startup: {proc.returncode}")
|
|
else:
|
|
stdout, stderr = proc.communicate(timeout=1)
|
|
pytest.fail(
|
|
f"Server died during startup. "
|
|
f"stdout: {stdout}, stderr: {stderr}"
|
|
)
|
|
else:
|
|
proc.terminate()
|
|
proc.wait(timeout=_PROCESS_EXIT_TIMEOUT)
|
|
pytest.fail(f"Server failed to start in {_SERVER_STARTUP_TIMEOUT} seconds")
|
|
|
|
# Kill server to simulate crash
|
|
proc.terminate()
|
|
time.sleep(1)
|
|
|
|
# Verify API calls now fail
|
|
with pytest.raises((openai.APIConnectionError, openai.APIStatusError)):
|
|
await client.completions.create(
|
|
model=MODEL_NAME, prompt="This should fail", max_tokens=1
|
|
)
|
|
|
|
return_code = proc.wait(timeout=_PROCESS_EXIT_TIMEOUT)
|
|
assert return_code is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wait_timeout_completes_requests():
|
|
"""Verify wait timeout: new requests rejected, in-flight requests complete."""
|
|
server_args = [
|
|
"--dtype",
|
|
"bfloat16",
|
|
"--max-model-len",
|
|
"256",
|
|
"--enforce-eager",
|
|
"--gpu-memory-utilization",
|
|
"0.05",
|
|
"--max-num-seqs",
|
|
"4",
|
|
"--shutdown-timeout",
|
|
"30",
|
|
]
|
|
|
|
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
|
|
client = remote_server.get_async_client()
|
|
proc = remote_server.proc
|
|
child_pids = _get_child_pids(proc.pid)
|
|
|
|
state = ShutdownState()
|
|
sigterm_sent = asyncio.Event()
|
|
|
|
request_task = asyncio.create_task(
|
|
_concurrent_request_loop(client, state, sigterm_sent, concurrency=10)
|
|
)
|
|
|
|
await asyncio.sleep(0.5)
|
|
proc.send_signal(signal.SIGTERM)
|
|
sigterm_sent.set()
|
|
|
|
try:
|
|
await asyncio.wait_for(request_task, timeout=_SHUTDOWN_DETECTION_TIMEOUT)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
finally:
|
|
state.stop_requesting = True
|
|
if not request_task.done():
|
|
request_task.cancel()
|
|
await asyncio.gather(request_task, return_exceptions=True)
|
|
|
|
# wait timeout should complete in-flight requests
|
|
assert state.requests_after_sigterm > 0, (
|
|
f"Wait timeout should complete in-flight requests. "
|
|
f"503: {state.got_503}, 500: {state.got_500}, "
|
|
f"conn_errors: {state.connection_errors}, errors: {state.errors}"
|
|
)
|
|
# server must stop accepting new requests (503, 500, or connection close)
|
|
assert state.got_503 or state.got_500 or state.connection_errors > 0, (
|
|
f"Server should stop accepting requests. "
|
|
f"completed: {state.requests_after_sigterm}, errors: {state.errors}"
|
|
)
|
|
|
|
await _assert_children_cleaned_up(child_pids)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("wait_for_engine_idle", [0.0, 2.0])
|
|
async def test_abort_timeout_exits_quickly(wait_for_engine_idle: float):
|
|
server_args = [
|
|
"--dtype",
|
|
"bfloat16",
|
|
"--max-model-len",
|
|
"256",
|
|
"--enforce-eager",
|
|
"--gpu-memory-utilization",
|
|
"0.05",
|
|
"--max-num-seqs",
|
|
"4",
|
|
"--shutdown-timeout",
|
|
"0",
|
|
]
|
|
|
|
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
|
|
proc = remote_server.proc
|
|
child_pids = _get_child_pids(proc.pid)
|
|
|
|
if wait_for_engine_idle > 0:
|
|
client = remote_server.get_async_client()
|
|
# Send requests to ensure engine is fully initialized
|
|
for _ in range(2):
|
|
await client.completions.create(
|
|
model=MODEL_NAME,
|
|
prompt="Test request: ",
|
|
max_tokens=10,
|
|
)
|
|
# Wait for engine to become idle
|
|
await asyncio.sleep(wait_for_engine_idle)
|
|
|
|
start_time = time.time()
|
|
proc.send_signal(signal.SIGTERM)
|
|
|
|
# abort timeout (0) should exit promptly
|
|
for _ in range(20):
|
|
if proc.poll() is not None:
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
if proc.poll() is None:
|
|
proc.kill()
|
|
proc.wait(timeout=5)
|
|
pytest.fail("Process did not exit after SIGTERM with abort timeout")
|
|
|
|
exit_time = time.time() - start_time
|
|
assert exit_time < 2, f"Default shutdown took too long: {exit_time:.1f}s"
|
|
assert proc.returncode in (0, -15, None), f"Unexpected: {proc.returncode}"
|
|
|
|
await _assert_children_cleaned_up(child_pids)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wait_timeout_with_short_duration():
|
|
"""Verify server exits cleanly with a short wait timeout."""
|
|
wait_timeout = 3
|
|
server_args = [
|
|
"--dtype",
|
|
"bfloat16",
|
|
"--max-model-len",
|
|
"256",
|
|
"--enforce-eager",
|
|
"--gpu-memory-utilization",
|
|
"0.05",
|
|
"--max-num-seqs",
|
|
"4",
|
|
"--shutdown-timeout",
|
|
str(wait_timeout),
|
|
]
|
|
|
|
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
|
|
client = remote_server.get_async_client()
|
|
proc = remote_server.proc
|
|
child_pids = _get_child_pids(proc.pid)
|
|
|
|
state = ShutdownState()
|
|
request_task = asyncio.create_task(
|
|
_concurrent_request_loop(client, state, concurrency=3)
|
|
)
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
start_time = time.time()
|
|
proc.send_signal(signal.SIGTERM)
|
|
|
|
# server should exit within wait_timeout + buffer
|
|
max_wait = wait_timeout + 15
|
|
for _ in range(int(max_wait * 10)):
|
|
if proc.poll() is not None:
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
exit_time = time.time() - start_time
|
|
|
|
state.stop_requesting = True
|
|
if not request_task.done():
|
|
request_task.cancel()
|
|
await asyncio.gather(request_task, return_exceptions=True)
|
|
|
|
if proc.poll() is None:
|
|
proc.kill()
|
|
proc.wait(timeout=5)
|
|
pytest.fail(f"Process did not exit within {max_wait}s after SIGTERM")
|
|
|
|
assert exit_time < wait_timeout + 10, (
|
|
f"Took too long to exit ({exit_time:.1f}s), expected <{wait_timeout + 10}s"
|
|
)
|
|
assert proc.returncode in (0, -15, None), f"Unexpected: {proc.returncode}"
|
|
|
|
await _assert_children_cleaned_up(child_pids)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_abort_timeout_fails_inflight_requests():
|
|
"""Verify abort timeout (0) immediately aborts in-flight requests."""
|
|
server_args = [
|
|
"--dtype",
|
|
"bfloat16",
|
|
"--max-model-len",
|
|
"256",
|
|
"--enforce-eager",
|
|
"--gpu-memory-utilization",
|
|
"0.05",
|
|
"--max-num-seqs",
|
|
"4",
|
|
"--shutdown-timeout",
|
|
"0",
|
|
]
|
|
|
|
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
|
|
client = remote_server.get_async_client()
|
|
proc = remote_server.proc
|
|
child_pids = _get_child_pids(proc.pid)
|
|
|
|
state = ShutdownState()
|
|
sigterm_sent = asyncio.Event()
|
|
|
|
request_task = asyncio.create_task(
|
|
_concurrent_request_loop(client, state, sigterm_sent, concurrency=10)
|
|
)
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
proc.send_signal(signal.SIGTERM)
|
|
sigterm_sent.set()
|
|
|
|
try:
|
|
await asyncio.wait_for(request_task, timeout=5)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
finally:
|
|
state.stop_requesting = True
|
|
if not request_task.done():
|
|
request_task.cancel()
|
|
await asyncio.gather(request_task, return_exceptions=True)
|
|
|
|
# With abort timeout (0), requests should be aborted (finish_reason='abort')
|
|
# or rejected (connection errors or API errors)
|
|
assert (
|
|
state.aborted_requests > 0
|
|
or state.connection_errors > 0
|
|
or state.got_500
|
|
or state.got_503
|
|
), (
|
|
f"Abort timeout should cause request aborts or failures. "
|
|
f"aborted: {state.aborted_requests}, "
|
|
f"503: {state.got_503}, 500: {state.got_500}, "
|
|
f"conn_errors: {state.connection_errors}, "
|
|
f"completed: {state.requests_after_sigterm}"
|
|
)
|
|
|
|
# Verify fast shutdown
|
|
start_time = time.time()
|
|
for _ in range(100):
|
|
if proc.poll() is not None:
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
exit_time = time.time() - start_time
|
|
assert exit_time < 10, f"Abort timeout shutdown took too long: {exit_time:.1f}s"
|
|
|
|
await _assert_children_cleaned_up(child_pids)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_request_rejection_during_shutdown():
|
|
"""Verify new requests are rejected with error during shutdown."""
|
|
server_args = [
|
|
"--dtype",
|
|
"bfloat16",
|
|
"--max-model-len",
|
|
"256",
|
|
"--enforce-eager",
|
|
"--gpu-memory-utilization",
|
|
"0.05",
|
|
"--max-num-seqs",
|
|
"4",
|
|
"--shutdown-timeout",
|
|
"30",
|
|
]
|
|
|
|
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
|
|
client = remote_server.get_async_client()
|
|
proc = remote_server.proc
|
|
child_pids = _get_child_pids(proc.pid)
|
|
|
|
proc.send_signal(signal.SIGTERM)
|
|
|
|
await asyncio.sleep(1.0)
|
|
|
|
# Try to send new requests - they should be rejected
|
|
rejected_count = 0
|
|
for _ in range(10):
|
|
try:
|
|
await client.completions.create(
|
|
model=MODEL_NAME, prompt="Hello", max_tokens=10
|
|
)
|
|
except (
|
|
openai.APIStatusError,
|
|
openai.APIConnectionError,
|
|
httpx.RemoteProtocolError,
|
|
):
|
|
rejected_count += 1
|
|
await asyncio.sleep(0.1)
|
|
|
|
assert rejected_count > 0, (
|
|
f"Expected requests to be rejected during shutdown, "
|
|
f"but {rejected_count} were rejected out of 10"
|
|
)
|
|
|
|
await _assert_children_cleaned_up(child_pids)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multi_api_server_shutdown():
|
|
"""Verify shutdown works with multiple API servers."""
|
|
server_args = [
|
|
"--dtype",
|
|
"bfloat16",
|
|
"--max-model-len",
|
|
"256",
|
|
"--enforce-eager",
|
|
"--gpu-memory-utilization",
|
|
"0.05",
|
|
"--max-num-seqs",
|
|
"4",
|
|
"--shutdown-timeout",
|
|
"30",
|
|
"--api-server-count",
|
|
"2",
|
|
]
|
|
|
|
with RemoteOpenAIServer(MODEL_NAME, server_args, auto_port=True) as remote_server:
|
|
client = remote_server.get_async_client()
|
|
proc = remote_server.proc
|
|
child_pids = _get_child_pids(proc.pid)
|
|
|
|
assert len(child_pids) >= 2, (
|
|
f"Expected at least 2 child processes, got {len(child_pids)}"
|
|
)
|
|
|
|
state = ShutdownState()
|
|
sigterm_sent = asyncio.Event()
|
|
|
|
# Start concurrent requests across both API servers
|
|
request_task = asyncio.create_task(
|
|
_concurrent_request_loop(client, state, sigterm_sent, concurrency=8)
|
|
)
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
# Send SIGTERM to parent - should propagate to all children
|
|
proc.send_signal(signal.SIGTERM)
|
|
sigterm_sent.set()
|
|
|
|
try:
|
|
await asyncio.wait_for(request_task, timeout=_SHUTDOWN_DETECTION_TIMEOUT)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
finally:
|
|
state.stop_requesting = True
|
|
if not request_task.done():
|
|
request_task.cancel()
|
|
await asyncio.gather(request_task, return_exceptions=True)
|
|
|
|
for _ in range(300): # up to 30 seconds
|
|
if proc.poll() is not None:
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
if proc.poll() is None:
|
|
proc.kill()
|
|
proc.wait(timeout=5)
|
|
pytest.fail("Process did not exit after SIGTERM")
|
|
|
|
await _assert_children_cleaned_up(child_pids)
|