Similar setup to other model

Dockerfile

@@ -1,34 +1,41 @@
 FROM python:3.10-slim
 
 ENV DEBIAN_FRONTEND=noninteractive \
+    RUSTUP_HOME=/root/.rustup \
+    CARGO_HOME=/root/.cargo \
+    PATH=/root/.cargo/bin:$PATH \
     MODEL_REPO="unsloth/Qwen3-0.6B-GGUF" \
     MODEL_FILE="Qwen3-0.6B-Q4_K_M.gguf" \
-    HF_HOME=/app/cache
+    HF_HOME=/app/cache \
+    TRANSFORMERS_CACHE=/app/cache
 
-# system deps + rust for llama-cpp
+WORKDIR /app
+
+# install system deps + Rust toolchain
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
-      build-essential cmake git curl wget libgomp1 ca-certificates && \
+      build-essential cmake git curl wget ninja-build libgomp1 ca-certificates \
+      gcc g++ libffi-dev libopenblas-dev libstdc++6 libgcc-s1 && \
     rm -rf /var/lib/apt/lists/* && \
     curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && \
-    . "$HOME/.cargo/env" && rustup default stable
+    rustup default stable
 
-WORKDIR /app
-
-RUN mkdir -p /app/cache /app/logs /app/workspace /app/pretrained_models/llm \
- && chmod -R 777 /app/cache /app/logs /app/workspace /app/pretrained_models/llm
+# cache dirs
+RUN mkdir -p /app/cache /app/pretrained_models/llm && chmod -R 777 /app/cache /app/pretrained_models/llm
 
+# Python deps (except llama-cpp-python)
 COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
-
-RUN git clone --recurse-submodules https://github.com/abetlen/llama-cpp-python.git /tmp/llama-cpp-python \
- && cd /tmp/llama-cpp-python \
- && FORCE_CMAKE=1 pip install --no-cache-dir . \
- && pip install --no-cache-dir "llama-cpp-python[server]" huggingface_hub \
- && rm -rf /tmp/llama-cpp-python
+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)
+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 LLM server code
+# copy the server code
 COPY llm_server.py /app/llm_server.py
 
 EXPOSE 7860

llm_server.py

@@ -1,11 +1,21 @@
-import os, time, logging
+import os
+import time
+import logging
+import 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 llama_cpp import Llama
-from huggingface_hub import hf_hub_download, login
 
+# ─── logging setup ────────────────────────────────────────────────────────────
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger("llm_server")
+
+# ─── FastAPI setup ────────────────────────────────────────────────────────────
 app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
@@ -14,54 +24,95 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-MODEL_DIR = Path("/app/pretrained_models/llm")
+# ─── model paths & env ────────────────────────────────────────────────────────
+MODEL_DIR  = Path("/app/pretrained_models/llm")
 MODEL_DIR.mkdir(parents=True, exist_ok=True)
-MODEL_FILE = os.getenv("MODEL_FILE")
+
+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():
-    logging.info("Starting LLM service…")
+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():
-        token = os.getenv("HF_TOKEN")
-        if token:
-            login(token=token)
+        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)
+            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_threads=2,
+        n_threads=os.cpu_count(),    # use all available vCPUs
         n_gpu_layers=0,
         use_mlock=True,
         f16_kv=True,
+        batch_prefill=True,
+        prefill_logits=False,
+        verbose=False,
     )
-    logging.info("LLM loaded.")
+    logger.info("✔ Llama model loaded")
+
+    # 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)
+
+# ─── OpenAI‐compatible endpoint ───────────────────────────────────────────────
 @app.post("/v1/chat/completions")
-async def chat(req: dict):
+async def chat_completions(req: dict):
     if req.get("model") != "llama-cpp":
         raise HTTPException(status_code=404, detail="Model not found")
+
+    # call into llama.cpp
     resp = llm.create_chat_completion(
         messages=req["messages"],
         max_tokens=req.get("max_tokens", 256),
         temperature=req.get("temperature", 0.7),
         top_p=req.get("top_p", 1.0),
-        stream=False
+        stream=False,
     )
+
+    # repackage into OpenAI JSON
     return JSONResponse({
-        "id": resp["id"],
-        "object": "chat.completion",
-        "created": resp.get("created", int(time.time())),
-        "model": "llama-cpp",
+        "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"],
+                "role":    resp["choices"][0]["message"]["role"],
                 "content": resp["choices"][0]["message"]["content"],
             },
             "finish_reason": resp["choices"][0].get("finish_reason", "stop"),