diff --git a/scripts/start_vllm.py b/scripts/start_vllm.py index 8c849a8..27206b9 100644 --- a/scripts/start_vllm.py +++ b/scripts/start_vllm.py @@ -3,428 +3,199 @@ import sys import os import json import shutil -import tempfile import subprocess from pathlib import Path -# Add script dir to path +# Configuration 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") -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") - -# 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") PORT = os.getenv("PORT", "8000") -def check_dependencies(): - if not shutil.which("dialog"): - print("Error: 'dialog' is required. Please install it (apt-get install dialog).") +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(): - """Detects AMD GPUs via rocm-smi or /dev/dri.""" + """Detect AMD GPUs""" 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 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"]) - } - + 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: - return default_config + 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 run_dialog(args): - """Runs dialog and returns stderr (selection).""" - with tempfile.NamedTemporaryFile(mode="w+") as tf: - cmd = ["dialog"] + args - try: - subprocess.run(cmd, stderr=tf, check=True) - tf.seek(0) - return tf.read().strip() - except subprocess.CalledProcessError: - return None # User cancelled +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 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.") - time.sleep(1) - except Exception as e: - print(f" Failed: {e}") - -def configure_and_launch(model_idx, gpu_count, use_default=False): - print(f"DEBUG: configure_and_launch called with model_idx={model_idx}, gpu_count={gpu_count}, use_default={use_default}") - model_id = MODELS_TO_RUN[model_idx] - config = MODEL_TABLE[model_id] - print(f"DEBUG: model_id={model_id}") - print(f"DEBUG: LOCAL_MODEL_DIR={LOCAL_MODEL_DIR}") +def launch_model(model_id, config, model_path, gpu_count): + """Launch vLLM server""" + log(f"Launching model: {model_id}") - # Determine whether we have a local copy to serve. Try multiple fallbacks: - # 1) LOCAL_MODEL_DIR// - # 2) LOCAL_MODEL_DIR/ - # 3) case-insensitive match of 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) - candidate_full = os.path.join(LOCAL_MODEL_DIR, model_id) - print(f"DEBUG: Checking candidate_full: {candidate_full}") - print(f"DEBUG: candidate_full exists: {os.path.isdir(candidate_full)}") - if os.path.isdir(candidate_full): - 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}") + # Get configuration + valid_tp = config.get("valid_tp", [1]) + max_tp = max(valid_tp) if valid_tp else 1 + tp_size = min(gpu_count, max_tp) - print(f"DEBUG: Final model_path: {model_path}") - print(f"DEBUG: model_path == model_id: {model_path == model_id}") + ctx = int(config.get("ctx", 8192)) + max_seqs = int(config.get("max_num_seqs", 64)) + gpu_util = float(config.get("gpu_util", 0.98)) - # if LOCAL_MODEL_DIR is specified, refuse to fall back to remote - if LOCAL_MODEL_DIR and model_path == model_id: - print(f"Error: model '{model_id}' not found under LOCAL_MODEL_DIR={LOCAL_MODEL_DIR}") - print("Off‑line mode active; network downloads are disabled.") - sys.exit(1) - - # Static Config - valid_tps = config.get("valid_tp", [1]) - max_tp = max(valid_tps) if valid_tps else 1 - - # Defaults - current_tp = min(gpu_count, max_tp) - current_seqs = 1 # Default to 1 concurrent user/request for stability - - # Initial Lookup - verified = get_verified_config(model_id, current_tp, current_seqs) - current_ctx = verified["ctx"] - current_util = verified["util"] - - clear_cache = False - use_eager = config.get("enforce_eager", False) # Default to model config, usually False - use_rocm_attn = False # Default to Triton - - name = model_id.split("/")[-1] - - # 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() + log(f"Config: TP={tp_size}, Ctx={ctx}, Seqs={max_seqs}, Util={gpu_util}") + # Build command cmd = [ "vllm", "serve", model_path, "--host", HOST, "--port", PORT, - "--tensor-parallel-size", str(current_tp), - "--max-num-seqs", str(current_seqs), - "--max-model-len", str(current_ctx), - "--gpu-memory-utilization", str(current_util), + "--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 use_eager: cmd.append("--enforce-eager") + if config.get("trust_remote"): + 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.update(config.get("env", {})) - if use_rocm_attn: - env["VLLM_V1_USE_PREFILL_DECODE_ATTENTION"] = "1" - 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 + # Launch vLLM + log("Starting vLLM server...") try: - result = subprocess.run(cmd, env=env, check=False) + result = subprocess.run(cmd, env=env) 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) 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) def main(): - check_dependencies() - gpu_count = detect_gpus() + """Main entry point""" + 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" if use_default: - # Find the index of default model - try: - 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) + log("Using default model mode") + model_id = default_model else: - while True: - # Build Model Menu - menu_items = [] - 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([ - "--clear", "--backtitle", f"AMD R9700 vLLM Launcher (GPUs: {gpu_count})", - "--title", "Select Model", - "--menu", "Choose a model to serve:", "20", "60", "10" - ] + menu_items) - - if not choice: - subprocess.run(["clear"]) - print("Selection cancelled.") - sys.exit(0) - - configure_and_launch(int(choice), gpu_count) + # 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()