Harvad Li
bc688a6211
The engine opens/mmaps every safetensors shard of the model (144+ files for GLM-5.2). On macOS the default soft fd limit is 256, so a stock terminal session fails while loading with: <model>/out-00125.safetensors: Too many open files the engine exited while loading Raise the soft limit toward the hard limit (capped at 65536) in the coli launcher before spawning the engine, so users don't need a manual 'ulimit -n' in every shell. No-op on Windows and on shells whose limit is already sufficient. Co-authored-by: Harvad Lee <hongyanab@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
638 lines
32 KiB
Python
Executable File
638 lines
32 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
|
||
|
||
# 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__))
|
||
|
||
# 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}v1.0{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)
|
||
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 stream_turn(p, sentinel, on_bytes):
|
||
"""legge fino alla sentinella; on_bytes riceve i chunk della risposta. Poi legge la riga STAT."""
|
||
pend=b""
|
||
while True:
|
||
b=p.stdout.read(1)
|
||
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)
|
||
return {"tok":int(m.group(1)),"tps":float(m.group(2)),"hit":float(m.group(3)),"rss":float(m.group(4))} if m else {}
|
||
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 cmd_chat(a):
|
||
need_model(a.model)
|
||
banner(f"chat · {os.path.basename(a.model)} · ram {a.ram or '-'}GB · topp {a.topp or 'off'}")
|
||
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.seek(0); print(errlog.read()[-1500:]); sys.exit("the engine exited while loading")
|
||
# 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 · :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: print(f"\n {C.yel}[engine terminated]{C.r}"); 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["tok"]>=a.ngen:
|
||
print(f" {C.yel}…stopped at --ngen ({a.ngen}); type :more to continue the response{C.r}")
|
||
print()
|
||
else:
|
||
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 cmd_serve(a):
|
||
need_model(a.model)
|
||
from openai_server import serve
|
||
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)
|
||
|
||
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))])
|
||
cmd=[py, os.path.join(TOOLS,"eval_glm.py"), "--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")
|
||
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="*")
|
||
sub.add_parser("chat", parents=[common])
|
||
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")))
|
||
pb=sub.add_parser("bench", parents=[common]); pb.add_argument("tasks", nargs="*")
|
||
pb.add_argument("--limit",type=int,default=40); pb.add_argument("--data",default=os.path.join(HERE,"bench"))
|
||
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,"bench":cmd_bench,
|
||
"convert":cmd_convert}.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()
|