import argparseimport atexitimport jsonimport osimport shutilimport subprocessimport sysimport timeimport urllib.errorimport urllib.request # ConfigurationDEFAULTS = { "model": "YCWTG/gemma-4-31B-it-NVFP4A16-GPTQ", "served_model_name": "YCWTG/gemma-4-31B-it-NVFP4A16-GPTQ", "host": "localhost", "port": 8000, "max_model_len": 63590, "enable_auto_tool_choice": True, "tool_call_parser": "gemma4", "max_num_seqs": 1, "reasoning_parser": "gemma4", "default_chat_template_kwargs": '{"enable_thinking": true}', "allowed_local_media_path": "/home/ycwtg/image",} RUNTIME = { "gpu_memory_utilization": 0.98, "startup_timeout_sec": 1800, "healthcheck_timeout_sec": 3, "healthcheck_interval_sec": 1, "chat_timeout_sec": 600,} # The API is always local; global HTTP_PROXY settings must not intercept it.LOCAL_HTTP = urllib.request.build_opener(urllib.request.ProxyHandler({})) SERVE_VALUE_ARGS = ( "served_model_name", "host", "port", "max_model_len", "tool_call_parser", "max_num_seqs", "reasoning_parser", "default_chat_template_kwargs",)CLIENT_VALUE_ARGS = ("model", *SERVE_VALUE_ARGS)BOOL_ARGS = ("enable_auto_tool_choice",) def cli_flag(name): return "--" + ("max_num_seqs" if name == "max_num_seqs" else name.replace("_", "-")) def value_options(args, names): return [part for name in names for part in (cli_flag(name), str(getattr(args, name)))] def boolean_options(args, explicit_false=False): return [ cli_flag(name) if getattr(args, name) else "--no-" + cli_flag(name)[2:] for name in BOOL_ARGS if explicit_false or getattr(args, name) ] def multiline_input(): print('User (type "END" on a single line to send, "exit" to quit):') lines = [] while True: line = input() text = line.strip() if text.lower() in {"exit", "quit"}: return None if text == "END": break lines.append(line) return "\n".join(lines) def resolve_client_host(host): return "127.0.0.1" if host in {"0.0.0.0", "::"} else host def launch_vllm(args): vllm = shutil.which("vllm", path=os.path.dirname(sys.executable)) or shutil.which("vllm") if not vllm: raise RuntimeError("vllm command not found. Activate an environment that has vllm installed.") cmd = [vllm, "serve", args.model, *value_options(args, SERVE_VALUE_ARGS)] media_path = args.allowed_local_media_path if media_path is not None and (not isinstance(media_path, str) or media_path.strip()): cmd += ["--allowed-local-media-path", str(media_path)] cmd += ["--gpu-memory-utilization", str(RUNTIME["gpu_memory_utilization"]), *boolean_options(args)] print("Launching vLLM:") print(" ".join(cmd)) env = os.environ.copy() env["PATH"] = os.path.dirname(vllm) + os.pathsep + env.get("PATH", "") try: return subprocess.Popen(cmd, env=env) except FileNotFoundError as e: raise RuntimeError("vllm command not found. Activate an environment that has vllm installed.") from e def stop_vllm(proc): if proc and proc.poll() is None: proc.terminate() try: proc.wait(timeout=10) except subprocess.TimeoutExpired: proc.kill() def wait_vllm_ready(base_url, timeout_sec=RUNTIME["startup_timeout_sec"], proc=None): deadline = time.time() + timeout_sec req = urllib.request.Request(url=f"{base_url}/v1/models") while time.time() < deadline: if proc and proc.poll() is not None: return False try: with LOCAL_HTTP.open(req, timeout=RUNTIME["healthcheck_timeout_sec"]) as resp: if resp.status == 200: return True except urllib.error.URLError: pass time.sleep(RUNTIME["healthcheck_interval_sec"]) return False def chat_once(base_url, model_name, messages): payload = {"model": model_name, "messages": messages, "skip_special_tokens": False} req = urllib.request.Request( url=f"{base_url}/v1/chat/completions", data=json.dumps(payload, ensure_ascii=False).encode("utf-8"), headers={"Content-Type": "application/json"}, method="POST", ) with LOCAL_HTTP.open(req, timeout=RUNTIME["chat_timeout_sec"]) as resp: data = json.loads(resp.read().decode("utf-8")) return data["choices"][0]["message"] def chat_loop(base_url, model_name): print("\n===== Chat Started =====\n") messages = [] while True: user_text = multiline_input() if user_text is None: break messages.append({"role": "user", "content": user_text}) try: assistant_msg = chat_once(base_url, model_name, messages) except Exception as e: print(f"\nRequest failed: {e}\n") messages.pop() continue content = assistant_msg.get("content") tool_calls = assistant_msg.get("tool_calls") if content: print(f"\nAssistant:\n{content}\n") elif tool_calls: print("\nAssistant(tool_calls):") print(json.dumps(tool_calls, ensure_ascii=False, indent=2)) print() else: print("\nAssistant:\n(empty response)\n") normalized_msg = {"role": "assistant", "content": content or ""} if tool_calls: normalized_msg["tool_calls"] = tool_calls messages.append(normalized_msg) def build_client_command(args): return [ sys.executable, os.path.abspath(__file__), "--_client", *value_options(args, CLIENT_VALUE_ARGS), *boolean_options(args, explicit_false=True), ] def spawn_chat_terminal(args): client_cmd = build_client_command(args) if os.name == "nt": terminal_cmd = ["cmd", "/c", "start", "", "cmd", "/k", subprocess.list2cmdline(client_cmd)] elif shutil.which("ptyxis"): terminal_cmd = ["ptyxis", "--standalone", "--new-window", "--title=vLLM Chat", "--", *client_cmd] elif shutil.which("gnome-terminal"): terminal_cmd = ["gnome-terminal", "--", *client_cmd] elif shutil.which("x-terminal-emulator"): terminal_cmd = ["x-terminal-emulator", "-e", *client_cmd] else: return False try: terminal_proc = subprocess.Popen(terminal_cmd) if terminal_cmd[0] == "ptyxis": time.sleep(0.5) if terminal_proc.poll() is not None: print(f"Failed to open Ptyxis (exit code {terminal_proc.returncode}).") return False return True except Exception as e: print(f"Failed to open a new terminal automatically: {e}") return False def parse_args(): parser = argparse.ArgumentParser(description="Minimal local vLLM chat script") parser.add_argument("--_client", action="store_true", help=argparse.SUPPRESS) def add(name, *flags, **kwargs): parser.add_argument( *(flags or (f"--{name.replace('_', '-')}",)), dest=name, default=DEFAULTS[name], **kwargs ) add("model") add("served_model_name") add("host") add("port", type=int) add("max_model_len", type=int) add("max_num_seqs", "--max-num-seqs", "--max_num_seqs", type=int) add("enable_auto_tool_choice", action=argparse.BooleanOptionalAction) add("allowed_local_media_path", help="Optional local media path. Leave empty to disable.") add("tool_call_parser") add("reasoning_parser") add("default_chat_template_kwargs") return parser.parse_args() def main(): args = parse_args() base_url = f"http://{resolve_client_host(args.host)}:{args.port}" if args._client: print(f"Waiting for model service: {base_url}") if wait_vllm_ready(base_url): chat_loop(base_url, args.served_model_name) else: print("Model service did not become ready.") return proc = launch_vllm(args) atexit.register(stop_vllm, proc) terminal_opened = spawn_chat_terminal(args) print(f"Waiting for service to become ready: {base_url}") if not wait_vllm_ready(base_url, proc=proc): print(f"vLLM failed to become ready (exit code: {proc.poll()}). Check server logs above.") stop_vllm(proc) sys.exit(1) if terminal_opened: print("Model is ready. Opened a new terminal for chat; this terminal keeps server logs.") print("Press Ctrl+C here to stop vLLM.") try: proc.wait() except KeyboardInterrupt: print("\nInterrupted. Stopping vLLM...") else: print("No supported terminal found. Falling back to chat in this terminal.") chat_loop(base_url, args.served_model_name) if __name__ == "__main__": main()