|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""Hotfix: emit a generation header when a request effectively ends with a closed assistant turn, including when a trailing latest_reminder annotation follows it. |
| 3 | +
|
| 4 | +Symptom |
| 5 | +------- |
| 6 | +An agent harness gets stuck emitting empty turns: 1-2 generated tokens with |
| 7 | +`finish_reason: "stop"`, no tool call, and fragments of hallucinated markup |
| 8 | +(`<result observation="no-op"></content>`, stray `</parameter>` / `</name>` / |
| 9 | +`<value>`), or a sentence fragment that continues the previous turn instead of |
| 10 | +answering. Server metrics during a live incident: 6 of 37 requests generated |
| 11 | +<= 10 tokens, all `stop`, zero `length` — the model chose to stop rather than |
| 12 | +being truncated. |
| 13 | +
|
| 14 | +Cause |
| 15 | +----- |
| 16 | +`render_message()` appends the generation header (`ASSISTANT_SP_TOKEN` plus the |
| 17 | +thinking token) only when the trailing message is `user` or `developer`: |
| 18 | +
|
| 19 | + elif messages[index].get("role") in ["user", "developer"]: |
| 20 | +
|
| 21 | +When the request's `messages` array ends with an **assistant** message, that |
| 22 | +branch is skipped, the turn is closed with EOS, and the prompt ends on a bare |
| 23 | +EOS with no header. The model is asked to generate from a closed, turnless |
| 24 | +state. Measured on this recipe: |
| 25 | +
|
| 26 | + assistant-final: '...to the reviewer for a fresh verdict:<|end of sentence|>' |
| 27 | + user-final: '...verdict:<EOS><|User|>continue<|Assistant|><think>' |
| 28 | +
|
| 29 | +It is self-sustaining: the harness records the resulting empty turn, so the next |
| 30 | +request also ends with an assistant message. One bad turn locks the loop; |
| 31 | +inserting any user-role message breaks it instantly. Requests reach this shape |
| 32 | +whenever a harness retries after a mid-stream error and re-sends the partial |
| 33 | +assistant turn, which is why the same prompt works for a long time and then |
| 34 | +does not. |
| 35 | +
|
| 36 | +The same retry shape can carry a trailing `latest_reminder` annotation after |
| 37 | +the re-sent partial assistant turn (the harness appends fresh context as a |
| 38 | +reminder message). The reminder defeats the fix twice over: stock closes the |
| 39 | +assistant turn with EOS, renders the bare reminder after it, and the prompt |
| 40 | +still ends with no generation header — the model reads the reminder from the |
| 41 | +same dead state. |
| 42 | +
|
| 43 | +Why a header and not `wo_eos` |
| 44 | +----------------------------- |
| 45 | +`encoding_dsv4.py` also has `assistant_msg_wo_eos_template`, and reopening the |
| 46 | +trailing turn looks like a tempting one-line fix. It is wrong. Measured on one |
| 47 | +prompt via /v1/completions, generated tokens: |
| 48 | +
|
| 49 | + trailing turn | stock (EOS) | wo_eos (reopen) | EOS + header |
| 50 | + partial | 32 | 208 | 260 |
| 51 | + complete | 16 (fragment)| 1 (EMPTY, dead)| 140 (healthy) |
| 52 | +
|
| 53 | +Reopening a *complete* assistant turn just moves the dead state: the model has |
| 54 | +nothing left to add and emits EOS immediately. Appending a fresh generation |
| 55 | +header after the closed assistant turn is correct for both shapes and matches |
| 56 | +the checkpoint encoder's existing generation transition. |
| 57 | +
|
| 58 | +Scope |
| 59 | +----- |
| 60 | +Only the final message of a request is affected, and only when it is an |
| 61 | +assistant turn, or a `latest_reminder` whose immediate predecessor is an |
| 62 | +assistant turn. Every other rendering path is untouched, including: |
| 63 | +consecutive assistant messages mid-transcript; reminders mid-transcript; and |
| 64 | +reminder tails directly after a `user`/`developer` message — those already |
| 65 | +end inside the pending generation slot (the checkpoint emits |
| 66 | +`ASSISTANT_SP_TOKEN` + thinking token *before* such a reminder) and must stay |
| 67 | +byte-identical. |
| 68 | +
|
| 69 | +Gating and fail-closed operation |
| 70 | +-------------------------------- |
| 71 | +The compose entrypoint invokes this script only when |
| 72 | +`DSPARK_ENABLE_ASSISTANT_FINAL_HOTFIX` is exactly `1` (default `0` = stock |
| 73 | +renderer, this script never runs), and chains it with `|| exit 1`. Because an |
| 74 | +invocation means the operator asked for the fix, everything fails nonzero: |
| 75 | +
|
| 76 | +- encoder file missing (the prerequisite `encoding_dsv4.py` copy did not |
| 77 | + happen) — a gated-ON boot must not silently serve the buggy stock renderer; |
| 78 | +- anchor text missing (upstream encoder drifted) — nothing is written; |
| 79 | +- post-write self-check failure (patched module does not import, a fixed |
| 80 | + shape still renders without a generation header — assistant-final, or |
| 81 | + assistant-final plus trailing `latest_reminder` — or a |
| 82 | + user->latest_reminder tail gains a second header) — the original file |
| 83 | + bytes are restored first, then exit 1. |
| 84 | +
|
| 85 | +Idempotent: an already-patched encoder is not rewritten, but it must still pass |
| 86 | +the self-check. Patches the encoding module the server actually loads, so it |
| 87 | +must run AFTER the compose entrypoint copies `encoding_dsv4.py` into place. |
| 88 | +
|
| 89 | +Usage (inside container, after encoder copy): |
| 90 | + python3 hotfix-dsv4-assistant-final-continuation.py |
| 91 | + python3 hotfix-dsv4-assistant-final-continuation.py /path/to/deepseek_v4_encoding.py |
| 92 | +""" |
| 93 | +from __future__ import annotations |
| 94 | + |
| 95 | +import importlib.util |
| 96 | +import sys |
| 97 | +from pathlib import Path |
| 98 | + |
| 99 | +DEFAULT_TARGET = Path( |
| 100 | + "/usr/local/lib/python3.12/dist-packages/vllm/tokenizers/deepseek_v4_encoding.py" |
| 101 | +) |
| 102 | +MARK = "[assistant-final-hotfix]" |
| 103 | + |
| 104 | +OLD = ( |
| 105 | + " elif messages[index].get(\"role\") in [\"user\", \"developer\"]:\n" |
| 106 | + " # Normal generation: append Assistant + thinking token\n" |
| 107 | +) |
| 108 | +NEW = ( |
| 109 | + " elif messages[index].get(\"role\") in [\"user\", \"developer\"] or (\n" |
| 110 | + f" # {MARK} A request may legitimately end with an assistant turn\n" |
| 111 | + " # (harness retry, continuation), optionally annotated by a\n" |
| 112 | + " # trailing latest_reminder harness message. Without a generation\n" |
| 113 | + " # header the prompt ends on a bare EOS (or a bare reminder after\n" |
| 114 | + " # the closed turn) and the model generates from a dead state:\n" |
| 115 | + " # immediate EOS, or raw DSML markup emitted as text. A reminder\n" |
| 116 | + " # tail directly after user/developer already ends inside the\n" |
| 117 | + " # pending generation slot and must stay byte-identical.\n" |
| 118 | + " messages[index].get(\"role\") == \"assistant\"\n" |
| 119 | + " and index == len(messages) - 1\n" |
| 120 | + " ) or (\n" |
| 121 | + " messages[index].get(\"role\") == \"latest_reminder\"\n" |
| 122 | + " and index == len(messages) - 1\n" |
| 123 | + " and index > 0\n" |
| 124 | + " and messages[index - 1].get(\"role\") == \"assistant\"\n" |
| 125 | + " ):\n" |
| 126 | + " # Normal generation: append Assistant + thinking token\n" |
| 127 | +) |
| 128 | + |
| 129 | + |
| 130 | +def _self_check(target: Path) -> tuple[bool, str]: |
| 131 | + """Import the patched encoder and confirm the fix actually renders. |
| 132 | +
|
| 133 | + A trailing-assistant transcript must end with the generation header |
| 134 | + (assistant speaker token + thinking token) instead of a bare EOS, the |
| 135 | + same shape annotated by a trailing latest_reminder must regain that |
| 136 | + fresh header after the reminder, and a user->latest_reminder tail must |
| 137 | + keep its stock single in-slot header (no second header appended). |
| 138 | + """ |
| 139 | + spec = importlib.util.spec_from_file_location("enc_check", target) |
| 140 | + if spec is None or spec.loader is None: |
| 141 | + return False, f"cannot load module spec from {target}" |
| 142 | + enc = importlib.util.module_from_spec(spec) |
| 143 | + try: |
| 144 | + spec.loader.exec_module(enc) |
| 145 | + base = [ |
| 146 | + {"role": "system", "content": "s"}, |
| 147 | + {"role": "user", "content": "u"}, |
| 148 | + {"role": "assistant", "content": "A finished answer."}, |
| 149 | + ] |
| 150 | + rendered = enc.encode_messages(base, "thinking", reasoning_effort="high") |
| 151 | + reminded = enc.encode_messages( |
| 152 | + base + [{"role": "latest_reminder", "content": "Fresh context."}], |
| 153 | + "thinking", |
| 154 | + reasoning_effort="high", |
| 155 | + ) |
| 156 | + intact = enc.encode_messages( |
| 157 | + [ |
| 158 | + {"role": "system", "content": "s"}, |
| 159 | + {"role": "user", "content": "u"}, |
| 160 | + {"role": "latest_reminder", "content": "Fresh context."}, |
| 161 | + ], |
| 162 | + "thinking", |
| 163 | + reasoning_effort="high", |
| 164 | + ) |
| 165 | + speaker = getattr( |
| 166 | + enc, "ASSISTANT_SP_TOKEN", getattr(enc, "assistant_sp_token", None) |
| 167 | + ) |
| 168 | + thinking = getattr(enc, "thinking_start_token", None) |
| 169 | + except Exception as err: # broken/unimportable patch must fail closed |
| 170 | + return False, f"self-check raised {type(err).__name__}: {err}" |
| 171 | + if not speaker or not thinking: |
| 172 | + return False, "generation-header tokens are unavailable" |
| 173 | + |
| 174 | + def _ends_with_header(out): |
| 175 | + if isinstance(out, str): |
| 176 | + return out.endswith(speaker + thinking) |
| 177 | + if isinstance(out, (list, tuple)): |
| 178 | + return list(out[-2:]) == [speaker, thinking] |
| 179 | + raise TypeError(f"unexpected encoder output type: {type(out).__name__}") |
| 180 | + |
| 181 | + try: |
| 182 | + valid = _ends_with_header(rendered) and _ends_with_header(reminded) |
| 183 | + # The user->latest_reminder tail ends inside the pending generation |
| 184 | + # slot (exactly one header, before the reminder); an over-broad |
| 185 | + # transition would append a second header there. |
| 186 | + untouched = intact.count(speaker) == 1 and not _ends_with_header(intact) |
| 187 | + except TypeError as err: |
| 188 | + return False, str(err) |
| 189 | + if not valid: |
| 190 | + return False, f"generation header does not terminate render: {rendered[-80:]!r}" |
| 191 | + if not untouched: |
| 192 | + return False, "transition widened too far: user->latest_reminder tail changed" |
| 193 | + return True, "generation headers terminate assistant-final renders" |
| 194 | + |
| 195 | + |
| 196 | +def main(argv: list[str]) -> int: |
| 197 | + target = Path(argv[1]) if len(argv) > 1 else DEFAULT_TARGET |
| 198 | + |
| 199 | + if not target.is_file(): |
| 200 | + # Invoked == gated ON: a missing encoder is a prerequisite failure, |
| 201 | + # not a skip (compose chains this with `|| exit 1`). |
| 202 | + print(f"[FAIL] {MARK} encoder file not found: {target}", file=sys.stderr) |
| 203 | + return 1 |
| 204 | + |
| 205 | + src = target.read_text(encoding="utf-8") |
| 206 | + |
| 207 | + if MARK in src: |
| 208 | + ok, why = _self_check(target) |
| 209 | + if ok: |
| 210 | + print(f"[OK] {MARK} already applied and verified: {target}") |
| 211 | + return 0 |
| 212 | + print( |
| 213 | + f"[FAIL] {MARK} already applied but self-check failed: {why}", |
| 214 | + file=sys.stderr, |
| 215 | + ) |
| 216 | + return 1 |
| 217 | + |
| 218 | + if OLD not in src: |
| 219 | + print(f"[FAIL] {MARK} anchor not found in {target}", file=sys.stderr) |
| 220 | + return 1 |
| 221 | + |
| 222 | + target.write_text(src.replace(OLD, NEW, 1), encoding="utf-8") |
| 223 | + ok, why = _self_check(target) |
| 224 | + if ok: |
| 225 | + print(f"[OK] {MARK} patched and verified: {target}") |
| 226 | + return 0 |
| 227 | + |
| 228 | + # Fail closed: never leave a written-but-unverified encoder behind. |
| 229 | + target.write_text(src, encoding="utf-8") |
| 230 | + print( |
| 231 | + f"[FAIL] {MARK} self-check failed, original restored ({target}): {why}", |
| 232 | + file=sys.stderr, |
| 233 | + ) |
| 234 | + return 1 |
| 235 | + |
| 236 | + |
| 237 | +if __name__ == "__main__": |
| 238 | + raise SystemExit(main(sys.argv)) |
0 commit comments