test(ci): wire P-EAGLE and Domino into the example test runner

ci/run_example_test.sh

@@ -15,6 +15,9 @@ case "${platform}/${backend}/${drafter}" in
   gpu/vllm/dspark)
     example="examples/run_qwen3-8b_drafter_dspark_vllm.sh"
     ;;
+  gpu/vllm/peagle|gpu/vllm/domino)
+    example="examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh"
+    ;;
   gpu/sglang/eagle3)
     example="examples/run_qwen3-8b_drafter_eagle3_sglang.sh"
     ;;
@@ -40,7 +43,7 @@ case "${platform}/${backend}/${drafter}" in
     example="examples/run_qwen3-8b_drafter_dflash_sglang.sh"
     ;;
   *)
-    echo "usage: $0 {gpu|npu} {vllm|sglang} {eagle3|megatron-eagle3|dflash|dspark}" >&2
+    echo "usage: $0 {gpu|npu} {vllm|sglang} {eagle3|megatron-eagle3|dflash|dspark|peagle|domino}" >&2
     exit 2
     ;;
 esac
@@ -58,6 +61,15 @@ for name in "${required_vars[@]}"; do
   fi
 done
 
+# Domino and P-EAGLE are not engine-level speculative algorithms, so neither can
+# be trained inside the rollout loop. They use the two-stage separate-training
+# workflow instead: stage 1 rolls out with the engine algorithm whose
+# hidden-state layout the drafter consumes and writes a feature store, stage 2
+# trains the drafter offline from that same store.
+# The example owns the collect-to-train algorithm pairing; the runner only needs
+# the collect side of it to pick the matching cached draft model.
+separate_training="false"
+
 case "${drafter}" in
   eagle3|megatron-eagle3)
     draft_model="${SPECO_EAGLE3_DRAFT_MODEL:-}"
@@ -71,9 +83,19 @@ case "${drafter}" in
     draft_model="${SPECO_DSPARK_DRAFT_MODEL:-}"
     draft_algorithm="DSPARK"
     ;;
+  peagle)
+    draft_model="${SPECO_EAGLE3_DRAFT_MODEL:-}"
+    draft_algorithm="EAGLE3"
+    separate_training="true"
+    ;;
+  domino)
+    draft_model="${SPECO_DFLASH_DRAFT_MODEL:-}"
+    draft_algorithm="DFLASH"
+    separate_training="true"
+    ;;
 esac
 if [[ -z "${draft_model}" ]]; then
-  echo "required ${drafter} draft model environment variable is not set" >&2
+  echo "required ${drafter} collect-stage draft model environment variable is not set" >&2
   exit 2
 fi
 
@@ -216,9 +238,62 @@ if [[ "${drafter}" == "dspark" ]]; then
   )
 fi
 
+train_overrides=()
+if [[ "${separate_training}" == "true" ]]; then
+  feature_store_dir="${SPECO_CKPT_DIR}/${drafter}_features"
+  # A draft init path that does not exist yet cold-starts the drafter from the
+  # target config, so this lane needs no extra model in the CI cache.
+  draft_init_path="${SPECO_CKPT_DIR}/${drafter}_draft_init"
+
+  # The example already sets the two-stage shape (collect_only/offline modes,
+  # the standalone launcher, the feature-store flags and the per-algorithm
+  # hyperparameters). Override only what CI has to control: where the stages
+  # meet on disk, and the sizes that keep a smoke run cheap.
+  # A configured SPECO_CKPT_DIR persists between jobs on a self-hosted runner,
+  # so drop any earlier shards: otherwise a collect stage that produced nothing
+  # still trains, and the job goes green on stale features.
+  rm -rf "${feature_store_dir}" "${draft_init_path}"
+
+  overrides+=(
+    "actor_rollout_ref.rollout.drafter.training.feature_store.path=${feature_store_dir}"
+    "actor_rollout_ref.rollout.drafter.training.feature_store.max_samples_per_shard=32"
+  )
+
+  train_overrides=(
+    "actor_rollout_ref.model.path=${SPECO_TARGET_MODEL}"
+    "actor_rollout_ref.rollout.drafter.model_path=${draft_init_path}"
+    "actor_rollout_ref.rollout.drafter.checkpoint_path=${SPECO_CKPT_DIR}/${drafter}_draft_ckpts"
+    "actor_rollout_ref.rollout.drafter.training.max_steps=1"
+    "actor_rollout_ref.rollout.drafter.training.save_interval_steps=1"
+    "actor_rollout_ref.rollout.drafter.training.batch_size_per_gpu=${SPECO_DRAFTER_BATCH_SIZE_PER_GPU:-1}"
+    "actor_rollout_ref.rollout.drafter.training.feature_store.path=${feature_store_dir}"
+  )
+
+  case "${drafter}" in
+    peagle)
+      train_overrides+=(
+        "actor_rollout_ref.rollout.drafter.training.peagle_num_draft_layers=1"
+        "actor_rollout_ref.rollout.drafter.training.peagle_num_depths=2"
+      )
+      ;;
+    domino)
+      train_overrides+=(
+        "actor_rollout_ref.rollout.drafter.training.domino_block_size=4"
+        "actor_rollout_ref.rollout.drafter.training.domino_num_anchors=8"
+        "actor_rollout_ref.rollout.drafter.training.domino_max_window=64"
+        "actor_rollout_ref.rollout.drafter.training.domino_emb_dim=64"
+        "actor_rollout_ref.rollout.drafter.training.domino_gru_hidden_dim=128"
+        "actor_rollout_ref.rollout.drafter.training.domino_lambda_base_decay_steps=10"
+      )
+      ;;
+  esac
+fi
+
 if [[ -n "${SPECO_EXTRA_HYDRA_ARGS:-}" ]]; then
   while IFS= read -r extra_arg; do
     [[ -z "${extra_arg}" ]] && continue
