main.py

"""
KVInfer — FastAPI Backend  v4.1
2 vCPU · 16 GB RAM HuggingFace Space ke liye optimize kiya hua

RAM estimate:
  2 engines × 4 GB (Llama 1B float32) = 8.0 GB
  2 engines × 8 sess × ~48 MB KV      = 0.8 GB
  Python + FastAPI + tokenizer         = ~0.7 GB
  ────────────────────────────────────────────
  TOTAL                               ≈ 9.5 GB  ✓ (16 GB mein safe)
"""
import asyncio, json, os, time, uuid
from contextlib import asynccontextmanager
from pathlib import Path

import psutil
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel, Field
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer

# ── Config ───────────────────────────────────────────────
BASE_DIR      = Path(__file__).parent
INFERENCE_EXE = BASE_DIR / "inference"
MODEL_BIN     = BASE_DIR / "model_llama.bin"
HF_REPO_ID    = os.environ.get("HF_REPO_ID", "YOUR_HF_USERNAME/YOUR_REPO")

BLOCK_SIZE      = 2048
MAX_GEN_CEILING = 500
SAFETY_MARGIN   = 50
MAX_SESS_TOKENS = BLOCK_SIZE - MAX_GEN_CEILING - SAFETY_MARGIN  # 1498

# 2 vCPU → 2 engines, each pinned to 1 thread
N_ENGINES = int(os.environ.get("N_ENGINES", "2"))

# Llama 3 special tokens
EOS_IDS  = [128001, 128009]   # <|end_of_text|>  <|eot_id|>
EOT_STR  = "<|eot_id|>"
SYS_H    = "<|start_header_id|>system<|end_header_id|>\n\n"
USR_H    = "<|start_header_id|>user<|end_header_id|>\n\n"
AST_H    = "<|start_header_id|>assistant<|end_header_id|>\n\n"
STOP_STR = ["<|eot_id|>", "<|start_header_id|>user", "<|start_header_id|>system"]

tokenizer = None

def load_tokenizer():
    global tokenizer
    local = BASE_DIR / "tokenizer_files"
    src   = str(local) if local.exists() else "unsloth/Llama-3.2-1B-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(src)
    print(f"[tok] vocab={tokenizer.vocab_size}")

def enc(text: str) -> list[int]:
    return tokenizer.encode(text, add_special_tokens=False)

def dec(ids: list[int]) -> str:
    return tokenizer.decode(ids, skip_special_tokens=False)

# ── Engine ───────────────────────────────────────────────
class Engine:
    def __init__(self, eid):
        self.eid = eid; self._proc = None; self._ready = False

    async def start(self):
        if not INFERENCE_EXE.exists(): raise RuntimeError("Binary not found")
        if not MODEL_BIN.exists():     raise RuntimeError("model_llama.bin not found")
        env = os.environ.copy()
        env["OMP_NUM_THREADS"] = "1"   # 1 thread per engine = 2 threads total on 2vCPU
        self._proc = await asyncio.create_subprocess_exec(
            str(INFERENCE_EXE),
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.DEVNULL,
            cwd=str(BASE_DIR), env=env,
        )
        while True:
            line = (await self._proc.stdout.readline()).decode().strip()
            if line.startswith("[engine]"): print(f"[E{self.eid}] {line}")
            elif line == "READY":
                self._ready = True
                print(f"[E{self.eid}] READY pid={self._proc.pid}")
                break
            elif line.startswith("ERROR"): raise RuntimeError(line)

    async def stop(self):
        if not self._proc: return
        try:
            self._proc.stdin.write(b"QUIT\n"); await self._proc.stdin.drain()
            await asyncio.wait_for(self._proc.wait(), 3.0)
        except: self._proc.kill()

    async def reset(self, sid):
        self._proc.stdin.write(f"RESET|{sid}\n".encode())
        await self._proc.stdin.drain()
        while True:
            raw = await self._proc.stdout.readline()
            if not raw or raw.decode().strip() == "RESET_OK": break

    async def generate(self, sid, tokens, max_new, temp, top_k):
        if not self._ready: yield {"type":"error","message":"not ready"}; return
        cmd = f"REQUEST|{sid}|{','.join(map(str,tokens))}|{max_new}|{temp}|{top_k}|{','.join(map(str,EOS_IDS))}\n"
        self._proc.stdin.write(cmd.encode()); await self._proc.stdin.drain()
        try:
            while True:
                raw = await self._proc.stdout.readline()
                if not raw: break
                line = raw.decode("utf-8","replace").strip()
                if not line: continue
                if line.startswith("TOKEN"):
                    p = line.split(); yield {"type":"token","id":int(p[1]),"text":dec([int(p[1])]),"elapsed_ms":float(p[2])}
                elif line.startswith("DONE"):
                    p = line.split(); t=int(p[1]); ms=float(p[2])
                    yield {"type":"done","total_tokens":t,"total_ms":ms,
                           "tps": round(t/(ms/1000),2) if ms>0 else 0}; break
                elif line.startswith("ERROR"):
                    yield {"type":"error","message":line}; break
        except asyncio.CancelledError:
            while True:
                raw = await self._proc.stdout.readline()
                if not raw or raw.decode().strip().startswith(("DONE","ERROR")): break
            raise

    @property
    def pid(self): return self._proc.pid if self._proc else None

