coli plan: read-only Disk/RAM/VRAM placement planner (mirrors engine budget math, versioned --json, --auto-tier) (#27)
* Add tiered resource planner * Apply resource plans to runner configuration --------- Co-authored-by: JustVugg <JustVugg@users.noreply.github.com>
This commit is contained in:
ZacharyZcR
@@ -91,6 +91,25 @@ cd c
|
||||
COLI_MODEL=/nvme/glm52_i4 ./coli chat
|
||||
```
|
||||
|
||||
Inspect the planned storage hierarchy before loading the model:
|
||||
|
||||
```bash
|
||||
COLI_MODEL=/nvme/glm52_i4 ./coli plan
|
||||
COLI_MODEL=/nvme/glm52_i4 ./coli plan --gpu 0,1 --ram 128 --vram 48 --json
|
||||
|
||||
# apply the bounded plan to the normal runner
|
||||
COLI_MODEL=/nvme/glm52_i4 ./coli chat --auto-tier
|
||||
```
|
||||
|
||||
`coli plan` reads only safetensors headers and reports the model's exact dense/expert
|
||||
footprint, runtime RAM reserve, safe expert-cache cap, and bounded VRAM hot tier. Its
|
||||
versioned JSON output is intended to be shared by the CLI, API server, Web UI, and
|
||||
desktop shell; it does not allocate model tensors or start inference.
|
||||
`--auto-tier` applies the same plan to `chat`, `run`, `serve`, and benchmarks. It
|
||||
sets the RAM budget and context immediately; the VRAM tier is enabled only when
|
||||
the current `glm` binary is linked with CUDA. Explicit flags and environment
|
||||
variables keep precedence over automatic values.
|
||||
|
||||
The engine at runtime is pure C — python is only used by the one-time converter.
|
||||
|
||||
### OpenAI-compatible API
|
||||
|
||||
@@ -7,6 +7,7 @@ CLI per far girare GLM-5.2 (744B) in locale, su CPU, in ~15-26 GB di RAM.
|
||||
coli serve API HTTP compatibile OpenAI (motore persistente)
|
||||
coli run "prompt" generazione singola
|
||||
coli info stato: modello, RAM, disco, config
|
||||
coli plan piano risorse Disk / RAM / VRAM
|
||||
coli bench [task...] benchmark di qualità (MMLU/HellaSwag/...)
|
||||
coli convert converte GLM-5.2-FP8 -> int4 (streaming)
|
||||
coli build compila il motore
|
||||
@@ -98,6 +99,25 @@ def need_model(model):
|
||||
if not os.path.exists(GLM):
|
||||
sys.exit(f"{C.yel}motore non compilato.{C.r} Esegui: 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 a.ram: e["RAM_GB"]=str(a.ram)
|
||||
@@ -106,6 +126,28 @@ def env_for(a):
|
||||
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)
|
||||
except (OSError,ValueError,json.JSONDecodeError) as error:
|
||||
sys.exit(f"{C.yel}piano risorse non valido:{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)
|
||||
return e
|
||||
|
||||
# ---------- rendering markdown in STREAMING per il terminale ----------
|
||||
@@ -267,6 +309,22 @@ def cmd_info(a):
|
||||
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 deve essere positivo")
|
||||
if a.vram<0: raise ValueError("--vram non puo essere negativo")
|
||||
plan=build_plan(a.model,ram,ctx,devices,vram)
|
||||
except (OSError, ValueError, json.JSONDecodeError) as error:
|
||||
sys.exit(f"{C.yel}impossibile creare il piano:{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_run(a):
|
||||
need_model(a.model)
|
||||
prompt=" ".join(a.prompt) if a.prompt else sys.exit('uso: coli run "il tuo prompt"')
|
||||
@@ -400,9 +458,7 @@ def cmd_bench(a):
|
||||
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=dict(os.environ)
|
||||
if a.topp: e["TOPP"]=str(a.topp)
|
||||
if a.topk: e["TOPK"]=str(a.topk)
|
||||
e=env_for(a)
|
||||
print(f" {C.dim}decode disk-bound: su hardware lento questo richiede ORE. Alza --limit su macchine veloci.{C.r}\n")
|
||||
sys.exit(subprocess.call(cmd, env=e))
|
||||
|
||||
@@ -428,6 +484,10 @@ def cmd_convert(a):
|
||||
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="applica automaticamente il piano RAM/VRAM")
|
||||
common.add_argument("--ctx",type=int,default=0)
|
||||
common.add_argument("--gpu",default=None,help="auto, none oppure lista device, es. 0,1")
|
||||
common.add_argument("--vram",type=float,default=0,help="budget VRAM totale in GB (0=auto)")
|
||||
common.add_argument("--repin", type=int, default=0, help="adatta gli expert RAM/VRAM ogni N token")
|
||||
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)
|
||||
@@ -435,6 +495,8 @@ def main():
|
||||
ap=argparse.ArgumentParser(prog="coli", parents=[common], description="colibrì — GLM-5.2 in locale")
|
||||
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")
|
||||
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])
|
||||
@@ -448,7 +510,7 @@ def main():
|
||||
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="salta la testa MTP (niente draft speculativi)")
|
||||
a=ap.parse_args()
|
||||
{"build":cmd_build,"info":cmd_info,"run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"bench":cmd_bench,
|
||||
{"build":cmd_build,"info":cmd_info,"plan":cmd_plan,"run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"bench":cmd_bench,
|
||||
"convert":cmd_convert}.get(a.cmd, lambda _:(banner(),print(__doc__)))(a)
|
||||
|
||||
if __name__=="__main__":
|
||||
|
||||
@@ -0,0 +1,216 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Hardware and model placement planning for colibri's disk/RAM/VRAM tiers."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import statistics
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
GB = 1_000_000_000
|
||||
EXPERT_RE = re.compile(r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.")
|
||||
|
||||
|
||||
def _tensor_sizes(path):
|
||||
file_size = path.stat().st_size
|
||||
with path.open("rb") as stream:
|
||||
raw = stream.read(8)
|
||||
if len(raw) != 8:
|
||||
raise ValueError(f"short safetensors header: {path}")
|
||||
length = int.from_bytes(raw, "little")
|
||||
if length < 2 or length > file_size - 8:
|
||||
raise ValueError(f"invalid safetensors header length: {path}")
|
||||
header = json.loads(stream.read(length))
|
||||
for name, meta in header.items():
|
||||
if name == "__metadata__":
|
||||
continue
|
||||
start, end = meta["data_offsets"]
|
||||
if not 0 <= start <= end <= file_size - 8 - length:
|
||||
raise ValueError(f"invalid tensor offsets for {name}: {path}")
|
||||
yield name, end - start
|
||||
|
||||
|
||||
def analyze_model(model):
|
||||
model = Path(model).resolve()
|
||||
config_path = model / "config.json"
|
||||
if not config_path.is_file():
|
||||
raise ValueError(f"missing config.json: {model}")
|
||||
config = json.loads(config_path.read_text())
|
||||
shards = sorted(model.glob("*.safetensors"))
|
||||
if not shards:
|
||||
raise ValueError(f"no safetensors shards: {model}")
|
||||
|
||||
dense_bytes = 0
|
||||
expert_groups = {}
|
||||
for shard in shards:
|
||||
for name, size in _tensor_sizes(shard):
|
||||
match = EXPERT_RE.search(name)
|
||||
if match:
|
||||
key = tuple(map(int, match.groups()))
|
||||
expert_groups[key] = expert_groups.get(key, 0) + size
|
||||
else:
|
||||
dense_bytes += size
|
||||
|
||||
layer_sizes = {}
|
||||
for (layer, _), size in expert_groups.items():
|
||||
layer_sizes.setdefault(layer, []).append(size)
|
||||
per_layer = {layer: int(statistics.median(sizes)) for layer, sizes in layer_sizes.items()}
|
||||
per_cap_bytes = sum(per_layer.values())
|
||||
typical_expert_bytes = int(statistics.median(per_layer.values())) if per_layer else 0
|
||||
model_bytes = sum(shard.stat().st_size for shard in shards)
|
||||
return {
|
||||
"path": str(model),
|
||||
"shards": len(shards),
|
||||
"model_bytes": model_bytes,
|
||||
"dense_bytes": dense_bytes,
|
||||
"expert_bytes": sum(expert_groups.values()),
|
||||
"expert_count": len(expert_groups),
|
||||
"expert_layers": len(per_layer),
|
||||
"typical_expert_bytes": typical_expert_bytes,
|
||||
"per_cap_bytes": per_cap_bytes,
|
||||
"config": config,
|
||||
}
|
||||
|
||||
|
||||
def memory_available():
|
||||
try:
|
||||
text = Path("/proc/meminfo").read_text()
|
||||
return int(re.search(r"MemAvailable:\s+(\d+)", text).group(1)) * 1024
|
||||
except (OSError, AttributeError):
|
||||
return 0
|
||||
|
||||
|
||||
def discover_gpus():
|
||||
command = ["nvidia-smi", "--query-gpu=index,name,memory.total,memory.free",
|
||||
"--format=csv,noheader,nounits"]
|
||||
try:
|
||||
result = subprocess.run(command, text=True, capture_output=True, check=True, timeout=5)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return []
|
||||
devices = []
|
||||
for line in result.stdout.splitlines():
|
||||
fields = [field.strip() for field in line.split(",", 3)]
|
||||
if len(fields) != 4:
|
||||
continue
|
||||
try:
|
||||
index, total, free = int(fields[0]), int(fields[2]), int(fields[3])
|
||||
except ValueError:
|
||||
continue
|
||||
devices.append({"index": index, "name": fields[1],
|
||||
"total_bytes": total * 1024 * 1024,
|
||||
"free_bytes": free * 1024 * 1024})
|
||||
return devices
|
||||
|
||||
|
||||
def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
|
||||
available_memory=None, available_disk=None, gpus=None):
|
||||
info = analyze_model(model)
|
||||
cfg = info["config"]
|
||||
available_memory = memory_available() if available_memory is None else available_memory
|
||||
if available_disk is None:
|
||||
fs = os.statvfs(info["path"])
|
||||
available_disk = fs.f_bavail * fs.f_frsize
|
||||
gpus = discover_gpus() if gpus is None else gpus
|
||||
if gpu_indices is not None:
|
||||
wanted = set(gpu_indices)
|
||||
gpus = [gpu for gpu in gpus if gpu["index"] in wanted]
|
||||
|
||||
ram_budget = int(ram_gb * GB) if ram_gb > 0 else int(available_memory * 0.88)
|
||||
if ram_budget < 4 * GB:
|
||||
ram_budget = 8 * GB
|
||||
typical = info["typical_expert_bytes"]
|
||||
layers = int(cfg.get("num_hidden_layers", 0)) + 1
|
||||
kv_bytes = layers * context * (int(cfg.get("kv_lora_rank", 0)) +
|
||||
int(cfg.get("qk_rope_head_dim", 0))) * 4
|
||||
kv_buffer = context * int(cfg.get("num_attention_heads", 0)) * (
|
||||
int(cfg.get("qk_nope_head_dim", 0)) + int(cfg.get("v_head_dim", 0))) * 4
|
||||
runtime_bytes = int(1.2 * GB + 2.5 * GB + 64 * typical + kv_bytes + kv_buffer)
|
||||
cache_bytes = max(0, ram_budget - info["dense_bytes"] - runtime_bytes)
|
||||
per_cap = info["per_cap_bytes"]
|
||||
configured_experts = int(cfg.get("n_routed_experts", 0))
|
||||
cap = int(cache_bytes // per_cap) if per_cap else 0
|
||||
if configured_experts:
|
||||
cap = min(cap, configured_experts)
|
||||
|
||||
reserve = 2 * GB
|
||||
gpu_plan = []
|
||||
safe_vram = 0
|
||||
for gpu in gpus:
|
||||
usable = max(0, gpu["free_bytes"] - reserve)
|
||||
safe_vram += usable
|
||||
gpu_plan.append(dict(gpu, reserve_bytes=reserve, usable_bytes=usable))
|
||||
requested_vram = int(vram_gb * GB) if vram_gb > 0 else safe_vram
|
||||
vram_budget = min(requested_vram, safe_vram, cache_bytes)
|
||||
vram_experts = int(vram_budget // typical) if typical else 0
|
||||
|
||||
warnings = []
|
||||
if cap < 1:
|
||||
warnings.append("RAM budget cannot hold one expert slot per sparse layer")
|
||||
if gpu_indices is not None and len(gpus) != len(set(gpu_indices)):
|
||||
warnings.append("one or more requested GPUs were not detected")
|
||||
if gpus and vram_budget < requested_vram:
|
||||
warnings.append("VRAM tier was clamped by free VRAM or its required RAM backing")
|
||||
|
||||
return {
|
||||
"version": 1,
|
||||
"model": {key: value for key, value in info.items() if key != "config"},
|
||||
"tiers": {
|
||||
"disk": {"role": "backing", "model_bytes": info["model_bytes"],
|
||||
"available_bytes": available_disk},
|
||||
"ram": {"role": "resident+cache", "available_bytes": available_memory,
|
||||
"budget_bytes": ram_budget, "dense_bytes": info["dense_bytes"],
|
||||
"runtime_bytes": runtime_bytes, "expert_cache_bytes": cache_bytes,
|
||||
"cache_slots_per_layer": cap},
|
||||
"vram": {"role": "hot-experts", "devices": gpu_plan,
|
||||
"budget_bytes": vram_budget, "expert_capacity": vram_experts},
|
||||
},
|
||||
"warnings": warnings,
|
||||
}
|
||||
|
||||
|
||||
def environment_for_plan(plan, env=None, cuda_enabled=True):
|
||||
"""Apply a plan without overriding explicit user environment settings."""
|
||||
result = dict(env or {})
|
||||
ram = plan["tiers"]["ram"]
|
||||
result.setdefault("RAM_GB", f"{ram['budget_bytes'] / GB:.3f}")
|
||||
|
||||
vram = plan["tiers"]["vram"]
|
||||
devices = [device["index"] for device in vram["devices"]]
|
||||
if not cuda_enabled or not devices or vram["budget_bytes"] <= 0:
|
||||
return result
|
||||
if result.get("COLI_CUDA", "1") == "0":
|
||||
return result
|
||||
|
||||
result.setdefault("COLI_CUDA", "1")
|
||||
if "COLI_GPU" not in result and "COLI_GPUS" not in result:
|
||||
key = "COLI_GPU" if len(devices) == 1 else "COLI_GPUS"
|
||||
result[key] = ",".join(map(str, devices))
|
||||
result.setdefault("CUDA_EXPERT_GB", f"{vram['budget_bytes'] / GB:.3f}")
|
||||
if result.get("PIN"):
|
||||
result.setdefault("PIN_GB", f"{vram['budget_bytes'] / GB:.3f}")
|
||||
return result
|
||||
|
||||
|
||||
def format_bytes(value):
|
||||
return f"{value / GB:.1f} GB"
|
||||
|
||||
|
||||
def format_plan(plan):
|
||||
model, tiers = plan["model"], plan["tiers"]
|
||||
lines = [f"model {model['shards']} shards · {format_bytes(model['model_bytes'])}",
|
||||
f"disk backing store · {format_bytes(tiers['disk']['available_bytes'])} free",
|
||||
f"RAM {format_bytes(tiers['ram']['budget_bytes'])} budget · "
|
||||
f"{format_bytes(tiers['ram']['dense_bytes'])} dense · "
|
||||
f"{format_bytes(tiers['ram']['runtime_bytes'])} runtime · "
|
||||
f"cap {tiers['ram']['cache_slots_per_layer']}/layer"]
|
||||
vram = tiers["vram"]
|
||||
if vram["devices"]:
|
||||
names = ", ".join(f"{gpu['index']}:{gpu['name']}" for gpu in vram["devices"])
|
||||
lines.append(f"VRAM {format_bytes(vram['budget_bytes'])} hot tier · "
|
||||
f"~{vram['expert_capacity']} experts · {names}")
|
||||
else:
|
||||
lines.append("VRAM no NVIDIA device detected · CPU path")
|
||||
lines.extend(f"warn {warning}" for warning in plan["warnings"])
|
||||
return "\n".join(lines)
|
||||
@@ -0,0 +1,111 @@
|
||||
import json
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from resource_plan import GB, analyze_model, build_plan, environment_for_plan, format_plan
|
||||
|
||||
|
||||
def write_shard(path, tensors):
|
||||
offset = 0
|
||||
header = {}
|
||||
payload = b""
|
||||
for name, size in tensors:
|
||||
header[name] = {"dtype": "U8", "shape": [size], "data_offsets": [offset, offset + size]}
|
||||
payload += b"\0" * size
|
||||
offset += size
|
||||
raw = json.dumps(header).encode()
|
||||
path.write_bytes(struct.pack("<Q", len(raw)) + raw + payload)
|
||||
|
||||
|
||||
class ResourcePlanTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.model = Path(self.tmp.name)
|
||||
(self.model / "config.json").write_text(json.dumps({
|
||||
"num_hidden_layers": 2,
|
||||
"n_routed_experts": 2,
|
||||
"kv_lora_rank": 4,
|
||||
"qk_rope_head_dim": 2,
|
||||
"qk_nope_head_dim": 3,
|
||||
"v_head_dim": 5,
|
||||
"num_attention_heads": 2,
|
||||
}))
|
||||
write_shard(self.model / "model.safetensors", [
|
||||
("model.embed_tokens.weight", 100),
|
||||
("model.layers.0.self_attn.q_a_proj.weight", 200),
|
||||
("model.layers.1.mlp.experts.0.gate_proj.weight", 30),
|
||||
("model.layers.1.mlp.experts.0.up_proj.weight", 30),
|
||||
("model.layers.1.mlp.experts.1.gate_proj.weight", 30),
|
||||
("model.layers.1.mlp.experts.1.up_proj.weight", 30),
|
||||
])
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
|
||||
def test_analyzes_dense_and_expert_storage(self):
|
||||
info = analyze_model(self.model)
|
||||
self.assertEqual(info["dense_bytes"], 300)
|
||||
self.assertEqual(info["expert_bytes"], 120)
|
||||
self.assertEqual(info["expert_count"], 2)
|
||||
self.assertEqual(info["per_cap_bytes"], 60)
|
||||
|
||||
def test_builds_bounded_three_tier_plan(self):
|
||||
gpus = [{"index": 0, "name": "test-gpu", "total_bytes": 12 * GB,
|
||||
"free_bytes": 10 * GB}]
|
||||
plan = build_plan(self.model, ram_gb=16, context=32, vram_gb=20,
|
||||
available_memory=32 * GB, available_disk=100 * GB, gpus=gpus)
|
||||
self.assertEqual(plan["version"], 1)
|
||||
self.assertEqual(plan["tiers"]["ram"]["budget_bytes"], 16 * GB)
|
||||
self.assertLessEqual(plan["tiers"]["vram"]["budget_bytes"], 8 * GB)
|
||||
self.assertIn("required RAM backing", plan["warnings"][0])
|
||||
self.assertIn("0:test-gpu", format_plan(plan))
|
||||
|
||||
def test_filters_requested_devices(self):
|
||||
gpus = [{"index": 0, "name": "a", "total_bytes": 8 * GB, "free_bytes": 8 * GB}]
|
||||
plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
|
||||
gpus=gpus, gpu_indices=[1])
|
||||
self.assertEqual(plan["tiers"]["vram"]["devices"], [])
|
||||
self.assertIn("not detected", plan["warnings"][0])
|
||||
|
||||
def test_cli_emits_versioned_json(self):
|
||||
cli = Path(__file__).parents[1] / "coli"
|
||||
run = subprocess.run([
|
||||
sys.executable, str(cli), "plan", "--model", str(self.model),
|
||||
"--gpu", "none", "--json",
|
||||
], text=True, capture_output=True, check=True)
|
||||
plan = json.loads(run.stdout)
|
||||
self.assertEqual(plan["version"], 1)
|
||||
self.assertEqual(plan["model"]["expert_count"], 2)
|
||||
|
||||
def test_applies_plan_without_overriding_explicit_settings(self):
|
||||
gpus = [
|
||||
{"index": 0, "name": "a", "total_bytes": 12 * GB, "free_bytes": 10 * GB},
|
||||
{"index": 1, "name": "b", "total_bytes": 12 * GB, "free_bytes": 10 * GB},
|
||||
]
|
||||
plan = build_plan(self.model, ram_gb=16, available_memory=32 * GB,
|
||||
available_disk=1, gpus=gpus)
|
||||
env = environment_for_plan(plan, {"RAM_GB": "12", "PIN": "stats.txt",
|
||||
"COLI_GPUS": "1"})
|
||||
self.assertEqual(env["RAM_GB"], "12")
|
||||
self.assertEqual(env["COLI_CUDA"], "1")
|
||||
self.assertEqual(env["COLI_GPUS"], "1")
|
||||
self.assertEqual(env["PIN_GB"], env["CUDA_EXPERT_GB"])
|
||||
|
||||
def test_cpu_binary_does_not_apply_gpu_tier(self):
|
||||
plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
|
||||
gpus=[{"index": 0, "name": "a", "total_bytes": 8 * GB,
|
||||
"free_bytes": 8 * GB}])
|
||||
env = environment_for_plan(plan, cuda_enabled=False)
|
||||
self.assertIn("RAM_GB", env)
|
||||
self.assertNotIn("COLI_CUDA", env)
|
||||
disabled = environment_for_plan(plan, {"COLI_CUDA": "0"}, cuda_enabled=True)
|
||||
self.assertNotIn("COLI_GPU", disabled)
|
||||
self.assertNotIn("CUDA_EXPERT_GB", disabled)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user