+    # Stage 2 runs a different entrypoint with a different config tree, so these
+    # rollout-oriented overrides stay on the collect stage.
     overrides+=("${extra_arg}")
   done <<< "${SPECO_EXTRA_HYDRA_ARGS}"
 fi
@@ -229,10 +304,26 @@ if [[ "${SPECO_DRY_RUN:-false}" == "true" ]]; then
   echo "drafter=${drafter}"
   echo "example=${example}"
   echo "draft_algorithm=${draft_algorithm}"
+  echo "separate_training=${separate_training}"
   echo "ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-}"
   printf 'Hydra overrides:\n'
   printf '  %q\n' "${overrides[@]}"
+  if [[ "${separate_training}" == "true" ]]; then
+    printf 'Hydra train overrides:\n'
+    printf '  %q\n' "${train_overrides[@]}"
+  fi
   exit 0
 fi
 
-bash "${example}" "${overrides[@]}"
+if [[ "${separate_training}" == "true" ]]; then
+  DRAFT_ALGO="${drafter}" RUN_STAGE=collect bash "${example}" "${overrides[@]}"
+  if [[ "${enable_training}" == "true" ]]; then
+    DRAFT_ALGO="${drafter}" RUN_STAGE=train DRAFT_TRAIN_GPUS_PER_NODE="${accelerator_count}" bash "${example}" "${train_overrides[@]}"
+  else
+    # Generation-only runs capture no hidden states, so the feature store the
+    # offline trainer would read is empty.
+    echo "SPECO_ENABLE_TRAINING is not true; skipping the offline train stage"
+  fi
+else
+  bash "${example}" "${overrides[@]}"
+fi

examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh

@@ -27,7 +27,10 @@ RUN_STAGE=${RUN_STAGE:-both}
 gen_tp=2
 train_sp=1
 ppo_gpus_per_node=8
-draft_train_gpus_per_node=8
+# The standalone launcher resolves its device count from the FIRST matching
+# override, so appending one after "$@" cannot lower it. Read it from the
+# environment instead, which is what a smaller box (or CI) needs.
+draft_train_gpus_per_node=${DRAFT_TRAIN_GPUS_PER_NODE:-8}
 
 MODEL_PATH=/path/to/model
 CKPTS_DIR=/path/to/checkpoint

tests/examples/test_ci_example_runner.py

@@ -13,10 +13,12 @@
 # limitations under the License.
 from __future__ import annotations
 
+import functools
 import os
 import shlex
 import subprocess
 import shutil
+import tempfile
 from pathlib import Path
 
 import pytest
@@ -38,6 +40,7 @@ def _workflow_source(name: str) -> str:
     return (WORKFLOWS / name).read_text(encoding="utf-8")
 
 
+@functools.lru_cache(maxsize=None)
 def _require_working_bash() -> str:
     bash = shutil.which("bash")
     if bash is None:
@@ -58,6 +61,7 @@ def _bash_path(path: Path, bash: str) -> str:
     return path.as_posix()
 
 
+@functools.lru_cache(maxsize=None)
 def _runner_script() -> str:
     return "\n".join(RUNNER.read_text(encoding="utf-8").splitlines()) + "\n"
 
