"""第 39 章：有界本地 HTTP 流式服务；只有 --serve 才加载模型。"""

import argparse
from collections import deque
from datetime import datetime, timezone
import hashlib
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
from pathlib import Path
import platform
import socket
import threading
import time
import uuid

from local_application_models import MODEL, REVISION

MAX_INPUT, MAX_OUTPUT, MAX_BODY = 768, 64, 32768
SYSTEM = "你是一位耐心的中文学习助手。请直接回答问题。"


class ServiceTimeout(RuntimeError):
    pass


def source_hash():
    return hashlib.sha256(Path(__file__).read_bytes()).hexdigest()


def validate_request(value):
    if not isinstance(value, dict) or set(value) - {"prompt", "max_new_tokens"}:
        raise ValueError("只接受 prompt 与可选 max_new_tokens")
    prompt, limit = value.get("prompt"), value.get("max_new_tokens", 24)
    if not isinstance(prompt, str) or not prompt.strip() or len(prompt) > 6000:
        raise ValueError("prompt 必须是 1～6000 字符的非空字符串")
    if type(limit) is not int or not 1 <= limit <= MAX_OUTPUT:
        raise ValueError(f"max_new_tokens 必须是 1～{MAX_OUTPUT} 的整数")
    return prompt, limit


def greedy_tokens(model, input_ids, limit, cached, eos_ids, deadline):
    """单条、无填充的请求；每次调用新建缓存，不跨请求共享。"""
    import torch
    full = input_ids
    cache = None
    with torch.inference_mode():
        for step in range(limit):
            if time.perf_counter() > deadline:
                raise ServiceTimeout("生成达到时间上限")
            current = full[:, -1:] if cached and step else full
            position = torch.arange(full.shape[1]-current.shape[1], full.shape[1])
            output = model(input_ids=current, attention_mask=torch.ones_like(full),
                           cache_position=position, past_key_values=cache,
                           use_cache=cached, logits_to_keep=1)
            token = output.logits[:, -1].argmax(-1, keepdim=True)
            cache = output.past_key_values if cached else None
            token_id = int(token[0, 0])
            full = torch.cat([full, token], dim=1)
            del output
            yield token_id
            if token_id in eos_ids:
                break


class Engine:
    def __init__(self, cached=True, local_only=False):
        import torch
        import transformers
        from transformers import AutoTokenizer, AutoModelForCausalLM
        torch.set_num_threads(4)
        torch.manual_seed(2026)
        self.cached = cached
        self.tokenizer = AutoTokenizer.from_pretrained(
            MODEL, revision=REVISION, trust_remote_code=False,
            local_files_only=local_only)
        self.model = AutoModelForCausalLM.from_pretrained(
            MODEL, revision=REVISION, dtype=torch.float32,
            attn_implementation="eager", trust_remote_code=False,
            local_files_only=local_only).eval()
        eos = self.model.generation_config.eos_token_id
        self.eos_ids = {eos} if isinstance(eos, int) else set(eos)
        self.metadata = dict(model=MODEL, revision=REVISION, cached=cached,
                             release="qwen-cached-v1" if cached else "qwen-full-v1",
                             source_sha256=source_hash(), system_prompt=SYSTEM,
                             max_input=MAX_INPUT, max_output=MAX_OUTPUT,
                             device="cpu", dtype="float32", attention="eager", threads=4,
                             decoding="greedy_without_penalties", eos_ids=sorted(self.eos_ids),
                             logits_to_keep=1,
                             python=platform.python_version(), torch=torch.__version__,
                             transformers=transformers.__version__, platform=platform.platform())

    def prepare(self, prompt, limit):
        ids = self.tokenizer.apply_chat_template(
            [{"role": "system", "content": SYSTEM}, {"role": "user", "content": prompt}],
            add_generation_prompt=True, return_tensors="pt")
        if ids.shape[1] > MAX_INPUT:
            raise ValueError(f"加上模板后的输入超过 {MAX_INPUT} 词元；不静默截断")
        if ids.shape[1]+limit > self.model.config.max_position_embeddings:
            raise ValueError("输入与预留输出超过模型上下文")
        return ids

    def stream(self, ids, limit, deadline):
        generated = []
        for token in greedy_tokens(self.model, ids, limit, self.cached, self.eos_ids, deadline):
            generated.append(token)
            # 累计解码可能修正之前未完整的 Unicode 字符；客户端应替换文本，不逐块追加。
            yield token, self.tokenizer.decode(generated, skip_special_tokens=True)


