This commit is contained in:
2026-03-27 01:32:30 +08:00
parent 90fad86337
commit d5fbf5e5b4
+140 -369
View File
@@ -3,428 +3,199 @@ import sys
import os import os
import json import json
import shutil import shutil
import tempfile
import subprocess import subprocess
from pathlib import Path from pathlib import Path
# Add script dir to path # Configuration
SCRIPT_DIR = Path("/opt/script") SCRIPT_DIR = Path("/opt/script")
OPT_DIR = Path("/opt")
# Config file path (check container path first, then local path)
CONFIG_PATH = Path("/config/config.json") CONFIG_PATH = Path("/config/config.json")
if not CONFIG_PATH.exists():
CONFIG_PATH = Path("/config.json")
# Local model directory (container path)
LOCAL_MODEL_DIR = os.getenv("LOCAL_MODEL_DIR", "/opt/model") LOCAL_MODEL_DIR = os.getenv("LOCAL_MODEL_DIR", "/opt/model")
# Load configuration from config.json
print(f"DEBUG: Looking for config at {CONFIG_PATH}")
print(f"DEBUG: Config file exists: {CONFIG_PATH.exists()}")
if not CONFIG_PATH.exists():
print(f"ERROR: Config file not found at {CONFIG_PATH}")
print(f"ERROR: Please mount config.json to /config/config.json")
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())
print(f"DEBUG: Loaded {len(MODELS_TO_RUN)} models from config")
print(f"DEBUG: Default model: {DEFAULT_MODEL}")
except Exception as e:
print(f"Error: Could not load config.json: {e}")
import traceback
traceback.print_exc()
sys.exit(1)
# Results file path
if (OPT_DIR / "max_context_results.json").exists():
RESULTS_FILE = OPT_DIR / "max_context_results.json"
else:
RESULTS_FILE = SCRIPT_DIR / "max_context_results.json"
HOST = os.getenv("HOST", "0.0.0.0") HOST = os.getenv("HOST", "0.0.0.0")
PORT = os.getenv("PORT", "8000") PORT = os.getenv("PORT", "8000")
def check_dependencies(): def log(msg):
if not shutil.which("dialog"): """Print log message with timestamp"""
print("Error: 'dialog' is required. Please install it (apt-get install dialog).") 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) sys.exit(1)
def detect_gpus(): def detect_gpus():
"""Detects AMD GPUs via rocm-smi or /dev/dri.""" """Detect AMD GPUs"""
try: try:
# Try rocm-smi first result = subprocess.run(
res = subprocess.run(["rocm-smi", "--showid", "--csv"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) ["rocm-smi", "--showid", "--csv"],
if res.returncode == 0: capture_output=True,
count = res.stdout.count("GPU") text=True,
if count > 0: return count timeout=10
except: pass )
if result.returncode == 0:
# Fallback to /dev/dri/render* count = result.stdout.count("GPU")
try: if count > 0:
return len(list(Path("/dev/dri").glob("renderD*"))) return count
except:
return 1
def get_verified_config(model_id, tp_size, max_seqs):
"""
Reads max_context_results.json to find the best verified configuration.
Returns dict: {'ctx': int, 'util': float}
"""
default_config = {
"ctx": int(MODEL_TABLE.get(model_id, {}).get("ctx", 8192)),
"util": 0.90 # Safe default
}
if not RESULTS_FILE.exists():
return default_config
try:
with open(RESULTS_FILE, "r") as f:
data = json.load(f)
# Filter for Model + TP + Sequences
matches = [r for r in data
if r["model"] == model_id
and r["tp"] == tp_size
and r["max_seqs"] == max_seqs
and r["status"] == "success"]
if not matches:
# Fallback 1: Try finding match with SAME TP but ANY Sequences (e.g. 1) to get base context?
# Actually, safer to fallback to default or try finding nearest sequence?
# Let's try finding exact match first. If fail, return default.
return default_config
# Sort by Util desc, then Context desc
# We prefer higher utilization if available (performance), as long as it is verified success
matches.sort(key=lambda x: (float(x["util"]), x["max_context_1_user"]), reverse=True)
best = matches[0]
return {
"ctx": best["max_context_1_user"],
"util": float(best["util"])
}
except Exception as e: except Exception as e:
return default_config log(f"Warning: rocm-smi failed: {e}")
def run_dialog(args): # Fallback to /dev/dri
"""Runs dialog and returns stderr (selection).""" try:
with tempfile.NamedTemporaryFile(mode="w+") as tf: render_devices = list(Path("/dev/dri").glob("renderD*"))
cmd = ["dialog"] + args if render_devices:
try: return len(render_devices)
subprocess.run(cmd, stderr=tf, check=True) except Exception:
tf.seek(0) pass
return tf.read().strip()
except subprocess.CalledProcessError:
return None # User cancelled
def nuke_vllm_cache(): log("Warning: Could not detect GPUs, assuming 1 GPU")
"""Removes vLLM cache directory to fix potential graph/incompatibility issues.""" return 1
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.")
time.sleep(1)
except Exception as e:
print(f" Failed: {e}")
def configure_and_launch(model_idx, gpu_count, use_default=False): def find_model_path(model_id):
print(f"DEBUG: configure_and_launch called with model_idx={model_idx}, gpu_count={gpu_count}, use_default={use_default}") """Find local model path"""
model_id = MODELS_TO_RUN[model_idx] log(f"Looking for model: {model_id}")
config = MODEL_TABLE[model_id] log(f"LOCAL_MODEL_DIR: {LOCAL_MODEL_DIR}")
print(f"DEBUG: model_id={model_id}")
print(f"DEBUG: LOCAL_MODEL_DIR={LOCAL_MODEL_DIR}")
# Determine whether we have a local copy to serve. Try multiple fallbacks: if not os.path.exists(LOCAL_MODEL_DIR):
# 1) LOCAL_MODEL_DIR/<owner>/<repo> log(f"ERROR: LOCAL_MODEL_DIR does not exist: {LOCAL_MODEL_DIR}")
# 2) LOCAL_MODEL_DIR/<repo> return None
# 3) case-insensitive match of <repo> in LOCAL_MODEL_DIR
model_path = model_id
print(f"DEBUG: Starting model path lookup...")
if LOCAL_MODEL_DIR:
print(f"DEBUG: LOCAL_MODEL_DIR is set to: {LOCAL_MODEL_DIR}")
print(f"DEBUG: LOCAL_MODEL_DIR exists: {os.path.exists(LOCAL_MODEL_DIR)}")
if os.path.exists(LOCAL_MODEL_DIR):
print(f"DEBUG: LOCAL_MODEL_DIR contents: {os.listdir(LOCAL_MODEL_DIR)}")
# Full repo path (owner/repo) # Try exact match
candidate_full = os.path.join(LOCAL_MODEL_DIR, model_id) candidate = os.path.join(LOCAL_MODEL_DIR, model_id)
print(f"DEBUG: Checking candidate_full: {candidate_full}") if os.path.isdir(candidate):
print(f"DEBUG: candidate_full exists: {os.path.isdir(candidate_full)}") log(f"Found model at: {candidate}")
if os.path.isdir(candidate_full): return candidate
model_path = candidate_full
print(f"DEBUG: Found model at: {model_path}")
else:
# Repo-name only (last segment)
repo_name = model_id.split('/')[-1]
candidate_repo = os.path.join(LOCAL_MODEL_DIR, repo_name)
print(f"DEBUG: Checking candidate_repo: {candidate_repo}")
print(f"DEBUG: candidate_repo exists: {os.path.isdir(candidate_repo)}")
if os.path.isdir(candidate_repo):
model_path = candidate_repo
print(f"DEBUG: Found model at: {model_path}")
else:
# Fallback: try to find a directory in LOCAL_MODEL_DIR that matches repo_name case-insensitively
print(f"DEBUG: Trying case-insensitive match for: {repo_name}")
try:
for entry in os.listdir(LOCAL_MODEL_DIR):
print(f"DEBUG: Checking entry: {entry}")
if entry.lower() == repo_name.lower():
entry_path = os.path.join(LOCAL_MODEL_DIR, entry)
if os.path.isdir(entry_path):
model_path = entry_path
print(f"DEBUG: Found model at: {model_path}")
break
except Exception as e:
print(f"DEBUG: Exception during case-insensitive lookup: {e}")
print(f"DEBUG: Final model_path: {model_path}") # Try without prefix
print(f"DEBUG: model_path == model_id: {model_path == model_id}") 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
# if LOCAL_MODEL_DIR is specified, refuse to fall back to remote # Try case-insensitive match
if LOCAL_MODEL_DIR and model_path == model_id: try:
print(f"Error: model '{model_id}' not found under LOCAL_MODEL_DIR={LOCAL_MODEL_DIR}") for entry in os.listdir(LOCAL_MODEL_DIR):
print("Off‑line mode active; network downloads are disabled.") if entry.lower() == repo_name.lower():
sys.exit(1) 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}")
# Static Config log(f"ERROR: Model not found: {model_id}")
valid_tps = config.get("valid_tp", [1]) log(f"Available models: {os.listdir(LOCAL_MODEL_DIR)}")
max_tp = max(valid_tps) if valid_tps else 1 return None
# Defaults def launch_model(model_id, config, model_path, gpu_count):
current_tp = min(gpu_count, max_tp) """Launch vLLM server"""
current_seqs = 1 # Default to 1 concurrent user/request for stability log(f"Launching model: {model_id}")
# Initial Lookup # Get configuration
verified = get_verified_config(model_id, current_tp, current_seqs) valid_tp = config.get("valid_tp", [1])
current_ctx = verified["ctx"] max_tp = max(valid_tp) if valid_tp else 1
current_util = verified["util"] tp_size = min(gpu_count, max_tp)
clear_cache = False ctx = int(config.get("ctx", 8192))
use_eager = config.get("enforce_eager", False) # Default to model config, usually False max_seqs = int(config.get("max_num_seqs", 64))
use_rocm_attn = False # Default to Triton gpu_util = float(config.get("gpu_util", 0.98))
name = model_id.split("/")[-1] log(f"Config: TP={tp_size}, Ctx={ctx}, Seqs={max_seqs}, Util={gpu_util}")
# If use_default is True, skip interactive menu and launch directly
if use_default:
print(f"DEBUG: use_default=True, skipping interactive menu")
print(f"DEBUG: Using default config: TP={current_tp}, Seqs={current_seqs}, Ctx={current_ctx}")
# Jump directly to launch
launch_server = True
else:
launch_server = False
while True:
cache_status = "YES" if clear_cache else "NO"
eager_status = "YES" if use_eager else "NO"
attn_backend = "ROCm" if use_rocm_attn else "Triton"
menu_args = [
"--clear", "--backtitle", f"AMD R9700 vLLM Launcher (GPUs: {gpu_count})",
"--title", f"Configuration: {name}",
"--menu", "Customize Launch Parameters:", "22", "65", "9",
"1", f"Tensor Parallelism: {current_tp}",
"2", f"Concurrent Requests: {current_seqs}",
"3", f"Context Length: {current_ctx} (Verified)",
"4", f"GPU Utilization: {current_util} (Verified)",
"5", f"Attention Backend: {attn_backend}",
"6", f"Erase vLLM Cache: {cache_status}",
"7", f"Force Eager Mode: {eager_status}",
"8", "LAUNCH SERVER"
]
choice = run_dialog(menu_args)
if not choice: return False # Back/Cancel
if choice == "1":
# TP Selection
new_tp = run_dialog([
"--title", "Tensor Parallelism",
"--rangebox", f"Set TP Size (1-{max_tp})", "10", "40", "1", str(max_tp), str(current_tp)
])
if new_tp:
new_tp_int = int(new_tp)
if new_tp_int != current_tp:
current_tp = new_tp_int
# RE-CALCULATE Config
verified = get_verified_config(model_id, current_tp, current_seqs)
current_ctx = verified["ctx"]
current_util = verified["util"]
elif choice == "2":
# Max Seqs Selection
new_seqs = run_dialog([
"--title", "Concurrent Requests",
"--menu", "Select Max Concurrent Requests:", "12", "40", "4",
"1", "1 (Latency Focus)",
"4", "4 (Balanced)",
"8", "8 (Throughput)",
"16", "16 (Max Load)"
])
if new_seqs:
current_seqs = int(new_seqs)
# RE-CALCULATE Config based on new concurrency
verified = get_verified_config(model_id, current_tp, current_seqs)
current_ctx = verified["ctx"]
current_util = verified["util"]
elif choice == "3":
# Configured Length Override
new_ctx = run_dialog([
"--title", "Context Length",
"--inputbox", f"Override verified limit ({current_ctx}):", "10", "40", str(current_ctx)
])
if new_ctx: current_ctx = int(new_ctx)
elif choice == "4":
# Util Override
pass
elif choice == "5":
# Toggle Attention Backend
use_rocm_attn = not use_rocm_attn
elif choice == "6":
# Toggle Cache
if not clear_cache:
# Enabling it -> Show Warning
warn_msg = (
"WARNING: Erasing the vLLM cache will remove the compiled compute graphs.\n\n"
"This is useful if you are experiencing crashes, 'invalid graph' errors,\n"
"or have switched vLLM versions recently.\n\n"
"However, the next startup will take longer as graphs are re-compiled.\n\n"
"Are you sure you want to enable this?"
)
confirm = run_dialog([
"--title", "Erase Cache Warning",
"--yesno", warn_msg, "12", "60"
])
# If confirm is not None (exit 0), it is YES.
if confirm is not None:
clear_cache = True
else:
# Disabling it -> No warning needed
clear_cache = False
elif choice == "7":
# Toggle Eager Mode
use_eager = not use_eager
elif choice == "8":
# Launch
launch_server = True
break
if launch_server:
break
# Build Command
subprocess.run(["clear"])
if clear_cache:
nuke_vllm_cache()
# Build command
cmd = [ cmd = [
"vllm", "serve", model_path, "vllm", "serve", model_path,
"--host", HOST, "--host", HOST,
"--port", PORT, "--port", PORT,
"--tensor-parallel-size", str(current_tp), "--tensor-parallel-size", str(tp_size),
"--max-num-seqs", str(current_seqs), "--max-num-seqs", str(max_seqs),
"--max-model-len", str(current_ctx), "--max-model-len", str(ctx),
"--gpu-memory-utilization", str(current_util), "--gpu-memory-utilization", str(gpu_util),
"--dtype", "auto" "--dtype", "auto"
] ]
if config.get("trust_remote"): cmd.append("--trust-remote-code") if config.get("trust_remote"):
if use_eager: cmd.append("--enforce-eager") cmd.append("--trust-remote-code")
# Env Vars if config.get("enforce_eager"):
cmd.append("--enforce-eager")
log(f"Command: {' '.join(cmd)}")
# Set environment
env = os.environ.copy() env = os.environ.copy()
env.update(config.get("env", {})) env.update(config.get("env", {}))
if use_rocm_attn: # Launch vLLM
env["VLLM_V1_USE_PREFILL_DECODE_ATTENTION"] = "1" log("Starting vLLM server...")
env["VLLM_USE_TRITON_FLASH_ATTN"] = "0"
# Optional: Explicitly mention these in print
print("\n" + "="*60)
print(f" Launching: {name}")
if model_path != model_id:
print(f" (using local model at {model_path})")
print(f" Config: TP={current_tp} | Seqs={current_seqs} | Ctx={current_ctx} | Util={current_util}")
print(f" Backend: {'ROCm' if use_rocm_attn else 'Triton'}")
if clear_cache:
print(f" Action: Clearing vLLM Cache (~/.cache/vllm)")
print(f" Command: {' '.join(cmd)}")
print("="*60 + "\n")
# Check if model path exists
if not os.path.exists(model_path):
print(f"ERROR: Model path does not exist: {model_path}")
print(f"Please ensure the model is mounted at {model_path}")
sys.exit(1)
# Run vllm serve
try: try:
result = subprocess.run(cmd, env=env, check=False) result = subprocess.run(cmd, env=env)
if result.returncode != 0: if result.returncode != 0:
print(f"\nERROR: vllm serve exited with code {result.returncode}") log(f"ERROR: vLLM exited with code {result.returncode}")
sys.exit(result.returncode) sys.exit(result.returncode)
except Exception as e: except Exception as e:
print(f"\nERROR: Failed to start vllm serve: {e}") log(f"ERROR: Failed to start vLLM: {e}")
import traceback
traceback.print_exc()
sys.exit(1) sys.exit(1)
def main(): def main():
check_dependencies() """Main entry point"""
gpu_count = detect_gpus() log("Starting vLLM launcher...")
# Check if we should use default model (for docker startup) # 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" use_default = os.getenv("USE_DEFAULT_MODEL", "false").lower() == "true"
if use_default: if use_default:
# Find the index of default model log("Using default model mode")
try: model_id = default_model
default_idx = MODELS_TO_RUN.index(DEFAULT_MODEL)
print(f"Using default model: {DEFAULT_MODEL}")
configure_and_launch(default_idx, gpu_count, use_default=True)
except ValueError:
print(f"Error: Default model {DEFAULT_MODEL} not found in configuration")
sys.exit(1)
else: else:
while True: # Interactive mode - for now just use default
# Build Model Menu log("Interactive mode not supported in container, using default model")
menu_items = [] model_id = default_model
for i, m_id in enumerate(MODELS_TO_RUN):
name = m_id.split("/")[-1]
# Mark default model
if m_id == DEFAULT_MODEL:
name += " (Default)"
menu_items.extend([str(i), name])
choice = run_dialog([ # Check if model is in config
"--clear", "--backtitle", f"AMD R9700 vLLM Launcher (GPUs: {gpu_count})", if model_id not in model_table:
"--title", "Select Model", log(f"ERROR: Model {model_id} not found in config")
"--menu", "Choose a model to serve:", "20", "60", "10" sys.exit(1)
] + menu_items)
if not choice: config = model_table[model_id]
subprocess.run(["clear"])
print("Selection cancelled.")
sys.exit(0)
configure_and_launch(int(choice), gpu_count) # 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__": if __name__ == "__main__":
main() main()