@@ -289,6 +293,15 @@ def test_example_runner_covers_gpu_and_npu_backend_matrix() -> None:
     assert "examples/run_qwen3-8b_drafter_dflash_sglang.sh" in source
     assert "examples/run_qwen3-8b_drafter_dspark_vllm.sh" in source
     assert "examples/run_qwen3-8b_drafter_dspark_vllm_npu.sh" in source
+    assert "gpu/vllm/peagle" in source
+    assert "gpu/vllm/domino" in source
+    assert (
+        "examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh" in source
+    )
+    # P-EAGLE attends through torch flex_attention, which the NPU runtime does
+    # not provide, so the separate-training lane stays GPU-only.
+    assert "npu/vllm/peagle" not in source
+    assert "npu/vllm/domino" not in source
 
 
 def test_example_runner_exposes_required_hydra_overrides() -> None:
@@ -319,28 +332,7 @@ def test_example_runner_exposes_required_hydra_overrides() -> None:
 
 
 def test_example_runner_dry_run_covers_npu_dspark() -> None:
-    bash = _require_working_bash()
-    env = {
-        "SPECO_DRY_RUN": "true",
-        "SPECO_TARGET_MODEL": "/models/target",
-        "SPECO_DSPARK_DRAFT_MODEL": "/models/dspark",
-        "SPECO_TRAIN_FILE": "/data/train.parquet",
-        "SPECO_TEST_FILE": "/data/test.parquet",
-        "SPECO_CKPT_DIR": "/tmp/speco",
-        "SPECO_ACCELERATOR_COUNT": "1",
-    }
-    script = "".join(
-        f"export {name}={shlex.quote(value)}\n" for name, value in env.items()
-    )
-    script += _runner_script()
-    result = subprocess.run(
-        [bash, "-s", "--", "npu", "vllm", "dspark"],
-        env=os.environ.copy(),
-        input=script.encode("utf-8"),
-        capture_output=True,
-        check=True,
-    )
-    stdout = result.stdout.decode("utf-8", errors="replace")
+    stdout = _dry_run("dspark", platform="npu")
 
     assert "example=examples/run_qwen3-8b_drafter_dspark_vllm_npu.sh" in stdout
     assert "draft_algorithm=DSPARK" in stdout
@@ -350,29 +342,148 @@ def test_example_runner_dry_run_covers_npu_dspark() -> None:
 
 
 def test_example_runner_dry_run_omits_ulysses_overrides_for_npu_megatron() -> None:
+    stdout = _dry_run("megatron-eagle3", platform="npu", accelerator_count="8")
+
+    assert "example=examples/run_qwen3-4b_actor_megatron_drafter_eagle3_vllm_npu.sh" in stdout
+    assert "draft_algorithm=EAGLE3" in stdout
+    assert "ulysses_sequence_parallel_size" not in stdout
+
+
+@functools.lru_cache(maxsize=None)
+def _dry_run(
+    drafter: str,
+    platform: str = "gpu",
+    backend: str = "vllm",
+    accelerator_count: str = "1",
+    extra_hydra_args: str | None = None,
+) -> str:
+    """Dump the Hydra overrides the runner would pass, without launching a job."""
     bash = _require_working_bash()
     env = {
         "SPECO_DRY_RUN": "true",
         "SPECO_TARGET_MODEL": "/models/target",
         "SPECO_EAGLE3_DRAFT_MODEL": "/models/eagle3",
+        "SPECO_DFLASH_DRAFT_MODEL": "/models/dflash",
+        "SPECO_DSPARK_DRAFT_MODEL": "/models/dspark",
         "SPECO_TRAIN_FILE": "/data/train.parquet",
         "SPECO_TEST_FILE": "/data/test.parquet",
         "SPECO_CKPT_DIR": "/tmp/speco",
-        "SPECO_ACCELERATOR_COUNT": "8",
+        "SPECO_ACCELERATOR_COUNT": accelerator_count,
     }
+    if extra_hydra_args is not None:
+        env["SPECO_EXTRA_HYDRA_ARGS"] = extra_hydra_args
     script = "".join(
         f"export {name}={shlex.quote(value)}\n" for name, value in env.items()
     )
     script += _runner_script()