class State:
    def __init__(self, engine):
        self.engine = engine
        self.capacity = threading.BoundedSemaphore(2)  # 一个执行、最多一个等待
        self.worker = threading.Lock()  # 共享模型的生成串行；这不是连续批处理
        self.lock = threading.Lock()
        self.inflight = 0
        self.counts = dict(completed=0, failed=0, disconnected=0, rejected=0)
        self.recent = deque(maxlen=50)

    def snapshot(self):
        with self.lock:
            return dict(**self.engine.metadata, inflight=self.inflight,
                        capacity=2, counters=dict(self.counts), recent=list(self.recent))

    def reject(self):
        with self.lock:
            self.counts["rejected"] += 1


class Server(ThreadingHTTPServer):
    daemon_threads = True
    allow_reuse_address = True

    def __init__(self, address, engine):
        self.state = State(engine)
        self.connections = threading.BoundedSemaphore(8)
        super().__init__(address, Handler)

    def process_request(self, request, address):
        if not self.connections.acquire(blocking=False):
            self.state.reject()
            self.shutdown_request(request)
            return
        try:
            super().process_request(request, address)
        except BaseException:
            self.connections.release()
            raise

    def process_request_thread(self, request, address):
        try:
            super().process_request_thread(request, address)
        finally:
            self.connections.release()


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def setup(self):
        super().setup()
        self.connection.settimeout(5)

    def log_message(self, *_):
        pass  # 不记录原始提示或回答；运行指标见 /metrics

    def send_json(self, status, value):
        data = (json.dumps(value, ensure_ascii=False)+"\n").encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Connection", "close")
        self.end_headers()
        self.close_connection = True
        self.wfile.write(data)

    def event(self, value):
        self.wfile.write((json.dumps(value, ensure_ascii=False)+"\n").encode())
        self.wfile.flush()

    def do_GET(self):
        if self.path == "/health":
            self.send_json(200, dict(ready=True, **self.server.state.snapshot()))
        elif self.path == "/metrics":
            self.send_json(200, self.server.state.snapshot())
        else:
            self.send_json(404, {"error": "unknown_path"})

    def do_POST(self):
        started = time.perf_counter()
        if self.path != "/generate":
            self.send_json(404, {"error": "unknown_path"})
            return
        try:
            if self.headers.get("Transfer-Encoding") is not None:
                raise ValueError("本例不接受分块请求体")
            if self.headers.get("Content-Type", "").split(";")[0].strip() != "application/json":
                raise ValueError("需要 application/json")
            length = int(self.headers.get("Content-Length", "0"))
            if not 0 < length <= MAX_BODY:
                raise ValueError("请求体为空或超过 32768 字节")
            raw = self.rfile.read(length)
            if len(raw) != length:
                raise ValueError("请求体不完整")
            prompt, limit = validate_request(json.loads(raw.decode("utf-8")))
        except (ValueError, UnicodeError):
            self.send_json(400, {"error": "invalid_request"})
            return
        except (TimeoutError, OSError):
            self.close_connection = True
            return
        state = self.server.state
        if not state.capacity.acquire(blocking=False):
            state.reject()
            self.send_json(429, {"error": "busy", "retryable": True})
            return
        with state.lock:
            state.inflight += 1
        acquired, streaming = False, False
        status, queue_seconds, input_tokens, generated = "failed", None, 0, []
        request_id = uuid.uuid4().hex
        waiting = time.perf_counter()
        try:
            acquired = state.worker.acquire(timeout=30)
            if not acquired:
                raise ServiceTimeout("排队超过 30 秒")
            queue_seconds = time.perf_counter()-waiting
            ids = state.engine.prepare(prompt, limit)
            input_tokens = ids.shape[1]
            self.send_response(200)
            self.send_header("Content-Type", "application/x-ndjson; charset=utf-8")
            self.send_header("Connection", "close")
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.close_connection = True
            streaming = True
            self.event(dict(type="start", request_id=request_id,
                            release=state.engine.metadata["release"],
                            input_tokens=input_tokens, queue_seconds=queue_seconds))
            deadline = time.perf_counter()+60
            for token, text in state.engine.stream(ids, limit, deadline):
                generated.append(token)
                self.event(dict(type="token", index=len(generated), token_id=token, text=text))
            self.event(dict(type="done", request_id=request_id, output_tokens=len(generated),
                            stop_reason="eos" if generated[-1] in state.engine.eos_ids else "limit",
                            server_seconds=time.perf_counter()-started))
            status = "completed"
        except (BrokenPipeError, ConnectionResetError, socket.timeout):
            status = "disconnected"
        except Exception as error:
            code = ("invalid_input" if isinstance(error, ValueError) else
                    "deadline_exceeded" if isinstance(error, ServiceTimeout) else "generation_failed")
            try:
                if streaming:
                    self.event(dict(type="error", request_id=request_id, error=code))
                else:
                    self.send_json(400 if isinstance(error, ValueError) else 503, {"error": code})
            except OSError:
                status = "disconnected"
        finally:
            if acquired:
                state.worker.release()
            record = dict(request_id=request_id, status=status, input_tokens=input_tokens,
                          output_tokens=len(generated), queue_seconds=queue_seconds,
                          server_seconds=time.perf_counter()-started)
            with state.lock:
                state.inflight -= 1
                state.counts[status] += 1
                state.recent.append(record)
            state.capacity.release()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--serve", action="store_true")
    parser.add_argument("--port", type=int, default=8123)
    parser.add_argument("--no-cache", action="store_true")
    parser.add_argument("--local-only", action="store_true")
    args = parser.parse_args()
    if not args.serve:
        print("默认不下载模型、不启动端口。")
        print("启动：python3 src/examples/model_service.py --serve --port 8123")
        print("仅监听 127.0.0.1；/health、/metrics 与 POST /generate。")
        return
    engine = Engine(cached=not args.no_cache, local_only=args.local_only)
    server = Server(("127.0.0.1", args.port), engine)
    print(json.dumps(dict(type="ready", port=server.server_port,
                          started_at=datetime.now(timezone.utc).isoformat(),
                          **engine.metadata), ensure_ascii=False), flush=True)
    try:
        server.serve_forever(poll_interval=.1)
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()


