import os import fnmatch import urllib.request import json import time import threading import asyncio import multiprocessing as mp from concurrent.futures import ProcessPoolExecutor from huggingface_hub import hf_hub_download, HfApi import uvicorn from fastapi import FastAPI, Request, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel # ── 1. Curated Model Catalogue ───────────────────────────────── # Each entry is also a role's fallback: the model this Space's role loads # on boot and falls back to if a Hub refresh (see discover_best_model_for_role) # doesn't find anything usable. One role per deployed Space — see SPACE_ROLE. CURATED_MODELS = { "deepseek": { "id": "deepseek", "role": "reasoning", "label": "DeepSeek-R1-Distill Qwen 7B", "emoji": "🧠", "category": "Reasoning / Coding", "repo": "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF", "file": "*Q4_K_M.gguf", "ram_gb": 4.8, "n_ctx": 8192, "description": "Distilled DeepSeek reasoning model — punches far above its size. Best for logic, math, coding.", }, "wizardvicuna": { "id": "wizardvicuna", "role": "uncensored-alt", "label": "Wizard-Vicuna 7B Uncensored", "emoji": "🧙", "category": "Uncensored", "repo": "TheBloke/Wizard-Vicuna-7B-Uncensored-GGUF", "file": "*Q4_K_M.gguf", "ram_gb": 4.8, "n_ctx": 4096, "description": "Classic uncensored Llama-2 7B — reliable, well-tested, same size class as the other curated models so it stays within the free CPU tier's time budget.", }, "dolphin": { "id": "dolphin", "role": "uncensored", "label": "Dolphin 2.9.3 Mistral 7B", "emoji": "🐬", "category": "Uncensored", # Was Dolphin 2.9.4 Llama 3.1 8B — measurably too slow on the free # CPU tier (regularly blew past the server's own 90s self-heal # timeout even for a 10-token reply, unlike the two 7B models here # which consistently finish in under 35s). This 7B Mistral-based # Dolphin variant matches the size class of the working models. "repo": "bartowski/dolphin-2.9.3-mistral-7B-32k-GGUF", "file": "*Q4_K_M.gguf", "ram_gb": 4.8, "n_ctx": 8192, "description": "Eric Hartford's Dolphin series — gold standard uncensored model. Follows all instructions without refusal.", }, "qwen": { "id": "qwen", "role": "general", "label": "Qwen 2.5 7B Instruct", "emoji": "⚡", "category": "Best General", "repo": "bartowski/Qwen2.5-7B-Instruct-GGUF", "file": "*Q4_K_M.gguf", "ram_gb": 4.8, "n_ctx": 8192, "description": "Alibaba's flagship 7B model — extremely capable, multilingual, fast. Best all-rounder for daily use.", }, # NOT the default (role deliberately not "reasoning") — live-tested after # deploy and it timed out at 90s on a trivial 2-token-answer prompt. # "Thinking"-tuned models ruminate at length before answering regardless # of base model size, which is exactly the failure mode that got Dolphin # downsized before. Smaller parameter count didn't save it. Left in the # catalogue as a manual, opt-in choice only — "deepseek" stays the # default reasoning model since it's the one actually proven to finish # in time on this infra. "qwen3thinking": { "id": "qwen3thinking", "role": "reasoning-alt", "label": "Qwen3 4B Thinking (2507)", "emoji": "🤔", "category": "Reasoning / Coding", "repo": "unsloth/Qwen3-4B-Thinking-2507-GGUF", "file": "*Q4_K_M.gguf", "ram_gb": 2.8, "n_ctx": 8192, "description": "Purpose-built reasoning/thinking model, newer generation than DeepSeek-R1-Distill-Qwen-7B and roughly half the size — faster on CPU with comparable reasoning quality.", }, # Same pattern as above for the uncensored role: newer official Dolphin # release (3.0) on a much smaller 3B Llama-3.2 base, replacing the 7B # Mistral-based Dolphin 2.9.3 as the uncensored Space's default. "dolphin" # and "wizardvicuna" above remain available as manual fallbacks. "dolphin3": { "id": "dolphin3", "role": "uncensored", "label": "Dolphin 3.0 Llama 3.2 3B", "emoji": "🐬", "category": "Uncensored", "repo": "bartowski/Dolphin3.0-Llama3.2-3B-GGUF", "file": "*Q4_K_M.gguf", "ram_gb": 2.2, "n_ctx": 8192, "description": "Official Dolphin 3.0 release from the cognitivecomputations/dphn team — newer generation, less than half the size of Dolphin 2.9.3, follows instructions without refusal.", }, } ROLE_TO_MODEL_ID = {spec["role"]: mid for mid, spec in CURATED_MODELS.items()} # Hub search terms used by discover_best_model_for_role() when you hit # /jolly/refresh-model. "prefer" keywords just re-rank same-popularity # candidates; they don't gate eligibility. ROLE_SEARCH = { "reasoning": {"query": "reasoning instruct gguf", "prefer": ["r1", "reasoning", "qwq", "thinking", "o1"]}, "uncensored": {"query": "uncensored abliterated gguf", "prefer": ["uncensored", "abliterated", "dolphin"]}, "general": {"query": "instruct chat gguf", "prefer": []}, } # Skip repos naming a parameter count above what the free CPU tier (16GB RAM) # can load at Q4 — keeps auto-discovery from picking something that OOMs. _SIZE_BLOCKLIST = ("70b", "72b", "65b", "34b", "32b", "30b", "27b", "24b", "22b", "20b", "14b", "13b", "405b") # ── 2. Config (env vars override defaults) ───────────────────── # SPACE_ROLE is the one thing you set per deployed Space — it picks which # curated model this instance boots into automatically, no MODEL_ID needed. SPACE_ROLE = os.environ.get("SPACE_ROLE", "general").strip().lower() if SPACE_ROLE not in ROLE_TO_MODEL_ID: print(f"[Config] Unknown SPACE_ROLE '{SPACE_ROLE}', defaulting to 'general'") SPACE_ROLE = "general" CONFIG_URL = os.environ.get("CONFIG_URL", "") MODEL_ID = os.environ.get("MODEL_ID", ROLE_TO_MODEL_ID[SPACE_ROLE]) # role's model unless overridden MODEL_REPO = os.environ.get("MODEL_REPO", "") # override repo if set MODEL_FILE = os.environ.get("MODEL_FILE", "") # override file if set N_CTX = int(os.environ.get("N_CTX", "8192")) N_THREADS = int(os.environ.get("N_THREADS", "4")) API_KEY = os.environ.get("API_KEY", "") MODEL_SOURCE = "fallback" # "fallback" (curated catalogue) or "hub-discovered" (via /jolly/refresh-model) LAST_REFRESHED = None # epoch seconds of last successful /jolly/refresh-model # Pull optional cluster config if CONFIG_URL: print(f"[Config] Fetching cluster config from: {CONFIG_URL}") try: with urllib.request.urlopen(urllib.request.Request(CONFIG_URL), timeout=10) as r: cfg = json.loads(r.read().decode()) MODEL_ID = cfg.get("MODEL_ID", cfg.get("model_id", MODEL_ID)) MODEL_REPO = cfg.get("MODEL_REPO", cfg.get("model_repo", MODEL_REPO)) MODEL_FILE = cfg.get("MODEL_FILE", cfg.get("model_file", MODEL_FILE)) if "N_CTX" in cfg or "n_ctx" in cfg: N_CTX = int(cfg.get("N_CTX", cfg.get("n_ctx", N_CTX))) if "N_THREADS" in cfg or "n_threads" in cfg: N_THREADS = int(cfg.get("N_THREADS", cfg.get("n_threads", N_THREADS))) API_KEY = cfg.get("API_KEY", cfg.get("api_key", API_KEY)) print("[Config] Cluster config synced.") except Exception as e: print(f"[Config] Fetch failed ({e}), using local env vars.") # Resolve active model spec from catalogue or raw env overrides def resolve_model_spec(model_id, repo_override="", file_override=""): if repo_override: return repo_override, file_override or "*.gguf", N_CTX spec = CURATED_MODELS.get(model_id, CURATED_MODELS[ROLE_TO_MODEL_ID[SPACE_ROLE]]) return spec["repo"], spec["file"], spec.get("n_ctx", N_CTX) ACTIVE_MODEL_ID = MODEL_ID if MODEL_ID in CURATED_MODELS else ROLE_TO_MODEL_ID[SPACE_ROLE] active_repo, active_file, active_ctx = resolve_model_spec(ACTIVE_MODEL_ID, MODEL_REPO, MODEL_FILE) # ── 3. Download helper ───────────────────────────────────────── def download_model(repo, file_pattern): try: if "*" not in file_pattern: return hf_hub_download(repo_id=repo, filename=file_pattern) except Exception as e: print(f"[Download] Direct download failed: {e}, searching…") api = HfApi() files = list(api.list_repo_files(repo_id=repo)) matches = [f for f in files if fnmatch.fnmatch(f.lower(), file_pattern.lower())] if not matches: matches = [f for f in files if f.endswith(".gguf")] if not matches: raise FileNotFoundError(f"No GGUF found in {repo}") print(f"[Download] Selected: {matches[0]}") return hf_hub_download(repo_id=repo, filename=matches[0]) # ── 3b. Hub auto-discovery (manual trigger via /jolly/refresh-model) ──── def discover_best_model_for_role(role): """Search the Hub for the most-downloaded GGUF model matching this role. Returns (repo, file, label) or None if nothing usable is found — callers must leave the currently-loaded model running on None rather than clobber it.""" spec = ROLE_SEARCH.get(role, ROLE_SEARCH["general"]) api = HfApi() try: candidates = list(api.list_models(search=spec["query"], sort="downloads", direction=-1, limit=25)) except Exception as e: print(f"[Discover] Hub search failed: {e}") return None def score(m): name = m.id.lower() return sum(1 for kw in spec["prefer"] if kw in name) candidates.sort(key=score, reverse=True) for m in candidates: name = m.id.lower() if any(b in name for b in _SIZE_BLOCKLIST): continue try: files = list(api.list_repo_files(m.id)) except Exception: continue match = next((f for f in files if fnmatch.fnmatch(f.lower(), "*q4_k_m.gguf")), None) if not match: match = next((f for f in files if f.lower().endswith(".gguf") and "q4" in f.lower()), None) if not match: continue label = m.id.split("/")[-1] print(f"[Discover] Role '{role}' -> {m.id} / {match}") return m.id, match, label print(f"[Discover] No suitable model found for role '{role}'") return None # ── 4. Inference runs in a dedicated child process ────────────── # llama.cpp's generation call does not hand control back to the asyncio # event loop while it runs — even in a background thread, it holds the GIL # long enough that the ENTIRE app (every route, including /jolly/health) # freezes for the whole duration of generation. A real OS process has its # own GIL, so the main server stays responsive no matter how slow a # generation is. The tradeoff: no token-by-token streaming animation # (results come back as one completed chunk), which is a fair price for # "the app never appears to hang." PUBLIC_PORT = 7860 _mp_ctx = mp.get_context("spawn") # Only meaningful inside the worker process. _worker_llama = None def _init_worker(model_path, ctx, n_threads): global _worker_llama from llama_cpp import Llama print(f"[Worker] Loading {model_path}") _worker_llama = Llama(model_path=model_path, n_ctx=ctx, n_threads=n_threads, verbose=False) print(f"[Worker] Ready: {model_path}") def _worker_generate(messages, max_tokens, temperature): return _worker_llama.create_chat_completion( messages=messages, max_tokens=max_tokens, temperature=temperature, stream=False ) def _worker_noop(): # Submitted lambdas aren't picklable (required for ProcessPoolExecutor # with the spawn start method) — a plain module-level function is. return True class WorkerBusy(Exception): pass class _Worker: """One persistent worker process holding the loaded model. Recreated on model switch.""" def __init__(self): self.executor = None self.ready = threading.Event() self.model_path = "" self.ctx = N_CTX self.lock = threading.Lock() self.queue_lock = asyncio.Lock() def start(self, model_path, ctx): self.ready.clear() self.ctx = ctx old = self.executor self.executor = ProcessPoolExecutor( max_workers=1, mp_context=_mp_ctx, initializer=_init_worker, initargs=(model_path, ctx, N_THREADS), ) self.model_path = model_path if old: old.shutdown(wait=False, cancel_futures=True) # ProcessPoolExecutor's initializer runs in the child as soon as the # pool spins up; submit a no-op through it so we only flip `ready` # once that child process has actually finished loading the model. def _mark_ready(): try: self.executor.submit(_worker_noop).result() self.ready.set() print(f"[Worker] Model ready: {model_path}") except Exception as e: print(f"[Worker] Failed to become ready: {e}") threading.Thread(target=_mark_ready, daemon=True).start() async def generate(self, messages, max_tokens, temperature): # Queue behind whatever's already generating instead of rejecting # outright — a single worker process only ever handles one request # at a time anyway, so a second request arriving mid-generation was # getting an instant 429 for no good reason (previously: fail-fast # via a busy flag). Waiting in line is strictly better for clients # that don't retry cleverly. Bounded so a truly stuck queue still # surfaces an error instead of hanging forever. try: await asyncio.wait_for(self.queue_lock.acquire(), timeout=100) except asyncio.TimeoutError: raise WorkerBusy("Server has a long queue right now — please wait and try again.") try: loop = asyncio.get_running_loop() with self.lock: executor = self.executor try: # A long system prompt (diary + user context can be several # thousand tokens) makes prompt processing itself slow, on top # of generation — a client giving up and retrying doesn't stop # the actual work on the worker process, which otherwise holds # the queue locked indefinitely for whoever asks next. Bound it, # and if it's blown through, the worker process is unrecoverable # for this request — kill and replace it rather than leave a # zombie generation running forever. return await asyncio.wait_for( loop.run_in_executor(executor, _worker_generate, messages, max_tokens, temperature), timeout=90, ) except asyncio.TimeoutError: print("[Worker] Generation exceeded 90s — killing and restarting the worker process") self.start(self.model_path, self.ctx) raise RuntimeError("Generation timed out (prompt may be too long for this model's context window) — the worker was restarted, try again with a shorter message.") finally: self.queue_lock.release() worker = _Worker() # ── 5. FastAPI app ────────────────────────────────────────────── proxy = FastAPI(title="JollyDay LLM Space — Multi-Model") proxy.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) def _check_api_key(request: Request): if not API_KEY: return got = request.headers.get("authorization", "") if got != f"Bearer {API_KEY}": raise HTTPException(status_code=401, detail="Invalid API key") # ── 6. /jolly/* control endpoints ───────────────────────────── @proxy.get("/jolly/health") async def jolly_health(): return JSONResponse({ "status": "ready" if worker.ready.is_set() else "loading", "role": SPACE_ROLE, "model_id": ACTIVE_MODEL_ID, "model_file": os.path.basename(worker.model_path), "model_source": MODEL_SOURCE, "n_ctx": N_CTX, "n_threads": N_THREADS, }) @proxy.get("/jolly/info") async def jolly_info(): spec = CURATED_MODELS.get(ACTIVE_MODEL_ID, {}) return JSONResponse({ "role": SPACE_ROLE, "model_id": ACTIVE_MODEL_ID, "model_repo": active_repo, "model_file": active_file, "model_path": os.path.basename(worker.model_path), "model_source": MODEL_SOURCE, "last_refreshed": LAST_REFRESHED, "label": spec.get("label", ""), "n_ctx": N_CTX, "n_threads": N_THREADS, "api_key_set": bool(API_KEY), "config_url": CONFIG_URL, "server_ready": worker.ready.is_set(), }) @proxy.get("/jolly/models") async def jolly_models(): return JSONResponse({ "role": SPACE_ROLE, "active_id": ACTIVE_MODEL_ID, "models": list(CURATED_MODELS.values()), }) @proxy.get("/jolly/active") async def jolly_active(): spec = CURATED_MODELS.get(ACTIVE_MODEL_ID, {}) return JSONResponse({ "id": ACTIVE_MODEL_ID, "role": SPACE_ROLE, "label": spec.get("label", os.path.basename(worker.model_path)), "emoji": spec.get("emoji", "🤖"), "category": spec.get("category", ""), "description": spec.get("description", ""), "model_file": os.path.basename(worker.model_path), "model_source": MODEL_SOURCE, "server_ready": worker.ready.is_set(), }) @proxy.get("/jolly/wake") async def jolly_wake(): """Lightweight wake endpoint — just touching this URL wakes a sleeping HF Space.""" spec = CURATED_MODELS.get(ACTIVE_MODEL_ID, {}) return JSONResponse({ "status": "ready" if worker.ready.is_set() else "loading", "model_id": ACTIVE_MODEL_ID, "label": spec.get("label", ACTIVE_MODEL_ID), "emoji": spec.get("emoji", "🤖"), }) class SwitchRequest(BaseModel): model_id: str @proxy.post("/jolly/switch") async def jolly_switch(req: SwitchRequest): global ACTIVE_MODEL_ID, active_repo, active_file, active_ctx, MODEL_SOURCE model_id = req.model_id.lower().strip() if model_id not in CURATED_MODELS: raise HTTPException( status_code=400, detail=f"Unknown model_id '{model_id}'. Valid options: {list(CURATED_MODELS.keys())}" ) if model_id == ACTIVE_MODEL_ID and worker.ready.is_set(): spec = CURATED_MODELS[model_id] return JSONResponse({"status": "already_active", "model_id": model_id, "label": spec["label"]}) spec = CURATED_MODELS[model_id] ACTIVE_MODEL_ID = model_id active_repo = spec["repo"] active_file = spec["file"] active_ctx = spec.get("n_ctx", N_CTX) MODEL_SOURCE = "fallback" def _switch(): try: path = download_model(active_repo, active_file) worker.start(path, active_ctx) except Exception as e: print(f"[Switch] ERROR: {e}") threading.Thread(target=_switch, daemon=True).start() return JSONResponse({ "status": "switching", "model_id": model_id, "label": spec["label"], "emoji": spec["emoji"], "description": spec["description"], "note": "Model is downloading/loading. Poll /jolly/health until status is 'ready'.", }) @proxy.post("/jolly/refresh-model") async def jolly_refresh_model(): """Manual 'update models' button: re-query the Hub for the current best model for this Space's role (SPACE_ROLE) and hot-swap to it if found. Leaves the running model untouched if the search comes up empty.""" global ACTIVE_MODEL_ID, active_repo, active_file, active_ctx, MODEL_SOURCE, LAST_REFRESHED found = discover_best_model_for_role(SPACE_ROLE) if not found: return JSONResponse({ "status": "no_change", "role": SPACE_ROLE, "detail": "No suitable Hub model found for this role — keeping the current model.", "active_model_id": ACTIVE_MODEL_ID, }) repo, file, label = found ACTIVE_MODEL_ID = f"hub:{label}" active_repo = repo active_file = file active_ctx = N_CTX MODEL_SOURCE = "hub-discovered" LAST_REFRESHED = time.time() def _switch(): try: path = download_model(active_repo, active_file) worker.start(path, active_ctx) except Exception as e: print(f"[Refresh] ERROR: {e}") threading.Thread(target=_switch, daemon=True).start() return JSONResponse({ "status": "switching", "role": SPACE_ROLE, "model_repo": repo, "model_file": file, "label": label, "note": "Discovered from Hub search. Poll /jolly/health until status is 'ready'.", }) # ── 7. OpenAI-style routes ────────────────────────────────────── # No true token streaming (see note above) — a streaming request still gets # a single SSE chunk with the full text, then [DONE]. Obsidian's client # just accumulates delta.content, so this renders identically to a real # stream, just without the progressive animation. @proxy.get("/v1/models") async def list_models(request: Request): _check_api_key(request) return JSONResponse({ "object": "list", "data": [{"id": ACTIVE_MODEL_ID, "object": "model", "owned_by": "me"}], }) @proxy.post("/v1/chat/completions") async def chat_completions(request: Request): _check_api_key(request) if not worker.ready.is_set(): return JSONResponse( {"error": "Model is loading, please retry in a moment.", "status": "loading"}, status_code=503, ) try: body = await request.json() except Exception: return JSONResponse({"error": "Invalid JSON body"}, status_code=400) messages = body.get("messages", []) # Hard-capped regardless of what the client asks for. Measured throughput # on this free CPU tier is ~5 tokens/sec, and responses no longer stream # token-by-token (see note above) — the client waits the full duration # with zero visible feedback. 180 tokens (~35s) was still landing "All # providers failed" in practice — almost certainly some client-side # network timeout (independent of our own AbortController) kicking in # around there. 80 tokens (~15s) leaves real margin. max_tokens = min(int(body.get("max_tokens", 80) or 80), 80) temperature = body.get("temperature", 0.7) stream = bool(body.get("stream", False)) try: result = await worker.generate(messages, max_tokens, temperature) except WorkerBusy as e: return JSONResponse({"error": str(e), "status": "busy"}, status_code=429) except Exception as e: return JSONResponse({"error": str(e)}, status_code=500) if stream: async def _gen(): content = result.get("choices", [{}])[0].get("message", {}).get("content", "") chunk = { "id": result.get("id", "chatcmpl"), "object": "chat.completion.chunk", "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}], } yield f"data: {json.dumps(chunk)}\n\n" yield "data: [DONE]\n\n" return StreamingResponse(_gen(), media_type="text/event-stream") return JSONResponse(result) # ── 8. Boot + launch ───────────────────────────────────────────── # Guarded so the `spawn` multiprocessing start method (re-imports this file # fresh in each worker process) doesn't recursively re-run the boot/serve # logic — only route/function definitions above this line execute in the # worker; nothing here does. if __name__ == "__main__": print(f"[Boot] Space role : {SPACE_ROLE}") print(f"[Boot] Model ID : {ACTIVE_MODEL_ID}") print(f"[Boot] Model repo : {active_repo}") print(f"[Boot] Model file : {active_file}") print(f"[Boot] Context : {active_ctx} Threads: {N_THREADS}") print(f"[Boot] API key set: {'yes' if API_KEY else 'no'}") try: boot_path = download_model(active_repo, active_file) except Exception as e: print(f"[WARN] {e} — falling back to Qwen 2.5 1.5B") boot_path = hf_hub_download( repo_id="Qwen/Qwen2.5-1.5B-Instruct-GGUF", filename="qwen2.5-1.5b-instruct-q4_k_m.gguf" ) worker.start(boot_path, active_ctx) print(f"[Proxy] Starting on :{PUBLIC_PORT}…") uvicorn.run(proxy, host="0.0.0.0", port=PUBLIC_PORT, log_level="info")