-    result = subprocess.run(
-        [bash, "-s", "--", "npu", "vllm", "megatron-eagle3"],
-        env=os.environ.copy(),
-        input=script.encode("utf-8"),
-        capture_output=True,
-        check=True,
+
+    # Run from a file rather than piping the script into `bash -s`: the runner
+    # exits early in dry-run mode, and on a script this size that races the
+    # writer filling the stdin pipe and can deadlock.
+    with tempfile.TemporaryDirectory() as tmp:
+        entry = Path(tmp) / "dry_run.sh"
+        entry.write_text(script, encoding="utf-8")
+        result = subprocess.run(
+            [bash, _bash_path(entry, bash), platform, backend, drafter],
+            cwd=ROOT,
+            env=os.environ.copy(),
+            stdin=subprocess.DEVNULL,
+            capture_output=True,
+            check=True,
+        )
+    return result.stdout.decode("utf-8", errors="replace")
+
+
+def _sections(stdout: str) -> tuple[str, str]:
+    """Split the dry-run dump into its collect and offline-train override blocks."""
+    marker = "Hydra train overrides:"
+    assert marker in stdout, stdout
+    collect, _, train = stdout.partition(marker)
+    _, _, collect = collect.partition("Hydra overrides:")
+    return collect, train
+
+
+@pytest.mark.parametrize(
+    ("drafter", "collect_algorithm", "shrunk_key"),
+    (
+        ("peagle", "EAGLE3", "training.peagle_num_depths=2"),
+        ("domino", "DFLASH", "training.domino_block_size=4"),
+    ),
+)
+def test_example_runner_dry_run_covers_separate_training_stages(
+    drafter: str, collect_algorithm: str, shrunk_key: str
+) -> None:
+    stdout = _dry_run(drafter)
+
+    assert (
+        "example=examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh"
+        in stdout
     )
-    stdout = result.stdout.decode("utf-8", errors="replace")
+    assert "separate_training=true" in stdout
+    # Stage 1 rolls out under the engine-servable algorithm whose hidden-state
+    # layout this drafter consumes, not under the drafter's own algorithm.
+    assert f"draft_algorithm={collect_algorithm}" in stdout
 
-    assert "example=examples/run_qwen3-4b_actor_megatron_drafter_eagle3_vllm_npu.sh" in stdout
-    assert "draft_algorithm=EAGLE3" in stdout
-    assert "ulysses_sequence_parallel_size" not in stdout
+    collect, train = _sections(stdout)
+
+    feature_store = (
+        "actor_rollout_ref.rollout.drafter.training.feature_store.path="
+        f"/tmp/speco/{drafter}_features"
+    )
+    assert (
+        "actor_rollout_ref.rollout.drafter.speculative_algorithm="
+        f"{collect_algorithm}" in collect
+    )
+    assert feature_store in collect
+
+    # Stage 2 reads that same store through the standalone offline trainer.
+    assert feature_store in train
+    assert (
+        f"actor_rollout_ref.rollout.drafter.model_path=/tmp/speco/{drafter}_draft_init"
+        in train
+    )
+    assert (
+        "actor_rollout_ref.rollout.drafter.checkpoint_path="
+        f"/tmp/speco/{drafter}_draft_ckpts" in train
+    )
+    assert "actor_rollout_ref.rollout.drafter.training.max_steps=1" in train
+    assert shrunk_key in train
+
+    # The standalone launcher takes the first matching override, so the device
+    # count has to reach the example through the environment instead.
+    assert "speco.draft_training.num_gpus_per_node" not in train
+    assert "DRAFT_TRAIN_GPUS_PER_NODE" in _runner_script()
+
+
+def test_example_runner_dry_run_keeps_algorithm_knobs_out_of_the_other_lane() -> None:
+    peagle = _sections(_dry_run("peagle"))[1]
+    domino = _sections(_dry_run("domino"))[1]
+
+    assert "domino_" not in peagle
+    assert "peagle_" not in domino
+
+
+def test_example_runner_dry_run_keeps_extra_hydra_args_on_the_collect_stage() -> None:
+    collect, train = _sections(_dry_run("peagle", extra_hydra_args="trainer.nnodes=1"))
+
+    assert "trainer.nnodes=1" in collect
+    # Stage 2 is a different entrypoint with a different config tree.
+    assert "trainer.nnodes=1" not in train
+
+
+def test_example_runner_skips_the_offline_stage_without_training() -> None:
+    source = _runner_script()
+
+    assert 'if [[ "${enable_training}" == "true" ]]; then' in source
+    assert "skipping the offline train stage" in source
+    assert 'rm -rf "${feature_store_dir}"' in source
+
+
+def test_separate_training_example_takes_its_device_count_from_the_env() -> None:
+    example = (
+        ROOT
+        / "examples"
+        / "run_qwen3-8b_drafter_domino_peagle_separate_training.sh"
+    ).read_text(encoding="utf-8")
+
+    assert "draft_train_gpus_per_node=${DRAFT_TRAIN_GPUS_PER_NODE:-8}" in example