fix(executors): honour --ep on single-node servers through the argv seal

src/hyperloom/inference_optimizer/cli/parser.py

@@ -387,10 +387,11 @@ def _build_parser() -> argparse.ArgumentParser:
         default=None,
         help="Expert-parallel size for MoE inference. 1 (default) keeps "
         "experts sharded by TP (legacy behaviour). >=2 enables true "
-        "expert parallelism: sglang adds `--expert-parallel-size N`, "
-        "vllm adds `--enable-expert-parallel`. Typical: EP=TP for "
-        "DSr1/DSv3 on multi-node. Default: 1. "
-        "EP > TP is rejected at server-restart time.",
+        "expert parallelism: sglang adds `--ep-size N`, vllm and atom "
+        "add `--enable-expert-parallel` (EP then spans all TP ranks). "
+        "Typical: EP=TP. Default: 1. On a single node an EP that does "
+        "not divide TP is dropped with a warning; on multi-node EP > TP "
+        "is rejected at server-restart time.",
     )
     opt.add_argument(
         "--pd-mode",

src/hyperloom/inference_optimizer/framework_registry.py

@@ -82,6 +82,15 @@ class FrameworkSpec:
 
 DEFAULT_FRAMEWORK = "sglang"
 
+# Server option that turns on expert parallelism: (flag to write, every
+# argparse alias of it, whether it takes the EP size). vLLM and ATOM take a
+# switch and derive the size from TP x DP. Frameworks absent here have none.
+EXPERT_PARALLEL_FLAGS: dict[str, tuple[str, tuple[str, ...], bool]] = {
+    "sglang": ("--ep-size", ("--ep-size", "--expert-parallel-size", "--ep"), True),
+    "vllm": ("--enable-expert-parallel", ("--enable-expert-parallel", "-ep"), False),
+    "atom": ("--enable-expert-parallel", ("--enable-expert-parallel",), False),
+}
+
 
 def names() -> tuple[str, ...]:
     """Return the canonical tuple of supported framework names."""
@@ -150,6 +159,15 @@ def server_args_env_name(framework: str | None) -> str:
     return extra_args_env(DEFAULT_FRAMEWORK)
 
 
+def expert_parallel_flag(framework: str | None) -> tuple[str, tuple[str, ...], bool] | None:
+    """Return the :data:`EXPERT_PARALLEL_FLAGS` row for the server reading ``framework``'s server args."""
+    env_name = server_args_env_name(framework)
+    return next(
+        (row for name, row in EXPERT_PARALLEL_FLAGS.items() if FRAMEWORKS[name].extra_args_env == env_name),
+        None,
+    )
+
+
 def throughput_unit(framework: str | None) -> str:
     """Return the throughput unit string for ``framework``."""
     return _spec_or_default(framework).throughput_unit

src/hyperloom/inference_optimizer/tests/test_server_argv_seal.py

@@ -13,7 +13,9 @@
 
 from hyperloom.inference_optimizer.framework_registry import names, server_args_env_name
 from hyperloom.orchestrator.actions.executors import _grid_runner, _workload_envs
+from hyperloom.orchestrator.actions.executors._grid_base import GridVariant
 from hyperloom.orchestrator.actions.executors._server_argv import (
+    ServerArgv,
     add_server_arg_unless_pinned,
     config_launch_env,
     config_server_argv,
@@ -269,3 +271,143 @@ def test_a_repaired_argv_reaches_the_launch_through_the_rendered_config(tmp_path
     assert resealed.argv == ("--tp", "8")
     # Read back off disk, because the file is what the launch opens.
     assert config_server_argv(config).argv == ("--tp", "8")
+
+
+def _seal_with_ep(framework: str, args: str, ep: object) -> ServerArgv:
+    """Seal ``args`` for ``framework`` with ``ep`` in the envs; ``None`` leaves EP unset."""
+    envs: dict[str, object] = {server_args_env_name(framework): args}
+    if ep is not None:
+        envs["EP"] = ep
+    return seal_server_argv(envs, framework)
+
+
+@pytest.mark.parametrize(
+    ("framework", "args", "ep", "expected"),
+    [
+        ("sglang", "--tp 8", 4, ("--tp", "8", "--ep-size", "4")),
+        ("sglang", "", "8", ("--ep-size", "8")),
+        ("sglang", "--ep-size 4", 4, ("--ep-size", "4")),
+        ("sglang", "--expert-parallel-size=4", 4, ("--expert-parallel-size=4",)),
+        ("sglang", "--ep 4", "4", ("--ep", "4")),
+        # A flag that merely starts with an alias is a different option.
+        ("sglang", "--ep-num-redundant-experts 32", 4, ("--ep-num-redundant-experts", "32", "--ep-size", "4")),
+        ("vllm", "--foo", 8, ("--foo", "--enable-expert-parallel")),
+        ("vllm", "", 8, ("--enable-expert-parallel",)),
+        ("vllm", "--enable-expert-parallel", 8, ("--enable-expert-parallel",)),
+        ("vllm", "-ep", 8, ("-ep",)),
+        ("atom", "--level 2", 8, ("--level", "2", "--enable-expert-parallel")),
+        ("vllm", "--foo", 1, ("--foo",)),
+        ("vllm", "--foo", None, ("--foo",)),
+        # EP 1 is a no-op, so a recipe that pins its own EP size is left standing.
+        ("sglang", "--ep-size 4", 1, ("--ep-size", "4")),
+        ("xdit", "--foo", 8, ("--foo",)),
+    ],
+)
+def test_the_seal_hands_the_server_its_expert_parallel_flag(framework, args, ep, expected):
+    """EP > 1 reaches the server's argv exactly once, in the framework's own spelling."""
+    assert _seal_with_ep(framework, args, ep).argv == expected
+
+
+@pytest.mark.parametrize("pinned", ["--ep-size 2", "--expert-parallel-size=8", "--ep 1", "--ep-size"])
+def test_an_expert_parallel_pin_that_contradicts_ep_fails_the_seal(pinned):
+    """Neither the operator's pin nor the CLI's EP may win silently."""
+    with pytest.raises(ValueError, match="EP=4"):
+        _seal_with_ep("sglang", pinned, 4)
+
+
+def test_a_contradicting_pin_fails_the_seal_even_when_ep_does_not_divide_tp():
+    """The skip for an EP that TP cannot hold must not let a different pinned size launch unreported."""
+    with pytest.raises(ValueError, match="EP=3"):
+        seal_server_argv({"EXTRA_SGLANG_ARGS": "--ep-size 2", "EP": 3, "TP": 8}, "sglang")
+
+
+def test_sealing_twice_adds_the_expert_parallel_flag_once():
+    """The seal reads its own output back as a pin."""
+    envs = {"EXTRA_SGLANG_ARGS": "--tp 8", "EP": 8}
+    seal_server_argv(envs, "sglang")
+
+    assert seal_server_argv(envs, "sglang").argv == ("--tp", "8", "--ep-size", "8")
+
+
+def test_a_resealed_argv_keeps_the_expert_parallel_flag(tmp_path):
+    """A repair rewrites the argument string, not the run's EP."""
+    config = tmp_path / "bench.yaml"
+    config.write_text(
+        yaml.safe_dump(
+            {"benchmark": {"framework": "sglang", "model": "/m", "envs": {"EP": 4, "EXTRA_SGLANG_ARGS": "--bogus x"}}}
+        ),
+        encoding="utf-8",
+    )
+
+    assert reseal_config_argv(config, "--tp 4").argv == ("--tp", "4", "--ep-size", "4")
+    assert config_server_argv(config).argv == ("--tp", "4", "--ep-size", "4")
+
+
+@pytest.mark.parametrize(("variant_mode", "base_mode"), [("replace", "append"), ("append", "replace")])
+def test_a_replacing_grid_variant_keeps_the_expert_parallel_flag(tmp_path, variant_mode, base_mode):
+    """Replacing the argument string does not replace the run's EP."""
+    base = tmp_path / "base.yaml"
+    base.write_text(
+        yaml.safe_dump(
+            {
+                "benchmark": {
+                    "framework": "vllm",
+                    "model": "/m",
+                    "envs": {"TP": 8, "EP": 8, "EXTRA_VLLM_ARGS": "--enable-expert-parallel --block-size 64"},
+                }
+            }
+        ),
+        encoding="utf-8",
+    )
+    variant = GridVariant(name="v", extra_server_args="--max-num-seqs 64", args_mode=variant_mode)
+
+    config = _grid_runner._build_variant_yaml(
+        base, "--max-num-seqs 32", variant, output_subdir=tmp_path / "v", base_args_mode=base_mode
+    )
+
+    argv = config_server_argv(config).argv
+    assert "--block-size" not in argv
+    assert argv.count("--enable-expert-parallel") == 1
+
+
+@pytest.mark.parametrize(
+    ("nodes", "tp", "ep"),
+    [("1", 8, 3), ("2", 16, 16)],
+    ids=["ep-does-not-divide-tp", "multi-node"],
+)
+def test_a_resealed_argv_gets_no_expert_parallel_flag_the_launch_cannot_take(tmp_path, monkeypatch, nodes, tp, ep):
+    """A repair goes through the same EP rule as the first seal."""
+    monkeypatch.setenv("INFERENCE_OPTIMIZER_NODES", nodes)
+    config = tmp_path / "bench.yaml"
+    config.write_text(
+        yaml.safe_dump(
+            {"benchmark": {"framework": "sglang", "model": "/m", "envs": {"TP": tp, "EP": ep, "EXTRA_SGLANG_ARGS": ""}}}
+        ),
+        encoding="utf-8",
+    )
+
+    assert reseal_config_argv(config, "--mem-fraction-static 0.9").argv == ("--mem-fraction-static", "0.9")
+
+
+def test_a_grid_variant_ep_that_does_not_divide_tp_adds_no_flag(tmp_path, caplog):
+    """A variant's extra_envs reach the seal unfiltered, so the seal is where EP meets TP."""
+    base = tmp_path / "base.yaml"
+    base.write_text(
+        yaml.safe_dump(
+            {
+                "benchmark": {
+                    "framework": "sglang",
+                    "model": "/m",
+                    "envs": {"TP": 8, "EXTRA_SGLANG_ARGS": "--mem-fraction-static 0.9"},
+                }
+            }
+        ),
+        encoding="utf-8",
+    )
+    variant = GridVariant(name="v", extra_envs={"EP": "3"})
+
+    with caplog.at_level("WARNING"):
+        config = _grid_runner._build_variant_yaml(base, "", variant, output_subdir=tmp_path / "v")
+
+    assert "--ep-size" not in config_server_argv(config).argv
+    assert "EP=3 does not divide the launched TP=8" in caplog.text

src/hyperloom/inference_optimizer/tests/test_workload_envs.py

@@ -5,12 +5,15 @@
 
 from __future__ import annotations
 
+import json
 import sys
 from types import SimpleNamespace
 
 import pytest
 import yaml
 
+from hyperloom.common.env_safety import OPERATOR_EXTRA_ENV_VAR
+from hyperloom.inference_optimizer.framework_registry import server_args_env_name
 from hyperloom.orchestrator.actions.executors import _workload_envs as we
 
 
@@ -272,79 +275,118 @@
     assert "chunked-prefill-size" in merged
 
 
-def test_vllm_ep_injects_expert_parallel(monkeypatch, tmp_path):
+def _ep_materialize(monkeypatch, tmp_path, *, framework, args, tp, ep):
+    """Materialize ``framework`` at ``tp``/``ep`` and return its benchmark envs."""
     _clear_env(monkeypatch)
     monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1")
-    monkeypatch.setenv("EP", "8")
-    src = _write(tmp_path / "cfg.yaml", framework="vllm", envs={"EXTRA_VLLM_ARGS": "--trust-remote-code"})
+    monkeypatch.setenv("TP", str(tp))
+    monkeypatch.setenv("EP", str(ep))
+    src = _write(tmp_path / "cfg.yaml", framework=framework, envs={server_args_env_name(framework): args})
+    return _materialize(src, tmp_path / "out")["envs"]
 
-    bench = _materialize(src, tmp_path / "out")
 
-    args = bench["envs"]["EXTRA_VLLM_ARGS"]
-    assert "--trust-remote-code" in args
-    assert "--enable-expert-parallel" in args
+def test_materialize_hands_sglang_its_expert_parallel_size(monkeypatch, tmp_path):
+    envs = _ep_materialize(monkeypatch, tmp_path, framework="sglang", args="--mem-fraction-static 0.9", tp=8, ep=4)
 
+    argv = envs["EXTRA_SGLANG_ARGS"].split()
+    assert envs["EP"] == 4
+    assert "--mem-fraction-static" in argv
+    assert argv.count("--ep-size") == 1
+    assert argv[argv.index("--ep-size") + 1] == "4"
 
-def test_vllm_ep_one_does_not_inject_expert_parallel(monkeypatch, tmp_path):
-    _clear_env(monkeypatch)
-    monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1")
-    monkeypatch.setenv("EP", "1")
-    src = _write(tmp_path / "cfg.yaml", framework="vllm", envs={"EXTRA_VLLM_ARGS": "--trust-remote-code"})
 
-    bench = _materialize(src, tmp_path / "out")
+@pytest.mark.parametrize("pinned", ["--ep-size 4", "--expert-parallel-size=4"])
+def test_materialize_keeps_an_operator_pinned_expert_parallel_size(monkeypatch, tmp_path, pinned):
+    monkeypatch.setenv("INFERENCE_OPTIMIZER_SERVER_ARGS", pinned)
 
-    args = bench["envs"]["EXTRA_VLLM_ARGS"]
-    assert "--trust-remote-code" in args
-    assert "--enable-expert-parallel" not in args
+    envs = _ep_materialize(monkeypatch, tmp_path, framework="sglang", args="", tp=8, ep=4)
 
+    flags = [tok.split("=", 1)[0] for tok in envs["EXTRA_SGLANG_ARGS"].split()]
+    assert sum(flag in {"--ep-size", "--expert-parallel-size", "--ep"} for flag in flags) == 1
 
-def test_vllm_ep_preserves_profile_args(monkeypatch, tmp_path):
-    _clear_env(monkeypatch)
-    monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1")
-    monkeypatch.setenv("EP", "8")
-    src = _write(
-        tmp_path / "cfg.yaml",
-        framework="vllm",
-        envs={"EXTRA_VLLM_ARGS": "--profiler-config.ignore_frontend True"},
-    )
 
-    bench = _materialize(src, tmp_path / "out")
+@pytest.mark.parametrize(
+    "args",
+    ["--trust-remote-code", "--profiler-config.ignore_frontend True", "--enable-expert-parallel"],
+)
+def test_materialize_enables_vllm_expert_parallel_once(monkeypatch, tmp_path, args):
+    envs = _ep_materialize(monkeypatch, tmp_path, framework="vllm", args=args, tp=8, ep=8)
+
+    assert args in envs["EXTRA_VLLM_ARGS"]
+    assert envs["EXTRA_VLLM_ARGS"].split().count("--enable-expert-parallel") == 1
+
+
+def test_materialize_ep_one_adds_no_expert_parallel_flag(monkeypatch, tmp_path):
+    envs = _ep_materialize(monkeypatch, tmp_path, framework="vllm", args="--trust-remote-code", tp=8, ep=1)
 
-    args = bench["envs"]["EXTRA_VLLM_ARGS"]
-    assert "--profiler-config.ignore_frontend True" in args
-    assert "--enable-expert-parallel" in args
+    assert envs["EXTRA_VLLM_ARGS"] == "--trust-remote-code"
 
 
-def test_vllm_ep_does_not_duplicate_existing_flag(monkeypatch, tmp_path):
+def test_materialize_drops_an_ep_the_clamped_tp_cannot_hold(monkeypatch, tmp_path, caplog):
     _clear_env(monkeypatch)
-    monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1")
+    monkeypatch.setenv("INFERENCE_OPTIMIZER_VISIBLE_GPU_COUNT", "4")
+    monkeypatch.setenv("TP", "8")
     monkeypatch.setenv("EP", "8")
-    src = _write(
-        tmp_path / "cfg.yaml",
-        framework="vllm",
-        envs={"EXTRA_VLLM_ARGS": "--enable-expert-parallel"},
-    )
+    src = _write(tmp_path / "cfg.yaml", envs={"EXTRA_SGLANG_ARGS": "--mem-fraction-static 0.9"})
 
-    bench = _materialize(src, tmp_path / "out")
+    with caplog.at_level("WARNING"):
+        envs = _materialize(src, tmp_path / "out")["envs"]
+
+    assert envs["TP"] == 4
+    assert envs["EP"] == 8
+    assert "--ep-size" not in envs["EXTRA_SGLANG_ARGS"]
+    assert "EP=8 does not divide the launched TP=4" in caplog.text
 
-    args = bench["envs"]["EXTRA_VLLM_ARGS"]
-    assert args.count("--enable-expert-parallel") == 1
+
+def test_materialize_drops_an_ep_that_does_not_divide_tp(monkeypatch, tmp_path, caplog):
+    with caplog.at_level("WARNING"):
+        envs = _ep_materialize(monkeypatch, tmp_path, framework="sglang", args="--mem-fraction-static 0.9", tp=8, ep=3)
+
+    assert envs["EP"] == 3
+    assert "--ep-size" not in envs["EXTRA_SGLANG_ARGS"]
+    assert "EP=3 does not divide the launched TP=8" in caplog.text
 
 
 @pytest.mark.parametrize(
-    "server_args, framework, ep, expected",
+    ("environ", "extra_envs"),
     [
-        ("--foo", "vllm", 8, "--foo --enable-expert-parallel"),
-        ("--foo", "vllm", 1, "--foo"),
-        ("--foo", "sglang", 8, "--foo"),
-        ("--foo", "vllm", "bad", "--foo"),
-        ("--foo", "vllm", None, "--foo"),
-        ("--enable-expert-parallel", "vllm", 8, "--enable-expert-parallel"),
-        ("", "vllm", 8, "--enable-expert-parallel"),
+        # ``--extra-env EP=3`` serializes the pin and also exports it under its own name.
+        ({OPERATOR_EXTRA_ENV_VAR: json.dumps({"EP": "3"}), "EP": "3"}, None),
+        ({}, {"EP": "3"}),
     ],
+    ids=["operator-pin", "extra-envs"],
 )
-def test_inject_vllm_expert_parallel_unit(server_args, framework, ep, expected):
-    assert we.inject_vllm_expert_parallel(server_args, framework, ep) == expected
+def test_an_extra_env_ep_that_does_not_divide_tp_adds_no_flag(monkeypatch, tmp_path, caplog, environ, extra_envs):
+    """An EP written by the extra-env merge meets the same TP rule as the run's own EP."""
+    _clear_env(monkeypatch)
+    monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1")
+    monkeypatch.setenv("TP", "8")
+    for name, value in environ.items():
+        monkeypatch.setenv(name, value)
+    src = _write(tmp_path / "cfg.yaml", envs={"EXTRA_SGLANG_ARGS": "--mem-fraction-static 0.9"})
+
+    with caplog.at_level("WARNING"):
+        envs = _materialize(src, tmp_path / "out", extra_envs=extra_envs)["envs"]
+
+    assert "--ep-size" not in envs["EXTRA_SGLANG_ARGS"].split()
+    assert "EP=3 does not divide the launched TP=8" in caplog.text
+
+
+@pytest.mark.parametrize("operator_pin", [False, True])
+def test_materialize_leaves_multi_node_ep_to_the_launcher(monkeypatch, tmp_path, operator_pin):
+    """The multi-node launcher passes EP itself; this pod's GPU count is not the cluster's TP."""
+    _clear_env(monkeypatch)
+    monkeypatch.setenv("INFERENCE_OPTIMIZER_NODES", "2")
+    monkeypatch.setenv("INFERENCE_OPTIMIZER_VISIBLE_GPU_COUNT", "8")
+    monkeypatch.setenv("TP", "16")
+    monkeypatch.setenv("EP", "16")
+    if operator_pin:
+        monkeypatch.setenv(OPERATOR_EXTRA_ENV_VAR, json.dumps({"EP": "16"}))
+    src = _write(tmp_path / "cfg.yaml", envs={"EXTRA_SGLANG_ARGS": "--mem-fraction-static 0.9"})
+
+    envs = _materialize(src, tmp_path / "out")["envs"]
+
+    assert "--ep-size" not in envs["EXTRA_SGLANG_ARGS"].split()
 
 
 def test_mimo_v2_injects_triton_attention(monkeypatch, tmp_path):
@@ -508,7 +550,7 @@
 
 def _profile_num_steps(bench) -> int:
     """Captured-step count the sglang profile path writes into PROFILE_EXTRA_BODY."""
     import json
 
     return int(json.loads(bench["envs"]["PROFILE_EXTRA_BODY"])["num_steps"])
 
@@ -773,7 +815,7 @@
 
 
 def _write_model(tmp_path, model_type):
     import json
 
     d = tmp_path / "m"
     d.mkdir(exist_ok=True)