55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
import os, sys
|
|
import asyncio
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
|
from vllm.entrypoints.openai.api_server import run_server
|
|
from vllm.entrypoints.openai.cli_args import make_arg_parser
|
|
|
|
|
|
def _pick_free_gpu():
|
|
if "CUDA_VISIBLE_DEVICES" in os.environ:
|
|
return os.environ["CUDA_VISIBLE_DEVICES"]
|
|
try:
|
|
|
|
out = subprocess.check_output(
|
|
["nvidia-smi", "--query-gpu=index,memory.free", "--format=csv,noheader"],
|
|
text=True,
|
|
)
|
|
lines = [l.strip() for l in out.strip().split("\n") if l.strip()]
|
|
best = max(lines, key=lambda x: int(x.split(",")[1].strip().split()[0]))
|
|
gpu = best.split(",")[0].strip()
|
|
except Exception:
|
|
gpu = "0"
|
|
|
|
return gpu
|
|
|
|
|
|
if __name__ == "__main__":
|
|
gpu = _pick_free_gpu()
|
|
if len(sys.argv) > 1:
|
|
gpu = sys.argv[1]
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
|
|
os.environ["FLASHINFER_DISABLE_VERSION_CHECK"] = "1"
|
|
|
|
model = str(Path(__file__).resolve().parent / "models/Qwen2.5-7B-Instruct")
|
|
|
|
parser = make_arg_parser(FlexibleArgumentParser())
|
|
args = parser.parse_args([
|
|
"--model", model,
|
|
"--host", "0.0.0.0",
|
|
"--port", "8000",
|
|
"--served-model-name", "Qwen2.5-7B-Instruct",
|
|
"--trust-remote-code",
|
|
"--max-model-len", "32768",
|
|
"--dtype", "bfloat16",
|
|
"--gpu-memory-utilization", "0.85",
|
|
"--enforce-eager",
|
|
"--enable-auto-tool-choice",
|
|
"--tool-call-parser", "hermes",
|
|
])
|
|
|
|
asyncio.run(run_server(args))
|