253 lines
8.2 KiB
Python
253 lines
8.2 KiB
Python
#!/usr/bin/env python3
|
|
import sys
|
|
import os
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
# Configuration
|
|
SCRIPT_DIR = Path("/opt/script")
|
|
CONFIG_PATH = Path("/config/config.json")
|
|
LOCAL_MODEL_DIR = os.getenv("LOCAL_MODEL_DIR", "/opt/model")
|
|
HOST = os.getenv("HOST", "0.0.0.0")
|
|
PORT = os.getenv("PORT", "8000")
|
|
|
|
def log(msg):
|
|
"""Print log message with timestamp"""
|
|
print(f"[START-VLLM] {msg}", flush=True)
|
|
|
|
def load_config():
|
|
"""Load configuration from config.json"""
|
|
log(f"Loading config from {CONFIG_PATH}")
|
|
|
|
if not CONFIG_PATH.exists():
|
|
log(f"ERROR: Config file not found at {CONFIG_PATH}")
|
|
sys.exit(1)
|
|
|
|
try:
|
|
with open(CONFIG_PATH, "r") as f:
|
|
config_data = json.load(f)
|
|
|
|
model_table = config_data["models"]
|
|
default_model = config_data["default_model"]
|
|
models_to_run = list(model_table.keys())
|
|
|
|
log(f"Loaded {len(models_to_run)} models from config")
|
|
log(f"Default model: {default_model}")
|
|
|
|
return model_table, default_model, models_to_run
|
|
except Exception as e:
|
|
log(f"ERROR: Failed to load config: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
sys.exit(1)
|
|
|
|
def detect_gpus():
|
|
"""Detect AMD GPUs"""
|
|
try:
|
|
result = subprocess.run(
|
|
["rocm-smi", "--showid", "--csv"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=10
|
|
)
|
|
if result.returncode == 0:
|
|
count = result.stdout.count("GPU")
|
|
if count > 0:
|
|
return count
|
|
except Exception as e:
|
|
log(f"Warning: rocm-smi failed: {e}")
|
|
|
|
# Fallback to /dev/dri
|
|
try:
|
|
render_devices = list(Path("/dev/dri").glob("renderD*"))
|
|
if render_devices:
|
|
return len(render_devices)
|
|
except Exception:
|
|
pass
|
|
|
|
log("Warning: Could not detect GPUs, assuming 1 GPU")
|
|
return 1
|
|
|
|
def find_model_path(model_id):
|
|
"""Find local model path"""
|
|
log(f"Looking for model: {model_id}")
|
|
log(f"LOCAL_MODEL_DIR: {LOCAL_MODEL_DIR}")
|
|
|
|
if not os.path.exists(LOCAL_MODEL_DIR):
|
|
log(f"ERROR: LOCAL_MODEL_DIR does not exist: {LOCAL_MODEL_DIR}")
|
|
return None
|
|
|
|
# Try exact match
|
|
candidate = os.path.join(LOCAL_MODEL_DIR, model_id)
|
|
if os.path.isdir(candidate):
|
|
log(f"Found model at: {candidate}")
|
|
return candidate
|
|
|
|
# Try without prefix
|
|
repo_name = model_id.split('/')[-1] if '/' in model_id else model_id
|
|
candidate = os.path.join(LOCAL_MODEL_DIR, repo_name)
|
|
if os.path.isdir(candidate):
|
|
log(f"Found model at: {candidate}")
|
|
return candidate
|
|
|
|
# Try case-insensitive match
|
|
try:
|
|
for entry in os.listdir(LOCAL_MODEL_DIR):
|
|
if entry.lower() == repo_name.lower():
|
|
entry_path = os.path.join(LOCAL_MODEL_DIR, entry)
|
|
if os.path.isdir(entry_path):
|
|
log(f"Found model at: {entry_path}")
|
|
return entry_path
|
|
except Exception as e:
|
|
log(f"ERROR: Failed to list directory: {e}")
|
|
|
|
log(f"ERROR: Model not found: {model_id}")
|
|
log(f"Available models: {os.listdir(LOCAL_MODEL_DIR)}")
|
|
return None
|
|
|
|
def launch_model(model_id, config, model_path, gpu_count):
|
|
"""Launch vLLM server"""
|
|
log(f"Launching model: {model_id}")
|
|
|
|
# Get configuration
|
|
valid_tp = config.get("valid_tp", [1])
|
|
max_tp = max(valid_tp) if valid_tp else 1
|
|
|
|
# Check for manual TP_SIZE override
|
|
tp_size_env = os.getenv("TP_SIZE")
|
|
if tp_size_env:
|
|
tp_size = int(tp_size_env)
|
|
log(f"TP_SIZE environment variable set: {tp_size}")
|
|
if tp_size not in valid_tp:
|
|
log(f"WARNING: TP_SIZE={tp_size} is not in valid_tp={valid_tp}, proceeding anyway")
|
|
else:
|
|
tp_size = min(gpu_count, max_tp)
|
|
|
|
ctx = int(config.get("ctx", 8192))
|
|
max_seqs = int(config.get("max_num_seqs", 64))
|
|
gpu_util = float(config.get("gpu_util", 0.98))
|
|
served_model_name = config.get("served_model_name", model_id)
|
|
|
|
log(f"Config: TP={tp_size}, Ctx={ctx}, Seqs={max_seqs}, Util={gpu_util}")
|
|
|
|
# Build command
|
|
cmd = [
|
|
"vllm", "serve", model_path,
|
|
"--served-model-name", served_model_name,
|
|
"--host", HOST,
|
|
"--port", PORT,
|
|
"--tensor-parallel-size", str(tp_size),
|
|
"--max-num-seqs", str(max_seqs),
|
|
"--max-model-len", str(ctx),
|
|
"--gpu-memory-utilization", str(gpu_util),
|
|
"--dtype", "auto"
|
|
]
|
|
|
|
if config.get("trust_remote"):
|
|
cmd.append("--trust-remote-code")
|
|
|
|
if config.get("enforce_eager"):
|
|
cmd.append("--enforce-eager")
|
|
|
|
# Add Qwen3.5 specific parameters
|
|
if "qwen3.5" in model_id.lower() or "qwen3_5" in model_id.lower():
|
|
cmd.extend(["--quantization", "moe_wna16"])
|
|
cmd.extend(["--reasoning-parser", "qwen3"])
|
|
log("Added Qwen3.5 specific parameters: --quantization moe_wna16 --reasoning-parser qwen3")
|
|
|
|
# Add tool call parser if specified
|
|
tool_call_parser = config.get("tool_call_parser")
|
|
openclaw_compat = os.getenv("OPENCLAW_COMPAT", "false").lower() == "true"
|
|
if openclaw_compat and not tool_call_parser:
|
|
tool_call_parser = os.getenv("OPENCLAW_TOOL_CALL_PARSER", "qwen3_xml")
|
|
if tool_call_parser:
|
|
cmd.extend(["--tool-call-parser", tool_call_parser])
|
|
enable_auto_tool_choice = config.get("enable_auto_tool_choice")
|
|
auto_tool_choice_env = os.getenv("ENABLE_AUTO_TOOL_CHOICE")
|
|
if auto_tool_choice_env is not None:
|
|
enable_auto_tool_choice = auto_tool_choice_env.lower() == "true"
|
|
if enable_auto_tool_choice is None:
|
|
enable_auto_tool_choice = True
|
|
if openclaw_compat:
|
|
enable_auto_tool_choice = True
|
|
if enable_auto_tool_choice:
|
|
cmd.extend(["--enable-auto-tool-choice"])
|
|
log("Added auto tool choice enabled")
|
|
else:
|
|
log("Auto tool choice disabled")
|
|
log(f"Added tool call parser: {tool_call_parser}")
|
|
|
|
log(f"Command: {' '.join(cmd)}")
|
|
|
|
# Set environment
|
|
env = os.environ.copy()
|
|
env.update(config.get("env", {}))
|
|
|
|
# Add AMD GPU tensor parallel environment variables
|
|
if tp_size > 1:
|
|
env["NCCL_DEBUG"] = "INFO"
|
|
env["NCCL_SOCKET_IFNAME"] = "^lo,docker0"
|
|
env["NCCL_P2P_DISABLE"] = "0"
|
|
env["NCCL_SHM_DISABLE"] = "0"
|
|
env["NCCL_IB_DISABLE"] = "1"
|
|
env["NCCL_P2P_LEVEL"] = "SYS"
|
|
env["HIP_VISIBLE_DEVICES"] = os.getenv("HIP_VISIBLE_DEVICES", "0,1")
|
|
log("Added NCCL environment variables for AMD GPU tensor parallel")
|
|
log(f"HIP_VISIBLE_DEVICES: {env['HIP_VISIBLE_DEVICES']}")
|
|
|
|
# Launch vLLM
|
|
log("Starting vLLM server...")
|
|
try:
|
|
result = subprocess.run(cmd, env=env)
|
|
if result.returncode != 0:
|
|
log(f"ERROR: vLLM exited with code {result.returncode}")
|
|
sys.exit(result.returncode)
|
|
except Exception as e:
|
|
log(f"ERROR: Failed to start vLLM: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
sys.exit(1)
|
|
|
|
def main():
|
|
"""Main entry point"""
|
|
log("Starting vLLM launcher...")
|
|
|
|
# Load configuration
|
|
model_table, default_model, models_to_run = load_config()
|
|
|
|
# Detect GPUs
|
|
gpu_count = detect_gpus()
|
|
log(f"Detected {gpu_count} GPU(s)")
|
|
|
|
# Check if we should use default model
|
|
use_default = os.getenv("USE_DEFAULT_MODEL", "false").lower() == "true"
|
|
|
|
if use_default:
|
|
log("Using default model mode")
|
|
model_id = default_model
|
|
else:
|
|
# Interactive mode - for now just use default
|
|
log("Interactive mode not supported in container, using default model")
|
|
model_id = default_model
|
|
|
|
# Check if model is in config
|
|
if model_id not in model_table:
|
|
log(f"ERROR: Model {model_id} not found in config")
|
|
sys.exit(1)
|
|
|
|
config = model_table[model_id]
|
|
|
|
# Find model path
|
|
model_path = find_model_path(model_id)
|
|
if not model_path:
|
|
log("ERROR: Could not find local model. Offline mode is active.")
|
|
sys.exit(1)
|
|
|
|
# Launch model
|
|
launch_model(model_id, config, model_path, gpu_count)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|