# ── Pool ─────────────────────────────────────────────────
class Pool:
    def __init__(self, n):
        self.n=n; self.engines=[Engine(i) for i in range(n)]
        self._locks=[]; self._smap={}; self._load=[]; self._ml=None

    async def start(self):
        self._ml=asyncio.Lock(); self._locks=[asyncio.Lock() for _ in range(self.n)]
        self._load=[0]*self.n
        await asyncio.gather(*(e.start() for e in self.engines))
        print(f"[pool] {self.n} engines up")

    async def stop(self):
        await asyncio.gather(*(e.stop() for e in self.engines),return_exceptions=True)

    async def _assign(self, sid):
        async with self._ml:
            if sid not in self._smap:
                idx=min(range(self.n),key=lambda i:self._load[i])
                self._smap[sid]=idx; self._load[idx]+=1
            return self._smap[sid]

    async def _drop(self, sid):
        async with self._ml:
            if sid in self._smap:
                idx=self._smap.pop(sid); self._load[idx]=max(0,self._load[idx]-1)

    async def generate(self, sid, tokens, max_new, temp, top_k):
        idx=await self._assign(sid)
        async with self._locks[idx]:
            async for c in self.engines[idx].generate(sid,tokens,max_new,temp,top_k): yield c

    async def reset(self, sid):
        async with self._ml: idx=self._smap.get(sid)
        if idx is not None:
            async with self._locks[idx]: await self.engines[idx].reset(sid)
        await self._drop(sid)

    def pids(self): return [e.pid for e in self.engines if e.pid]

    def status(self):
        return [{"id":i,"pid":self.engines[i].pid,"sessions":self._load[i],
                 "busy":self._locks[i].locked(),"ready":self.engines[i]._ready}
                for i in range(self.n)]

pool = Pool(N_ENGINES)

# ── Session ──────────────────────────────────────────────
class Sess:
    def __init__(self, sys_p):
        self.sys_p=sys_p; self.history=[]; self.n_cached=0

    def push_user(self, m): self.history.append({"role":"user","content":m})
    def push_asst(self, m): self.history.append({"role":"assistant","content":m})

    def new_tokens(self, msg):
        if self.n_cached == 0:
            text = f"<|begin_of_text|>{SYS_H}{self.sys_p}{EOT_STR}{USR_H}{msg}{EOT_STR}{AST_H}"
        else:
            text = f"{USR_H}{msg}{EOT_STR}{AST_H}"
        return enc(text)

sessions: dict[str, Sess] = {}
metrics = {"req":0,"tok":0,"ms":0.0,"err":0,"t0":time.time()}

def total_ram():
    try:
        mb=psutil.Process(os.getpid()).memory_info().rss
        for p in pool.pids():
            try: mb+=psutil.Process(p).memory_info().rss
            except: pass
        return round(mb/1e6,1)
    except: return 0.0

# ── Lifespan ─────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app):
    print("[start] Loading tokenizer…")
    load_tokenizer()
    if not MODEL_BIN.exists():
        try:
            print("[start] Downloading model_llama.bin from HF…")
            hf_hub_download(repo_id=HF_REPO_ID,filename="model_llama.bin",local_dir=str(BASE_DIR))
        except Exception as e: print(f"[warn] download failed: {e}")
    try: await pool.start()
    except Exception as e: print(f"[error] pool start: {e}")
    yield
    await pool.stop()

app = FastAPI(title="KVInfer",version="4.1",lifespan=lifespan)
app.add_middleware(CORSMiddleware,allow_origins=["*"],allow_methods=["*"],allow_headers=["*"])

