Swap to pre-model packing

Dockerfile

@@ -1,5 +1,11 @@
+# 1) Base
 FROM python:3.10-slim
 
+# 2) Build‐time HF token
+ARG HF_TOKEN
+ENV HF_TOKEN=${HF_TOKEN}
+
+# 3) Environment
 ENV DEBIAN_FRONTEND=noninteractive \
     RUSTUP_HOME=/root/.rustup \
     CARGO_HOME=/root/.cargo \
@@ -11,7 +17,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
 
 WORKDIR /app
 
-# install system deps + Rust toolchain
+# 4) System deps + Rust
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
       build-essential cmake git curl wget ninja-build libgomp1 ca-certificates \
@@ -20,22 +26,49 @@ RUN apt-get update && \
     curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && \
     rustup default stable
 
-# cache dirs
-RUN mkdir -p /app/cache /app/pretrained_models/llm && chmod -R 777 /app/cache /app/pretrained_models/llm
+# 5) Prepare directories
+RUN mkdir -p /app/cache /app/pretrained_models/llm && \
+    chmod -R 777 /app/cache /app/pretrained_models/llm
 
-# Python deps (except llama-cpp-python)
+# 6) Install Python deps (except llama-cpp-python)
 COPY requirements.txt .
 RUN sed -i '/llama-cpp-python/d' requirements.txt && \
     pip install --no-cache-dir -r requirements.txt
 
-# build llama-cpp-python from source (with its llama.cpp submodule)
+# 7) Build llama-cpp-python from source
 RUN git clone --recursive https://github.com/abetlen/llama-cpp-python.git /tmp/llama-cpp-python && \
     cd /tmp/llama-cpp-python && \
     git submodule update --init --recursive && \
     python -m pip install --no-cache-dir . && \
     rm -rf /tmp/llama-cpp-python
 
-# copy the server code
+# 8) Install huggingface_hub CLI
+RUN pip install --no-cache-dir huggingface_hub
+
+# 9) Clone llama.cpp for quantize tool
+RUN git clone https://github.com/ggerganov/llama.cpp.git /tmp/llama.cpp && \
+    cd /tmp/llama.cpp && make quantize
+
+# 10) Download & pre-quantize model (no runtime repack)
+RUN python3 - <<EOF
+import os
+from huggingface_hub import login, hf_hub_download
+if os.getenv('HF_TOKEN'): login(token=os.getenv('HF_TOKEN'))
+hf_hub_download(repo_id=os.getenv('MODEL_REPO'),
+                filename=os.getenv('MODEL_FILE'),
+                local_dir='/app/pretrained_models/llm',
+                token=os.getenv('HF_TOKEN'))
+EOF
+
+RUN /tmp/llama.cpp/build/quantize \
+      /app/pretrained_models/llm/${MODEL_FILE} \
+      /app/pretrained_models/llm/${MODEL_FILE}.packed.gguf \
+      q4_K_M && \
+    mv /app/pretrained_models/llm/${MODEL_FILE}.packed.gguf \
+       /app/pretrained_models/llm/${MODEL_FILE} && \
+    rm -rf /tmp/llama.cpp
+
+# 11) Copy server
 COPY llm_server.py /app/llm_server.py
 
 EXPOSE 7860

llm_server.py

@@ -1,114 +1,50 @@
-import os
-import time
-import logging
-import asyncio
+import os, time, logging, asyncio
 from pathlib import Path
-
 from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
-
-from huggingface_hub import login, hf_hub_download
+from huggingface_hub import login
 from llama_cpp import Llama
 
-# ─── logging setup ────────────────────────────────────────────────────────────
+# ─── Logging ────────────────────────────────────────────────────────────────
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 logger = logging.getLogger("llm_server")
 
-# ─── FastAPI setup ────────────────────────────────────────────────────────────
+# ─── FastAPI setup ─────────────────────────────────────────────────────────
 app = FastAPI()
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# ─── model paths & env ────────────────────────────────────────────────────────
-MODEL_DIR  = Path("/app/pretrained_models/llm")
-MODEL_DIR.mkdir(parents=True, exist_ok=True)
 
+# ─── Model paths & token ────────────────────────────────────────────────────
+MODEL_DIR = Path("/app/pretrained_models/llm")
+MODEL_DIR.mkdir(exist_ok=True)
 MODEL_FILE = os.getenv("MODEL_FILE", "Qwen3-0.6B-Q4_K_M.gguf")
 MODEL_PATH = MODEL_DIR / MODEL_FILE
