952 lines
49 KiB
Python
Executable File
952 lines
49 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
colibrì — tiny engine, immense model.
|
||
Run GLM-5.2 (744B) locally on CPU with roughly 15-26 GB of RAM.
|
||
|
||
coli chat interactive chat (loads the model once)
|
||
coli serve OpenAI-compatible HTTP API (persistent engine)
|
||
coli run "prompt" one-shot generation
|
||
coli info model, RAM, disk, and configuration status
|
||
coli plan Disk / RAM / VRAM resource plan
|
||
coli doctor installation and execution-plan diagnostics
|
||
coli bench [task...] quality benchmarks (MMLU/HellaSwag/...)
|
||
coli convert convert GLM-5.2-FP8 to int4, one shard at a time
|
||
coli build build the engine
|
||
|
||
Configuration through environment variables or flags (also valid after the subcommand):
|
||
COLI_MODEL=<dir> model directory (default /home/vincenzo/glm52_i4)
|
||
--ram N RAM budget in GB (automatically sizes the expert cache)
|
||
--repin N adapt RAM/VRAM experts every N tokens
|
||
--topp P adaptive expert top-p --topk N fixed top-k
|
||
--ngen N maximum response tokens --cap N cache slots/layer
|
||
"""
|
||
import os, sys, subprocess, argparse, json, time, signal, shutil, threading, re, codecs, tempfile, textwrap, struct
|
||
|
||
# The engine mmaps every shard (144+ files); macOS default RLIMIT_NOFILE is 256.
|
||
if sys.platform != "win32":
|
||
try:
|
||
import resource
|
||
_soft, _hard = resource.getrlimit(resource.RLIMIT_NOFILE)
|
||
_want = min(65536 if _hard == resource.RLIM_INFINITY else _hard, 65536)
|
||
if _soft < _want:
|
||
resource.setrlimit(resource.RLIMIT_NOFILE, (_want, _hard))
|
||
except (ImportError, ValueError, OSError):
|
||
pass
|
||
|
||
# Windows: forza output UTF-8 (console cp1252 tronca Unicode box-drawing/emoji)
|
||
if sys.platform == "win32":
|
||
for s in (sys.stdout, sys.stderr):
|
||
try: s.reconfigure(encoding="utf-8")
|
||
except (AttributeError, OSError): pass
|
||
|
||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||
sys.path.insert(0, HERE)
|
||
from version import __version__ as _version
|
||
|
||
# Run-in-place (source checkout, "cd c && ./coli ..."): the engine, the
|
||
# support modules (resource_plan.py, doctor.py, openai_server.py) and
|
||
# tools/ all live next to this script — unchanged from before.
|
||
#
|
||
# Installed layout ("make install"): this script is $(PREFIX)/bin/coli,
|
||
# while the engine binaries and support files live in
|
||
# $(PREFIX)/libexec/colibri, since they aren't meant to be run directly
|
||
# by users. COLI_ENGINE overrides the engine path explicitly if neither
|
||
# guess is right (e.g. a custom packaging layout).
|
||
_EXE = ".exe" if sys.platform == "win32" else ""
|
||
_LIBEXEC = os.path.join(os.path.dirname(HERE), "libexec", "colibri")
|
||
_here_glm = os.path.join(HERE, "glm" + _EXE)
|
||
|
||
if os.environ.get("COLI_ENGINE"):
|
||
GLM = os.environ["COLI_ENGINE"]
|
||
TOOLS = os.path.join(os.path.dirname(GLM), "tools")
|
||
elif os.path.exists(_here_glm):
|
||
GLM = _here_glm
|
||
TOOLS = os.path.join(HERE, "tools")
|
||
else:
|
||
GLM = os.path.join(_LIBEXEC, "glm" + _EXE)
|
||
TOOLS = os.path.join(_LIBEXEC, "tools")
|
||
sys.path.insert(0, _LIBEXEC) # so `import resource_plan`, `doctor`, `openai_server` still resolve
|
||
|
||
DEF_MODEL = os.environ.get("COLI_MODEL", "/home/vincenzo/glm52_i4")
|
||
END = b"\x01\x01END\x01\x01\n"
|
||
READY = b"\x01\x01READY\x01\x01\n"
|
||
|
||
# ---------- palette & stile ----------
|
||
def _c(n): return f"\033[38;5;{n}m"
|
||
class C:
|
||
teal=_c(37); cyan=_c(80); mag=_c(170); org=_c(208); grn=_c(78); yel=_c(179)
|
||
dim="\033[2m"; b="\033[1m"; r="\033[0m"; gray=_c(242); dgray=_c(238)
|
||
@staticmethod
|
||
def off():
|
||
for k,v in vars(C).items():
|
||
if isinstance(v,str) and v.startswith("\033"): setattr(C,k,"")
|
||
TTY = sys.stdout.isatty() or os.environ.get("COLI_COLOR")=="1"
|
||
if not TTY: C.off()
|
||
|
||
# ---------- colibrì 8-bit (pixel art, 2 pixel verticali per carattere) ----------
|
||
SPRITE = [
|
||
"....MMM.........",
|
||
"...MMMMM..w.....",
|
||
"....MMMM.ww.....",
|
||
"OOOOTTeTCC......",
|
||
"....TTTTTCC.....",
|
||
".....TTTTCC.....",
|
||
"......TTCC......",
|
||
".......TC.......",
|
||
"........C.......",
|
||
"................",
|
||
]
|
||
PAL = {"M":170, "T":37, "C":80, "O":208, "e":231, "w":80, ".":None}
|
||
|
||
def sprite_lines():
|
||
if not TTY:
|
||
return [" (\\ ", " )·> ", " / \\ ", " ", " "]
|
||
out=[]
|
||
for y in range(0,len(SPRITE),2):
|
||
top, bot = SPRITE[y], SPRITE[y+1] if y+1<len(SPRITE) else "."*len(SPRITE[y])
|
||
row=""
|
||
for x in range(len(top)):
|
||
ct, cb = PAL.get(top[x]), PAL.get(bot[x])
|
||
if ct is None and cb is None: row+= "\033[0m "
|
||
elif ct is not None and cb is None: row+= f"\033[38;5;{ct}m\033[49m▀"
|
||
elif ct is None and cb is not None: row+= f"\033[38;5;{cb}m\033[49m▄"
|
||
else: row+= f"\033[38;5;{ct}m\033[48;5;{cb}m▀"
|
||
out.append(row+"\033[0m")
|
||
return out
|
||
|
||
def banner(sub=""):
|
||
sp=sprite_lines()
|
||
txt=[
|
||
f"{C.teal}{C.b}colibrì{C.r} {C.dim}v{_version}{C.r}",
|
||
f"{C.dim}tiny engine, immense model{C.r}",
|
||
f"{C.gray}GLM-5.2 · 744B MoE · int4 · streaming CPU{C.r}",
|
||
f"{C.dgray}{sub}{C.r}" if sub else "",
|
||
"",
|
||
]
|
||
print()
|
||
for i,s in enumerate(sp):
|
||
t = txt[i] if i<len(txt) else ""
|
||
print(f" {s} {t}")
|
||
print(f" {C.dgray}{'─'*58}{C.r}")
|
||
|
||
def hline(w): return f"{C.dgray}{'─'*w}{C.r}"
|
||
|
||
# ---------- util ----------
|
||
def term_w(): return min(shutil.get_terminal_size((80,20)).columns, 100)
|
||
|
||
def need_model(model):
|
||
if not os.path.isdir(model):
|
||
sys.exit(f"{C.yel}model not found:{C.r} {model}\n set COLI_MODEL or use --model")
|
||
if not os.path.exists(os.path.join(model,"tokenizer.json")):
|
||
sys.exit(f"{C.yel}tokenizer.json is missing from {model}{C.r}")
|
||
if not os.path.exists(GLM):
|
||
sys.exit(f"{C.yel}engine is not built.{C.r} Run: coli build")
|
||
|
||
def cuda_binary():
|
||
if not os.path.exists(GLM) or sys.platform != "linux": return False
|
||
try:
|
||
linked=subprocess.run(["ldd",GLM],capture_output=True,text=True,timeout=3)
|
||
return any("libcudart" in line and "not found" not in line
|
||
for line in linked.stdout.splitlines())
|
||
except (OSError,subprocess.SubprocessError): return False
|
||
|
||
def resource_request(a, env):
|
||
ctx=a.ctx or int(env.get("CTX",4096))
|
||
ram=a.ram or float(env.get("RAM_GB",0))
|
||
vram=a.vram or float(env.get("CUDA_EXPERT_GB",0))
|
||
gpu=a.gpu
|
||
if gpu is None:
|
||
gpu=env.get("COLI_GPUS",env.get("COLI_GPU","auto"))
|
||
devices=None if gpu=="auto" else ([] if gpu=="none" else
|
||
[int(value) for value in gpu.split(",")])
|
||
return ram,ctx,devices,vram
|
||
|
||
def env_for(a):
|
||
e = dict(os.environ, SNAP=a.model)
|
||
if sys.platform == "win32":
|
||
# COLI_NO_OMP_TUNE spegne SOLO il blocco OMP (stesso perimetro del
|
||
# self-exec di glm.c; presence-based come nel motore: impostarla a
|
||
# qualsiasi valore, anche 0, disattiva). I default I/O piu' sotto
|
||
# restano attivi: kill-switch dedicati = le var stesse (DIRECT=0 ecc.)
|
||
if not e.get("COLI_NO_OMP_TUNE"):
|
||
# parita' col tuning OMP self-exec di glm.c (solo Linux/FreeBSD, e
|
||
# comunque saltato sotto CUDA/Metal): libgomp legge queste variabili
|
||
# prima di main, quindi su Windows vanno nell'ambiente del figlio.
|
||
# niente OMP_PROC_BIND/OMP_PLACES: la libgomp di MinGW non supporta
|
||
# l'affinity su Windows ("Affinity not supported on this configuration")
|
||
from resource_plan import physical_cpu_count
|
||
for k, v in (("OMP_WAIT_POLICY", "active"),
|
||
("GOMP_SPINCOUNT", "200000"),
|
||
("OMP_DYNAMIC", "FALSE"),
|
||
("OMP_NUM_THREADS", str(physical_cpu_count()))):
|
||
e.setdefault(k, v)
|
||
# Default Windows misurati sul box di riferimento (docs/tuning-9950x3d-5090.md),
|
||
# tutti lossless e tutti setdefault (un override esplicito vince sempre):
|
||
# - DIRECT=1: 10.7 GB/s O_DIRECT vs 9.0 buffered (iobench, anche a cache
|
||
# calda); nel motore 0.48 -> 1.02 tok/s. Upstream #162: 1.47x.
|
||
# - PIPE=1: overlap load/matmul, +8% sopra DIRECT (byte-identico, riordina
|
||
# solo l'I/O). PIPE_WORKERS resta al default 8 (sweep 4/8/16 piatto).
|
||
# - PILOT_REAL=1: prefetch cross-layer con load veri (unico prefetch
|
||
# funzionante su Windows: fadvise e' no-op), +11%, hit rate +19 punti.
|
||
e.setdefault("DIRECT", "1")
|
||
e.setdefault("PIPE", "1")
|
||
e.setdefault("PILOT_REAL", "1")
|
||
e["COLI_POLICY"]=a.policy
|
||
if a.ram: e["RAM_GB"]=str(a.ram)
|
||
if a.ngen: e["NGEN"]=str(a.ngen)
|
||
if a.topp: e["TOPP"]=str(a.topp)
|
||
if a.topk: e["TOPK"]=str(a.topk)
|
||
if a.temp is not None: e["TEMP"]=str(a.temp) # 0 = greedy; default motore: 1.0 + nucleus 0.95
|
||
if a.repin: e["REPIN"]=str(a.repin)
|
||
if a.ctx: e["CTX"]=str(a.ctx)
|
||
if a.auto_tier:
|
||
from resource_plan import build_plan, environment_for_plan, format_bytes
|
||
if a.gpu is not None:
|
||
e.pop("COLI_GPU",None); e.pop("COLI_GPUS",None)
|
||
if a.gpu=="none":
|
||
e["COLI_CUDA"]="0"; e.pop("CUDA_EXPERT_GB",None); e.pop("CUDA_DENSE",None)
|
||
else: e.pop("COLI_CUDA",None)
|
||
elif e.get("COLI_CUDA")=="0":
|
||
e.pop("COLI_GPU",None); e.pop("COLI_GPUS",None)
|
||
e.pop("CUDA_EXPERT_GB",None); e.pop("CUDA_DENSE",None)
|
||
if a.vram and a.gpu!="none": e["CUDA_EXPERT_GB"]=str(a.vram)
|
||
try:
|
||
ram,ctx,devices,vram=resource_request(a,e)
|
||
plan=build_plan(a.model,ram,ctx,devices,vram,policy=a.policy)
|
||
except (OSError,ValueError,json.JSONDecodeError) as error:
|
||
sys.exit(f"{C.yel}invalid resource plan:{C.r} {error}")
|
||
has_cuda=cuda_binary()
|
||
e=environment_for_plan(plan,e,has_cuda)
|
||
rt=plan["tiers"]["ram"]; vt=plan["tiers"]["vram"]
|
||
gpu=f" · VRAM {format_bytes(vt['budget_bytes'])}" if has_cuda and vt["devices"] else " · CPU"
|
||
print(f" {C.dim}[PLAN] RAM {format_bytes(rt['budget_bytes'])} · cap {rt['cache_slots_per_layer']}/layer{gpu}{C.r}",file=sys.stderr)
|
||
else:
|
||
# --gpu/--vram SENZA --auto-tier: prima venivano ignorati in silenzio e il run
|
||
# partiva CPU-only senza alcun avviso — benchmark "GPU" pubblicati per errore (#121).
|
||
if a.gpu is not None:
|
||
e.pop("COLI_GPU",None); e.pop("COLI_GPUS",None)
|
||
if a.gpu=="none":
|
||
e["COLI_CUDA"]="0"; e.pop("CUDA_EXPERT_GB",None); e.pop("CUDA_DENSE",None)
|
||
else:
|
||
if not cuda_binary():
|
||
sys.exit(f"{C.yel}--gpu needs the CUDA build:{C.r} make glm CUDA=1 (this binary is CPU-only)")
|
||
e["COLI_CUDA"]="1"
|
||
if a.gpu!="auto": e["COLI_GPUS"]=a.gpu
|
||
e.setdefault("CUDA_DENSE","1")
|
||
if a.vram and a.gpu!="none":
|
||
if not cuda_binary():
|
||
sys.exit(f"{C.yel}--vram needs the CUDA build:{C.r} make glm CUDA=1 (this binary is CPU-only)")
|
||
e["COLI_CUDA"]="1"; e["CUDA_EXPERT_GB"]=str(a.vram)
|
||
return e
|
||
|
||
# ---------- rendering markdown in STREAMING per il terminale ----------
|
||
class MDStream:
|
||
"""Interpreta il markdown della risposta mentre arriva: i ``` diventano riquadri,
|
||
**x** grassetto vero, `x` colorato, # titoli, - puntini. I marker non si vedono mai.
|
||
Regge i chunk spezzati a meta' marker (hold-back) e l'output sporco (``` doppi)."""
|
||
def __init__(self, indent=" "):
|
||
self.ind=indent
|
||
self.cur="" # riga parziale non ancora emessa
|
||
self.code=False; self.lang=""
|
||
self.bold=False; self.icode=False
|
||
self.justclosed=False # l'ultima riga era una chiusura ```? (anti ``` doppi)
|
||
self.printed=0 # caratteri della riga corrente gia' emessi
|
||
def _fence(self, line):
|
||
lang=line.strip()[3:].strip().strip("`")
|
||
if not self.code:
|
||
if not lang and self.justclosed: return # ``` orfano dopo una chiusura: rumore, ignora
|
||
self.code=True; self.lang=lang
|
||
sys.stdout.write(f"{self.ind}{C.dgray}\u256d\u2500 {lang or 'code'}{C.r}\n")
|
||
elif lang: # ```lang mentre siamo GIA' in code: chiudi e riapri
|
||
sys.stdout.write(f"{self.ind}{C.dgray}\u2570\u2500{C.r}\n{self.ind}{C.dgray}\u256d\u2500 {lang}{C.r}\n")
|
||
self.lang=lang
|
||
else:
|
||
self.code=False; self.justclosed=True
|
||
sys.stdout.write(f"{self.ind}{C.dgray}\u2570\u2500{C.r}\n")
|
||
def _inline(self, txt, out):
|
||
i=0
|
||
while i<len(txt):
|
||
ch=txt[i]
|
||
if ch=="`":
|
||
self.icode=not self.icode
|
||
out.append(C.org if self.icode else C.r); i+=1; continue
|
||
if ch=="*":
|
||
j=i
|
||
while j<len(txt) and txt[j]=="*": j+=1
|
||
if j-i>=2: # **/***: grassetto on/off, gli asterischi spariscono
|
||
self.bold=not self.bold
|
||
out.append(C.b if self.bold else C.r)
|
||
else: out.append("*") # * singolo: lascialo (moltiplicazioni ecc.)
|
||
i=j; continue
|
||
out.append(ch); i+=1
|
||
def _line(self, line, partial=False):
|
||
if not partial and line.lstrip().startswith("```"):
|
||
self._fence(line); self.printed=0; return
|
||
if line.strip(): self.justclosed=False
|
||
seg=line[self.printed:] # emetti solo la parte nuova della riga
|
||
out=[]
|
||
if self.code:
|
||
if self.printed==0: out.append(f"{self.ind}{C.dgray}\u2502{C.r} {C.cyan}")
|
||
out.append(seg)
|
||
else:
|
||
if self.printed==0:
|
||
out.append(self.ind)
|
||
st=seg.lstrip()
|
||
if st.startswith("#"): # titolo: via i #, grassetto teal
|
||
seg=st.lstrip("#").strip(); out.append(f"{C.teal}{C.b}"); self.bold=True
|
||
elif st.startswith(("- ","* ")): # lista: puntino vero
|
||
seg=st[2:]; out.append(f"{C.teal}\u2022{C.r} ")
|
||
self._inline(seg,out)
|
||
sys.stdout.write("".join(out)); sys.stdout.flush()
|
||
self.printed=len(line)
|
||
if not partial: # fine riga: reset stati inline (robusto ai marker orfani)
|
||
sys.stdout.write(C.r+"\n"); sys.stdout.flush()
|
||
self.bold=self.icode=False
|
||
self.printed=0
|
||
def feed(self, s):
|
||
self.cur+=s
|
||
while "\n" in self.cur:
|
||
line,self.cur=self.cur.split("\n",1)
|
||
self._line(line)
|
||
st=self.cur.lstrip() # riga parziale: possibile fence? aspetta il newline
|
||
if st and (st.startswith("```") or (len(st)<3 and "```".startswith(st))):
|
||
return
|
||
if st.startswith("#") and self.printed==0:
|
||
return # titolo: rendi la riga intera al newline
|
||
hold=0 # trattieni marker potenzialmente spezzati in coda
|
||
while hold<len(self.cur) and self.cur[-1-hold] in "*`": hold+=1
|
||
safe=self.cur[:len(self.cur)-hold] if hold else self.cur
|
||
if len(safe)>self.printed: self._line(safe, partial=True)
|
||
def close(self):
|
||
if self.cur: self._line(self.cur); self.cur=""
|
||
if self.code:
|
||
sys.stdout.write(f"\n{self.ind}{C.dgray}\u2570\u2500{C.r}"); self.code=False
|
||
sys.stdout.write(C.r); sys.stdout.flush()
|
||
|
||
class Spinner:
|
||
FRAMES=["⠋","⠙","⠹","⠸","⠼","⠴","⠦","⠧","⠇","⠏"]
|
||
def __init__(self,label,tick=None):
|
||
self.label=label; self.tick=tick; self.suffix=""
|
||
self.stop_evt=threading.Event(); self.t0=time.time(); self.th=None
|
||
def start(self):
|
||
if not TTY: return
|
||
def run():
|
||
i=0
|
||
while not self.stop_evt.is_set():
|
||
el=time.time()-self.t0
|
||
if self.tick and i%8==0: # ~1 Hz: legge il progresso dal log
|
||
try: self.suffix=self.tick() or self.suffix
|
||
except Exception: pass
|
||
suf=f" {C.dgray}· {self.suffix}{C.r}" if self.suffix else ""
|
||
sys.stdout.write(f"\r {C.teal}{self.FRAMES[i%10]}{C.r} {C.dim}{self.label} {el:.0f}s{C.r}{suf}\033[K")
|
||
sys.stdout.flush(); i+=1; time.sleep(0.12)
|
||
self.th=threading.Thread(target=run,daemon=True); self.th.start()
|
||
def stop(self):
|
||
self.stop_evt.set()
|
||
if self.th: self.th.join(timeout=0.4)
|
||
if TTY: sys.stdout.write("\r\033[K"); sys.stdout.flush()
|
||
|
||
def engine_diag(p, errlog=None):
|
||
"""Lo stdout del motore ha chiuso: il processo e' morto. Dire PERCHE'.
|
||
Un SIGKILL dell'OOM-killer del kernel e' altrimenti invisibile — niente errore,
|
||
niente exit code, il motore muore muto e sembra un bug nostro (issue #305).
|
||
Quasi sempre e' memoria: il picco reale ha sforato la RAM della macchina."""
|
||
try: rc=p.wait(timeout=5)
|
||
except Exception: rc=p.poll()
|
||
if rc is None: why="its output closed but the process is still alive"
|
||
elif rc<0: # POSIX: morte per segnale
|
||
try: why=f"killed by {signal.Signals(-rc).name}"
|
||
except Exception: why=f"killed by signal {-rc}"
|
||
elif rc>0: why=f"exit code {rc}"
|
||
else: why="exited cleanly"
|
||
print(f"\n {C.yel}[engine terminated: {why}]{C.r}")
|
||
if rc is not None and rc<0 and -rc==getattr(signal,"SIGKILL",-1):
|
||
print(f" {C.yel}nothing in the engine sends SIGKILL to itself: this is the kernel's\n"
|
||
f" OOM-killer. The peak RSS exceeded the machine's free memory.\n"
|
||
f" Lower --ram, lower PIN_GB, or shorten the context.{C.r}")
|
||
if errlog is not None:
|
||
try:
|
||
errlog.flush()
|
||
for _ in range(20): # il drain thread scrive su errlog quando il child chiude
|
||
tail=open(errlog.name).read().strip()
|
||
if tail: break
|
||
time.sleep(0.05)
|
||
if tail: print(f" {C.dgray}" + "\n ".join(tail.splitlines()[-6:]) + f"{C.r}")
|
||
except Exception: pass
|
||
|
||
def stream_turn(p, sentinel, on_bytes):
|
||
"""legge fino alla sentinella; on_bytes riceve i chunk della risposta. Poi legge la riga STAT.
|
||
Il PRIMO Ctrl-C durante lo stream non chiude la sessione: il motore (handler SIGINT)
|
||
chiude il turno per la via del tetto NGEN e noi dreniamo fino alla sentinella.
|
||
Un SECONDO Ctrl-C esce davvero."""
|
||
pend=b""; interrupted=False
|
||
while True:
|
||
try:
|
||
b=p.stdout.read(1)
|
||
except KeyboardInterrupt:
|
||
if interrupted or p.poll() is not None: raise
|
||
interrupted=True
|
||
try: p.send_signal(signal.SIGINT) # non-TTY: il motore potrebbe non aver visto il Ctrl-C
|
||
except Exception: pass
|
||
print(f"\n {C.yel}⏹ stopping… (Ctrl-C again to quit){C.r}", flush=True)
|
||
continue
|
||
if b==b"": return None
|
||
pend+=b
|
||
if pend.endswith(sentinel):
|
||
rest=pend[:-len(sentinel)]
|
||
if rest: on_bytes(rest)
|
||
line=p.stdout.readline().decode("utf-8","replace").strip() # STAT tok tps hit rss
|
||
m=re.match(r"STAT (\S+) (\S+) (\S+) (\S+)", line)
|
||
st={"tok":int(m.group(1)),"tps":float(m.group(2)),"hit":float(m.group(3)),"rss":float(m.group(4))} if m else {}
|
||
if interrupted: st["interrupted"]=True
|
||
return st
|
||
if len(pend)>len(sentinel):
|
||
out=pend[:-len(sentinel)]; pend=pend[-len(sentinel):]
|
||
on_bytes(out)
|
||
|
||
# ---------- comandi ----------
|
||
def cmd_build(a):
|
||
banner("build")
|
||
if not os.path.exists(os.path.join(HERE, "Makefile")):
|
||
sys.exit(f"{C.yel}coli build{C.r} only works from a source checkout (this is an installed copy).\n"
|
||
f" Clone https://github.com/JustVugg/colibri and run ./setup.sh, or make -C c glm.")
|
||
sys.exit(subprocess.call(["make","-C",HERE,"glm"]))
|
||
|
||
def cmd_info(a):
|
||
banner("info")
|
||
cfgp=os.path.join(a.model,"config.json")
|
||
def row(k,v): print(f" {C.gray}{k:<10}{C.r} {v}")
|
||
if os.path.exists(cfgp):
|
||
c=json.load(open(cfgp))
|
||
row("model", a.model)
|
||
row("arch", f"hidden {c.get('hidden_size')} · {c.get('num_hidden_layers')} layer · "
|
||
f"{c.get('n_routed_experts')} expert/layer · top-{c.get('num_experts_per_tok')}")
|
||
sts=[x for x in os.listdir(a.model) if x.endswith('.safetensors')]
|
||
sz=sum(os.path.getsize(os.path.join(a.model,x)) for x in sts)
|
||
row("shards", f"{len(sts)} files · {sz/1e9:.0f} GB on disk")
|
||
else:
|
||
print(f" {C.yel}config.json is missing (incomplete conversion?){C.r}")
|
||
try:
|
||
mi=open('/proc/meminfo').read()
|
||
tot=int(re.search(r'MemTotal:\s+(\d+)',mi).group(1))/1e6
|
||
av=int(re.search(r'MemAvailable:\s+(\d+)',mi).group(1))/1e6
|
||
row("RAM", f"{tot:.0f} GB total · {av:.1f} GB available")
|
||
except Exception: pass
|
||
try:
|
||
fs = shutil.disk_usage(a.model if os.path.isdir(a.model) else HERE)
|
||
row("disk", f"{fs.free/1e9:.0f} GB free")
|
||
except OSError:
|
||
row("disk", "? GB (unavailable)")
|
||
row("engine", "ready ✓" if os.path.exists(GLM) else "not built (coli build)")
|
||
knobs=[]
|
||
if a.ram: knobs.append(f"ram {a.ram}GB")
|
||
if a.topp: knobs.append(f"topp {a.topp}")
|
||
if a.topk: knobs.append(f"topk {a.topk}")
|
||
if knobs: row("tuning", " · ".join(knobs))
|
||
print()
|
||
|
||
def cmd_plan(a):
|
||
from resource_plan import build_plan, format_plan
|
||
try:
|
||
ram,ctx,devices,vram=resource_request(a,os.environ)
|
||
if ctx<1: raise ValueError("--ctx must be positive")
|
||
if a.vram<0: raise ValueError("--vram cannot be negative")
|
||
plan=build_plan(a.model,ram,ctx,devices,vram,policy=a.policy)
|
||
except (OSError, ValueError, json.JSONDecodeError) as error:
|
||
sys.exit(f"{C.yel}cannot create resource plan:{C.r} {error}")
|
||
if a.json:
|
||
print(json.dumps(plan,indent=2))
|
||
return
|
||
banner("plan · Disk / RAM / VRAM")
|
||
print(textwrap.indent(format_plan(plan)," "))
|
||
print()
|
||
|
||
def cmd_doctor(a):
|
||
from doctor import exit_code, format_doctor, run_doctor
|
||
try:
|
||
ram,ctx,devices,vram=resource_request(a,os.environ)
|
||
if ctx<1: raise ValueError("--ctx must be positive")
|
||
if ram<0: raise ValueError("--ram cannot be negative")
|
||
if vram<0: raise ValueError("--vram cannot be negative")
|
||
except ValueError as error:
|
||
report={"schema_version":1,"status":"error","model":os.path.abspath(a.model),
|
||
"checks":[{"id":"config.arguments","status":"fail","summary":str(error)}],
|
||
"plan":None}
|
||
print(json.dumps(report,indent=2) if a.json else format_doctor(report))
|
||
return 2
|
||
report=run_doctor(a.model,ram,ctx,devices,vram,engine_path=GLM)
|
||
print(json.dumps(report,indent=2) if a.json else format_doctor(report))
|
||
return exit_code(report)
|
||
|
||
def cmd_run(a):
|
||
need_model(a.model)
|
||
prompt=" ".join(a.prompt) if a.prompt else sys.exit('usage: coli run "your prompt"')
|
||
banner("run")
|
||
# template ufficiale GLM-5.2: niente \n dopo i ruoli; <think></think> = risposta diretta (nothink)
|
||
e=env_for(a); e["PROMPT"]=f"[gMASK]<sop><|user|>{prompt}<|assistant|><think></think>"
|
||
sys.exit(subprocess.call([GLM, str(a.cap)], env=e))
|
||
|
||
def server_probe(base, api_key=None, timeout=1.5):
|
||
"""Is a coli serve alive at `base`? Returns its model_id, or None.
|
||
Probes /health then /v1/models — both cheap, neither touches the engine."""
|
||
import urllib.request, urllib.error
|
||
def get(path):
|
||
req=urllib.request.Request(base.rstrip("/")+path)
|
||
if api_key: req.add_header("Authorization", f"Bearer {api_key}")
|
||
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||
return json.loads(r.read().decode("utf-8","replace"))
|
||
try:
|
||
if get("/health").get("status")!="ok": return None
|
||
data=get("/v1/models").get("data") or []
|
||
return data[0]["id"] if data else None
|
||
except Exception:
|
||
return None
|
||
|
||
def chat_attached(a, base, model_id):
|
||
"""The chat REPL over HTTP against a running `coli serve`.
|
||
|
||
Why this exists (the cold-chat cost, measured): spawning a private engine
|
||
pays 34-136 s of resident load on EVERY start, and begins with an empty
|
||
expert cache — hit rate 4% cold vs 55% warm, a ~10x on early decode. A
|
||
resident server pays load once and keeps the LRU warm across sessions;
|
||
its KV slots reuse the conversation prefix, so a continued chat skips
|
||
re-prefill too. The engine byte-protocol stays untouched — this is plain
|
||
OpenAI SSE over localhost, stdlib only."""
|
||
import urllib.request
|
||
print(f" {C.grn}✦ attached{C.r} {C.dim}to {base} · model {model_id} · the engine stays warm after you quit{C.r}")
|
||
print(f" {C.dim}type and press Enter · Ctrl-C stops the answer · :reset starts a new conversation · :q exits{C.r}\n")
|
||
msgs=[]
|
||
w=term_w()-4
|
||
while True:
|
||
if TTY:
|
||
print(f" {C.dgray}╭{'─'*w}╮{C.r}")
|
||
try: msg=input(f" {C.dgray}│{C.r} {C.teal}{C.b}›{C.r} ")
|
||
except EOFError: print(); break
|
||
print(f" {C.dgray}╰{'─'*w}╯{C.r}")
|
||
else:
|
||
try: msg=input()
|
||
except EOFError: break
|
||
msg=msg.strip()
|
||
if msg in (":q",":quit","exit"): break
|
||
if not msg: continue
|
||
if msg==":reset": msgs=[]; print(f" {C.dim}✦ new conversation{C.r}\n"); continue
|
||
msgs.append({"role":"user","content":msg})
|
||
body=json.dumps({"model":model_id,"messages":msgs,"stream":True,
|
||
"max_tokens":a.ngen}).encode()
|
||
req=urllib.request.Request(base.rstrip("/")+"/v1/chat/completions", data=body,
|
||
headers={"Content-Type":"application/json"})
|
||
if a.api_key: req.add_header("Authorization", f"Bearer {a.api_key}")
|
||
print(f"\n {C.teal}◆ colibrì{C.r}")
|
||
sp=Spinner("thinking…"); sp.start()
|
||
md=MDStream(" "); reply=[]; first=True; t0=time.time(); interrupted=False
|
||
try:
|
||
with urllib.request.urlopen(req) as r:
|
||
for raw in r:
|
||
line=raw.decode("utf-8","replace").strip()
|
||
if not line.startswith("data: "): continue
|
||
data=line[6:]
|
||
if data=="[DONE]": break
|
||
try: ev=json.loads(data)
|
||
except ValueError: continue
|
||
for ch in ev.get("choices",[]):
|
||
d=ch.get("delta",{})
|
||
txt=d.get("content")
|
||
if not txt: continue # ping/ruolo/reasoning: non è testo
|
||
if first: sp.stop(); first=False
|
||
md.feed(txt); reply.append(txt)
|
||
except KeyboardInterrupt:
|
||
interrupted=True # il server annulla la richiesta alla disconnessione
|
||
except OSError as e:
|
||
sp.stop()
|
||
print(f"\n {C.yel}[server unreachable: {e}]{C.r}"); break
|
||
sp.stop(); md.close()
|
||
if reply: msgs.append({"role":"assistant","content":"".join(reply)})
|
||
else: msgs.pop() # turno vuoto: non sporcare la history
|
||
el=time.time()-t0
|
||
note=" · ⏹ interrupted" if interrupted else ""
|
||
print(f"\r {C.dgray}└─ ~{len(''.join(reply))//4} tok · {el:.0f}s{note}{C.r}\n")
|
||
print(f" {C.dim}goodbye — the engine keeps running for the next chat 🐦{C.r}")
|
||
|
||
def kv_resume_notice(model_dir):
|
||
"""SERVE mode silently resumes .coli_kv from disk (glm.c kv_disk_load): a chat
|
||
started today continues a conversation from days ago, with `first=0` so the
|
||
turn is appended WITHOUT the [gMASK]<sop> prefix. The engine does announce it
|
||
on stderr — but nothing here ever shows that: the drain thread's
|
||
p.stderr.read() blocks until EOF, so on a healthy start errlog is still empty
|
||
when the status lines are printed. The warning only appeared once the engine
|
||
DIED, which is exactly when it no longer mattered.
|
||
|
||
Measured cost of the silence: a chat inherited 670 tokens of an old Italian
|
||
session ("il mio numero preferito e 7, ricordalo!"). Every later reply came
|
||
back in Italian, and "explain fibonacci in short" was answered about the
|
||
number 7 — the model was being coherent with a context nobody could see, and
|
||
it read as a quantization bug for a day.
|
||
|
||
So say it here, in Python, from the file itself: no pipe, no thread, no
|
||
Windows deadlock risk (see the stderr comment below)."""
|
||
p=os.path.join(model_dir, ".coli_kv")
|
||
try:
|
||
with open(p,"rb") as f:
|
||
if f.read(8)!=b"COLIKV1\0": return
|
||
h=struct.unpack("<8i", f.read(32))
|
||
n=h[6]
|
||
if n<1: return
|
||
age=time.time()-os.path.getmtime(p)
|
||
when=f"{age/86400:.0f}d ago" if age>86400 else f"{age/3600:.0f}h ago" if age>3600 else "just now"
|
||
print(f" {C.yel}↺ resuming a saved conversation: {n} tokens, last written {when}{C.r}")
|
||
print(f" {C.dgray} it steers tone, language and topic. :reset clears it · "
|
||
f"KVSAVE=0 disables saving · delete {p} to start clean{C.r}")
|
||
except (OSError, struct.error): pass
|
||
|
||
def cmd_chat(a):
|
||
# ATTACH: a running `coli serve` beats a private engine every time — the load
|
||
# (34-136 s) and the cache warmth survive between sessions. Explicit --attach
|
||
# wins; otherwise probe localhost quietly and use it if it's there. --no-attach
|
||
# forces the old behaviour. The probe costs ~1 ms when nothing is listening.
|
||
if not getattr(a,"no_attach",False):
|
||
base=getattr(a,"attach",None) or "http://127.0.0.1:8000"
|
||
mid=server_probe(base, getattr(a,"api_key",None))
|
||
if mid:
|
||
banner(f"chat · {mid} · attached")
|
||
chat_attached(a, base, mid); return
|
||
if getattr(a,"attach",None):
|
||
sys.exit(f"--attach: no coli serve answering at {base} (start one with: coli serve --model <dir>)")
|
||
need_model(a.model)
|
||
banner(f"chat · {os.path.basename(a.model)} · ram {a.ram or '-'}GB · topp {a.topp or 'off'}")
|
||
kv_resume_notice(a.model)
|
||
errlog=tempfile.NamedTemporaryFile(mode="w+", suffix=".log", delete=False)
|
||
e=env_for(a); e["SERVE"]="1"
|
||
# stderr -> PIPE, NOT stderr=errlog (file). On Windows/MinGW, pointing the
|
||
# child's stderr at a file/DEVNULL handle stalls the CRT so stdout (the byte
|
||
# protocol coli reads one byte at a time) never flushes and chat hangs at
|
||
# ~10 GB resident. A PIPE whose read end nobody drains still works: the
|
||
# engine emits only ~400 bytes of status to stderr, which fits comfortably
|
||
# in the OS pipe buffer, so it never blocks. We snapshot stderr into errlog
|
||
# once the READY sentinel arrives, so the status-line display below works
|
||
# exactly as before. (Do NOT add a concurrent stderr drain thread: on
|
||
# Windows, reading two child pipes simultaneously deadlocks CPython's IO.)
|
||
p=subprocess.Popen([GLM,str(a.cap)], env=e, stdin=subprocess.PIPE,
|
||
stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=0)
|
||
sp=Spinner("waking the giant (744B)…"); sp.start()
|
||
st=stream_turn(p, READY, lambda b: None)
|
||
sp.stop()
|
||
if st is None:
|
||
try: errlog.write(p.stderr.read().decode("utf-8","replace"))
|
||
except (OSError, ValueError): pass
|
||
errlog.flush(); errlog.seek(0); print(errlog.read()[-1500:])
|
||
engine_diag(p) # perche' e' morto (OOM-kill compreso); errlog e' gia' stampato sopra
|
||
sys.exit("the engine exited while loading")
|
||
p.stdout.readline() # TIERS line (web-dashboard protocol): emitted once right after STAT;
|
||
# left unread it leaks into the first answer's text
|
||
# READY received. Drain the child's stderr into errlog without blocking:
|
||
# the engine is still alive (blocked on stdin), so a plain read() would
|
||
# hang forever waiting for EOF. A short bounded drain grabs the ~400 bytes
|
||
# of load-time status ([RAM_GB], [MTP], ...) that were already emitted.
|
||
_drain_box={"done":False}
|
||
def _drain():
|
||
try: errlog.write(p.stderr.read().decode("utf-8","replace"))
|
||
except (OSError, ValueError): pass
|
||
_drain_box["done"]=True
|
||
threading.Thread(target=_drain, daemon=True).start()
|
||
_drain_box["th"]=threading.current_thread()
|
||
for _ in range(20): # up to ~1s for the load-status lines
|
||
if _drain_box["done"]: break
|
||
time.sleep(0.05)
|
||
errlog.flush()
|
||
try:
|
||
elog=open(errlog.name).read()
|
||
mload=re.search(r"loaded in ([0-9.]+)s \| resident dense: ([0-9.]+) MB", elog)
|
||
if mload: print(f" {C.grn}✓{C.r} ready in {mload.group(1)}s {C.dim}· resident {float(mload.group(2))/1000:.1f} GB · RSS {st.get('rss','?')} GB{C.r}")
|
||
for l in elog.splitlines(): # una riga di stato per riga, senza path
|
||
if l.startswith(("[RAM_GB","[PIN]","[MTP]","[USAGE]","[DSA]","[KV]")):
|
||
l=re.sub(r" ?\(?/[^ )]+\)?","",l.strip()) # via i percorsi lunghi
|
||
l=re.sub(r" from$","",l)
|
||
for chunk in textwrap.wrap(l, term_w()-4) or [l]:
|
||
print(f" {C.dgray}{chunk}{C.r}")
|
||
except Exception: pass
|
||
print(f" {C.dim}type and press Enter · Ctrl-C stops the answer · :more continues · :reset clears memory · :q exits{C.r}\n")
|
||
w=term_w()-4
|
||
def user_box(msg):
|
||
"""ri-disegna il messaggio dentro una box che si ADATTA su piu' righe:
|
||
l'input grezzo (che sborda) viene cancellato e sostituito dal testo avvolto."""
|
||
cols=shutil.get_terminal_size((80,20)).columns
|
||
used=max(1, (6+len(msg)+cols-1)//cols) # righe occupate dall'input (" │ › "+msg)
|
||
sys.stdout.write(f"\x1b[{used}A\x1b[0J") # su di N righe e pulisci fino in fondo
|
||
inner=w-3 # spazio utile: " │ › "+testo+"│" = w+4 colonne
|
||
lines=textwrap.wrap(msg, inner) or [""]
|
||
for i,ln in enumerate(lines):
|
||
pre = f"{C.teal}{C.b}›{C.r}" if i==0 else " "
|
||
print(f" {C.dgray}│{C.r} {pre} {ln}{' '*(inner-len(ln))}{C.dgray}│{C.r}")
|
||
print(f" {C.dgray}╰{'─'*w}╯{C.r}")
|
||
try:
|
||
while True:
|
||
if TTY:
|
||
print(f" {C.dgray}╭{'─'*w}╮{C.r}")
|
||
try: msg=input(f" {C.dgray}│{C.r} {C.teal}{C.b}›{C.r} ")
|
||
except EOFError: print(); break
|
||
try: user_box(msg.strip())
|
||
except Exception: print(f" {C.dgray}╰{'─'*w}╯{C.r}")
|
||
else:
|
||
try: msg=input()
|
||
except EOFError: break
|
||
msg=msg.strip()
|
||
if msg in (":q",":quit","exit"): break
|
||
if not msg: continue
|
||
if msg==":reset":
|
||
p.stdin.write(b"\x02RESET\n"); p.stdin.flush()
|
||
stream_turn(p, END, lambda b: None)
|
||
print(f" {C.dim}✦ memory cleared{C.r}\n"); continue
|
||
if msg in (":piu",":più",":more",":continua"):
|
||
p.stdin.write(b"\x02MORE\n"); p.stdin.flush()
|
||
else:
|
||
p.stdin.write((msg.replace("\n"," ")+"\n").encode()); p.stdin.flush()
|
||
print(f"\n {C.teal}◆ colibrì{C.r}")
|
||
dec=codecs.getincrementaldecoder("utf-8")("replace")
|
||
state={"first":True}
|
||
def prefill_tick(path=errlog.name):
|
||
try:
|
||
with open(path) as f:
|
||
f.seek(max(0, os.path.getsize(path)-1500)); tail=f.read()
|
||
pl=[l for l in tail.splitlines() if l.startswith("[prefill]")]
|
||
return pl[-1].replace("[prefill] ","prefill ") if pl else ""
|
||
except Exception: return ""
|
||
sp2=Spinner("thinking…", tick=prefill_tick); sp2.start()
|
||
md=MDStream(" ") # markdown -> terminale, in streaming
|
||
raw=os.environ.get("COLI_RAW")=="1"
|
||
def echo(bs, _dec=dec, _st=state):
|
||
if _st["first"]:
|
||
sp2.stop(); _st["first"]=False
|
||
if raw: sys.stdout.write(" ")
|
||
s=_dec.decode(bs)
|
||
if not s: return
|
||
if raw: sys.stdout.write(s.replace("\n","\n ")); sys.stdout.flush()
|
||
else: md.feed(s)
|
||
t0=time.time()
|
||
st=stream_turn(p, END, echo)
|
||
if not raw: md.close()
|
||
sp2.stop()
|
||
if st is None: engine_diag(p, errlog); break
|
||
el=time.time()-t0
|
||
if st.get("tok"):
|
||
print(f"\r {C.dgray}└─ {st['tok']} tok · {st['tps']:.2f} tok/s · hit {st['hit']:.0f}% · RSS {st['rss']:.1f} GB · {el:.0f}s{C.r}")
|
||
if st.get("interrupted"):
|
||
print(f" {C.yel}⏹ interrupted; type :more to continue the response{C.r}")
|
||
elif st["tok"]>=a.ngen:
|
||
print(f" {C.yel}…stopped at --ngen ({a.ngen}); type :more to continue the response{C.r}")
|
||
print()
|
||
else:
|
||
if st.get("interrupted"): print(f" {C.yel}⏹ interrupted{C.r}")
|
||
print()
|
||
except KeyboardInterrupt:
|
||
print(f"\n {C.dim}interrupted{C.r}")
|
||
finally:
|
||
try: p.stdin.close(); p.terminate()
|
||
except Exception: pass
|
||
try: os.unlink(errlog.name)
|
||
except Exception: pass
|
||
print(f" {C.teal}goodbye{C.r} {C.dim}— the hummingbird returns to its nest{C.r} 🐦\n")
|
||
|
||
def serve_pidfile(port): return os.path.join(tempfile.gettempdir(), f"coli-serve-{port}.pid")
|
||
|
||
def cmd_serve(a):
|
||
need_model(a.model)
|
||
# pidfile: cosi' `coli stop` spegne tutto con un comando, senza pkill a mano.
|
||
# EN: pidfile so `coli stop` can shut everything down without manual pkill.
|
||
try:
|
||
with open(serve_pidfile(a.port),"w") as f: f.write(f"{os.getpid()} {a.model}\n")
|
||
except OSError: pass
|
||
from openai_server import serve
|
||
try:
|
||
serve(a.model, a.host, a.port, a.model_id, a.api_key,
|
||
a.cap,a.ngen,GLM,env_for(a),a.cors_origin,
|
||
a.max_queue,a.queue_timeout,a.kv_slots)
|
||
finally:
|
||
try: os.unlink(serve_pidfile(a.port))
|
||
except OSError: pass
|
||
|
||
def cmd_stop(a):
|
||
"""Shut down a running `coli serve` AND its engine — one command, no pkill.
|
||
The engine re-execs itself for OMP tuning, so its process is named `exe`,
|
||
not `glm`: every `pkill -x glm` in history silently killed nothing (that is
|
||
how two 17+5 GB ghost engines OOM'd this box on 2026-07-16). This finds the
|
||
real processes: the pidfile first, then /proc by cmdline/environ — only
|
||
processes that are demonstrably ours (SERVE=1 + our SNAP, or `coli serve`
|
||
in the command line)."""
|
||
banner("stop")
|
||
targets=[] # (pid, descrizione)
|
||
pf=serve_pidfile(a.port)
|
||
try:
|
||
pid=int(open(pf).read().split()[0])
|
||
os.kill(pid,0); targets.append((pid,f"coli serve (pidfile, port {a.port})"))
|
||
except (OSError,ValueError,IndexError): pass
|
||
for pd in os.listdir("/proc"):
|
||
if not pd.isdigit(): continue
|
||
pid=int(pd)
|
||
try:
|
||
cmd=open(f"/proc/{pd}/cmdline","rb").read().replace(b"\0",b" ").decode("utf-8","replace")
|
||
if "coli" in cmd and " serve" in cmd and pid!=os.getpid():
|
||
if not any(p==pid for p,_ in targets): targets.append((pid,"coli serve (cmdline)"))
|
||
comm=open(f"/proc/{pd}/comm").read().strip()
|
||
if comm in ("glm","exe","olmoe"):
|
||
env=open(f"/proc/{pd}/environ","rb").read().replace(b"\0",b"\n").decode("utf-8","replace")
|
||
if "SERVE=1" in env: targets.append((pid,f"engine `{comm}` (SERVE=1)"))
|
||
except (OSError,PermissionError): continue
|
||
if not targets:
|
||
print(f" nothing running — no serve on port {a.port}, no SERVE engines"); return
|
||
for pid,desc in targets: print(f" {'would stop' if a.dry_run else 'stopping'} {pid}: {desc}")
|
||
if a.dry_run: return
|
||
for pid,_ in targets:
|
||
try: os.kill(pid, signal.SIGTERM)
|
||
except OSError: pass
|
||
time.sleep(2.0)
|
||
for pid,_ in targets:
|
||
try: os.kill(pid, signal.SIGKILL); print(f" {pid}: forced (SIGKILL)")
|
||
except OSError: pass # gia' morto: bene
|
||
try: os.unlink(pf)
|
||
except OSError: pass
|
||
print(f" {C.grn}✓ stopped{C.r} — RAM released")
|
||
|
||
def cmd_web(a):
|
||
"""serve + open the dashboard in the browser once the API answers."""
|
||
need_model(a.model)
|
||
dist = os.path.join(os.path.dirname(os.path.abspath(HERE)), "web", "dist")
|
||
if not os.path.exists(os.path.join(dist, "index.html")):
|
||
print(f"{C.yel}web UI not built:{C.r} run cd web && npm install && npm run build first;")
|
||
print("serving the API anyway (the dashboard will 404 until built).")
|
||
url = f"http://{a.host}:{a.port}/"
|
||
if not getattr(a, "no_browser", False):
|
||
import threading, urllib.request, webbrowser
|
||
def opener():
|
||
for _ in range(600): # the 744B engine takes minutes to load
|
||
time.sleep(2)
|
||
try:
|
||
urllib.request.urlopen(f"http://{a.host}:{a.port}/health", timeout=2)
|
||
webbrowser.open(url); return
|
||
except OSError:
|
||
continue
|
||
threading.Thread(target=opener, daemon=True).start()
|
||
print(f"dashboard: {url} (opens automatically when the engine is ready)")
|
||
cmd_serve(a)
|
||
|
||
def cmd_bench(a):
|
||
need_model(a.model)
|
||
banner("bench")
|
||
# python con `tokenizers`: l'ambiente del progetto se c'e', altrimenti quello corrente
|
||
venv_py = os.path.join(HERE, "mio_env", "Scripts" if sys.platform == "win32" else "bin", "python3")
|
||
py = venv_py if os.path.exists(venv_py) else sys.executable
|
||
tasks = ",".join(a.tasks) if a.tasks else "hellaswag,arc_challenge,mmlu"
|
||
# dataset mancanti -> li scarica una volta (fetch_benchmarks.py li mette in --data come JSONL)
|
||
missing=[t for t in tasks.split(",") if not os.path.exists(os.path.join(a.data,f"{t}.jsonl"))]
|
||
if missing:
|
||
print(f" {C.dim}downloading missing datasets: {', '.join(missing)}{C.r}")
|
||
subprocess.call([py, os.path.join(TOOLS,"fetch_benchmarks.py"),
|
||
"--out", a.data, "--tasks", ",".join(missing), "--limit", str(max(a.limit,200))])
|
||
# il fetch riprova da solo (#304), ma se l'hub resta giu' il bench gira sui task
|
||
# disponibili invece di passare a eval file inesistenti.
|
||
# EN: the fetch retries on its own (#304), but if the hub stays down the bench
|
||
# runs on the available tasks instead of handing eval nonexistent files.
|
||
still=[t for t in tasks.split(",") if not os.path.exists(os.path.join(a.data,f"{t}.jsonl"))]
|
||
if still:
|
||
tasks=",".join(t for t in tasks.split(",") if t not in still)
|
||
print(f" {C.yel}skipping (download failed, rerun later): {', '.join(still)}{C.r}")
|
||
if not tasks:
|
||
print(f" {C.yel}no datasets available — nothing to bench{C.r}"); sys.exit(1)
|
||
cmd=[py, os.path.join(TOOLS,"eval_glm.py"), "--glm", GLM, "--snap",a.model,
|
||
"--tasks", tasks, "--limit", str(a.limit), "--data", a.data]
|
||
if a.ram: cmd+=["--ram",str(a.ram)]
|
||
e=env_for(a)
|
||
print(f" {C.dim}decode is disk-bound: this takes HOURS on slow hardware. Raise --limit on faster machines.{C.r}\n")
|
||
sys.exit(subprocess.call(cmd, env=e))
|
||
|
||
def cmd_convert(a):
|
||
banner("convert")
|
||
# python con torch/safetensors: l'ambiente del progetto se c'e', altrimenti quello corrente
|
||
venv_py = os.path.join(HERE, "mio_env", "Scripts" if sys.platform == "win32" else "bin", "python3")
|
||
py = venv_py if os.path.exists(venv_py) else sys.executable
|
||
base=[py, os.path.join(TOOLS,"convert_fp8_to_int4.py"),
|
||
"--repo", a.repo, "--outdir", a.model, "--ebits", str(a.ebits), "--io-bits", str(a.io_bits)]
|
||
if a.xbits: base+=["--xbits",str(a.xbits)]
|
||
# passo 1: modello principale (78 layer). Resumabile: riparte dagli shard mancanti.
|
||
print(f" {C.dim}[1/2] model: {' '.join(base)}{C.r}")
|
||
rc=subprocess.call(base)
|
||
if rc!=0: sys.exit(rc)
|
||
if a.no_mtp: sys.exit(0)
|
||
# passo 2: testa MTP (layer 78). SEMPRE int8: a int4 i draft sbagliano quasi sempre
|
||
# (acceptance 0-4% vs 39-59%, misurato — issue #8) e la speculazione non parte mai.
|
||
mtp_cmd=list(base); i=mtp_cmd.index("--ebits"); mtp_cmd[i+1]=str(max(8,a.ebits))
|
||
print(f" {C.dim}[2/2] int8 MTP head (speculative drafts){C.r}")
|
||
sys.exit(subprocess.call(mtp_cmd+["--mtp"]))
|
||
|
||
def main():
|
||
common=argparse.ArgumentParser(add_help=False)
|
||
common.add_argument("--model", default=DEF_MODEL); common.add_argument("--ram", type=int, default=0) # 0 = auto (il motore usa l'88% della RAM disponibile)
|
||
common.add_argument("--auto-tier",action="store_true",help="automatically apply the RAM/VRAM plan")
|
||
common.add_argument("--ctx",type=int,default=0)
|
||
common.add_argument("--gpu",default=None,help="auto, none, or a device list such as 0,1")
|
||
common.add_argument("--vram",type=float,default=0,help="total VRAM budget in GB (0=auto)")
|
||
common.add_argument("--policy",choices=("quality","balanced","experimental-fast"),
|
||
default=os.environ.get("COLI_POLICY","quality"),
|
||
help="resource policy (explicit --topk/--topp overrides warn and proceed)")
|
||
common.add_argument("--repin", type=int, default=0, help="adapt RAM/VRAM experts every N tokens")
|
||
common.add_argument("--cap", type=int, default=8); common.add_argument("--ngen", type=int, default=1024) # rete di sicurezza: la fine vera la decidono gli stop token
|
||
common.add_argument("--topp", type=float, default=0); common.add_argument("--topk", type=int, default=0)
|
||
common.add_argument("--temp", type=float, default=None) # temperatura token (0=greedy, default 1.0+nucleus .95)
|
||
ap=argparse.ArgumentParser(prog="coli", parents=[common], description="colibrì — run GLM-5.2 locally")
|
||
ap.add_argument("--version", action="version", version=f"colibrì {_version}")
|
||
sub=ap.add_subparsers(dest="cmd")
|
||
sub.add_parser("build", parents=[common]); sub.add_parser("info", parents=[common])
|
||
pp=sub.add_parser("plan",parents=[common])
|
||
pp.add_argument("--json",action="store_true")
|
||
pd=sub.add_parser("doctor",parents=[common])
|
||
pd.add_argument("--json",action="store_true",help="emit a versioned JSON report")
|
||
pr=sub.add_parser("run", parents=[common]); pr.add_argument("prompt", nargs="*")
|
||
pc=sub.add_parser("chat", parents=[common])
|
||
pc.add_argument("--attach", nargs="?", const="http://127.0.0.1:8000", default=None,
|
||
help="chat against a running `coli serve` instead of spawning an engine "
|
||
"(keeps the model loaded and the expert cache warm across chat sessions). "
|
||
"Bare --attach probes localhost:8000.")
|
||
pc.add_argument("--no-attach", action="store_true",
|
||
help="never auto-attach, always spawn a private engine")
|
||
pc.add_argument("--api-key", default=os.environ.get("COLI_API_KEY"))
|
||
ps=sub.add_parser("serve", parents=[common])
|
||
ps.add_argument("--host",default="127.0.0.1"); ps.add_argument("--port",type=int,default=8000)
|
||
ps.add_argument("--model-id",default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))
|
||
ps.add_argument("--api-key",default=os.environ.get("COLI_API_KEY"))
|
||
ps.add_argument("--cors-origin",action="append",default=None)
|
||
ps.add_argument("--max-queue",type=int,default=int(os.environ.get("COLI_MAX_QUEUE","8")))
|
||
ps.add_argument("--queue-timeout",type=float,default=float(os.environ.get("COLI_QUEUE_TIMEOUT","300")))
|
||
ps.add_argument("--kv-slots",type=int,default=int(os.environ.get("COLI_KV_SLOTS","1")))
|
||
pst=sub.add_parser("stop", parents=[common], help="shut down a running coli serve and its engine")
|
||
pst.add_argument("--port",type=int,default=8000); pst.add_argument("--dry-run",action="store_true")
|
||
pw=sub.add_parser("web", parents=[common], help="serve + open the dashboard in a browser")
|
||
for arg,kw in (("--host",dict(default="127.0.0.1")),("--port",dict(type=int,default=8000)),
|
||
("--model-id",dict(default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))),
|
||
("--api-key",dict(default=os.environ.get("COLI_API_KEY"))),
|
||
("--cors-origin",dict(action="append",default=None)),
|
||
("--max-queue",dict(type=int,default=int(os.environ.get("COLI_MAX_QUEUE","8")))),
|
||
("--queue-timeout",dict(type=float,default=float(os.environ.get("COLI_QUEUE_TIMEOUT","300")))),
|
||
("--kv-slots",dict(type=int,default=int(os.environ.get("COLI_KV_SLOTS","1"))))):
|
||
pw.add_argument(arg,**kw)
|
||
pw.add_argument("--no-browser",action="store_true",help="don't auto-open the browser")
|
||
pb=sub.add_parser("bench", parents=[common]); pb.add_argument("tasks", nargs="*")
|
||
pb.add_argument("--limit",type=int,default=40)
|
||
if sys.platform == "win32":
|
||
_cache_root = os.environ.get("LOCALAPPDATA", os.path.expanduser("~\\AppData\\Local"))
|
||
else:
|
||
_cache_root = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
|
||
_bench_cache = os.path.join(_cache_root, "colibri", "bench")
|
||
pb.add_argument("--data",default=_bench_cache)
|
||
pc=sub.add_parser("convert", parents=[common]); pc.add_argument("--repo",default="zai-org/GLM-5.2-FP8")
|
||
pc.add_argument("--ebits",type=int,default=4); pc.add_argument("--io-bits",type=int,default=8); pc.add_argument("--xbits",type=int,default=0)
|
||
pc.add_argument("--no-mtp",action="store_true",help="skip the MTP head (no speculative drafts)")
|
||
a=ap.parse_args()
|
||
handler={"build":cmd_build,"info":cmd_info,"plan":cmd_plan,"doctor":cmd_doctor,
|
||
"run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"stop":cmd_stop,"bench":cmd_bench,
|
||
"convert":cmd_convert,"web":cmd_web}.get(a.cmd)
|
||
if handler: sys.exit(handler(a) or 0)
|
||
banner(); print(__doc__)
|
||
|
||
if __name__=="__main__":
|
||
signal.signal(signal.SIGINT, signal.default_int_handler)
|
||
main()
|