ToolAgent/server.py

71 lines
2.0 KiB
Python

import os, sys
import asyncio
import subprocess
import threading
from pathlib import Path
import transformers
from transformers import AddedToken, PreTrainedTokenizerBase
# ── monkey-patch: make tokenizer encode thread-safe ──
_tokenizer_lock = threading.Lock()
_orig_encode = PreTrainedTokenizerBase.encode
def _safe_encode(self, *args, **kwargs):
with _tokenizer_lock:
return _orig_encode(self, *args, **kwargs)
PreTrainedTokenizerBase.encode = _safe_encode
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/Qwen3-VL-8B-Instruct")
parser = make_arg_parser(FlexibleArgumentParser())
args = parser.parse_args([
"--model", model,
"--host", "0.0.0.0",
"--port", "8000",
"--served-model-name", "Qwen3-VL-8B-Instruct",
"--trust-remote-code",
"--max-model-len", "32768",
"--dtype", "bfloat16",
"--gpu-memory-utilization", "0.70",
"--enforce-eager",
"--enable-auto-tool-choice",
"--tool-call-parser", "hermes",
"--disable-frontend-multiprocessing",
"--skip-mm-profiling",
])
asyncio.run(run_server(args))