-
 HF_TOKEN   = os.getenv("HF_TOKEN")
 
-# ─── startup: download, init, warm-up, schedule keep-alive ────────────────────
-@app.on_event("startup")
-async def startup_event():
-    logger.info("🔧 Starting LLM service…")
-
-    # login if we have a token
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    # download weights if missing
-    if not MODEL_PATH.exists():
-        logger.info(f"→ Downloading {MODEL_FILE} from {os.getenv('MODEL_REPO')}")
-        hf_hub_download(
-            repo_id=os.getenv("MODEL_REPO"),
-            filename=MODEL_FILE,
-            local_dir=str(MODEL_DIR),
-            token=HF_TOKEN,
-        )
-        logger.info("✔ Download complete")
-
-    # instantiate llama.cpp
-    global llm
-    llm = Llama(
-        model_path=str(MODEL_PATH),
-        n_ctx=1024,
-        n_batch=64,
-        n_threads=os.cpu_count(),    # use all available vCPUs
-        n_gpu_layers=0,
-        use_mlock=True,
-        use_mmap=True,
-        f16_kv=True,
-        batch_prefill=True,
-        prefill_logits=False,
-        verbose=True,
-    )
-    logger.info("✔ Llama model loaded")
-
-    # warm-up the model
-    try:
-        logger.info("🔄 Warming up model…")
-        llm.create_chat_completion(
-            messages=[{"role": "user", "content": "/no_think ok"}],
-            max_tokens=1,
-            stream=False,
-        )
-        logger.info("✔ Warm-up complete")
-    except Exception as e:
-        logger.error(f"Warm-up failed: {e}")
-
-    # schedule periodic keep-alive so the Space never goes idle
-    asyncio.create_task(_keep_model_warm())
-    logger.info("🔄 Keep-alive warm-up task scheduled (every 12 min)")
-
-async def _keep_model_warm():
-    while True:
-        try:
-            logger.debug("…warm-up ping")
-            # a 1-token echo
-            llm.create_chat_completion(
-                messages=[{"role": "user", "content": "/no_think ok"}],
-                max_tokens=1,
-                stream=False,
-            )
-            logger.debug("…ping done")
-        except Exception as e:
-            logger.warning(f"Warm-up ping failed: {e}")
-        # HF Spaces idle timeout is ~15 min; ping every 12
-        await asyncio.sleep(12 * 60)
+# ─── Instantiate & warm at import time ─────────────────────────────────────
+if HF_TOKEN:
+    login(token=HF_TOKEN)
+
+logger.info("Loading & warming model (this may take a moment)...")
+llm = Llama(
+    model_path=str(MODEL_PATH),
+    n_ctx=1024,
+    n_threads=os.cpu_count(),  # saturate all vCPUs
+    n_gpu_layers=0,
+    use_mlock=True,
+    f16_kv=True,
+)
+# one‐token warm-up to avoid any first‐request penalty
+llm.create_chat_completion(
+    messages=[{"role":"user","content":"/no_think ok"}],
+    max_tokens=1,
+    stream=False,
+)
+logger.info("Model ready")
 
-# ─── OpenAI‐compatible endpoint ───────────────────────────────────────────────
+# ─── OpenAI‐compatible endpoint ─────────────────────────────────────────────
 @app.post("/v1/chat/completions")
-async def chat_completions(req: dict):
+async def chat(req: dict):
     if req.get("model") != "llama-cpp":
-        raise HTTPException(status_code=404, detail="Model not found")
-
-    # call into llama.cpp
+        raise HTTPException(404, "Model not found")
     resp = llm.create_chat_completion(
         messages=req["messages"],
         max_tokens=req.get("max_tokens", 256),
@@ -116,15 +52,13 @@ async def chat_completions(req: dict):
         top_p=req.get("top_p", 1.0),
         stream=False,
     )
-
-    # repackage into OpenAI JSON
     return JSONResponse({
         "id":       resp["id"],
         "object":   "chat.completion",
         "created":  resp.get("created", int(time.time())),
         "model":    "llama-cpp",
         "choices": [{
-            "index":   0,
+            "index": 0,
             "message": {
                 "role":    resp["choices"][0]["message"]["role"],
                 "content": resp["choices"][0]["message"]["content"],