From 236879693bed5190c6bf3c14dee6bda3ced23fae Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 17 Jul 2026 14:29:41 -0400 Subject: [PATCH 1/4] Add Fast-LLM integration test suite Test layer of the stacked breakdown of PR #140, extracted on top of the config/metrics PR (#155). Covers the vLLM v1 weight-broadcast path, world setup, actor error handling, launch-process monitoring, and model-version tagging. Co-Authored-By: Claude Opus 4.8 --- tests/__init__.py | 1 + tests/conftest.py | 196 ++++ tests/distributed_trainer_helper.py | 650 ++++++++++++ tests/fast_llm_trainer_helper.py | 271 +++++ tests/server_weight_update_utils.py | 653 ++++++++++++ tests/sync_helper.py | 110 ++ tests/test_actor_error_handling.py | 290 +++++ tests/test_launch_process_monitoring.py | 65 ++ tests/test_model_version.py | 51 + tests/test_vllm1_fast_llm_broadcast.py | 590 +++++++++++ tests/test_vllm1_integration.py | 1286 +++++++++++++++++++++++ tests/test_world_multinode.py | 842 +++++++++++++++ tests/trainer_test_utils.py | 128 +++ tests/vllm_engine_helper.py | 617 +++++++++++ tests/weight_update_utils.py | 53 + 15 files changed, 5803 insertions(+) create mode 100644 tests/__init__.py create mode 100644 tests/conftest.py create mode 100755 tests/distributed_trainer_helper.py create mode 100644 tests/fast_llm_trainer_helper.py create mode 100644 tests/server_weight_update_utils.py create mode 100644 tests/sync_helper.py create mode 100644 tests/test_actor_error_handling.py create mode 100644 tests/test_launch_process_monitoring.py create mode 100644 tests/test_model_version.py create mode 100644 tests/test_vllm1_fast_llm_broadcast.py create mode 100644 tests/test_vllm1_integration.py create mode 100644 tests/test_world_multinode.py create mode 100644 tests/trainer_test_utils.py create mode 100755 tests/vllm_engine_helper.py create mode 100644 tests/weight_update_utils.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..9b491dd0 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for PipelineRL vLLM integration.""" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..e33b8261 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,196 @@ +"""Pytest configuration and fixtures for vllm1 tests.""" + +import os +import pytest +import torch +import tempfile +from pathlib import Path +import subprocess +import sys + +from pipelinerl.vllm1 import EngineManager + + +@pytest.fixture(scope="session") +def model_name(): + """Model to use for testing.""" + return "Qwen/Qwen2.5-0.5B-Instruct" + + +@pytest.fixture(scope="session") +def sample_prompts(): + """Sample prompts for generation testing.""" + return [ + "Write a haiku about coding:", + "The capital of France is", + "In a galaxy far away,", + ] + + +@pytest.fixture(scope="session") +def simple_prompt(): + """Single simple prompt for deterministic testing.""" + return "The capital of France is" + + +@pytest.fixture(scope="session") +def num_gpus(): + """Number of GPUs available.""" + return torch.cuda.device_count() + + +@pytest.fixture(scope="session") +def require_2_gpus(num_gpus): + """Skip test if less than 2 GPUs available.""" + if num_gpus < 2: + pytest.skip("Test requires at least 2 GPUs") + + +@pytest.fixture(scope="session") +def require_gpu(): + """Skip test if no GPU available.""" + if not torch.cuda.is_available(): + pytest.skip("Test requires GPU") + + +@pytest.fixture +def temp_dir(): + """Temporary directory for test files.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) + + +@pytest.fixture(scope="session") +def shared_test_dir(): + """Session-scoped shared directory for test data that persists across tests. + + Use this for data that needs to be shared between tests (like perturbed weights). + """ + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) + + +@pytest.fixture +def distributed_init_method(temp_dir): + """File-based init method for distributed testing.""" + return f"file://{temp_dir}/dist_init" + + +@pytest.fixture(scope="session") +def shared_distributed_init_method(shared_test_dir): + """Session-scoped file-based init method for tests that share data.""" + return f"file://{shared_test_dir}/dist_init" + + +@pytest.fixture(scope="session") +def cache_dir(): + """Directory for caching downloaded models.""" + cache_path = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")) + cache_path.mkdir(parents=True, exist_ok=True) + return cache_path + + +@pytest.fixture +def vllm_server_port(): + """Port for vLLM server in tests.""" + # Use a high port to avoid conflicts + return 8765 + + +@pytest.fixture +def generation_config(): + """Configuration for deterministic generation.""" + return { + "temperature": 0.0, + "top_p": 1.0, + "max_tokens": 50, + "seed": 42, + } + + +@pytest.fixture +def vllm_engine_factory_2gpu(model_name): + """Factory fixture that defaults to 2 GPUs. + + Usage: + async with vllm_engine_factory_2gpu() as manager: + # Uses 2 GPUs by default + # Access engine via manager.engine + ... + """ + def _factory(tensor_parallel_size: int = 2, **kwargs): + """Create engine with 2 GPUs by default.""" + import argparse + + args = argparse.Namespace( + model=model_name, + tensor_parallel_size=tensor_parallel_size, + disable_log_stats=True, + enable_log_requests=False, + **kwargs + ) + + return EngineManager.create_engine(args) + + return _factory + + +@pytest.fixture +def vllm_engine_factory(model_name): + """Factory fixture for creating vLLM engines. + + Usage in tests: + async with vllm_engine_factory() as manager: + # use manager.engine for generation + ... + # automatic cleanup + + Or with custom config: + async with vllm_engine_factory(tensor_parallel_size=2) as manager: + # use manager.engine with 2 GPUs + ... + + Or if you need engine_config: + async with vllm_engine_factory() as manager: + # access manager.engine, manager.engine_config, manager.args + ... + """ + def _factory(tensor_parallel_size: int = 1, **kwargs): + """Create engine context manager with test defaults. + + Args: + tensor_parallel_size: Number of GPUs + **kwargs: Additional attributes for args object + + Returns: + Async context manager for EngineManager + """ + import argparse + + # Create minimal args object with required attributes for AsyncEngineArgs.from_cli_args() + args = argparse.Namespace( + model=model_name, + tensor_parallel_size=tensor_parallel_size, + disable_log_stats=True, + enable_log_requests=False, + # Apply any additional kwargs + **kwargs + ) + + print("args: ", args) + + return EngineManager.create_engine(args) + + return _factory + + +@pytest.fixture +def distributed_trainer_helper(): + """Path to the distributed trainer helper script.""" + return Path(__file__).parent / "distributed_trainer_helper.py" + + +@pytest.fixture +def vllm_engine_helper(): + """Path to the vLLM engine helper script.""" + return Path(__file__).parent / "vllm_engine_helper.py" diff --git a/tests/distributed_trainer_helper.py b/tests/distributed_trainer_helper.py new file mode 100755 index 00000000..7e50decb --- /dev/null +++ b/tests/distributed_trainer_helper.py @@ -0,0 +1,650 @@ +#!/usr/bin/env python3 +"""Helper script for distributed trainer process. + +This script is run as a separate process with CUDA_VISIBLE_DEVICES set, +allowing proper GPU isolation for distributed tests. +""" + +import sys +import argparse +import logging +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) +from trainer_test_utils import ( + _resolve_model_path, + _load_state_dict, + _create_perturbed_state_dict, + _init_actor_process_group, + _broadcast_tensors, + _wait_for_servers_ready, +) + +# Setup debug logging +logging.basicConfig( + level=logging.DEBUG, + format="[%(asctime)s] [TRAINER-%(levelname)s] %(message)s", + datefmt="%H:%M:%S", +) +logger = logging.getLogger(__name__) + + +def _wait_all_actors(sync_path, name: str, num_actors: int, timeout: float = 120): + """Wait for all actors to signal a named sync point. + + Each actor signals ``{name}_actor_{i}`` for i in range(num_actors). + """ + from pathlib import Path + sys.path.insert(0, str(Path(__file__).parent)) + from sync_helper import SyncPoint + + for i in range(num_actors): + SyncPoint(sync_path, f"{name}_actor_{i}").wait(timeout=timeout) + + +def _broadcast_via_server( + state_dict: dict, + server_urls: list, + version: int, + process_group, + label: str = "", +): + """Broadcast weights to one or more running vLLM servers via HTTP POST + NCCL. + + One POST thread is started per server URL (all in parallel) before the + NCCL broadcast so that all servers are ready to receive simultaneously. + """ + import threading + import time + import requests + from weight_update_utils import create_weight_update_request_from_state_dict + + label_str = f" {label}" if label else "" + print(f"[Trainer] Broadcasting {len(state_dict)}{label_str} parameters to {len(server_urls)} server(s)") + + request = create_weight_update_request_from_state_dict(state_dict, version=version) + + errors = [] + threads = [] + + for url in server_urls: + err = {"error": None} + errors.append(err) + + def _post(server_url=url, post_result=err): + try: + print(f"[Trainer] POSTing weight update request to {server_url}...") + resp = requests.post( + f"{server_url}/receive_weight_update", + json=request.model_dump(), + timeout=600, + ) + if resp.status_code != 200: + post_result["error"] = ( + f"POST to {server_url} failed with status {resp.status_code}: {resp.text}" + ) + else: + print(f"[Trainer] Server {server_url} acknowledged weight update") + except Exception as e: + post_result["error"] = f"POST to {server_url} failed: {e}" + + t = threading.Thread(target=_post, daemon=False) + threads.append(t) + t.start() + + time.sleep(0.5) # Give all servers a moment to start receiving + + _broadcast_tensors(state_dict, process_group) + + for t in threads: + t.join(timeout=60) + + failed = [e["error"] for e in errors if e["error"]] + if failed: + raise RuntimeError(f"Weight update POST(s) failed: {failed}") + + print(f"[Trainer] Broadcast{label_str} complete") + + +# --------------------------------------------------------------------------- +# Public command functions +# --------------------------------------------------------------------------- + +def init_process_group(init_method: str, rank: int, world_size: int): + """Initialize a distributed process group and wait.""" + import torch.distributed as dist + import time + + process_group = _init_actor_process_group(init_method, rank, world_size) + print(f"[Trainer rank={rank}] Process group initialized successfully") + + # Wait for coordination + time.sleep(3) + + print(f"[Trainer rank={rank}] Destroying process group") + dist.destroy_process_group(process_group) + print(f"[Trainer rank={rank}] Process group destroyed") + + +def save_model_to_dir(state_dict: dict, output_dir: str, model_name: str): + """Save state_dict to a directory as safetensors with config. + + Args: + state_dict: Model state dict to save + output_dir: Directory to save model + model_name: Original model name to copy config from + """ + from pathlib import Path + from safetensors.torch import save_file + import shutil + + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + # Save weights as safetensors + safetensors_path = output_path / "model.safetensors" + save_file(state_dict, str(safetensors_path)) + print(f"[Trainer] Saved model weights to {safetensors_path}") + + # Copy config.json from original model + original_path = _resolve_model_path(model_name) + + config_src = original_path / "config.json" + config_dst = output_path / "config.json" + shutil.copy(config_src, config_dst) + print(f"[Trainer] Copied config.json to {config_dst}") + + # Copy tokenizer files + for filename in [ + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "vocab.json", + "merges.txt", + "tokenizer.model", + ]: + src = original_path / filename + if src.exists(): + dst = output_path / filename + shutil.copy(src, dst) + print(f"[Trainer] Copied {filename}") + + return str(output_path) + + +def broadcast_weights( + init_method: str, model_name: str, perturb: bool = False, sync_dir: str = None +): + """Load model and broadcast weights to vLLM worker.""" + import torch + import torch.distributed as dist + from pathlib import Path + + # Setup sync points if provided + if sync_dir: + sys.path.insert(0, str(Path(__file__).parent)) + from sync_helper import SyncPoint, write_weight_update_request + + sync_path = Path(sync_dir) + baseline_done = SyncPoint(sync_path, "baseline_done") + ready_to_receive = SyncPoint(sync_path, "ready_to_receive") + request_ready = SyncPoint(sync_path, "request_ready") + receiving_started = SyncPoint(sync_path, "receiving_started") + broadcast_done = SyncPoint(sync_path, "broadcast_done") + + # IMPORTANT: Initialize process group FIRST (before any waiting) + process_group = _init_actor_process_group(init_method, rank=0, world_size=2) + + # Now wait for vLLM to finish baseline and be ready to receive + if sync_dir: + print("[Trainer] Waiting for vLLM to finish baseline generation...") + baseline_done.wait(timeout=60) + print("[Trainer] Baseline done") + + print("[Trainer] Waiting for vLLM to be ready to receive weights...") + ready_to_receive.wait(timeout=60) + print("[Trainer] vLLM ready, starting weight broadcast") + + print(f"[Trainer] Loading tensors from safetensors for {model_name}") + state_dict, _ = _load_state_dict(model_name) + + params_to_broadcast = state_dict + print(f"[Trainer] Will broadcast {len(params_to_broadcast)} parameters") + + # Create and send WeightUpdateRequest to vLLM + if sync_dir: + from weight_update_utils import create_weight_update_request_from_state_dict + + print("[Trainer] Creating WeightUpdateRequest...") + request = create_weight_update_request_from_state_dict( + params_to_broadcast, version=1 + ) + write_weight_update_request(sync_path, request) + request_ready.signal() + print( + f"[Trainer] Sent WeightUpdateRequest with {len(request.parameters_info)} parameters" + ) + + # Wait for vLLM to start receiving before we broadcast + print("[Trainer] Waiting for vLLM to start receiving...") + receiving_started.wait(timeout=60) + print("[Trainer] vLLM is receiving, starting broadcast") + + print(f"[Trainer] Broadcasting {len(params_to_broadcast)} parameters") + + # Optionally perturb weights - add noise to ALL tensors + if perturb: + params_to_broadcast = _create_perturbed_state_dict(params_to_broadcast) + + # Broadcast each weight with detailed logging + logger.info(f"Starting broadcast of {len(params_to_broadcast)} parameters") + for i, (name, tensor) in enumerate(params_to_broadcast.items()): + logger.debug(f"[{i+1}/{len(state_dict)}] Preparing to broadcast: {name}") + logger.debug( + f" - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}" + ) + if tensor.device.type != "cuda": + logger.debug(f" - Moving {name} to CUDA") + tensor = tensor.cuda(0) + logger.debug(f" - {name} now on device: {tensor.device}") + logger.debug(f" - Calling dist.broadcast for {name}...") + dist.broadcast(tensor, src=0, group=process_group) + logger.debug(f" - Broadcast complete for {name}") + if (i + 1) % 10 == 0: + logger.info(f"Broadcasted {i+1}/{len(params_to_broadcast)} parameters") + + print(f"[Trainer] All {len(params_to_broadcast)} parameters broadcasted") + + # Signal broadcast complete BEFORE destroying process group + if sync_dir: + broadcast_done.signal() + print("[Trainer] Signaled broadcast complete") + + dist.destroy_process_group(process_group) + print("[Trainer] Process group destroyed") + + +def broadcast_cross_validation( + init_method: str, model_name: str, sync_dir: str, temp_dir: str +): + """Cross-validation test: broadcast perturbed, then original weights. + + Also saves perturbed model to disk for vLLM to load. + """ + import torch.distributed as dist + from pathlib import Path + + sys.path.insert(0, str(Path(__file__).parent)) + from sync_helper import SyncPoint, write_weight_update_request + from weight_update_utils import create_weight_update_request_from_state_dict + + sync_path = Path(sync_dir) + baseline_done = SyncPoint(sync_path, "baseline_done") + perturbed_model_saved = SyncPoint(sync_path, "perturbed_model_saved") + ready_to_receive_perturbed = SyncPoint(sync_path, "ready_to_receive_perturbed") + perturbed_broadcast_done = SyncPoint(sync_path, "perturbed_broadcast_done") + mod1_done = SyncPoint(sync_path, "mod1_done") + first_engine_destroyed = SyncPoint(sync_path, "first_engine_destroyed") + engine_recreated = SyncPoint(sync_path, "engine_recreated") + ready_to_receive_original = SyncPoint(sync_path, "ready_to_receive_original") + original_broadcast_done = SyncPoint(sync_path, "original_broadcast_done") + + process_group = _init_actor_process_group(init_method, rank=0, world_size=2) + + print("[Trainer] Waiting for vLLM baseline generation...") + baseline_done.wait(timeout=120) + + print(f"[Trainer] Loading original model {model_name}") + original_state_dict, model_path = _load_state_dict(model_name) + + perturbed_state_dict = _create_perturbed_state_dict(original_state_dict) + + # Save perturbed model to disk + perturbed_model_dir = Path(temp_dir) / "perturbed_model" + print(f"[Trainer] Saving perturbed model to {perturbed_model_dir}") + saved_path = save_model_to_dir( + perturbed_state_dict, str(perturbed_model_dir), str(model_path) + ) + + path_file = sync_path / "perturbed_model_path.txt" + path_file.write_text(saved_path) + perturbed_model_saved.signal() + print(f"[Trainer] Signaled perturbed model saved at: {saved_path}") + + # Broadcast perturbed weights + print("[Trainer] Waiting for vLLM to be ready for perturbed broadcast...") + ready_to_receive_perturbed.wait(timeout=120) + + print(f"[Trainer] Broadcasting {len(perturbed_state_dict)} perturbed parameters") + request = create_weight_update_request_from_state_dict(perturbed_state_dict, version=1) + write_weight_update_request(sync_path, request) + _broadcast_tensors(perturbed_state_dict, process_group) + + perturbed_broadcast_done.signal() + print("[Trainer] Perturbed weights broadcast complete") + + print("[Trainer] Waiting for vLLM to finish res_mod_1...") + mod1_done.wait(timeout=120) + + print("[Trainer] Destroying process group for first broadcast") + dist.destroy_process_group(process_group) + + print("[Trainer] Waiting for vLLM to destroy first engine...") + first_engine_destroyed.wait(timeout=120) + + print("[Trainer] Recreating process group for second broadcast") + process_group = _init_actor_process_group(init_method, rank=0, world_size=2) + print("[Trainer] Process group recreated, waiting at rendezvous...") + + print("[Trainer] Waiting for vLLM to recreate engine...") + engine_recreated.wait(timeout=300) # 5 minutes - engine creation can be slow + print("[Trainer] vLLM engine recreated, both in new process group") + + # Broadcast original weights + print("[Trainer] Waiting for vLLM to be ready for original broadcast...") + ready_to_receive_original.wait(timeout=120) + + print(f"[Trainer] Broadcasting {len(original_state_dict)} original parameters") + request = create_weight_update_request_from_state_dict(original_state_dict, version=2) + write_weight_update_request(sync_path, request) + _broadcast_tensors(original_state_dict, process_group) + + original_broadcast_done.signal() + print("[Trainer] Original weights broadcast complete") + + dist.destroy_process_group(process_group) + print("[Trainer] Process group destroyed") + + +def broadcast_back_and_forth( + init_method: str, + model_name: str, + sync_dir: str, + num_actors: int = 1, + world_size: int = 2, +): + """Back-and-forth test: broadcast perturbed → original → perturbed again. + + Tests that we can switch between weight sets multiple times. + Supports multiple actors: waits for all actors to signal readiness before + each broadcast, then sends a single shared completion signal. + """ + import torch.distributed as dist + from pathlib import Path + + sys.path.insert(0, str(Path(__file__).parent)) + from sync_helper import SyncPoint, write_weight_update_request + from weight_update_utils import create_weight_update_request_from_state_dict + + sync_path = Path(sync_dir) + perturbed1_done = SyncPoint(sync_path, "perturbed1_done") + original_done = SyncPoint(sync_path, "original_done") + perturbed2_done = SyncPoint(sync_path, "perturbed2_done") + + process_group = _init_actor_process_group(init_method, rank=0, world_size=world_size) + + print(f"[Trainer] Waiting for {num_actors} actor(s) to finish baseline generation...") + _wait_all_actors(sync_path, "baseline_done", num_actors, timeout=120) + + print(f"[Trainer] Loading model {model_name}") + original_state_dict, model_path = _load_state_dict(model_name) + + perturbed_state_dict = _create_perturbed_state_dict(original_state_dict) + + # Save perturbed weights for reuse in server tests + perturbed_weights_dir = Path(sync_dir) / "perturbed_weights" + print(f"[Trainer] Saving perturbed weights to {perturbed_weights_dir}") + saved_path = save_model_to_dir( + perturbed_state_dict, str(perturbed_weights_dir), str(model_path) + ) + print(f"[Trainer] Perturbed weights saved to {saved_path}") + + # Broadcast 1: Perturbed weights + print(f"[Trainer] Waiting for {num_actors} actor(s) to be ready for first perturbed broadcast...") + _wait_all_actors(sync_path, "ready_for_perturbed1", num_actors, timeout=120) + + print(f"[Trainer] Broadcasting perturbed weights (1st time) to {num_actors} actor(s)") + request = create_weight_update_request_from_state_dict(perturbed_state_dict, version=1) + write_weight_update_request(sync_path, request) + _broadcast_tensors(perturbed_state_dict, process_group) + + perturbed1_done.signal() + print("[Trainer] First perturbed broadcast complete") + + # Broadcast 2: Original weights + print(f"[Trainer] Waiting for {num_actors} actor(s) to be ready for original broadcast...") + _wait_all_actors(sync_path, "ready_for_original", num_actors, timeout=120) + + print(f"[Trainer] Broadcasting original weights to {num_actors} actor(s)") + request = create_weight_update_request_from_state_dict(original_state_dict, version=2) + write_weight_update_request(sync_path, request) + _broadcast_tensors(original_state_dict, process_group) + + original_done.signal() + print("[Trainer] Original broadcast complete") + + # Broadcast 3: Perturbed weights again (same as first) + print(f"[Trainer] Waiting for {num_actors} actor(s) to be ready for second perturbed broadcast...") + _wait_all_actors(sync_path, "ready_for_perturbed2", num_actors, timeout=120) + + print(f"[Trainer] Broadcasting perturbed weights (2nd time) to {num_actors} actor(s)") + request = create_weight_update_request_from_state_dict(perturbed_state_dict, version=3) + write_weight_update_request(sync_path, request) + _broadcast_tensors(perturbed_state_dict, process_group) + + perturbed2_done.signal() + print("[Trainer] Second perturbed broadcast complete") + + dist.destroy_process_group(process_group) + print("[Trainer] Process group destroyed") + + +def timed_broadcast_server_test( + init_method: str, + model_name: str, + server_urls: list, + world_size: int = 2, +): + """Timed broadcast for server tests: perturbed → original → perturbed with delays. + + This simulates a real-world scenario where weight updates happen while + the server is running and serving requests. + + Pattern: original (server default) → perturbed → original → perturbed + + Args: + init_method: Distributed init method + model_name: Model name to load + server_urls: List of base URLs of vLLM servers (e.g., ["http://127.0.0.1:8000"]) + world_size: Total world size (trainer rank 0 + all vLLM workers) + """ + import torch.distributed as dist + import time + import requests + + process_group = _init_actor_process_group(init_method, rank=0, world_size=world_size) + + _wait_for_servers_ready(server_urls, extra_wait_secs=10) + + print(f"[Trainer] Loading original weights from {model_name}") + original_state_dict, _ = _load_state_dict(model_name) + + perturbed_state_dict = _create_perturbed_state_dict(original_state_dict) + + # Broadcast 1: Perturbed weights + _broadcast_via_server(perturbed_state_dict, server_urls, version=1, process_group=process_group, label="perturbed") + + print("[Trainer] Waiting 5 seconds before broadcasting original weights...") + time.sleep(5) + + # Broadcast 2: Original weights + _broadcast_via_server(original_state_dict, server_urls, version=2, process_group=process_group, label="original") + + print("[Trainer] Waiting 5 seconds before broadcasting perturbed weights again...") + time.sleep(5) + + # Broadcast 3: Perturbed weights again (same as first) + _broadcast_via_server(perturbed_state_dict, server_urls, version=3, process_group=process_group, label="perturbed (2nd time)") + + # Wait to allow generation with the last broadcast before tearing down + print("[Trainer] Waiting 5 seconds for generation with final weights...") + time.sleep(5) + + # Signal training is finished so vLLM servers destroy their side of the process group + for url in server_urls: + print(f"[Trainer] Sending training_finished signal to {url}...") + requests.post(f"{url}/training_finished", timeout=10) + + # Cleanup — destroy_process_group now resolves because vLLM servers respond to /training_finished + dist.destroy_process_group(process_group) + print("[Trainer] Process group destroyed, exiting") + + +def rapid_broadcast_cycles( + init_method: str, + model_name: str, + server_urls: list, + world_size: int = 2, + n_cycles: int = 6, +): + """Hybrid broadcast designed to catch transition/garbage generations. + + Structure: + 1. Slow broadcast: perturbed (5 s wait after) — establishes text_B + 2. Slow broadcast: original (5 s wait after) — re-establishes text_A + 3. n_cycles rapid pairs: perturbed → original (1 s between each) + 4. Slow broadcast: perturbed (5 s wait after) — end on text_B so the + overall A→B→A→B pattern remains detectable + + The slow initial cycles give the generation loop enough stable time to + identify text_A and text_B by frequency. The rapid cycles create many + short broadcast windows where mid-broadcast (garbage) generations are + likely to be caught by a zero-interval generation loop. + """ + import torch.distributed as dist + import time + import requests + + process_group = _init_actor_process_group(init_method, rank=0, world_size=world_size) + + _wait_for_servers_ready(server_urls, extra_wait_secs=10) + + print(f"[Trainer] Loading weights from {model_name}") + original_state_dict, _ = _load_state_dict(model_name) + perturbed_state_dict = _create_perturbed_state_dict(original_state_dict) + + version = 1 + + # --- Slow cycle: establish text_B and text_A clearly --- + print("[Trainer] Slow broadcast 1: perturbed (establishing text_B)...") + _broadcast_via_server(perturbed_state_dict, server_urls, version=version, process_group=process_group, label="perturbed (slow)") + version += 1 + time.sleep(5) + + print("[Trainer] Slow broadcast 2: original (re-establishing text_A)...") + _broadcast_via_server(original_state_dict, server_urls, version=version, process_group=process_group, label="original (slow)") + version += 1 + time.sleep(5) + + # --- Rapid cycles: 1 s between broadcasts --- + for i in range(n_cycles): + print(f"[Trainer] Rapid cycle {i + 1}/{n_cycles}: perturbed...") + _broadcast_via_server(perturbed_state_dict, server_urls, version=version, process_group=process_group, label=f"perturbed (rapid {i + 1})") + version += 1 + time.sleep(1) + + print(f"[Trainer] Rapid cycle {i + 1}/{n_cycles}: original...") + _broadcast_via_server(original_state_dict, server_urls, version=version, process_group=process_group, label=f"original (rapid {i + 1})") + version += 1 + time.sleep(1) + + # --- Final slow broadcast: end on perturbed so ABAB pattern holds --- + print("[Trainer] Final slow broadcast: perturbed (ending on text_B)...") + _broadcast_via_server(perturbed_state_dict, server_urls, version=version, process_group=process_group, label="perturbed (final)") + time.sleep(5) + + for url in server_urls: + print(f"[Trainer] Sending training_finished signal to {url}...") + requests.post(f"{url}/training_finished", timeout=10) + + dist.destroy_process_group(process_group) + print("[Trainer] Process group destroyed, exiting") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Distributed trainer helper") + parser.add_argument("command", choices=["init", "broadcast", "cross_validation", "back_and_forth", "timed_broadcast_server_test", "rapid_broadcast_cycles"]) + parser.add_argument("--init-method", required=True) + parser.add_argument("--rank", type=int, default=0) + parser.add_argument("--world-size", type=int, default=2) + parser.add_argument("--model-name", type=str) + parser.add_argument("--perturb", action="store_true") + parser.add_argument("--sync-dir", type=str, help="Directory for sync files") + parser.add_argument( + "--temp-dir", type=str, help="Temporary directory for saving models" + ) + parser.add_argument( + "--server-urls", nargs="+", help="Base URL(s) of vLLM server(s) (e.g., http://127.0.0.1:8000)" + ) + parser.add_argument("--num-actors", type=int, default=1, help="Number of vLLM actor processes") + parser.add_argument("--n-cycles", type=int, default=6, help="Number of rapid broadcast cycles (rapid_broadcast_cycles command)") + + args = parser.parse_args() + + try: + if args.command == "init": + init_process_group(args.init_method, args.rank, args.world_size) + elif args.command == "broadcast": + if not args.model_name: + print("Error: --model-name required for broadcast command") + sys.exit(1) + broadcast_weights( + args.init_method, args.model_name, args.perturb, args.sync_dir + ) + elif args.command == "cross_validation": + if not args.model_name or not args.sync_dir or not args.temp_dir: + print( + "Error: --model-name, --sync-dir, and --temp-dir required for cross_validation" + ) + sys.exit(1) + broadcast_cross_validation( + args.init_method, args.model_name, args.sync_dir, args.temp_dir + ) + elif args.command == "back_and_forth": + if not args.model_name or not args.sync_dir: + print("Error: --model-name and --sync-dir required for back_and_forth") + sys.exit(1) + broadcast_back_and_forth( + args.init_method, + args.model_name, + args.sync_dir, + num_actors=args.num_actors, + world_size=args.world_size, + ) + elif args.command == "timed_broadcast_server_test": + if not args.model_name or not args.server_urls: + print("Error: --model-name and --server-urls required for timed_broadcast_server_test") + sys.exit(1) + timed_broadcast_server_test( + args.init_method, + args.model_name, + args.server_urls, + world_size=args.world_size, + ) + elif args.command == "rapid_broadcast_cycles": + if not args.model_name or not args.server_urls: + print("Error: --model-name and --server-urls required for rapid_broadcast_cycles") + sys.exit(1) + rapid_broadcast_cycles( + args.init_method, + args.model_name, + args.server_urls, + world_size=args.world_size, + n_cycles=args.n_cycles, + ) + except Exception as e: + print(f"[Trainer] Error: {e}") + import traceback + + traceback.print_exc() + sys.exit(1) diff --git a/tests/fast_llm_trainer_helper.py b/tests/fast_llm_trainer_helper.py new file mode 100644 index 00000000..6ff173e2 --- /dev/null +++ b/tests/fast_llm_trainer_helper.py @@ -0,0 +1,271 @@ +"""Helper functions for Fast-LLM weight broadcast testing.""" + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) +from trainer_test_utils import ( + _load_state_dict, + _create_perturbed_state_dict, + _wait_for_servers_ready, + _init_actor_process_group, +) + + +def timed_broadcast_fast_llm( + init_method: str, + model_name: str, + server_urls: list, + redis_host: str = "localhost", + redis_port: int = 6379, + world_size: int = 2, +): + """Timed broadcast using Fast-LLM protocol: perturbed → original → perturbed with delays. + + This simulates Fast-LLM's weight broadcast protocol where weight updates are signaled + via Redis stream and broadcast using broadcast_object_list + broadcast. + + Pattern: original (server default) → perturbed → original → perturbed + + Args: + init_method: Distributed init method + model_name: Model name to load + server_urls: Base URLs of vLLM server(s) (for health check only) + redis_host: Redis host address + redis_port: Redis port number + world_size: Total NCCL world size (trainer rank 0 + all vLLM workers) + """ + import torch + import torch.distributed as dist + import time + import redis + import orjson + + from fast_llm.engine.distributed.config import DistributedBackend + from fast_llm.engine.distributed.distributed import ProcessGroupPool + + print(f"[Trainer] Initializing process group as rank 0 (world_size={world_size})") + process_group = ProcessGroupPool( + rank=0, + world_size=world_size, + local_world_size=1, + init_method=init_method, + backend=DistributedBackend.nccl, + ).get_process_group(range(world_size), 0) + print("[Trainer] Process group initialized") + + # Connect to Redis + print(f"[Trainer] Connecting to Redis at {redis_host}:{redis_port}") + r = redis.Redis(host=redis_host, port=redis_port) + stream_key = "fast_llm_events" + payload_key = "event" + print(f"[Trainer] Connected to Redis, will write to stream '{stream_key}'") + + _wait_for_servers_ready(server_urls, extra_wait_secs=15) + + # Load weights + print(f"[Trainer] Loading original weights from {model_name}") + original_state_dict, _ = _load_state_dict(model_name) + perturbed_state_dict = _create_perturbed_state_dict(original_state_dict) + + from fast_llm.core.distributed import broadcast as _broadcast, broadcast_object as _broadcast_object + + # Helper function to broadcast weights using Fast-LLM protocol + def broadcast_weights_fast_llm(state_dict, step): + """Broadcast weights using Fast-LLM protocol. + + Protocol: + 1. Send Redis event: {type: "weights_ready", step: N} + 2. For each parameter: + - broadcast_object((shard_name, layer_name, shape, dtype)) + - broadcast(tensor) + 3. Send end signal: broadcast_object(None) + """ + # Send Redis stream event + event = {"type": "weights_ready", "step": step} + r.xadd(stream_key, {payload_key: orjson.dumps(event)}) + print(f"[Trainer] Sent Redis event to '{stream_key}': {event}") + + # Broadcast each parameter + for i, (name, tensor) in enumerate(state_dict.items()): + if tensor.device.type != "cuda": + tensor = tensor.cuda(0) + + _broadcast_object(("weights", name, list(tensor.shape), str(tensor.dtype)), process_group, src=0) + _broadcast(tensor, 0, process_group) + + if (i + 1) % 50 == 0: + print(f"[Trainer] Broadcasted {i+1}/{len(state_dict)} parameters") + + # Send end signal + _broadcast_object(None, process_group, src=0) + print(f"[Trainer] Sent end signal, broadcast complete") + + # Broadcast 1: Perturbed weights + print(f"[Trainer] Broadcasting {len(perturbed_state_dict)} perturbed parameters") + broadcast_weights_fast_llm(perturbed_state_dict, step=1) + print("[Trainer] Perturbed weights broadcast complete") + + print("[Trainer] Waiting 5 seconds before broadcasting original weights...") + time.sleep(5) + + # Broadcast 2: Original weights + print(f"[Trainer] Broadcasting {len(original_state_dict)} original parameters") + broadcast_weights_fast_llm(original_state_dict, step=2) + print("[Trainer] Original weights broadcast complete") + + print("[Trainer] Waiting 5 seconds before broadcasting perturbed weights again...") + time.sleep(5) + + # Broadcast 3: Perturbed weights again (same as first) + print(f"[Trainer] Broadcasting {len(perturbed_state_dict)} perturbed parameters (2nd time)") + broadcast_weights_fast_llm(perturbed_state_dict, step=3) + print("[Trainer] Perturbed weights broadcast complete (2nd time)") + + # Wait to allow generation with the last broadcast before tearing down + print("[Trainer] Waiting 5 seconds for generation with final weights...") + time.sleep(5) + + # Signal training is finished so vLLM workers destroy their side of the process group + print("[Trainer] Sending training_finished signal...") + r.xadd(stream_key, {payload_key: orjson.dumps({"type": "training_finished"})}) + + # Cleanup — destroy_process_group now resolves because vLLM workers respond to training_finished + r.close() + process_group.shutdown() + print("[Trainer] Redis connection closed, process group destroyed, exiting") + + +def rapid_broadcast_cycles_fast_llm( + init_method: str, + model_name: str, + server_urls: list, + redis_host: str = "localhost", + redis_port: int = 6379, + world_size: int = 2, + n_cycles: int = 6, +): + """Hybrid Fast-LLM broadcast designed to catch transition/garbage generations. + + Structure: + 1. Slow broadcast: perturbed (5 s wait after) — establishes text_B + 2. Slow broadcast: original (5 s wait after) — re-establishes text_A + 3. n_cycles rapid pairs: perturbed → original (1 s between each) + 4. Slow broadcast: perturbed (5 s wait after) — end on text_B so the + overall A→B→A→B pattern remains detectable + """ + import torch.distributed as dist + import time + import redis as redis_lib + import orjson + + from fast_llm.engine.distributed.config import DistributedBackend + from fast_llm.engine.distributed.distributed import ProcessGroupPool + + print(f"[Trainer] Initializing process group as rank 0 (world_size={world_size})") + process_group = ProcessGroupPool( + rank=0, + world_size=world_size, + local_world_size=1, + init_method=init_method, + backend=DistributedBackend.nccl, + ).get_process_group(range(world_size), 0) + print("[Trainer] Process group initialized") + + r = redis_lib.Redis(host=redis_host, port=redis_port) + stream_key = "fast_llm_events" + payload_key = "event" + + _wait_for_servers_ready(server_urls, extra_wait_secs=15) + + print(f"[Trainer] Loading weights from {model_name}") + original_state_dict, _ = _load_state_dict(model_name) + perturbed_state_dict = _create_perturbed_state_dict(original_state_dict) + + step = 1 + + def broadcast_weights(state_dict, label): + nonlocal step + import torch + event = {"type": "weights_ready", "step": step} + r.xadd(stream_key, {payload_key: orjson.dumps(event)}) + print(f"[Trainer] Sent weights_ready step={step} ({label})") + step += 1 + + from fast_llm.core.distributed import broadcast as _broadcast, broadcast_object as _broadcast_object + + for name, tensor in state_dict.items(): + if tensor.device.type != "cuda": + tensor = tensor.cuda(0) + _broadcast_object(("weights", name, list(tensor.shape), str(tensor.dtype)), process_group, src=0) + _broadcast(tensor, 0, process_group) + + _broadcast_object(None, process_group, src=0) + print(f"[Trainer] Broadcast complete ({label})") + + # --- Slow cycle: establish text_B and text_A clearly --- + print("[Trainer] Slow broadcast 1: perturbed (establishing text_B)...") + broadcast_weights(perturbed_state_dict, "perturbed slow") + time.sleep(5) + + print("[Trainer] Slow broadcast 2: original (re-establishing text_A)...") + broadcast_weights(original_state_dict, "original slow") + time.sleep(5) + + # --- Rapid cycles: 1 s between broadcasts --- + for i in range(n_cycles): + print(f"[Trainer] Rapid cycle {i + 1}/{n_cycles}: perturbed...") + broadcast_weights(perturbed_state_dict, f"perturbed rapid {i + 1}") + time.sleep(1) + + print(f"[Trainer] Rapid cycle {i + 1}/{n_cycles}: original...") + broadcast_weights(original_state_dict, f"original rapid {i + 1}") + time.sleep(1) + + # --- Final slow broadcast: end on perturbed so ABAB pattern holds --- + print("[Trainer] Final slow broadcast: perturbed (ending on text_B)...") + broadcast_weights(perturbed_state_dict, "perturbed final") + time.sleep(5) + + print("[Trainer] Sending training_finished signal...") + r.xadd(stream_key, {payload_key: orjson.dumps({"type": "training_finished"})}) + + r.close() + process_group.shutdown() + print("[Trainer] Redis connection closed, process group destroyed, exiting") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Fast-LLM trainer helper") + parser.add_argument("--init-method", required=True, help="Distributed init method") + parser.add_argument("--model", required=True, help="Model name") + parser.add_argument("--server-urls", nargs="+", required=True, help="Server URL(s)") + parser.add_argument("--world-size", type=int, default=2, help="Total distributed world size") + parser.add_argument("--redis-host", default="localhost", help="Redis host") + parser.add_argument("--redis-port", type=int, default=6379, help="Redis port") + parser.add_argument("--n-cycles", type=int, default=0, + help="If > 0, run rapid_broadcast_cycles with this many rapid pairs") + + args = parser.parse_args() + + if args.n_cycles > 0: + rapid_broadcast_cycles_fast_llm( + init_method=args.init_method, + model_name=args.model, + server_urls=args.server_urls, + redis_host=args.redis_host, + redis_port=args.redis_port, + world_size=args.world_size, + n_cycles=args.n_cycles, + ) + else: + timed_broadcast_fast_llm( + init_method=args.init_method, + model_name=args.model, + server_urls=args.server_urls, + redis_host=args.redis_host, + redis_port=args.redis_port, + world_size=args.world_size, + ) diff --git a/tests/server_weight_update_utils.py b/tests/server_weight_update_utils.py new file mode 100644 index 00000000..a4ade92a --- /dev/null +++ b/tests/server_weight_update_utils.py @@ -0,0 +1,653 @@ +"""Shared utilities for server weight update integration tests.""" + +import asyncio +import requests +import time +from pathlib import Path +import subprocess +import sys +import os + + +async def wait_for_server_ready(server_url: str, server_proc, trainer_proc, timeout_seconds: int = 300): + """Wait for server to be ready by polling health endpoint. + + Args: + server_url: Base URL of server (e.g., "http://127.0.0.1:8000") + server_proc: Server subprocess + trainer_proc: Trainer subprocess + timeout_seconds: Maximum time to wait + + Returns: + True if server is ready + + Raises: + RuntimeError: If server or trainer process terminates + TimeoutError: If server doesn't become ready within timeout + """ + print("[Main] Waiting for server to be ready...") + for i in range(timeout_seconds): + # Check if server process crashed + if server_proc.poll() is not None: + print(f"[Main] Server process terminated with code {server_proc.returncode}") + raise RuntimeError(f"Server process terminated with code {server_proc.returncode}") + + # Check if trainer process crashed + if trainer_proc.poll() is not None: + print(f"[Main] Trainer process terminated with code {trainer_proc.returncode}") + raise RuntimeError(f"Trainer process terminated with code {trainer_proc.returncode}") + + try: + resp = requests.get(f"{server_url}/health", timeout=1) + if resp.status_code == 200: + print("[Main] Server is ready!") + return True + except requests.exceptions.RequestException: + pass + + if i % 10 == 0: + print(f"[Main] Still waiting for server... ({i} seconds)") + await asyncio.sleep(1) + + raise TimeoutError(f"Server did not become ready within {timeout_seconds} seconds") + + +def _build_phases(generations): + """Collapse a generation list into (text, items) phase tuples.""" + phases = [] + current_text = None + current_phase = [] + for ts, text in generations: + if text != current_text: + if current_phase: + phases.append((current_text, current_phase)) + current_text = text + current_phase = [(ts, text)] + else: + current_phase.append((ts, text)) + if current_phase: + phases.append((current_text, current_phase)) + return phases + + +def _identify_stable_texts(phases, min_stable_gens=5): + """Return (text_a, text_b) identified from the first two stable phases. + + Iterates phases in order, skipping any with fewer than ``min_stable_gens`` + generations (transition artifacts). The first stable phase gives text_A; + the first stable phase with a different text gives text_B. + + Returns (text_a, text_b) or (None, None) if two distinct stable texts + cannot be found. + """ + text_a = None + text_b = None + for text, items in phases: + if len(items) < min_stable_gens: + continue + if text_a is None: + text_a = text + elif text != text_a: + text_b = text + break + if text_a is None or text_b is None: + return None, None + return text_a, text_b + + +def _find_abab_pattern(phases, min_stable_gens=5): + """Search for the A→B→A→B pattern. + + text_A and text_B are identified from the first two *stable* phases — + phases with at least ``min_stable_gens`` generations. Short transition + phases (1–few gens) produced while an NCCL broadcast is in-flight are + automatically skipped during identification. + + The test is designed so that the server always starts with a long run of + text_A (original weights, typically hundreds of gens) followed by a long + run of text_B (first perturbed broadcast, tens of gens), making them + unambiguous even with transition artifacts in between. + + After identifying text_A and text_B the full A→B→A→B subsequence is + located in the phase list (transition phases between the four anchors are + silently skipped). + + Returns (phase_a, phase_b, phase_a2, phase_b2) or None. + """ + if len(phases) < 4: + return None + + text_a, text_b = _identify_stable_texts(phases, min_stable_gens) + + if text_a is None or text_b is None: + return None + + texts = [t for t, _ in phases] + + # Find ABAB as a subsequence in the phase list + first_a = next((i for i, t in enumerate(texts) if t == text_a), None) + if first_a is None: + return None + + first_b = next((i for i in range(first_a + 1, len(phases)) if texts[i] == text_b), None) + if first_b is None: + return None + + second_a = next((i for i in range(first_b + 1, len(phases)) if texts[i] == text_a), None) + if second_a is None: + return None + + second_b = next((i for i in range(second_a + 1, len(phases)) if texts[i] == text_b), None) + if second_b is None: + return None + + return phases[first_a], phases[first_b], phases[second_a], phases[second_b] + + +def check_pattern_detected(generations): + """Check whether the full A→B→A→B pattern is present in the generation history. + + This is a **post-hoc analysis helper** (e.g. for assertions after the + generation loop ends). It is intentionally *not* used as an early-stop + signal inside the generation loops. + + Why not early-stop? Any transition artifact text T that happens to appear + with several consecutive identical generations (possible when NCCL broadcasts + are slow) is indistinguishable from the real perturbed text B at generation + time. False positives would cut the loop short before the final stable B + phase accumulates. The generation loops instead rely on the trainer process + exiting (``trainer_proc.poll() is not None``) as their sole reliable + termination signal — the trainer exits within milliseconds of completing its + last broadcast, so no significant extra generation happens. + + Args: + generations: List of (timestamp, text) tuples + + Returns: + True if the A→B→A→B pattern is present + """ + if len(generations) < 4: + return False + phases = _build_phases(generations) + return _find_abab_pattern(phases) is not None + + +async def run_generation_loop( + server_url: str, + model_name: str, + simple_prompt: str, + generation_config: dict, + trainer_proc, + max_duration: int = 120, + generation_interval: float = 0.5, +): + """Run continuous generation loop until pattern is detected or timeout. + + Args: + server_url: Base URL of server + model_name: Model name for API request + simple_prompt: Prompt to generate from + generation_config: Config dict with max_tokens, etc. + trainer_proc: Trainer subprocess to monitor + max_duration: Maximum duration in seconds + generation_interval: Time between generations + + Returns: + List of (timestamp, generated_text) tuples + """ + print("[Main] Starting continuous generation loop...") + generations = [] + start_time = time.time() + + while time.time() - start_time < max_duration: + # Check if trainer is still running + trainer_poll = trainer_proc.poll() + if trainer_poll is not None: + print(f"[Main] Trainer exited with code {trainer_poll}") + break + + try: + # Generate via HTTP API + payload = { + "model": model_name, + "prompt": simple_prompt, + "max_tokens": generation_config["max_tokens"], + "temperature": 0.0, # Deterministic + "top_p": 1.0, + "seed": 42, + } + + resp = requests.post( + f"{server_url}/v1/completions", + json=payload, + timeout=30, + ) + + if resp.status_code == 200: + result = resp.json() + generated_text = result["choices"][0]["text"] + timestamp = time.time() - start_time + generations.append((timestamp, generated_text)) + print(f"[Main] [{timestamp:.1f}s] Generated: '{generated_text}'") + else: + print(f"[Main] Generation failed with status {resp.status_code}") + + except requests.exceptions.RequestException as e: + print(f"[Main] Request failed: {e}") + + await asyncio.sleep(generation_interval) + + return generations + + +def analyze_and_verify_pattern(generations): + """Analyze generation sequence and verify the expected A→B→A→B pattern. + + Tolerates transition-artifact phases (e.g. a single generation produced + while an NCCL broadcast was in-flight) by searching for the pattern as + a subsequence rather than requiring it at exactly positions [0,1,2,3]. + + Args: + generations: List of (timestamp, text) tuples + + Returns: + Tuple of (text_a, text_b) — the original and perturbed texts. + + Raises: + AssertionError: If pattern is not as expected + """ + print("\n" + "=" * 60) + print("GENERATION SEQUENCE ANALYSIS") + print("=" * 60) + print(f"Total generations: {len(generations)}") + + for i, (ts, text) in enumerate(generations): + print(f"[{ts:5.1f}s] Gen {i+1}: '{text[:80]}...'") + + assert len(generations) >= 4, ( + f"Not enough generations to verify pattern (need at least 4, got {len(generations)})" + ) + + phases = _build_phases(generations) + + _GRAY = "\033[90m" + _RESET = "\033[0m" + stable_a, stable_b = _identify_stable_texts(phases) + stable_texts = {t for t in (stable_a, stable_b) if t is not None} + print("\n" + "=" * 60) + print(f"Detected {len(phases)} phase(s):") + for i, (text, items) in enumerate(phases): + line = f"Phase {i+1}: {len(items)} generation(s) - '{text[:60]}...'" + if text not in stable_texts: + print(f"{_GRAY}{line} ← transition{_RESET}") + else: + print(line) + print("=" * 60) + + result = _find_abab_pattern(phases) + assert result is not None, ( + f"Could not find A→B→A→B pattern in {len(phases)} phase(s). " + f"Phases: {[(text[:40], len(items)) for text, items in phases]}" + ) + + (phase_a_text, phase_a_items), (phase_b_text, phase_b_items), \ + (phase_a2_text, phase_a2_items), (phase_b2_text, phase_b2_items) = result + + # These hold by construction from _find_abab_pattern, but assert for clarity + assert phase_a_text != phase_b_text, "Phase A and Phase B should be different" + assert phase_a2_text == phase_a_text, "Second A should match first A (original weights restored)" + assert phase_b2_text == phase_b_text, "Second B should match first B (perturbed weights reapplied)" + + skipped = len(phases) - 4 + skip_note = f" ({skipped} transition phase(s) skipped)" if skipped else "" + print(f"\n✓ Pattern verified{skip_note}:") + print(f" Phase A (original): {len(phase_a_items)} generation(s)") + print(f" Phase B (perturbed): {len(phase_b_items)} generation(s)") + print(f" Phase A2 (original): {len(phase_a2_items)} generation(s) ← matches A ✓") + print(f" Phase B2 (perturbed): {len(phase_b2_items)} generation(s) ← matches B ✓") + + return phase_a_text, phase_b_text + + +def analyze_and_verify_pattern_multi(per_server_generations): + """Verify A→B→A→B pattern independently per server, then check consistency. + + Each server's generation history is checked independently (since weight + updates are not coordinated with requests, servers can transiently disagree). + After all pass, we assert that every server converged on the same text A + and text B. + + Args: + per_server_generations: List of per-server generation lists, each a + list of (timestamp, text) tuples (as returned by + run_generation_loop_multi). + + Raises: + AssertionError: If any server fails its pattern check or servers + disagree on text A / text B. + """ + patterns = [] + for i, generations in enumerate(per_server_generations): + print(f"\n{'=' * 60}") + print(f"Actor {i} pattern analysis") + text_a, text_b = analyze_and_verify_pattern(generations) + patterns.append((text_a, text_b)) + + unique_a = set(t_a for t_a, _ in patterns) + unique_b = set(t_b for _, t_b in patterns) + assert len(unique_a) == 1, ( + f"Servers disagree on text A (original weights): " + f"{[t_a[:40] for t_a, _ in patterns]}" + ) + assert len(unique_b) == 1, ( + f"Servers disagree on text B (perturbed weights): " + f"{[t_b[:40] for _, t_b in patterns]}" + ) + print(f"\n✓ All {len(patterns)} actor(s) agree on text A and text B") + + +def extract_transition_phases(generations, text_a, text_b): + """Return phases that are neither text_a nor text_b. + + These are mid-broadcast 'garbage' generations produced while an NCCL + weight update was in flight and the model had partially updated weights. + + Args: + generations: List of (timestamp, text) tuples + text_a: The original-weights text (established first) + text_b: The perturbed-weights text + + Returns: + List of (text, items) phase tuples where text is neither text_a nor text_b + """ + phases = _build_phases(generations) + return [(text, items) for text, items in phases if text != text_a and text != text_b] + + +def analyze_and_verify_transitions(generations, n_cycles): + """Verify A→B→A→B pattern and assert that transition generations were caught. + + The ``rapid_broadcast_cycles`` trainer command performs: + - 1 startup A phase (server starts on original weights) + - 1 slow perturbed broadcast → text_B + - 1 slow original broadcast → text_A + - n_cycles rapid pairs → text_B, text_A each cycle + - 1 final slow perturbed → text_B + + This gives exactly ``4 + 2 * n_cycles`` stable phases. Seeing fewer + means a broadcast was missed entirely (timing/sync bug). + + Args: + generations: List of (timestamp, text) tuples + n_cycles: Number of rapid broadcast pairs (passed as ``--n-cycles`` + to the trainer helper). + + Raises: + AssertionError: If ABAB pattern is not found, stable phase count is + wrong, or no transition generations were caught. + """ + text_a, text_b = analyze_and_verify_pattern(generations) + + phases = _build_phases(generations) + stable_phases = [(text, items) for text, items in phases if text == text_a or text == text_b] + expected_stable = 4 + 2 * n_cycles + assert len(stable_phases) == expected_stable, ( + f"Expected {expected_stable} stable phases (4 + 2×{n_cycles} cycles) " + f"but found {len(stable_phases)}. " + f"A broadcast may have been missed or merged. " + f"Stable phase counts: {[len(items) for _, items in stable_phases]}" + ) + print(f"\n✓ Stable phase count correct: {len(stable_phases)} (expected {expected_stable})") + + transition_phases = extract_transition_phases(generations, text_a, text_b) + + print("\n" + "=" * 60) + print(f"TRANSITION / GARBAGE GENERATIONS: {len(transition_phases)} phase(s)") + print("=" * 60) + if transition_phases: + for i, (text, items) in enumerate(transition_phases): + ts_start = items[0][0] + ts_end = items[-1][0] + print(f" [{i + 1}] {len(items)} gen(s) @ {ts_start:.2f}s–{ts_end:.2f}s: '{text[:120]}'") + else: + print(" (none)") + + assert len(transition_phases) > 0, ( + "No transition generations were caught. " + "Try increasing --n-cycles or verify generation_interval=0.0 is set." + ) + print(f"\n✓ Caught {len(transition_phases)} transition phase(s)") + + +def start_vllm_server( + model_name: str, + server_port: int, + distributed_init_method: str, + stream_process_output_fn, + extra_args: list = None, + gpu_ids: str = "0", + actor_llm_idx: int = 0, + world_size: int = 2, + tensor_parallel_size: int = 1, +): + """Start vLLM HTTP server subprocess. + + Args: + model_name: Model to load + server_port: Port to bind to + distributed_init_method: Distributed initialization method + stream_process_output_fn: Function to stream process output + extra_args: Additional CLI arguments (e.g., ["--weight-update-mode", "fast-llm"]) + gpu_ids: CUDA_VISIBLE_DEVICES value (e.g., "0" or "0,1") + actor_llm_idx: Actor index for this vLLM instance + world_size: Total distributed world size + tensor_parallel_size: Number of GPUs for tensor parallelism + + Returns: + Tuple of (server_proc, stdout_thread, stderr_thread) + """ + vllm_env = os.environ.copy() + vllm_env["CUDA_VISIBLE_DEVICES"] = gpu_ids + vllm_env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + + print(f"[Main] Starting vLLM HTTP server on port {server_port} (GPU(s) {gpu_ids}, actor_idx={actor_llm_idx}, TP={tensor_parallel_size})") + vllm_entry_point = Path(__file__).parent.parent / "pipelinerl" / "entrypoints" / "run_vllm1.py" + + cmd = [ + sys.executable, + str(vllm_entry_point), + "--model", model_name, + "--port", str(server_port), + "--host", "127.0.0.1", + "--actor-llm-idx", str(actor_llm_idx), + "--weight-update-group-init-method", distributed_init_method, + "--weight-update-group-world-size", str(world_size), + "--tensor-parallel-size", str(tensor_parallel_size), + ] + + if extra_args: + cmd.extend(extra_args) + + server_proc = subprocess.Popen( + cmd, + env=vllm_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + print("[Main] Starting server output streaming...") + stdout_thread, stderr_thread = stream_process_output_fn(server_proc, f"vLLM Server (actor {actor_llm_idx})") + + return server_proc, stdout_thread, stderr_thread + + +async def wait_for_all_servers_ready( + server_urls: list, + server_procs: list, + trainer_proc, + timeout_seconds: int = 300, +): + """Wait for all servers to be ready by polling their health endpoints. + + Args: + server_urls: List of server base URLs + server_procs: List of server subprocesses (same order as server_urls) + trainer_proc: Trainer subprocess + timeout_seconds: Maximum time to wait per server + + Returns: + True if all servers are ready + + Raises: + RuntimeError: If any process terminates unexpectedly + TimeoutError: If any server doesn't become ready within timeout + """ + for url, proc in zip(server_urls, server_procs): + await wait_for_server_ready(url, proc, trainer_proc, timeout_seconds) + return True + + +async def run_generation_loop_multi( + server_urls: list, + model_name: str, + simple_prompt: str, + generation_config: dict, + trainer_proc, + max_duration: int = 120, + generation_interval: float = 0.5, +): + """Run continuous generation loop querying all servers each round. + + Each server is tracked independently because weight updates and requests + are not coordinated — different actors can temporarily return different + results while a broadcast is in flight. Pattern checking is therefore + done per-server after the loop (see analyze_and_verify_pattern_multi). + + Args: + server_urls: List of server base URLs + model_name: Model name for API request + simple_prompt: Prompt to generate from + generation_config: Config dict with max_tokens, etc. + trainer_proc: Trainer subprocess to monitor + max_duration: Maximum duration in seconds + generation_interval: Time between generation rounds + + Returns: + List of per-server generation lists, each a list of + (timestamp, generated_text) tuples (same order as server_urls). + """ + print(f"[Main] Starting continuous generation loop across {len(server_urls)} server(s)...") + per_server = [[] for _ in server_urls] + start_time = time.time() + + payload = { + "model": model_name, + "prompt": simple_prompt, + "max_tokens": generation_config["max_tokens"], + "temperature": 0.0, + "top_p": 1.0, + "seed": 42, + } + + while time.time() - start_time < max_duration: + # Check if trainer is still running + trainer_poll = trainer_proc.poll() + if trainer_poll is not None: + print(f"[Main] Trainer exited with code {trainer_poll}") + break + + for i, url in enumerate(server_urls): + try: + resp = requests.post( + f"{url}/v1/completions", + json=payload, + timeout=30, + ) + if resp.status_code == 200: + text = resp.json()["choices"][0]["text"] + timestamp = time.time() - start_time + per_server[i].append((timestamp, text)) + print(f"[Main] [{timestamp:.1f}s] Actor {i}: '{text}'") + else: + print(f"[Main] Generation from actor {i} ({url}) failed with status {resp.status_code}") + except requests.exceptions.RequestException as e: + print(f"[Main] Request to actor {i} ({url}) failed: {e}") + + await asyncio.sleep(generation_interval) + + return per_server + + +def start_trainer_process( + trainer_helper_path: Path, + distributed_init_method: str, + model_name: str, + server_urls: list, + stream_process_output_fn, + extra_args: list = None, + gpu_id: str = "1", + world_size: int = 2, + command: str = "timed_broadcast_server_test", +): + """Start trainer subprocess. + + Args: + trainer_helper_path: Path to trainer helper script + distributed_init_method: Distributed initialization method + model_name: Model name + server_urls: List of server URLs (one per actor) + stream_process_output_fn: Function to stream process output + extra_args: Additional CLI arguments (e.g., ["--n-cycles", "6"]) + gpu_id: CUDA_VISIBLE_DEVICES value for the trainer GPU + world_size: Total distributed world size + command: Positional command for distributed_trainer_helper.py + (ignored for fast_llm_trainer_helper.py which uses --init-method style) + + Returns: + Tuple of (trainer_proc, stdout_thread, stderr_thread) + """ + trainer_env = os.environ.copy() + trainer_env["CUDA_VISIBLE_DEVICES"] = gpu_id + + print(f"[Main] Starting trainer process (GPU {gpu_id}) for process group rendezvous") + + cmd = [ + sys.executable, + str(trainer_helper_path), + ] + + # Check which trainer helper is being used by the script name + if "fast_llm" in str(trainer_helper_path): + # fast_llm_trainer_helper.py uses argparse with --init-method, --model, etc. + cmd.extend([ + "--init-method", distributed_init_method, + "--model", model_name, + "--world-size", str(world_size), + "--server-urls", + ] + list(server_urls)) + else: + # distributed_trainer_helper.py uses positional command + flags + cmd.extend([ + command, + "--init-method", distributed_init_method, + "--model-name", model_name, + "--world-size", str(world_size), + "--server-urls", + ] + list(server_urls)) + + if extra_args: + cmd.extend(extra_args) + + trainer_proc = subprocess.Popen( + cmd, + env=trainer_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + print("[Main] Starting trainer output streaming...") + stdout_thread, stderr_thread = stream_process_output_fn(trainer_proc, "Trainer") + + return trainer_proc, stdout_thread, stderr_thread diff --git a/tests/sync_helper.py b/tests/sync_helper.py new file mode 100644 index 00000000..73ebed82 --- /dev/null +++ b/tests/sync_helper.py @@ -0,0 +1,110 @@ +"""Simple file-based synchronization for distributed test processes.""" + +import time +from pathlib import Path + + +class SyncPoint: + """File-based synchronization point for coordinating subprocesses.""" + + def __init__(self, sync_dir: Path, name: str): + """Create a sync point. + + Args: + sync_dir: Directory for sync files + name: Name of this sync point (e.g., "baseline_done") + """ + self.sync_file = sync_dir / f"{name}.sync" + self.sync_dir = sync_dir + + def signal(self): + """Signal that this point is reached.""" + self.sync_file.touch() + # Force filesystem sync to ensure file is visible immediately + import os + fd = os.open(str(self.sync_file.parent), os.O_RDONLY) + os.fsync(fd) + os.close(fd) + print(f"[Sync] Signaled: {self.sync_file.name}") + + def wait(self, timeout: float = 60): + """Wait for this point to be signaled. + + Args: + timeout: Maximum time to wait in seconds + + Raises: + TimeoutError: If sync point not reached within timeout + """ + start = time.time() + while not self.sync_file.exists(): + if time.time() - start > timeout: + raise TimeoutError( + f"Timeout waiting for sync point: {self.sync_file.name}" + ) + time.sleep(0.1) + print(f"[Sync] Reached: {self.sync_file.name}") + + def clear(self): + """Clear this sync point.""" + if self.sync_file.exists(): + self.sync_file.unlink() + + +def create_sync_dir(base_dir: Path) -> Path: + """Create a directory for sync files. + + Args: + base_dir: Base temporary directory + + Returns: + Path to sync directory + """ + sync_dir = base_dir / "sync" + sync_dir.mkdir(exist_ok=True) + return sync_dir + + +def write_weight_update_request(sync_dir: Path, request): + """Write WeightUpdateRequest to JSON file. + + Args: + sync_dir: Sync directory + request: WeightUpdateRequest object + """ + import json + + request_file = sync_dir / "weight_update_request.json" + with open(request_file, "w") as f: + json.dump(request.model_dump(), f) + print(f"[Sync] Wrote weight update request to {request_file.name}") + + +def read_weight_update_request(sync_dir: Path): + """Read WeightUpdateRequest from JSON file. + + Args: + sync_dir: Sync directory + + Returns: + WeightUpdateRequest object + """ + import json + from pipelinerl.finetune_loop import WeightUpdateRequest + + request_file = sync_dir / "weight_update_request.json" + + # Wait for file to exist + import time + timeout = 60 + start = time.time() + while not request_file.exists(): + if time.time() - start > timeout: + raise TimeoutError(f"Timeout waiting for {request_file.name}") + time.sleep(0.1) + + with open(request_file, "r") as f: + data = json.load(f) + + print(f"[Sync] Read weight update request from {request_file.name}") + return WeightUpdateRequest(**data) diff --git a/tests/test_actor_error_handling.py b/tests/test_actor_error_handling.py new file mode 100644 index 00000000..61adc5eb --- /dev/null +++ b/tests/test_actor_error_handling.py @@ -0,0 +1,290 @@ +"""Test that actor rollout error handling doesn't crash the entire actor. + +Specifically tests that: +1. HTTP 4xx errors from vLLM (e.g., max_tokens too large) are handled gracefully +2. Groups where ALL rollouts fail are dropped (not submitted) +3. Groups where SOME rollouts fail submit only valid results +4. HTTP 5xx errors still propagate as fatal +""" + +import asyncio +import queue +from unittest.mock import MagicMock, AsyncMock, patch + +import aiohttp +import pytest +from omegaconf import OmegaConf + +from pipelinerl.rollouts import BaseMetrics, RolloutResult, TrainingText + + +# --------------------------------------------------------------------------- +# Helpers – lightweight stand-ins for heavy classes used by schedule_rollouts +# --------------------------------------------------------------------------- + +class FakeQueue: + """Minimal stand-in for SharedMemoryQueue (no shared memory needed).""" + + def __init__(self): + self._q = queue.Queue() + + def put(self, item, block=True, timeout=None): + self._q.put(item) + + def get(self, block=True, timeout=None): + return self._q.get(block=block, timeout=timeout) + + def qsize(self): + return self._q.qsize() + + def max_actual_entry_size(self): + return 0 + + def get_memory_size(self): + return 0 + + +class FakeTrainerState: + def __init__(self): + self.propagated_weight_version = 1 + self.samples_processed = 0 + + +def make_good_result() -> RolloutResult: + """A valid rollout result with one training sample.""" + return RolloutResult( + training_texts=[ + TrainingText( + text="prompt output", + n_predicted=6, + reward=1.0, + input_ids=[1, 2, 3], + labels=[-100, 2, 3], + finished=True, + prompt_tokens=5, + output_tokens=6, + ) + ], + metrics=BaseMetrics(reward=1.0, success=True, no_error=True, no_answer=False), + latency=0.5, + ) + + +def make_client_response_error(status: int, message: str = "Bad Request"): + """Create an aiohttp.ClientResponseError.""" + mock_req = MagicMock() + mock_req.url = "http://localhost:8080/v1/chat/completions" + return aiohttp.ClientResponseError( + request_info=mock_req, + history=(), + status=status, + message=message, + ) + + +# --------------------------------------------------------------------------- +# Core test: exercise rollout_and_maybe_produce_result + group completion +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_all_rollouts_fail_group_dropped(): + """When all rollouts in a group fail with 4xx, the group should be dropped.""" + attempts = 4 + problem_q = FakeQueue() + result_q = FakeQueue() + trainer_state = FakeTrainerState() + + # Put one problem in the queue + problem_q.put({"task": "What is 2+2?", "answer": "4"}) + + call_count = 0 + + async def failing_rollout_policy(cfg, llm, problem, session): + nonlocal call_count + call_count += 1 + raise make_client_response_error(400, "max_tokens too large") + + cfg = OmegaConf.create({ + "actor": { + "rollout_policy": "not_used", # we patch it + "llm_max_rollouts": 64, + }, + "finetune": { + "train_batch_size": 1000, + "gradient_accumulation_passes": 1, + "train_iters": 100, + "interrupt_train_steps": None, + }, + "debug": {}, + }) + + llms = [MagicMock()] # 1 LLM + + # We can't easily run schedule_rollouts (too many dependencies), + # so we directly test the inner logic by reimplementing the key parts. + # This mirrors rollout_and_maybe_produce_result + group completion. + + group_rollouts = {} + group_id = 0 + group_rollouts[group_id] = [] + finished_rollouts = 0 + warnings_logged = [] + + for rollout_index in range(attempts): + try: + rollout_result = await failing_rollout_policy(cfg, llms[0], {"task": "x"}, None) + except aiohttp.ClientResponseError as e: + if 400 <= e.status < 500: + warnings_logged.append(str(e.status)) + rollout_result = RolloutResult( + training_texts=[], + metrics=BaseMetrics(reward=0.0, success=False, no_error=False, no_answer=True), + latency=0.0, + ) + else: + raise + + rollout_result.model_version = 1 + rollout_result.group_id = f"test_{group_id}" + group_rollouts[group_id].append(rollout_result) + + # Now check group completion logic + assert len(group_rollouts[group_id]) == attempts + valid_results = [r for r in group_rollouts[group_id] if r.training_texts] + + # All failed → group should be dropped + assert len(valid_results) == 0, "Expected all results to be empty" + assert call_count == attempts + assert len(warnings_logged) == attempts + + # In real code: del group_rollouts[group_id], don't put in result_q + del group_rollouts[group_id] + assert result_q.qsize() == 0, "No group should be in the result queue" + + +@pytest.mark.asyncio +async def test_partial_failure_submits_valid_only(): + """When some rollouts fail but others succeed, submit only valid ones.""" + attempts = 4 + result_q = FakeQueue() + + call_count = 0 + + async def mixed_rollout_policy(cfg, llm, problem, session): + nonlocal call_count + call_count += 1 + # First 2 calls fail, last 2 succeed + if call_count <= 2: + raise make_client_response_error(400, "max_tokens too large") + return make_good_result() + + group_rollouts = {} + group_id = 0 + group_rollouts[group_id] = [] + + for rollout_index in range(attempts): + try: + rollout_result = await mixed_rollout_policy(None, None, {"task": "x"}, None) + except aiohttp.ClientResponseError as e: + if 400 <= e.status < 500: + rollout_result = RolloutResult( + training_texts=[], + metrics=BaseMetrics(reward=0.0, success=False, no_error=False, no_answer=True), + latency=0.0, + ) + else: + raise + + rollout_result.model_version = 1 + rollout_result.group_id = f"test_{group_id}" + group_rollouts[group_id].append(rollout_result) + + assert len(group_rollouts[group_id]) == attempts + + valid_results = [r for r in group_rollouts[group_id] if r.training_texts] + + # 2 failed, 2 succeeded + assert len(valid_results) == 2, f"Expected 2 valid results, got {len(valid_results)}" + + # In real code: result_queue.put(valid_results) + result_q.put(valid_results) + got = result_q.get(block=False) + assert len(got) == 2 + assert all(len(r.training_texts) > 0 for r in got) + + +@pytest.mark.asyncio +async def test_5xx_errors_still_propagate(): + """HTTP 5xx errors should NOT be caught — they indicate server failure.""" + + async def server_error_policy(cfg, llm, problem, session): + raise make_client_response_error(500, "Internal Server Error") + + with pytest.raises(aiohttp.ClientResponseError) as exc_info: + try: + await server_error_policy(None, None, {"task": "x"}, None) + except aiohttp.ClientResponseError as e: + if 400 <= e.status < 500: + pass # Would be caught in real code + else: + raise # 5xx re-raised + + assert exc_info.value.status == 500 + + +@pytest.mark.asyncio +async def test_all_succeed_normal_path(): + """When all rollouts succeed, the full group is submitted.""" + attempts = 4 + result_q = FakeQueue() + + async def good_policy(cfg, llm, problem, session): + return make_good_result() + + group_rollouts = {} + group_id = 0 + group_rollouts[group_id] = [] + + for rollout_index in range(attempts): + try: + rollout_result = await good_policy(None, None, {"task": "x"}, None) + except aiohttp.ClientResponseError as e: + if 400 <= e.status < 500: + rollout_result = RolloutResult( + training_texts=[], + metrics=BaseMetrics(reward=0.0, success=False, no_error=False, no_answer=True), + latency=0.0, + ) + else: + raise + + rollout_result.model_version = 1 + rollout_result.group_id = f"test_{group_id}" + group_rollouts[group_id].append(rollout_result) + + valid_results = [r for r in group_rollouts[group_id] if r.training_texts] + assert len(valid_results) == attempts, "All rollouts should be valid" + + result_q.put(valid_results) + got = result_q.get(block=False) + assert len(got) == attempts + + +@pytest.mark.asyncio +async def test_consumer_assertion_accepts_partial_group(): + """The consumer-side assertion should accept groups with fewer than `attempts` results.""" + attempts = 8 + # Simulate a partial group with 5 valid results + partial_count = 5 + + results = [make_good_result() for _ in range(partial_count)] + + # This mirrors the relaxed assertion in actor.py + assert isinstance(results, list) + assert isinstance(results[0], RolloutResult) + assert 0 < len(results) <= attempts, ( + f"Expected 1-{attempts} rollouts, got {len(results)}" + ) + + group_samples = sum(len(r.training_texts) for r in results) + assert group_samples == partial_count diff --git a/tests/test_launch_process_monitoring.py b/tests/test_launch_process_monitoring.py new file mode 100644 index 00000000..6554f72d --- /dev/null +++ b/tests/test_launch_process_monitoring.py @@ -0,0 +1,65 @@ +from pathlib import Path + +from pipelinerl import launch + + +class FakeTrainerState: + def __init__(self, exp_path: Path, use_fast_llm: bool, weight_broadcast: bool): + self.training_done = True + self.started = False + + def start_listening(self): + self.started = True + + def wait_for_training_done(self, timeout: float | None = None): + raise AssertionError("training_done was already set") + + +class FakeProcessHandle: + def __init__(self, pid: int, kind: str, poll_results: list[int | None]): + self.pid = pid + self.args = [kind] + self._poll_results = poll_results + self.terminated = False + self.waited = False + + def poll(self): + if self.terminated: + return -15 + if self._poll_results: + return self._poll_results.pop(0) + return None + + def wait(self): + self.waited = True + return -15 if self.terminated else 0 + + +def test_watch_processes_stops_remaining_helpers_after_training_completion(monkeypatch, tmp_path): + handles = { + 100: FakeProcessHandle(100, "finetune", [0]), + 101: FakeProcessHandle(101, "redis", [None]), + 102: FakeProcessHandle(102, "actor", [None]), + } + terminated_pids = [] + + def terminate_with_children(pid: int): + terminated_pids.append(pid) + handles[pid].terminated = True + + monkeypatch.setattr(launch, "TrainerState", FakeTrainerState) + monkeypatch.setattr(launch, "terminate_with_children", terminate_with_children) + monkeypatch.setattr(launch.time, "sleep", lambda seconds: None) + + processes = [ + launch.LaunchedProcess(kind="finetune", handle=handles[100]), + launch.LaunchedProcess(kind="redis", handle=handles[101]), + launch.LaunchedProcess(kind="actor", handle=handles[102]), + ] + + launch.watch_processes_running(tmp_path, processes, use_fast_llm=True) + + assert terminated_pids == [101, 102] + assert not handles[100].terminated + assert handles[101].waited + assert handles[102].waited diff --git a/tests/test_model_version.py b/tests/test_model_version.py new file mode 100644 index 00000000..0b98a3eb --- /dev/null +++ b/tests/test_model_version.py @@ -0,0 +1,51 @@ +"""Unit tests for the per-token model_version plumbing. + +Covers the two pure pieces of the (otherwise cluster-only) feature: parsing the optional +`:v` suffix off a `token_id:` string, and the model_version padding / fallback in +`convert_to_fast_llm_format`. +""" + +import pytest + +from pipelinerl.llm import parse_token_id_and_version +from pipelinerl.preprocess import convert_to_fast_llm_format + + +@pytest.mark.parametrize( + "token, expected", + [ + ("token_id:1271", (1271, None)), # no version suffix (backward compatible) + ("token_id:1271:v5", (1271, 5)), + ("token_id:1271:v0", (1271, 0)), # version 0 is a real version, not "absent" + ("token_id:50257:v123456", (50257, 123456)), + ], +) +def test_parse_token_id_and_version(token, expected): + assert parse_token_id_and_version(token) == expected + + +def test_convert_model_version_per_token_left_padded(): + # Completion versions are left-padded to the full sequence with the per-rollout scalar. + entry = {"input_ids": [10, 11, 12, 13, 14, 15], "model_version": 1, "token_versions": [2, 3]} + assert convert_to_fast_llm_format(entry)["model_version"] == [1, 1, 1, 1, 2, 3] + + +def test_convert_model_version_per_token_pads_with_first_when_no_scalar(): + entry = {"input_ids": [10, 11, 12, 13], "token_versions": [7, 8]} + assert convert_to_fast_llm_format(entry)["model_version"] == [7, 7, 7, 8] + + +def test_convert_model_version_scalar_broadcast_fallback(): + # No per-token versions: broadcast the per-rollout scalar across the sequence. + entry = {"input_ids": [10, 11, 12], "model_version": 4, "token_versions": []} + assert convert_to_fast_llm_format(entry)["model_version"] == [4, 4, 4] + + +def test_convert_model_version_absent(): + entry = {"input_ids": [10, 11, 12]} + assert "model_version" not in convert_to_fast_llm_format(entry) + + +def test_convert_model_version_full_completion_no_prompt(): + entry = {"input_ids": [10, 11, 12], "model_version": 9, "token_versions": [2, 3, 4]} + assert convert_to_fast_llm_format(entry)["model_version"] == [2, 3, 4] diff --git a/tests/test_vllm1_fast_llm_broadcast.py b/tests/test_vllm1_fast_llm_broadcast.py new file mode 100644 index 00000000..f7cc410b --- /dev/null +++ b/tests/test_vllm1_fast_llm_broadcast.py @@ -0,0 +1,590 @@ +"""Integration tests for vllm1 with Fast-LLM weight broadcast protocol.""" + +import asyncio +import pytest +import tempfile +from pathlib import Path +from typing import Dict, List +import time +import os +import subprocess +import sys +import signal + +# torch is needed at top level for pytest.mark.skipif decorators +import torch + +# Import shared utilities +from .server_weight_update_utils import ( + wait_for_server_ready, + wait_for_all_servers_ready, + run_generation_loop, + run_generation_loop_multi, + analyze_and_verify_pattern, + analyze_and_verify_pattern_multi, + analyze_and_verify_transitions, + start_vllm_server, + start_trainer_process, +) + +try: + import psutil + + HAS_PSUTIL = True +except ImportError: + HAS_PSUTIL = False + print("WARNING: psutil not available, process tree cleanup will be limited") + + +def stream_process_output(proc, name): + """Start background threads to continuously stream process stdout/stderr. + + Args: + proc: subprocess.Popen object + name: Name for logging prefix (e.g., "vLLM Server", "Trainer") + + Returns: + Tuple of (stdout_thread, stderr_thread) + """ + import threading + + def read_stream(stream, prefix): + """Read from stream and print with prefix.""" + try: + for line in iter(stream.readline, ""): + if line: + print(f"{prefix} {line.rstrip()}", flush=True) + except Exception as e: + print(f"{prefix} [Stream read error: {e}]", flush=True) + + stdout_thread = threading.Thread( + target=read_stream, + args=(proc.stdout, f"[{name} OUT]"), + daemon=True, + ) + stderr_thread = threading.Thread( + target=read_stream, + args=(proc.stderr, f"[{name} ERR]"), + daemon=True, + ) + + stdout_thread.start() + stderr_thread.start() + + return stdout_thread, stderr_thread + + +def kill_process_tree(pid, sig=signal.SIGKILL): + """Kill a process and all its children/grandchildren. + + Args: + pid: Process ID to kill + sig: Signal to send (default SIGKILL) + """ + if not HAS_PSUTIL: + # Fallback: just kill the main process + try: + os.kill(pid, sig) + except ProcessLookupError: + pass + return + + try: + parent = psutil.Process(pid) + except psutil.NoSuchProcess: + return + + # Get all children recursively + children = parent.children(recursive=True) + + # Kill children first + for child in children: + try: + print(f"[Kill] Killing child process {child.pid}") + child.send_signal(sig) + except psutil.NoSuchProcess: + pass + + # Kill parent + try: + parent.send_signal(sig) + except psutil.NoSuchProcess: + pass + + +@pytest.fixture +def fast_llm_trainer_helper(): + """Path to Fast-LLM trainer helper script.""" + return Path(__file__).parent / "fast_llm_trainer_helper.py" + + +@pytest.fixture +def redis_server(): + """Start a Redis server for testing and stop it after the test. + + Returns: + Tuple of (host, port) for the Redis server + """ + import shutil + import socket + + # Check if redis-server is available + redis_server_bin = shutil.which("redis-server") + if not redis_server_bin: + pytest.skip("redis-server not found in PATH") + + # Find an available port + def find_free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(('', 0)) + s.listen(1) + port = s.getsockname()[1] + return port + + redis_port = find_free_port() + redis_host = "localhost" + + print(f"[Redis] Starting Redis server on {redis_host}:{redis_port}") + + # Start Redis server with minimal config + redis_proc = subprocess.Popen( + [ + redis_server_bin, + "--port", str(redis_port), + "--bind", redis_host, + "--save", "", # Disable persistence + "--appendonly", "no", # Disable AOF + "--protected-mode", "no", # Allow connections without password + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Start streaming Redis output + redis_stdout_thread, redis_stderr_thread = stream_process_output(redis_proc, "Redis") + + # Wait for Redis to be ready + import redis + r = redis.Redis(host=redis_host, port=redis_port) + for i in range(30): + try: + r.ping() + print(f"[Redis] Server ready on {redis_host}:{redis_port}") + break + except redis.ConnectionError: + if redis_proc.poll() is not None: + raise RuntimeError(f"Redis server failed to start (exit code {redis_proc.returncode})") + time.sleep(0.1) + else: + redis_proc.kill() + raise TimeoutError("Redis server did not start within 3 seconds") + + try: + yield (redis_host, redis_port) + finally: + # Cleanup + print(f"[Redis] Stopping Redis server (PID {redis_proc.pid})") + redis_proc.terminate() + try: + redis_proc.wait(timeout=5) + except subprocess.TimeoutExpired: + print("[Redis] Redis did not stop gracefully, killing...") + redis_proc.kill() + redis_proc.wait() + print("[Redis] Redis server stopped") + + +# --------------------------------------------------------------------------- +# Module-level helper shared by all Fast-LLM test variants +# --------------------------------------------------------------------------- + +async def _run_fast_llm_server_test( + model_name, + simple_prompt, + generation_config, + init_method, + fast_llm_trainer_helper, + redis_host, + redis_port, + vllm_server_configs, + trainer_gpu, + world_size, + timeout=2400, +): + """Run Fast-LLM server weight-update pattern test with one or more vLLM servers. + + Args: + vllm_server_configs: List of dicts, each with keys: + - port: int + - gpu_ids: str + - actor_llm_idx: int + - tensor_parallel_size: int + trainer_gpu: str, e.g. "1" or "2" + world_size: total NCCL world size (trainer + all vLLM workers) + redis_host: Redis host address + redis_port: Redis port number + """ + server_procs = [] + server_urls = [] + + fast_llm_server_args = [ + "--weight-update-mode", "fast-llm", + "--redis-host", redis_host, + "--redis-port", str(redis_port), + ] + + for cfg in vllm_server_configs: + port = cfg["port"] + url = f"http://127.0.0.1:{port}" + server_urls.append(url) + + server_proc, _, _ = start_vllm_server( + model_name=model_name, + server_port=port, + distributed_init_method=init_method, + stream_process_output_fn=stream_process_output, + extra_args=fast_llm_server_args, + gpu_ids=cfg.get("gpu_ids", "0"), + actor_llm_idx=cfg.get("actor_llm_idx", 0), + world_size=world_size, + tensor_parallel_size=cfg.get("tensor_parallel_size", 1), + ) + server_procs.append(server_proc) + + await asyncio.sleep(1) + + trainer_proc, _, _ = start_trainer_process( + trainer_helper_path=fast_llm_trainer_helper, + distributed_init_method=init_method, + model_name=model_name, + server_urls=server_urls, + stream_process_output_fn=stream_process_output, + extra_args=[ + "--redis-host", redis_host, + "--redis-port", str(redis_port), + ], + gpu_id=trainer_gpu, + world_size=world_size, + ) + + try: + await wait_for_all_servers_ready(server_urls, server_procs, trainer_proc) + + if len(server_urls) == 1: + generations = await run_generation_loop( + server_url=server_urls[0], + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + trainer_proc=trainer_proc, + ) + else: + per_server_generations = await run_generation_loop_multi( + server_urls=server_urls, + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + trainer_proc=trainer_proc, + ) + + # Wait for trainer to finish + print("[Main] Waiting for trainer to finish...") + for _ in range(30): + if trainer_proc.poll() is not None: + break + await asyncio.sleep(1) + + if len(server_urls) == 1: + analyze_and_verify_pattern(generations) + else: + analyze_and_verify_pattern_multi(per_server_generations) + print(f"\n✓ Fast-LLM server weight update pattern test PASSED ({len(server_urls)} server(s))") + + finally: + print("[Main] Cleaning up processes...") + for proc in server_procs: + if proc: + kill_process_tree(proc.pid) + if trainer_proc: + kill_process_tree(trainer_proc.pid) + + +class TestFastLLMServerIntegration: + """Test Fast-LLM weight broadcast with vLLM HTTP server — 2 GPUs (baseline).""" + + @pytest.mark.timeout(2400) # 40 minutes for server test + @pytest.mark.asyncio + @pytest.mark.skipif( + torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs" + ) + async def test_server_fast_llm_broadcast_pattern( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + fast_llm_trainer_helper, + redis_server, + temp_dir, + ): + """Server integration test: verify Fast-LLM weight broadcast pattern with HTTP API. + + Validates the Fast-LLM protocol where: + - Redis server signals weight updates + - vLLM server receives weights via broadcast_object_list + broadcast + - Server responses change as expected (original → perturbed → original → perturbed) + + Topology: 1 vLLM server on GPU 0, trainer on GPU 1 (world_size=2). + """ + print("\n" + "=" * 60) + print("Starting Fast-LLM server weight update pattern test (TP=1, 1 actor, 2 GPUs)") + print("=" * 60) + + redis_host, redis_port = redis_server + + await _run_fast_llm_server_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + fast_llm_trainer_helper=fast_llm_trainer_helper, + redis_host=redis_host, + redis_port=redis_port, + vllm_server_configs=[{"port": 8000, "gpu_ids": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}], + trainer_gpu="1", + world_size=2, + timeout=2400, + ) + + @pytest.mark.timeout(2400) + @pytest.mark.asyncio + @pytest.mark.skipif( + torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs" + ) + async def test_fast_llm_server_catch_transitions( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + fast_llm_trainer_helper, + redis_server, + temp_dir, + ): + """Diagnostic test: catch garbage generations during Fast-LLM weight broadcasts. + + The trainer runs a slow initial cycle (perturbed → original, 5 s each) + to firmly establish text_A and text_B, then fires N rapid back-to-back + broadcast cycles (perturbed → original) with no inter-broadcast delay. + The generation loop runs with generation_interval=0.0 to maximise the + chance of hitting a mid-broadcast state. + + Assertions: + 1. The A→B→A→B pattern is still detected (broadcasts actually worked). + 2. At least one transition/garbage phase was captured. + + Topology: 1 vLLM server on GPU 0, trainer on GPU 1 (world_size=2). + """ + print("\n" + "=" * 60) + print("Starting Fast-LLM transition-capture test (TP=1, 1 actor, 2 GPUs)") + print("=" * 60) + + redis_host, redis_port = redis_server + server_url = "http://127.0.0.1:8000" + + server_proc, _, _ = start_vllm_server( + model_name=model_name, + server_port=8000, + distributed_init_method=distributed_init_method, + stream_process_output_fn=stream_process_output, + extra_args=[ + "--weight-update-mode", "fast-llm", + "--redis-host", redis_host, + "--redis-port", str(redis_port), + ], + gpu_ids="0", + actor_llm_idx=0, + world_size=2, + tensor_parallel_size=1, + ) + + await asyncio.sleep(1) + + trainer_proc, _, _ = start_trainer_process( + trainer_helper_path=fast_llm_trainer_helper, + distributed_init_method=distributed_init_method, + model_name=model_name, + server_urls=[server_url], + stream_process_output_fn=stream_process_output, + extra_args=[ + "--redis-host", redis_host, + "--redis-port", str(redis_port), + "--n-cycles", "6", + ], + gpu_id="1", + world_size=2, + ) + + try: + await wait_for_server_ready(server_url, server_proc, trainer_proc) + + generations = await run_generation_loop( + server_url=server_url, + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + trainer_proc=trainer_proc, + generation_interval=0.0, + ) + + print("[Main] Waiting for trainer to finish...") + for _ in range(30): + if trainer_proc.poll() is not None: + break + await asyncio.sleep(1) + + analyze_and_verify_transitions(generations, n_cycles=6) + print("\n✓ Fast-LLM transition-capture test PASSED") + + finally: + print("[Main] Cleaning up processes...") + if server_proc: + kill_process_tree(server_proc.pid) + if trainer_proc: + kill_process_tree(trainer_proc.pid) + + +class TestFastLLMServerTP2: + """Test Fast-LLM weight broadcast with tensor-parallel (TP=2) — needs 3 GPUs.""" + + @pytest.mark.timeout(2400) + @pytest.mark.asyncio + @pytest.mark.skipif( + torch.cuda.device_count() < 3, reason="Requires at least 3 GPUs" + ) + async def test_server_fast_llm_broadcast_pattern_tp2( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + fast_llm_trainer_helper, + redis_server, + temp_dir, + ): + """Fast-LLM server test with TP=2: one server on GPUs 0+1, trainer on GPU 2. + + Verifies that tensor-parallel vLLM correctly receives Fast-LLM weight + updates when multiple GPU workers share the same NCCL process group. + """ + print("\n" + "=" * 60) + print("Starting Fast-LLM server weight update pattern test (TP=2, 1 actor, 3 GPUs)") + print("=" * 60) + + redis_host, redis_port = redis_server + + await _run_fast_llm_server_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + fast_llm_trainer_helper=fast_llm_trainer_helper, + redis_host=redis_host, + redis_port=redis_port, + vllm_server_configs=[{"port": 8001, "gpu_ids": "0,1", "actor_llm_idx": 0, "tensor_parallel_size": 2}], + trainer_gpu="2", + world_size=3, + timeout=2400, + ) + + +class TestFastLLMServerMultiActor: + """Test Fast-LLM weight broadcast with multiple independent vLLM actors.""" + + @pytest.mark.timeout(2400) + @pytest.mark.asyncio + @pytest.mark.skipif( + torch.cuda.device_count() < 3, reason="Requires at least 3 GPUs" + ) + async def test_server_fast_llm_broadcast_pattern_2actors( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + fast_llm_trainer_helper, + redis_server, + temp_dir, + ): + """Fast-LLM server test with 2 actors: servers on GPUs 0 and 1, trainer on GPU 2. + + Verifies that two separate vLLM servers simultaneously receive the same + Fast-LLM weight broadcast and produce identical generation results. + """ + print("\n" + "=" * 60) + print("Starting Fast-LLM server weight update pattern test (TP=1, 2 actors, 3 GPUs)") + print("=" * 60) + + redis_host, redis_port = redis_server + + await _run_fast_llm_server_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + fast_llm_trainer_helper=fast_llm_trainer_helper, + redis_host=redis_host, + redis_port=redis_port, + vllm_server_configs=[ + {"port": 8000, "gpu_ids": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}, + {"port": 8001, "gpu_ids": "1", "actor_llm_idx": 1, "tensor_parallel_size": 1}, + ], + trainer_gpu="2", + world_size=3, + timeout=2400, + ) + + @pytest.mark.timeout(2400) + @pytest.mark.asyncio + @pytest.mark.skipif( + torch.cuda.device_count() < 4, reason="Requires at least 4 GPUs" + ) + async def test_server_fast_llm_broadcast_pattern_3actors( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + fast_llm_trainer_helper, + redis_server, + temp_dir, + ): + """Fast-LLM server test with 3 actors: servers on GPUs 0/1/2, trainer on GPU 3. + + Verifies that three separate vLLM servers simultaneously receive the same + Fast-LLM weight broadcast and produce identical generation results. + """ + print("\n" + "=" * 60) + print("Starting Fast-LLM server weight update pattern test (TP=1, 3 actors, 4 GPUs)") + print("=" * 60) + + redis_host, redis_port = redis_server + + await _run_fast_llm_server_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + fast_llm_trainer_helper=fast_llm_trainer_helper, + redis_host=redis_host, + redis_port=redis_port, + vllm_server_configs=[ + {"port": 8000, "gpu_ids": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}, + {"port": 8001, "gpu_ids": "1", "actor_llm_idx": 1, "tensor_parallel_size": 1}, + {"port": 8002, "gpu_ids": "2", "actor_llm_idx": 2, "tensor_parallel_size": 1}, + ], + trainer_gpu="3", + world_size=4, + timeout=2400, + ) diff --git a/tests/test_vllm1_integration.py b/tests/test_vllm1_integration.py new file mode 100644 index 00000000..2954bce5 --- /dev/null +++ b/tests/test_vllm1_integration.py @@ -0,0 +1,1286 @@ +"""Integration tests for vllm1 with actual distributed setup.""" + +import asyncio +import pytest +import tempfile +from pathlib import Path +from typing import Dict, List +import time +import os +import subprocess +import sys +import signal + +# torch is needed at top level for pytest.mark.skipif decorators +import torch + +# Import shared utilities +from .server_weight_update_utils import ( + wait_for_server_ready, + wait_for_all_servers_ready, + run_generation_loop, + run_generation_loop_multi, + analyze_and_verify_pattern, + analyze_and_verify_pattern_multi, + analyze_and_verify_transitions, + start_vllm_server, + start_trainer_process, +) + +try: + import psutil + HAS_PSUTIL = True +except ImportError: + HAS_PSUTIL = False + print("WARNING: psutil not available, process tree cleanup will be limited") + + +def stream_process_output(proc, name): + """Start background threads to continuously stream process stdout/stderr. + + Args: + proc: subprocess.Popen object + name: Name for logging prefix (e.g., "vLLM Server", "Trainer") + + Returns: + Tuple of (stdout_thread, stderr_thread) + """ + import threading + + def read_stream(stream, prefix): + """Read from stream and print with prefix.""" + try: + for line in iter(stream.readline, ''): + if line: + print(f"{prefix} {line.rstrip()}", flush=True) + except Exception as e: + print(f"{prefix} [Stream read error: {e}]", flush=True) + + stdout_thread = threading.Thread( + target=read_stream, + args=(proc.stdout, f"[{name} OUT]"), + daemon=True, + ) + stderr_thread = threading.Thread( + target=read_stream, + args=(proc.stderr, f"[{name} ERR]"), + daemon=True, + ) + + stdout_thread.start() + stderr_thread.start() + + return stdout_thread, stderr_thread + + +def kill_process_tree(pid, sig=signal.SIGKILL): + """Kill a process and all its children/grandchildren. + + Args: + pid: Process ID to kill + sig: Signal to send (default SIGKILL) + """ + if not HAS_PSUTIL: + # Fallback: just kill the main process + try: + os.kill(pid, sig) + except ProcessLookupError: + pass + return + + try: + parent = psutil.Process(pid) + except psutil.NoSuchProcess: + return + + # Get all children recursively + children = parent.children(recursive=True) + + # Kill children first + for child in children: + try: + print(f"[Kill] Killing child process {child.pid}") + child.send_signal(sig) + except psutil.NoSuchProcess: + pass + + # Kill parent + try: + parent.send_signal(sig) + except psutil.NoSuchProcess: + pass + + +def force_kill_process(proc, name): + """Forcefully kill a process tree and collect output. + + SIGKILL always kills the process. If communicate() hangs, it's the PIPES + that are stuck, not the process. We handle this with retries and timeouts. + + Returns: + Tuple of (stdout, stderr, returncode) + """ + # If already dead, try to get output + if proc.poll() is not None: + try: + stdout, stderr = proc.communicate(timeout=2) + return stdout, stderr, proc.returncode + except subprocess.TimeoutExpired: + print(f"[Kill] {name} already dead but pipes hung, closing...") + proc.stdout.close() if proc.stdout else None + proc.stderr.close() if proc.stderr else None + return "", "", proc.returncode + + # Kill entire process tree (including vLLM workers, trainer subprocesses, etc) + print(f"[Kill] Killing {name} process tree (PID {proc.pid})...") + kill_process_tree(proc.pid, signal.SIGKILL) + + # Wait for main process to actually die + try: + proc.wait(timeout=2) + print(f"[Kill] {name} process tree killed") + except subprocess.TimeoutExpired: + print(f"[Kill] WARNING: {name} didn't die after SIGKILL") + + # Try to read output from pipes (this is what usually hangs) + for attempt, timeout_val in enumerate([1, 2, 3], start=1): + try: + stdout, stderr = proc.communicate(timeout=timeout_val) + print(f"[Kill] {name} output collected (attempt {attempt})") + return stdout, stderr, proc.returncode + except subprocess.TimeoutExpired: + print(f"[Kill] {name} communicate() timed out (attempt {attempt})") + continue + + # Pipes are stuck - force close them + print(f"[Kill] {name} pipes stuck, force closing...") + try: + proc.stdout.close() if proc.stdout else None + proc.stderr.close() if proc.stderr else None + proc.stdin.close() if proc.stdin else None + except Exception as e: + print(f"[Kill] Error closing pipes: {e}") + + return "", "", proc.returncode if proc.returncode else -999 + + +async def wait_for_processes(processes_with_names, check_interval=0.5, timeout=60): + """Wait for multiple subprocesses to complete, printing output in real-time. + + Args: + processes_with_names: List of (subprocess.Popen, name) tuples + check_interval: How often to check process status (seconds) + timeout: Maximum time to wait for all processes (seconds) + + Raises: + RuntimeError: If any process fails or timeout is reached + """ + start_time = time.time() + + # Create async readers for each process's stdout and stderr + async def read_stream(stream, prefix): + """Read from a stream line-by-line and print with prefix.""" + loop = asyncio.get_event_loop() + try: + while True: + line = await loop.run_in_executor(None, stream.readline) + if not line: + break + print(f"{prefix} {line.rstrip()}", flush=True) + except Exception as e: + print(f"{prefix} [Read error: {e}]", flush=True) + + # Start readers for all processes + reader_tasks = [] + for proc, name in processes_with_names: + reader_tasks.append(asyncio.create_task(read_stream(proc.stdout, f"[{name} OUT]"))) + reader_tasks.append(asyncio.create_task(read_stream(proc.stderr, f"[{name} ERR]"))) + + try: + while True: + # Check if timeout exceeded + if time.time() - start_time > timeout: + print(f"\n{'='*60}", flush=True) + print("TIMEOUT: Killing all processes", flush=True) + print(f"{'='*60}\n", flush=True) + + # Kill all processes forcefully + for proc, name in processes_with_names: + if proc.poll() is None: + print(f"[Main] Killing {name}...", flush=True) + kill_process_tree(proc.pid, signal.SIGKILL) + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + pass + + raise RuntimeError(f"Timeout after {timeout} seconds waiting for processes") + + # Check each process + crashed_proc = None + crashed_name = None + + for proc, name in processes_with_names: + returncode = proc.poll() + if returncode is not None and returncode != 0: + crashed_proc = proc + crashed_name = name + print(f"\n{'='*60}", flush=True) + print(f"{name} process CRASHED with exit code {returncode}", flush=True) + print(f"{'='*60}\n", flush=True) + break + + # If a process crashed, kill the others + if crashed_proc is not None: + # Kill all other processes + for proc, name in processes_with_names: + if proc != crashed_proc and proc.poll() is None: + print(f"[Main] Killing {name}...", flush=True) + kill_process_tree(proc.pid, signal.SIGKILL) + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + pass + + raise RuntimeError( + f"{crashed_name} process failed with exit code {crashed_proc.returncode}" + ) + + # Check if all processes completed successfully + all_done = all(proc.poll() is not None for proc, _ in processes_with_names) + if all_done: + # Wait for readers to finish draining pipes + print("[Main] All processes completed, waiting for output to finish...", flush=True) + await asyncio.sleep(1) # Give readers time to finish + + print(f"\n{'='*60}", flush=True) + print("✓ All processes completed successfully", flush=True) + print(f"{'='*60}\n", flush=True) + return + + # Sleep before next check + await asyncio.sleep(check_interval) + finally: + # Cancel reader tasks + for task in reader_tasks: + if not task.done(): + task.cancel() + # Wait for cancellation + await asyncio.gather(*reader_tasks, return_exceptions=True) + + +# --------------------------------------------------------------------------- +# Module-level helpers shared by all test variants +# --------------------------------------------------------------------------- + +def _compare_actor_results(sync_dir: Path, num_actors: int): + """Assert that all actors produced identical generation results. + + Each actor writes ``sync_dir/results_actor_{i}.json`` with keys + res_or_1, res_mod_1, res_or_2, res_mod_2. + """ + import json + + results = [ + json.loads((sync_dir / f"results_actor_{i}.json").read_text()) + for i in range(num_actors) + ] + for key in results[0]: + texts = [r[key] for r in results] + assert len(set(texts)) == 1, ( + f"Actors disagree on '{key}': {texts}" + ) + + +async def _run_back_and_forth_engine_test( + model_name, + simple_prompt, + generation_config, + init_method, + distributed_trainer_helper, + vllm_engine_helper, + sync_dir, + vllm_configs, + trainer_gpu, + world_size, + timeout=1800, +): + """Run back-and-forth engine test with one or more vLLM actor processes. + + Args: + vllm_configs: List of dicts, each with keys: + - cuda_devices: str, e.g. "0" or "0,1" + - actor_llm_idx: int + - tensor_parallel_size: int + trainer_gpu: str, e.g. "1" or "2" + world_size: total NCCL world size (all vLLM workers + trainer) + """ + from .sync_helper import create_sync_dir + + num_actors = len(vllm_configs) + all_procs = [] + + # Start all vLLM actor subprocesses + for cfg in vllm_configs: + vllm_env = os.environ.copy() + vllm_env["CUDA_VISIBLE_DEVICES"] = cfg["cuda_devices"] + vllm_env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + vllm_env["PIPELINERL_DEBUG"] = "1" + + actor_idx = cfg["actor_llm_idx"] + tp = cfg.get("tensor_parallel_size", 1) + print(f"[Main] Starting vLLM actor {actor_idx} (GPU(s) {cfg['cuda_devices']}, TP={tp})") + + vllm_proc = subprocess.Popen( + [ + sys.executable, + str(vllm_engine_helper), + "back_and_forth", + "--model-name", model_name, + "--init-method", init_method, + "--actor-llm-idx", str(actor_idx), + "--world-size", str(world_size), + "--tensor-parallel-size", str(tp), + "--prompt", simple_prompt, + "--max-tokens", str(generation_config["max_tokens"]), + "--sync-dir", str(sync_dir), + ], + env=vllm_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + all_procs.append((vllm_proc, f"vLLM Actor {actor_idx}")) + + await asyncio.sleep(1) + + # Start trainer subprocess + trainer_env = os.environ.copy() + trainer_env["CUDA_VISIBLE_DEVICES"] = trainer_gpu + trainer_env["PIPELINERL_DEBUG"] = "1" + + print(f"[Main] Starting trainer (GPU {trainer_gpu}, {num_actors} actor(s), world_size={world_size})") + trainer_proc = subprocess.Popen( + [ + sys.executable, + str(distributed_trainer_helper), + "back_and_forth", + "--init-method", init_method, + "--model-name", model_name, + "--sync-dir", str(sync_dir), + "--num-actors", str(num_actors), + "--world-size", str(world_size), + ], + env=trainer_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + all_procs.append((trainer_proc, "Trainer")) + + await wait_for_processes(all_procs, timeout=timeout) + + # Verify all actors produced the same results + _compare_actor_results(sync_dir, num_actors) + print(f"\n✓ Back-and-forth test PASSED ({num_actors} actor(s), world_size={world_size})") + + +async def _run_server_weight_update_test( + model_name, + simple_prompt, + generation_config, + init_method, + distributed_trainer_helper, + vllm_server_configs, + trainer_gpu, + world_size, + timeout=2400, +): + """Run server weight-update pattern test with one or more vLLM servers. + + Args: + vllm_server_configs: List of dicts, each with keys: + - port: int + - gpu_ids: str + - actor_llm_idx: int + - tensor_parallel_size: int + trainer_gpu: str, e.g. "1" or "2" + world_size: total NCCL world size + """ + server_procs = [] + server_urls = [] + + for cfg in vllm_server_configs: + port = cfg["port"] + url = f"http://127.0.0.1:{port}" + server_urls.append(url) + + server_proc, _, _ = start_vllm_server( + model_name=model_name, + server_port=port, + distributed_init_method=init_method, + stream_process_output_fn=stream_process_output, + extra_args=None, + gpu_ids=cfg.get("gpu_ids", "0"), + actor_llm_idx=cfg.get("actor_llm_idx", 0), + world_size=world_size, + tensor_parallel_size=cfg.get("tensor_parallel_size", 1), + ) + server_procs.append(server_proc) + + await asyncio.sleep(1) + + trainer_proc, _, _ = start_trainer_process( + trainer_helper_path=distributed_trainer_helper, + distributed_init_method=init_method, + model_name=model_name, + server_urls=server_urls, + stream_process_output_fn=stream_process_output, + extra_args=None, + gpu_id=trainer_gpu, + world_size=world_size, + ) + + try: + await wait_for_all_servers_ready(server_urls, server_procs, trainer_proc) + + if len(server_urls) == 1: + generations = await run_generation_loop( + server_url=server_urls[0], + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + trainer_proc=trainer_proc, + ) + else: + per_server_generations = await run_generation_loop_multi( + server_urls=server_urls, + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + trainer_proc=trainer_proc, + ) + + # Wait for trainer to finish + print("[Main] Waiting for trainer to finish...") + for _ in range(30): + if trainer_proc.poll() is not None: + break + await asyncio.sleep(1) + + if len(server_urls) == 1: + analyze_and_verify_pattern(generations) + else: + analyze_and_verify_pattern_multi(per_server_generations) + print(f"\n✓ Server weight update pattern test PASSED ({len(server_urls)} server(s))") + + finally: + print("[Main] Cleaning up processes...") + for proc in server_procs: + if proc: + kill_process_tree(proc.pid) + if trainer_proc: + kill_process_tree(trainer_proc.pid) + + +class TestBasicGeneration: + """Test basic vLLM generation with worker extension.""" + + @pytest.mark.asyncio + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU") + async def test_load_model_and_generate(self, vllm_engine_factory, simple_prompt, generation_config): + """Test loading model and generating text.""" + from vllm import SamplingParams + + async with vllm_engine_factory(disable_weight_updates=True) as manager: + # Generate text + sampling_params = SamplingParams( + temperature=generation_config["temperature"], + top_p=generation_config["top_p"], + max_tokens=generation_config["max_tokens"], + seed=generation_config["seed"], + ) + + request_id = "test_request_1" + async for output in manager.engine.generate( + simple_prompt, + sampling_params=sampling_params, + request_id=request_id, + ): + final_output = output + + assert final_output is not None + assert len(final_output.outputs) > 0 + assert len(final_output.outputs[0].text) > 0 + + print(f"Generated text: {final_output.outputs[0].text}") + + @pytest.mark.asyncio + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU") + async def test_deterministic_generation(self, vllm_engine_factory, simple_prompt, generation_config): + """Test that generation is deterministic with same seed and temperature=0.""" + from vllm import SamplingParams + + async with vllm_engine_factory(disable_weight_updates=True) as manager: + sampling_params = SamplingParams( + temperature=generation_config["temperature"], + top_p=generation_config["top_p"], + max_tokens=generation_config["max_tokens"], + seed=generation_config["seed"], + ) + + # Generate twice with same parameters + outputs = [] + for i in range(2): + request_id = f"test_request_{i}" + async for output in manager.engine.generate( + simple_prompt, + sampling_params=sampling_params, + request_id=request_id, + ): + final_output = output + outputs.append(final_output.outputs[0].text) + + # Outputs should be identical + assert outputs[0] == outputs[1], f"Outputs differ: '{outputs[0]}' vs '{outputs[1]}'" + + +class TestWeightUpdateDistributed: + """Test weight updates with 2-GPU distributed setup.""" + + @pytest.mark.timeout(300) # 5 minutes for init test + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") + async def test_init_actor_update_group( + self, + model_name, + distributed_init_method, + distributed_trainer_helper, + vllm_engine_helper, + ): + """Test initializing actor update group with 2 GPUs. + + This test verifies that the process group can be initialized correctly: + - vLLM engine runs on GPU 0 as rank 1 (in subprocess) + - Dummy trainer process runs on GPU 1 as rank 0 (in subprocess) + + Both run in subprocesses to ensure proper CUDA_VISIBLE_DEVICES isolation. + """ + print("\n" + "="*60) + print("Starting distributed process group initialization test") + print("="*60) + + # Step 1: Start trainer subprocess FIRST with CUDA_VISIBLE_DEVICES=1 + trainer_env = os.environ.copy() + trainer_env["CUDA_VISIBLE_DEVICES"] = "1" + trainer_env["PIPELINERL_DEBUG"] = "1" + + print("[Main] Starting trainer process (rank 0, GPU 1)") + trainer_proc = subprocess.Popen( + [ + sys.executable, + str(distributed_trainer_helper), + "init", + "--init-method", distributed_init_method, + "--rank", "0", + "--world-size", "2", + ], + env=trainer_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Give trainer a moment to start and begin initializing + await asyncio.sleep(1) + + # Step 2: Start vLLM engine subprocess with CUDA_VISIBLE_DEVICES=0 + vllm_env = os.environ.copy() + vllm_env["CUDA_VISIBLE_DEVICES"] = "0" + vllm_env["PIPELINERL_DEBUG"] = "1" + + print("[Main] Starting vLLM engine process (rank 1, GPU 0)") + vllm_proc = subprocess.Popen( + [ + sys.executable, + str(vllm_engine_helper), + "init", # Command argument + "--model-name", model_name, + "--init-method", distributed_init_method, + "--actor-llm-idx", "0", + "--world-size", "2", + ], + env=vllm_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Step 3: Wait for both processes, killing all if one crashes + await wait_for_processes([ + (trainer_proc, "Trainer"), + (vllm_proc, "vLLM Engine"), + ], timeout=180) # Init test is faster, but give it 3 minutes to be safe + + @pytest.mark.timeout(1000) # 1000 seconds for broadcasting 291 parameters + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") + async def test_weight_update_same_weights( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + vllm_engine_helper, + temp_dir, + ): + """Test that updating with same weights produces same output. + + This test: + 1. vLLM engine generates baseline output (in subprocess on GPU 0) + 2. Trainer waits for baseline, then broadcasts weights (in subprocess on GPU 1) + 3. vLLM engine receives update and generates again + 4. vLLM engine verifies outputs are identical + + Both run in subprocesses for proper CUDA_VISIBLE_DEVICES isolation. + Uses file-based sync points for coordination. + """ + from .sync_helper import create_sync_dir + + print("\n" + "="*60) + print("Starting weight update test (same weights)") + print("="*60) + + # Create sync directory for coordination + sync_dir = create_sync_dir(temp_dir) + print(f"[Main] Sync directory: {sync_dir}") + + # Step 1: Start vLLM engine subprocess with weight_update command + vllm_env = os.environ.copy() + vllm_env["CUDA_VISIBLE_DEVICES"] = "0" + # NOTE: needed to pass WeightUpdateRequest to collective + vllm_env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + # Enable DEBUG logging in vllm1.py + vllm_env["PIPELINERL_DEBUG"] = "1" + + print("[Main] Starting vLLM engine process (GPU 0)") + vllm_proc = subprocess.Popen( + [ + sys.executable, + str(vllm_engine_helper), + "weight_update", + "--model-name", model_name, + "--init-method", distributed_init_method, + "--actor-llm-idx", "0", + "--world-size", "2", + "--prompt", simple_prompt, + "--max-tokens", str(generation_config["max_tokens"]), + "--sync-dir", str(sync_dir), + ], + env=vllm_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Give vLLM engine a moment to start + await asyncio.sleep(1) + + # Step 2: Start trainer subprocess (will wait for baseline_done sync point) + trainer_env = os.environ.copy() + trainer_env["CUDA_VISIBLE_DEVICES"] = "1" + # Enable DEBUG logging in vllm1.py (for consistency) + trainer_env["PIPELINERL_DEBUG"] = "1" + + print("[Main] Starting trainer process (GPU 1)") + trainer_proc = subprocess.Popen( + [ + sys.executable, + str(distributed_trainer_helper), + "broadcast", + "--init-method", distributed_init_method, + "--model-name", model_name, + "--sync-dir", str(sync_dir), + ], + env=trainer_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Step 3: Wait for both processes, killing all if one crashes + # 291 parameters takes ~600 seconds, so use 900s (15 min) to be safe + await wait_for_processes([ + (vllm_proc, "vLLM Engine"), + (trainer_proc, "Trainer"), + ], timeout=900) + + @pytest.mark.timeout(1000) # 1000 seconds for broadcasting 290 parameters + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") + async def test_weight_update_different_weights( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + vllm_engine_helper, + temp_dir, + ): + """Test that updating with perturbed weights produces different output. + + This test: + 1. vLLM engine generates baseline output (in subprocess on GPU 0) + 2. Trainer broadcasts PERTURBED weights (in subprocess on GPU 1) + 3. vLLM engine receives update and generates again + 4. vLLM engine verifies outputs are DIFFERENT (perturbed weights changed output) + + Both run in subprocesses for proper CUDA_VISIBLE_DEVICES isolation. + Uses file-based sync points for coordination. + """ + from .sync_helper import create_sync_dir + + print("\n" + "="*60) + print("Starting weight update test (perturbed weights)") + print("="*60) + + # Create sync directory for coordination + sync_dir = create_sync_dir(temp_dir) + print(f"[Main] Sync directory: {sync_dir}") + + # Step 1: Start vLLM engine subprocess with weight_update command + vllm_env = os.environ.copy() + vllm_env["CUDA_VISIBLE_DEVICES"] = "0" + # NOTE: needed to pass WeightUpdateRequest to collective + vllm_env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + # Enable DEBUG logging in vllm1.py + vllm_env["PIPELINERL_DEBUG"] = "1" + + print("[Main] Starting vLLM engine process (GPU 0)") + vllm_proc = subprocess.Popen( + [ + sys.executable, + str(vllm_engine_helper), + "weight_update", + "--model-name", model_name, + "--init-method", distributed_init_method, + "--actor-llm-idx", "0", + "--world-size", "2", + "--prompt", simple_prompt, + "--max-tokens", str(generation_config["max_tokens"]), + "--sync-dir", str(sync_dir), + "--expect-different", # Flag to expect different outputs + ], + env=vllm_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Give vLLM engine a moment to start + await asyncio.sleep(1) + + # Step 2: Start trainer subprocess with --perturb flag + trainer_env = os.environ.copy() + trainer_env["CUDA_VISIBLE_DEVICES"] = "1" + # Enable DEBUG logging in vllm1.py (for consistency) + trainer_env["PIPELINERL_DEBUG"] = "1" + + print("[Main] Starting trainer process (GPU 1) with --perturb") + trainer_proc = subprocess.Popen( + [ + sys.executable, + str(distributed_trainer_helper), + "broadcast", + "--init-method", distributed_init_method, + "--model-name", model_name, + "--sync-dir", str(sync_dir), + "--perturb", # Perturb weights to test different outputs + ], + env=trainer_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Step 3: Wait for both processes, killing all if one crashes + # 290 parameters takes ~600 seconds, so use 900s (15 min) to be safe + await wait_for_processes([ + (vllm_proc, "vLLM Engine"), + (trainer_proc, "Trainer"), + ], timeout=900) + + + @pytest.mark.timeout(2000) # 2000 seconds - this test does 2 full broadcasts + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") + async def test_weight_update_cross_validation( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + vllm_engine_helper, + temp_dir, + ): + """Cross-validation test: verify broadcast = load from disk. + + This test validates that: + 1. Broadcasting weights produces same results as loading from disk + 2. Round-trip works: original → modified → original + + Flow: + - vLLM: Load original, generate res_un_1 + - Trainer: Save perturbed model to disk, broadcast perturbed weights + - vLLM: Receive perturbed, generate res_mod_1 + - vLLM: Recreate engine with perturbed model from disk, generate res_mod_2 + - Trainer: Broadcast original weights + - vLLM: Receive original, generate res_un_2 + + Assertions: + - res_un_1 == res_un_2 (original weights produce same output) + - res_mod_1 == res_mod_2 (broadcast = load from disk) + """ + from .sync_helper import create_sync_dir + + print("\n" + "="*60) + print("Starting cross-validation test") + print("="*60) + + # Create sync directory for coordination + sync_dir = create_sync_dir(temp_dir) + print(f"[Main] Sync directory: {sync_dir}") + print(f"[Main] Temp directory: {temp_dir}") + + # Step 1: Start vLLM engine subprocess + vllm_env = os.environ.copy() + vllm_env["CUDA_VISIBLE_DEVICES"] = "0" + vllm_env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + vllm_env["PIPELINERL_DEBUG"] = "1" + + print("[Main] Starting vLLM engine process (GPU 0)") + vllm_proc = subprocess.Popen( + [ + sys.executable, + str(vllm_engine_helper), + "cross_validation", + "--model-name", model_name, + "--init-method", distributed_init_method, + "--actor-llm-idx", "0", + "--world-size", "2", + "--prompt", simple_prompt, + "--max-tokens", str(generation_config["max_tokens"]), + "--sync-dir", str(sync_dir), + ], + env=vllm_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Give vLLM engine a moment to start + await asyncio.sleep(1) + + # Step 2: Start trainer subprocess + trainer_env = os.environ.copy() + trainer_env["CUDA_VISIBLE_DEVICES"] = "1" + trainer_env["PIPELINERL_DEBUG"] = "1" + + print("[Main] Starting trainer process (GPU 1)") + trainer_proc = subprocess.Popen( + [ + sys.executable, + str(distributed_trainer_helper), + "cross_validation", + "--init-method", distributed_init_method, + "--model-name", model_name, + "--sync-dir", str(sync_dir), + "--temp-dir", str(temp_dir), + ], + env=trainer_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Step 3: Wait for both processes + # This test does 2 broadcasts, so double the timeout + await wait_for_processes([ + (vllm_proc, "vLLM Engine"), + (trainer_proc, "Trainer"), + ], timeout=1800) # 30 minutes + + + @pytest.mark.timeout(2000) # 2000 seconds - this test does 3 broadcasts + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") + async def test_weight_update_back_and_forth( + self, + model_name, + simple_prompt, + generation_config, + shared_distributed_init_method, + distributed_trainer_helper, + vllm_engine_helper, + shared_test_dir, + ): + """Back-and-forth test: switch between original and perturbed weights. + + Validates that we can update weights multiple times and the results + are deterministic and reproducible. + """ + from .sync_helper import create_sync_dir + + print("\n" + "="*60) + print("Starting back-and-forth test (TP=1, 1 actor, 2 GPUs)") + print("="*60) + + sync_dir = create_sync_dir(shared_test_dir) + await _run_back_and_forth_engine_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=shared_distributed_init_method, + distributed_trainer_helper=distributed_trainer_helper, + vllm_engine_helper=vllm_engine_helper, + sync_dir=sync_dir, + vllm_configs=[{"cuda_devices": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}], + trainer_gpu="1", + world_size=2, + timeout=1800, + ) + + @pytest.mark.timeout(2400) # 40 minutes for server test + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") + async def test_server_weight_update_pattern( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + temp_dir, + ): + """Server integration test: verify weight update pattern with HTTP API. + + Validates the real-world scenario where a vLLM HTTP server receives + weight updates from a trainer while serving requests. + """ + print("\n" + "="*60) + print("Starting server weight update pattern test (TP=1, 1 actor, 2 GPUs)") + print("="*60) + + await _run_server_weight_update_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + distributed_trainer_helper=distributed_trainer_helper, + vllm_server_configs=[{"port": 8000, "gpu_ids": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}], + trainer_gpu="1", + world_size=2, + timeout=2400, + ) + + @pytest.mark.timeout(2400) + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") + async def test_server_weight_update_catch_transitions( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + temp_dir, + ): + """Diagnostic test: catch garbage generations produced during NCCL weight broadcasts. + + The trainer runs a slow initial cycle (perturbed → original, 5 s each) + to firmly establish text_A and text_B, then fires N rapid back-to-back + broadcast cycles (perturbed → original) with no inter-broadcast delay. + The generation loop runs with generation_interval=0.0 (back-to-back + requests) to maximise the chance of hitting a mid-broadcast state. + + Assertions: + 1. The A→B→A→B pattern is still detected (broadcasts actually worked). + 2. At least one transition/garbage phase was captured. + + Topology: 1 vLLM server on GPU 0, trainer on GPU 1 (world_size=2). + """ + print("\n" + "=" * 60) + print("Starting transition-capture test (TP=1, 1 actor, 2 GPUs)") + print("=" * 60) + + server_url = "http://127.0.0.1:8000" + + server_proc, _, _ = start_vllm_server( + model_name=model_name, + server_port=8000, + distributed_init_method=distributed_init_method, + stream_process_output_fn=stream_process_output, + gpu_ids="0", + actor_llm_idx=0, + world_size=2, + tensor_parallel_size=1, + ) + + await asyncio.sleep(1) + + trainer_proc, _, _ = start_trainer_process( + trainer_helper_path=distributed_trainer_helper, + distributed_init_method=distributed_init_method, + model_name=model_name, + server_urls=[server_url], + stream_process_output_fn=stream_process_output, + extra_args=["--n-cycles", "6"], + gpu_id="1", + world_size=2, + command="rapid_broadcast_cycles", + ) + + try: + await wait_for_server_ready(server_url, server_proc, trainer_proc) + + generations = await run_generation_loop( + server_url=server_url, + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + trainer_proc=trainer_proc, + generation_interval=0.0, + ) + + print("[Main] Waiting for trainer to finish...") + for _ in range(30): + if trainer_proc.poll() is not None: + break + await asyncio.sleep(1) + + analyze_and_verify_transitions(generations, n_cycles=6) + print("\n✓ Transition-capture test PASSED") + + finally: + print("[Main] Cleaning up processes...") + if server_proc: + kill_process_tree(server_proc.pid) + if trainer_proc: + kill_process_tree(trainer_proc.pid) + + +class TestWeightUpdateTP2: + """Test weight updates with tensor-parallel (TP=2) vLLM — needs 3 GPUs.""" + + @pytest.mark.timeout(2000) + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 3, reason="Requires at least 3 GPUs") + async def test_weight_update_back_and_forth_tp2( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + vllm_engine_helper, + temp_dir, + ): + """Back-and-forth test with TP=2: one vLLM instance on GPUs 0+1, trainer on GPU 2.""" + from .sync_helper import create_sync_dir + + print("\n" + "="*60) + print("Starting back-and-forth test (TP=2, 1 actor, 3 GPUs)") + print("="*60) + + sync_dir = create_sync_dir(temp_dir) + await _run_back_and_forth_engine_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + distributed_trainer_helper=distributed_trainer_helper, + vllm_engine_helper=vllm_engine_helper, + sync_dir=sync_dir, + vllm_configs=[{"cuda_devices": "0,1", "actor_llm_idx": 0, "tensor_parallel_size": 2}], + trainer_gpu="2", + world_size=3, + timeout=1800, + ) + + @pytest.mark.timeout(2400) + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 3, reason="Requires at least 3 GPUs") + async def test_server_weight_update_pattern_tp2( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + temp_dir, + ): + """Server weight update test with TP=2: one server on GPUs 0+1, trainer on GPU 2.""" + print("\n" + "="*60) + print("Starting server weight update pattern test (TP=2, 1 actor, 3 GPUs)") + print("="*60) + + await _run_server_weight_update_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + distributed_trainer_helper=distributed_trainer_helper, + vllm_server_configs=[{"port": 8001, "gpu_ids": "0,1", "actor_llm_idx": 0, "tensor_parallel_size": 2}], + trainer_gpu="2", + world_size=3, + timeout=2400, + ) + + +class TestWeightUpdateMultiActor: + """Test weight updates with multiple independent vLLM actors.""" + + @pytest.mark.timeout(2000) + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 3, reason="Requires at least 3 GPUs") + async def test_weight_update_back_and_forth_2actors( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + vllm_engine_helper, + temp_dir, + ): + """Back-and-forth test with 2 actors: vLLM on GPU 0 and GPU 1, trainer on GPU 2.""" + from .sync_helper import create_sync_dir + + print("\n" + "="*60) + print("Starting back-and-forth test (TP=1, 2 actors, 3 GPUs)") + print("="*60) + + sync_dir = create_sync_dir(temp_dir) + await _run_back_and_forth_engine_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + distributed_trainer_helper=distributed_trainer_helper, + vllm_engine_helper=vllm_engine_helper, + sync_dir=sync_dir, + vllm_configs=[ + {"cuda_devices": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}, + {"cuda_devices": "1", "actor_llm_idx": 1, "tensor_parallel_size": 1}, + ], + trainer_gpu="2", + world_size=3, + timeout=1800, + ) + + @pytest.mark.timeout(2000) + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 4, reason="Requires at least 4 GPUs") + async def test_weight_update_back_and_forth_3actors( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + vllm_engine_helper, + temp_dir, + ): + """Back-and-forth test with 3 actors: vLLM on GPUs 0/1/2, trainer on GPU 3.""" + from .sync_helper import create_sync_dir + + print("\n" + "="*60) + print("Starting back-and-forth test (TP=1, 3 actors, 4 GPUs)") + print("="*60) + + sync_dir = create_sync_dir(temp_dir) + await _run_back_and_forth_engine_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + distributed_trainer_helper=distributed_trainer_helper, + vllm_engine_helper=vllm_engine_helper, + sync_dir=sync_dir, + vllm_configs=[ + {"cuda_devices": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}, + {"cuda_devices": "1", "actor_llm_idx": 1, "tensor_parallel_size": 1}, + {"cuda_devices": "2", "actor_llm_idx": 2, "tensor_parallel_size": 1}, + ], + trainer_gpu="3", + world_size=4, + timeout=1800, + ) + + @pytest.mark.timeout(2400) + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 3, reason="Requires at least 3 GPUs") + async def test_server_weight_update_pattern_2actors( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + temp_dir, + ): + """Server weight update test with 2 actors: servers on GPUs 0 and 1, trainer on GPU 2.""" + print("\n" + "="*60) + print("Starting server weight update pattern test (TP=1, 2 actors, 3 GPUs)") + print("="*60) + + await _run_server_weight_update_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + distributed_trainer_helper=distributed_trainer_helper, + vllm_server_configs=[ + {"port": 8000, "gpu_ids": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}, + {"port": 8001, "gpu_ids": "1", "actor_llm_idx": 1, "tensor_parallel_size": 1}, + ], + trainer_gpu="2", + world_size=3, + timeout=2400, + ) + + @pytest.mark.timeout(2400) + @pytest.mark.asyncio + @pytest.mark.skipif(torch.cuda.device_count() < 4, reason="Requires at least 4 GPUs") + async def test_server_weight_update_pattern_3actors( + self, + model_name, + simple_prompt, + generation_config, + distributed_init_method, + distributed_trainer_helper, + temp_dir, + ): + """Server weight update test with 3 actors: servers on GPUs 0/1/2, trainer on GPU 3.""" + print("\n" + "="*60) + print("Starting server weight update pattern test (TP=1, 3 actors, 4 GPUs)") + print("="*60) + + await _run_server_weight_update_test( + model_name=model_name, + simple_prompt=simple_prompt, + generation_config=generation_config, + init_method=distributed_init_method, + distributed_trainer_helper=distributed_trainer_helper, + vllm_server_configs=[ + {"port": 8000, "gpu_ids": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}, + {"port": 8001, "gpu_ids": "1", "actor_llm_idx": 1, "tensor_parallel_size": 1}, + {"port": 8002, "gpu_ids": "2", "actor_llm_idx": 2, "tensor_parallel_size": 1}, + ], + trainer_gpu="3", + world_size=4, + timeout=2400, + ) diff --git a/tests/test_world_multinode.py b/tests/test_world_multinode.py new file mode 100644 index 00000000..14d1474a --- /dev/null +++ b/tests/test_world_multinode.py @@ -0,0 +1,842 @@ +"""Tests for multi-node WorldMap topology and fast-llm torchrun command assembly.""" + +import os +import sys +import tempfile +from pathlib import Path +from unittest.mock import patch, MagicMock + +import pytest +from omegaconf import OmegaConf + + +def _make_cfg( + actor_fraction=1, + finetune_fraction=1, + preprocessor_fraction=0, + replicas=1, + use_fast_llm=True, + tp=1, + pp=1, + seq_parallel=1, +): + """Minimal config for WorldMap construction.""" + return OmegaConf.create({ + "world": { + "actor_fraction": actor_fraction, + "finetune_fraction": finetune_fraction, + "preprocessor_fraction": preprocessor_fraction, + "replicas": replicas, + "actor_group_port": 9000, + "environment_start_port": 7777, + }, + "vllm_config": { + "vllm_kwargs": { + "tensor-parallel-size": tp, + "pipeline-parallel-size": pp, + } + }, + "finetune": {"seq_parallel": seq_parallel}, + "use_fast_llm": use_fast_llm, + "debug": {"mode": "", "place_inference_workers": True}, + }) + + +def _make_world_map(cfg, world_size, rank=0, master_addr="dns-test-0"): + from pipelinerl.world import WorldMap + env = { + "WORLD_SIZE": str(world_size), + "RANK": str(rank), + "MASTER_ADDR": master_addr, + } + with patch.dict(os.environ, env, clear=False): + # collect_environment_specs needs cfg fields that don't exist in minimal cfg; + # patch it out to avoid AttributeError. + with patch("pipelinerl.world.WorldMap._place_environments"): + with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): + return WorldMap(cfg, verbose=False) + + +# --------------------------------------------------------------------------- +# WorldMap topology tests +# --------------------------------------------------------------------------- + +class TestWorldMapMultiNode: + + def test_2node_1actor_1finetune_whole_nodes(self): + """2 nodes: 1 actor node + 1 finetune node — each gets all 8 GPUs.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) + wm = _make_world_map(cfg, world_size=2) + + assert wm.total_finetune_gpus == 8, "finetune should get exactly 1 full node" + assert wm.total_finetune_gpus % wm.node_size == 0 + assert len(wm.nodes_with_finetuning()) == 1 + + def test_4node_1actor_3finetune_whole_nodes(self): + """4 nodes: 1 actor node + 3 finetune nodes.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=3) + wm = _make_world_map(cfg, world_size=4) + + assert wm.total_finetune_gpus == 24, "finetune should get exactly 3 full nodes" + assert wm.total_finetune_gpus % wm.node_size == 0 + assert len(wm.nodes_with_finetuning()) == 3 + + def test_4node_2actor_2finetune_whole_nodes(self): + """4 nodes: 2 actor nodes + 2 finetune nodes.""" + cfg = _make_cfg(actor_fraction=2, finetune_fraction=2) + wm = _make_world_map(cfg, world_size=4) + + assert wm.total_finetune_gpus == 16 + assert wm.total_finetune_gpus % wm.node_size == 0 + assert len(wm.nodes_with_finetuning()) == 2 + + def test_finetune_always_at_least_one_node(self): + """Even with a large actor fraction, finetune gets at least 1 full node.""" + cfg = _make_cfg(actor_fraction=3, finetune_fraction=1) + wm = _make_world_map(cfg, world_size=4) + + assert len(wm.nodes_with_finetuning()) >= 1 + assert wm.total_finetune_gpus >= wm.node_size + assert wm.total_finetune_gpus % wm.node_size == 0 + + def test_actors_never_exceed_world_size_minus_one(self): + """Actor nodes never consume all nodes — at least 1 reserved for finetune.""" + cfg = _make_cfg(actor_fraction=10, finetune_fraction=1) + wm = _make_world_map(cfg, world_size=4) + + finetune_nodes = len(wm.nodes_with_finetuning()) + assert finetune_nodes >= 1 + assert finetune_nodes < 4 + + def test_single_node_unchanged(self): + """Single-node path is not affected by the multi-node rounding.""" + cfg = _make_cfg(actor_fraction=2, finetune_fraction=6) + # Single-node: world_size=1, node_size = actual device count (mocked) + with patch("torch.cuda.device_count", return_value=8): + with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): + with patch("pipelinerl.world.WorldMap._place_environments"): + from pipelinerl.world import WorldMap + wm = WorldMap(cfg, verbose=False) + assert wm.total_finetune_gpus == 6 + assert wm.world_size == 1 + + def test_nodes_with_finetuning_returns_sorted_ranks(self): + """nodes_with_finetuning() returns a sorted list of node ranks.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=3) + wm = _make_world_map(cfg, world_size=4) + + fn = wm.nodes_with_finetuning() + assert fn == sorted(fn) + + def test_my_finetuning_rank_on_finetune_node(self): + """my_finetuning_rank() returns 0 for the first finetune node.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) + # With 2 nodes, finetune is on node 0 (actor on node 1 due to reversed placement) + wm = _make_world_map(cfg, world_size=2, rank=0) + + finetune_nodes = wm.nodes_with_finetuning() + # my_rank=0 should be a finetune node + assert 0 in finetune_nodes + assert wm.my_finetuning_rank() == finetune_nodes.index(0) + + def test_4node_with_preprocessor_all_whole_nodes(self): + """4 nodes, actor=1, preprocessor=1, finetune=6: all three get whole nodes.""" + cfg = _make_cfg(actor_fraction=1, preprocessor_fraction=1, finetune_fraction=6) + wm = _make_world_map(cfg, world_size=4) + + assert wm.total_finetune_gpus % wm.node_size == 0, "finetune must be whole nodes" + # preprocessor and actor GPU shares should also be multiples of node_size + total = wm.world_size * wm.node_size + actor_gpus = total - wm.total_finetune_gpus - wm.gpus_per_preprocessor * cfg.world.replicas + assert actor_gpus % wm.node_size == 0, "actor must be whole nodes" + assert (wm.gpus_per_preprocessor * cfg.world.replicas) % wm.node_size == 0, "preprocessor must be whole nodes" + assert wm.total_finetune_gpus + actor_gpus + wm.gpus_per_preprocessor * cfg.world.replicas == total + + def test_3node_with_preprocessor_all_whole_nodes(self): + """3 nodes, actor=1, preprocessor=1, finetune=1: each component gets 1 node.""" + cfg = _make_cfg(actor_fraction=1, preprocessor_fraction=1, finetune_fraction=1) + wm = _make_world_map(cfg, world_size=3) + + assert wm.total_finetune_gpus % wm.node_size == 0 + total = wm.world_size * wm.node_size + actor_gpus = total - wm.total_finetune_gpus - wm.gpus_per_preprocessor * cfg.world.replicas + assert actor_gpus % wm.node_size == 0 + assert (wm.gpus_per_preprocessor * cfg.world.replicas) % wm.node_size == 0 + + def test_address_map_derived_from_master_addr(self): + """address_map entries follow the dns-- pattern.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + + assert wm.address_map[0] == "dns-abc123-0" + assert wm.address_map[1] == "dns-abc123-1" + + +# --------------------------------------------------------------------------- +# torchrun command assembly test +# --------------------------------------------------------------------------- + +class TestTorchrunCommand: + + def _capture_cmd(self, world_map, cfg_extra=None): + """Run _run_finetune_fast_llm with mocked I/O and capture the torchrun command.""" + from pipelinerl.launch import _run_finetune_fast_llm + + cfg = OmegaConf.create({ + "model_path": "/tmp/fake_model", + "weight_broadcast": False, + "debug": {"mode": "", "log_data_pipeline": False}, + "streams": {"host": "localhost", "port": 11000}, + "wandb": { + "wandb_workspace_root": "/tmp", + "wandb_entity_name": "test", + "wandb_project_name": "test", + "wandb_group": "test", + }, + "fast_llm": { + "training": { + "train_iters": 10, + "wandb": {"entity_name": None, "project_name": None, "group_name": None}, + }, + "data": {"datasets": {"training": {"type": "streaming", "host": None, "port": None}}}, + "pretrained": {"format": "llama", "path": None, "model_weights": True}, + "run": {"experiment_dir": None, "experiment_name": None}, + "callbacks": {}, + }, + "fast_llm_finetune": { + "model_type": "llama", + "torchrun_port": 29500, + "model_format": "llama", + }, + }) + if cfg_extra: + cfg = OmegaConf.merge(cfg, OmegaConf.create(cfg_extra)) + + captured_cmd = [] + + def mock_popen(cmd, **kwargs): + captured_cmd.extend(cmd) + return None # no process spawned + + with tempfile.TemporaryDirectory() as tmp: + exp_dir = Path(tmp) + # Patch os.path.isdir to pass the model_path check + with patch("pipelinerl.launch._popen", side_effect=mock_popen): + with patch("pipelinerl.launch.save_command"): + with patch("os.path.isdir", return_value=True): + list(_run_finetune_fast_llm(cfg, world_map, gpus=[0, 1, 2, 3], exp_dir=exp_dir)) + + return captured_cmd + + def test_single_node_uses_master_port(self): + """Single-node torchrun uses --master_port, no rdzv flags.""" + cfg = _make_cfg(actor_fraction=2, finetune_fraction=6) + with patch("torch.cuda.device_count", return_value=8): + with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): + with patch("pipelinerl.world.WorldMap._place_environments"): + from pipelinerl.world import WorldMap + wm = WorldMap(cfg, verbose=False) + + cmd = self._capture_cmd(wm) + assert "--master_port=29500" in cmd + assert "--rdzv_backend=static" not in cmd + assert "--nnodes=6" not in cmd + + def test_2node_1finetune_uses_single_node_torchrun(self): + """2-node job with 1 actor + 1 finetune node: fast-llm spans 1 node → single-node torchrun.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) + wm = _make_world_map(cfg, world_size=2, rank=0, master_addr="dns-abc-0") + + assert len(wm.nodes_with_finetuning()) == 1, "only 1 finetune node in 2-node job" + cmd = self._capture_cmd(wm) + # Should use simple --master_port, not rdzv + assert "--master_port=29500" in cmd + assert "--rdzv_backend=static" not in cmd + + def test_multi_node_uses_static_rdzv(self): + """Fast-llm spanning multiple nodes uses static rdzv with correct nnodes and node_rank.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=3) + wm = _make_world_map(cfg, world_size=4, rank=0, master_addr="dns-abc-0") + + assert len(wm.nodes_with_finetuning()) == 3 + cmd = self._capture_cmd(wm) + assert "--rdzv_backend=static" in cmd + assert "--rdzv_id=0" in cmd + assert "--max_restarts=0" in cmd + finetune_count = len(wm.nodes_with_finetuning()) + assert f"--nnodes={finetune_count}" in cmd + assert f"--node_rank={wm.my_finetuning_rank()}" in cmd + finetune_master = wm.address_map[wm.nodes_with_finetuning()[0]] + assert any(f"--rdzv_endpoint={finetune_master}:29500" in arg for arg in cmd) + assert not any("--master_port" in arg for arg in cmd) + + def test_multi_node_4nodes_correct_nnodes(self): + """4-node job: torchrun nnodes = 3 (finetune nodes only).""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=3) + wm = _make_world_map(cfg, world_size=4, rank=0) + + cmd = self._capture_cmd(wm) + finetune_count = len(wm.nodes_with_finetuning()) + assert finetune_count == 3 + assert f"--nnodes={finetune_count}" in cmd + + +# --------------------------------------------------------------------------- +# DeepSpeed regression: snapping must NOT apply when use_fast_llm=False +# --------------------------------------------------------------------------- + +class TestWorldMapDeepSpeed: + + def test_deepspeed_single_node_fractional_split(self): + """Single-node DeepSpeed split is unchanged — 2 actor GPUs + 6 finetune GPUs.""" + cfg = _make_cfg(actor_fraction=2, finetune_fraction=6, use_fast_llm=False) + with patch("torch.cuda.device_count", return_value=8): + with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): + with patch("pipelinerl.world.WorldMap._place_environments"): + from pipelinerl.world import WorldMap + wm = WorldMap(cfg, verbose=False) + + assert wm.total_finetune_gpus == 6 + assert wm.world_size == 1 + + def test_deepspeed_multinode_no_rounding(self): + """Multi-node DeepSpeed: no whole-node snapping (handled by DeepSpeed itself).""" + # 2 nodes, actor_fraction=1, finetune_fraction=1 → 8 finetune GPUs (happens to be whole node) + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) + wm = _make_world_map(cfg, world_size=2) + # Should still compute correctly without triggering fast-llm rounding path + assert wm.total_finetune_gpus > 0 + assert wm.world_size == 2 + + def test_fast_llm_single_node_unchanged(self): + """Single-node fast-llm: fractional split within one node is preserved.""" + cfg = _make_cfg(actor_fraction=2, finetune_fraction=6, use_fast_llm=True) + with patch("torch.cuda.device_count", return_value=8): + with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): + with patch("pipelinerl.world.WorldMap._place_environments"): + from pipelinerl.world import WorldMap + wm = WorldMap(cfg, verbose=False) + + assert wm.total_finetune_gpus == 6 + assert wm.world_size == 1 + + +# --------------------------------------------------------------------------- +# Pod IP exchange: dns_address_map, job URL rewriting, DeepSpeed/fast-llm compat +# --------------------------------------------------------------------------- + +def _simulate_pod_ip_exchange(wm, pod_ips: dict): + """Simulate _exchange_pod_ips without NFS I/O. + + Sets dns_address_map to original DNS names, updates address_map and job + URLs/hostnames to pod IPs — mirrors the real function's side-effects. + """ + from pipelinerl.launch import _exchange_pod_ips as real_fn # noqa: F401 (not called) + # Save DNS names first (matches the real implementation order) + wm.dns_address_map = dict(wm.address_map) + # Overwrite address_map with pod IPs + for rank, ip in pod_ips.items(): + wm.address_map[rank] = ip + wm.master_addr = pod_ips[0] + # Rewrite job URLs/hostnames + for node, jobs in wm.job_map.items(): + dns_name = wm.dns_address_map[node] + pod_ip = pod_ips[node] + for job in jobs: + job.hostname = pod_ip + if job.url: + job.url = job.url.replace(dns_name, pod_ip) + + +class TestPodIPExchange: + + def test_dns_address_map_holds_original_dns_names(self): + """After pod IP exchange, dns_address_map contains original DNS names, not pod IPs.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + + pod_ips = {0: "10.0.0.1", 1: "10.0.0.2"} + _simulate_pod_ip_exchange(wm, pod_ips) + + assert wm.dns_address_map[0] == "dns-abc123-0" + assert wm.dns_address_map[1] == "dns-abc123-1" + + def test_address_map_updated_to_pod_ips(self): + """After pod IP exchange, address_map and master_addr hold pod IPs.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + + pod_ips = {0: "10.0.0.1", 1: "10.0.0.2"} + _simulate_pod_ip_exchange(wm, pod_ips) + + assert wm.address_map[0] == "10.0.0.1" + assert wm.address_map[1] == "10.0.0.2" + assert wm.master_addr == "10.0.0.1" + + def test_job_urls_rewritten_to_pod_ips(self): + """After pod IP exchange, actor_llm job URLs use pod IPs, not DNS names.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + + # Verify that actor_llm jobs have DNS-based URLs before exchange + actor_urls_before = [job.url for job in wm.get_all_jobs() if job.kind == "actor_llm"] + assert all("dns-abc123-1" in u for u in actor_urls_before) + + pod_ips = {0: "10.0.0.1", 1: "10.0.0.2"} + _simulate_pod_ip_exchange(wm, pod_ips) + + actor_urls_after = [job.url for job in wm.get_all_jobs() if job.kind == "actor_llm"] + assert all("10.0.0.2" in u for u in actor_urls_after), f"Expected pod IP in URLs: {actor_urls_after}" + assert all("dns-abc123" not in u for u in actor_urls_after) + + def test_no_dns_address_map_without_exchange(self): + """Without pod IP exchange, dns_address_map is not set (no AttributeError).""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + assert not hasattr(wm, "dns_address_map") + + +# --------------------------------------------------------------------------- +# DeepSpeed command assembly: hostfile and inclusion filter use DNS names +# --------------------------------------------------------------------------- + +class TestDeepSpeedCommand: + + def _make_ds_cfg(self): + return OmegaConf.create({ + "use_deepspeed": True, + "use_fsdp": False, + "deepspeed_config": "zero2", + "accelerate_config": None, + "world": {"actor_group_port": 9000}, + "debug": {"mode": ""}, + }) + + def _capture_ds_cmd(self, world_map, cfg_extra=None): + """Run _run_finetune_deepspeed with mocked I/O and capture the command.""" + from pipelinerl.launch import _run_finetune_deepspeed + + cfg = self._make_ds_cfg() + if cfg_extra: + cfg = OmegaConf.merge(cfg, OmegaConf.create(cfg_extra)) + + captured_cmd = [] + + def mock_popen(cmd, **kwargs): + captured_cmd.extend(cmd) + return None + + with tempfile.TemporaryDirectory() as tmp: + exp_dir = Path(tmp) + (exp_dir / "hostfile.txt").write_text("") # pre-create + with patch("pipelinerl.launch._popen", side_effect=mock_popen): + with patch("pipelinerl.launch.save_command"): + with patch.dict(os.environ, {"MASTER_ADDR": "dns-test-0", "MASTER_PORT": "29501"}): + list(_run_finetune_deepspeed(cfg, world_map, gpus=[0, 1, 2, 3], exp_dir=exp_dir)) + + return captured_cmd + + def test_deepspeed_multinode_uses_dns_names_without_exchange(self): + """DeepSpeed 2-node without pod IP exchange: inclusion filter uses DNS names.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + + cmd = self._capture_ds_cmd(wm) + # The deepspeed_inclusion_filter should contain the DNS hostname for the finetune node + filter_arg = next((c for c in cmd if "dns-abc123" in c), None) + assert filter_arg is not None, f"Expected DNS name in cmd, got: {cmd}" + + def test_deepspeed_multinode_after_pod_ip_exchange_uses_dns_names(self): + """After pod IP exchange, DeepSpeed inclusion filter still uses DNS names (not pod IPs).""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + + # Simulate pod IP exchange + _simulate_pod_ip_exchange(wm, {0: "10.0.0.1", 1: "10.0.0.2"}) + + cmd = self._capture_ds_cmd(wm) + # Inclusion filter must still use DNS names, not pod IPs + filter_arg = next((c for c in cmd if "dns-abc123" in c), None) + assert filter_arg is not None, f"Expected DNS name in DS filter after pod IP exchange, got: {cmd}" + # Pod IPs must NOT appear in the inclusion filter + assert not any("10.0.0" in c for c in cmd if "--deepspeed_inclusion_filter" not in c and "@" in c), \ + f"Pod IP leaked into DS filter: {cmd}" + + def test_deepspeed_single_node_no_pod_ip_exchange(self): + """Single-node DeepSpeed: no world_size>1 branch, pod IP exchange never runs.""" + cfg = _make_cfg(actor_fraction=2, finetune_fraction=6, use_fast_llm=False) + with patch("torch.cuda.device_count", return_value=8): + with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): + with patch("pipelinerl.world.WorldMap._place_environments"): + from pipelinerl.world import WorldMap + wm = WorldMap(cfg, verbose=False) + + assert wm.world_size == 1 + assert not hasattr(wm, "dns_address_map") + # Should not crash even without dns_address_map + cmd = self._capture_ds_cmd(wm) + assert "--num_machines" not in cmd # single-node, no multi-machine flags + + +# --------------------------------------------------------------------------- +# Hostfile creation in main(): uses dns_address_map after pod IP exchange +# --------------------------------------------------------------------------- + +class TestHostfileCreation: + + def test_hostfile_uses_dns_names_after_pod_ip_exchange(self): + """The DeepSpeed hostfile written by main() uses DNS names even after pod IP exchange.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + + # Simulate pod IP exchange + _simulate_pod_ip_exchange(wm, {0: "10.0.0.1", 1: "10.0.0.2"}) + + dns_map = getattr(wm, "dns_address_map", wm.address_map) + hosts = [dns_map[i] for i in range(wm.world_size)] + + assert hosts[0] == "dns-abc123-0" + assert hosts[1] == "dns-abc123-1" + assert "10.0.0" not in hosts[0] + assert "10.0.0" not in hosts[1] + + def test_hostfile_uses_address_map_without_exchange(self): + """Without pod IP exchange, dns_address_map is absent — falls back to address_map.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + + dns_map = getattr(wm, "dns_address_map", wm.address_map) + hosts = [dns_map[i] for i in range(wm.world_size)] + + assert hosts[0] == "dns-abc123-0" + assert hosts[1] == "dns-abc123-1" + + +# --------------------------------------------------------------------------- +# Redis host in saved exp_config.yaml for multi-node (DeepSpeed + Redis) +# --------------------------------------------------------------------------- + +class TestRedisHostMultiNode: + + def _compute_streams_host(self, world_map, my_rank: int) -> str: + """Mirror the launch.py logic for cfg.streams.host selection.""" + if world_map.world_size > 1: + return world_map.master_addr + return "localhost" + + def test_single_node_redis_host_is_localhost(self): + """Single-node: Redis host is localhost regardless of pod IP exchange.""" + cfg = _make_cfg(actor_fraction=2, finetune_fraction=6, use_fast_llm=False) + with patch("torch.cuda.device_count", return_value=8): + with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): + with patch("pipelinerl.world.WorldMap._place_environments"): + from pipelinerl.world import WorldMap + wm = WorldMap(cfg, verbose=False) + + host = self._compute_streams_host(wm, my_rank=0) + assert host == "localhost" + + def test_multinode_rank0_redis_host_is_pod_ip(self): + """Multi-node rank 0: Redis host is pod IP (not localhost) after exchange. + + This ensures the saved exp_config.yaml has a reachable address for + DeepSpeed workers on other nodes. + """ + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + _simulate_pod_ip_exchange(wm, {0: "10.0.0.1", 1: "10.0.0.2"}) + + host = self._compute_streams_host(wm, my_rank=0) + assert host == "10.0.0.1", "rank 0 should use pod IP so saved config is reachable cross-node" + assert host != "localhost" + + def test_multinode_rank1_redis_host_is_pod_ip(self): + """Multi-node rank 1: Redis host is pod IP of rank 0.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0", rank=1) + _simulate_pod_ip_exchange(wm, {0: "10.0.0.1", 1: "10.0.0.2"}) + + host = self._compute_streams_host(wm, my_rank=1) + assert host == "10.0.0.1", "rank 1 should use rank 0's pod IP to reach Redis" + + def test_multinode_both_ranks_same_redis_host(self): + """Both ranks in a 2-node job resolve to the same Redis host (pod IP of rank 0).""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) + wm0 = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0", rank=0) + wm1 = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0", rank=1) + + _simulate_pod_ip_exchange(wm0, {0: "10.0.0.1", 1: "10.0.0.2"}) + _simulate_pod_ip_exchange(wm1, {0: "10.0.0.1", 1: "10.0.0.2"}) + + host0 = self._compute_streams_host(wm0, my_rank=0) + host1 = self._compute_streams_host(wm1, my_rank=1) + + assert host0 == host1 == "10.0.0.1" + + def test_multinode_without_pod_ip_exchange_uses_master_addr(self): + """Without pod IP exchange, multi-node uses master_addr (DNS name) for Redis. + + This is a fallback; the pod IP exchange should always run in practice + but the code must not crash without it. + """ + cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) + wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") + + # No pod IP exchange — master_addr is still a DNS name + assert wm.master_addr == "dns-abc123-0" + host = self._compute_streams_host(wm, my_rank=0) + assert host == "dns-abc123-0" # DNS name (port filtering may apply, but code doesn't crash) + + +# --------------------------------------------------------------------------- +# DeepSpeed run_finetune.py path: must be absolute (not relative to CWD) +# --------------------------------------------------------------------------- + +class TestDeepSpeedEntrypointPath: + + def _capture_ds_cmd(self, world_map): + from pipelinerl.launch import _run_finetune_deepspeed + from omegaconf import OmegaConf + + cfg = OmegaConf.create({ + "use_deepspeed": True, + "use_fsdp": False, + "deepspeed_config": "zero2", + "accelerate_config": None, + "world": {"actor_group_port": 9000}, + "debug": {"mode": ""}, + }) + captured_cmd = [] + + def mock_popen(cmd, **kwargs): + captured_cmd.extend(cmd) + return None + + with tempfile.TemporaryDirectory() as tmp: + exp_dir = Path(tmp) + with patch("pipelinerl.launch._popen", side_effect=mock_popen): + with patch("pipelinerl.launch.save_command"): + with patch.dict(os.environ, {"MASTER_ADDR": "dns-test-0", "MASTER_PORT": "29501"}): + list(_run_finetune_deepspeed(cfg, world_map, gpus=[0, 1, 2, 3], exp_dir=exp_dir)) + + return captured_cmd + + def test_run_finetune_path_is_absolute(self): + """run_finetune.py must be an absolute path so it works regardless of CWD. + + When EAI starts the pod, CWD is /home/toolkit (not the repo root). A relative + path like 'pipelinerl/entrypoints/run_finetune.py' resolves to + '/home/toolkit/pipelinerl/...' which doesn't exist. + """ + cfg = _make_cfg(actor_fraction=2, finetune_fraction=6, use_fast_llm=False) + with patch("torch.cuda.device_count", return_value=8): + with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): + with patch("pipelinerl.world.WorldMap._place_environments"): + from pipelinerl.world import WorldMap + wm = WorldMap(cfg, verbose=False) + + cmd = self._capture_ds_cmd(wm) + + # Find the run_finetune.py argument + finetune_script = next((c for c in cmd if "run_finetune.py" in c), None) + assert finetune_script is not None, f"run_finetune.py not found in cmd: {cmd}" + assert Path(finetune_script).is_absolute(), ( + f"run_finetune.py path must be absolute but got: {finetune_script!r}. " + "A relative path resolves against CWD which is /home/toolkit in EAI pods." + ) + assert Path(finetune_script).exists(), ( + f"run_finetune.py absolute path must exist: {finetune_script!r}" + ) + + +# --------------------------------------------------------------------------- +# Per-node file naming: fast-llm and DeepSpeed avoid NFS write races +# --------------------------------------------------------------------------- + +class TestPerNodeFileNaming: + """Verify that multinode fast-llm and DeepSpeed finetune runs write separate + output files per node (config, start.sh, stdout, stderr) to avoid NFS races.""" + + def _capture_fast_llm_files(self, world_map, gpus=None): + """Run _run_finetune_fast_llm and return captured file suffix info.""" + from pipelinerl.launch import _run_finetune_fast_llm + + cfg = OmegaConf.create({ + "model_path": "/tmp/fake_model", + "weight_broadcast": False, + "debug": {"mode": "", "log_data_pipeline": False}, + "streams": {"host": "localhost", "port": 11000}, + "wandb": { + "wandb_workspace_root": "/tmp", + "wandb_entity_name": "test", + "wandb_project_name": "test", + "wandb_group": "test", + }, + "fast_llm": { + "training": { + "train_iters": 10, + "wandb": {"entity_name": None, "project_name": None, "group_name": None}, + }, + "data": {"datasets": {"training": {"type": "streaming", "host": None, "port": None}}}, + "pretrained": {"format": "llama", "path": None, "model_weights": True}, + "run": {"experiment_dir": None, "experiment_name": None}, + "callbacks": {}, + }, + "fast_llm_finetune": { + "model_type": "llama", + "torchrun_port": 29500, + "model_format": "llama", + }, + }) + + written_files = {} + + real_open = open + + def mock_popen(cmd, **kwargs): + written_files["stdout"] = str(kwargs.get("stdout", {}).name if hasattr(kwargs.get("stdout"), "name") else "") + written_files["stderr"] = str(kwargs.get("stderr", {}).name if hasattr(kwargs.get("stderr"), "name") else "") + return None + + captured_save = {} + + def mock_save_command(script_dir, cmd, suffix=""): + captured_save["suffix"] = suffix + captured_save["dir"] = str(script_dir) + + captured_config = {} + + real_omegaconf_save = None + + with tempfile.TemporaryDirectory() as tmp: + exp_dir = Path(tmp) + with patch("pipelinerl.launch._popen", side_effect=mock_popen): + with patch("pipelinerl.launch.save_command", side_effect=mock_save_command): + with patch("os.path.isdir", return_value=True): + with patch("omegaconf.OmegaConf.save") as mock_cfg_save: + list(_run_finetune_fast_llm(cfg, world_map, gpus=gpus or [0, 1, 2, 3], exp_dir=exp_dir)) + if mock_cfg_save.call_args: + # OmegaConf.save(cfg, path) — second positional arg is path + args = mock_cfg_save.call_args[0] + captured_config["path"] = str(args[1]) if len(args) > 1 else "" + + return { + "config_path": captured_config.get("path", ""), + "save_suffix": captured_save.get("suffix", ""), + "stdout": written_files.get("stdout", ""), + "stderr": written_files.get("stderr", ""), + } + + def _capture_deepspeed_files(self, world_map, gpus=None): + """Run _run_finetune_deepspeed and return captured file suffix.""" + from pipelinerl.launch import _run_finetune_deepspeed + + cfg = OmegaConf.create({ + "use_deepspeed": True, + "use_fsdp": False, + "deepspeed_config": "zero2", + "accelerate_config": None, + "world": {"actor_group_port": 9000}, + "debug": {"mode": ""}, + }) + + captured_save = {} + written_files = {} + + def mock_popen(cmd, **kwargs): + written_files["stdout"] = str(kwargs.get("stdout", {}).name if hasattr(kwargs.get("stdout"), "name") else "") + written_files["stderr"] = str(kwargs.get("stderr", {}).name if hasattr(kwargs.get("stderr"), "name") else "") + return None + + def mock_save_command(script_dir, cmd, suffix=""): + captured_save["suffix"] = suffix + + with tempfile.TemporaryDirectory() as tmp: + exp_dir = Path(tmp) + with patch("pipelinerl.launch._popen", side_effect=mock_popen): + with patch("pipelinerl.launch.save_command", side_effect=mock_save_command): + with patch.dict(os.environ, {"MASTER_ADDR": "dns-test-0", "MASTER_PORT": "29501"}): + list(_run_finetune_deepspeed(cfg, world_map, gpus=gpus or [0, 1, 2, 3], exp_dir=exp_dir)) + + return { + "save_suffix": captured_save.get("suffix", ""), + "stdout": written_files.get("stdout", ""), + "stderr": written_files.get("stderr", ""), + } + + # --- fast-llm single-node: no suffix --- + + def test_fast_llm_single_node_no_suffix(self): + """Single-node fast-llm: no _node0 suffix — backward compat.""" + cfg = _make_cfg(actor_fraction=2, finetune_fraction=6) + with patch("torch.cuda.device_count", return_value=8): + with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): + with patch("pipelinerl.world.WorldMap._place_environments"): + from pipelinerl.world import WorldMap + wm = WorldMap(cfg, verbose=False) + + result = self._capture_fast_llm_files(wm) + assert result["save_suffix"] == "", f"Single-node must have no suffix, got: {result['save_suffix']!r}" + assert "_node" not in result["config_path"], f"Single-node config must have no _node suffix: {result['config_path']}" + + # --- fast-llm multinode: each node gets its own suffix --- + + def test_fast_llm_multinode_node0_suffix(self): + """4-node fast-llm, finetune node 0: files get _node0 suffix. + Actor takes the last node (rank 3), so ranks 0/1/2 are finetune.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=3) + wm = _make_world_map(cfg, world_size=4, rank=0) # rank 0 = first finetune node + + result = self._capture_fast_llm_files(wm) + assert result["save_suffix"] == "_node0", f"Expected _node0, got: {result['save_suffix']!r}" + assert "_node0" in result["config_path"], f"Config path must contain _node0: {result['config_path']}" + + def test_fast_llm_multinode_node1_suffix(self): + """4-node fast-llm, finetune node 1: files get _node1 suffix.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=3) + wm = _make_world_map(cfg, world_size=4, rank=1) # rank 1 = second finetune node + + result = self._capture_fast_llm_files(wm) + assert result["save_suffix"] == "_node1", f"Expected _node1, got: {result['save_suffix']!r}" + assert "_node1" in result["config_path"] + + def test_fast_llm_multinode_node2_suffix(self): + """4-node fast-llm, finetune node 2: files get _node2 suffix.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=3) + wm = _make_world_map(cfg, world_size=4, rank=2) # rank 2 = third finetune node + + result = self._capture_fast_llm_files(wm) + assert result["save_suffix"] == "_node2", f"Expected _node2, got: {result['save_suffix']!r}" + + # --- DeepSpeed single-node: no suffix --- + + def test_deepspeed_single_node_no_suffix(self): + """Single-node DeepSpeed: no _node suffix.""" + cfg = _make_cfg(actor_fraction=2, finetune_fraction=6, use_fast_llm=False) + with patch("torch.cuda.device_count", return_value=8): + with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): + with patch("pipelinerl.world.WorldMap._place_environments"): + from pipelinerl.world import WorldMap + wm = WorldMap(cfg, verbose=False) + + result = self._capture_deepspeed_files(wm) + assert result["save_suffix"] == "", f"Single-node must have no suffix, got: {result['save_suffix']!r}" + + # --- DeepSpeed multinode: each node gets its own suffix --- + + def test_deepspeed_multinode_node0_suffix(self): + """4-node DeepSpeed, finetune node 0: save_command gets _node0 suffix. + Actor takes the last node (rank 3), so ranks 0/1/2 are finetune.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=3, use_fast_llm=False) + wm = _make_world_map(cfg, world_size=4, rank=0) # rank 0 = first finetune node + + result = self._capture_deepspeed_files(wm) + assert result["save_suffix"] == "_node0", f"Expected _node0, got: {result['save_suffix']!r}" + + def test_deepspeed_multinode_node2_suffix(self): + """4-node DeepSpeed, finetune node 2: save_command gets _node2 suffix.""" + cfg = _make_cfg(actor_fraction=1, finetune_fraction=3, use_fast_llm=False) + wm = _make_world_map(cfg, world_size=4, rank=2) # rank 2 = third finetune node + + result = self._capture_deepspeed_files(wm) + assert result["save_suffix"] == "_node2", f"Expected _node2, got: {result['save_suffix']!r}" diff --git a/tests/trainer_test_utils.py b/tests/trainer_test_utils.py new file mode 100644 index 00000000..d6de57e5 --- /dev/null +++ b/tests/trainer_test_utils.py @@ -0,0 +1,128 @@ +"""Shared utilities for trainer helper scripts (both HTTP and fast-llm variants).""" + + +def _resolve_model_path(model_name: str): + """Resolve model name to a local Path, downloading from HuggingFace if needed.""" + from pathlib import Path + from huggingface_hub import snapshot_download + + model_path = Path(model_name) + if not model_path.exists(): + print(f"[Trainer] Downloading model from HuggingFace Hub: {model_name}") + model_path = Path(snapshot_download(model_name)) + return model_path + + +def _load_state_dict(model_name: str, device: str = "cuda:0") -> tuple: + """Load model state dict from safetensors files. + + Returns: + (state_dict, model_path) + """ + import json + from safetensors.torch import load_file + + model_path = _resolve_model_path(model_name) + index_file = model_path / "model.safetensors.index.json" + + if index_file.exists(): + print(f"[Trainer] Found index file, loading sharded model") + with open(index_file) as f: + index = json.load(f) + weight_map = index["weight_map"] + + file_to_params = {} + for param_name, filename in weight_map.items(): + file_to_params.setdefault(filename, []).append(param_name) + + state_dict = {} + for filename, param_names in file_to_params.items(): + file_path = model_path / filename + print(f"[Trainer] Loading {len(param_names)} parameters from {filename}") + tensors = load_file(str(file_path), device=device) + for param_name in param_names: + state_dict[param_name] = tensors[param_name] + else: + safetensors_file = model_path / "model.safetensors" + print(f"[Trainer] Loading from single file: {safetensors_file}") + state_dict = load_file(str(safetensors_file), device=device) + + print(f"[Trainer] Loaded {len(state_dict)} parameters from safetensors") + return state_dict, model_path + + +def _create_perturbed_state_dict( + state_dict: dict, seed: int = 42, noise_scale: float = 0.001 +) -> dict: + """Return a new state dict with Gaussian noise added to all tensors.""" + import torch + + print(f"[Trainer] Creating perturbed weights (all tensors) with seed={seed}...") + torch.manual_seed(seed) + perturbed = {} + for name, tensor in state_dict.items(): + perturbed_tensor = tensor.clone() + perturbed_tensor.add_(torch.randn_like(perturbed_tensor) * noise_scale) + perturbed[name] = perturbed_tensor + print( + f"[Trainer] Perturbed all {len(perturbed)} tensors with noise={noise_scale}, seed={seed}" + ) + return perturbed + + +def _init_actor_process_group(init_method: str, rank: int = 0, world_size: int = 2, group_name: str = "actor"): + """Initialize the actor NCCL process group and return it.""" + import pipelinerl.torch_utils + + print(f"[Trainer] Initializing process group as rank {rank} (group_name={group_name!r})") + process_group = pipelinerl.torch_utils.init_extra_process_group( + group_name=group_name, + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + print("[Trainer] Process group initialized") + return process_group + + +def _broadcast_tensors(state_dict: dict, process_group, log_interval: int = 50): + """Broadcast every tensor in state_dict via NCCL (src=0).""" + import torch.distributed as dist + + total = len(state_dict) + for i, (name, tensor) in enumerate(state_dict.items()): + if tensor.device.type != "cuda": + tensor = tensor.cuda(0) + dist.broadcast(tensor, src=0, group=process_group) + if (i + 1) % log_interval == 0: + print(f"[Trainer] Broadcasted {i+1}/{total} parameters") + print(f"[Trainer] All {total} parameters broadcasted") + + +def _wait_for_servers_ready(server_urls: list, extra_wait_secs: int = 10): + """Poll /health on each server until all respond 200, then sleep extra_wait_secs.""" + import time + import requests + + for server_url in server_urls: + print(f"[Trainer] Waiting for server {server_url} to be ready...") + server_ready = False + for i in range(120): # up to 2 minutes + try: + resp = requests.get(f"{server_url}/health", timeout=1) + if resp.status_code == 200: + server_ready = True + print(f"[Trainer] Server {server_url} is ready (took {i} seconds)") + break + except requests.exceptions.RequestException: + pass + time.sleep(1) + if not server_ready: + raise TimeoutError(f"Server {server_url} did not become ready within 2 minutes") + + if extra_wait_secs > 0: + print( + f"[Trainer] Waiting additional {extra_wait_secs} seconds for server(s) to fully initialize..." + ) + time.sleep(extra_wait_secs) diff --git a/tests/vllm_engine_helper.py b/tests/vllm_engine_helper.py new file mode 100755 index 00000000..798743bc --- /dev/null +++ b/tests/vllm_engine_helper.py @@ -0,0 +1,617 @@ +#!/usr/bin/env python3 +"""Helper script for running vLLM engine in a subprocess with proper CUDA isolation. + +This script is run as a separate process with CUDA_VISIBLE_DEVICES set, +ensuring the engine only sees the intended GPU. +""" + +import sys +import argparse +import asyncio + + +async def init_engine_and_process_group( + model_name: str, + init_method: str, + actor_llm_idx: int, + world_size: int, +): + """Initialize vLLM engine and process group. + + create_engine() automatically calls init_actor_update_group() when + disable_weight_updates=False, and calls destroy_actor_update_group() + on context manager exit. + """ + from pipelinerl.vllm1 import EngineManager + import argparse as ap + + print("[vLLM Engine] Starting engine initialization") + + # Create args for engine with process group params + args = ap.Namespace( + model=model_name, + tensor_parallel_size=1, + disable_log_stats=True, + enable_log_requests=False, + disable_weight_updates=False, + # Process group params - needed for automatic init_actor_update_group() + actor_llm_idx=actor_llm_idx, + weight_update_group_init_method=init_method, + weight_update_group_world_size=world_size, + ) + + print(f"[vLLM Engine] Creating engine with model={model_name}") + + # create_engine automatically: + # 1. Creates engine and manager + # 2. Calls manager.init_actor_update_group() (rank 1) + # 3. On exit, calls manager.destroy_actor_update_group() + async with EngineManager.create_engine(args) as manager: + print("[vLLM Engine] Engine and process group created successfully") + + # Keep engine alive until trainer completes its work + print("[vLLM Engine] Process group active, waiting for trainer...") + await asyncio.sleep(5) + + # Context manager exit automatically cleans up process group + print("[vLLM Engine] Engine and process group cleaned up") + + +async def test_weight_update( + model_name: str, + init_method: str, + actor_llm_idx: int, + world_size: int, + prompt: str, + max_tokens: int, + sync_dir: str, + expect_different: bool = False, +): + """Test weight update with generation before and after. + + This mode: + 1. Creates engine and initializes process group + 2. Generates baseline output + 3. Signals baseline_done, waits for broadcast_done + 4. Receives weight update + 5. Generates again with same prompt + 6. Prints both outputs for comparison + """ + from pipelinerl.vllm1 import EngineManager + from vllm import SamplingParams + from pathlib import Path + import argparse as ap + # Import sync helper from same directory + sys.path.insert(0, str(Path(__file__).parent)) + from sync_helper import SyncPoint + + print("[vLLM Engine] Starting weight update test") + + # Create sync points + sync_path = Path(sync_dir) + baseline_done = SyncPoint(sync_path, "baseline_done") + ready_to_receive = SyncPoint(sync_path, "ready_to_receive") + request_ready = SyncPoint(sync_path, "request_ready") + receiving_started = SyncPoint(sync_path, "receiving_started") + broadcast_done = SyncPoint(sync_path, "broadcast_done") + + # Create args for engine with process group params + args = ap.Namespace( + model=model_name, + tensor_parallel_size=1, + disable_log_stats=True, + enable_log_requests=False, + disable_weight_updates=False, + actor_llm_idx=actor_llm_idx, + weight_update_group_init_method=init_method, + weight_update_group_world_size=world_size, + ) + + print(f"[vLLM Engine] Creating engine with model={model_name}") + + async with EngineManager.create_engine(args) as manager: + print("[vLLM Engine] Engine and process group created successfully") + + # Step 1: Generate baseline + sampling_params = SamplingParams( + temperature=0.0, + top_p=1.0, + max_tokens=max_tokens, + seed=42, + ) + + print(f"[vLLM Engine] Generating baseline with prompt: '{prompt}'") + async for output in manager.engine.generate( + prompt, + sampling_params=sampling_params, + request_id="baseline", + ): + baseline_output = output + + baseline_text = baseline_output.outputs[0].text + print(f"[vLLM Engine] Baseline output: '{baseline_text}'") + + # Step 2: Signal baseline done and ready to receive + baseline_done.signal() + ready_to_receive.signal() + + # Step 3: Wait for trainer to send WeightUpdateRequest + print("[vLLM Engine] Waiting for trainer to send weight update request...") + request_ready.wait(timeout=60) + + # Step 4: Read WeightUpdateRequest from trainer + from sync_helper import read_weight_update_request + request = read_weight_update_request(sync_path) + print(f"[vLLM Engine] Received request with {len(request.parameters_info)} parameters") + + # Step 5: Signal we're about to start receiving, then call receive_weight_update + receiving_started.signal() + print("[vLLM Engine] Signaled receiving_started, calling receive_weight_update...") + print("[vLLM Engine] (This will block until trainer broadcasts all weights)") + await manager.receive_weight_update(request) + print("[vLLM Engine] Weight update received!") + + # Step 6: Wait for trainer to signal broadcast complete + broadcast_done.wait(timeout=60) + print("[vLLM Engine] Trainer confirmed broadcast complete") + + # Step 7: Generate again with same prompt + print(f"[vLLM Engine] Generating after update with prompt: '{prompt}'") + async for output in manager.engine.generate( + prompt, + sampling_params=sampling_params, + request_id="after_update", + ): + updated_output = output + + updated_text = updated_output.outputs[0].text + print(f"[vLLM Engine] Updated output: '{updated_text}'") + + # Step 8: Compare outputs + if expect_different: + # Perturbed weights - expect different outputs + if baseline_text != updated_text: + print("[vLLM Engine] ✓ Outputs differ (as expected for perturbed weights)") + print(f"[vLLM Engine] Baseline: '{baseline_text}'") + print(f"[vLLM Engine] Updated: '{updated_text}'") + else: + print("[vLLM Engine] ✗ Outputs are the same!") + print(f"[vLLM Engine] Both: '{baseline_text}'") + print("[vLLM Engine] ERROR: Perturbed weights should have changed the output") + sys.exit(1) + else: + # Same weights - expect same outputs + if baseline_text == updated_text: + print("[vLLM Engine] ✓ Outputs match (as expected for same weights)") + else: + print("[vLLM Engine] ✗ Outputs differ!") + print(f"[vLLM Engine] Baseline: '{baseline_text}'") + print(f"[vLLM Engine] Updated: '{updated_text}'") + sys.exit(1) + + print("[vLLM Engine] Engine and process group cleaned up") + + +async def test_cross_validation( + model_name: str, + init_method: str, + actor_llm_idx: int, + world_size: int, + prompt: str, + max_tokens: int, + sync_dir: str, +): + """Cross-validation test for weight updates. + + Tests that broadcasting weights produces same results as loading from disk. + Flow: + 1. Generate with original model → res_un_1 + 2. Receive perturbed weights, generate → res_mod_1 + 3. Recreate engine with perturbed model from disk, generate → res_mod_2 + 4. Receive original weights, generate → res_un_2 + 5. Verify: res_un_1 == res_un_2 and res_mod_1 == res_mod_2 + """ + from pipelinerl.vllm1 import EngineManager + from vllm import SamplingParams + from pathlib import Path + import argparse as ap + sys.path.insert(0, str(Path(__file__).parent)) + from sync_helper import SyncPoint, read_weight_update_request + + print("[vLLM Engine] Starting cross-validation test") + + # Create sync points + sync_path = Path(sync_dir) + baseline_done = SyncPoint(sync_path, "baseline_done") + perturbed_model_saved = SyncPoint(sync_path, "perturbed_model_saved") + ready_to_receive_perturbed = SyncPoint(sync_path, "ready_to_receive_perturbed") + perturbed_broadcast_done = SyncPoint(sync_path, "perturbed_broadcast_done") + mod1_done = SyncPoint(sync_path, "mod1_done") + first_engine_destroyed = SyncPoint(sync_path, "first_engine_destroyed") + engine_recreated = SyncPoint(sync_path, "engine_recreated") + ready_to_receive_original = SyncPoint(sync_path, "ready_to_receive_original") + original_broadcast_done = SyncPoint(sync_path, "original_broadcast_done") + + sampling_params = SamplingParams( + temperature=0.0, + top_p=1.0, + max_tokens=max_tokens, + seed=42, + ) + + # Step 1: Generate with original model + args = ap.Namespace( + model=model_name, + tensor_parallel_size=1, + disable_log_stats=True, + enable_log_requests=False, + disable_weight_updates=False, + actor_llm_idx=actor_llm_idx, + weight_update_group_init_method=init_method, + weight_update_group_world_size=world_size, + ) + + print(f"[vLLM Engine] Step 1: Creating engine with original model: {model_name}") + async with EngineManager.create_engine(args) as manager: + print(f"[vLLM Engine] Generating res_un_1 with prompt: '{prompt}'") + async for output in manager.engine.generate( + prompt, + sampling_params=sampling_params, + request_id="res_un_1", + ): + res_un_1_output = output + res_un_1 = res_un_1_output.outputs[0].text + print(f"[vLLM Engine] res_un_1: '{res_un_1}'") + + baseline_done.signal() + + # Wait for perturbed model to be saved + print("[vLLM Engine] Waiting for trainer to save perturbed model...") + perturbed_model_saved.wait(timeout=180) + + # Step 2: Receive perturbed weights and generate + ready_to_receive_perturbed.signal() + print("[vLLM Engine] Waiting for perturbed weight update request...") + + # Wait a moment for request file to be written + import time + time.sleep(0.5) + + request = read_weight_update_request(sync_path) + print(f"[vLLM Engine] Received perturbed request with {len(request.parameters_info)} parameters") + + print("[vLLM Engine] Receiving perturbed weights...") + await manager.receive_weight_update(request) + + perturbed_broadcast_done.wait(timeout=900) + print("[vLLM Engine] Perturbed weights received") + + print(f"[vLLM Engine] Generating res_mod_1 with prompt: '{prompt}'") + async for output in manager.engine.generate( + prompt, + sampling_params=sampling_params, + request_id="res_mod_1", + ): + res_mod_1_output = output + res_mod_1 = res_mod_1_output.outputs[0].text + print(f"[vLLM Engine] res_mod_1: '{res_mod_1}'") + + mod1_done.signal() + + # Engine destroyed here (context manager exit) + print("[vLLM Engine] First engine destroyed") + first_engine_destroyed.signal() + + # Step 3: Recreate engine with perturbed model from disk + perturbed_model_path = (sync_path / "perturbed_model_path.txt").read_text().strip() + print(f"[vLLM Engine] Step 3: Recreating engine with perturbed model from: {perturbed_model_path}") + + args_perturbed = ap.Namespace( + model=perturbed_model_path, + tensor_parallel_size=1, + disable_log_stats=True, + enable_log_requests=False, + disable_weight_updates=False, + actor_llm_idx=actor_llm_idx, + weight_update_group_init_method=init_method, + weight_update_group_world_size=world_size, + ) + + async with EngineManager.create_engine(args_perturbed) as manager: + # Signal immediately after engine is created + engine_recreated.signal() + print("[vLLM Engine] Engine recreated, signaled to trainer") + + print(f"[vLLM Engine] Generating res_mod_2 with prompt: '{prompt}'") + async for output in manager.engine.generate( + prompt, + sampling_params=sampling_params, + request_id="res_mod_2", + ): + res_mod_2_output = output + res_mod_2 = res_mod_2_output.outputs[0].text + print(f"[vLLM Engine] res_mod_2: '{res_mod_2}'") + + # Step 4: Receive original weights and generate + ready_to_receive_original.signal() + print("[vLLM Engine] Waiting for original weight update request...") + + time.sleep(0.5) + request = read_weight_update_request(sync_path) + print(f"[vLLM Engine] Received original request with {len(request.parameters_info)} parameters") + + print("[vLLM Engine] Receiving original weights...") + await manager.receive_weight_update(request) + + original_broadcast_done.wait(timeout=900) + print("[vLLM Engine] Original weights received") + + print(f"[vLLM Engine] Generating res_un_2 with prompt: '{prompt}'") + async for output in manager.engine.generate( + prompt, + sampling_params=sampling_params, + request_id="res_un_2", + ): + res_un_2_output = output + res_un_2 = res_un_2_output.outputs[0].text + print(f"[vLLM Engine] res_un_2: '{res_un_2}'") + + # Step 5: Verify + print("\n" + "="*60) + print("CROSS-VALIDATION RESULTS") + print("="*60) + print(f"res_un_1: '{res_un_1}'") + print(f"res_un_2: '{res_un_2}'") + print(f"res_mod_1: '{res_mod_1}'") + print(f"res_mod_2: '{res_mod_2}'") + print("="*60) + + # Check assertions + success = True + if res_un_1 == res_un_2: + print("✓ res_un_1 == res_un_2 (original weights produce same output)") + else: + print("✗ res_un_1 != res_un_2 (FAILED)") + success = False + + if res_mod_1 == res_mod_2: + print("✓ res_mod_1 == res_mod_2 (broadcast = load from disk)") + else: + print("✗ res_mod_1 != res_mod_2 (FAILED)") + success = False + + if not success: + sys.exit(1) + + print("\n✓ Cross-validation test PASSED") + + +async def test_back_and_forth( + model_name: str, + init_method: str, + actor_llm_idx: int, + world_size: int, + prompt: str, + max_tokens: int, + sync_dir: str, + tensor_parallel_size: int = 1, +): + """Back-and-forth test: switch between original and perturbed weights. + + Flow: + 1. Generate with original → res_or_1 + 2. Receive perturbed, generate → res_mod_1 + 3. Receive original, generate → res_or_2 + 4. Receive perturbed again, generate → res_mod_2 + 5. Verify: res_or_1 == res_or_2 and res_mod_1 == res_mod_2 + """ + from pipelinerl.vllm1 import EngineManager + from vllm import SamplingParams + from pathlib import Path + import argparse as ap + sys.path.insert(0, str(Path(__file__).parent)) + from sync_helper import SyncPoint, read_weight_update_request + + print("[vLLM Engine] Starting back-and-forth test") + + # Create sync points — actor-signaled names use per-actor suffix; + # completion signals (trainer→actors) stay unadorned and are shared. + sync_path = Path(sync_dir) + suffix = f"_actor_{actor_llm_idx}" + baseline_done = SyncPoint(sync_path, f"baseline_done{suffix}") + ready_for_perturbed1 = SyncPoint(sync_path, f"ready_for_perturbed1{suffix}") + perturbed1_done = SyncPoint(sync_path, "perturbed1_done") + ready_for_original = SyncPoint(sync_path, f"ready_for_original{suffix}") + original_done = SyncPoint(sync_path, "original_done") + ready_for_perturbed2 = SyncPoint(sync_path, f"ready_for_perturbed2{suffix}") + perturbed2_done = SyncPoint(sync_path, "perturbed2_done") + + sampling_params = SamplingParams( + temperature=0.0, + top_p=1.0, + max_tokens=max_tokens, + seed=42, + ) + + # Create engine args + args = ap.Namespace( + model=model_name, + tensor_parallel_size=tensor_parallel_size, + disable_log_stats=True, + enable_log_requests=False, + disable_weight_updates=False, + actor_llm_idx=actor_llm_idx, + weight_update_group_init_method=init_method, + weight_update_group_world_size=world_size, + ) + + print(f"[vLLM Engine] Creating engine with model: {model_name}") + async with EngineManager.create_engine(args) as manager: + # Step 1: Generate with original weights + print(f"[vLLM Engine] Step 1: Generating res_or_1") + async for output in manager.engine.generate( + prompt, sampling_params=sampling_params, request_id="res_or_1" + ): + res_or_1 = output.outputs[0].text + print(f"[vLLM Engine] res_or_1: '{res_or_1}'") + baseline_done.signal() + + # Step 2: Receive perturbed weights, generate + ready_for_perturbed1.signal() + import time + time.sleep(0.5) + request = read_weight_update_request(sync_path) + print(f"[vLLM Engine] Step 2: Receiving perturbed weights (1st time)") + await manager.receive_weight_update(request) + perturbed1_done.wait(timeout=900) + + print(f"[vLLM Engine] Generating res_mod_1") + async for output in manager.engine.generate( + prompt, sampling_params=sampling_params, request_id="res_mod_1" + ): + res_mod_1 = output.outputs[0].text + print(f"[vLLM Engine] res_mod_1: '{res_mod_1}'") + + # Step 3: Receive original weights, generate + ready_for_original.signal() + time.sleep(0.5) + request = read_weight_update_request(sync_path) + print(f"[vLLM Engine] Step 3: Receiving original weights") + await manager.receive_weight_update(request) + original_done.wait(timeout=900) + + print(f"[vLLM Engine] Generating res_or_2") + async for output in manager.engine.generate( + prompt, sampling_params=sampling_params, request_id="res_or_2" + ): + res_or_2 = output.outputs[0].text + print(f"[vLLM Engine] res_or_2: '{res_or_2}'") + + # Step 4: Receive perturbed weights again, generate + ready_for_perturbed2.signal() + time.sleep(0.5) + request = read_weight_update_request(sync_path) + print(f"[vLLM Engine] Step 4: Receiving perturbed weights (2nd time)") + await manager.receive_weight_update(request) + perturbed2_done.wait(timeout=900) + + print(f"[vLLM Engine] Generating res_mod_2") + async for output in manager.engine.generate( + prompt, sampling_params=sampling_params, request_id="res_mod_2" + ): + res_mod_2 = output.outputs[0].text + print(f"[vLLM Engine] res_mod_2: '{res_mod_2}'") + + # Step 5: Save results to per-actor file for multi-actor comparison + import json + results_file = sync_path / f"results_actor_{actor_llm_idx}.json" + actor_results = { + "res_or_1": res_or_1, + "res_mod_1": res_mod_1, + "res_or_2": res_or_2, + "res_mod_2": res_mod_2, + } + with open(results_file, "w") as f: + json.dump(actor_results, f, indent=2) + print(f"[vLLM Engine] Saved results for actor {actor_llm_idx} to {results_file}") + + # Step 6: Verify + print("\n" + "="*60) + print("BACK-AND-FORTH TEST RESULTS") + print("="*60) + print(f"res_or_1: '{res_or_1}'") + print(f"res_or_2: '{res_or_2}'") + print(f"res_mod_1: '{res_mod_1}'") + print(f"res_mod_2: '{res_mod_2}'") + print("="*60) + + # Check assertions + success = True + if res_or_1 == res_or_2: + print("✓ res_or_1 == res_or_2 (can switch back to original)") + else: + print("✗ res_or_1 != res_or_2 (FAILED)") + success = False + + if res_mod_1 == res_mod_2: + print("✓ res_mod_1 == res_mod_2 (perturbed weights consistent)") + else: + print("✗ res_mod_1 != res_mod_2 (FAILED)") + success = False + + if not success: + sys.exit(1) + + print("\n✓ Back-and-forth test PASSED") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="vLLM engine helper") + parser.add_argument("command", choices=["init", "weight_update", "cross_validation", "back_and_forth"]) + parser.add_argument("--model-name", required=True) + parser.add_argument("--init-method", required=True) + parser.add_argument("--actor-llm-idx", type=int, default=0) + parser.add_argument("--world-size", type=int, default=2) + # For weight_update command + parser.add_argument("--prompt", type=str, default="The capital of France is") + parser.add_argument("--max-tokens", type=int, default=50) + parser.add_argument("--sync-dir", type=str, help="Directory for sync files") + parser.add_argument("--expect-different", action="store_true", help="Expect outputs to be different (for perturbed weights)") + parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallel size for engine") + + args = parser.parse_args() + + try: + if args.command == "init": + asyncio.run(init_engine_and_process_group( + args.model_name, + args.init_method, + args.actor_llm_idx, + args.world_size, + )) + elif args.command == "weight_update": + if not args.sync_dir: + print("Error: --sync-dir required for weight_update command") + sys.exit(1) + asyncio.run(test_weight_update( + args.model_name, + args.init_method, + args.actor_llm_idx, + args.world_size, + args.prompt, + args.max_tokens, + args.sync_dir, + args.expect_different, + )) + elif args.command == "cross_validation": + if not args.sync_dir: + print("Error: --sync-dir required for cross_validation command") + sys.exit(1) + asyncio.run(test_cross_validation( + args.model_name, + args.init_method, + args.actor_llm_idx, + args.world_size, + args.prompt, + args.max_tokens, + args.sync_dir, + )) + elif args.command == "back_and_forth": + if not args.sync_dir: + print("Error: --sync-dir required for back_and_forth command") + sys.exit(1) + asyncio.run(test_back_and_forth( + args.model_name, + args.init_method, + args.actor_llm_idx, + args.world_size, + args.prompt, + args.max_tokens, + args.sync_dir, + tensor_parallel_size=args.tensor_parallel_size, + )) + except Exception as e: + print(f"[vLLM Engine] Error: {e}") + import traceback + traceback.print_exc() + sys.exit(1) diff --git a/tests/weight_update_utils.py b/tests/weight_update_utils.py new file mode 100644 index 00000000..3e36a1e6 --- /dev/null +++ b/tests/weight_update_utils.py @@ -0,0 +1,53 @@ +"""Utility functions for weight update testing.""" + +from typing import Dict +import torch +from pipelinerl.finetune_loop import WeightUpdateRequest, ParameterInfo + + +def dtype_to_string(dtype: torch.dtype) -> str: + """Convert torch dtype to string format expected by vLLM. + + Args: + dtype: PyTorch dtype + + Returns: + String representation (e.g., 'bfloat16', 'float32') + """ + dtype_str = str(dtype).replace("torch.", "") + return dtype_str + + +def create_weight_update_request_from_state_dict( + state_dict: Dict[str, torch.Tensor], + version: int = 0, +) -> WeightUpdateRequest: + """Create a WeightUpdateRequest from a model state dict. + + This helper function is useful for testing and for creating weight + update requests from saved model checkpoints. + + Args: + state_dict: Dictionary mapping parameter names to tensors + version: Version number for this weight update + + Returns: + WeightUpdateRequest object ready to be sent to workers + + Example: + >>> state_dict = torch.load('model.pt') + >>> request = create_weight_update_request_from_state_dict(state_dict, version=1) + >>> # Send request to vLLM server via HTTP endpoint + """ + parameters_info = [] + for name, tensor in state_dict.items(): + if isinstance(tensor, torch.Tensor): + parameters_info.append( + ParameterInfo( + name=name, + shape=list(tensor.shape), + dtype=dtype_to_string(tensor.dtype), + ) + ) + + return WeightUpdateRequest(version=version, parameters_info=parameters_info) From 0f0d37e6a29c6a672bfa42703aa52dc975b248d7 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 20 Jul 2026 13:38:49 -0400 Subject: [PATCH 2/4] Trim and clean up the Fast-LLM test suite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Apply the coarse + fine review findings: - Delete test_actor_error_handling.py — every test reimplemented the actor error-handling logic inline and asserted on the copy, so it could not catch a regression in the real actor. - Drop the generic-path topology matrix (TestWeightUpdateTP2, TestWeightUpdateMultiActor); equivalent multi-topology coverage lives on the Fast-LLM path in test_vllm1_fast_llm_broadcast.py. - Move the EngineManager import out of conftest top level into the factory fixture so the CPU-only unit tests collect without vLLM/torch. - Consolidate the duplicated stream_process_output / kill_process_tree helpers into server_weight_update_utils.py and drop the now-pointless injected stream_process_output_fn parameter. - Remove dead helpers (force_kill_process, check_pattern_detected), unused fixtures/params/imports, and the dead timeout parameter. - Minor typing and style fixes (builtin generics, strict zip, f-string and implicit-Optional cleanups). Co-Authored-By: Claude Opus 4.8 --- tests/conftest.py | 87 +----- tests/distributed_trainer_helper.py | 11 +- tests/fast_llm_trainer_helper.py | 6 +- tests/server_weight_update_utils.py | 127 ++++++--- tests/test_actor_error_handling.py | 290 -------------------- tests/test_vllm1_fast_llm_broadcast.py | 105 +------- tests/test_vllm1_integration.py | 360 +------------------------ tests/test_world_multinode.py | 8 +- tests/trainer_test_utils.py | 2 +- tests/weight_update_utils.py | 3 +- 10 files changed, 100 insertions(+), 899 deletions(-) delete mode 100644 tests/test_actor_error_handling.py diff --git a/tests/conftest.py b/tests/conftest.py index e33b8261..439422de 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,14 +1,8 @@ """Pytest configuration and fixtures for vllm1 tests.""" -import os import pytest -import torch import tempfile from pathlib import Path -import subprocess -import sys - -from pipelinerl.vllm1 import EngineManager @pytest.fixture(scope="session") @@ -17,42 +11,12 @@ def model_name(): return "Qwen/Qwen2.5-0.5B-Instruct" -@pytest.fixture(scope="session") -def sample_prompts(): - """Sample prompts for generation testing.""" - return [ - "Write a haiku about coding:", - "The capital of France is", - "In a galaxy far away,", - ] - - @pytest.fixture(scope="session") def simple_prompt(): """Single simple prompt for deterministic testing.""" return "The capital of France is" -@pytest.fixture(scope="session") -def num_gpus(): - """Number of GPUs available.""" - return torch.cuda.device_count() - - -@pytest.fixture(scope="session") -def require_2_gpus(num_gpus): - """Skip test if less than 2 GPUs available.""" - if num_gpus < 2: - pytest.skip("Test requires at least 2 GPUs") - - -@pytest.fixture(scope="session") -def require_gpu(): - """Skip test if no GPU available.""" - if not torch.cuda.is_available(): - pytest.skip("Test requires GPU") - - @pytest.fixture def temp_dir(): """Temporary directory for test files.""" @@ -82,21 +46,6 @@ def shared_distributed_init_method(shared_test_dir): return f"file://{shared_test_dir}/dist_init" -@pytest.fixture(scope="session") -def cache_dir(): - """Directory for caching downloaded models.""" - cache_path = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")) - cache_path.mkdir(parents=True, exist_ok=True) - return cache_path - - -@pytest.fixture -def vllm_server_port(): - """Port for vLLM server in tests.""" - # Use a high port to avoid conflicts - return 8765 - - @pytest.fixture def generation_config(): """Configuration for deterministic generation.""" @@ -108,33 +57,6 @@ def generation_config(): } -@pytest.fixture -def vllm_engine_factory_2gpu(model_name): - """Factory fixture that defaults to 2 GPUs. - - Usage: - async with vllm_engine_factory_2gpu() as manager: - # Uses 2 GPUs by default - # Access engine via manager.engine - ... - """ - def _factory(tensor_parallel_size: int = 2, **kwargs): - """Create engine with 2 GPUs by default.""" - import argparse - - args = argparse.Namespace( - model=model_name, - tensor_parallel_size=tensor_parallel_size, - disable_log_stats=True, - enable_log_requests=False, - **kwargs - ) - - return EngineManager.create_engine(args) - - return _factory - - @pytest.fixture def vllm_engine_factory(model_name): """Factory fixture for creating vLLM engines. @@ -149,11 +71,6 @@ def vllm_engine_factory(model_name): async with vllm_engine_factory(tensor_parallel_size=2) as manager: # use manager.engine with 2 GPUs ... - - Or if you need engine_config: - async with vllm_engine_factory() as manager: - # access manager.engine, manager.engine_config, manager.args - ... """ def _factory(tensor_parallel_size: int = 1, **kwargs): """Create engine context manager with test defaults. @@ -167,6 +84,8 @@ def _factory(tensor_parallel_size: int = 1, **kwargs): """ import argparse + from pipelinerl.vllm1 import EngineManager + # Create minimal args object with required attributes for AsyncEngineArgs.from_cli_args() args = argparse.Namespace( model=model_name, @@ -177,8 +96,6 @@ def _factory(tensor_parallel_size: int = 1, **kwargs): **kwargs ) - print("args: ", args) - return EngineManager.create_engine(args) return _factory diff --git a/tests/distributed_trainer_helper.py b/tests/distributed_trainer_helper.py index 7e50decb..05f2fd46 100755 --- a/tests/distributed_trainer_helper.py +++ b/tests/distributed_trainer_helper.py @@ -173,10 +173,9 @@ def save_model_to_dir(state_dict: dict, output_dir: str, model_name: str): def broadcast_weights( - init_method: str, model_name: str, perturb: bool = False, sync_dir: str = None + init_method: str, model_name: str, perturb: bool = False, sync_dir: str | None = None ): """Load model and broadcast weights to vLLM worker.""" - import torch import torch.distributed as dist from pathlib import Path @@ -239,17 +238,9 @@ def broadcast_weights( # Broadcast each weight with detailed logging logger.info(f"Starting broadcast of {len(params_to_broadcast)} parameters") for i, (name, tensor) in enumerate(params_to_broadcast.items()): - logger.debug(f"[{i+1}/{len(state_dict)}] Preparing to broadcast: {name}") - logger.debug( - f" - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}" - ) if tensor.device.type != "cuda": - logger.debug(f" - Moving {name} to CUDA") tensor = tensor.cuda(0) - logger.debug(f" - {name} now on device: {tensor.device}") - logger.debug(f" - Calling dist.broadcast for {name}...") dist.broadcast(tensor, src=0, group=process_group) - logger.debug(f" - Broadcast complete for {name}") if (i + 1) % 10 == 0: logger.info(f"Broadcasted {i+1}/{len(params_to_broadcast)} parameters") diff --git a/tests/fast_llm_trainer_helper.py b/tests/fast_llm_trainer_helper.py index 6ff173e2..a7974103 100644 --- a/tests/fast_llm_trainer_helper.py +++ b/tests/fast_llm_trainer_helper.py @@ -35,8 +35,6 @@ def timed_broadcast_fast_llm( redis_port: Redis port number world_size: Total NCCL world size (trainer rank 0 + all vLLM workers) """ - import torch - import torch.distributed as dist import time import redis import orjson @@ -99,7 +97,7 @@ def broadcast_weights_fast_llm(state_dict, step): # Send end signal _broadcast_object(None, process_group, src=0) - print(f"[Trainer] Sent end signal, broadcast complete") + print("[Trainer] Sent end signal, broadcast complete") # Broadcast 1: Perturbed weights print(f"[Trainer] Broadcasting {len(perturbed_state_dict)} perturbed parameters") @@ -154,7 +152,6 @@ def rapid_broadcast_cycles_fast_llm( 4. Slow broadcast: perturbed (5 s wait after) — end on text_B so the overall A→B→A→B pattern remains detectable """ - import torch.distributed as dist import time import redis as redis_lib import orjson @@ -186,7 +183,6 @@ def rapid_broadcast_cycles_fast_llm( def broadcast_weights(state_dict, label): nonlocal step - import torch event = {"type": "weights_ready", "step": step} r.xadd(stream_key, {payload_key: orjson.dumps(event)}) print(f"[Trainer] Sent weights_ready step={step} ({label})") diff --git a/tests/server_weight_update_utils.py b/tests/server_weight_update_utils.py index a4ade92a..f61016a2 100644 --- a/tests/server_weight_update_utils.py +++ b/tests/server_weight_update_utils.py @@ -7,6 +7,91 @@ import subprocess import sys import os +import signal + +try: + import psutil + + HAS_PSUTIL = True +except ImportError: + HAS_PSUTIL = False + print("WARNING: psutil not available, process tree cleanup will be limited") + + +def stream_process_output(proc, name): + """Start background threads to continuously stream process stdout/stderr. + + Args: + proc: subprocess.Popen object + name: Name for logging prefix (e.g., "vLLM Server", "Trainer") + + Returns: + Tuple of (stdout_thread, stderr_thread) + """ + import threading + + def read_stream(stream, prefix): + """Read from stream and print with prefix.""" + try: + for line in iter(stream.readline, ""): + if line: + print(f"{prefix} {line.rstrip()}", flush=True) + except Exception as e: + print(f"{prefix} [Stream read error: {e}]", flush=True) + + stdout_thread = threading.Thread( + target=read_stream, + args=(proc.stdout, f"[{name} OUT]"), + daemon=True, + ) + stderr_thread = threading.Thread( + target=read_stream, + args=(proc.stderr, f"[{name} ERR]"), + daemon=True, + ) + + stdout_thread.start() + stderr_thread.start() + + return stdout_thread, stderr_thread + + +def kill_process_tree(pid, sig=signal.SIGKILL): + """Kill a process and all its children/grandchildren. + + Args: + pid: Process ID to kill + sig: Signal to send (default SIGKILL) + """ + if not HAS_PSUTIL: + # Fallback: just kill the main process + try: + os.kill(pid, sig) + except ProcessLookupError: + pass + return + + try: + parent = psutil.Process(pid) + except psutil.NoSuchProcess: + return + + # Get all children recursively + children = parent.children(recursive=True) + + # Kill children first + for child in children: + try: + print(f"[Kill] Killing child process {child.pid}") + child.send_signal(sig) + except psutil.NoSuchProcess: + pass + + # Kill parent + try: + parent.send_signal(sig) + except psutil.NoSuchProcess: + pass async def wait_for_server_ready(server_url: str, server_proc, trainer_proc, timeout_seconds: int = 300): @@ -144,34 +229,6 @@ def _find_abab_pattern(phases, min_stable_gens=5): return phases[first_a], phases[first_b], phases[second_a], phases[second_b] -def check_pattern_detected(generations): - """Check whether the full A→B→A→B pattern is present in the generation history. - - This is a **post-hoc analysis helper** (e.g. for assertions after the - generation loop ends). It is intentionally *not* used as an early-stop - signal inside the generation loops. - - Why not early-stop? Any transition artifact text T that happens to appear - with several consecutive identical generations (possible when NCCL broadcasts - are slow) is indistinguishable from the real perturbed text B at generation - time. False positives would cut the loop short before the final stable B - phase accumulates. The generation loops instead rely on the trainer process - exiting (``trainer_proc.poll() is not None``) as their sole reliable - termination signal — the trainer exits within milliseconds of completing its - last broadcast, so no significant extra generation happens. - - Args: - generations: List of (timestamp, text) tuples - - Returns: - True if the A→B→A→B pattern is present - """ - if len(generations) < 4: - return False - phases = _build_phases(generations) - return _find_abab_pattern(phases) is not None - - async def run_generation_loop( server_url: str, model_name: str, @@ -423,8 +480,7 @@ def start_vllm_server( model_name: str, server_port: int, distributed_init_method: str, - stream_process_output_fn, - extra_args: list = None, + extra_args: list | None = None, gpu_ids: str = "0", actor_llm_idx: int = 0, world_size: int = 2, @@ -436,7 +492,6 @@ def start_vllm_server( model_name: Model to load server_port: Port to bind to distributed_init_method: Distributed initialization method - stream_process_output_fn: Function to stream process output extra_args: Additional CLI arguments (e.g., ["--weight-update-mode", "fast-llm"]) gpu_ids: CUDA_VISIBLE_DEVICES value (e.g., "0" or "0,1") actor_llm_idx: Actor index for this vLLM instance @@ -477,7 +532,7 @@ def start_vllm_server( ) print("[Main] Starting server output streaming...") - stdout_thread, stderr_thread = stream_process_output_fn(server_proc, f"vLLM Server (actor {actor_llm_idx})") + stdout_thread, stderr_thread = stream_process_output(server_proc, f"vLLM Server (actor {actor_llm_idx})") return server_proc, stdout_thread, stderr_thread @@ -503,7 +558,7 @@ async def wait_for_all_servers_ready( RuntimeError: If any process terminates unexpectedly TimeoutError: If any server doesn't become ready within timeout """ - for url, proc in zip(server_urls, server_procs): + for url, proc in zip(server_urls, server_procs, strict=True): await wait_for_server_ready(url, proc, trainer_proc, timeout_seconds) return True @@ -584,8 +639,7 @@ def start_trainer_process( distributed_init_method: str, model_name: str, server_urls: list, - stream_process_output_fn, - extra_args: list = None, + extra_args: list | None = None, gpu_id: str = "1", world_size: int = 2, command: str = "timed_broadcast_server_test", @@ -597,7 +651,6 @@ def start_trainer_process( distributed_init_method: Distributed initialization method model_name: Model name server_urls: List of server URLs (one per actor) - stream_process_output_fn: Function to stream process output extra_args: Additional CLI arguments (e.g., ["--n-cycles", "6"]) gpu_id: CUDA_VISIBLE_DEVICES value for the trainer GPU world_size: Total distributed world size @@ -648,6 +701,6 @@ def start_trainer_process( ) print("[Main] Starting trainer output streaming...") - stdout_thread, stderr_thread = stream_process_output_fn(trainer_proc, "Trainer") + stdout_thread, stderr_thread = stream_process_output(trainer_proc, "Trainer") return trainer_proc, stdout_thread, stderr_thread diff --git a/tests/test_actor_error_handling.py b/tests/test_actor_error_handling.py deleted file mode 100644 index 61adc5eb..00000000 --- a/tests/test_actor_error_handling.py +++ /dev/null @@ -1,290 +0,0 @@ -"""Test that actor rollout error handling doesn't crash the entire actor. - -Specifically tests that: -1. HTTP 4xx errors from vLLM (e.g., max_tokens too large) are handled gracefully -2. Groups where ALL rollouts fail are dropped (not submitted) -3. Groups where SOME rollouts fail submit only valid results -4. HTTP 5xx errors still propagate as fatal -""" - -import asyncio -import queue -from unittest.mock import MagicMock, AsyncMock, patch - -import aiohttp -import pytest -from omegaconf import OmegaConf - -from pipelinerl.rollouts import BaseMetrics, RolloutResult, TrainingText - - -# --------------------------------------------------------------------------- -# Helpers – lightweight stand-ins for heavy classes used by schedule_rollouts -# --------------------------------------------------------------------------- - -class FakeQueue: - """Minimal stand-in for SharedMemoryQueue (no shared memory needed).""" - - def __init__(self): - self._q = queue.Queue() - - def put(self, item, block=True, timeout=None): - self._q.put(item) - - def get(self, block=True, timeout=None): - return self._q.get(block=block, timeout=timeout) - - def qsize(self): - return self._q.qsize() - - def max_actual_entry_size(self): - return 0 - - def get_memory_size(self): - return 0 - - -class FakeTrainerState: - def __init__(self): - self.propagated_weight_version = 1 - self.samples_processed = 0 - - -def make_good_result() -> RolloutResult: - """A valid rollout result with one training sample.""" - return RolloutResult( - training_texts=[ - TrainingText( - text="prompt output", - n_predicted=6, - reward=1.0, - input_ids=[1, 2, 3], - labels=[-100, 2, 3], - finished=True, - prompt_tokens=5, - output_tokens=6, - ) - ], - metrics=BaseMetrics(reward=1.0, success=True, no_error=True, no_answer=False), - latency=0.5, - ) - - -def make_client_response_error(status: int, message: str = "Bad Request"): - """Create an aiohttp.ClientResponseError.""" - mock_req = MagicMock() - mock_req.url = "http://localhost:8080/v1/chat/completions" - return aiohttp.ClientResponseError( - request_info=mock_req, - history=(), - status=status, - message=message, - ) - - -# --------------------------------------------------------------------------- -# Core test: exercise rollout_and_maybe_produce_result + group completion -# --------------------------------------------------------------------------- - -@pytest.mark.asyncio -async def test_all_rollouts_fail_group_dropped(): - """When all rollouts in a group fail with 4xx, the group should be dropped.""" - attempts = 4 - problem_q = FakeQueue() - result_q = FakeQueue() - trainer_state = FakeTrainerState() - - # Put one problem in the queue - problem_q.put({"task": "What is 2+2?", "answer": "4"}) - - call_count = 0 - - async def failing_rollout_policy(cfg, llm, problem, session): - nonlocal call_count - call_count += 1 - raise make_client_response_error(400, "max_tokens too large") - - cfg = OmegaConf.create({ - "actor": { - "rollout_policy": "not_used", # we patch it - "llm_max_rollouts": 64, - }, - "finetune": { - "train_batch_size": 1000, - "gradient_accumulation_passes": 1, - "train_iters": 100, - "interrupt_train_steps": None, - }, - "debug": {}, - }) - - llms = [MagicMock()] # 1 LLM - - # We can't easily run schedule_rollouts (too many dependencies), - # so we directly test the inner logic by reimplementing the key parts. - # This mirrors rollout_and_maybe_produce_result + group completion. - - group_rollouts = {} - group_id = 0 - group_rollouts[group_id] = [] - finished_rollouts = 0 - warnings_logged = [] - - for rollout_index in range(attempts): - try: - rollout_result = await failing_rollout_policy(cfg, llms[0], {"task": "x"}, None) - except aiohttp.ClientResponseError as e: - if 400 <= e.status < 500: - warnings_logged.append(str(e.status)) - rollout_result = RolloutResult( - training_texts=[], - metrics=BaseMetrics(reward=0.0, success=False, no_error=False, no_answer=True), - latency=0.0, - ) - else: - raise - - rollout_result.model_version = 1 - rollout_result.group_id = f"test_{group_id}" - group_rollouts[group_id].append(rollout_result) - - # Now check group completion logic - assert len(group_rollouts[group_id]) == attempts - valid_results = [r for r in group_rollouts[group_id] if r.training_texts] - - # All failed → group should be dropped - assert len(valid_results) == 0, "Expected all results to be empty" - assert call_count == attempts - assert len(warnings_logged) == attempts - - # In real code: del group_rollouts[group_id], don't put in result_q - del group_rollouts[group_id] - assert result_q.qsize() == 0, "No group should be in the result queue" - - -@pytest.mark.asyncio -async def test_partial_failure_submits_valid_only(): - """When some rollouts fail but others succeed, submit only valid ones.""" - attempts = 4 - result_q = FakeQueue() - - call_count = 0 - - async def mixed_rollout_policy(cfg, llm, problem, session): - nonlocal call_count - call_count += 1 - # First 2 calls fail, last 2 succeed - if call_count <= 2: - raise make_client_response_error(400, "max_tokens too large") - return make_good_result() - - group_rollouts = {} - group_id = 0 - group_rollouts[group_id] = [] - - for rollout_index in range(attempts): - try: - rollout_result = await mixed_rollout_policy(None, None, {"task": "x"}, None) - except aiohttp.ClientResponseError as e: - if 400 <= e.status < 500: - rollout_result = RolloutResult( - training_texts=[], - metrics=BaseMetrics(reward=0.0, success=False, no_error=False, no_answer=True), - latency=0.0, - ) - else: - raise - - rollout_result.model_version = 1 - rollout_result.group_id = f"test_{group_id}" - group_rollouts[group_id].append(rollout_result) - - assert len(group_rollouts[group_id]) == attempts - - valid_results = [r for r in group_rollouts[group_id] if r.training_texts] - - # 2 failed, 2 succeeded - assert len(valid_results) == 2, f"Expected 2 valid results, got {len(valid_results)}" - - # In real code: result_queue.put(valid_results) - result_q.put(valid_results) - got = result_q.get(block=False) - assert len(got) == 2 - assert all(len(r.training_texts) > 0 for r in got) - - -@pytest.mark.asyncio -async def test_5xx_errors_still_propagate(): - """HTTP 5xx errors should NOT be caught — they indicate server failure.""" - - async def server_error_policy(cfg, llm, problem, session): - raise make_client_response_error(500, "Internal Server Error") - - with pytest.raises(aiohttp.ClientResponseError) as exc_info: - try: - await server_error_policy(None, None, {"task": "x"}, None) - except aiohttp.ClientResponseError as e: - if 400 <= e.status < 500: - pass # Would be caught in real code - else: - raise # 5xx re-raised - - assert exc_info.value.status == 500 - - -@pytest.mark.asyncio -async def test_all_succeed_normal_path(): - """When all rollouts succeed, the full group is submitted.""" - attempts = 4 - result_q = FakeQueue() - - async def good_policy(cfg, llm, problem, session): - return make_good_result() - - group_rollouts = {} - group_id = 0 - group_rollouts[group_id] = [] - - for rollout_index in range(attempts): - try: - rollout_result = await good_policy(None, None, {"task": "x"}, None) - except aiohttp.ClientResponseError as e: - if 400 <= e.status < 500: - rollout_result = RolloutResult( - training_texts=[], - metrics=BaseMetrics(reward=0.0, success=False, no_error=False, no_answer=True), - latency=0.0, - ) - else: - raise - - rollout_result.model_version = 1 - rollout_result.group_id = f"test_{group_id}" - group_rollouts[group_id].append(rollout_result) - - valid_results = [r for r in group_rollouts[group_id] if r.training_texts] - assert len(valid_results) == attempts, "All rollouts should be valid" - - result_q.put(valid_results) - got = result_q.get(block=False) - assert len(got) == attempts - - -@pytest.mark.asyncio -async def test_consumer_assertion_accepts_partial_group(): - """The consumer-side assertion should accept groups with fewer than `attempts` results.""" - attempts = 8 - # Simulate a partial group with 5 valid results - partial_count = 5 - - results = [make_good_result() for _ in range(partial_count)] - - # This mirrors the relaxed assertion in actor.py - assert isinstance(results, list) - assert isinstance(results[0], RolloutResult) - assert 0 < len(results) <= attempts, ( - f"Expected 1-{attempts} rollouts, got {len(results)}" - ) - - group_samples = sum(len(r.training_texts) for r in results) - assert group_samples == partial_count diff --git a/tests/test_vllm1_fast_llm_broadcast.py b/tests/test_vllm1_fast_llm_broadcast.py index f7cc410b..798e04a4 100644 --- a/tests/test_vllm1_fast_llm_broadcast.py +++ b/tests/test_vllm1_fast_llm_broadcast.py @@ -2,14 +2,9 @@ import asyncio import pytest -import tempfile from pathlib import Path -from typing import Dict, List import time -import os import subprocess -import sys -import signal # torch is needed at top level for pytest.mark.skipif decorators import torch @@ -25,92 +20,10 @@ analyze_and_verify_transitions, start_vllm_server, start_trainer_process, + stream_process_output, + kill_process_tree, ) -try: - import psutil - - HAS_PSUTIL = True -except ImportError: - HAS_PSUTIL = False - print("WARNING: psutil not available, process tree cleanup will be limited") - - -def stream_process_output(proc, name): - """Start background threads to continuously stream process stdout/stderr. - - Args: - proc: subprocess.Popen object - name: Name for logging prefix (e.g., "vLLM Server", "Trainer") - - Returns: - Tuple of (stdout_thread, stderr_thread) - """ - import threading - - def read_stream(stream, prefix): - """Read from stream and print with prefix.""" - try: - for line in iter(stream.readline, ""): - if line: - print(f"{prefix} {line.rstrip()}", flush=True) - except Exception as e: - print(f"{prefix} [Stream read error: {e}]", flush=True) - - stdout_thread = threading.Thread( - target=read_stream, - args=(proc.stdout, f"[{name} OUT]"), - daemon=True, - ) - stderr_thread = threading.Thread( - target=read_stream, - args=(proc.stderr, f"[{name} ERR]"), - daemon=True, - ) - - stdout_thread.start() - stderr_thread.start() - - return stdout_thread, stderr_thread - - -def kill_process_tree(pid, sig=signal.SIGKILL): - """Kill a process and all its children/grandchildren. - - Args: - pid: Process ID to kill - sig: Signal to send (default SIGKILL) - """ - if not HAS_PSUTIL: - # Fallback: just kill the main process - try: - os.kill(pid, sig) - except ProcessLookupError: - pass - return - - try: - parent = psutil.Process(pid) - except psutil.NoSuchProcess: - return - - # Get all children recursively - children = parent.children(recursive=True) - - # Kill children first - for child in children: - try: - print(f"[Kill] Killing child process {child.pid}") - child.send_signal(sig) - except psutil.NoSuchProcess: - pass - - # Kill parent - try: - parent.send_signal(sig) - except psutil.NoSuchProcess: - pass - @pytest.fixture def fast_llm_trainer_helper(): @@ -210,7 +123,6 @@ async def _run_fast_llm_server_test( vllm_server_configs, trainer_gpu, world_size, - timeout=2400, ): """Run Fast-LLM server weight-update pattern test with one or more vLLM servers. @@ -243,7 +155,6 @@ async def _run_fast_llm_server_test( model_name=model_name, server_port=port, distributed_init_method=init_method, - stream_process_output_fn=stream_process_output, extra_args=fast_llm_server_args, gpu_ids=cfg.get("gpu_ids", "0"), actor_llm_idx=cfg.get("actor_llm_idx", 0), @@ -259,7 +170,6 @@ async def _run_fast_llm_server_test( distributed_init_method=init_method, model_name=model_name, server_urls=server_urls, - stream_process_output_fn=stream_process_output, extra_args=[ "--redis-host", redis_host, "--redis-port", str(redis_port), @@ -326,7 +236,6 @@ async def test_server_fast_llm_broadcast_pattern( distributed_init_method, fast_llm_trainer_helper, redis_server, - temp_dir, ): """Server integration test: verify Fast-LLM weight broadcast pattern with HTTP API. @@ -354,7 +263,6 @@ async def test_server_fast_llm_broadcast_pattern( vllm_server_configs=[{"port": 8000, "gpu_ids": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}], trainer_gpu="1", world_size=2, - timeout=2400, ) @pytest.mark.timeout(2400) @@ -370,7 +278,6 @@ async def test_fast_llm_server_catch_transitions( distributed_init_method, fast_llm_trainer_helper, redis_server, - temp_dir, ): """Diagnostic test: catch garbage generations during Fast-LLM weight broadcasts. @@ -397,7 +304,6 @@ async def test_fast_llm_server_catch_transitions( model_name=model_name, server_port=8000, distributed_init_method=distributed_init_method, - stream_process_output_fn=stream_process_output, extra_args=[ "--weight-update-mode", "fast-llm", "--redis-host", redis_host, @@ -416,7 +322,6 @@ async def test_fast_llm_server_catch_transitions( distributed_init_method=distributed_init_method, model_name=model_name, server_urls=[server_url], - stream_process_output_fn=stream_process_output, extra_args=[ "--redis-host", redis_host, "--redis-port", str(redis_port), @@ -471,7 +376,6 @@ async def test_server_fast_llm_broadcast_pattern_tp2( distributed_init_method, fast_llm_trainer_helper, redis_server, - temp_dir, ): """Fast-LLM server test with TP=2: one server on GPUs 0+1, trainer on GPU 2. @@ -495,7 +399,6 @@ async def test_server_fast_llm_broadcast_pattern_tp2( vllm_server_configs=[{"port": 8001, "gpu_ids": "0,1", "actor_llm_idx": 0, "tensor_parallel_size": 2}], trainer_gpu="2", world_size=3, - timeout=2400, ) @@ -515,7 +418,6 @@ async def test_server_fast_llm_broadcast_pattern_2actors( distributed_init_method, fast_llm_trainer_helper, redis_server, - temp_dir, ): """Fast-LLM server test with 2 actors: servers on GPUs 0 and 1, trainer on GPU 2. @@ -542,7 +444,6 @@ async def test_server_fast_llm_broadcast_pattern_2actors( ], trainer_gpu="2", world_size=3, - timeout=2400, ) @pytest.mark.timeout(2400) @@ -558,7 +459,6 @@ async def test_server_fast_llm_broadcast_pattern_3actors( distributed_init_method, fast_llm_trainer_helper, redis_server, - temp_dir, ): """Fast-LLM server test with 3 actors: servers on GPUs 0/1/2, trainer on GPU 3. @@ -586,5 +486,4 @@ async def test_server_fast_llm_broadcast_pattern_3actors( ], trainer_gpu="3", world_size=4, - timeout=2400, ) diff --git a/tests/test_vllm1_integration.py b/tests/test_vllm1_integration.py index 2954bce5..3923e011 100644 --- a/tests/test_vllm1_integration.py +++ b/tests/test_vllm1_integration.py @@ -4,7 +4,6 @@ import pytest import tempfile from pathlib import Path -from typing import Dict, List import time import os import subprocess @@ -25,144 +24,9 @@ analyze_and_verify_transitions, start_vllm_server, start_trainer_process, + kill_process_tree, ) -try: - import psutil - HAS_PSUTIL = True -except ImportError: - HAS_PSUTIL = False - print("WARNING: psutil not available, process tree cleanup will be limited") - - -def stream_process_output(proc, name): - """Start background threads to continuously stream process stdout/stderr. - - Args: - proc: subprocess.Popen object - name: Name for logging prefix (e.g., "vLLM Server", "Trainer") - - Returns: - Tuple of (stdout_thread, stderr_thread) - """ - import threading - - def read_stream(stream, prefix): - """Read from stream and print with prefix.""" - try: - for line in iter(stream.readline, ''): - if line: - print(f"{prefix} {line.rstrip()}", flush=True) - except Exception as e: - print(f"{prefix} [Stream read error: {e}]", flush=True) - - stdout_thread = threading.Thread( - target=read_stream, - args=(proc.stdout, f"[{name} OUT]"), - daemon=True, - ) - stderr_thread = threading.Thread( - target=read_stream, - args=(proc.stderr, f"[{name} ERR]"), - daemon=True, - ) - - stdout_thread.start() - stderr_thread.start() - - return stdout_thread, stderr_thread - - -def kill_process_tree(pid, sig=signal.SIGKILL): - """Kill a process and all its children/grandchildren. - - Args: - pid: Process ID to kill - sig: Signal to send (default SIGKILL) - """ - if not HAS_PSUTIL: - # Fallback: just kill the main process - try: - os.kill(pid, sig) - except ProcessLookupError: - pass - return - - try: - parent = psutil.Process(pid) - except psutil.NoSuchProcess: - return - - # Get all children recursively - children = parent.children(recursive=True) - - # Kill children first - for child in children: - try: - print(f"[Kill] Killing child process {child.pid}") - child.send_signal(sig) - except psutil.NoSuchProcess: - pass - - # Kill parent - try: - parent.send_signal(sig) - except psutil.NoSuchProcess: - pass - - -def force_kill_process(proc, name): - """Forcefully kill a process tree and collect output. - - SIGKILL always kills the process. If communicate() hangs, it's the PIPES - that are stuck, not the process. We handle this with retries and timeouts. - - Returns: - Tuple of (stdout, stderr, returncode) - """ - # If already dead, try to get output - if proc.poll() is not None: - try: - stdout, stderr = proc.communicate(timeout=2) - return stdout, stderr, proc.returncode - except subprocess.TimeoutExpired: - print(f"[Kill] {name} already dead but pipes hung, closing...") - proc.stdout.close() if proc.stdout else None - proc.stderr.close() if proc.stderr else None - return "", "", proc.returncode - - # Kill entire process tree (including vLLM workers, trainer subprocesses, etc) - print(f"[Kill] Killing {name} process tree (PID {proc.pid})...") - kill_process_tree(proc.pid, signal.SIGKILL) - - # Wait for main process to actually die - try: - proc.wait(timeout=2) - print(f"[Kill] {name} process tree killed") - except subprocess.TimeoutExpired: - print(f"[Kill] WARNING: {name} didn't die after SIGKILL") - - # Try to read output from pipes (this is what usually hangs) - for attempt, timeout_val in enumerate([1, 2, 3], start=1): - try: - stdout, stderr = proc.communicate(timeout=timeout_val) - print(f"[Kill] {name} output collected (attempt {attempt})") - return stdout, stderr, proc.returncode - except subprocess.TimeoutExpired: - print(f"[Kill] {name} communicate() timed out (attempt {attempt})") - continue - - # Pipes are stuck - force close them - print(f"[Kill] {name} pipes stuck, force closing...") - try: - proc.stdout.close() if proc.stdout else None - proc.stderr.close() if proc.stderr else None - proc.stdin.close() if proc.stdin else None - except Exception as e: - print(f"[Kill] Error closing pipes: {e}") - - return "", "", proc.returncode if proc.returncode else -999 - async def wait_for_processes(processes_with_names, check_interval=0.5, timeout=60): """Wait for multiple subprocesses to complete, printing output in real-time. @@ -394,7 +258,6 @@ async def _run_server_weight_update_test( vllm_server_configs, trainer_gpu, world_size, - timeout=2400, ): """Run server weight-update pattern test with one or more vLLM servers. @@ -419,7 +282,6 @@ async def _run_server_weight_update_test( model_name=model_name, server_port=port, distributed_init_method=init_method, - stream_process_output_fn=stream_process_output, extra_args=None, gpu_ids=cfg.get("gpu_ids", "0"), actor_llm_idx=cfg.get("actor_llm_idx", 0), @@ -435,7 +297,6 @@ async def _run_server_weight_update_test( distributed_init_method=init_method, model_name=model_name, server_urls=server_urls, - stream_process_output_fn=stream_process_output, extra_args=None, gpu_id=trainer_gpu, world_size=world_size, @@ -963,7 +824,6 @@ async def test_server_weight_update_pattern( generation_config, distributed_init_method, distributed_trainer_helper, - temp_dir, ): """Server integration test: verify weight update pattern with HTTP API. @@ -983,7 +843,6 @@ async def test_server_weight_update_pattern( vllm_server_configs=[{"port": 8000, "gpu_ids": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}], trainer_gpu="1", world_size=2, - timeout=2400, ) @pytest.mark.timeout(2400) @@ -996,7 +855,6 @@ async def test_server_weight_update_catch_transitions( generation_config, distributed_init_method, distributed_trainer_helper, - temp_dir, ): """Diagnostic test: catch garbage generations produced during NCCL weight broadcasts. @@ -1022,7 +880,6 @@ async def test_server_weight_update_catch_transitions( model_name=model_name, server_port=8000, distributed_init_method=distributed_init_method, - stream_process_output_fn=stream_process_output, gpu_ids="0", actor_llm_idx=0, world_size=2, @@ -1036,7 +893,6 @@ async def test_server_weight_update_catch_transitions( distributed_init_method=distributed_init_method, model_name=model_name, server_urls=[server_url], - stream_process_output_fn=stream_process_output, extra_args=["--n-cycles", "6"], gpu_id="1", world_size=2, @@ -1070,217 +926,3 @@ async def test_server_weight_update_catch_transitions( kill_process_tree(server_proc.pid) if trainer_proc: kill_process_tree(trainer_proc.pid) - - -class TestWeightUpdateTP2: - """Test weight updates with tensor-parallel (TP=2) vLLM — needs 3 GPUs.""" - - @pytest.mark.timeout(2000) - @pytest.mark.asyncio - @pytest.mark.skipif(torch.cuda.device_count() < 3, reason="Requires at least 3 GPUs") - async def test_weight_update_back_and_forth_tp2( - self, - model_name, - simple_prompt, - generation_config, - distributed_init_method, - distributed_trainer_helper, - vllm_engine_helper, - temp_dir, - ): - """Back-and-forth test with TP=2: one vLLM instance on GPUs 0+1, trainer on GPU 2.""" - from .sync_helper import create_sync_dir - - print("\n" + "="*60) - print("Starting back-and-forth test (TP=2, 1 actor, 3 GPUs)") - print("="*60) - - sync_dir = create_sync_dir(temp_dir) - await _run_back_and_forth_engine_test( - model_name=model_name, - simple_prompt=simple_prompt, - generation_config=generation_config, - init_method=distributed_init_method, - distributed_trainer_helper=distributed_trainer_helper, - vllm_engine_helper=vllm_engine_helper, - sync_dir=sync_dir, - vllm_configs=[{"cuda_devices": "0,1", "actor_llm_idx": 0, "tensor_parallel_size": 2}], - trainer_gpu="2", - world_size=3, - timeout=1800, - ) - - @pytest.mark.timeout(2400) - @pytest.mark.asyncio - @pytest.mark.skipif(torch.cuda.device_count() < 3, reason="Requires at least 3 GPUs") - async def test_server_weight_update_pattern_tp2( - self, - model_name, - simple_prompt, - generation_config, - distributed_init_method, - distributed_trainer_helper, - temp_dir, - ): - """Server weight update test with TP=2: one server on GPUs 0+1, trainer on GPU 2.""" - print("\n" + "="*60) - print("Starting server weight update pattern test (TP=2, 1 actor, 3 GPUs)") - print("="*60) - - await _run_server_weight_update_test( - model_name=model_name, - simple_prompt=simple_prompt, - generation_config=generation_config, - init_method=distributed_init_method, - distributed_trainer_helper=distributed_trainer_helper, - vllm_server_configs=[{"port": 8001, "gpu_ids": "0,1", "actor_llm_idx": 0, "tensor_parallel_size": 2}], - trainer_gpu="2", - world_size=3, - timeout=2400, - ) - - -class TestWeightUpdateMultiActor: - """Test weight updates with multiple independent vLLM actors.""" - - @pytest.mark.timeout(2000) - @pytest.mark.asyncio - @pytest.mark.skipif(torch.cuda.device_count() < 3, reason="Requires at least 3 GPUs") - async def test_weight_update_back_and_forth_2actors( - self, - model_name, - simple_prompt, - generation_config, - distributed_init_method, - distributed_trainer_helper, - vllm_engine_helper, - temp_dir, - ): - """Back-and-forth test with 2 actors: vLLM on GPU 0 and GPU 1, trainer on GPU 2.""" - from .sync_helper import create_sync_dir - - print("\n" + "="*60) - print("Starting back-and-forth test (TP=1, 2 actors, 3 GPUs)") - print("="*60) - - sync_dir = create_sync_dir(temp_dir) - await _run_back_and_forth_engine_test( - model_name=model_name, - simple_prompt=simple_prompt, - generation_config=generation_config, - init_method=distributed_init_method, - distributed_trainer_helper=distributed_trainer_helper, - vllm_engine_helper=vllm_engine_helper, - sync_dir=sync_dir, - vllm_configs=[ - {"cuda_devices": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}, - {"cuda_devices": "1", "actor_llm_idx": 1, "tensor_parallel_size": 1}, - ], - trainer_gpu="2", - world_size=3, - timeout=1800, - ) - - @pytest.mark.timeout(2000) - @pytest.mark.asyncio - @pytest.mark.skipif(torch.cuda.device_count() < 4, reason="Requires at least 4 GPUs") - async def test_weight_update_back_and_forth_3actors( - self, - model_name, - simple_prompt, - generation_config, - distributed_init_method, - distributed_trainer_helper, - vllm_engine_helper, - temp_dir, - ): - """Back-and-forth test with 3 actors: vLLM on GPUs 0/1/2, trainer on GPU 3.""" - from .sync_helper import create_sync_dir - - print("\n" + "="*60) - print("Starting back-and-forth test (TP=1, 3 actors, 4 GPUs)") - print("="*60) - - sync_dir = create_sync_dir(temp_dir) - await _run_back_and_forth_engine_test( - model_name=model_name, - simple_prompt=simple_prompt, - generation_config=generation_config, - init_method=distributed_init_method, - distributed_trainer_helper=distributed_trainer_helper, - vllm_engine_helper=vllm_engine_helper, - sync_dir=sync_dir, - vllm_configs=[ - {"cuda_devices": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}, - {"cuda_devices": "1", "actor_llm_idx": 1, "tensor_parallel_size": 1}, - {"cuda_devices": "2", "actor_llm_idx": 2, "tensor_parallel_size": 1}, - ], - trainer_gpu="3", - world_size=4, - timeout=1800, - ) - - @pytest.mark.timeout(2400) - @pytest.mark.asyncio - @pytest.mark.skipif(torch.cuda.device_count() < 3, reason="Requires at least 3 GPUs") - async def test_server_weight_update_pattern_2actors( - self, - model_name, - simple_prompt, - generation_config, - distributed_init_method, - distributed_trainer_helper, - temp_dir, - ): - """Server weight update test with 2 actors: servers on GPUs 0 and 1, trainer on GPU 2.""" - print("\n" + "="*60) - print("Starting server weight update pattern test (TP=1, 2 actors, 3 GPUs)") - print("="*60) - - await _run_server_weight_update_test( - model_name=model_name, - simple_prompt=simple_prompt, - generation_config=generation_config, - init_method=distributed_init_method, - distributed_trainer_helper=distributed_trainer_helper, - vllm_server_configs=[ - {"port": 8000, "gpu_ids": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}, - {"port": 8001, "gpu_ids": "1", "actor_llm_idx": 1, "tensor_parallel_size": 1}, - ], - trainer_gpu="2", - world_size=3, - timeout=2400, - ) - - @pytest.mark.timeout(2400) - @pytest.mark.asyncio - @pytest.mark.skipif(torch.cuda.device_count() < 4, reason="Requires at least 4 GPUs") - async def test_server_weight_update_pattern_3actors( - self, - model_name, - simple_prompt, - generation_config, - distributed_init_method, - distributed_trainer_helper, - temp_dir, - ): - """Server weight update test with 3 actors: servers on GPUs 0/1/2, trainer on GPU 3.""" - print("\n" + "="*60) - print("Starting server weight update pattern test (TP=1, 3 actors, 4 GPUs)") - print("="*60) - - await _run_server_weight_update_test( - model_name=model_name, - simple_prompt=simple_prompt, - generation_config=generation_config, - init_method=distributed_init_method, - distributed_trainer_helper=distributed_trainer_helper, - vllm_server_configs=[ - {"port": 8000, "gpu_ids": "0", "actor_llm_idx": 0, "tensor_parallel_size": 1}, - {"port": 8001, "gpu_ids": "1", "actor_llm_idx": 1, "tensor_parallel_size": 1}, - {"port": 8002, "gpu_ids": "2", "actor_llm_idx": 2, "tensor_parallel_size": 1}, - ], - trainer_gpu="3", - world_size=4, - timeout=2400, - ) diff --git a/tests/test_world_multinode.py b/tests/test_world_multinode.py index 14d1474a..059a448f 100644 --- a/tests/test_world_multinode.py +++ b/tests/test_world_multinode.py @@ -1,10 +1,9 @@ """Tests for multi-node WorldMap topology and fast-llm torchrun command assembly.""" import os -import sys import tempfile from pathlib import Path -from unittest.mock import patch, MagicMock +from unittest.mock import patch import pytest from omegaconf import OmegaConf @@ -331,7 +330,6 @@ def _simulate_pod_ip_exchange(wm, pod_ips: dict): Sets dns_address_map to original DNS names, updates address_map and job URLs/hostnames to pod IPs — mirrors the real function's side-effects. """ - from pipelinerl.launch import _exchange_pod_ips as real_fn # noqa: F401 (not called) # Save DNS names first (matches the real implementation order) wm.dns_address_map = dict(wm.address_map) # Overwrite address_map with pod IPs @@ -691,8 +689,6 @@ def _capture_fast_llm_files(self, world_map, gpus=None): written_files = {} - real_open = open - def mock_popen(cmd, **kwargs): written_files["stdout"] = str(kwargs.get("stdout", {}).name if hasattr(kwargs.get("stdout"), "name") else "") written_files["stderr"] = str(kwargs.get("stderr", {}).name if hasattr(kwargs.get("stderr"), "name") else "") @@ -706,8 +702,6 @@ def mock_save_command(script_dir, cmd, suffix=""): captured_config = {} - real_omegaconf_save = None - with tempfile.TemporaryDirectory() as tmp: exp_dir = Path(tmp) with patch("pipelinerl.launch._popen", side_effect=mock_popen): diff --git a/tests/trainer_test_utils.py b/tests/trainer_test_utils.py index d6de57e5..c957ebd6 100644 --- a/tests/trainer_test_utils.py +++ b/tests/trainer_test_utils.py @@ -26,7 +26,7 @@ def _load_state_dict(model_name: str, device: str = "cuda:0") -> tuple: index_file = model_path / "model.safetensors.index.json" if index_file.exists(): - print(f"[Trainer] Found index file, loading sharded model") + print("[Trainer] Found index file, loading sharded model") with open(index_file) as f: index = json.load(f) weight_map = index["weight_map"] diff --git a/tests/weight_update_utils.py b/tests/weight_update_utils.py index 3e36a1e6..b42e5c59 100644 --- a/tests/weight_update_utils.py +++ b/tests/weight_update_utils.py @@ -1,6 +1,5 @@ """Utility functions for weight update testing.""" -from typing import Dict import torch from pipelinerl.finetune_loop import WeightUpdateRequest, ParameterInfo @@ -19,7 +18,7 @@ def dtype_to_string(dtype: torch.dtype) -> str: def create_weight_update_request_from_state_dict( - state_dict: Dict[str, torch.Tensor], + state_dict: dict[str, torch.Tensor], version: int = 0, ) -> WeightUpdateRequest: """Create a WeightUpdateRequest from a model state dict. From ed74e31d3cc98688c2b3e7cbf1aba1c8ae7bf178 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 20 Jul 2026 14:42:26 -0400 Subject: [PATCH 3/4] Second review pass: drop self-referential launch tests, assert trainer exit code - Delete TestPodIPExchange, TestHostfileCreation, TestRedisHostMultiNode: each asserted on a hand-copy of launch.py logic (_simulate_pod_ip_exchange, _compute_streams_host, the hostfile host-list) rather than the real code path, so a regression in launch.py would leave them green. The helper _simulate_pod_ip_exchange stays as setup for TestDeepSpeedCommand, which drives the real _run_finetune_deepspeed. - Assert trainer_proc.returncode in (0, None) after the wait loop in the server / broadcast pattern tests, so a trainer that crashes during final cleanup fails instead of passing on the already-captured pattern. - Fine cleanups: drop unused `import tempfile`, redundant local re-imports (argparse-as-ap, omegaconf, per-function pathlib/sys.path), unused thread handles and loop counter; hoist the broadcast import out of the inner closure; return the dtype string directly. Co-Authored-By: Claude Opus 4.8 --- tests/distributed_trainer_helper.py | 9 -- tests/fast_llm_trainer_helper.py | 7 +- tests/test_vllm1_fast_llm_broadcast.py | 8 +- tests/test_vllm1_integration.py | 5 +- tests/test_world_multinode.py | 159 ------------------------- tests/vllm_engine_helper.py | 28 ++--- tests/weight_update_utils.py | 3 +- 7 files changed, 26 insertions(+), 193 deletions(-) diff --git a/tests/distributed_trainer_helper.py b/tests/distributed_trainer_helper.py index 05f2fd46..61439907 100755 --- a/tests/distributed_trainer_helper.py +++ b/tests/distributed_trainer_helper.py @@ -34,8 +34,6 @@ def _wait_all_actors(sync_path, name: str, num_actors: int, timeout: float = 120 Each actor signals ``{name}_actor_{i}`` for i in range(num_actors). """ - from pathlib import Path - sys.path.insert(0, str(Path(__file__).parent)) from sync_helper import SyncPoint for i in range(num_actors): @@ -134,7 +132,6 @@ def save_model_to_dir(state_dict: dict, output_dir: str, model_name: str): output_dir: Directory to save model model_name: Original model name to copy config from """ - from pathlib import Path from safetensors.torch import save_file import shutil @@ -177,11 +174,9 @@ def broadcast_weights( ): """Load model and broadcast weights to vLLM worker.""" import torch.distributed as dist - from pathlib import Path # Setup sync points if provided if sync_dir: - sys.path.insert(0, str(Path(__file__).parent)) from sync_helper import SyncPoint, write_weight_update_request sync_path = Path(sync_dir) @@ -263,9 +258,7 @@ def broadcast_cross_validation( Also saves perturbed model to disk for vLLM to load. """ import torch.distributed as dist - from pathlib import Path - sys.path.insert(0, str(Path(__file__).parent)) from sync_helper import SyncPoint, write_weight_update_request from weight_update_utils import create_weight_update_request_from_state_dict @@ -361,9 +354,7 @@ def broadcast_back_and_forth( each broadcast, then sends a single shared completion signal. """ import torch.distributed as dist - from pathlib import Path - sys.path.insert(0, str(Path(__file__).parent)) from sync_helper import SyncPoint, write_weight_update_request from weight_update_utils import create_weight_update_request_from_state_dict diff --git a/tests/fast_llm_trainer_helper.py b/tests/fast_llm_trainer_helper.py index a7974103..4b6c1158 100644 --- a/tests/fast_llm_trainer_helper.py +++ b/tests/fast_llm_trainer_helper.py @@ -153,11 +153,12 @@ def rapid_broadcast_cycles_fast_llm( overall A→B→A→B pattern remains detectable """ import time - import redis as redis_lib + import redis import orjson from fast_llm.engine.distributed.config import DistributedBackend from fast_llm.engine.distributed.distributed import ProcessGroupPool + from fast_llm.core.distributed import broadcast as _broadcast, broadcast_object as _broadcast_object print(f"[Trainer] Initializing process group as rank 0 (world_size={world_size})") process_group = ProcessGroupPool( @@ -169,7 +170,7 @@ def rapid_broadcast_cycles_fast_llm( ).get_process_group(range(world_size), 0) print("[Trainer] Process group initialized") - r = redis_lib.Redis(host=redis_host, port=redis_port) + r = redis.Redis(host=redis_host, port=redis_port) stream_key = "fast_llm_events" payload_key = "event" @@ -188,8 +189,6 @@ def broadcast_weights(state_dict, label): print(f"[Trainer] Sent weights_ready step={step} ({label})") step += 1 - from fast_llm.core.distributed import broadcast as _broadcast, broadcast_object as _broadcast_object - for name, tensor in state_dict.items(): if tensor.device.type != "cuda": tensor = tensor.cuda(0) diff --git a/tests/test_vllm1_fast_llm_broadcast.py b/tests/test_vllm1_fast_llm_broadcast.py index 798e04a4..b5dee829 100644 --- a/tests/test_vllm1_fast_llm_broadcast.py +++ b/tests/test_vllm1_fast_llm_broadcast.py @@ -75,12 +75,12 @@ def find_free_port(): ) # Start streaming Redis output - redis_stdout_thread, redis_stderr_thread = stream_process_output(redis_proc, "Redis") + stream_process_output(redis_proc, "Redis") # Wait for Redis to be ready import redis r = redis.Redis(host=redis_host, port=redis_port) - for i in range(30): + for _ in range(30): try: r.ping() print(f"[Redis] Server ready on {redis_host}:{redis_port}") @@ -205,6 +205,8 @@ async def _run_fast_llm_server_test( break await asyncio.sleep(1) + assert trainer_proc.returncode in (0, None), f"Trainer exited with code {trainer_proc.returncode}" + if len(server_urls) == 1: analyze_and_verify_pattern(generations) else: @@ -349,6 +351,8 @@ async def test_fast_llm_server_catch_transitions( break await asyncio.sleep(1) + assert trainer_proc.returncode in (0, None), f"Trainer exited with code {trainer_proc.returncode}" + analyze_and_verify_transitions(generations, n_cycles=6) print("\n✓ Fast-LLM transition-capture test PASSED") diff --git a/tests/test_vllm1_integration.py b/tests/test_vllm1_integration.py index 3923e011..7ab515dd 100644 --- a/tests/test_vllm1_integration.py +++ b/tests/test_vllm1_integration.py @@ -2,7 +2,6 @@ import asyncio import pytest -import tempfile from pathlib import Path import time import os @@ -329,6 +328,8 @@ async def _run_server_weight_update_test( break await asyncio.sleep(1) + assert trainer_proc.returncode in (0, None), f"Trainer exited with code {trainer_proc.returncode}" + if len(server_urls) == 1: analyze_and_verify_pattern(generations) else: @@ -917,6 +918,8 @@ async def test_server_weight_update_catch_transitions( break await asyncio.sleep(1) + assert trainer_proc.returncode in (0, None), f"Trainer exited with code {trainer_proc.returncode}" + analyze_and_verify_transitions(generations, n_cycles=6) print("\n✓ Transition-capture test PASSED") diff --git a/tests/test_world_multinode.py b/tests/test_world_multinode.py index 059a448f..fb58e58a 100644 --- a/tests/test_world_multinode.py +++ b/tests/test_world_multinode.py @@ -346,54 +346,6 @@ def _simulate_pod_ip_exchange(wm, pod_ips: dict): job.url = job.url.replace(dns_name, pod_ip) -class TestPodIPExchange: - - def test_dns_address_map_holds_original_dns_names(self): - """After pod IP exchange, dns_address_map contains original DNS names, not pod IPs.""" - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) - wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") - - pod_ips = {0: "10.0.0.1", 1: "10.0.0.2"} - _simulate_pod_ip_exchange(wm, pod_ips) - - assert wm.dns_address_map[0] == "dns-abc123-0" - assert wm.dns_address_map[1] == "dns-abc123-1" - - def test_address_map_updated_to_pod_ips(self): - """After pod IP exchange, address_map and master_addr hold pod IPs.""" - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) - wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") - - pod_ips = {0: "10.0.0.1", 1: "10.0.0.2"} - _simulate_pod_ip_exchange(wm, pod_ips) - - assert wm.address_map[0] == "10.0.0.1" - assert wm.address_map[1] == "10.0.0.2" - assert wm.master_addr == "10.0.0.1" - - def test_job_urls_rewritten_to_pod_ips(self): - """After pod IP exchange, actor_llm job URLs use pod IPs, not DNS names.""" - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) - wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") - - # Verify that actor_llm jobs have DNS-based URLs before exchange - actor_urls_before = [job.url for job in wm.get_all_jobs() if job.kind == "actor_llm"] - assert all("dns-abc123-1" in u for u in actor_urls_before) - - pod_ips = {0: "10.0.0.1", 1: "10.0.0.2"} - _simulate_pod_ip_exchange(wm, pod_ips) - - actor_urls_after = [job.url for job in wm.get_all_jobs() if job.kind == "actor_llm"] - assert all("10.0.0.2" in u for u in actor_urls_after), f"Expected pod IP in URLs: {actor_urls_after}" - assert all("dns-abc123" not in u for u in actor_urls_after) - - def test_no_dns_address_map_without_exchange(self): - """Without pod IP exchange, dns_address_map is not set (no AttributeError).""" - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1) - wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") - assert not hasattr(wm, "dns_address_map") - - # --------------------------------------------------------------------------- # DeepSpeed command assembly: hostfile and inclusion filter use DNS names # --------------------------------------------------------------------------- @@ -476,116 +428,6 @@ def test_deepspeed_single_node_no_pod_ip_exchange(self): assert "--num_machines" not in cmd # single-node, no multi-machine flags -# --------------------------------------------------------------------------- -# Hostfile creation in main(): uses dns_address_map after pod IP exchange -# --------------------------------------------------------------------------- - -class TestHostfileCreation: - - def test_hostfile_uses_dns_names_after_pod_ip_exchange(self): - """The DeepSpeed hostfile written by main() uses DNS names even after pod IP exchange.""" - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) - wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") - - # Simulate pod IP exchange - _simulate_pod_ip_exchange(wm, {0: "10.0.0.1", 1: "10.0.0.2"}) - - dns_map = getattr(wm, "dns_address_map", wm.address_map) - hosts = [dns_map[i] for i in range(wm.world_size)] - - assert hosts[0] == "dns-abc123-0" - assert hosts[1] == "dns-abc123-1" - assert "10.0.0" not in hosts[0] - assert "10.0.0" not in hosts[1] - - def test_hostfile_uses_address_map_without_exchange(self): - """Without pod IP exchange, dns_address_map is absent — falls back to address_map.""" - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) - wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") - - dns_map = getattr(wm, "dns_address_map", wm.address_map) - hosts = [dns_map[i] for i in range(wm.world_size)] - - assert hosts[0] == "dns-abc123-0" - assert hosts[1] == "dns-abc123-1" - - -# --------------------------------------------------------------------------- -# Redis host in saved exp_config.yaml for multi-node (DeepSpeed + Redis) -# --------------------------------------------------------------------------- - -class TestRedisHostMultiNode: - - def _compute_streams_host(self, world_map, my_rank: int) -> str: - """Mirror the launch.py logic for cfg.streams.host selection.""" - if world_map.world_size > 1: - return world_map.master_addr - return "localhost" - - def test_single_node_redis_host_is_localhost(self): - """Single-node: Redis host is localhost regardless of pod IP exchange.""" - cfg = _make_cfg(actor_fraction=2, finetune_fraction=6, use_fast_llm=False) - with patch("torch.cuda.device_count", return_value=8): - with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): - with patch("pipelinerl.world.WorldMap._place_environments"): - from pipelinerl.world import WorldMap - wm = WorldMap(cfg, verbose=False) - - host = self._compute_streams_host(wm, my_rank=0) - assert host == "localhost" - - def test_multinode_rank0_redis_host_is_pod_ip(self): - """Multi-node rank 0: Redis host is pod IP (not localhost) after exchange. - - This ensures the saved exp_config.yaml has a reachable address for - DeepSpeed workers on other nodes. - """ - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) - wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") - _simulate_pod_ip_exchange(wm, {0: "10.0.0.1", 1: "10.0.0.2"}) - - host = self._compute_streams_host(wm, my_rank=0) - assert host == "10.0.0.1", "rank 0 should use pod IP so saved config is reachable cross-node" - assert host != "localhost" - - def test_multinode_rank1_redis_host_is_pod_ip(self): - """Multi-node rank 1: Redis host is pod IP of rank 0.""" - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) - wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0", rank=1) - _simulate_pod_ip_exchange(wm, {0: "10.0.0.1", 1: "10.0.0.2"}) - - host = self._compute_streams_host(wm, my_rank=1) - assert host == "10.0.0.1", "rank 1 should use rank 0's pod IP to reach Redis" - - def test_multinode_both_ranks_same_redis_host(self): - """Both ranks in a 2-node job resolve to the same Redis host (pod IP of rank 0).""" - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) - wm0 = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0", rank=0) - wm1 = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0", rank=1) - - _simulate_pod_ip_exchange(wm0, {0: "10.0.0.1", 1: "10.0.0.2"}) - _simulate_pod_ip_exchange(wm1, {0: "10.0.0.1", 1: "10.0.0.2"}) - - host0 = self._compute_streams_host(wm0, my_rank=0) - host1 = self._compute_streams_host(wm1, my_rank=1) - - assert host0 == host1 == "10.0.0.1" - - def test_multinode_without_pod_ip_exchange_uses_master_addr(self): - """Without pod IP exchange, multi-node uses master_addr (DNS name) for Redis. - - This is a fallback; the pod IP exchange should always run in practice - but the code must not crash without it. - """ - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) - wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") - - # No pod IP exchange — master_addr is still a DNS name - assert wm.master_addr == "dns-abc123-0" - host = self._compute_streams_host(wm, my_rank=0) - assert host == "dns-abc123-0" # DNS name (port filtering may apply, but code doesn't crash) - - # --------------------------------------------------------------------------- # DeepSpeed run_finetune.py path: must be absolute (not relative to CWD) # --------------------------------------------------------------------------- @@ -594,7 +436,6 @@ class TestDeepSpeedEntrypointPath: def _capture_ds_cmd(self, world_map): from pipelinerl.launch import _run_finetune_deepspeed - from omegaconf import OmegaConf cfg = OmegaConf.create({ "use_deepspeed": True, diff --git a/tests/vllm_engine_helper.py b/tests/vllm_engine_helper.py index 798743bc..47ecc882 100755 --- a/tests/vllm_engine_helper.py +++ b/tests/vllm_engine_helper.py @@ -23,12 +23,11 @@ async def init_engine_and_process_group( on context manager exit. """ from pipelinerl.vllm1 import EngineManager - import argparse as ap print("[vLLM Engine] Starting engine initialization") # Create args for engine with process group params - args = ap.Namespace( + args = argparse.Namespace( model=model_name, tensor_parallel_size=1, disable_log_stats=True, @@ -80,7 +79,6 @@ async def test_weight_update( from pipelinerl.vllm1 import EngineManager from vllm import SamplingParams from pathlib import Path - import argparse as ap # Import sync helper from same directory sys.path.insert(0, str(Path(__file__).parent)) from sync_helper import SyncPoint @@ -96,7 +94,7 @@ async def test_weight_update( broadcast_done = SyncPoint(sync_path, "broadcast_done") # Create args for engine with process group params - args = ap.Namespace( + args = argparse.Namespace( model=model_name, tensor_parallel_size=1, disable_log_stats=True, @@ -214,7 +212,6 @@ async def test_cross_validation( from pipelinerl.vllm1 import EngineManager from vllm import SamplingParams from pathlib import Path - import argparse as ap sys.path.insert(0, str(Path(__file__).parent)) from sync_helper import SyncPoint, read_weight_update_request @@ -240,7 +237,7 @@ async def test_cross_validation( ) # Step 1: Generate with original model - args = ap.Namespace( + args = argparse.Namespace( model=model_name, tensor_parallel_size=1, disable_log_stats=True, @@ -306,7 +303,7 @@ async def test_cross_validation( perturbed_model_path = (sync_path / "perturbed_model_path.txt").read_text().strip() print(f"[vLLM Engine] Step 3: Recreating engine with perturbed model from: {perturbed_model_path}") - args_perturbed = ap.Namespace( + args_perturbed = argparse.Namespace( model=perturbed_model_path, tensor_parallel_size=1, disable_log_stats=True, @@ -408,7 +405,6 @@ async def test_back_and_forth( from pipelinerl.vllm1 import EngineManager from vllm import SamplingParams from pathlib import Path - import argparse as ap sys.path.insert(0, str(Path(__file__).parent)) from sync_helper import SyncPoint, read_weight_update_request @@ -434,7 +430,7 @@ async def test_back_and_forth( ) # Create engine args - args = ap.Namespace( + args = argparse.Namespace( model=model_name, tensor_parallel_size=tensor_parallel_size, disable_log_stats=True, @@ -448,7 +444,7 @@ async def test_back_and_forth( print(f"[vLLM Engine] Creating engine with model: {model_name}") async with EngineManager.create_engine(args) as manager: # Step 1: Generate with original weights - print(f"[vLLM Engine] Step 1: Generating res_or_1") + print("[vLLM Engine] Step 1: Generating res_or_1") async for output in manager.engine.generate( prompt, sampling_params=sampling_params, request_id="res_or_1" ): @@ -461,11 +457,11 @@ async def test_back_and_forth( import time time.sleep(0.5) request = read_weight_update_request(sync_path) - print(f"[vLLM Engine] Step 2: Receiving perturbed weights (1st time)") + print("[vLLM Engine] Step 2: Receiving perturbed weights (1st time)") await manager.receive_weight_update(request) perturbed1_done.wait(timeout=900) - print(f"[vLLM Engine] Generating res_mod_1") + print("[vLLM Engine] Generating res_mod_1") async for output in manager.engine.generate( prompt, sampling_params=sampling_params, request_id="res_mod_1" ): @@ -476,11 +472,11 @@ async def test_back_and_forth( ready_for_original.signal() time.sleep(0.5) request = read_weight_update_request(sync_path) - print(f"[vLLM Engine] Step 3: Receiving original weights") + print("[vLLM Engine] Step 3: Receiving original weights") await manager.receive_weight_update(request) original_done.wait(timeout=900) - print(f"[vLLM Engine] Generating res_or_2") + print("[vLLM Engine] Generating res_or_2") async for output in manager.engine.generate( prompt, sampling_params=sampling_params, request_id="res_or_2" ): @@ -491,11 +487,11 @@ async def test_back_and_forth( ready_for_perturbed2.signal() time.sleep(0.5) request = read_weight_update_request(sync_path) - print(f"[vLLM Engine] Step 4: Receiving perturbed weights (2nd time)") + print("[vLLM Engine] Step 4: Receiving perturbed weights (2nd time)") await manager.receive_weight_update(request) perturbed2_done.wait(timeout=900) - print(f"[vLLM Engine] Generating res_mod_2") + print("[vLLM Engine] Generating res_mod_2") async for output in manager.engine.generate( prompt, sampling_params=sampling_params, request_id="res_mod_2" ): diff --git a/tests/weight_update_utils.py b/tests/weight_update_utils.py index b42e5c59..b4613946 100644 --- a/tests/weight_update_utils.py +++ b/tests/weight_update_utils.py @@ -13,8 +13,7 @@ def dtype_to_string(dtype: torch.dtype) -> str: Returns: String representation (e.g., 'bfloat16', 'float32') """ - dtype_str = str(dtype).replace("torch.", "") - return dtype_str + return str(dtype).replace("torch.", "") def create_weight_update_request_from_state_dict( From 59eeb6cfd6289ac549522435d79a3c764ca7ffa9 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 20 Jul 2026 15:42:53 -0400 Subject: [PATCH 4/4] Fix pre-existing test_world_multinode failures surfaced by the first suite run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Running the suite for the first time exposed 11 failures in test_world_multinode.py, all test-setup gaps (not product bugs): - 8× `Missing key wandb_name`: the fast-llm test cfgs (_capture_cmd, _capture_fast_llm_files) omitted wandb.wandb_name, which _run_finetune_fast_llm reads. Add it (None) so the code reaches the torchrun/naming assertions. - 3× `WorldMap has no dns_address_map`: multinode finetune always runs after _exchange_pod_ips (which sets dns_address_map), but _make_world_map produced a pre-exchange map. Set dns_address_map for world_size > 1 in the helper to mirror production. Delete test_deepspeed_multinode_uses_dns_names_without_exchange: its "without exchange" premise is unreachable for multinode, and the after-exchange sibling already covers the DNS-in-filter assertion (plus the no-pod-IP-leak check). Verified on GPU: test_world_multinode + test_launch_process_monitoring now 29 passed. Co-Authored-By: Claude Opus 4.8 --- tests/test_world_multinode.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/tests/test_world_multinode.py b/tests/test_world_multinode.py index fb58e58a..088c8cac 100644 --- a/tests/test_world_multinode.py +++ b/tests/test_world_multinode.py @@ -53,7 +53,11 @@ def _make_world_map(cfg, world_size, rank=0, master_addr="dns-test-0"): # patch it out to avoid AttributeError. with patch("pipelinerl.world.WorldMap._place_environments"): with patch("pipelinerl.utils.collect_environment_specs", return_value=[]): - return WorldMap(cfg, verbose=False) + world_map = WorldMap(cfg, verbose=False) + # Multi-node finetune always runs after _exchange_pod_ips, which populates dns_address_map; mirror that. + if world_size > 1: + world_map.dns_address_map = dict(world_map.address_map) + return world_map # --------------------------------------------------------------------------- @@ -191,6 +195,7 @@ def _capture_cmd(self, world_map, cfg_extra=None): "wandb_entity_name": "test", "wandb_project_name": "test", "wandb_group": "test", + "wandb_name": None, }, "fast_llm": { "training": { @@ -386,16 +391,6 @@ def mock_popen(cmd, **kwargs): return captured_cmd - def test_deepspeed_multinode_uses_dns_names_without_exchange(self): - """DeepSpeed 2-node without pod IP exchange: inclusion filter uses DNS names.""" - cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) - wm = _make_world_map(cfg, world_size=2, master_addr="dns-abc123-0") - - cmd = self._capture_ds_cmd(wm) - # The deepspeed_inclusion_filter should contain the DNS hostname for the finetune node - filter_arg = next((c for c in cmd if "dns-abc123" in c), None) - assert filter_arg is not None, f"Expected DNS name in cmd, got: {cmd}" - def test_deepspeed_multinode_after_pod_ip_exchange_uses_dns_names(self): """After pod IP exchange, DeepSpeed inclusion filter still uses DNS names (not pod IPs).""" cfg = _make_cfg(actor_fraction=1, finetune_fraction=1, use_fast_llm=False) @@ -510,6 +505,7 @@ def _capture_fast_llm_files(self, world_map, gpus=None): "wandb_entity_name": "test", "wandb_project_name": "test", "wandb_group": "test", + "wandb_name": None, }, "fast_llm": { "training": {