# ── Models ───────────────────────────────────────────────
class ChatReq(BaseModel):
    message:        str
    session_id:     str   = Field(default_factory=lambda: str(uuid.uuid4()))
    system_prompt:  str   = "You are a helpful, concise assistant."
    max_new_tokens: int   = Field(default=256, ge=1, le=500)
    temperature:    float = Field(default=0.7, ge=0.01, le=2.0)
    top_k:          int   = Field(default=40, ge=1, le=200)

class ResetReq(BaseModel):
    session_id: str

# ── Routes ───────────────────────────────────────────────
@app.get("/")
async def ui(): return FileResponse(BASE_DIR/"index.html")

@app.get("/health")
async def health():
    mem=psutil.virtual_memory()
    return {"status":"ok" if any(e._ready for e in pool.engines) else "starting",
            "engines_ready":sum(1 for e in pool.engines if e._ready),
            "engines_total":N_ENGINES,"active_sessions":len(sessions),
            "process_ram_mb":total_ram(),"system_ram_used_pct":mem.percent,
            "uptime_seconds":round(time.time()-metrics["t0"],1)}

@app.get("/metrics")
async def get_metrics():
    n,tok,ms=metrics["req"],metrics["tok"],metrics["ms"]
    mem=psutil.virtual_memory()
    return {"total_requests":n,"total_tokens":tok,"total_errors":metrics["err"],
            "avg_tps":round(tok/(ms/1000),2) if ms>0 else 0,
            "active_sessions":len(sessions),"n_engines":N_ENGINES,
            "engines_ready":sum(1 for e in pool.engines if e._ready),
            "engines_busy":sum(1 for lk in pool._locks if lk.locked()),
            "process_ram_mb":total_ram(),"system_ram_used_pct":mem.percent,
            "uptime_s":round(time.time()-metrics["t0"],1)}

@app.post("/chat")
async def chat(req: ChatReq):
    if not any(e._ready for e in pool.engines):
        raise HTTPException(503,"No engines ready yet — please wait a moment.")
    sess=sessions.setdefault(req.session_id, Sess(req.system_prompt))
    toks=sess.new_tokens(req.message)
    if sess.n_cached+len(toks)+req.max_new_tokens > MAX_SESS_TOKENS:
        await pool.reset(req.session_id); sess.n_cached=0; toks=sess.new_tokens(req.message)
    sess.push_user(req.message); metrics["req"]+=1

    async def stream():
        parts=[]; t0=time.time(); stopped=False
        try:
            async for c in pool.generate(req.session_id,toks,req.max_new_tokens,req.temperature,req.top_k):
                if c["type"]=="token" and not stopped:
                    parts.append(c["text"]); joined="".join(parts)
                    for s in STOP_STR:
                        if s in joined: parts=[joined[:joined.find(s)]]; stopped=True; break
                    if not stopped: yield f"data:{json.dumps(c)}\n\n"
                elif c["type"]=="done":
                    reply="".join(parts).strip()
                    for s in STOP_STR: reply=reply.split(s)[0]
                    reply=reply.strip()
                    sess.push_asst(reply)
                    sess.n_cached+=len(toks)+c["total_tokens"]
                    metrics["tok"]+=c["total_tokens"]; metrics["ms"]+=(time.time()-t0)*1000
                    yield f"data:{json.dumps({**c,'session_id':req.session_id,'full_response':reply})}\n\n"
                elif c["type"]=="error":
                    metrics["err"]+=1; yield f"data:{json.dumps(c)}\n\n"
        except Exception as e:
            metrics["err"]+=1; yield f"data:{json.dumps({'type':'error','message':str(e)})}\n\n"
        finally: yield "data:[DONE]\n\n"

    return StreamingResponse(stream(),media_type="text/event-stream",
        headers={"Cache-Control":"no-cache","X-Accel-Buffering":"no"})

@app.post("/chat/reset")
async def reset(req: ResetReq):
    sessions.pop(req.session_id, None)
    await pool.reset(req.session_id)
    return {"status":"ok","session_id":req.session_id}

@app.get("/chat/history")
async def history(session_id: str):
    s=sessions.get(session_id)
    if not s: return {"session_id":session_id,"turns":0,"history":[]}
    return {"session_id":session_id,"tokens_in_engine":s.n_cached,
            "turns":sum(1 for m in s.history if m["role"]=="user"),"history":s.history}

@app.get("/pool/status")
async def pool_status(): return {"n_engines":N_ENGINES,"engines":pool.status(),"sessions":len(sessions)}

if __name__=="__main__":
    import uvicorn; uvicorn.run("main:app",host="0.0.0.0",port=7860,reload=False)