fix(server-args): keep +custom_ops JSON valid across variant composition

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)

src/hyperloom/orchestrator/actions/executors/_grid_server_args.py

@@ -100,29 +100,74 @@ 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.
 
     ``remove_args`` entries are flag-oriented. ``"--foo"`` removes ``--foo`` and
     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 "<missing>",
+            )
+
+
+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: