diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 7f2e94bf1..5218123d2 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -186,7 +186,7 @@ jobs: extra_deps: "" env_vars: "SHOW_ROUTER_LOGS=1" reruns: "--reruns 3 --reruns-delay 2" - - name: router-embeddings + - name: e2e timeout: 45 test_dirs: "e2e_test/router e2e_test/embeddings" extra_deps: "" diff --git a/sgl-model-gateway/e2e_test/conftest.py b/sgl-model-gateway/e2e_test/conftest.py index 96eac4db5..8a06a831b 100644 --- a/sgl-model-gateway/e2e_test/conftest.py +++ b/sgl-model-gateway/e2e_test/conftest.py @@ -175,15 +175,25 @@ from infra import ( PARAM_MODEL, PARAM_SETUP_BACKEND, ConnectionMode, + WorkerIdentity, + WorkerType, ) # --------------------------------------------------------------------------- -# Test collection: scan for required backends +# Test collection: scan for required workers # --------------------------------------------------------------------------- -# Global storage for scanned requirements -_scanned_backends: set[str] = set() # {"grpc", "http", "openai", ...} -_scanned_models: set[str] = set() # Models needed by tests +# Track max worker counts: (model_id, mode, worker_type) -> max_count +# This unified approach handles regular, prefill, and decode workers the same way +_worker_counts: dict[tuple[str, ConnectionMode, WorkerType], int] = {} + +# Track first-seen order to preserve test collection order +_first_seen_order: list[tuple[str, ConnectionMode, WorkerType]] = [] + +# Track max GPU requirement for any single test (for validation) +_max_test_gpu_requirement: int = 0 +_max_test_name: str = "" + _needs_default_model: bool = False # True if any e2e test lacks explicit model marker @@ -192,93 +202,259 @@ def pytest_collection_modifyitems( config: pytest.Config, items: list[pytest.Item], ) -> None: - """Scan collected tests to determine required backends and models. + """Scan collected tests to determine required workers. This runs after test collection but before tests execute. - It extracts backend requirements from @pytest.mark.parametrize markers. + It extracts worker requirements from markers in test collection order, + tracking the max count needed for each (model, mode, worker_type) combination. + + Also tracks the max GPU requirement for any single test for validation. """ - global _scanned_backends, _scanned_models, _needs_default_model + global _worker_counts, _first_seen_order, _needs_default_model + global _max_test_gpu_requirement, _max_test_name + + from infra import MODEL_SPECS + + def track_worker( + model_id: str, mode: ConnectionMode, worker_type: WorkerType, count: int + ) -> None: + """Track a worker requirement, updating max count if needed.""" + key = (model_id, mode, worker_type) + if key not in _worker_counts: + _first_seen_order.append(key) + _worker_counts[key] = count + else: + _worker_counts[key] = max(_worker_counts[key], count) + + def calculate_test_gpus( + model_id: str, prefill: int, decode: int, regular: int + ) -> int: + """Calculate GPU requirement for a single test.""" + if model_id not in MODEL_SPECS: + return 0 + tp = MODEL_SPECS[model_id].get("tp", 1) + return tp * (prefill + decode + regular) for item in items: - # Track if this test has an explicit model marker - has_model_marker = False + # Extract model from marker or use default + model_marker = item.get_closest_marker(PARAM_MODEL) + model_id = model_marker.args[0] if model_marker and model_marker.args else None - # Scan parametrize markers for setup_backend + # Check parametrize for model + if model_id is None: + for marker in item.iter_markers("parametrize"): + if marker.args and len(marker.args) >= 2: + param_name = marker.args[0] + if param_name == PARAM_MODEL or PARAM_MODEL in param_name: + param_values = marker.args[1] + if isinstance(param_values, (list, tuple)) and param_values: + model_id = param_values[0] # First model in parametrize + break + + # Extract backends from parametrize + backends: list[str] = [] for marker in item.iter_markers("parametrize"): if marker.args and len(marker.args) >= 2: param_name = marker.args[0] param_values = marker.args[1] - if param_name == PARAM_SETUP_BACKEND: - # Extract backend names from parametrize values if isinstance(param_values, (list, tuple)): - _scanned_backends.update(param_values) + backends.extend(param_values) - elif param_name == PARAM_MODEL or PARAM_MODEL in param_name: - # Extract model names from parametrize - if isinstance(param_values, (list, tuple)): - _scanned_models.update(param_values) - has_model_marker = True + # Check for workers marker (@pytest.mark.workers(...)) + workers_marker = item.get_closest_marker("workers") + prefill_count = 0 + decode_count = 0 + regular_count = 1 # Default to 1 regular worker + if workers_marker: + prefill_count = workers_marker.kwargs.get("prefill") or 0 + decode_count = workers_marker.kwargs.get("decode") or 0 + regular_count = workers_marker.kwargs.get("count") or 1 - # Check for @pytest.mark.model("name") markers - model_marker = item.get_closest_marker(PARAM_MODEL) - if model_marker and model_marker.args: - model_name = model_marker.args[0] - _scanned_models.add(model_name) - has_model_marker = True - - # Check if this is an e2e test without an explicit model marker - # Such tests need the DEFAULT_MODEL - if not has_model_marker and item.get_closest_marker("e2e"): + # Track if this test needs default model + is_e2e = item.get_closest_marker("e2e") is not None + if model_id is None and is_e2e: _needs_default_model = True + model_id = DEFAULT_MODEL - logger.info( - "Scanned test requirements - backends: %s, models: %s, needs default: %s", - _scanned_backends or {"(none)"}, - _scanned_models or {"(none)"}, - _needs_default_model, - ) + # Track worker requirements and calculate this test's GPU requirement + test_gpus = 0 + if model_id and backends: + for backend in backends: + # "pd" backend means PD workers + if backend == "pd": + mode = ConnectionMode.HTTP # PD uses HTTP mode + # Default to 1 prefill + 1 decode if not specified + p_count = prefill_count if prefill_count > 0 else 1 + d_count = decode_count if decode_count > 0 else 1 + track_worker(model_id, mode, WorkerType.PREFILL, p_count) + track_worker(model_id, mode, WorkerType.DECODE, d_count) + test_gpus = max( + test_gpus, calculate_test_gpus(model_id, p_count, d_count, 0) + ) + else: + try: + mode = ConnectionMode(backend) + except ValueError: + # Cloud backend (openai, xai, etc.) - skip + continue + + # Check if this backend also has PD workers + if prefill_count > 0 or decode_count > 0: + track_worker(model_id, mode, WorkerType.PREFILL, prefill_count) + track_worker(model_id, mode, WorkerType.DECODE, decode_count) + test_gpus = max( + test_gpus, + calculate_test_gpus( + model_id, prefill_count, decode_count, 0 + ), + ) + else: + # Regular worker + track_worker(model_id, mode, WorkerType.REGULAR, regular_count) + test_gpus = max( + test_gpus, + calculate_test_gpus(model_id, 0, 0, regular_count), + ) + + elif model_id and is_e2e: + # E2E test without explicit backend - will use HTTP by default + track_worker(model_id, ConnectionMode.HTTP, WorkerType.REGULAR, 1) + test_gpus = calculate_test_gpus(model_id, 0, 0, 1) + + # Track max GPU requirement across all tests + if test_gpus > _max_test_gpu_requirement: + _max_test_gpu_requirement = test_gpus + _max_test_name = item.nodeid + + # Log results + if _worker_counts: + summary = [] + for key in _first_seen_order: + model_id, mode, worker_type = key + count = _worker_counts[key] + if worker_type == WorkerType.REGULAR: + summary.append(f"{model_id}:{mode.value}x{count}") + else: + summary.append(f"{model_id}:{mode.value}:{worker_type.value}x{count}") + logger.info("Scanned worker requirements (in test order): %s", summary) + logger.info( + "Max GPU requirement for single test: %d (%s)", + _max_test_gpu_requirement, + _max_test_name, + ) + else: + logger.info("Scanned worker requirements: (none)") -def get_pool_requirements() -> list[tuple[str, ConnectionMode]]: +def get_pool_requirements() -> list[WorkerIdentity]: """Build pool requirements from scanned test markers. Returns: - List of (model_id, ConnectionMode) tuples to try to pre-launch. - Models that don't fit will be launched on-demand. + List of WorkerIdentity objects to pre-launch. + Each WorkerIdentity has (model_id, mode, worker_type, index). + Requirements are ordered by first appearance in test collection order, + so workers needed by earlier tests are launched first. + + Note: + If a model's first test needs PD workers (prefill/decode), we skip + pre-launching regular workers for that model (they'd be evicted + immediately when PD workers are launched). """ - models = set(_scanned_models) + # Track which models have PD workers as their first requirement + # These models shouldn't have regular workers pre-launched + models_with_pd_first: set[str] = set() + first_worker_type_per_model: dict[str, WorkerType] = {} - # Add DEFAULT_MODEL if any e2e test lacks an explicit model marker, - # or if no models were specified at all - if _needs_default_model or not models: - models.add(DEFAULT_MODEL) + for model_id, mode, worker_type in _first_seen_order: + if model_id not in first_worker_type_per_model: + first_worker_type_per_model[model_id] = worker_type + if worker_type in (WorkerType.PREFILL, WorkerType.DECODE): + models_with_pd_first.add(model_id) + logger.info( + "Model %s has PD test first - skipping regular worker pre-launch", + model_id, + ) - # Convert scanned string backends to ConnectionMode enums - # Filter to local backends only (grpc, http) - cloud backends don't need workers - local_modes: set[ConnectionMode] = set() - for backend in _scanned_backends: - try: - mode = ConnectionMode(backend) - if mode in LOCAL_MODES: - local_modes.add(mode) - except ValueError: - # Not a ConnectionMode (e.g., "openai", "xai", "pd") - skip - pass + # Generate individual WorkerIdentity objects in first-seen order + requirements: list[WorkerIdentity] = [] + for model_id, mode, worker_type in _first_seen_order: + # Skip regular workers for models that have PD first + if model_id in models_with_pd_first and worker_type == WorkerType.REGULAR: + continue - # Default to HTTP if no local backends specified - if not local_modes: - local_modes = {ConnectionMode.HTTP} + count = _worker_counts.get((model_id, mode, worker_type), 1) + for i in range(count): + requirements.append(WorkerIdentity(model_id, mode, worker_type, i)) - # Build requirements: each model needs each mode - requirements: list[tuple[str, ConnectionMode]] = [] - for model in models: - for mode in local_modes: - requirements.append((model, mode)) + # Add default if no requirements + if not requirements: + requirements.append(WorkerIdentity(DEFAULT_MODEL, ConnectionMode.HTTP)) return requirements +def validate_gpu_requirements() -> tuple[int, int]: + """Check if there are enough GPUs for any single test. + + Returns: + Tuple of (max_required_gpus, available_gpus). + + Note: + We check the max requirement for any single test, not the sum. + Workers can be evicted between tests, so we only need enough GPUs + for the most demanding test. + """ + # Count available GPUs + available_gpus = 0 + try: + import torch + + if torch.cuda.is_available(): + available_gpus = torch.cuda.device_count() + except ImportError: + pass + + return _max_test_gpu_requirement, available_gpus + + +def pytest_collection_finish(session: pytest.Session) -> None: + """Validate GPU requirements after test collection. + + This runs after all tests are collected but before any tests execute. + Fails fast if any single test requires more GPUs than available. + """ + if not _worker_counts: + return + + # Skip validation if model pool is disabled + if os.environ.get(ENV_SKIP_MODEL_POOL, "").lower() in ("1", "true", "yes"): + return + + max_required, available_gpus = validate_gpu_requirements() + + if max_required > available_gpus: + raise pytest.UsageError( + f"\n{'='*60}\n" + f"GPU REQUIREMENTS EXCEEDED\n" + f"{'='*60}\n" + f"Test '{_max_test_name}' requires {max_required} GPUs\n" + f"Available: {available_gpus} GPUs\n" + f"\nOptions:\n" + f" 1. Run tests that fit: pytest -k 'not {_max_test_name.split('::')[0]}'\n" + f" 2. Reduce workers: @pytest.mark.workers(prefill=1, decode=1)\n" + f" 3. Skip GPU tests: SKIP_MODEL_POOL=1 pytest\n" + f"{'='*60}" + ) + + logger.info( + "GPU validation passed: max %d required (by %s), %d available", + max_required, + _max_test_name, + available_gpus, + ) + + # --------------------------------------------------------------------------- # Custom pytest markers # --------------------------------------------------------------------------- @@ -333,13 +509,15 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool": routers (~1-2s) pointing to these workers. Startup behavior: - - Scans test markers to determine required (model, mode) combinations - - Launches workers sequentially, but they boot up concurrently + - Scans test markers to determine required workers (model, mode, type, count) + - Launches workers in test collection order - Waits for all workers to become healthy before returning Test requirements are auto-detected from: - - @pytest.mark.parametrize("setup_backend", ["grpc", "http"]) + - @pytest.mark.parametrize("setup_backend", ["grpc", "http", "pd"]) - @pytest.mark.model("model-name") + - @pytest.mark.workers(count=N) for regular workers + - @pytest.mark.workers(prefill=N, decode=N) for PD workers Environment variable overrides: - E2E_MODELS: Comma-separated model IDs (e.g., "llama-8b,qwen-7b") @@ -388,15 +566,20 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool": if not backend_modes: backend_modes = {ConnectionMode.HTTP} - requirements = [(m, b) for m in models for b in backend_modes] - logger.info("Using env var requirements: %s", requirements) + # Create WorkerIdentity objects (regular workers only from env vars) + requirements = [ + WorkerIdentity(m, b, WorkerType.REGULAR, 0) + for m in models + for b in backend_modes + ] + logger.info("Using env var requirements: %s", [str(r) for r in requirements]) else: # Use scanned requirements from test markers requirements = get_pool_requirements() - logger.info("Using scanned requirements: %s", requirements) + logger.info("Using scanned requirements: %s", [str(r) for r in requirements]) # Filter to valid models - requirements = [(m, b) for m, b in requirements if m in MODEL_SPECS] + requirements = [r for r in requirements if r.model_id in MODEL_SPECS] if not requirements: logger.warning("No valid requirements, model pool will be empty") @@ -408,7 +591,10 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool": _model_pool = ModelPool(allocator) startup_timeout = int(os.environ.get(ENV_STARTUP_TIMEOUT, "300")) - _model_pool.startup(requirements=requirements, startup_timeout=startup_timeout) + _model_pool.startup( + requirements=requirements, + startup_timeout=startup_timeout, + ) # Log final GPU allocation summary logger.info(_model_pool.allocator.summary()) @@ -624,15 +810,16 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): f"({num_prefill} prefill + {num_decode} decode), found {gpu_count}" ) - # Try to use pre-launched PD workers, or launch new ones if needed + # Try to use pre-launched PD workers, or launch additional ones if needed existing_prefills = model_pool.get_workers_by_type(model_id, WorkerType.PREFILL) existing_decodes = model_pool.get_workers_by_type(model_id, WorkerType.DECODE) - if ( - len(existing_prefills) >= num_prefill - and len(existing_decodes) >= num_decode - ): - # Use pre-launched workers + # Calculate how many more we need (if any) + missing_prefill = max(0, num_prefill - len(existing_prefills)) + missing_decode = max(0, num_decode - len(existing_decodes)) + + if missing_prefill == 0 and missing_decode == 0: + # Use pre-launched workers (we have enough) prefills = existing_prefills[:num_prefill] decodes = existing_decodes[:num_decode] logger.info( @@ -641,13 +828,48 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): len(decodes), ) else: - # Launch new PD workers (custom config or not pre-launched) - prefills, decodes = model_pool.launch_pd_workers( - model_id=model_id, - num_prefill=num_prefill, - num_decode=num_decode, - startup_timeout=300, + # Build WorkerIdentity list for missing workers + workers_to_launch: list[WorkerIdentity] = [] + for i in range(missing_prefill): + workers_to_launch.append( + WorkerIdentity( + model_id, + ConnectionMode.HTTP, + WorkerType.PREFILL, + len(existing_prefills) + i, + ) + ) + for i in range(missing_decode): + workers_to_launch.append( + WorkerIdentity( + model_id, + ConnectionMode.HTTP, + WorkerType.DECODE, + len(existing_decodes) + i, + ) + ) + + logger.info( + "Have %d/%d prefill, %d/%d decode. Launching %d more workers", + len(existing_prefills), + num_prefill, + len(existing_decodes), + num_decode, + len(workers_to_launch), ) + new_instances = model_pool.launch_workers( + workers_to_launch, startup_timeout=300 + ) + + # Combine existing + newly launched + new_prefills = [ + w for w in new_instances if w.worker_type == WorkerType.PREFILL + ] + new_decodes = [ + w for w in new_instances if w.worker_type == WorkerType.DECODE + ] + prefills = existing_prefills + new_prefills + decodes = existing_decodes + new_decodes model_path = prefills[0].model_path if prefills else None @@ -698,17 +920,31 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): try: if num_workers > 1: - # Launch multiple workers on separate GPUs - instances = model_pool.launch_regular_workers( - model_id=model_id, - num_workers=num_workers, - mode=connection_mode, - startup_timeout=300, - ) - if not instances: - pytest.fail( - f"Failed to launch {num_workers} workers for {model_id}" + # Check existing workers + existing = model_pool.get_workers_by_type(model_id, WorkerType.REGULAR) + existing_for_mode = [w for w in existing if w.mode == connection_mode] + + if len(existing_for_mode) >= num_workers: + instances = existing_for_mode[:num_workers] + else: + # Launch missing workers + missing = num_workers - len(existing_for_mode) + workers_to_launch = [ + WorkerIdentity( + model_id, + connection_mode, + WorkerType.REGULAR, + len(existing_for_mode) + i, + ) + for i in range(missing) + ] + new_instances = model_pool.launch_workers( + workers_to_launch, startup_timeout=300 ) + instances = existing_for_mode + new_instances + + if not instances: + pytest.fail(f"Failed to get {num_workers} workers for {model_id}") worker_urls = [inst.worker_url for inst in instances] model_path = instances[0].model_path else: diff --git a/sgl-model-gateway/e2e_test/infra/__init__.py b/sgl-model-gateway/e2e_test/infra/__init__.py index 6bb4d8448..53ddf9eef 100644 --- a/sgl-model-gateway/e2e_test/infra/__init__.py +++ b/sgl-model-gateway/e2e_test/infra/__init__.py @@ -37,7 +37,7 @@ from .gpu_allocator import ( ) from .gpu_monitor import GPUMonitor from .gpu_monitor import should_monitor as should_monitor_gpu -from .model_pool import ModelInstance, ModelPool +from .model_pool import ModelInstance, ModelPool, WorkerIdentity from .model_specs import ( # Default model paths; Model groups CHAT_MODELS, DEFAULT_EMBEDDING_MODEL_PATH, @@ -63,10 +63,11 @@ from .process_utils import ( from .run_eval import run_eval __all__ = [ - # Enums + # Enums and Identity "ConnectionMode", "WorkerType", "Runtime", + "WorkerIdentity", # Convenience sets "LOCAL_MODES", "LOCAL_RUNTIMES", diff --git a/sgl-model-gateway/e2e_test/infra/constants.py b/sgl-model-gateway/e2e_test/infra/constants.py index 60a73f744..2fcd85c38 100644 --- a/sgl-model-gateway/e2e_test/infra/constants.py +++ b/sgl-model-gateway/e2e_test/infra/constants.py @@ -1,6 +1,6 @@ """Constants and enums for E2E test infrastructure.""" -from enum import Enum, auto +from enum import Enum class ConnectionMode(str, Enum): diff --git a/sgl-model-gateway/e2e_test/infra/gpu_allocator.py b/sgl-model-gateway/e2e_test/infra/gpu_allocator.py index 9e5d63371..c6a690dc4 100644 --- a/sgl-model-gateway/e2e_test/infra/gpu_allocator.py +++ b/sgl-model-gateway/e2e_test/infra/gpu_allocator.py @@ -244,19 +244,27 @@ class GPUAllocator: logger.warning("Failed to detect GPUs: %s", e) return [] - def allocate_slots(self, model_specs: dict[str, dict]) -> list[GPUSlot]: + def allocate_slots( + self, model_specs: dict[str, dict], preserve_order: bool = False + ) -> list[GPUSlot]: """Allocate GPU slots based on model memory requirements. - Uses a first-fit decreasing bin-packing algorithm: + Uses a first-fit decreasing bin-packing algorithm by default: 1. Sort models by memory requirement (largest first) 2. For each model, find the first GPU(s) that can fit it 3. For multi-GPU models, find consecutive GPUs + When preserve_order=True, processes models in dict insertion order + (test collection order) instead of sorting by memory. This ensures + models needed by earlier tests are allocated first. + Note: This method tracks used GPUs across multiple calls, so subsequent allocations will use different GPUs than previous ones. Args: model_specs: Dict of model_id -> spec dict with 'memory_gb' and 'tp' keys + preserve_order: If True, allocate in dict order (test order) instead + of sorting by memory size. Default False. Returns: List of GPUSlots with assigned models (only the newly allocated slots) @@ -265,17 +273,21 @@ class GPUAllocator: logger.warning("No GPUs available for allocation") return [] - # Sort models by memory requirement (largest first for better packing) - sorted_models = sorted( - model_specs.items(), - key=lambda x: x[1].get("memory_gb", 0), - reverse=True, - ) + if preserve_order: + # Process in dict insertion order (test collection order) + ordered_models = list(model_specs.items()) + else: + # Sort models by memory requirement (largest first for better packing) + ordered_models = sorted( + model_specs.items(), + key=lambda x: x[1].get("memory_gb", 0), + reverse=True, + ) # Track new slots allocated in this call new_slots: list[GPUSlot] = [] - for model_id, spec in sorted_models: + for model_id, spec in ordered_models: memory_gb = spec.get("memory_gb", 16) tp_size = spec.get("tp", 1) diff --git a/sgl-model-gateway/e2e_test/infra/model_pool.py b/sgl-model-gateway/e2e_test/infra/model_pool.py index 5b01abf12..30b1e1368 100644 --- a/sgl-model-gateway/e2e_test/infra/model_pool.py +++ b/sgl-model-gateway/e2e_test/infra/model_pool.py @@ -31,9 +31,63 @@ from .process_utils import detect_ib_device logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class WorkerIdentity: + """Unique identity for a single worker instance. + + Each worker is uniquely identified by (model_id, mode, worker_type, index). + For example: + - llama-8b:http (regular worker, index 0) + - llama-8b:http:prefill_0 (first prefill worker) + - llama-8b:http:prefill_1 (second prefill worker) + - llama-8b:http:decode_0 (first decode worker) + + Frozen/hashable so it can be used in sets and as dict keys for deduplication. + """ + + model_id: str + mode: ConnectionMode = ConnectionMode.HTTP + worker_type: WorkerType = WorkerType.REGULAR + index: int = 0 + + @property + def is_prefill(self) -> bool: + """Check if this is a prefill worker.""" + return self.worker_type == WorkerType.PREFILL + + @property + def is_decode(self) -> bool: + """Check if this is a decode worker.""" + return self.worker_type == WorkerType.DECODE + + @property + def is_regular(self) -> bool: + """Check if this is a regular worker.""" + return self.worker_type == WorkerType.REGULAR + + @property + def key(self) -> str: + """Unique key for this worker instance.""" + if self.worker_type == WorkerType.REGULAR: + if self.index == 0: + return f"{self.model_id}:{self.mode.value}" + return f"{self.model_id}:{self.mode.value}:{self.index}" + return ( + f"{self.model_id}:{self.mode.value}:{self.worker_type.value}_{self.index}" + ) + + def __str__(self) -> str: + """String representation for logging.""" + return self.key + + @dataclass class ModelInstance: - """A running model instance.""" + """A running model instance. + + Contains both identity (model_id, mode, worker_type) and runtime state + (process, port, gpu_slot, etc.). + """ model_id: str mode: ConnectionMode @@ -42,21 +96,20 @@ class ModelInstance: port: int process: subprocess.Popen gpu_slot: GPUSlot | None + key: str # Unique instance key (e.g., "llama-8b:http:prefill_0") worker_type: WorkerType = WorkerType.REGULAR bootstrap_port: int | None = None # For prefill workers in PD mode last_used: float = 0.0 # Timestamp for MRU eviction _healthy: bool = False # Track if initial health check passed @property - def key(self) -> str: - """Unique key for this instance. - - Regular: 'model_id:mode' (e.g., 'llama-8b:http') - PD workers: 'model_id:mode:worker_type' (e.g., 'llama-8b:http:prefill') - """ - if self.worker_type == WorkerType.REGULAR: - return f"{self.model_id}:{self.mode.value}" - return f"{self.model_id}:{self.mode.value}:{self.worker_type.value}" + def identity(self) -> WorkerIdentity: + """Get the identity (model_id, mode, worker_type) of this instance.""" + return WorkerIdentity( + model_id=self.model_id, + mode=self.mode, + worker_type=self.worker_type, + ) @property def worker_url(self) -> str: @@ -200,86 +253,120 @@ class ModelPool: def startup( self, - requirements: list[tuple[str, ConnectionMode]] | None = None, + requirements: list[WorkerIdentity] | None = None, startup_timeout: int = DEFAULT_STARTUP_TIMEOUT, ) -> None: - """Start worker processes for the required models. + """Start worker processes for the required workers in order. Workers are launched sequentially (one Popen at a time) but boot up concurrently since model loading happens in parallel across processes. This method blocks until all workers pass health checks. + All worker types (regular, prefill, decode) are handled uniformly. + Each WorkerIdentity uniquely identifies a worker by (model_id, mode, + worker_type, index). + Args: - requirements: List of (model_id, mode) tuples specifying what to start. - mode is ConnectionMode.HTTP or ConnectionMode.GRPC. + requirements: List of WorkerIdentity specifying what to start. If None, starts default model in HTTP mode. startup_timeout: Timeout in seconds for all models to become healthy. """ self._startup_timeout = startup_timeout if requirements is None: - requirements = [(DEFAULT_MODEL, ConnectionMode.HTTP)] + requirements = [WorkerIdentity(DEFAULT_MODEL, ConnectionMode.HTTP)] - # Deduplicate and validate - requirements = list(set(requirements)) - valid_requirements = [] - for model_id, mode in requirements: - if model_id not in MODEL_SPECS: - logger.warning("Unknown model %s, skipping", model_id) + # Validate requirements + valid_requirements: list[WorkerIdentity] = [] + for identity in requirements: + if identity.model_id not in MODEL_SPECS: + logger.warning("Unknown model %s, skipping", identity.model_id) continue - if mode not in LOCAL_MODES: - logger.warning("Invalid mode %s for %s, skipping", mode, model_id) + if identity.mode not in LOCAL_MODES: + logger.warning( + "Invalid mode %s for %s, skipping", identity.mode, identity.model_id + ) continue - valid_requirements.append((model_id, mode)) + valid_requirements.append(identity) if not valid_requirements: logger.warning("No valid requirements to start") return - logger.info("Starting model pool with: %s", valid_requirements) + logger.info( + "Starting model pool with %d workers: %s", + len(valid_requirements), + [str(r) for r in valid_requirements], + ) - # Build allocation specs - each (model, mode) combo needs its own slot - # Use "model_id:mode" as the allocation key - allocation_specs = {} - for model_id, mode in valid_requirements: - spec = MODEL_SPECS[model_id] - key = f"{model_id}:{mode.value}" - allocation_specs[key] = { - "model": spec["model"], - "memory_gb": spec.get("memory_gb", 16), - "tp": spec.get("tp", 1), + # Detect IB device once for PD workers + has_pd = any(r.is_prefill or r.is_decode for r in valid_requirements) + ib_device = detect_ib_device() if has_pd else None + if ib_device: + logger.info("Detected InfiniBand device: %s", ib_device) + + # Track bootstrap ports for PD groups (all PD workers of same model/mode share one) + pd_bootstrap_ports: dict[tuple[str, ConnectionMode], int] = {} + + deferred: list[str] = [] + + # Process requirements in order - all workers treated uniformly + for identity in valid_requirements: + spec = get_model_spec(identity.model_id) + tp = spec.get("tp", 1) + + # Check if we have enough GPUs + available_gpus = self.allocator.available_gpus() + if len(available_gpus) < tp: + logger.info( + "Not enough GPUs for %s (need %d, have %d), deferring", + identity, + tp, + len(available_gpus), + ) + deferred.append(str(identity)) + continue + + # Allocate GPU slot + allocation_specs = { + identity.key: { + "model": spec["model"], + "memory_gb": spec.get("memory_gb", 16), + "tp": tp, + } } + slots = self.allocator.allocate_slots(allocation_specs, preserve_order=True) + if not slots: + deferred.append(str(identity)) + continue - # Allocate GPU slots - slots = self.allocator.allocate_slots(allocation_specs) + # Get bootstrap port for PD workers (shared within model/mode group) + bootstrap_port = None + if identity.is_prefill or identity.is_decode: + pd_key = (identity.model_id, identity.mode) + if pd_key not in pd_bootstrap_ports: + pd_bootstrap_ports[pd_key] = get_open_port() + bootstrap_port = pd_bootstrap_ports[pd_key] - # Track which models got slots - launched_keys = set() + # Launch the worker + self._launch_model( + model_id=identity.model_id, + mode=identity.mode, + gpu_slot=slots[0], + worker_type=identity.worker_type, + bootstrap_port=bootstrap_port if identity.is_prefill else None, + ib_device=( + ib_device if (identity.is_prefill or identity.is_decode) else None + ), + instance_key=identity.key, + ) - if not slots: - logger.warning("No GPU slots allocated, launching without GPU assignment") - # Fallback: launch without specific GPU assignment - for model_id, mode in valid_requirements: - self._launch_model(model_id, mode, gpu_slot=None) - launched_keys.add(f"{model_id}:{mode.value}") - else: - # Launch on allocated slots - for slot in slots: - if slot.assigned_model: - # Parse "model_id:mode" back - model_id, mode_str = slot.assigned_model.rsplit(":", 1) - mode = ConnectionMode(mode_str) - self._launch_model(model_id, mode, gpu_slot=slot) - launched_keys.add(slot.assigned_model) - - # Log models that will be launched on-demand (not enough GPUs to pre-launch) - all_keys = set(allocation_specs.keys()) - deferred_keys = all_keys - launched_keys - if deferred_keys: + # Log deferred workers + if deferred: logger.info( - "%d models deferred for on-demand launch: %s", - len(deferred_keys), - deferred_keys, + "%d workers deferred for on-demand launch: %s", + len(deferred), + deferred, ) # Wait for all launched models to be healthy @@ -391,6 +478,7 @@ class ModelPool: port=port, process=proc, gpu_slot=gpu_slot, + key=key, worker_type=worker_type, bootstrap_port=bootstrap_port, last_used=time.time(), @@ -714,230 +802,119 @@ class ModelPool: if inst.model_id == model_id and inst.worker_type == worker_type ] - def launch_regular_workers( + def launch_workers( self, - model_id: str, - num_workers: int, - mode: ConnectionMode = ConnectionMode.HTTP, + workers: list[WorkerIdentity], startup_timeout: int = DEFAULT_STARTUP_TIMEOUT, allow_eviction: bool = True, ) -> list[ModelInstance]: - """Launch multiple regular workers for load balancing. + """Launch workers of any type. + + This is the unified method for launching workers. It handles all worker + types (regular, prefill, decode) uniformly. Args: - model_id: Model identifier from MODEL_SPECS. - num_workers: Number of workers to launch. - mode: Connection mode (HTTP or GRPC). + workers: List of WorkerIdentity objects specifying workers to launch. startup_timeout: Timeout for workers to become healthy. allow_eviction: If True, evict MRU models to free GPUs. Returns: - List of ModelInstance objects. + List of launched ModelInstance objects. """ + if not workers: + return [] + self._startup_timeout = startup_timeout - if model_id not in MODEL_SPECS: - raise ValueError(f"Unknown model: {model_id}") + # Validate all workers + valid_workers: list[WorkerIdentity] = [] + for w in workers: + if w.model_id not in MODEL_SPECS: + logger.warning("Unknown model %s, skipping", w.model_id) + continue + if w.mode not in LOCAL_MODES: + logger.warning("Invalid mode %s, skipping", w.mode) + continue + valid_workers.append(w) - spec = get_model_spec(model_id) - tp = spec.get("tp", 1) - required_gpus = num_workers * tp + if not valid_workers: + return [] + + # Calculate total GPUs needed + total_gpus = 0 + for w in valid_workers: + spec = get_model_spec(w.model_id) + total_gpus += spec.get("tp", 1) # Check if we have enough GPUs available = self.allocator.available_gpus() - if len(available) < required_gpus: + if len(available) < total_gpus: if allow_eviction: logger.info( - "Need %d GPUs for %d workers, only %d available. Evicting MRU models...", - required_gpus, - num_workers, + "Need %d GPUs for %d workers, only %d available. Evicting...", + total_gpus, + len(valid_workers), len(available), ) - # Exclude REGULAR workers of same model/mode from eviction - self._evict_for_gpus( - required_gpus, - exclude_model_id=model_id, - exclude_mode=mode, - exclude_worker_types={WorkerType.REGULAR}, - ) + self._evict_for_gpus(total_gpus) else: - logger.info( - "Need %d GPUs for %d workers, only %d available. " - "Skipping (eviction not allowed).", - required_gpus, - num_workers, + logger.warning( + "Need %d GPUs, only %d available. Skipping launch.", + total_gpus, len(available), ) return [] - # Build allocation specs for all workers + # Build allocation specs allocation_specs = {} - for i in range(num_workers): - key = f"{model_id}:{mode.value}:{i}" - allocation_specs[key] = { + for w in valid_workers: + spec = get_model_spec(w.model_id) + allocation_specs[w.key] = { "model": spec["model"], "memory_gb": spec.get("memory_gb", 16), - "tp": tp, + "tp": spec.get("tp", 1), } # Allocate GPU slots - slots = self.allocator.allocate_slots(allocation_specs) - slot_map = {slot.assigned_model: slot for slot in slots} + slots = self.allocator.allocate_slots(allocation_specs, preserve_order=True) + slot_map = {s.assigned_model: s for s in slots} if not slots: raise RuntimeError( - f"Failed to allocate GPU slots for {num_workers} workers after eviction. " - f"Need {required_gpus} GPUs." + f"Failed to allocate GPU slots for {len(valid_workers)} workers" ) + # Detect IB device for PD workers + has_pd = any(w.is_prefill or w.is_decode for w in valid_workers) + ib_device = detect_ib_device() if has_pd else None + + # Track bootstrap ports for PD groups (shared within model/mode) + pd_bootstrap_ports: dict[tuple[str, ConnectionMode], int] = {} + instances: list[ModelInstance] = [] + for w in valid_workers: + # Get bootstrap port for PD workers + bootstrap_port = None + if w.is_prefill or w.is_decode: + pd_key = (w.model_id, w.mode) + if pd_key not in pd_bootstrap_ports: + pd_bootstrap_ports[pd_key] = get_open_port() + bootstrap_port = pd_bootstrap_ports[pd_key] - # Launch workers - for i in range(num_workers): - key = f"{model_id}:{mode.value}:{i}" - gpu_slot = slot_map.get(key) instance = self._launch_model( - model_id=model_id, - mode=mode, - gpu_slot=gpu_slot, - worker_type=WorkerType.REGULAR, - instance_key=key, + model_id=w.model_id, + mode=w.mode, + gpu_slot=slot_map.get(w.key), + worker_type=w.worker_type, + bootstrap_port=bootstrap_port if w.is_prefill else None, + ib_device=ib_device if (w.is_prefill or w.is_decode) else None, + instance_key=w.key, ) instances.append(instance) - # Wait for all to be healthy self._wait_all_healthy() - return instances - def launch_pd_workers( - self, - model_id: str, - num_prefill: int = 1, - num_decode: int = 1, - mode: ConnectionMode = ConnectionMode.HTTP, - startup_timeout: int = DEFAULT_STARTUP_TIMEOUT, - allow_eviction: bool = True, - ) -> tuple[list[ModelInstance], list[ModelInstance]]: - """Launch prefill and decode workers for PD disaggregation. - - Args: - model_id: Model identifier from MODEL_SPECS. - num_prefill: Number of prefill workers to launch. Defaults to 1. - num_decode: Number of decode workers to launch. Defaults to 1. - mode: Connection mode (HTTP or GRPC). - startup_timeout: Timeout for workers to become healthy. - allow_eviction: If True, evict MRU models to free GPUs. If False, - return empty lists when not enough GPUs available. - - Returns: - Tuple of (prefill_instances, decode_instances). - """ - self._startup_timeout = startup_timeout - - if model_id not in MODEL_SPECS: - raise ValueError(f"Unknown model: {model_id}") - - spec = get_model_spec(model_id) - ib_device = detect_ib_device() - if ib_device: - logger.info("Detected InfiniBand device: %s", ib_device) - - # Calculate total GPUs needed for PD workers - tp = spec.get("tp", 1) - required_gpus = (num_prefill + num_decode) * tp - - # Check if we have enough GPUs - available = self.allocator.available_gpus() - if len(available) < required_gpus: - if allow_eviction: - logger.info( - "Need %d GPUs for PD workers, only %d available. Evicting MRU models...", - required_gpus, - len(available), - ) - # Exclude PD workers of same model/mode, but evict REGULAR workers - self._evict_for_gpus( - required_gpus, - exclude_model_id=model_id, - exclude_mode=mode, - exclude_worker_types={WorkerType.PREFILL, WorkerType.DECODE}, - ) - else: - logger.info( - "Need %d GPUs for PD workers, only %d available. " - "Skipping pre-launch (eviction not allowed).", - required_gpus, - len(available), - ) - return [], [] - - # Build allocation specs for all PD workers - # Each worker needs its own GPU slot - allocation_specs = {} - for i in range(num_prefill): - key = f"{model_id}:{mode.value}:prefill_{i}" - allocation_specs[key] = { - "model": spec["model"], - "memory_gb": spec.get("memory_gb", 16), - "tp": tp, - } - for i in range(num_decode): - key = f"{model_id}:{mode.value}:decode_{i}" - allocation_specs[key] = { - "model": spec["model"], - "memory_gb": spec.get("memory_gb", 16), - "tp": tp, - } - - # Allocate GPU slots - slots = self.allocator.allocate_slots(allocation_specs) - slot_map = {slot.assigned_model: slot for slot in slots} - - if not slots: - raise RuntimeError( - f"Failed to allocate GPU slots for PD workers after eviction. " - f"Need {required_gpus} GPUs." - ) - - prefill_instances: list[ModelInstance] = [] - decode_instances: list[ModelInstance] = [] - - # Launch prefill workers - for i in range(num_prefill): - key = f"{model_id}:{mode.value}:prefill_{i}" - gpu_slot = slot_map.get(key) - bootstrap_port = get_open_port() - instance = self._launch_model( - model_id=model_id, - mode=mode, - gpu_slot=gpu_slot, - worker_type=WorkerType.PREFILL, - bootstrap_port=bootstrap_port, - ib_device=ib_device, - instance_key=key, - ) - prefill_instances.append(instance) - - # Launch decode workers - for i in range(num_decode): - key = f"{model_id}:{mode.value}:decode_{i}" - gpu_slot = slot_map.get(key) - instance = self._launch_model( - model_id=model_id, - mode=mode, - gpu_slot=gpu_slot, - worker_type=WorkerType.DECODE, - ib_device=ib_device, - instance_key=key, - ) - decode_instances.append(instance) - - # Wait for all to be healthy - self._wait_all_healthy() - - return prefill_instances, decode_instances - def get_client( self, model_id: str, mode: ConnectionMode | str = ConnectionMode.HTTP ) -> "openai.OpenAI":