236 lines
7.5 KiB
Python
236 lines
7.5 KiB
Python
#!/usr/bin/env python3
|
|
import sys
|
|
import os
|
|
import json
|
|
import yaml
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
# Add benchmarks dir to path to import config
|
|
SCRIPT_DIR = Path(__file__).parent.resolve()
|
|
OPT_DIR = Path("/opt")
|
|
|
|
# Required environment variable pointing to a local models directory
|
|
LOCAL_MODEL_DIR = os.getenv("LOCAL_MODEL_DIR")
|
|
if not LOCAL_MODEL_DIR:
|
|
print("Error: LOCAL_MODEL_DIR environment variable is required.")
|
|
sys.exit(1)
|
|
|
|
# Configuration file path
|
|
CONFIG_FILE = os.getenv("VLLM_CONFIG_FILE", "/etc/vllm/model_config.yaml")
|
|
|
|
# Default configuration
|
|
DEFAULT_CONFIG = {
|
|
"default": "",
|
|
"models": {},
|
|
"server": {
|
|
"host": "0.0.0.0",
|
|
"log_level": "info"
|
|
}
|
|
}
|
|
|
|
def detect_gpus():
|
|
"""Detects AMD GPUs via rocm-smi or /dev/dri."""
|
|
try:
|
|
# Try rocm-smi first
|
|
res = subprocess.run(["rocm-smi", "--showid", "--csv"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
if res.returncode == 0:
|
|
count = res.stdout.count("GPU")
|
|
if count > 0: return count
|
|
except: pass
|
|
|
|
# Fallback to /dev/dri/render*
|
|
try:
|
|
return len(list(Path("/dev/dri").glob("renderD*")))
|
|
except:
|
|
return 1
|
|
|
|
def load_config():
|
|
"""Load configuration from YAML file."""
|
|
config = DEFAULT_CONFIG.copy()
|
|
|
|
if Path(CONFIG_FILE).exists():
|
|
try:
|
|
with open(CONFIG_FILE, "r", encoding="utf-8") as f:
|
|
user_config = yaml.safe_load(f)
|
|
if user_config:
|
|
config.update(user_config)
|
|
except Exception as e:
|
|
print(f"Warning: Failed to load config file: {e}")
|
|
print("Using default configuration.")
|
|
else:
|
|
print(f"Warning: Config file not found at {CONFIG_FILE}")
|
|
print("Using default configuration.")
|
|
|
|
return config
|
|
|
|
def nuke_vllm_cache():
|
|
"""Removes vLLM cache directory to fix potential graph/incompatibility issues."""
|
|
cache = Path.home() / ".cache" / "vllm"
|
|
if cache.exists():
|
|
try:
|
|
print(f"Clearing vLLM cache at {cache}...", end="", flush=True)
|
|
subprocess.run(["rm", "-rf", str(cache)], check=True)
|
|
cache.mkdir(parents=True, exist_ok=True)
|
|
print(" Done.")
|
|
except Exception as e:
|
|
print(f" Failed: {e}")
|
|
|
|
def get_local_model_path(model_name):
|
|
"""Get local model path from LOCAL_MODEL_DIR."""
|
|
# Try multiple paths:
|
|
# 1) LOCAL_MODEL_DIR/model_name
|
|
# 2) case-insensitive match in LOCAL_MODEL_DIR
|
|
|
|
# Exact match
|
|
candidate_exact = os.path.join(LOCAL_MODEL_DIR, model_name)
|
|
if os.path.isdir(candidate_exact):
|
|
return candidate_exact
|
|
|
|
# Case-insensitive match
|
|
try:
|
|
for entry in os.listdir(LOCAL_MODEL_DIR):
|
|
if entry.lower() == model_name.lower():
|
|
entry_path = os.path.join(LOCAL_MODEL_DIR, entry)
|
|
if os.path.isdir(entry_path):
|
|
return entry_path
|
|
except Exception as e:
|
|
print(f"Error searching for model: {e}")
|
|
|
|
return None
|
|
|
|
def verify_model_path(model_path):
|
|
"""Verify that the model path contains required files."""
|
|
required = ["config.json", "pytorch_model.bin", "model.safetensors"]
|
|
found = any(os.path.isfile(os.path.join(model_path, f)) for f in required)
|
|
if not found:
|
|
print(f"Error: local model dir {model_path} missing expected files {required}")
|
|
return False
|
|
return True
|
|
|
|
def main():
|
|
# Load configuration
|
|
config = load_config()
|
|
|
|
# Get model name from command line or use default
|
|
model_name = None
|
|
if len(sys.argv) > 1:
|
|
model_name = sys.argv[1]
|
|
|
|
if not model_name:
|
|
model_name = config.get("default", "")
|
|
if not model_name:
|
|
print("Error: Either specify a model name as argument or set 'default' in config.")
|
|
sys.exit(1)
|
|
|
|
# Get model configuration
|
|
model_config = config.get("models", {}).get(model_name)
|
|
if not model_config:
|
|
print(f"Error: Model '{model_name}' not found in configuration.")
|
|
sys.exit(1)
|
|
|
|
# Get model path
|
|
model_path = model_config.get("path")
|
|
if not model_path:
|
|
# Fallback to LOCAL_MODEL_DIR if path not specified
|
|
model_path = get_local_model_path(model_name)
|
|
if not model_path:
|
|
print(f"Error: model '{model_name}' not found under LOCAL_MODEL_DIR={LOCAL_MODEL_DIR}")
|
|
sys.exit(1)
|
|
|
|
# Verify model path
|
|
if not verify_model_path(model_path):
|
|
sys.exit(1)
|
|
|
|
# Detect GPU count
|
|
gpu_count = detect_gpus()
|
|
|
|
# Get server configuration
|
|
server_config = config.get("server", {})
|
|
host = server_config.get("host", "0.0.0.0")
|
|
|
|
# Build command
|
|
cmd = ["vllm", "serve", model_path]
|
|
|
|
# Add server parameters
|
|
cmd.extend(["--host", host])
|
|
if "port" in model_config:
|
|
cmd.extend(["--port", str(model_config["port"])])
|
|
|
|
# Add API key for OpenAI compatibility
|
|
api_key = model_config.get("api_key")
|
|
if api_key:
|
|
cmd.extend(["--api-key", api_key])
|
|
|
|
# Add model parameters
|
|
if "tensor_parallel_size" in model_config:
|
|
tp_size = min(model_config["tensor_parallel_size"], gpu_count)
|
|
cmd.extend(["--tensor-parallel-size", str(tp_size)])
|
|
|
|
if "max_num_seqs" in model_config:
|
|
cmd.extend(["--max-num-seqs", str(model_config["max_num_seqs"])])
|
|
|
|
if "max_model_len" in model_config:
|
|
cmd.extend(["--max-model-len", str(model_config["max_model_len"])])
|
|
|
|
if "gpu_memory_utilization" in model_config:
|
|
cmd.extend(["--gpu-memory-utilization", str(model_config["gpu_memory_utilization"])])
|
|
|
|
if "dtype" in model_config:
|
|
cmd.extend(["--dtype", model_config["dtype"]])
|
|
else:
|
|
cmd.extend(["--dtype", "auto"])
|
|
|
|
if "max_num_batched_tokens" in model_config:
|
|
cmd.extend(["--max-num-batched-tokens", str(model_config["max_num_batched_tokens"])])
|
|
|
|
if "block_size" in model_config:
|
|
cmd.extend(["--block-size", str(model_config["block_size"])])
|
|
|
|
if "swap_space" in model_config:
|
|
cmd.extend(["--swap-space", str(model_config["swap_space"])])
|
|
|
|
if model_config.get("enforce_eager", False):
|
|
cmd.append("--enforce-eager")
|
|
|
|
if model_config.get("trust_remote", False):
|
|
cmd.append("--trust-remote-code")
|
|
|
|
# Set environment variables
|
|
env = os.environ.copy()
|
|
|
|
if model_config.get("use_rocm_attn", False):
|
|
env["VLLM_V1_USE_PREFILL_DECODE_ATTENTION"] = "1"
|
|
env["VLLM_USE_TRITON_FLASH_ATTN"] = "0"
|
|
|
|
# Clear cache if requested
|
|
if model_config.get("clear_cache", False):
|
|
nuke_vllm_cache()
|
|
|
|
# Print configuration
|
|
print("\n" + "="*60)
|
|
print(f" Launching: {model_name}")
|
|
print(f" Model Path: {model_path}")
|
|
print(f" Host: {host}")
|
|
if "port" in model_config:
|
|
print(f" Port: {model_config['port']}")
|
|
if api_key:
|
|
print(f" API Key: {api_key[:8]}...{api_key[-4:]}")
|
|
if "tensor_parallel_size" in model_config:
|
|
print(f" TP Size: {tp_size}")
|
|
if "max_num_seqs" in model_config:
|
|
print(f" Max Seqs: {model_config['max_num_seqs']}")
|
|
if "max_model_len" in model_config:
|
|
print(f" Max Ctx: {model_config['max_model_len']}")
|
|
if "gpu_memory_utilization" in model_config:
|
|
print(f" GPU Util: {model_config['gpu_memory_utilization']}")
|
|
print(f" Command: {' '.join(cmd)}")
|
|
print("="*60 + "\n")
|
|
|
|
# Launch vLLM server
|
|
os.execvpe("vllm", cmd, env)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|