diff --git a/src/hyperloom/inference_optimizer/tests/test_grid_runner.py b/src/hyperloom/inference_optimizer/tests/test_grid_runner.py index 2cec7bf2d1..f48fdc22a1 100644 --- a/src/hyperloom/inference_optimizer/tests/test_grid_runner.py +++ b/src/hyperloom/inference_optimizer/tests/test_grid_runner.py @@ -2371,6 +2371,65 @@ def test_remove_noop_when_nothing_matches_keeps_json(self): out = _grid_runner.remove_server_args(raw, ["--port"]) assert json.loads(out.split("--compilation-config ", 1)[1].strip()) == {"cudagraph_mode": "FULL"} + def test_sign_prefixed_custom_op_survives_removal(self): + """A ``+``/``-`` prefixed custom-op value must survive the round trip. + + Live regression: every variant of an explore round died at vLLM argv + parse with ``Invalid JSON: key must be a string`` -- including a control + leg that added one env var and zero args, because + ``strip_benchmark_harness_flags`` routes EVERY composed variant through + ``remove_server_args`` with a non-empty denylist. The old POSIX + ``shlex.split``/rejoin stripped the JSON quotes, and the bareword repair + heuristic could not re-quote ``+fused_rms_norm_gated`` (its ``+`` was + outside the charset), so the corruption reached the server verbatim. + """ + raw = ( + "--compilation-config " + '{"mode":3,"custom_ops":["+fused_rms_norm_gated"],' + '"cudagraph_capture_sizes":[1,2,3]} --port 8888' + ) + out = _grid_runner.remove_server_args(raw, ["--port"]) + assert "--port" not in out + blob = json.loads(out.split("--compilation-config ", 1)[1].strip()) + assert blob["custom_ops"] == ["+fused_rms_norm_gated"] + assert blob["cudagraph_capture_sizes"] == [1, 2, 3] + + def test_compose_preserves_json_for_every_args_mode(self): + """All four variant shapes keep both JSON flags ``json.loads``-able.""" + cc = '{"mode":3,"custom_ops":["+fused_rms_norm_gated"]}' + sc = '{"method":"mtp","num_speculative_tokens":2}' + base = f"--max-num-seqs 20 --compilation-config {cc} --speculative-config {sc}" + + def _blob(text, flag): + toks = text.split() + return json.loads(toks[toks.index(flag) + 1]) + + for kwargs in ( + {"inherited_args": base}, + {"inherited_args": base, "remove_args": ["--max-num-seqs"]}, + {"base_extra_args": base, "args_mode": "replace"}, + ): + out = _grid_runner.compose_server_args(**kwargs) + assert _blob(out, "--compilation-config") == json.loads(cc) + assert _blob(out, "--speculative-config") == json.loads(sc) + # Removing the JSON flag and supplying a variant replacement. + repl = '{"mode":3,"custom_ops":["-rms_norm"]}' + out = _grid_runner.compose_server_args( + inherited_args=base, + remove_args=["--compilation-config"], + variant_extra_args=f"--compilation-config {repl}", + ) + assert out.count("--compilation-config") == 1 + assert _blob(out, "--compilation-config") == json.loads(repl) + assert _blob(out, "--speculative-config") == json.loads(sc) + + def test_shell_quoted_plain_operand_loses_its_wrappers(self): + """Magpie expands ``EXTRA_*_ARGS`` unquoted, so wrappers must not persist.""" + out = _grid_runner.remove_server_args( + "--tool-call-parser 'kimi_k3' --port 8888", ["--port"] + ) + assert out == "--tool-call-parser kimi_k3" + # --------------------------------------------------------------------------- # benchmark_report settling (real-run race) diff --git a/src/hyperloom/orchestrator/actions/executors/_grid_server_args.py b/src/hyperloom/orchestrator/actions/executors/_grid_server_args.py index 71c188cb92..61bb62df35 100644 --- a/src/hyperloom/orchestrator/actions/executors/_grid_server_args.py +++ b/src/hyperloom/orchestrator/actions/executors/_grid_server_args.py @@ -100,6 +100,46 @@ def merge_server_args(*parts: str | None) -> str: return " ".join(str(p).strip() for p in parts if str(p or "").strip()) +def _unwrap_shell_quotes(token: str) -> str: + """Drop one balanced pair of shell quotes wrapping a whitespace-free token. + + ``shlex.split(..., posix=False)`` keeps quote bytes in the token it returns, + which is exactly what protects a JSON value's inner double quotes. But a + plain operand written as ``--tool-call-parser 'kimi'`` must not keep its + wrappers, because Magpie expands ``EXTRA_*_ARGS`` unquoted and the wrappers + would reach argv literally. Only fully-wrapped, whitespace-free tokens are + unwrapped; anything whose content needs the quotes is returned verbatim so + token boundaries can never shift. A JSON blob starts with ``{``/``[`` and is + therefore never touched here. + """ + if len(token) >= 2 and token[0] == token[-1] and token[0] in "\"'": + inner = token[1:-1] + if inner and token[0] not in inner and not any(ch.isspace() for ch in inner): + return inner + return token + + +def _split_args_preserving_json(text: str) -> list[str] | None: + """Tokenize a server-arg string WITHOUT stripping JSON's inner double quotes. + + ``shlex.split(text)`` defaults to ``posix=True``, which consumes every quote + byte. Splitting a JSON-valued flag that way and space-joining the result + turns a stored-valid ``--compilation-config {"mode":3}`` into + ``{mode:3}``, and vLLM then aborts at argv parse with + ``Invalid JSON: key must be a string``. Valid blobs are compacted first (so + a blob is a single whitespace-free word) and then split in non-POSIX mode, + which preserves the quote bytes verbatim — a lossless round trip. + + Returns ``None`` when the string is not tokenizable at all, so callers can + leave the input untouched rather than guess. + """ + try: + tokens = shlex.split(_reserialize_json_blobs(text), posix=False) + except ValueError: + return None + return [_unwrap_shell_quotes(tok) for tok in tokens] + + def remove_server_args(server_args: str | None, remove_args: Any) -> str: """Remove flag specs from a server-arg string. @@ -107,22 +147,27 @@ def remove_server_args(server_args: str | None, remove_args: Any) -> str: its following value when one is present; ``"--foo=bar"`` removes that exact token shape; ``"--foo bar"`` removes the exact flag/value pair. Unknown / unparseable inputs are left untouched rather than guessed. + + Tokenization is quote-preserving (:func:`_split_args_preserving_json`), so + every flag this function does NOT remove survives byte-for-byte, JSON values + included. This matters far beyond explicit removals: + :func:`strip_benchmark_harness_flags` routes EVERY composed variant through + here with a non-empty denylist, so a lossy round trip would corrupt a + sibling ``--compilation-config`` even for a variant that removes nothing. """ args = str(server_args or "").strip() removes = to_str_list(remove_args) if not args or not removes: return args - try: - tokens = shlex.split(args) - except ValueError: + tokens = _split_args_preserving_json(args) + if tokens is None: return args remove_flags: set[str] = set() remove_pairs: set[tuple[str, str | None]] = set() for spec in removes: - try: - spec_tokens = shlex.split(spec) - except ValueError: + spec_tokens = _split_args_preserving_json(spec) + if spec_tokens is None: spec_tokens = spec.split() i = 0 while i < len(spec_tokens): @@ -161,11 +206,11 @@ def remove_server_args(server_args: str | None, remove_args: Any) -> str: continue out.append(tok) i += 1 - # ``shlex.split`` above strips the inner double quotes of any JSON-valued - # flag (``--compilation-config {"cudagraph_mode":"FULL"}`` -> - # ``{cudagraph_mode:FULL}``); re-quote/compact the JSON blobs so removal - # never corrupts a sibling flag that vLLM parses with ``json.loads``. - return _reserialize_json_blobs(" ".join(out)) + # Tokens are already JSON-compacted and quote-preserving, so a plain join is + # lossless: no after-the-fact re-quoting of a damaged blob is needed (and + # none was ever possible for a value like ``["+fused_rms_norm_gated"]``, + # whose ``+`` sign the repair heuristic could not reconstruct). + return " ".join(out) # Serving-ineligible harness flags. Enroll here; compose_server_args strips them @@ -199,14 +244,74 @@ def compose_server_args( combined_base = merge_server_args(inherited_args, base_extra_args) pruned = remove_server_args(combined_base, remove_args) composed = merge_server_args(pruned, variant_extra_args) - return strip_benchmark_harness_flags(composed) + result = strip_benchmark_harness_flags(composed) + _warn_on_damaged_json_values(composed, result) + return result + + +def _json_flag_values(args: str) -> dict[str, list[str]]: + """Map each :data:`SPACE_VALUE_FLAGS` occurrence to its raw value token.""" + found: dict[str, list[str]] = {} + for flag in SPACE_VALUE_FLAGS: + start = 0 + while True: + i = args.find(flag + " ", start) + if i < 0: + break + value = args[i + len(flag) :].strip().split(" ", 1)[0] + if value[:1] in ("{", "["): + found.setdefault(flag, []).append(value) + start = i + len(flag) + return found + + +def _warn_on_damaged_json_values(before: str, after: str) -> None: + """Log loudly when composition turned a parseable JSON flag value unparseable. + + This is a regression tripwire, not a repair. The composer damaged + ``--compilation-config`` for an entire optimization session by shlex + round-tripping it lossily: the value stayed a single shell word, so nothing + downstream looked wrong, and the only symptom was every variant server dying + at argv parse ~18s in while the baseline (which never routes through + :func:`compose_server_args`) ran clean for 4206s. The damage was silent + because ``_repair_unquoted_json`` "succeeded" on the sibling + ``--speculative-config`` and merely returned ``None`` for the one blob it + could not reconstruct. Emitting a loud, greppable line here converts that + class of failure from a multi-round mystery into one log grep. + """ + try: + was = _json_flag_values(before) + now = _json_flag_values(after) + except Exception: # never let a diagnostic break composition + return + for flag, values in now.items(): + healthy_before = any(_parses_as_json(v) for v in was.get(flag, [])) + if healthy_before and not any(_parses_as_json(v) for v in values): + log.error( + "server-arg composition CORRUPTED %s: its value parsed as JSON " + "before composition and does not after. The launched server will " + "abort at argv parse. Damaged value: %s", + flag, + values[0][:200] if values else "", + ) + + +def _parses_as_json(value: str) -> bool: + try: + json.loads(value) + except Exception: + return False + return True # A JSON "bareword": an identifier-like token that appears where a double-quoted # JSON key or string value should be (letters/digits/underscore plus the ``.``, -# ``/``, ``-`` common in model ids and paths). Numbers, ``true``/``false``/ -# ``null`` are handled separately so they stay unquoted. -_JSON_BAREWORD = r"[A-Za-z_][A-Za-z0-9_./-]*" +# ``/``, ``-`` common in model ids and paths). An optional leading ``+``/``-`` +# sign covers vLLM's custom-op toggles (``custom_ops:["+fused_rms_norm_gated"]``) +# — without it the repair silently failed on exactly those values. A sign is +# only accepted when a letter/underscore follows, so numbers (``-1``) and +# ``true``/``false``/``null`` are still handled separately and stay unquoted. +_JSON_BAREWORD = r"[+-]?[A-Za-z_][A-Za-z0-9_./-]*" _UNQUOTED_KEY_RE = re.compile(r"([{,]\s*)(" + _JSON_BAREWORD + r")(\s*:)") _UNQUOTED_VALUE_RE = re.compile(r"([:\[,]\s*)(" + _JSON_BAREWORD + r")") @@ -223,6 +328,14 @@ def _repair_unquoted_json(blob: str) -> str | None: This is a narrowly scoped recovery heuristic for known JSON-valued server flags after shlex damage, not a general parser for JSON-like syntax. + + It is NOT the fix for the composer: :func:`remove_server_args` no longer + damages JSON in the first place (it tokenizes quote-preservingly). This + remains only as a recovery layer for strings that were already persisted in + damaged form by the earlier lossy round trip, or that arrive damaged from + another producer. Never rely on it for newly composed args — a blob is only + repairable when every stripped-quote value happens to be re-quotable, which + is not decidable in general. """ def _quote_value(m: "re.Match[str]") -> str: