| """Measure standalone packed loading and a short System One request in a fresh process. |
| |
| Use an empty HF_HOME and HF_HUB_OFFLINE=1. The base loaders are explicitly |
| blocked. RSS alone misses Metal buffers, so also sample physical footprint. |
| """ |
| import argparse |
| import json |
| import platform |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| from kev.api import SystemOneRequest, to_record |
| from kev.checkpoint import Checkpoint, LoadOptions |
| from kev.model import SERVE_MAX_BRANCH, SERVE_MAX_STATE |
| from scripts.mlx_long_states import RSSPeak, footprint |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--run", required=True) |
| parser.add_argument("--out", required=True) |
| parser.add_argument("--mlx-limit-gb", type=float, default=3, help="MLX memory guideline, not a hard RAM cap") |
| parser.add_argument("--filler-repeats", type=int, default=0, help="append neutral filler to measure longer states; this is not an accuracy benchmark") |
| args = parser.parse_args() |
| import kev.mlx_model |
| import huggingface_hub |
|
|
| def forbidden(*unused, **kwargs): |
| raise AssertionError("standalone profiling must never call a base loader or download") |
|
|
| kev.mlx_model.load_base = kev.mlx_model.load_full = kev.mlx_model.merge_lora = forbidden |
| huggingface_hub.snapshot_download = forbidden |
| mx.set_memory_limit(int(args.mlx_limit_gb * 1e9)) |
| mx.reset_peak_memory() |
| start = time.perf_counter() |
| with RSSPeak() as load_peak: |
| ck = Checkpoint(args.run) |
| tok, model = ck.load("mps", LoadOptions()) |
| result = {"run": str(Path(args.run).resolve()), "backend": model.backend, "dtype": model.dtype, |
| "platform": platform.platform(), "original_loaders_blocked": True, |
| "mlx_limit_bytes": int(args.mlx_limit_gb * 1e9), "mlx_limit_is_hard_cap": False, "load_seconds": time.perf_counter() - start, |
| "mlx_load_peak_bytes": mx.get_peak_memory(), "mlx_after_load_bytes": mx.get_active_memory(), |
| "load_footprint_peak_bytes": max(load_peak.footprint, footprint()), "load_rss_peak_bytes": load_peak.peak} |
| request = SystemOneRequest(state="I was charged twice for order 8812. Please refund the duplicate charge." + " This is a routine account note." * args.filler_repeats, |
| questions={"team": {"type": "choice", "instructions": "Which team should handle this ticket?", |
| "criteria": {"billing": "Charges and payments", "shipping": "Delivery and delays"}}}) |
| record, _ = to_record(request) |
| enc = model.encode(tok, record, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH, strict=True) |
| start = time.perf_counter() |
| with RSSPeak() as infer_peak: |
| probs = model.probs(enc) |
| result.update(state_tokens=enc["seg"].count(0), input_tokens=len(enc["ids"]), inference_seconds=time.perf_counter() - start, |
| probabilities={"billing": float(probs[0][0]), "shipping": float(probs[0][1])}, |
| mlx_inference_peak_bytes=mx.get_peak_memory(), inference_footprint_peak_bytes=max(infer_peak.footprint, footprint()), |
| inference_rss_peak_bytes=infer_peak.peak) |
| output = Path(args.out) |
| output.parent.mkdir(parents=True, exist_ok=True) |
| output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") |
| print(json.dumps(result, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