if __name__ == "__main__":
    main()

"""第 39 章：真实 HTTP 客户端、服务子进程与 RSS 采样；默认复核随书记录。"""

import argparse
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import queue
import statistics
import subprocess
import sys
import tempfile
import threading
import time
from urllib.request import Request, urlopen

from model_service import MODEL, REVISION, SYSTEM, source_hash

DATA = Path(__file__).with_name("data")/"service-benchmark.json"
NEW_TOKENS, REPEATS, SAMPLE_SECONDS = 24, 3, .01
PROMPTS = {
    "short": "请分三点说明，为什么学习编程时应该亲手运行例子？",
    "long": ("学习日志：今天我们练习了变量、循环和函数。先猜程序输出，再运行核对，"
             "最后修改一个条件，记录结果怎样变化。\n")*6+
            "请根据这份日志，分三点建议明天怎样继续练习。"
}


def fingerprint(value):
    return hashlib.sha256(json.dumps(value, ensure_ascii=False, sort_keys=True).encode()).hexdigest()


def request_once(base, prompt, limit=NEW_TOKENS, barrier=None):
    payload = json.dumps(dict(prompt=prompt, max_new_tokens=limit), ensure_ascii=False).encode()
    request = Request(base+"/generate", data=payload, headers={"Content-Type": "application/json"})
    if barrier:
        barrier.wait(timeout=15)
    started = time.perf_counter()
    events, stamps, token_ids = [], [], []
    with urlopen(request, timeout=120) as response:
        if response.status != 200:
            raise RuntimeError(f"HTTP {response.status}")
        for line in response:
            stamp = time.perf_counter()-started
            event = json.loads(line)
            events.append(event)
            if event["type"] == "token":
                stamps.append(stamp)
                token_ids.append(event["token_id"])
    elapsed = time.perf_counter()-started
    if not events or events[0]["type"] != "start" or events[-1]["type"] != "done":
        raise RuntimeError("流式响应没有完整结束：" + repr(events[-1:]))
    start, done = events[0], events[-1]
    if not stamps or done["output_tokens"] != len(stamps):
        raise RuntimeError("词元数量与完成事件不一致")
    return dict(request_id=start["request_id"], release=start["release"],
                input_tokens=start["input_tokens"], output_tokens=len(stamps),
                queue_seconds=start["queue_seconds"], server_seconds=done["server_seconds"],
                total_seconds=elapsed, token_seconds=stamps, ids=token_ids,
                text=events[-2]["text"], stop_reason=done["stop_reason"])


