Add e2e smoke tests for the new datagen system

scripts/launch_vllm.py

@@ -1,6 +1,7 @@
 import argparse
 import json
 import os
+import sys
 
 
 def parse_args():
@@ -34,15 +35,13 @@ def parse_args():
         action="store_true",
         help="Print the command that would be executed without running it",
     )
-    parser.add_argument(
-        "vllm_args", nargs=argparse.REMAINDER, help="Arguments to be passed to vLLM"
-    )
-
-    return parser.parse_args()
+    return parser.parse_known_args()
 
 
 def main():
-    args = parse_args()
+    args, vllm_args = parse_args()
+    if "--" in vllm_args:
+        vllm_args.remove("--")
 
     if args.layers:
         layers = args.layers
@@ -71,14 +70,16 @@ def main():
     }
 
     cmd = [
-        "vllm",
+        sys.executable,
+        "-m",
+        "vllm.entrypoints.cli.main",
         "serve",
         args.model,
         "--speculative_config",
         json.dumps(speculative_config),
         "--kv_transfer_config",
         json.dumps(kv_transfer_config),
-        *args.vllm_args,
+        *vllm_args,
     ]
 
     print("Running command:")

scripts/train.py

@@ -257,10 +257,15 @@ def main(args: argparse.Namespace):
         )
         train_files, val_files = split_files(args.data_path, ratio=0.9)
         train_dataset: BaseEagle3Dataset = Eagle3SampleFileDataset(
-            file_list=train_files, max_len=args.total_seq_len, transform=noise_transform
+            file_list=train_files,
+            max_len=args.total_seq_len,
+            transform=noise_transform,
+            hidden_states_dtype=hidden_states_dtype,
         )
         val_dataset: BaseEagle3Dataset = Eagle3SampleFileDataset(
-            file_list=val_files, max_len=args.total_seq_len
+            file_list=val_files,
+            max_len=args.total_seq_len,
+            hidden_states_dtype=hidden_states_dtype,
         )
     else:
         train_dataset = Eagle3ArrowDataset(
@@ -273,6 +278,7 @@ def main(args: argparse.Namespace):
             transform=noise_transform,
             split_ratio=0.9,
             model=args.verifier_name_or_path,
+            hidden_states_dtype=hidden_states_dtype,
         )
         val_dataset = Eagle3ArrowDataset(
             datapath=args.data_path,
@@ -283,6 +289,7 @@ def main(args: argparse.Namespace):
             on_generate=args.on_generate,
             split_ratio=-0.1,
             model=args.verifier_name_or_path,
+            hidden_states_dtype=hidden_states_dtype,
         )
 
     train_loader = setup_dataloader(

tests/e2e/vllm/test_offline_training.py

@@ -0,0 +1,117 @@
+"""E2E test for the offline training workflow.
+
+Exercises the full offline pipeline:
+  1. Prepare data (scripts/prepare_data.py)
+  2. Launch a vLLM server for hidden-state extraction (scripts/launch_vllm.py)
+  3. Generate hidden states offline (scripts/data_generation_offline2.py)
+  4. Stop the vLLM server
+  5. Train a draft model using pre-generated hidden states (scripts/train.py)
+  6. Validate the trained checkpoint via vLLM inference (run_vllm_engine)
+"""
+
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+from loguru import logger
+
+from tests.e2e.vllm.utils import (
+    SCRIPTS_DIR,
+    launch_vllm_server,
+    prepare_data,
+    run_vllm_engine,
+    stop_vllm_server,
+)
+
+MODEL = "Qwen/Qwen3-0.6B"
+VLLM_PORT = 8322
+
+
+@pytest.fixture
+def vllm_server(tmp_path):
+    """Launch a vLLM server configured for hidden-state extraction."""
+    hidden_states_path = str(tmp_path / "hidden_states")
+    process = launch_vllm_server(MODEL, VLLM_PORT, hidden_states_path)
+
+    yield {
+        "port": VLLM_PORT,
+        "hidden_states_path": hidden_states_path,
+        "process": process,
+    }
+
+    stop_vllm_server(process)
+
+
+@pytest.mark.e2e
+@pytest.mark.slow
+def test_offline_training(
+    tmp_path: Path, prompts: list[list[dict[str, str]]], vllm_server
+):
+    data_path = tmp_path / "data"
+    hidden_states_path = tmp_path / "offline_hidden_states"
+    save_path = tmp_path / "checkpoints"
+    port = vllm_server["port"]
+
+    # Step 1: Prepare data
+    prepare_data(MODEL, data_path)
+
+    # Step 2: Generate hidden states offline
+    datagen_cmd = [
+        sys.executable,
+        str(SCRIPTS_DIR / "data_generation_offline2.py"),
+        "--preprocessed-data",
+        str(data_path),
+        "--endpoint",
+        f"http://localhost:{port}/v1",
+        "--output",
+        str(hidden_states_path),
+        "--max-samples",
+        "50",
+        "--concurrency",
+        "4",
+        "--validate-outputs",
+    ]
+    logger.info("Generating hidden states offline: {}", " ".join(datagen_cmd))
+    result = subprocess.run(  # noqa: S603
+        datagen_cmd, stderr=subprocess.PIPE, text=True, check=False
+    )
+    assert result.returncode == 0, (
+        f"data_generation_offline2.py failed:\n{result.stderr}"
+    )
+
+    # Step 3: Stop the vLLM server to free GPU memory before training
+    stop_vllm_server(vllm_server["process"])
+
+    # Step 4: Train using pre-generated hidden states (no live server needed)
+    train_cmd = [
+        sys.executable,
+        str(SCRIPTS_DIR / "train.py"),
+        "--verifier-name-or-path",
+        MODEL,
+        "--data-path",
+        str(data_path),
+        "--hidden-states-path",
+        str(hidden_states_path),
+        "--save-path",
+        str(save_path),
+        "--draft-vocab-size",
+        "8192",
+        "--epochs",
+        "1",
+        "--lr",
+        "3e-4",
+        "--total-seq-len",
+        "512",
+        "--on-missing",
+        "raise",
+    ]
+    logger.info("Running training: {}", " ".join(train_cmd))
+    result = subprocess.run(  # noqa: S603
+        train_cmd, stderr=subprocess.PIPE, text=True, check=False
+    )
+    assert result.returncode == 0, f"train.py failed:\n{result.stderr}"
+
+    # Step 5: Validate trained checkpoint with vLLM inference
+    checkpoint_path = str(save_path / "0")
+    run_vllm_engine(model_path=checkpoint_path, tmp_path=tmp_path, prompts=prompts)

tests/e2e/vllm/test_online_training.py

@@ -0,0 +1,92 @@
+"""E2E test for the online training workflow.
+
+Exercises the full pipeline documented in examples/ONLINE_TRAINING.md:
+  1. Prepare data (scripts/prepare_data.py)
+  2. Launch a vLLM server for hidden-state extraction (scripts/launch_vllm.py)
+  3. Train a draft model against the live server (scripts/train.py)
+  4. Validate the trained checkpoint via vLLM inference (run_vllm_engine)
+"""
+
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+from loguru import logger
+
+from tests.e2e.vllm.utils import (
+    SCRIPTS_DIR,
+    launch_vllm_server,
+    prepare_data,
+    run_vllm_engine,
+    stop_vllm_server,
+)
+
+MODEL = "Qwen/Qwen3-0.6B"
+VLLM_PORT = 8321
+
+
+@pytest.fixture
+def vllm_server(tmp_path):
+    """Launch a vLLM server configured for hidden-state extraction."""
+    hidden_states_path = str(tmp_path / "hidden_states")
+    process = launch_vllm_server(MODEL, VLLM_PORT, hidden_states_path)
+
+    yield {
+        "port": VLLM_PORT,
+        "hidden_states_path": hidden_states_path,
+        "process": process,
+    }
+
+    stop_vllm_server(process)
+
+
+@pytest.mark.e2e
+@pytest.mark.slow
+def test_online_training(
+    tmp_path: Path, prompts: list[list[dict[str, str]]], vllm_server
+):
+    data_path = tmp_path / "data"
+    save_path = tmp_path / "checkpoints"
+    port = vllm_server["port"]
+
+    # Step 1: Prepare data
+    prepare_data(MODEL, data_path)
+
+    # Step 2: Train against live vLLM server
+    train_cmd = [
+        sys.executable,
+        str(SCRIPTS_DIR / "train.py"),
+        "--verifier-name-or-path",
+        MODEL,
+        "--data-path",
+        str(data_path),
+        "--vllm-endpoint",
+        f"http://localhost:{port}/v1",
+        "--save-path",
+        str(save_path),
+        "--draft-vocab-size",
+        "8192",
+        "--epochs",
+        "1",
+        "--lr",
+        "3e-4",
+        "--total-seq-len",
+        "512",
+        "--on-missing",
+        "generate",
+        "--on-generate",
+        "delete",
+    ]
+    logger.info("Running training: {}", " ".join(train_cmd))
+    result = subprocess.run(  # noqa: S603
+        train_cmd, stderr=subprocess.PIPE, text=True, check=False
+    )
+    assert result.returncode == 0, f"train.py failed:\n{result.stderr}"
+
+    # Stop the vLLM server to free GPU memory before running inference
+    stop_vllm_server(vllm_server["process"])
+
+    # Step 3: Validate trained checkpoint with vLLM inference
+    checkpoint_path = str(save_path / "0")
+    run_vllm_engine(model_path=checkpoint_path, tmp_path=tmp_path, prompts=prompts)