[smg][ci] preserve model launch order with test collected (#16618)

.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: ""

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
+# Track max worker counts: (model_id, mode, worker_type) -> max_count
-_scanned_backends: set[str] = set()  # {"grpc", "http", "openai", ...}
+# This unified approach handles regular, prefill, and decode workers the same way
-_scanned_models: set[str] = set()  # Models needed by tests
+_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
+        # Extract model from marker or use default
-        has_model_marker = False
+        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:
+        # Check for workers marker (@pytest.mark.workers(...))
-                    # Extract model names from parametrize
+        workers_marker = item.get_closest_marker("workers")
-                    if isinstance(param_values, (list, tuple)):
+        prefill_count = 0
-                        _scanned_models.update(param_values)
+        decode_count = 0
-                        has_model_marker = True
+        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
+        # Track if this test needs default model
-        model_marker = item.get_closest_marker(PARAM_MODEL)
+        is_e2e = item.get_closest_marker("e2e") is not None
-        if model_marker and model_marker.args:
+        if model_id is None and is_e2e:
-            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"):
             _needs_default_model = True
+            model_id = DEFAULT_MODEL
 
-    logger.info(
+        # Track worker requirements and calculate this test's GPU requirement
-        "Scanned test requirements - backends: %s, models: %s, needs default: %s",
+        test_gpus = 0
-        _scanned_backends or {"(none)"},
+        if model_id and backends:
-        _scanned_models or {"(none)"},
+            for backend in backends:
-        _needs_default_model,
+                # "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)
 
-def get_pool_requirements() -> list[tuple[str, ConnectionMode]]:
+        # 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[WorkerIdentity]:
     """Build pool requirements from scanned test markers.
 
     Returns:
-        List of (model_id, ConnectionMode) tuples to try to pre-launch.
+        List of WorkerIdentity objects to pre-launch.
-        Models that don't fit will be launched on-demand.
+        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,
+    for model_id, mode, worker_type in _first_seen_order:
-    # or if no models were specified at all
+        if model_id not in first_worker_type_per_model:
-    if _needs_default_model or not models:
+            first_worker_type_per_model[model_id] = worker_type
-        models.add(DEFAULT_MODEL)
+            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
+    # Generate individual WorkerIdentity objects in first-seen order
-    # Filter to local backends only (grpc, http) - cloud backends don't need workers
+    requirements: list[WorkerIdentity] = []
-    local_modes: set[ConnectionMode] = set()
+    for model_id, mode, worker_type in _first_seen_order:
-    for backend in _scanned_backends:
+        # Skip regular workers for models that have PD first
-        try:
+        if model_id in models_with_pd_first and worker_type == WorkerType.REGULAR:
-            mode = ConnectionMode(backend)
+            continue
-            if mode in LOCAL_MODES:
-                local_modes.add(mode)
-        except ValueError:
-            # Not a ConnectionMode (e.g., "openai", "xai", "pd") - skip
-            pass
 
-    # Default to HTTP if no local backends specified
+        count = _worker_counts.get((model_id, mode, worker_type), 1)
-    if not local_modes:
+        for i in range(count):
-        local_modes = {ConnectionMode.HTTP}
+            requirements.append(WorkerIdentity(model_id, mode, worker_type, i))
 
-    # Build requirements: each model needs each mode
+    # Add default if no requirements
-    requirements: list[tuple[str, ConnectionMode]] = []
+    if not requirements:
-    for model in models:
+        requirements.append(WorkerIdentity(DEFAULT_MODEL, ConnectionMode.HTTP))
-        for mode in local_modes:
-            requirements.append((model, mode))
 
     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
+    - Scans test markers to determine required workers (model, mode, type, count)
-    - Launches workers sequentially, but they boot up concurrently
+    - 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]
+        # Create WorkerIdentity objects (regular workers only from env vars)
-        logger.info("Using env var requirements: %s", requirements)
+        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 (
+        # Calculate how many more we need (if any)
-            len(existing_prefills) >= num_prefill
+        missing_prefill = max(0, num_prefill - len(existing_prefills))
-            and len(existing_decodes) >= num_decode
+        missing_decode = max(0, num_decode - len(existing_decodes))
-        ):
+
-            # Use pre-launched workers
+        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)
+            # Build WorkerIdentity list for missing workers
-            prefills, decodes = model_pool.launch_pd_workers(
+            workers_to_launch: list[WorkerIdentity] = []
-                model_id=model_id,
+            for i in range(missing_prefill):
-                num_prefill=num_prefill,
+                workers_to_launch.append(
-                num_decode=num_decode,
+                    WorkerIdentity(
-                startup_timeout=300,
+                        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
+                # Check existing workers
-                instances = model_pool.launch_regular_workers(
+                existing = model_pool.get_workers_by_type(model_id, WorkerType.REGULAR)
-                    model_id=model_id,
+                existing_for_mode = [w for w in existing if w.mode == connection_mode]
-                    num_workers=num_workers,
+
-                    mode=connection_mode,
+                if len(existing_for_mode) >= num_workers:
-                    startup_timeout=300,
+                    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(
+                    pytest.fail(f"Failed to get {num_workers} workers for {model_id}")
-                        f"Failed to launch {num_workers} workers for {model_id}"
-                    )
                 worker_urls = [inst.worker_url for inst in instances]
                 model_path = instances[0].model_path
             else:

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",

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):

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,8 +273,12 @@ class GPUAllocator:
             logger.warning("No GPUs available for allocation")
             return []
 
+        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)
-        sorted_models = sorted(
+            ordered_models = sorted(
                 model_specs.items(),
                 key=lambda x: x[1].get("memory_gb", 0),
                 reverse=True,
@@ -275,7 +287,7 @@ class GPUAllocator:
         # 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)
 

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:
+    def identity(self) -> WorkerIdentity:
-        """Unique key for this instance.
+        """Get the identity (model_id, mode, worker_type) of this instance."""
-
+        return WorkerIdentity(
-        Regular: 'model_id:mode' (e.g., 'llama-8b:http')
+            model_id=self.model_id,
-        PD workers: 'model_id:mode:worker_type' (e.g., 'llama-8b:http:prefill')
+            mode=self.mode,
-        """
+            worker_type=self.worker_type,
-        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}"
 
     @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.
+            requirements: List of WorkerIdentity specifying what to start.
-                         mode is ConnectionMode.HTTP or ConnectionMode.GRPC.
                          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
+        # Validate requirements
-        requirements = list(set(requirements))
+        valid_requirements: list[WorkerIdentity] = []
-        valid_requirements = []
+        for identity in requirements:
-        for model_id, mode in requirements:
+            if identity.model_id not in MODEL_SPECS:
-            if model_id not in MODEL_SPECS:
+                logger.warning("Unknown model %s, skipping", identity.model_id)
-                logger.warning("Unknown model %s, skipping", model_id)
                 continue
-            if mode not in LOCAL_MODES:
+            if identity.mode not in LOCAL_MODES:
-                logger.warning("Invalid mode %s for %s, skipping", mode, model_id)
+                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
+        # Detect IB device once for PD workers
-        # Use "model_id:mode" as the allocation key
+        has_pd = any(r.is_prefill or r.is_decode for r in valid_requirements)
-        allocation_specs = {}
+        ib_device = detect_ib_device() if has_pd else None
-        for model_id, mode in valid_requirements:
+        if ib_device:
-            spec = MODEL_SPECS[model_id]
+            logger.info("Detected InfiniBand device: %s", ib_device)
-            key = f"{model_id}:{mode.value}"
+
-            allocation_specs[key] = {
+        # 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": spec.get("tp", 1),
+                    "tp": tp,
                 }
-
+            }
-        # Allocate GPU slots
+            slots = self.allocator.allocate_slots(allocation_specs, preserve_order=True)
-        slots = self.allocator.allocate_slots(allocation_specs)
-
-        # Track which models got slots
-        launched_keys = set()
-
             if not slots:
-            logger.warning("No GPU slots allocated, launching without GPU assignment")
+                deferred.append(str(identity))
-            # Fallback: launch without specific GPU assignment
+                continue
-            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)
+            # Get bootstrap port for PD workers (shared within model/mode group)
-        all_keys = set(allocation_specs.keys())
+            bootstrap_port = None
-        deferred_keys = all_keys - launched_keys
+            if identity.is_prefill or identity.is_decode:
-        if deferred_keys:
+                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]
+
+            # 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,
+            )
+
+        # Log deferred workers
+        if deferred:
             logger.info(
-                "%d models deferred for on-demand launch: %s",
+                "%d workers deferred for on-demand launch: %s",
-                len(deferred_keys),
+                len(deferred),
-                deferred_keys,
+                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,
+        workers: list[WorkerIdentity],
-        num_workers: int,
-        mode: ConnectionMode = ConnectionMode.HTTP,
         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.
+            workers: List of WorkerIdentity objects specifying workers to launch.
-            num_workers: Number of workers to launch.
-            mode: Connection mode (HTTP or GRPC).
             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:
+        # Validate all workers
-            raise ValueError(f"Unknown model: {model_id}")
+        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)
+        if not valid_workers:
-        tp = spec.get("tp", 1)
+            return []
-        required_gpus = num_workers * tp
+
+        # 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...",
+                    "Need %d GPUs for %d workers, only %d available. Evicting...",
-                    required_gpus,
+                    total_gpus,
-                    num_workers,
+                    len(valid_workers),
                     len(available),
                 )
-                # Exclude REGULAR workers of same model/mode from eviction
+                self._evict_for_gpus(total_gpus)
-                self._evict_for_gpus(
-                    required_gpus,
-                    exclude_model_id=model_id,
-                    exclude_mode=mode,
-                    exclude_worker_types={WorkerType.REGULAR},
-                )
             else:
-                logger.info(
+                logger.warning(
-                    "Need %d GPUs for %d workers, only %d available. "
+                    "Need %d GPUs, only %d available. Skipping launch.",
-                    "Skipping (eviction not allowed).",
+                    total_gpus,
-                    required_gpus,
-                    num_workers,
                     len(available),
                 )
                 return []
 
-        # Build allocation specs for all workers
+        # Build allocation specs
         allocation_specs = {}
-        for i in range(num_workers):
+        for w in valid_workers:
-            key = f"{model_id}:{mode.value}:{i}"
+            spec = get_model_spec(w.model_id)
-            allocation_specs[key] = {
+            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)
+        slots = self.allocator.allocate_slots(allocation_specs, preserve_order=True)
-        slot_map = {slot.assigned_model: slot for slot in slots}
+        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"Failed to allocate GPU slots for {len(valid_workers)} workers"
-                f"Need {required_gpus} GPUs."
             )
 
+        # 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,
+                model_id=w.model_id,
-                mode=mode,
+                mode=w.mode,
-                gpu_slot=gpu_slot,
+                gpu_slot=slot_map.get(w.key),
-                worker_type=WorkerType.REGULAR,
+                worker_type=w.worker_type,
-                instance_key=key,
+                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":