#!/usr/bin/env python3
"""Strict-limit RAM sweep for EPQ on LFM2.5-24B-A2B Q4_K_M.

Unlike the original memory_levels_ling.py (which only sets the pager's own
budget), every run here is launched under a *hard* OS limit:

    prlimit --data=<level> --nofile=65536

On Linux >= 4.7 RLIMIT_DATA covers brk() AND mmap'd anonymous memory, so it is
a real ceiling on the arena + heap + every later anonymous allocation. If the
process exceeds it, malloc fails and the run aborts -- we record that as a hard
failure instead of a silent overshoot.

Two caveats this harness records explicitly:

  * RLIMIT_DATA polices *anonymous* memory only. The ~1.3-2.1 GiB of
    non-expert weights are mmap'd file-backed pages and are therefore NOT
    policed here, so peak RSS can exceed the level. Both numbers are reported.
  * `-march=native` on this host defines __AMX_INT8__, which makes llama.cpp
    allocate a ~1.88 GiB AMX repack buffer before the arena-sized context. That
    alone pushes the floor past 4 GiB, so the sweep runs against the
    GGML_NATIVE=OFF (no-AMX) build unless --amx is passed.

Swap is proven absent at run time by sampling /proc/<pid>/status VmSwap; a
nonzero value fails the run.
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import re
import statistics
import subprocess
import sys
import threading
import time
from pathlib import Path

MIB = 1024 * 1024
GIB = 1024 ** 3

PROMPT = ("Write a detailed paragraph about how a mixture-of-experts model routes "
          "tokens to expert weights, and why paging those weights from SSD matters.")


def args_config(args):
    """JSON-safe view of the argparse namespace (drops the resolve helper)."""
    return {k: v for k, v in vars(args).items() if not callable(v)}


def host_facts() -> dict:
    cpu = ""
    for line in Path('/proc/cpuinfo').read_text().splitlines():
        if line.startswith('model name'):
            cpu = line.split(':', 1)[1].strip()
            break
    mem_total = 0
    for line in Path('/proc/meminfo').read_text().splitlines():
        if line.startswith('MemTotal'):
            mem_total = int(line.split()[1]) * 1024
            break
    cg = {}
    for k, f in (('memory.max', 'memory.max'),
                 ('memory.swap.max', 'memory.swap.max'),
                 ('memory.high', 'memory.high')):
        try:
            cg[k] = Path('/sys/fs/cgroup/' + f).read_text().strip()
        except Exception:
            cg[k] = None
    return {'cpu_model': cpu, 'cpu_threads': os.cpu_count(),
            'mem_total_bytes': mem_total, 'cgroup': cg}


class ProcSampler(threading.Thread):
    """Sample RSS / AnonHWM / VmSwap until told to stop.

    AnonHWM is smaps_rollup's own anonymous high-water mark and is the quantity
    RLIMIT_DATA actually polices, so it is what we compare against the level.
    """

    def __init__(self, pid: int, interval: float = 0.03):
        super().__init__(daemon=True)
        self.pid = pid
        self.interval = interval
        self.stop_flag = threading.Event()
        self.peak_rss = 0
        self.peak_anon = 0
        self.peak_swap = 0
        self.n = 0

    def _read(self) -> tuple[int, int, int]:
        rss = anon = swp = 0
        try:
            for line in Path(f'/proc/{self.pid}/status').read_text().splitlines():
                if line.startswith('VmRSS:'):
                    rss = int(line.split()[1]) * 1024
                elif line.startswith('VmSwap:'):
                    swp = int(line.split()[1]) * 1024
        except OSError:
            pass
        try:
            for line in Path(f'/proc/{self.pid}/smaps_rollup').read_text().splitlines():
                if line.startswith('Rss:'):
                    rss = max(rss, int(line.split()[1]) * 1024)
                elif line.startswith('Anonymous:'):
                    anon = int(line.split()[1]) * 1024
                elif line.startswith('AnonHWM:'):
                    anon = max(anon, int(line.split()[1]) * 1024)
        except OSError:
            pass
        return rss, anon, swp

    def run(self):
        while not self.stop_flag.is_set():
            r, a, s = self._read()
            self.n += 1
            self.peak_rss = max(self.peak_rss, r)
            self.peak_anon = max(self.peak_anon, a)
            self.peak_swap = max(self.peak_swap, s)
            time.sleep(self.interval)
        r, a, s = self._read()
        self.peak_rss = max(self.peak_rss, r)
        self.peak_anon = max(self.peak_anon, a)
        self.peak_swap = max(self.peak_swap, s)


def evict(shards: list[str], helper: str) -> None:
    subprocess.run(['sync'], check=False)
    if helper:
        subprocess.run([helper, '--evict'] + shards, check=False)


def run_once(args, level_mb: int, repeat: int) -> dict:
    stats_file = Path(args.output_dir) / f'level_{level_mb}_r{repeat}.jsonl'
    log_file = Path(args.output_dir) / f'level_{level_mb}_r{repeat}.log'
    stats_file.unlink(missing_ok=True)
    os.makedirs(args.output_dir, exist_ok=True)

    inner = [
        str(Path(args.bin).resolve()),
        '-m', str(Path(args.model).resolve()),
        '-p', PROMPT, '-n', str(args.n_predict), '-c', str(args.ctx),
        '-t', str(args.threads), '-b', str(args.ubatch), '--temp', '0',
        '--no-interactive', '--no-chat', '--stats-file', str(stats_file),
        '--ram-budget-mb', str(level_mb),
    ]
    # The strict ceiling: RLIMIT_DATA = arena + heap + all anon mmap.
    cmd = ['prlimit', f'--data={level_mb * MIB}', '--nofile=65536'] + inner

    env = dict(os.environ)
    env.update({
        'SDQ_RAM_BUDGET_MB': str(level_mb),
        'SDQ_RESERVE_MB': str(args.reserve_mb_for(level_mb)),
        'SDQ_PREFETCH': '0',
        'SDQ_FADV_DONTNEED': '1',
        'SDQ_BENCH_NO_EOG': '1',
    })

    evict(args.shards, args.evict_helper)

    t0 = time.perf_counter()
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                            text=True, encoding='utf-8', errors='replace', env=env)
    sampler = ProcSampler(proc.pid)
    sampler.start()
    try:
        out, err = proc.communicate(timeout=args.timeout)
    except subprocess.TimeoutExpired:
        proc.kill()
        out, err = proc.communicate()
        sampler.stop_flag.set()
        return {'level_mb': level_mb, 'repeat': repeat + 1,
                'error': f'timeout after {args.timeout}s'}
    sampler.stop_flag.set()
    sampler.join(timeout=2)
    wall = time.perf_counter() - t0
    both = out + '\n' + err
    log_file.write_text(both, encoding='utf-8', errors='replace')

    res = {
        'level_mb': level_mb,
        'level_gib': level_mb / 1024,
        'repeat': repeat + 1,
        'wall_s': round(wall, 2),
        'exit_code': proc.returncode,
        'peak_rss_mb': round(sampler.peak_rss / MIB, 1),
        'peak_anon_mb': round(sampler.peak_anon / MIB, 1),
        'peak_swap_mb': round(sampler.peak_swap / MIB, 1),
        'swap_free': sampler.peak_swap == 0,
        'anon_within_limit': sampler.peak_anon <= level_mb * MIB,
        'limit_enforced_by': 'prlimit --data (RLIMIT_DATA)',
    }

    if proc.returncode != 0:
        sig = -proc.returncode if proc.returncode < 0 else None
        res['signal'] = sig
        res['status'] = 'OOM_OR_ABORT'
        res['failure'] = ' | '.join(
            [l.strip() for l in err.splitlines()
             if re.search(r'bad_alloc|cannot allocate|MemoryError|std::terminate|'
                          r'abort|Killed|Expert arena|too small|error', l, re.I)][-5:])
        res['wall_s'] = round(wall, 2)
        return res

    res['status'] = 'OK'

    def grab(pat, text=both):
        m = re.search(pat, text, re.M)
        return m

    d = grab(r'decode:\s*(\d+)\s+tok\s*/\s*([\d.]+)s\s*=\s*([\d.]+)\s+tok/s')
    p = grab(r'prefill:\s*(\d+)\s+tok\s*/\s*([\d.]+)s\s*=\s*([\d.]+)\s+tok/s')
    mem = grab(r'總 RSS\s+(\d+)\s+MB\s*/\s+峰值\s+(\d+)\s+MB\s*/\s+swap\s+(\d+)\s+MB')
    split = grab(r'匿名 RSS\s+(\d+)\s+MB\s*/\s+檔案對映 RSS\s+(\d+)\s+MB\s*/\s+總 RSS\s+(\d+)\s+MB\s*/\s+峰值\s+(\d+)\s+MB')
    arena = grab(r'slots=\d+\s+arena=([\d.]+)\s+GiB')
    slots = grab(r'slots=(\d+)')
    ssd_tot = grab(r'全程 SSD 實際讀取:\s*([\d.]+)\s*GB')
    mib = grab(r'SSD 讀\s*([\d.]+)\s*MB\s*\(([\d.]+)\s*MB/tok\)')

    if d:
        res['decode_tok_s'] = float(d.group(3))
        res['decode_tok'] = int(d.group(1))
    if p:
        res['prefill_tok_s'] = float(p.group(3))
        res['prefill_tok'] = int(p.group(1))
    if mem:
        res['final_rss_mb'] = int(mem.group(1))
        res['selfpeak_rss_mb'] = int(mem.group(2))
        res['self_swap_mb'] = int(mem.group(3))
    if split:
        res['rss_anon_mb'] = int(split.group(1))
        res['rss_file_mb'] = int(split.group(2))
        res['rss_total_mb'] = int(split.group(3))
        res['selfpeak_rss_mb'] = int(split.group(4))
    if arena:
        res['arena_gib'] = float(arena.group(1))
    if slots:
        res['n_slots'] = int(slots.group(1))
    if ssd_tot:
        res['total_ssd_gb'] = float(ssd_tot.group(1))
    if mib:
        res['decode_ssd_mb'] = float(mib.group(1))
        res['ssd_mb_per_token'] = float(mib.group(2))

    rec = {}
    if stats_file.exists():
        for line in stats_file.read_text().splitlines():
            r = json.loads(line)
            if r.get('phase') == 'decode':
                rec = r
    stats_file.unlink(missing_ok=True)

    pl = rec.get('per_layer') or []
    rates = [x['hit_rate'] for x in pl if x.get('req', 0) > 0]
    hr = rec.get('hit_rate')
    res.update({
        'expert_requests': rec.get('expert_requests'),
        'expert_hits': rec.get('expert_hits'),
        'hot_rate_pct': hr,
        'cold_rate_pct': (100.0 - hr) if hr is not None else None,
        'evictions': rec.get('evictions'),
        'per_layer_hit_min': min(rates) if rates else None,
        'per_layer_hit_max': max(rates) if rates else None,
    })

    asst = ''
    started = False
    for line in out.splitlines():
        if line.startswith('assistant:'):
            started = True
            rest = line[len('assistant:'):].strip()
            if rest:
                asst = rest
            continue
        if started and line.strip() and not line.startswith('['):
            asst += ' ' + line.strip()
    res['assistant'] = asst.strip()[:200]
    return res


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', required=True)
    ap.add_argument('--shards', nargs='+', required=True)
    ap.add_argument('--bin', required=True)
    ap.add_argument('--levels-mb', type=int, nargs='+', default=[4096, 8192, 16384, 32768])
    ap.add_argument('--repeat', type=int, default=2)
    ap.add_argument('--ctx', type=int, default=2048)
    ap.add_argument('--ubatch', type=int, default=128)
    ap.add_argument('--threads', type=int, default=16)
    ap.add_argument('--n-predict', type=int, default=48)
    ap.add_argument('--reserve-mb', type=int, default=2400)
    ap.add_argument('--timeout', type=int, default=1800)
    ap.add_argument('--evict-helper', default='/tmp/pagecache_resident')
    ap.add_argument(
        '--reserve-map',
        default='4096:2400,8192:2400,16384:2400,32768:2400',
        help='per-level SDQ_RESERVE_MB, "level_mb:reserve_mb,...". Calibrated '
             'by measuring anon_hwm - arena, i.e. every anonymous allocation '
             'that happens after the arena is sized.')
    ap.add_argument('--output-dir', default='/root/app/ling/results_strict')
    ap.add_argument('--out-json', default='/root/app/ling/results_strict/strict_memory_levels.json')
    args = ap.parse_args()

    args._reserve = {}
    for item in args.reserve_map.split(','):
        k, v = item.split(':')
        args._reserve[int(k)] = int(v)
    args.reserve_mb_for = lambda lvl: args._reserve.get(lvl, 2400)

    Path(args.output_dir).mkdir(parents=True, exist_ok=True)

    print(f'host: {host_facts()}', file=sys.stderr)
    runs = []
    for lvl in args.levels_mb:
        for r in range(args.repeat):
            print(f'--- level {lvl} MB  repeat {r + 1}/{args.repeat} ---',
                  file=sys.stderr, flush=True)
            try:
                res = run_once(args, lvl, r)
            except Exception as e:
                res = {'level_mb': lvl, 'repeat': r + 1, 'error': repr(e)}
            runs.append(res)
            print(json.dumps(res, ensure_ascii=False)[:400], file=sys.stderr, flush=True)
            Path(args.out_json).write_text(json.dumps(
                {'host': host_facts(), 'config': args_config(args), 'runs': runs},
                indent=2, ensure_ascii=False))

    print('DONE', file=sys.stderr)
    Path(args.out_json).write_text(json.dumps(
        {'host': host_facts(), 'config': args_config(args), 'runs': runs},
        indent=2, ensure_ascii=False))


if __name__ == '__main__':
    main()