src/cism/cli.py

"""Command-line entry points. Imports optional HTTP dependencies only for serve."""

from __future__ import annotations

import argparse
import json
import statistics
import sys
import time


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(prog="cism")
    commands = parser.add_subparsers(dest="command", required=True)
    for name in ("inspect", "generate", "bench", "cache-report"):
        command = commands.add_parser(name)
        command.add_argument("model", help="HF repository ID or local safetensors directory")
        command.add_argument("--precision", choices=["fp32", "int8", "hybrid-int4", "hybrid-fp4"], default="fp32")
        command.add_argument("--act-precision", choices=["fp32", "int8"], default="fp32",
                             help="activation precision for quantized matrices (default fp32)")
        command.add_argument("--revision")
        command.add_argument("--trust-remote-code", action="store_true")
        command.add_argument("--local-files-only", action="store_true")
        command.add_argument("--lock-pages", action="store_true",
                             help="pin weight pages in physical RAM (VirtualLock/mlock)")
        command.add_argument("--threads", type=int, default=1, metavar="N",
                             help="native decode worker threads, including caller (1-64)")
        if name in ("generate", "bench"):
            command.add_argument("--prompt", default="The future of artificial intelligence is")
            command.add_argument("--max-tokens", type=int, default=64)
        if name == "generate":
            command.add_argument("--temperature", type=float, default=0)
            command.add_argument("--seed", type=int, default=0)
            command.add_argument("--spec", type=int, default=0, metavar="K",
                                  help="greedy prompt-lookup speculation (2-16, temperature must be 0)")
        if name == "bench":
            command.add_argument("--runs", type=int, default=5,
                                 help="measured repetitions (legacy name; --reps overrides when given)")
            command.add_argument("--reps", type=int, default=None,
                                 help="measured repetitions per setup; median reported (default: --runs value)")
            command.add_argument("--spec", type=int, default=0, metavar="K",
                                  help="greedy prompt-lookup speculation (2-16)")
        if name == "cache-report":
            command.add_argument("--decode-tokens", type=int, default=64)
            command.add_argument("--scrub-mb", type=int, default=96,
                                 help="cache-eviction buffer size in MiB (use ~2x L3)")
    serve = commands.add_parser("serve")
    serve.add_argument("models", nargs="+",
                       help="one or more HF repository IDs or local safetensors directories")
    serve.add_argument("--precision", choices=["fp32", "int8", "hybrid-int4", "hybrid-fp4"], default="fp32")
    serve.add_argument("--act-precision", choices=["fp32", "int8"], default="fp32",
                       help="activation precision for quantized matrices (default fp32)")
    serve.add_argument("--revision")
    serve.add_argument("--trust-remote-code", action="store_true")
    serve.add_argument("--local-files-only", action="store_true")
    serve.add_argument("--lock-pages", action="store_true",
                       help="pin weight pages in physical RAM (VirtualLock/mlock)")
    serve.add_argument("--threads", type=int, default=1, metavar="N",
                       help="native decode worker threads per model, including caller (1-64)")
    serve.add_argument("--keep-warm", type=float, default=0.0, metavar="SECONDS",
                       help="touch weights of every loaded model every SECONDS while idle")
    serve.add_argument("--allow-remote-code-load", action="store_true",
                       help="permit POST /v1/models/load with trust_remote_code=true")
    serve.add_argument("--host", default="127.0.0.1")
    serve.add_argument("--port", type=int, default=8000)
    serve.add_argument("--queue-size", type=int, default=8)
    tui = commands.add_parser("tui")
    tui.add_argument("model", help="HF repository ID or local safetensors directory")
    tui.add_argument("--precision", choices=["fp32", "int8", "hybrid-int4", "hybrid-fp4"], default="int8",
                     help="weight precision (default int8)")
    tui.add_argument("--threads", type=int, default=2, metavar="N",
                     help="native decode worker threads, including caller (1-64)")
    tui.add_argument("--max-tokens", type=int, default=128, metavar="N",
                     help="max new tokens per generation (>=1)")
    tui.add_argument("--temperature", type=float, default=0.0,
                     help="sampling temperature, finite >=0 (0 for greedy)")
    tui.add_argument("--top-p", type=float, default=1.0,
                     help="nucleus sampling mass in (0, 1]")
    tui.add_argument("--top-k", type=int, default=0, metavar="K",
                     help="top-k cutoff >=0 (0 disables)")
    tui.add_argument("--seed", type=int, default=0, metavar="S",
                     help="sampling seed in [0, 2**64)")
    tui.add_argument("--spec-k", type=int, default=0, metavar="K",
                     help="greedy prompt-lookup speculation (0 off or 2-16, requires temperature 0)")
    tui.add_argument("--trust-remote-code", action="store_true",
                     help="allow remote code when loading the model")
    tui.add_argument("--local-files-only", action="store_true",
                     help="load only from local cache, never download")
    args = parser.parse_args(argv)
    if args.command == "bench":
        _reps = args.reps if args.reps is not None else args.runs
        if _reps < 1 or args.runs < 1 or args.max_tokens < 2:
            parser.error("bench requires --runs/--reps >= 1 and --max-tokens >= 2")
    if args.command == "tui":
        import math as _math
        if not 1 <= args.threads <= 64:
            parser.error("tui requires --threads in [1, 64]")
        if args.max_tokens < 1:
            parser.error("tui requires --max-tokens >= 1")
        if not _math.isfinite(args.temperature) or args.temperature < 0:
            parser.error("tui requires --temperature finite and >= 0")
        if not _math.isfinite(args.top_p) or not 0 < args.top_p <= 1:
            parser.error("tui requires --top-p in (0, 1]")
        if args.top_k < 0:
            parser.error("tui requires --top-k >= 0")
        if not 0 <= args.seed < 2**64:
            parser.error("tui requires --seed in [0, 2**64)")
        if args.spec_k != 0 and not 2 <= args.spec_k <= 16:
            parser.error("tui requires --spec-k 0 or in [2, 16]")
        if args.spec_k >= 2 and args.temperature != 0:
            parser.error("tui requires --temperature 0 when --spec-k is used")
        from .tui import run as _tui_run

        class _TeeStdout:
            """Mirror tui 'error: ...' lines to stderr for CLI conventions."""

            def __init__(self, out, err):
                self._out = out
                self._err = err

            def write(self, s):
                try:
                    res = self._out.write(s)
                except Exception:
                    res = None
                try:
                    if isinstance(s, str) and s.startswith("error:"):
                        self._err.write(s)
                        try:
                            self._err.flush()
                        except Exception:
                            pass
                except Exception:
                    pass
                if res is not None:
                    return res
                return len(s) if isinstance(s, str) else 0

            def flush(self):
                try:
                    self._out.flush()
                except Exception:
                    pass

            def __getattr__(self, name):
                return getattr(self._out, name)

        _real_out, _real_err = sys.stdout, sys.stderr
        sys.stdout = _TeeStdout(_real_out, _real_err)
        try:
            try:
                return _tui_run(
                    args.model,
                    precision=args.precision,
                    threads=args.threads,
                    max_new_tokens=args.max_tokens,
                    temperature=args.temperature,
                    top_p=args.top_p,
                    top_k=args.top_k,
                    seed=args.seed,
                    spec_k=args.spec_k,
                    trust_remote_code=args.trust_remote_code,
                    local_files_only=args.local_files_only,
                )
            except (ValueError, RuntimeError, ImportError, OSError) as exc:
                print(f"error: {exc}", file=_real_err)
                return 1
        finally:
            sys.stdout = _real_out
    from .engine import Engine

    try:
        engines = {}
        for model in (args.models if args.command == "serve" else [args.model]):
            engine = Engine.from_pretrained(model, precision=args.precision, revision=args.revision,
                                            trust_remote_code=args.trust_remote_code,
                                            local_files_only=args.local_files_only,
                                            threads=args.threads,
                                            act_precision=args.act_precision)
            if args.lock_pages and engine.lock_pages() == 0:
                print("cism: page locking failed for this process; continuing unlocked", file=sys.stderr)
            engines[engine.model_id] = engine
        if args.command == "inspect":
            print(json.dumps(next(iter(engines.values())).info, indent=2))
        elif args.command == "generate":
            engine = next(iter(engines.values()))
            for chunk in engine.stream(args.prompt, max_new_tokens=args.max_tokens,
                                       temperature=args.temperature, seed=args.seed,
                                       spec_k=args.spec):
                print(chunk.text, end="", flush=True)
            print()
        elif args.command == "cache-report":
            import numpy as np
            engine = next(iter(engines.values()))
            weight_bytes = engine.info["weight_bytes"]

            def timed_scan(repeats: int, evictions: int) -> float:
                # Unique 2x-L3 writes evict residency; the scan is timed alone.
                scrub = np.empty(args.scrub_mb * 1024 * 1024, dtype=np.uint8)
                samples = []
                for _ in range(repeats):
                    for _ in range(evictions):
                        scrub[:] = 1
                    start = time.perf_counter()
                    engine.scan()
                    samples.append(time.perf_counter() - start)
                return statistics.median(samples)

            warm_seconds = timed_scan(20, evictions=0)
            cold_seconds = timed_scan(5, evictions=2)
            prompt = engine.encode("cache report probe")
            session = engine._native.create_session(prompt, max_new_tokens=args.decode_tokens,
                                                     eos_token_ids=[])
            session.next_tokens(1)
            start = time.perf_counter()
            generated = session.next_tokens(args.decode_tokens - 1)
            decode_seconds = time.perf_counter() - start
            tokens_per_second = len(generated) / decode_seconds if generated and decode_seconds > 0 else 0
            effective_gbps = weight_bytes * tokens_per_second / 1e9
            warm_gbps = weight_bytes / warm_seconds / 1e9
            cold_gbps = weight_bytes / cold_seconds / 1e9
            ratio = cold_seconds / warm_seconds if warm_seconds > 0 else 0
            if ratio >= 2:
                residency = "cache-resident when warm (cold scan is much slower)"
            elif ratio >= 1.3:
                residency = "partially cache-resident"
            else:
                residency = "no measurable residency benefit from this probe"
            if not tokens_per_second:
                bottleneck = "decode produced no tokens"
            elif effective_gbps <= cold_gbps:
                bottleneck = ("compute-bound: decode bandwidth is below even the cold DRAM scan; "
                              "the cache is not the bottleneck, the kernels are")
            elif effective_gbps < warm_gbps * 0.8:
                bottleneck = "DRAM-bandwidth-limited: weights are effectively streaming from RAM"
            else:
                bottleneck = "bandwidth-limited at cache-resident rates"
            print(json.dumps({**engine.info, "weight_bytes": weight_bytes,
                              "warm_scan_gbps": round(warm_gbps, 2),
                              "cold_scan_gbps": round(cold_gbps, 2),
                              "warm_to_cold_ratio": round(ratio, 2),
                              "decode_tokens_per_second": round(tokens_per_second, 1),
                              "decode_effective_gbps": round(effective_gbps, 2),
                              "residency": residency, "decode_bottleneck": bottleneck}, indent=2))
        elif args.command == "bench":
            from .autobench import _cpu_snapshot, _thermal_status
            engine = next(iter(engines.values()))
            prompt = engine.encode(args.prompt)
            reps = args.reps if args.reps is not None else args.runs
            # Deliberately ignore EOS for fixed-length comparison; report this.
            ttfts, decode_rates = [], []
            for run in range(reps + 1):
                start = time.perf_counter()
                session = engine._native.create_session(prompt, max_new_tokens=args.max_tokens,
                                                         eos_token_ids=[])
                first = session.next_tokens(1)
                after_first = time.perf_counter()
                tail = session.next_tokens(args.max_tokens - 1)
                end = time.perf_counter()
                if len(first) + len(tail) != args.max_tokens:
                    raise RuntimeError("Benchmark failed to generate the requested token count")
                if run:
                    ttfts.append((after_first - start) * 1000)
                    decode_rates.append(len(tail) / (end - after_first))
            throttled, drop_percent = _thermal_status(decode_rates)
            try:
                cpu_snapshot = _cpu_snapshot()
            except Exception:
                cpu_snapshot = {}
            if not isinstance(cpu_snapshot, dict):
                cpu_snapshot = {}
            print(json.dumps({**engine.info, "prompt_tokens": len(prompt),
                              "generated_tokens": args.max_tokens, "runs": reps, "reps": reps,
                              "eos_ignored": True, "warmup_runs": 1,
                              "median_prefill_first_token_ms": statistics.median(ttfts),
                              "median_decode_tokens_per_second": statistics.median(decode_rates),
                              "decode_tokens_per_second": decode_rates,
                              "decode_tokens_per_second_min": min(decode_rates),
                              "decode_tokens_per_second_max": max(decode_rates),
                              "decode_tokens_per_second_std": round(statistics.pstdev(decode_rates), 3)
                              if len(decode_rates) >= 2 else 0.0,
                              "prefill_first_token_ms_min": min(ttfts),
                              "prefill_first_token_ms_max": max(ttfts),
                              "thermal_throttled": throttled,
                              "thermal_drop_percent": drop_percent,
                              "cpu_snapshot": cpu_snapshot}, indent=2))
        else:
            import uvicorn
            from .server import create_app
            if args.keep_warm:
                for engine in engines.values():
                    engine.start_keep_warm(args.keep_warm)
            loaded = engines if len(engines) > 1 else next(iter(engines.values()))
            uvicorn.run(create_app(loaded, queue_size=args.queue_size,
                                   allow_remote_code_load=args.allow_remote_code_load,
                                   load_lock_pages=args.lock_pages,
                                   load_keep_warm=args.keep_warm,
                                   load_threads_default=args.threads,
                                   load_local_files_only=args.local_files_only,
                                   load_act_precision=args.act_precision),
                        host=args.host, port=args.port, workers=1)
    except (ValueError, RuntimeError, ImportError, OSError) as exc:
        print(f"cism: {exc}", file=sys.stderr)
        return 1
    return 0