init
This commit is contained in:
@@ -0,0 +1,235 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user