@contextmanager
def service_process(cached, local_only=False):
    script = Path(__file__).with_name("model_service.py")
    command = [sys.executable, str(script), "--serve", "--port", "0"]
    if not cached:
        command.append("--no-cache")
    if local_only:
        command.append("--local-only")
    with tempfile.TemporaryFile(mode="w+") as errors:
        process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=errors, text=True)
        ready_queue = queue.Queue()
        reader = threading.Thread(target=lambda: ready_queue.put(process.stdout.readline()), daemon=True)
        reader.start()
        try:
            line = ready_queue.get(timeout=180)
            if not line:
                errors.seek(0)
                raise RuntimeError("服务启动失败：" + errors.read()[-3000:])
            ready = json.loads(line)
            if ready["type"] != "ready" or ready["cached"] != cached:
                raise RuntimeError("启动信息不符合预期")
            yield process, "http://127.0.0.1:"+str(ready["port"]), ready
        finally:
            process.terminate()
            try:
                process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait(timeout=10)
            process.stdout.close()
            reader.join(timeout=1)


def measure_group(process, base, prompt, concurrency):
    import psutil
    monitor = psutil.Process(process.pid)
    samples = []
    stopped = threading.Event()
    baseline = monitor.memory_info().rss
    started = time.perf_counter()

    def sample():
        while not stopped.is_set():
            samples.append([time.perf_counter()-started, monitor.memory_info().rss])
            stopped.wait(SAMPLE_SECONDS)

    sampler = threading.Thread(target=sample)
    sampler.start()
    barrier = threading.Barrier(concurrency+1)
    try:
        with ThreadPoolExecutor(max_workers=concurrency) as pool:
            futures = [pool.submit(request_once, base, prompt, NEW_TOKENS, barrier)
                       for _ in range(concurrency)]
            batch_start = time.perf_counter()
            barrier.wait(timeout=15)
            rows = [future.result() for future in futures]
            elapsed = time.perf_counter()-batch_start
    finally:
        stopped.set()
        sampler.join()
    samples.append([time.perf_counter()-started, monitor.memory_info().rss])
    return dict(concurrency=concurrency, group_seconds=elapsed, requests=rows,
                baseline_rss=baseline, peak_sampled_rss=max(rss for _, rss in samples),
                rss_samples=samples)


