From df8499c670a65aa503bb91cbf4f2fd3e5bdafb55 Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Tue, 31 Mar 2026 18:58:49 +0000 Subject: [PATCH 1/6] Update launch_vllm.py to use calling python env Signed-off-by: Fynn Schmitt-Ulms --- scripts/launch_vllm.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/scripts/launch_vllm.py b/scripts/launch_vllm.py index b24386411..e3e3e682a 100644 --- a/scripts/launch_vllm.py +++ b/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:") From 34516ccbf888c0b33326b76dfdc08fc8763f5817 Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Tue, 31 Mar 2026 18:58:49 +0000 Subject: [PATCH 2/6] Add e2e online training smoke test Signed-off-by: Fynn Schmitt-Ulms --- tests/e2e/vllm/test_online_training.py | 161 +++++++++++++++++++++++++ 1 file changed, 161 insertions(+) create mode 100644 tests/e2e/vllm/test_online_training.py diff --git a/tests/e2e/vllm/test_online_training.py b/tests/e2e/vllm/test_online_training.py new file mode 100644 index 000000000..525d03774 --- /dev/null +++ b/tests/e2e/vllm/test_online_training.py @@ -0,0 +1,161 @@ +"""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 os +import subprocess +import sys +import time +import urllib.error +import urllib.request +from pathlib import Path + +import pytest +from loguru import logger + +from tests.e2e.vllm.utils import run_vllm_engine + +MODEL = "Qwen/Qwen3-0.6B" +VLLM_PORT = 8321 +VLLM_PYTHON = os.environ.get("VLLM_PYTHON", sys.executable) +SCRIPTS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "scripts" + + +def wait_for_server(port: int, timeout: float = 180.0, poll_interval: float = 2.0): + """Poll vLLM server health endpoint until ready or timeout.""" + url = f"http://localhost:{port}/health" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + with urllib.request.urlopen(url, timeout=5) as resp: # noqa: S310 + if resp.status == 200: + return + except (urllib.error.URLError, ConnectionError, OSError): + pass + time.sleep(poll_interval) + raise TimeoutError(f"vLLM server on port {port} not ready after {timeout}s") + + +@pytest.fixture +def vllm_server(tmp_path): + """Launch a vLLM server configured for hidden-state extraction.""" + hidden_states_path = str(tmp_path / "hidden_states") + + cmd = [ + VLLM_PYTHON, + str(SCRIPTS_DIR / "launch_vllm.py"), + MODEL, + "--hidden-states-path", + hidden_states_path, + "--", + "--port", + str(VLLM_PORT), + "--max-model-len", + "513", + "--gpu-memory-utilization", + "0.5", + ] + logger.info("Starting vLLM server: {}", " ".join(cmd)) + + process = subprocess.Popen(cmd) # noqa: S603 + + try: + wait_for_server(VLLM_PORT) + logger.info("vLLM server ready on port {}", VLLM_PORT) + except Exception: + process.terminate() + process.wait(timeout=30) + raise + + yield {"port": VLLM_PORT, "hidden_states_path": hidden_states_path, "process": process} + + # Ensure cleanup if the test didn't stop the server itself + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + logger.info("vLLM server stopped (fixture teardown)") + + +@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_cmd = [ + sys.executable, + str(SCRIPTS_DIR / "prepare_data.py"), + "--model", + MODEL, + "--data", + "sharegpt", + "--output", + str(data_path), + "--max-samples", + "50", + "--seq-length", + "512", + ] + logger.info("Preparing data: {}", " ".join(prepare_cmd)) + result = subprocess.run( # noqa: S603 + prepare_cmd, stderr=subprocess.PIPE, text=True, check=False + ) + assert result.returncode == 0, f"prepare_data.py failed:\n{result.stderr}" + + # 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 + server_process = vllm_server["process"] + server_process.terminate() + try: + server_process.wait(timeout=30) + except subprocess.TimeoutExpired: + server_process.kill() + server_process.wait(timeout=10) + logger.info("vLLM server stopped before inference validation") + + # 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) From c5209cb3b466c478ca6df2e893bc7ee46fedf532 Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Tue, 31 Mar 2026 18:58:49 +0000 Subject: [PATCH 3/6] Fix hidden states dtype issue Signed-off-by: Fynn Schmitt-Ulms --- scripts/train.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index e00018eda..1490bad40 100644 --- a/scripts/train.py +++ b/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( From 057efbddc9501f24af949f70413d63ab36bed765 Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Tue, 31 Mar 2026 18:58:49 +0000 Subject: [PATCH 4/6] Refactor to move logic from test_online_training into utils Signed-off-by: Fynn Schmitt-Ulms --- tests/e2e/vllm/test_online_training.py | 95 +++---------------- tests/e2e/vllm/utils.py | 126 ++++++++++++++++++++++++- 2 files changed, 136 insertions(+), 85 deletions(-) diff --git a/tests/e2e/vllm/test_online_training.py b/tests/e2e/vllm/test_online_training.py index 525d03774..b2b74266f 100644 --- a/tests/e2e/vllm/test_online_training.py +++ b/tests/e2e/vllm/test_online_training.py @@ -7,82 +7,34 @@ 4. Validate the trained checkpoint via vLLM inference (run_vllm_engine) """ -import os import subprocess import sys -import time -import urllib.error -import urllib.request from pathlib import Path import pytest from loguru import logger -from tests.e2e.vllm.utils import run_vllm_engine +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 -VLLM_PYTHON = os.environ.get("VLLM_PYTHON", sys.executable) -SCRIPTS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "scripts" - - -def wait_for_server(port: int, timeout: float = 180.0, poll_interval: float = 2.0): - """Poll vLLM server health endpoint until ready or timeout.""" - url = f"http://localhost:{port}/health" - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - try: - with urllib.request.urlopen(url, timeout=5) as resp: # noqa: S310 - if resp.status == 200: - return - except (urllib.error.URLError, ConnectionError, OSError): - pass - time.sleep(poll_interval) - raise TimeoutError(f"vLLM server on port {port} not ready after {timeout}s") @pytest.fixture def vllm_server(tmp_path): """Launch a vLLM server configured for hidden-state extraction.""" hidden_states_path = str(tmp_path / "hidden_states") - - cmd = [ - VLLM_PYTHON, - str(SCRIPTS_DIR / "launch_vllm.py"), - MODEL, - "--hidden-states-path", - hidden_states_path, - "--", - "--port", - str(VLLM_PORT), - "--max-model-len", - "513", - "--gpu-memory-utilization", - "0.5", - ] - logger.info("Starting vLLM server: {}", " ".join(cmd)) - - process = subprocess.Popen(cmd) # noqa: S603 - - try: - wait_for_server(VLLM_PORT) - logger.info("vLLM server ready on port {}", VLLM_PORT) - except Exception: - process.terminate() - process.wait(timeout=30) - raise + process = launch_vllm_server(MODEL, VLLM_PORT, hidden_states_path) yield {"port": VLLM_PORT, "hidden_states_path": hidden_states_path, "process": process} - # Ensure cleanup if the test didn't stop the server itself - if process.poll() is None: - process.terminate() - try: - process.wait(timeout=30) - except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout=10) - logger.info("vLLM server stopped (fixture teardown)") + stop_vllm_server(process) @pytest.mark.e2e @@ -95,25 +47,7 @@ def test_online_training( port = vllm_server["port"] # Step 1: Prepare data - prepare_cmd = [ - sys.executable, - str(SCRIPTS_DIR / "prepare_data.py"), - "--model", - MODEL, - "--data", - "sharegpt", - "--output", - str(data_path), - "--max-samples", - "50", - "--seq-length", - "512", - ] - logger.info("Preparing data: {}", " ".join(prepare_cmd)) - result = subprocess.run( # noqa: S603 - prepare_cmd, stderr=subprocess.PIPE, text=True, check=False - ) - assert result.returncode == 0, f"prepare_data.py failed:\n{result.stderr}" + prepare_data(MODEL, data_path) # Step 2: Train against live vLLM server train_cmd = [ @@ -147,14 +81,7 @@ def test_online_training( assert result.returncode == 0, f"train.py failed:\n{result.stderr}" # Stop the vLLM server to free GPU memory before running inference - server_process = vllm_server["process"] - server_process.terminate() - try: - server_process.wait(timeout=30) - except subprocess.TimeoutExpired: - server_process.kill() - server_process.wait(timeout=10) - logger.info("vLLM server stopped before inference validation") + stop_vllm_server(vllm_server["process"]) # Step 3: Validate trained checkpoint with vLLM inference checkpoint_path = str(save_path / "0") diff --git a/tests/e2e/vllm/utils.py b/tests/e2e/vllm/utils.py index c0cdd2711..a4a6193f9 100644 --- a/tests/e2e/vllm/utils.py +++ b/tests/e2e/vllm/utils.py @@ -2,13 +2,137 @@ import os import subprocess import sys +import time +import urllib.error +import urllib.request from collections.abc import Iterable from pathlib import Path from textwrap import indent from loguru import logger -__all__ = ["run_vllm_engine"] +__all__ = [ + "SCRIPTS_DIR", + "VLLM_PYTHON", + "launch_vllm_server", + "prepare_data", + "run_vllm_engine", + "stop_vllm_server", + "wait_for_server", +] + +VLLM_PYTHON = os.environ.get("VLLM_PYTHON", sys.executable) +SCRIPTS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "scripts" + + +def wait_for_server( + port: int, + timeout: float = 180.0, + poll_interval: float = 2.0, + process: subprocess.Popen | None = None, +): + """Poll vLLM server health endpoint until ready or timeout. + + If *process* is provided, checks whether it has exited between polls + so that startup failures are reported immediately instead of waiting + for the full timeout. + """ + + logger.info("Waiting for server") + url = f"http://localhost:{port}/health" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if process is not None and process.poll() is not None: + raise RuntimeError( + f"vLLM server process exited with code {process.returncode} " + "before becoming ready" + ) + try: + with urllib.request.urlopen(url, timeout=5) as resp: # noqa: S310 + if resp.status == 200: + return + except (urllib.error.URLError, ConnectionError, OSError): + pass + time.sleep(poll_interval) + raise TimeoutError(f"vLLM server on port {port} not ready after {timeout}s") + + +def launch_vllm_server( + model: str, + port: int, + hidden_states_path: str, + max_model_len: int = 513, + gpu_memory_utilization: float = 0.5, +) -> subprocess.Popen: + """Launch a vLLM server configured for hidden-state extraction. + + Returns the server subprocess. Caller is responsible for stopping it + via stop_vllm_server(). + """ + cmd = [ + VLLM_PYTHON, + str(SCRIPTS_DIR / "launch_vllm.py"), + model, + "--hidden-states-path", + hidden_states_path, + "--", + "--port", + str(port), + "--max-model-len", + str(max_model_len), + "--gpu-memory-utilization", + str(gpu_memory_utilization), + ] + logger.info("Starting vLLM server: {}", " ".join(cmd)) + + process = subprocess.Popen(cmd) # noqa: S603 + + try: + wait_for_server(port, process=process) + logger.info("vLLM server ready on port {}", port) + except Exception: + process.terminate() + process.wait(timeout=30) + raise + + return process + + +def stop_vllm_server(process: subprocess.Popen): + """Gracefully stop a vLLM server subprocess.""" + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + logger.info("vLLM server stopped") + + +def prepare_data( + model: str, data_path: Path, max_samples: int = 50, seq_length: int = 512 +): + """Tokenize ShareGPT data using prepare_data.py.""" + cmd = [ + sys.executable, + str(SCRIPTS_DIR / "prepare_data.py"), + "--model", + model, + "--data", + "sharegpt", + "--output", + str(data_path), + "--max-samples", + str(max_samples), + "--seq-length", + str(seq_length), + ] + logger.info("Preparing data: {}", " ".join(cmd)) + result = subprocess.run( # noqa: S603 + cmd, stderr=subprocess.PIPE, text=True, check=False + ) + assert result.returncode == 0, f"prepare_data.py failed:\n{result.stderr}" def run_vllm_engine( From 4ca1580afb95b00e14f4b9a279565ed50d97b861 Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Tue, 31 Mar 2026 18:58:49 +0000 Subject: [PATCH 5/6] Add e2e offline training smoke test Signed-off-by: Fynn Schmitt-Ulms --- tests/e2e/vllm/test_offline_training.py | 113 ++++++++++++++++++++++++ 1 file changed, 113 insertions(+) create mode 100644 tests/e2e/vllm/test_offline_training.py diff --git a/tests/e2e/vllm/test_offline_training.py b/tests/e2e/vllm/test_offline_training.py new file mode 100644 index 000000000..bb616bbed --- /dev/null +++ b/tests/e2e/vllm/test_offline_training.py @@ -0,0 +1,113 @@ +"""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) From 7d172ad0285524851b05820ca6137aae9497abae Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Fri, 3 Apr 2026 13:28:38 +0000 Subject: [PATCH 6/6] Format Signed-off-by: Fynn Schmitt-Ulms --- tests/e2e/vllm/test_offline_training.py | 12 ++++++++---- tests/e2e/vllm/test_online_training.py | 6 +++++- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/tests/e2e/vllm/test_offline_training.py b/tests/e2e/vllm/test_offline_training.py index bb616bbed..d05a77f55 100644 --- a/tests/e2e/vllm/test_offline_training.py +++ b/tests/e2e/vllm/test_offline_training.py @@ -34,7 +34,11 @@ def vllm_server(tmp_path): 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} + yield { + "port": VLLM_PORT, + "hidden_states_path": hidden_states_path, + "process": process, + } stop_vllm_server(process) @@ -72,9 +76,9 @@ def test_offline_training( 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}" + 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"]) diff --git a/tests/e2e/vllm/test_online_training.py b/tests/e2e/vllm/test_online_training.py index b2b74266f..dfd4910fc 100644 --- a/tests/e2e/vllm/test_online_training.py +++ b/tests/e2e/vllm/test_online_training.py @@ -32,7 +32,11 @@ def vllm_server(tmp_path): 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} + yield { + "port": VLLM_PORT, + "hidden_states_path": hidden_states_path, + "process": process, + } stop_vllm_server(process)