|
|
| import os |
| import time |
| import json |
| import shutil |
| import tarfile |
| import zipfile |
| import stat |
| import subprocess |
| import threading |
| from pathlib import Path |
|
|
| import requests |
| import gradio as gr |
| from huggingface_hub import hf_hub_download |
|
|
|
|
| APP_DIR = Path("/app") |
|
|
| MODEL_DIR = APP_DIR / "models" |
| MODEL_DIR.mkdir(parents=True, exist_ok=True) |
|
|
| MODEL_REPO = "Qwen/Qwen2.5-Coder-3B-Instruct-GGUF" |
| MODEL_FILE = "qwen2.5-coder-3b-instruct-q4_k_m.gguf" |
| MODEL_PATH = MODEL_DIR / MODEL_FILE |
|
|
| LLAMA_DIR = APP_DIR / "llama-server" |
| LLAMA_EXTRACT_DIR = APP_DIR / "llama-extracted" |
| LLAMA_SERVER = LLAMA_DIR / "llama-server" |
|
|
| SERVER_HOST = "127.0.0.1" |
| SERVER_PORT = 8080 |
|
|
| GRADIO_HOST = "0.0.0.0" |
| GRADIO_PORT = 7860 |
|
|
| API_KEY = "local-space-key" |
|
|
| PARALLEL_USERS = 4 |
| CONTEXT_PER_USER = 2048 |
| TOTAL_CONTEXT = PARALLEL_USERS * CONTEXT_PER_USER |
|
|
| CPU_THREADS = int( |
| os.environ.get("CPU_THREADS", "4") |
| ) |
|
|
| CPU_THREADS_BATCH = int( |
| os.environ.get( |
| "CPU_THREADS_BATCH", |
| str(CPU_THREADS) |
| ) |
| ) |
|
|
| BATCH_SIZE = int( |
| os.environ.get("BATCH_SIZE", "512") |
| ) |
|
|
| MAX_TOKENS = int( |
| os.environ.get("MAX_TOKENS", "512") |
| ) |
|
|
| TEMPERATURE = float( |
| os.environ.get("TEMPERATURE", "0.2") |
| ) |
|
|
| TOP_P = float( |
| os.environ.get("TOP_P", "0.9") |
| ) |
|
|
| server_process = None |
| server_lock = threading.Lock() |
|
|
|
|
| def download_model(): |
|
|
| if MODEL_PATH.exists(): |
|
|
| size_gb = ( |
| MODEL_PATH.stat().st_size |
| / 1024 |
| / 1024 |
| / 1024 |
| ) |
|
|
| print( |
| f"[MODEL] Already exists: {MODEL_PATH}" |
| ) |
|
|
| print( |
| f"[MODEL] Size: {size_gb:.2f} GB" |
| ) |
|
|
| return str(MODEL_PATH) |
|
|
| print( |
| "[MODEL] Downloading Qwen2.5-Coder-3B Q4_K_M..." |
| ) |
|
|
| hf_token = os.environ.get("HF_TOKEN") |
|
|
| kwargs = { |
| "repo_id": MODEL_REPO, |
| "filename": MODEL_FILE, |
| "local_dir": str(MODEL_DIR), |
| } |
|
|
| if hf_token: |
| kwargs["token"] = hf_token |
|
|
| downloaded = hf_hub_download( |
| **kwargs |
| ) |
|
|
| print( |
| f"[MODEL] Download complete: {downloaded}" |
| ) |
|
|
| return downloaded |
|
|
|
|
| def find_llama_server(directory): |
|
|
| directory = Path(directory) |
|
|
| for path in directory.rglob( |
| "llama-server" |
| ): |
|
|
| if path.is_file(): |
|
|
| try: |
| path.chmod( |
| path.stat().st_mode |
| | stat.S_IXUSR |
| | stat.S_IXGRP |
| | stat.S_IXOTH |
| ) |
| except Exception: |
| pass |
|
|
| return path |
|
|
| return None |
|
|
|
|
| def find_cpu_asset(): |
|
|
| api_url = ( |
| "https://api.github.com/" |
| "repos/ggml-org/llama.cpp/releases" |
| ) |
|
|
| headers = { |
| "Accept": |
| "application/vnd.github+json", |
| "User-Agent": |
| "hf-space-qwen-llama-server" |
| } |
|
|
| print( |
| "[LLAMA] Searching llama.cpp releases..." |
| ) |
|
|
| response = requests.get( |
| api_url, |
| params={ |
| "per_page": 20 |
| }, |
| headers=headers, |
| timeout=30 |
| ) |
|
|
| response.raise_for_status() |
|
|
| releases = response.json() |
|
|
| if not isinstance( |
| releases, |
| list |
| ): |
| raise RuntimeError( |
| "Invalid GitHub releases response." |
| ) |
|
|
| stable = [] |
|
|
| for release in releases: |
|
|
| if not release.get( |
| "draft", |
| False |
| ) and not release.get( |
| "prerelease", |
| False |
| ): |
| stable.append(release) |
|
|
| releases = stable + [ |
| r for r in releases |
| if r not in stable |
| ] |
|
|
| for release in releases: |
|
|
| tag = release.get( |
| "tag_name", |
| "unknown" |
| ) |
|
|
| assets = release.get( |
| "assets", |
| [] |
| ) |
|
|
| print( |
| f"[LLAMA] Checking {tag}" |
| ) |
|
|
| for asset in assets: |
|
|
| name = asset.get( |
| "name", |
| "" |
| ) |
|
|
| lower = name.lower() |
|
|
| special = any( |
| x in lower |
| for x in ( |
| "cuda", |
| "vulkan", |
| "rocm", |
| "openvino", |
| "sycl", |
| "metal" |
| ) |
| ) |
|
|
| current_cpu = ( |
| "ubuntu-x64" in lower |
| and lower.endswith( |
| ".tar.gz" |
| ) |
| ) |
|
|
| old_cpu = ( |
| "ubuntu" in lower |
| and "x64" in lower |
| and ( |
| "cpu" in lower |
| or "bin-" in lower |
| ) |
| and lower.endswith( |
| ( |
| ".tar.gz", |
| ".zip" |
| ) |
| ) |
| ) |
|
|
| if ( |
| not special |
| and ( |
| current_cpu |
| or old_cpu |
| ) |
| ): |
|
|
| print( |
| f"[LLAMA] Found CPU build: {name}" |
| ) |
|
|
| return asset |
|
|
| print( |
| "[LLAMA] No compatible CPU build found." |
| ) |
|
|
| for release in releases[:5]: |
|
|
| print( |
| "Release:", |
| release.get( |
| "tag_name", |
| "unknown" |
| ) |
| ) |
|
|
| for asset in release.get( |
| "assets", |
| [] |
| ): |
|
|
| print( |
| " ", |
| asset.get( |
| "name", |
| "" |
| ) |
| ) |
|
|
| raise RuntimeError( |
| "Could not find Ubuntu x64 CPU " |
| "llama.cpp release asset." |
| ) |
|
|
|
|
| def download_llama_server(): |
|
|
| if LLAMA_SERVER.exists(): |
|
|
| try: |
| LLAMA_SERVER.chmod( |
| LLAMA_SERVER.stat().st_mode |
| | stat.S_IXUSR |
| | stat.S_IXGRP |
| | stat.S_IXOTH |
| ) |
| except Exception: |
| pass |
|
|
| print( |
| f"[LLAMA] Existing server found: " |
| f"{LLAMA_SERVER}" |
| ) |
|
|
| return str( |
| LLAMA_SERVER |
| ) |
|
|
| print( |
| "[LLAMA] Downloading prebuilt llama.cpp server..." |
| ) |
|
|
| asset = find_cpu_asset() |
|
|
| asset_name = asset[ |
| "name" |
| ] |
|
|
| download_url = asset[ |
| "browser_download_url" |
| ] |
|
|
| archive_path = ( |
| APP_DIR / |
| asset_name |
| ) |
|
|
| print( |
| f"[LLAMA] Selected: {asset_name}" |
| ) |
|
|
| if not archive_path.exists(): |
|
|
| print( |
| "[LLAMA] Downloading archive..." |
| ) |
|
|
| headers = { |
| "User-Agent": |
| "hf-space-qwen-llama-server" |
| } |
|
|
| with requests.get( |
| download_url, |
| headers=headers, |
| stream=True, |
| timeout=120 |
| ) as response: |
|
|
| response.raise_for_status() |
|
|
| total = int( |
| response.headers.get( |
| "content-length", |
| 0 |
| ) |
| ) |
|
|
| downloaded = 0 |
| last_print = 0 |
|
|
| with open( |
| archive_path, |
| "wb" |
| ) as f: |
|
|
| for chunk in response.iter_content( |
| chunk_size=1024 * 1024 |
| ): |
|
|
| if not chunk: |
| continue |
|
|
| f.write(chunk) |
|
|
| downloaded += len( |
| chunk |
| ) |
|
|
| if total: |
|
|
| now = time.time() |
|
|
| if now - last_print >= 0.5: |
|
|
| percent = ( |
| downloaded |
| * 100 |
| / total |
| ) |
|
|
| print( |
| f"\r[LLAMA] " |
| f"{percent:5.1f}% " |
| f"{downloaded / 1024 / 1024:.1f} MB / " |
| f"{total / 1024 / 1024:.1f} MB", |
| end="", |
| flush=True |
| ) |
|
|
| last_print = now |
|
|
| print() |
|
|
| print( |
| "[LLAMA] Archive download complete." |
| ) |
|
|
| else: |
|
|
| print( |
| "[LLAMA] Archive already exists." |
| ) |
|
|
| if LLAMA_EXTRACT_DIR.exists(): |
|
|
| shutil.rmtree( |
| LLAMA_EXTRACT_DIR |
| ) |
|
|
| LLAMA_EXTRACT_DIR.mkdir( |
| parents=True, |
| exist_ok=True |
| ) |
|
|
| print( |
| "[LLAMA] Extracting..." |
| ) |
|
|
| if asset_name.lower().endswith( |
| ".tar.gz" |
| ): |
|
|
| with tarfile.open( |
| archive_path, |
| "r:gz" |
| ) as tar: |
|
|
| tar.extractall( |
| LLAMA_EXTRACT_DIR |
| ) |
|
|
| elif asset_name.lower().endswith( |
| ".zip" |
| ): |
|
|
| with zipfile.ZipFile( |
| archive_path, |
| "r" |
| ) as z: |
|
|
| z.extractall( |
| LLAMA_EXTRACT_DIR |
| ) |
|
|
| else: |
|
|
| raise RuntimeError( |
| f"Unsupported archive: {asset_name}" |
| ) |
|
|
| found = find_llama_server( |
| LLAMA_EXTRACT_DIR |
| ) |
|
|
| if found is None: |
|
|
| for path in LLAMA_EXTRACT_DIR.rglob( |
| "*" |
| ): |
|
|
| if path.is_file(): |
| print( |
| path |
| ) |
|
|
| raise RuntimeError( |
| "llama-server was not found " |
| "inside the archive." |
| ) |
|
|
| print( |
| f"[LLAMA] Found: {found}" |
| ) |
|
|
| source_dir = found.parent |
|
|
| LLAMA_DIR.mkdir( |
| parents=True, |
| exist_ok=True |
| ) |
|
|
| for item in source_dir.iterdir(): |
|
|
| destination = ( |
| LLAMA_DIR / |
| item.name |
| ) |
|
|
| if item.is_dir(): |
|
|
| if destination.exists(): |
| shutil.rmtree( |
| destination |
| ) |
|
|
| shutil.copytree( |
| item, |
| destination |
| ) |
|
|
| else: |
|
|
| shutil.copy2( |
| item, |
| destination |
| ) |
|
|
| final_server = ( |
| LLAMA_DIR / |
| found.name |
| ) |
|
|
| if not final_server.exists(): |
|
|
| raise RuntimeError( |
| "Failed to install llama-server." |
| ) |
|
|
| final_server.chmod( |
| final_server.stat().st_mode |
| | stat.S_IXUSR |
| | stat.S_IXGRP |
| | stat.S_IXOTH |
| ) |
|
|
| library_dirs = [] |
|
|
| for root, dirs, files in os.walk( |
| LLAMA_DIR |
| ): |
|
|
| for filename in files: |
|
|
| if ( |
| filename.endswith(".so") |
| or ".so." in filename |
| ): |
|
|
| library_dirs.append( |
| root |
| ) |
|
|
| break |
|
|
| library_dirs = list( |
| dict.fromkeys( |
| library_dirs |
| ) |
| ) |
|
|
| if library_dirs: |
|
|
| old_ld = os.environ.get( |
| "LD_LIBRARY_PATH", |
| "" |
| ) |
|
|
| new_ld = ":".join( |
| library_dirs |
| ) |
|
|
| if old_ld: |
|
|
| os.environ[ |
| "LD_LIBRARY_PATH" |
| ] = ( |
| new_ld |
| + ":" |
| + old_ld |
| ) |
|
|
| else: |
|
|
| os.environ[ |
| "LD_LIBRARY_PATH" |
| ] = new_ld |
|
|
| print( |
| "[LLAMA] Libraries configured." |
| ) |
|
|
| print( |
| f"[LLAMA] Server installed: " |
| f"{final_server}" |
| ) |
|
|
| return str( |
| final_server |
| ) |
|
|
|
|
| def start_llama_server(): |
|
|
| global server_process |
|
|
| with server_lock: |
|
|
| if server_process is not None: |
|
|
| if ( |
| server_process.poll() |
| is None |
| ): |
|
|
| return |
|
|
| server_path = ( |
| download_llama_server() |
| ) |
|
|
| command = [ |
| server_path, |
| "-m", |
| str(MODEL_PATH), |
| "-np", |
| str(PARALLEL_USERS), |
| "-c", |
| str(TOTAL_CONTEXT), |
| "--cont-batching", |
| "-ngl", |
| "0", |
| "-t", |
| str(CPU_THREADS), |
| "-tb", |
| str(CPU_THREADS_BATCH), |
| "-b", |
| str(BATCH_SIZE), |
| "--host", |
| SERVER_HOST, |
| "--port", |
| str(SERVER_PORT), |
| "--api-key", |
| API_KEY |
| ] |
|
|
| print( |
| "[LLAMA] Starting server..." |
| ) |
|
|
| server_process = subprocess.Popen( |
| command, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT, |
| text=True, |
| bufsize=1, |
| env=os.environ.copy() |
| ) |
|
|
| def read_logs(): |
|
|
| try: |
|
|
| for line in iter( |
| server_process.stdout.readline, |
| "" |
| ): |
|
|
| if line: |
| print( |
| "[llama.cpp]", |
| line.rstrip() |
| ) |
|
|
| except Exception: |
| pass |
|
|
| threading.Thread( |
| target=read_logs, |
| daemon=True |
| ).start() |
|
|
| health_url = ( |
| f"http://{SERVER_HOST}:" |
| f"{SERVER_PORT}/health" |
| ) |
|
|
| print( |
| "[LLAMA] Waiting for server..." |
| ) |
|
|
| started = time.time() |
|
|
| while True: |
|
|
| if ( |
| server_process is not None |
| and server_process.poll() |
| is not None |
| ): |
|
|
| raise RuntimeError( |
| "llama-server stopped " |
| f"with code " |
| f"{server_process.returncode}" |
| ) |
|
|
| try: |
|
|
| response = requests.get( |
| health_url, |
| timeout=2 |
| ) |
|
|
| if response.status_code in ( |
| 200, |
| 204 |
| ): |
|
|
| print( |
| "[LLAMA] Server is ready." |
| ) |
|
|
| return |
|
|
| except Exception: |
| pass |
|
|
| if ( |
| time.time() - started |
| > 180 |
| ): |
|
|
| raise RuntimeError( |
| "Timed out waiting for llama-server." |
| ) |
|
|
| time.sleep(1) |
|
|
|
|
| def clean_content(content): |
|
|
| if content is None: |
| return "" |
|
|
| if isinstance( |
| content, |
| str |
| ): |
| return content |
|
|
| if isinstance( |
| content, |
| list |
| ): |
|
|
| result = [] |
|
|
| for item in content: |
|
|
| if isinstance( |
| item, |
| str |
| ): |
|
|
| result.append( |
| item |
| ) |
|
|
| elif isinstance( |
| item, |
| dict |
| ): |
|
|
| if "text" in item: |
|
|
| result.append( |
| str( |
| item["text"] |
| ) |
| ) |
|
|
| elif "content" in item: |
|
|
| result.append( |
| str( |
| item["content"] |
| ) |
| ) |
|
|
| else: |
|
|
| result.append( |
| str(item) |
| ) |
|
|
| return "\n".join( |
| x for x in result |
| if x |
| ) |
|
|
| if isinstance( |
| content, |
| dict |
| ): |
|
|
| if "text" in content: |
|
|
| return str( |
| content["text"] |
| ) |
|
|
| if "content" in content: |
|
|
| return str( |
| content["content"] |
| ) |
|
|
| return json.dumps( |
| content, |
| ensure_ascii=False |
| ) |
|
|
| return str(content) |
|
|
|
|
| def build_messages( |
| message, |
| history |
| ): |
|
|
| messages = [] |
|
|
| history = history or [] |
|
|
| recent_history = history[-10:] |
|
|
| for item in recent_history: |
|
|
| if not isinstance( |
| item, |
| dict |
| ): |
| continue |
|
|
| role = item.get( |
| "role" |
| ) |
|
|
| content = clean_content( |
| item.get( |
| "content", |
| "" |
| ) |
| ) |
|
|
| if role in ( |
| "user", |
| "assistant", |
| "system" |
| ) and content: |
|
|
| messages.append({ |
| "role": |
| role, |
| "content": |
| content |
| }) |
|
|
| messages.append({ |
| "role": |
| "user", |
| "content": |
| clean_content( |
| message |
| ) |
| }) |
|
|
| return messages |
|
|
|
|
| def generate_response( |
| message, |
| history |
| ): |
|
|
| if not message: |
| return "Please enter a message." |
|
|
| messages = build_messages( |
| message, |
| history |
| ) |
|
|
| url = ( |
| f"http://{SERVER_HOST}:" |
| f"{SERVER_PORT}/v1/chat/completions" |
| ) |
|
|
| payload = { |
| "model": |
| "qwen2.5-coder-3b", |
| "messages": |
| messages, |
| "temperature": |
| TEMPERATURE, |
| "top_p": |
| TOP_P, |
| "max_tokens": |
| MAX_TOKENS, |
| "stream": |
| False |
| } |
|
|
| headers = { |
| "Content-Type": |
| "application/json", |
| "Authorization": |
| f"Bearer {API_KEY}" |
| } |
|
|
| try: |
|
|
| response = requests.post( |
| url, |
| headers=headers, |
| json=payload, |
| timeout=300 |
| ) |
|
|
| except requests.RequestException as e: |
|
|
| return ( |
| "Server connection error:\n\n" |
| f"{e}" |
| ) |
|
|
| if response.status_code != 200: |
|
|
| try: |
|
|
| data = response.json() |
|
|
| return ( |
| "llama.cpp error:\n\n" |
| + json.dumps( |
| data, |
| indent=2, |
| ensure_ascii=False |
| ) |
| ) |
|
|
| except Exception: |
|
|
| return ( |
| f"llama.cpp HTTP " |
| f"{response.status_code}\n\n" |
| f"{response.text}" |
| ) |
|
|
| try: |
|
|
| data = response.json() |
|
|
| choices = data.get( |
| "choices", |
| [] |
| ) |
|
|
| if not choices: |
|
|
| return "No response generated." |
|
|
| choice = choices[0] |
|
|
| message_data = choice.get( |
| "message", |
| {} |
| ) |
|
|
| result = message_data.get( |
| "content", |
| "" |
| ) |
|
|
| if not result: |
|
|
| result = choice.get( |
| "text", |
| "" |
| ) |
|
|
| return result.strip() |
|
|
| except Exception as e: |
|
|
| return ( |
| "Response parsing error:\n\n" |
| f"{e}\n\n" |
| f"{response.text}" |
| ) |
|
|
|
|
| def create_ui(): |
|
|
| demo = gr.ChatInterface( |
| fn=generate_response, |
| concurrency_limit=PARALLEL_USERS |
| ) |
|
|
| return demo |
|
|
|
|
| def initialize(): |
|
|
| print( |
| "============================================================" |
| ) |
|
|
| print( |
| "Qwen2.5-Coder-3B llama.cpp server" |
| ) |
|
|
| print( |
| "============================================================" |
| ) |
|
|
| print( |
| f"[CONFIG] Users: {PARALLEL_USERS}" |
| ) |
|
|
| print( |
| f"[CONFIG] Context/user: {CONTEXT_PER_USER}" |
| ) |
|
|
| print( |
| f"[CONFIG] CPU threads: {CPU_THREADS}" |
| ) |
|
|
| print( |
| f"[CONFIG] Batch threads: {CPU_THREADS_BATCH}" |
| ) |
|
|
| print( |
| f"[CONFIG] Batch size: {BATCH_SIZE}" |
| ) |
|
|
| print( |
| f"[CONFIG] Max tokens: {MAX_TOKENS}" |
| ) |
|
|
| print( |
| "============================================================" |
| ) |
|
|
| download_model() |
|
|
| download_llama_server() |
|
|
| start_llama_server() |
|
|
| print( |
| "============================================================" |
| ) |
|
|
| print( |
| "INITIALIZATION COMPLETE" |
| ) |
|
|
| print( |
| f"Model: {MODEL_FILE}" |
| ) |
|
|
| print( |
| f"Users: {PARALLEL_USERS}" |
| ) |
|
|
| print( |
| f"Context/user: {CONTEXT_PER_USER}" |
| ) |
|
|
| print( |
| f"CPU threads: {CPU_THREADS}" |
| ) |
|
|
| print( |
| "============================================================" |
| ) |
|
|
|
|
| def shutdown(): |
|
|
| global server_process |
|
|
| print( |
| "[SYSTEM] Shutting down..." |
| ) |
|
|
| if server_process is None: |
| return |
|
|
| try: |
|
|
| if ( |
| server_process.poll() |
| is None |
| ): |
|
|
| server_process.terminate() |
|
|
| try: |
|
|
| server_process.wait( |
| timeout=5 |
| ) |
|
|
| except subprocess.TimeoutExpired: |
|
|
| server_process.kill() |
|
|
| except Exception as e: |
|
|
| print( |
| f"[SYSTEM] Shutdown error: {e}" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
|
|
| try: |
|
|
| initialize() |
|
|
| demo = create_ui() |
|
|
| demo.queue( |
| default_concurrency_limit= |
| PARALLEL_USERS |
| ) |
|
|
| demo.launch( |
| server_name= |
| GRADIO_HOST, |
| server_port= |
| GRADIO_PORT, |
| show_error=True, |
| max_threads= |
| PARALLEL_USERS + 4, |
| share=False |
| ) |
|
|
| except KeyboardInterrupt: |
|
|
| print( |
| "[SYSTEM] Interrupted." |
| ) |
|
|
| except Exception as e: |
|
|
| print( |
| "============================================================" |
| ) |
|
|
| print( |
| "STARTUP FAILED" |
| ) |
|
|
| print( |
| repr(e) |
| ) |
|
|
| print( |
| "============================================================" |
| ) |
|
|
| raise |
|
|
| finally: |
|
|
| shutdown() |