def benchmark(path, local_only=False):
    import psutil
    report = dict(complete=False, measured_at=datetime.now(timezone.utc).isoformat(),
                  model=MODEL, revision=REVISION, system_prompt=SYSTEM, prompts=PROMPTS,
                  prompts_sha256=fingerprint(PROMPTS), new_tokens=NEW_TOKENS,
                  repeats=REPEATS, rss_sample_seconds=SAMPLE_SECONDS, psutil=psutil.__version__,
                  source_sha256=source_hash(), benchmark_sha256=hashlib.sha256(
                      Path(__file__).read_bytes()).hexdigest(), servers=[], rows=[])
    path.parent.mkdir(parents=True, exist_ok=True)
    # 每种计算路径使用全新服务进程；预热与加载不计入请求测量。
    for cached in [False, True]:
        with service_process(cached, local_only) as (process, base, metadata):
            warmups = {key: request_once(base, prompt) for key, prompt in PROMPTS.items()}
            metadata["warmups"] = warmups
            report["servers"].append(metadata)
            for trial in range(REPEATS):
                order = [1, 2] if trial % 2 == 0 else [2, 1]
                for case, prompt in PROMPTS.items():
                    for concurrency in order:
                        group = measure_group(process, base, prompt, concurrency)
                        group.update(cached=cached, case=case, trial=trial)
                        for row in group["requests"]:
                            assert row["ids"] == warmups[case]["ids"], "同条件的贪心输出发生变化"
                        report["rows"].append(group)
                        path.write_text(json.dumps(report, ensure_ascii=False, indent=2)+"\n")
                        print(f"cached={cached} {case} c={concurrency} trial={trial+1} "
                              f"group={group['group_seconds']:.2f}s", flush=True)
            metadata["final_metrics"] = json.loads(urlopen(base+"/metrics", timeout=10).read())
    for case in PROMPTS:
        assert report["servers"][0]["warmups"][case]["ids"] == \
               report["servers"][1]["warmups"][case]["ids"], "两种计算路径输出不同"
    report["complete"] = True
    path.write_text(json.dumps(report, ensure_ascii=False, indent=2)+"\n")
    return report


def checked_report(path=DATA):
    report = json.loads(path.read_text())
    assert report["complete"] and report["model"] == MODEL and report["revision"] == REVISION
    assert report["source_sha256"] == source_hash()
    assert report["system_prompt"] == SYSTEM and report["prompts"] == PROMPTS
    assert report["prompts_sha256"] == fingerprint(PROMPTS)
    assert report["new_tokens"] == NEW_TOKENS and report["repeats"] == REPEATS
    assert report["rss_sample_seconds"] == SAMPLE_SECONDS
    expected = {(cached, case, c, t) for cached in [False, True] for case in PROMPTS
                for c in [1, 2] for t in range(REPEATS)}
    actual = {(g["cached"], g["case"], g["concurrency"], g["trial"]) for g in report["rows"]}
    assert actual == expected and len(report["rows"]) == len(expected)
    assert len(report["servers"]) == 2
    for server, cached in zip(report["servers"], [False, True]):
        assert server["cached"] == cached and server["revision"] == REVISION
        assert server["source_sha256"] == report["source_sha256"]
        assert server["logits_to_keep"] == 1
        metrics = server["final_metrics"]
        assert metrics["inflight"] == 0
        assert metrics["counters"] == dict(completed=20, failed=0, disconnected=0, rejected=0)
    for group in report["rows"]:
        assert len(group["requests"]) == group["concurrency"]
        samples = group["rss_samples"]
        assert samples and all(t >= 0 and rss > 0 for t, rss in samples)
        assert [t for t, _ in samples] == sorted(t for t, _ in samples)
        assert group["peak_sampled_rss"] == max(rss for _, rss in samples)
        for row in group["requests"]:
            stamps = row["token_seconds"]
            assert 0 < stamps[0] <= stamps[-1] <= row["total_seconds"] <= group["group_seconds"]+.1
            assert stamps == sorted(stamps)
            assert len(stamps) == len(row["ids"]) == row["output_tokens"] <= NEW_TOKENS
            assert 0 <= row["queue_seconds"] <= row["server_seconds"] <= row["total_seconds"]
            reference = report["servers"][0]["warmups"][group["case"]]
            assert row["ids"] == reference["ids"]
    return report


