inferencekit-cli 0.1.0__py3-none-any.whl

inferencekit/mcp_server.py

@@ -0,0 +1,178 @@
+"""``inferencekit mcp serve``: every registered operation as an MCP tool, over stdio.
+
+Built on the low-level server of the official ``mcp`` Python SDK (v2), because the
+tool list and input schemas already exist in the registry; no per-tool code here.
+
+Tool results are trimmed to stay well under ~10k tokens: big lists inside a record
+are cut to their first items and replaced by a pointer that says how to page through
+the rest with the ``inspect_evidence`` tool (``path``, ``offset``, ``limit``). A result
+that is already a page (``list_evidence``, or ``inspect_evidence`` with a list path)
+is shortened to fewer items instead, with ``next_offset`` moved back to the first item
+that was cut, so following it skips nothing. The stored record is never modified; only
+the view sent to the model is shortened.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+from functools import partial
+
+from inferencekit import __version__, registry
+from inferencekit.registry import OpError
+
+# ~4 characters per token for JSON: 32k characters is about 8k tokens.
+MAX_CHARS = 32_000
+LIST_PREVIEW = 10
+TOOLS_PAGE = 50
+
+INSTRUCTIONS = (
+    "inferencekit measures inference performance and writes evidence records. Start with doctor; "
+    "use roofline (with gpu_bw from measure_bandwidth) for the ceiling. For an optimisation run: "
+    "intake_template / intake_validate (start the run dir), validate_manifest (the run manifest), "
+    "bench_rounds (every manifest engine, interleaved rounds) or bench_serve (one OpenAI-compatible server: "
+    "single-user TTFT/TPOT/decode tok/s, or load mode for N* and $/1M tokens) for published speed numbers; "
+    "bench_decode only for engines without a server. run_gates for the anti-cheat checks; quality_trajectory "
+    "+ quality_score (agreement with a reference), quality_gsm8k / quality_niah for task accuracy; "
+    "prompts_show / prompts_verify for the prompt sets. analyze_nsys / summarize_ncu to explain speed; "
+    "bench_image / bench_audio for client-timed media; compare with image_tradeoff=true only for a "
+    "labelled image step/precision tradeoff, with separate quality proof. AMD: capture_rocprof / "
+    "analyze_rocprof for supported ROCm kernel traces. "
+    "compare for A/B verdicts; report to write it up; "
+    "build_review for the run's review page; ledger_trailers / history for Perf-* commit trailers. "
+    "Pass chat_template_kwargs (e.g. {\"enable_thinking\": false}) identically to bench_serve, bench_rounds, "
+    "run_gates, quality_trajectory, quality_gsm8k and quality_niah (an intake pin in the run dir is applied and "
+    "enforced when run_dir is given). Long lists in results are truncated: page them with "
+    "inspect_evidence(id, path, offset, limit)."
+)
+
+
+def _size(obj) -> int:
+    return len(json.dumps(obj, ensure_ascii=False, default=str))
+
+
+def _lists(node, path=""):
+    """Yield (path, parent, key, list) for every list inside node, largest first by caller."""
+    if isinstance(node, dict):
+        for k, v in node.items():
+            p = f"{path}.{k}" if path else k
+            if isinstance(v, list):
+                yield p, node, k, v
+            yield from _lists(v, p)
+    elif isinstance(node, list):
+        for i, v in enumerate(node):
+            yield from _lists(v, f"{path}.{i}")
+
+
+def _trim_strings(node, limit: int) -> None:
+    """Cut every string longer than ``limit`` characters in place, saying how much was cut."""
+    items = node.items() if isinstance(node, dict) else enumerate(node) if isinstance(node, list) else []
+    for k, v in list(items):
+        if isinstance(v, str) and len(v) > limit:
+            node[k] = v[:limit] + f"\n... [truncated {len(v) - limit} characters]"
+        elif isinstance(v, (dict, list)):
+            _trim_strings(v, limit)
+
+
+def _is_page(view: dict) -> bool:
+    """A paged result from store.page(): {..., total, offset, items, next_offset}."""
+    return isinstance(view.get("items"), list) and "offset" in view and "next_offset" in view
+
+
+def _fit_page(view: dict, max_chars: int) -> dict:
+    """Keep the first k items of a page that fit, and point next_offset / the hint at item offset+k."""
+    items = view["items"]
+    k = len(items)
+    while k > 1 and _size({**view, "items": items[:k]}) > max_chars:
+        k = max(1, k * 3 // 4)
+    if k == len(items):
+        return view
+    nxt = view["offset"] + k
+    view["items"], view["next_offset"] = items[:k], nxt
+    if "path" in view:  # inspect_evidence: repeat the same path, from the first item not shown
+        view["more"] = (f"page shortened to {k} items: inspect_evidence(id='{view.get('evidence_id')}', "
+                        f"path='{view['path']}', offset={nxt}, limit={k})")
+    else:  # list_evidence
+        view["more"] = f"page shortened to {k} items: call again with offset={nxt}, limit={k}"
+    return view
+
+
+def fit(result: dict, max_chars: int = MAX_CHARS) -> dict:
+    """Shrink a result for a model's context window, with paging hints for what was cut."""
+    if _size(result) <= max_chars:
+        return result
+    view = copy.deepcopy(result)
+    if _is_page(view):
+        view = _fit_page(view, max_chars)
+        if _size(view) <= max_chars:
+            return view
+    record = view.get("record", view)  # inspect_evidence wraps the record
+    eid = record.get("evidence_id") if isinstance(record, dict) else None
+    prefix = "record." if "record" in view else ""
+    candidates = sorted(_lists(view), key=lambda t: -_size(t[3]))
+    for path, parent, key, items in candidates:
+        if _size(view) <= max_chars:
+            break
+        if len(items) <= LIST_PREVIEW:
+            continue
+        inner = path[len(prefix):] if path.startswith(prefix) else path
+        if _is_page(view) and "path" in view and inner.startswith("items."):
+            # A list inside one item of a page: "items.0.x" is item offset+0 of the paged path.
+            i, _, rest = inner[len("items."):].partition(".")
+            inner = f"{view['path']}.{view['offset'] + int(i)}" + (f".{rest}" if rest else "")
+        parent[key] = {"truncated": True, "total": len(items), "first": items[:LIST_PREVIEW],
+                       "more": f"inspect_evidence(id='{eid}', path='{inner}', offset={LIST_PREVIEW}, limit=50)"
+                       if eid else "call again with offset/limit"}
+    if _size(view) > max_chars:  # long text, e.g. a report's markdown (the result also gives its file path)
+        _trim_strings(view, max_chars // 2)
+    if _size(view) > max_chars and isinstance(record, dict) and "data" in record:
+        record["data"] = {"truncated": True,
+                          "more": f"inspect_evidence(id='{eid}', path='data.<key>') for one part at a time",
+                          "keys": sorted(result.get("record", result).get("data", {}).keys())}
+    if _size(view) > max_chars:  # last resort: a summary that always fits
+        view = {"truncated": True, "evidence_id": eid, "summary": record.get("summary") if isinstance(record, dict)
+                else None, "more": "use inspect_evidence with a path"}
+    return view
+
+
+def build_server():
+    try:
+        import anyio
+        import mcp.types as types
+        from mcp.server.lowlevel import Server
+    except ImportError as e:  # pragma: no cover - mcp is a declared dependency
+        raise SystemExit(f"inferencekit mcp serve needs the MCP SDK (pip install 'mcp>=2.3,<3'): {e}") from None
+
+    ops = registry.load_all()
+    tools = [types.Tool(name=op.name, description=f"{op.summary} (CLI: inferencekit {' '.join(op.cli)})",
+                        input_schema=op.input_schema()) for op in ops.values()]
+
+    async def list_tools(ctx, params):
+        start = int(params.cursor) if params is not None and params.cursor else 0
+        page = tools[start:start + TOOLS_PAGE]
+        nxt = start + TOOLS_PAGE
+        return types.ListToolsResult(tools=page, next_cursor=str(nxt) if nxt < len(tools) else None)
+
+    async def call_tool(ctx, params):
+        try:
+            # Operations block on subprocesses for minutes; keep the event loop free.
+            result = await anyio.to_thread.run_sync(partial(registry.call, params.name, params.arguments or {}))
+        except OpError as e:
+            return types.CallToolResult(content=[types.TextContent(text=f"error: {e}")], is_error=True)
+        view = fit(result)
+        return types.CallToolResult(content=[types.TextContent(text=json.dumps(view, ensure_ascii=False))])
+
+    server = Server("inferencekit", version=__version__, instructions=INSTRUCTIONS, on_list_tools=list_tools,
+                    on_call_tool=call_tool)
+    return server, anyio
+
+
+def serve() -> None:
+    server, anyio = build_server()
+    from mcp.server.stdio import stdio_server
+
+    async def main():
+        async with stdio_server() as (read_stream, write_stream):
+            await server.run(read_stream, write_stream, server.create_initialization_options())
+
+    anyio.run(main)

inferencekit/ops/__init__.py

@@ -0,0 +1,30 @@
+"""Every operation module registers itself on import. Order = order in ``inferencekit --help``."""
+
+from inferencekit.ops import (  # noqa: F401
+    doctor,
+    roofline,
+    ceiling,
+    bench,
+    nsys,
+    ncu,
+    rocprof,
+    compare,
+    correctness,
+    store,
+    report,
+    review,
+    ledger,
+    manifest,
+)
+from inferencekit.ops import intake  # noqa: F401,E402
+from inferencekit.ops import quality  # noqa: F401,E402
+from inferencekit.ops import bench_serve  # noqa: E402,F401
+from inferencekit.ops import bench_rounds  # noqa: F401,E402
+from inferencekit.ops import prompts  # noqa: F401,E402
+from inferencekit.ops import gates  # noqa: F401,E402
+from inferencekit.ops import bench_clip  # noqa: F401,E402  (video: bench clip)
+from inferencekit.ops import quality_clip  # noqa: F401,E402  (video: quality clip)
+from inferencekit.ops import quality_clip_steps  # noqa: F401,E402  (video: quality clip-steps)
+from inferencekit.ops import gates_clip  # noqa: F401,E402  (video: gates clip; needs bench_clip)
+from inferencekit.ops import bench_image  # noqa: F401,E402
+from inferencekit.ops import bench_audio  # noqa: F401,E402

inferencekit/ops/bench.py

@@ -0,0 +1,184 @@
+"""bench decode: batch-1 decode speed, TTFT and per-token latency, with the engine's real flags.
+
+Definitions (per round):
+  ttft_ms       wall time of a 1-token generation: prefill + first sample (offline, no network)
+  decode_tok_s  (tokens - 1) / (t_total - ttft), t_total = wall time of a gen_len generation
+  tpot_ms       (t_total - ttft) / (tokens - 1): time per output token after the first
+
+Reported value = median over rounds, with a 95% bootstrap CI over rounds. Rounds are
+the unit of repetition because tokens inside one generation are not independent.
+
+Interleaving: to compare A and B fairly, alternate them (A, B, A, B ...) so slow
+drift (clocks, temperature, neighbours) hits both equally. ``--extend <ev_id>`` runs
+more rounds of the same configuration and writes a new record holding the old and
+new rounds, each with wall-clock timestamps; ``compare`` checks the alternation.
+
+Environment: the engine runs with the caller's full environment (unlike profiler
+captures, which use an allowlist). The record lists the *names* of the engine-
+configuration variables that were set (``data.env_config_names``), so a profile taken
+with a different configuration can be spotted. Secret-looking keys inside
+``--engine-args`` (``hf_token``, ``api_key``...) are stored as <redacted>.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+
+from inferencekit import evidence as ev
+from inferencekit import stats
+from inferencekit.engines import adapters, gates
+from inferencekit.gpuinfo import gpu_identity
+from inferencekit.ops.doctor import gpu_processes
+from inferencekit.proc import config_env_names
+from inferencekit.registry import Arg, OpError, operation
+
+ENGINES = ("vllm", "sglang", "llamacpp", "cmd")
+# Arguments that must match for --extend: they define "the same configuration".
+SAME_CONFIG = ("engine", "model", "prompt_len", "gen_len", "seed", "engine_args", "cmd")
+
+
+def _engine_args(text: str) -> dict:
+    try:
+        value = json.loads(text or "{}")
+    except ValueError as e:
+        raise OpError(f"engine_args is not valid JSON: {e}") from None
+    if not isinstance(value, dict):
+        raise OpError("engine_args must be a JSON object, e.g. '{\"max_model_len\": 4096}'")
+    return value
+
+
+def _ci_metric(name: str, unit: str, values: list[float]) -> dict:
+    if not values:
+        return ev.metric(name, None, unit, "median", n=0)
+    ci = stats.bootstrap_ci(values, ps=(50,))[50]
+    return ev.metric(name, stats.summarize(values)["median"], unit, "median", n=len(values), ci95=ci)
+
+
+@operation("bench_decode", ("bench", "decode"), "Batch-1 decode tok/s, per-token latency and TTFT for vllm, "
+           "sglang, llama.cpp or any command; records effective engine flags and hard-gate violations.",
+           [Arg("engine", "string", "Engine adapter.", required=True, choices=ENGINES),
+            Arg("model", "string", "HF id or local path (a .gguf file for llamacpp). Optional for cmd."),
+            Arg("prompt_len", "integer", "Prompt length in tokens.", default=128),
+            Arg("gen_len", "integer", "Generated tokens per round (>= 2).", default=128),
+            Arg("rounds", "integer", "Timed rounds.", default=5),
+            Arg("warmup", "integer", "Untimed warm-up generations (the first also records reference tokens).",
+                default=1),
+            Arg("seed", "integer", "Seed for synthetic prompts and the engine.", default=0),
+            Arg("engine_args", "string", "JSON object passed to the engine (vLLM LLM(...), SGLang Engine(...), "
+                "or llama-bench flags such as {\"ngl\": 99, \"fa\": 1}). Values of secret-looking keys "
+                "(hf_token, api_key...) are stored as <redacted>.", default="{}"),
+            Arg("python", "string", "Interpreter that has the engine installed (vllm/sglang). Default: this one."),
+            Arg("cmd", "string", "For engine=cmd: command printing a JSON line {\"decode_tok_s\": ..., "
+                "\"ttft_ms\": ...}. Stored only as a digest.", sensitive=True),
+            Arg("extend", "string", "Evidence id of an earlier run of the same configuration; its rounds are kept "
+                "and new rounds appended (use to interleave A/B rounds)."),
+            Arg("record_tokens", "boolean", "Store the reference generation's token ids (for correctness).",
+                default=True),
+            Arg("timeout", "integer", "Seconds before the engine run (and every process it started) is killed.",
+                default=3600)])
+def bench_decode(call, engine, model, prompt_len, gen_len, rounds, warmup, seed, engine_args, python, cmd, extend,
+                 record_tokens, timeout):
+    if gen_len < 2:
+        raise OpError("gen_len must be >= 2 (decode speed needs tokens after the first)")
+    if rounds < 1:
+        raise OpError("rounds must be >= 1")
+    if engine != "cmd" and not model:
+        raise OpError(f"engine={engine} needs --model")
+    if engine == "cmd" and not cmd:
+        raise OpError("engine=cmd needs --cmd")
+    eargs = _engine_args(engine_args)
+    store = ev.Store()
+
+    prev, prev_rounds = None, []
+    if extend:
+        prev = store.load(extend)
+        if prev.get("kind") != "benchmark":
+            raise OpError(f"--extend {extend}: not a benchmark record")
+        problem = ev.integrity_problem(prev)
+        if problem:
+            raise OpError(f"--extend: {problem}; refusing to build on it")
+        prev_args = prev["data"].get("request", {}).get("args", {})
+        mine = ev.redact(call.request()["args"])  # stored args went through the same redaction
+        diff = [k for k in SAME_CONFIG if prev_args.get(k) != mine.get(k)]
+        if diff:
+            raise OpError(f"--extend needs the same configuration; these differ: {diff}")
+        prev_rounds = prev["data"].get("rounds", [])
+    offset = len(prev_rounds)
+
+    # Leftovers from an earlier run (an orphaned engine, someone else's job) steal bandwidth.
+    others = gpu_processes()
+    with store.staging_dir() as staging:  # removed again if the run fails
+        cfg = {"model": model, "prompt_len": prompt_len, "gen_len": gen_len, "rounds": rounds, "warmup": warmup,
+               "seed": seed, "engine_args": eargs, "round_offset": offset}
+        if engine in ("vllm", "sglang"):
+            result = adapters.run_python_worker(engine, python or sys.executable, cfg, staging, timeout)
+        elif engine == "llamacpp":
+            result = adapters.run_llamacpp(model, prompt_len, gen_len, rounds, eargs, staging, timeout)
+            for r in result["rounds"]:
+                r["index"] += offset
+        else:
+            result = adapters.run_cmd(cmd, rounds, offset, timeout)
+
+        flags = result.get("flags") or {}
+        backend = gates.attention_backend_from_log(result.get("log_text", "")) if engine == "vllm" else None
+        gate = gates.check(engine, flags, backend)
+        limitations = list(result.get("limitations", []))
+        if others:
+            limitations.append(f"{len(others)} other compute process(es) were on the GPU when the run started ("
+                               + ", ".join(f"{o['name']} {o['used_mib']} MiB" for o in others[:4])
+                               + "): timings may be contended")
+        limitations += [f"HARD GATE: {v}" for v in gate["violations"]]
+        limitations += gate["warnings"]
+        if gate["valid"] is None:
+            limitations.append("engine flags not visible to inferencekit: CUDA graphs / attention backend not checked")
+        if prev is not None:
+            old_flags = prev["identity"]["engine"].get("flags", {})
+            if old_flags != ev.redact(flags):
+                limitations.append(f"effective engine flags changed since {prev['evidence_id'][:15]}")
+
+        all_rounds = prev_rounds + result["rounds"]
+        decode = [r["decode_tok_s"] for r in all_rounds if r.get("decode_tok_s") is not None]
+        ttft = [r["ttft_ms"] for r in all_rounds if r.get("ttft_ms") is not None]
+        tpot = [r["tpot_ms"] for r in all_rounds if r.get("tpot_ms") is not None]
+        missing = []
+        if not ttft:
+            missing.append("ttft_ms")
+        if len(decode) < 2:
+            limitations.append("fewer than 2 rounds: no confidence interval")
+        metrics = [_ci_metric("decode_tok_s", "tok/s", decode), _ci_metric("ttft_ms", "ms", ttft),
+                   _ci_metric("tpot_ms", "ms", tpot)]
+
+        eng = result.get("engine") or {}
+        ident = ev.identity(
+            gpu=gpu_identity(),
+            engine={"name": eng.get("name") or engine, "version": eng.get("version"), "commit": eng.get("commit"),
+                    "flags": flags},
+            model={"id": model, "revision": eargs.get("revision"),
+                   "dtype": str(flags["dtype"]) if flags.get("dtype") is not None else None,
+                   "quant": str(flags.get("quantization")) if flags.get("quantization") else None},
+            workload={"prompt_len": prompt_len, "gen_len": gen_len, "batch": 1, "seed": seed})
+        med = metrics[0]["value"]
+        lo, hi = metrics[0]["ci95"] or (None, None)
+        summary = f"{ident['engine']['name']} {model or ''}".strip() + ": decode "
+        summary += f"{med:.1f} tok/s" if med is not None else "n/a"
+        if lo is not None and hi is not None:
+            summary += f" (95% CI {lo:.1f}-{hi:.1f}, n={len(decode)})"
+        if metrics[1]["value"] is not None:
+            summary += f", TTFT {metrics[1]['value']:.1f} ms"
+        if gate["violations"]:
+            summary += " [HARD GATE: " + "; ".join(gate["violations"]) + "]"
+        tokens = result.get("tokens") if record_tokens else None
+        if tokens is None and prev is not None:
+            tokens = prev["data"].get("tokens")
+        record = ev.build(
+            call, kind="benchmark", claim_class="observed", summary=summary, identity_=ident, metrics=metrics,
+            artifacts=[ev.artifact(p, staging, role) for p, role in result.get("files", []) if p.is_file()],
+            missing=missing, limitations=limitations,
+            data={"rounds": all_rounds, "tokens": tokens, "prompt": result.get("prompt"),
+                  "claim_validity": gate, "extends": prev["evidence_id"] if prev else None,
+                  "env_config_names": config_env_names(), "gpu_processes_at_start": len(others),
+                  "definitions": {"ttft_ms": "wall time of a 1-token generation (prefill + first sample)",
+                                  "decode_tok_s": "(tokens - 1) / (t_total - ttft) per round",
+                                  "tpot_ms": "(t_total - ttft) / (tokens - 1) per round"}})
+        return store.save(record, staging)