vllm-project--vllm-omni
wehub-resource-sync
135 行
5.2 KiB
Python
135 行
5.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import functools
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from collections.abc import Callable
|
|
from contextlib import ExitStack, suppress
|
|
from typing import Any, Literal
|
|
|
|
import cloudpickle
|
|
from typing_extensions import ParamSpec
|
|
from vllm.platforms import current_platform
|
|
|
|
_P = ParamSpec("_P")
|
|
|
|
|
|
def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
|
|
"""Decorator to fork a new process for each test function."""
|
|
|
|
@functools.wraps(func)
|
|
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
|
|
os.setpgrp()
|
|
from _pytest.outcomes import Skipped
|
|
|
|
with (
|
|
tempfile.NamedTemporaryFile(
|
|
delete=False, mode="w+b", prefix=f"vllm_test_{func.__name__}_{os.getpid()}_", suffix=".exc"
|
|
) as exc_file,
|
|
ExitStack() as delete_after,
|
|
):
|
|
exc_file_path = exc_file.name
|
|
delete_after.callback(os.remove, exc_file_path)
|
|
pid = os.fork()
|
|
if pid == 0:
|
|
delete_after.pop_all()
|
|
try:
|
|
func(*args, **kwargs)
|
|
except Skipped as e:
|
|
print(str(e))
|
|
os._exit(0)
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
tb_string = traceback.format_exc()
|
|
exc_to_serialize: dict[str, Any]
|
|
try:
|
|
exc_to_serialize = {"pickled_exception": e}
|
|
cloudpickle.dumps(exc_to_serialize)
|
|
except (Exception, KeyboardInterrupt):
|
|
exc_to_serialize = {
|
|
"exception_type": type(e).__name__,
|
|
"exception_msg": str(e),
|
|
"traceback": tb_string,
|
|
}
|
|
try:
|
|
with open(exc_file_path, "wb") as f:
|
|
cloudpickle.dump(exc_to_serialize, f)
|
|
except Exception:
|
|
print(tb_string)
|
|
os._exit(1)
|
|
else:
|
|
os._exit(0)
|
|
else:
|
|
pgid = os.getpgid(pid)
|
|
_pid, _exitcode = os.waitpid(pid, 0)
|
|
old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
|
|
os.killpg(pgid, signal.SIGTERM)
|
|
signal.signal(signal.SIGTERM, old_signal_handler)
|
|
if _exitcode != 0:
|
|
exc_info = {}
|
|
if os.path.exists(exc_file_path):
|
|
with suppress(Exception), open(exc_file_path, "rb") as f:
|
|
exc_info = cloudpickle.load(f)
|
|
if (original_exception := exc_info.get("pickled_exception")) is not None:
|
|
assert isinstance(original_exception, Exception)
|
|
raise original_exception
|
|
if (original_tb := exc_info.get("traceback")) is not None:
|
|
raise AssertionError(
|
|
f"Test {func.__name__} failed when called with args {args} and kwargs {kwargs}"
|
|
f" (exit code: {_exitcode}):\n{original_tb}"
|
|
) from None
|
|
raise AssertionError(
|
|
f"function {func.__name__} failed when called with args {args} and kwargs {kwargs}"
|
|
f" (exit code: {_exitcode})"
|
|
) from None
|
|
|
|
return wrapper
|
|
|
|
|
|
def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
|
|
"""Decorator to spawn a new process for each test function."""
|
|
|
|
@functools.wraps(f)
|
|
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
|
|
if os.environ.get("RUNNING_IN_SUBPROCESS") == "1":
|
|
return f(*args, **kwargs)
|
|
|
|
import torch.multiprocessing as mp
|
|
|
|
with suppress(RuntimeError):
|
|
mp.set_start_method("spawn")
|
|
|
|
module_name = f.__module__
|
|
env = os.environ.copy()
|
|
env["RUNNING_IN_SUBPROCESS"] = "1"
|
|
|
|
with tempfile.TemporaryDirectory() as tempdir:
|
|
output_filepath = os.path.join(tempdir, "new_process.tmp")
|
|
input_bytes = cloudpickle.dumps((f, output_filepath))
|
|
cmd = [sys.executable, "-m", f"{module_name}"]
|
|
returned = subprocess.run(cmd, input=input_bytes, capture_output=True, env=env)
|
|
try:
|
|
returned.check_returncode()
|
|
except Exception as e:
|
|
raise RuntimeError(f"Error raised in subprocess:\n{returned.stderr.decode()}") from e
|
|
|
|
return wrapper
|
|
|
|
|
|
def create_new_process_for_each_test(
|
|
method: Literal["spawn", "fork"] | None = None,
|
|
) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
|
|
"""Creates a decorator that runs each test function in a new process."""
|
|
if method is None:
|
|
use_spawn = current_platform.is_xpu()
|
|
method = "spawn" if use_spawn else "fork"
|
|
assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
|
|
if method == "fork":
|
|
return fork_new_process_for_each_test
|
|
return spawn_new_process_for_each_test
|