x

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")
+HOST = os.getenv("HOST", "0.0.0.0")
+PORT = os.getenv("PORT", "8000")
 
-# Load configuration from config.json
+def log(msg):
-print(f"DEBUG: Looking for config at {CONFIG_PATH}")
+    """Print log message with timestamp"""
-print(f"DEBUG: Config file exists: {CONFIG_PATH.exists()}")
+    print(f"[START-VLLM] {msg}", flush=True)
-if not CONFIG_PATH.exists():
+
-    print(f"ERROR: Config file not found at {CONFIG_PATH}")
+def load_config():
-    print(f"ERROR: Please mount config.json to /config/config.json")
+    """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:
+    try:
         with open(CONFIG_PATH, "r") as f:
             config_data = json.load(f)
-    MODEL_TABLE = config_data["models"]
+        
-    DEFAULT_MODEL = config_data["default_model"]
+        model_table = config_data["models"]
-    MODELS_TO_RUN = list(MODEL_TABLE.keys())
+        default_model = config_data["default_model"]
-    print(f"DEBUG: Loaded {len(MODELS_TO_RUN)} models from config")
+        models_to_run = list(model_table.keys())
-    print(f"DEBUG: Default model: {DEFAULT_MODEL}")
+        
-except Exception as e:
+        log(f"Loaded {len(models_to_run)} models from config")
-    print(f"Error: Could not load config.json: {e}")
+        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)
 
-# 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).")
-        sys.exit(1)
-
 def detect_gpus():
-    """Detects AMD GPUs via rocm-smi or /dev/dri."""
+    """Detect AMD GPUs"""
     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:
+            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/render*
+    # Fallback to /dev/dri
     try:
-        return len(list(Path("/dev/dri").glob("renderD*")))
+        render_devices = list(Path("/dev/dri").glob("renderD*"))
-    except:
+        if render_devices:
+            return len(render_devices)
+    except Exception:
+        pass
+    
+    log("Warning: Could not detect GPUs, assuming 1 GPU")
     return 1
 
-def get_verified_config(model_id, tp_size, max_seqs):
+def find_model_path(model_id):
-    """
+    """Find local model path"""
-    Reads max_context_results.json to find the best verified configuration.
+    log(f"Looking for model: {model_id}")
-    Returns dict: {'ctx': int, 'util': float}
+    log(f"LOCAL_MODEL_DIR: {LOCAL_MODEL_DIR}")
-    """
-    default_config = {
-        "ctx": int(MODEL_TABLE.get(model_id, {}).get("ctx", 8192)),
-        "util": 0.90 # Safe default
-    }
     
-    if not RESULTS_FILE.exists():
+    if not os.path.exists(LOCAL_MODEL_DIR):
-        return default_config
+        log(f"ERROR: LOCAL_MODEL_DIR does not exist: {LOCAL_MODEL_DIR}")
+        return None
     
-    try:
+    # Try exact match
-        with open(RESULTS_FILE, "r") as f:
+    candidate = os.path.join(LOCAL_MODEL_DIR, model_id)
-            data = json.load(f)
+    if os.path.isdir(candidate):
+        log(f"Found model at: {candidate}")
+        return candidate
     
-        # Filter for Model + TP + Sequences
+    # Try without prefix
-        matches = [r for r in data 
+    repo_name = model_id.split('/')[-1] if '/' in model_id else model_id
-                  if r["model"] == model_id 
+    candidate = os.path.join(LOCAL_MODEL_DIR, repo_name)
-                  and r["tp"] == tp_size 
+    if os.path.isdir(candidate):
-                  and r["max_seqs"] == max_seqs 
+        log(f"Found model at: {candidate}")
-                  and r["status"] == "success"]
+        return candidate
     
-        if not matches:
+    # Try case-insensitive match
-            # 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:
-        return default_config
-
-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 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}")
-    
-    # Determine whether we have a local copy to serve. Try multiple fallbacks:
-    # 1) LOCAL_MODEL_DIR/<owner>/<repo>
-    # 2) LOCAL_MODEL_DIR/<repo>
-    # 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)
-        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
+                    log(f"Found model at: {entry_path}")
-                                print(f"DEBUG: Found model at: {model_path}")
+                    return entry_path
-                                break
     except Exception as e:
-                    print(f"DEBUG: Exception during case-insensitive lookup: {e}")
+        log(f"ERROR: Failed to list directory: {e}")
     
-    print(f"DEBUG: Final model_path: {model_path}")
+    log(f"ERROR: Model not found: {model_id}")
-    print(f"DEBUG: model_path == model_id: {model_path == model_id}")
+    log(f"Available models: {os.listdir(LOCAL_MODEL_DIR)}")
+    return None
 
-    # if LOCAL_MODEL_DIR is specified, refuse to fall back to remote
+def launch_model(model_id, config, model_path, gpu_count):
-    if LOCAL_MODEL_DIR and model_path == model_id:
+    """Launch vLLM server"""
-        print(f"Error: model '{model_id}' not found under LOCAL_MODEL_DIR={LOCAL_MODEL_DIR}")
+    log(f"Launching model: {model_id}")
-        print("Off‑line mode active; network downloads are disabled.")
-        sys.exit(1)
     
-    # Static Config
+    # Get configuration
-    valid_tps = config.get("valid_tp", [1])
+    valid_tp = config.get("valid_tp", [1])
-    max_tp = max(valid_tps) if valid_tps else 1
+    max_tp = max(valid_tp) if valid_tp else 1
+    tp_size = min(gpu_count, max_tp)
     
-    # Defaults
+    ctx = int(config.get("ctx", 8192))
-    current_tp = min(gpu_count, max_tp)
+    max_seqs = int(config.get("max_num_seqs", 64))
-    current_seqs = 1 # Default to 1 concurrent user/request for stability
+    gpu_util = float(config.get("gpu_util", 0.98))
     
-    # Initial Lookup
+    log(f"Config: TP={tp_size}, Ctx={ctx}, Seqs={max_seqs}, Util={gpu_util}")
-    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()
     
+    # Build command
     cmd = [
         "vllm", "serve", model_path,
         "--host", HOST,
         "--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"
     ]
     
-    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.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:
-        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()
+    """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"
     
     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:
-        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__":
     main()