def summarize(report):
    result = []
    for case in PROMPTS:
        for cached in [False, True]:
            for c in [1, 2]:
                groups = [g for g in report["rows"] if
                          (g["case"], g["cached"], g["concurrency"]) == (case, cached, c)]
                requests = [r for g in groups for r in g["requests"]]
                first = [r["token_seconds"][0] for r in requests]
                last = [r["total_seconds"] for r in requests]
                speed = [(r["output_tokens"]-1)/(r["token_seconds"][-1]-r["token_seconds"][0])
                         for r in requests if r["output_tokens"] > 1]
                result.append(dict(case=case, cached=cached, concurrency=c,
                                   input_tokens=requests[0]["input_tokens"],
                                   ttft=statistics.median(first), total=statistics.median(last),
                                   queue=statistics.median(r["queue_seconds"] for r in requests),
                                   decode_tps=statistics.median(speed),
                                   throughput=statistics.median(
                                       sum(r["output_tokens"] for r in g["requests"])/g["group_seconds"]
                                       for g in groups),
                                   max_sampled_rss=max(g["peak_sampled_rss"] for g in groups)))
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--out", type=Path)
    parser.add_argument("--local-only", action="store_true")
    args = parser.parse_args()
    if args.run and args.out is None:
        parser.error("请用 --out 保存自己的测量")
    report = benchmark(args.out, args.local_only) if args.run else checked_report()
    print("本次 HTTP 测量：" if args.run else "随书保存的 HTTP 测量（默认不加载模型）：")
    for row in summarize(report):
        print(f"{row['case']}；KV={row['cached']}；并发={row['concurrency']}；"
              f"首词元={row['ttft']:.3f}s；总时长={row['total']:.3f}s；"
              f"后续速度={row['decode_tps']:.1f} token/s；"
              f"合计吞吐={row['throughput']:.1f} token/s；"
              f"RSS采样最大值={row['max_sampled_rss']/2**20:.1f} MiB")
    print("两种路径、两种输入、并发 1/2，各测 3 组；同输入的输出 ID 一致。")
    print("RSS 只采样服务进程，不含客户端；采样最大值可能低于真实峰值。")


if __name__ == "__main__":
    main()

1. 客户端 TTFT 包含排队；只从模型开始处理提示时计时的 prefill 不包含。比较两份报告前，要先对齐起点。
2. 首词元之后有 23 个词元，耗时 $2.7-0.4=2.3$ 秒，所以是 10 token/s。
3. HTTP 头已经发送后，模型仍可能失败。完整结束事件用于区分正常完成与中途断流；本例测量脚本会拒绝没有 done 的记录。
4. 共享模型的执行被锁串行化，第二个请求主要在等待。连续批处理可以把多个活动请求组织到同一次模型计算中，但仍受计算、内存与延迟目标限制。
5. RSS 还包含激活、KV、运行库、临时张量和分配器保留的内存。某些低位实现还会临时恢复浮点权重，必须测实际运行。
6. 本例的计数只在进程内存里，重启就重新开始。长期监控需要外部保存指标，并记录版本、实例和重启时间。
7. 至少核对权重 revision、分词器与模板、输出事件、长度限制、解码规则、重要任务结果及性能。保留可恢复的旧运行环境；端口能连上只是其中一步。
8. 使用新的输出路径，并更新自己的实验配置与记录。若模型提前生成 EOS，实际输出数可能不到 48；仍按实际收到的词元计算。不要预设首词元一定不变，整机噪声和实现路径都可能影响结果。
9. 把接口开销与模型计算分开看，保留错误与拒绝请求。如果改变精度或输出内容，再补能力对照，不能只比较速度。