File size: 8,538 Bytes
3fd1a35
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
"""Dedicated-board regression for opt-in generated-state reuse; no concurrent inference."""
import argparse
import json
import time
import urllib.request
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--budget-test", action="store_true")
    args = parser.parse_args()
    client = urllib.request.build_opener(urllib.request.ProxyHandler({}))

    def open_request(path, body=None):
        return client.open(urllib.request.Request(args.base + path,
            data=None if body is None else json.dumps(body).encode(),
            headers={"Content-Type": "application/json"}), timeout=180)

    def call(path, body=None):
        with open_request(path, body) as response:
            return json.load(response)

    def infer(messages, session="a", reuse=True, count=16, **extra):
        body = {"model": "mindnano-ling3-tiny", "user": "generated-test", "messages": messages,
                "temperature": 0, "max_tokens": count, "reuse_generated_state": reuse}
        if session is not None:
            body["session_id"] = session
        return call("/v1/chat/completions", dict(body, **extra))

    def equal(a, b):
        assert a["choices"] == b["choices"], "identical numerical path changed output"

    def follow(seed, result, question="请简短说明刚才的内容。"):
        return seed + [result["choices"][0]["message"], {"role": "user", "content": question}]

    def clear():
        call("/v1/cache/clear", {})

    results = {"cases": [], "passed": False}
    assert not call("/v1/generation/status")["active"]
    assert call("/health")["capabilities"]["generated_state_reuse"]
    clear()
    try:
        seed = [{"role": "user", "content": "1" * 107}]  # 128-token stable prefix
        if args.budget_test:
            first = infer(seed, count=64)
            assert first["mindnano_metrics"]["continuation_skipped_budget"]
            assert first["mindnano_metrics"]["continuation_stored"]
            later = infer(follow(seed, first))
            assert later["mindnano_metrics"]["generated_cached_tokens"] == 63
            status = call("/v1/cache/status")
            assert status["snapshot_bytes"] <= status["budget_bytes"]
            # Fair single-session comparison: the old path also keeps its
            # original prompt cache. Compare against it, not only cold prefill.
            clear()
            first = infer(seed, session=None, count=64)
            messages = follow(seed, first)
            incremental = infer(messages, session=None)
            clear()
            reference_seed = infer(seed, session=None, reuse=False, count=64)
            equal(first, reference_seed)
            baseline = infer(messages, session=None, reuse=False)
            assert baseline["mindnano_metrics"]["cached_tokens"] == 128
            assert incremental["mindnano_metrics"]["generated_cached_tokens"] == 63
            results["cached_baseline"] = {"incremental": incremental, "original_prompt_cache": baseline,
                                          "same_reply": incremental["choices"] == baseline["choices"]}
            clear()
            long_seed = infer(seed, session=None, count=512)
            long_messages = follow(seed, long_seed)
            long_incremental = infer(long_messages, session=None)
            assert long_incremental["mindnano_metrics"]["generated_cached_tokens"] == 511
            clear()
            long_reference = infer(seed, session=None, reuse=False, count=512)
            equal(long_seed, long_reference)
            long_baseline = infer(long_messages, session=None, reuse=False)
            assert long_baseline["mindnano_metrics"]["cached_tokens"] == 128
            results["long_cached_baseline"] = {"incremental": long_incremental, "original_prompt_cache": long_baseline,
                                               "same_reply": long_incremental["choices"] == long_baseline["choices"]}
            results.update(passed=True, budget=status)
            return
        for session in (None, "a"):
            clear()
            first = infer(seed, session, count=64)
            assert first["mindnano_metrics"]["continuation_stored"]
            messages = follow(seed, first)
            if session:
                infer([{"role": "user", "content": "你好。"}], "other")
            warm = infer(messages, session)
            m = warm["mindnano_metrics"]
            assert m["generated_cached_tokens"] == 63
            assert m["cache_status"] == "generated_prefix_hit"
            assert m["cached_tokens"] + m["prompt_evaluated_tokens"] == warm["usage"]["prompt_tokens"]
            again = infer(messages, session)
            equal(warm, again)
            assert again["mindnano_metrics"]["prompt_evaluated_tokens"] == 0
            # Turning it off must also reject exact caches derived from decode state.
            strict = infer(messages, session, reuse=False)
            assert strict["mindnano_metrics"]["cached_tokens"] == 0
            cold = infer(messages, session, reuse=False, cache_prompt=False)
            equal(strict, cold)
            results["cases"].append({"session": session, "warm": warm, "cold": cold,
                                     "same_reply_as_cold": warm["choices"] == cold["choices"]})

        # Fork and cancel preserve the last committed generated prefix.
        clear()
        first = infer(seed, count=64)
        call("/v1/cache/fork", {"user": "generated-test", "source_session_id": "a", "target_session_id": "fork"})
        messages = follow(seed, first)
        branch = infer(messages, "fork")
        assert branch["mindnano_metrics"]["generated_cached_tokens"] == 63
        body = {"model": "mindnano-ling3-tiny", "user": "generated-test", "session_id": "a",
                "messages": [{"role": "user", "content": "讲个故事。"}], "stream": True,
                "flow_control": "ack", "reuse_generated_state": True, "max_tokens": 64}
        with open_request("/v1/chat/completions", body) as response:
            for line in response:
                if line.startswith(b"data:"):
                    event = json.loads(line[5:])
                    if event.get("mindnano_flow"):
                        call("/v1/cancel", {"request_id": event["id"]})
                        break
            for line in response:
                if line.strip() == b"data: [DONE]":
                    break
        deadline = time.monotonic() + 10
        while call("/v1/generation/status")["active"]:
            assert time.monotonic() < deadline
            time.sleep(.05)
        resumed = infer(messages)
        assert resumed["mindnano_metrics"]["generated_cached_tokens"] == 63
        equal(branch, resumed)

        # Changed/shorter histories and stop-filtered output cannot reuse hidden tokens.
        for variant in ("edited", "shorter", "stop"):
            clear()
            first = infer(seed, count=64, **({"stop": "11111"} if variant == "stop" else {}))
            messages = follow(seed, first)
            if variant == "edited":
                messages[1]["content"] = "修改过的回答。"
            elif variant == "shorter":
                messages = [{"role": "user", "content": "你好。"}]
            response = infer(messages)
            assert response["mindnano_metrics"]["generated_cached_tokens"] == 0
        # Sample both language modes. Different cold output is measured, not hidden.
        for think in (False, True):
            clear()
            seed = [{"role": "user", "content": "用中文介绍一次会议记录应包含什么。"}]
            first = infer(seed, count=128, enable_thinking=think)
            messages = follow(seed, first, "请简短总结。")
            warm = infer(messages, enable_thinking=think)
            cold = infer(messages, reuse=False, cache_prompt=False, enable_thinking=think)
            results["cases"].append({"thinking": think, "warm": warm, "cold": cold,
                                     "same_reply_as_cold": warm["choices"] == cold["choices"]})
        results["passed"] = True
    finally:
        clear()
        args.output.write_text(json.dumps(results, ensure_ascii=False, indent=2) + "\n")
        print(json.dumps({"generated_cache": "PASS" if results["passed"] else "FAIL",
                          "results": str(args.output)}), flush=True)


if __name__ == "__main__":
    main()