new strategy

Author: OliverBryant

xinference/core/launch_strategy.py

@@ -16,7 +16,7 @@
 from dataclasses import dataclass
 from typing import Dict, List, Mapping, Optional, Set, Tuple, Union
 
-from ..device_utils import initialize_gpu_memory_info, update_gpu_memory_info
+from ..device_utils import update_gpu_memory_info
 from .utils import parse_replica_model_uid
 
 logger = logging.getLogger(__name__)
@@ -34,12 +34,36 @@ class LaunchModelSpec:
     quantization: Optional[str] = None
 
 
-class IdleFirstLaunchStrategy:
+class LaunchStrategy:
+    """
+    Base class for launch strategies.
+    Concrete implementations should override allocate/release/is_idle.
+    """
+
+    def allocate(
+        self,
+        spec: LaunchModelSpec,
+        total_gpu_devices: List[int],
+        user_specified_allocated_devices: Set[int],
+        allocated_gpus: Mapping[int, Set[str]],
+    ) -> List[int]:
+        raise NotImplementedError
+
+    def release(self, model_uid: str, devices: List[int]) -> None:
+        raise NotImplementedError
+
+    def is_idle(self) -> bool:
+        raise NotImplementedError
+
+
+class IdleFirstLaunchStrategy(LaunchStrategy):
     """
     Prefer the GPU running Xinference, otherwise keep allocating onto the emptiest
     remaining GPU.
     """
 
+    _DEFAULT_BOOKED_MB = 1024  # logical reservation per replica
+
     def __init__(
         self,
         total_gpu_devices: List[int],
@@ -50,9 +74,9 @@ def __init__(
     ):
         self._allowed_devices = allowed_devices
         self._total_gpu_devices = self._filter_allowed(total_gpu_devices)
-        self._gpu_memory_info = gpu_memory_info or initialize_gpu_memory_info(
-            self._total_gpu_devices, logger=logger
-        )
+        if gpu_memory_info is None:
+            raise ValueError("gpu_memory_info must be provided for launch strategy")
+        self._gpu_memory_info = gpu_memory_info
         # Track which GPUs have been used in the first round for each model
         self._model_spread_used_gpus: Dict[str, Set[int]] = (
             model_spread_used_gpus if model_spread_used_gpus is not None else {}
@@ -61,6 +85,8 @@ def __init__(
         self._active_model_counts: Dict[str, int] = (
             active_model_counts if active_model_counts is not None else {}
         )
+        # Logical reservations (MB) per GPU for this strategy's base model
+        self._reserved_memory_mb: Dict[int, float] = {}
 
     def _filter_allowed(self, total_gpu_devices: List[int]) -> List[int]:
         if self._allowed_devices is None:
@@ -80,13 +106,15 @@ def _select_emptiest_gpu(
         for dev in candidates:
             update_gpu_memory_info(self._gpu_memory_info, dev, logger=logger)
             available = self._gpu_memory_info.get(dev, {}).get("available", 0)
+            # Deduct logical reservations to avoid stacking replicas too quickly
+            available -= self._reserved_memory_mb.get(dev, 0)
             # Penalize GPUs already planned/allocated to avoid stacking too early
             penalty = pending_gpu_counts.get(dev, 0) + len(
                 allocated_gpus.get(dev, set())
             )
             scored.append((dev, available - penalty))
 
-        scored.sort(key=lambda item: item[1], reverse=True)
+        scored.sort(key=lambda item: (-item[1], item[0]))
         return scored[0][0] if scored else None
 
     def allocate(
@@ -112,20 +140,23 @@ def allocate(
         selected: List[int] = []
 
         while len(selected) < n_gpu:
-            # If some GPUs haven't received a replica for this model yet, try them first
-            if len(used_in_spread) < len(available_total):
+            # Prefer truly idle GPUs first: those without existing allocations
+            unoccupied_gpus = [
+                dev
+                for dev in available_total
+                if dev not in user_specified_allocated_devices
+                and not allocated_gpus.get(dev)
+            ]
+            spreading_phase = bool(unoccupied_gpus) and len(used_in_spread) < len(
+                unoccupied_gpus
+            )
+            if spreading_phase:
+                # First round: try to place replicas on distinct, unoccupied GPUs
                 candidate_pool = [
-                    dev
-                    for dev in available_total
-                    if dev not in user_specified_allocated_devices
-                    and dev not in used_in_spread
+                    dev for dev in unoccupied_gpus if dev not in used_in_spread
                 ]
                 if not candidate_pool:
-                    candidate_pool = [
-                        dev
-                        for dev in available_total
-                        if dev not in user_specified_allocated_devices
-                    ]
+                    candidate_pool = [dev for dev in unoccupied_gpus]
             else:
                 candidate_pool = [
                     dev
@@ -149,6 +180,11 @@ def allocate(
         self._active_model_counts[base_model_uid] = (
             self._active_model_counts.get(base_model_uid, 0) + 1
         )
+        # Reserve logical memory for selected GPUs
+        for dev in selected:
+            self._reserved_memory_mb[dev] = (
+                self._reserved_memory_mb.get(dev, 0.0) + self._DEFAULT_BOOKED_MB
+            )
         return selected
 
     def release(self, model_uid: str, devices: List[int]) -> None:
@@ -160,8 +196,22 @@ def release(self, model_uid: str, devices: List[int]) -> None:
         if count <= 1:
             self._active_model_counts.pop(base_model_uid, None)
             self._model_spread_used_gpus.pop(base_model_uid, None)
+            for dev in devices:
+                if dev in self._reserved_memory_mb:
+                    self._reserved_memory_mb[dev] -= self._DEFAULT_BOOKED_MB
+                    if self._reserved_memory_mb[dev] <= 0:
+                        self._reserved_memory_mb.pop(dev, None)
         else:
             self._active_model_counts[base_model_uid] = count - 1
+            for dev in devices:
+                if dev in self._reserved_memory_mb:
+                    self._reserved_memory_mb[dev] -= self._DEFAULT_BOOKED_MB
+                    if self._reserved_memory_mb[dev] <= 0:
+                        self._reserved_memory_mb.pop(dev, None)
+
+    def is_idle(self) -> bool:
+        """Return True when no active models are tracked by this strategy."""
+        return not self._active_model_counts
 
 
 def create_launch_strategy(

xinference/core/supervisor.py

@@ -1096,6 +1096,9 @@ async def _launch_one_model(worker_ref, _replica_model_uid, rank: int):
             model_type = model_type or "LLM"
 
             try:
+                # Ensure per-base-model launch strategy is ready on worker before concurrent launches
+                await worker_ref.ensure_launch_strategy(model_uid)
+
                 subpool_address = await worker_ref.launch_builtin_model(
                     model_uid=_replica_model_uid,
                     model_name=model_name,

xinference/core/tests/test_worker.py

@@ -44,7 +44,7 @@ def _select_emptiest_gpu(
             )
             scored.append((dev, available - penalty))
 
-        # Prefer higher available memory, then the lowest GPU index.
+        # Prefer higher available memory, then lowest GPU index.
         scored.sort(key=lambda item: (-item[1], item[0]))
         return scored[0][0]
 
@@ -62,12 +62,13 @@ def __init__(
             for idx in cuda_devices
         }
 
-    def _create_launch_strategy_instance(self):
+    def _gather_initial_gpu_memory_info(self):
+        return self._test_gpu_memory_info
+
+    def _create_launch_strategy_instance(self, gpu_memory_info=None):
         return DeterministicIdleFirstLaunchStrategy(
             self._total_gpu_devices,
-            gpu_memory_info=self._test_gpu_memory_info,
-            model_spread_used_gpus=self._model_spread_used_gpus,
-            active_model_counts=self._active_model_counts,
+            gpu_memory_info=gpu_memory_info or self._test_gpu_memory_info,
         )
 
     async def __post_create__(self):

xinference/core/worker.py

@@ -158,6 +158,10 @@ def __init__(
         # Share launch spread/replica counts across strategy instances
         self._model_spread_used_gpus: Dict[str, Set[int]] = {}
         self._active_model_counts: Dict[str, int] = {}
+        # Cached launch strategies per base model
+        self._launch_strategies: Dict[str, Any] = {}
+        # Protect concurrent allocations/releases so bookings stay consistent
+        self._allocation_lock = threading.Lock()
         from ..constants import (
             XINFERENCE_LAUNCH_ALLOWED_GPUS,
             XINFERENCE_LAUNCH_STRATEGY,
@@ -559,7 +563,7 @@ def _collect_user_specified_devices(self) -> Set[int]:
                 user_specified_allocated_devices.add(dev)
         return user_specified_allocated_devices
 
-    def _create_launch_strategy_instance(self):
+    def _gather_initial_gpu_memory_info(self) -> Optional[Dict[int, Dict[str, float]]]:
         # Try to seed strategy with current GPU memory snapshot from NVML
         initial_gpu_memory_info: Optional[Dict[int, Dict[str, float]]] = None
         try:
@@ -576,27 +580,64 @@ def _create_launch_strategy_instance(self):
             initial_gpu_memory_info = gpu_info or None
         except Exception:
             initial_gpu_memory_info = None
+        return initial_gpu_memory_info
 
+    def _create_launch_strategy_instance(
+        self, gpu_memory_info: Optional[Dict[int, Dict[str, float]]] = None
+    ):
+        if gpu_memory_info is None:
+            raise ValueError("gpu_memory_info is required to create launch strategy")
         return create_launch_strategy(
             strategy_name=self._launch_strategy_name,
             total_gpu_devices=self._total_gpu_devices,
             allowed_devices=self._launch_allowed_gpus,
-            gpu_memory_info=initial_gpu_memory_info,
-            model_spread_used_gpus=self._model_spread_used_gpus,
-            active_model_counts=self._active_model_counts,
+            gpu_memory_info=gpu_memory_info,
+        )
+
+    def _get_base_model_uid(self, model_uid: str) -> str:
+        try:
+            base_model_uid, _ = parse_replica_model_uid(model_uid)
+            return base_model_uid
+        except Exception:
+            return model_uid
+
+    def _get_or_create_launch_strategy(self, model_uid: str):
+        base_model_uid = self._get_base_model_uid(model_uid)
+        strategy = self._launch_strategies.get(base_model_uid)
+        if strategy is not None:
+            return strategy
+        strategy = self._create_launch_strategy_instance(
+            gpu_memory_info=self._gather_initial_gpu_memory_info()
         )
+        self._launch_strategies[base_model_uid] = strategy
+        return strategy
+
+    def ensure_launch_strategy(self, model_uid: str):
+        """
+        Ensure a launch strategy exists for the given base model.
+        This is intended to be triggered from supervisor before concurrent launches.
+        """
+        base_model_uid = self._get_base_model_uid(model_uid)
+        with self._allocation_lock:
+            if base_model_uid in self._launch_strategies:
+                return
+            strategy = self._create_launch_strategy_instance(
+                gpu_memory_info=self._gather_initial_gpu_memory_info()
+            )
+            self._launch_strategies[base_model_uid] = strategy
 
     def allocate_devices(self, model_uid: str, n_gpu: int) -> List[int]:
         spec = LaunchModelSpec(model_uid=model_uid, n_gpu=n_gpu)
-        strategy = self._create_launch_strategy_instance()
-        devices = strategy.allocate(
-            spec=spec,
-            total_gpu_devices=self._total_gpu_devices,
-            user_specified_allocated_devices=self._collect_user_specified_devices(),
-            allocated_gpus=self._gpu_to_model_uid,
-        )
-        for dev in devices:
-            self._gpu_to_model_uid[int(dev)].add(model_uid)
+        strategy = self._get_or_create_launch_strategy(model_uid)
+        with self._allocation_lock:
+            devices = strategy.allocate(
+                spec=spec,
+                total_gpu_devices=self._total_gpu_devices,
+                user_specified_allocated_devices=self._collect_user_specified_devices(),
+                allocated_gpus=self._gpu_to_model_uid,
+            )
+            for dev in devices:
+                self._gpu_to_model_uid[int(dev)].add(model_uid)
         return sorted(devices)
 
     def allocate_devices_for_model(
@@ -616,15 +657,16 @@ def allocate_devices_for_model(
             model_format=model_format,
             quantization=quantization,
         )
-        strategy = self._create_launch_strategy_instance()
-        devices = strategy.allocate(
-            spec=spec,
-            total_gpu_devices=self._total_gpu_devices,
-            user_specified_allocated_devices=self._collect_user_specified_devices(),
-            allocated_gpus=self._gpu_to_model_uid,
-        )
-        for dev in devices:
-            self._gpu_to_model_uid[int(dev)].add(model_uid)
+        strategy = self._get_or_create_launch_strategy(model_uid)
+        with self._allocation_lock:
+            devices = strategy.allocate(
+                spec=spec,
+                total_gpu_devices=self._total_gpu_devices,
+                user_specified_allocated_devices=self._collect_user_specified_devices(),
+                allocated_gpus=self._gpu_to_model_uid,
+            )
+            for dev in devices:
+                self._gpu_to_model_uid[int(dev)].add(model_uid)
         return sorted(devices)
 
     async def allocate_devices_with_gpu_idx(
@@ -666,35 +708,40 @@ async def allocate_devices_with_gpu_idx(
         return sorted(gpu_idx)
 
     def release_devices(self, model_uid: str):
-        devices = [
-            dev for dev, uids in self._gpu_to_model_uid.items() if model_uid in uids
-        ]
-        for dev in devices:
-            if model_uid in self._gpu_to_model_uid[dev]:
-                self._gpu_to_model_uid[dev].remove(model_uid)
-            if not self._gpu_to_model_uid[dev]:
-                del self._gpu_to_model_uid[dev]
-
-        # check embedding
-        for dev in self._gpu_to_embedding_model_uids:
-            if model_uid in self._gpu_to_embedding_model_uids[dev]:
-                self._gpu_to_embedding_model_uids[dev].remove(model_uid)
-
-        # check user-specified slots
-        for dev in list(self._user_specified_gpu_to_model_uids):
-            model_infos = [
-                info
-                for info in self._user_specified_gpu_to_model_uids[dev]
-                if info[0] == model_uid
+        base_model_uid = self._get_base_model_uid(model_uid)
+        strategy = self._launch_strategies.get(base_model_uid)
+        with self._allocation_lock:
+            devices = [
+                dev for dev, uids in self._gpu_to_model_uid.items() if model_uid in uids
             ]
-            for model_info in model_infos:
-                self._user_specified_gpu_to_model_uids[dev].remove(model_info)
-            if not self._user_specified_gpu_to_model_uids[dev]:
-                del self._user_specified_gpu_to_model_uids[dev]
-
-        # Keep strategy bookkeeping in sync for spread逻辑
-        strategy = self._create_launch_strategy_instance()
-        strategy.release(model_uid, devices)
+            for dev in devices:
+                if model_uid in self._gpu_to_model_uid[dev]:
+                    self._gpu_to_model_uid[dev].remove(model_uid)
+                if not self._gpu_to_model_uid[dev]:
+                    del self._gpu_to_model_uid[dev]
+
+            # check embedding
+            for dev in self._gpu_to_embedding_model_uids:
+                if model_uid in self._gpu_to_embedding_model_uids[dev]:
+                    self._gpu_to_embedding_model_uids[dev].remove(model_uid)
+
+            # check user-specified slots
+            for dev in list(self._user_specified_gpu_to_model_uids):
+                model_infos = [
+                    info
+                    for info in self._user_specified_gpu_to_model_uids[dev]
+                    if info[0] == model_uid
+                ]
+                for model_info in model_infos:
+                    self._user_specified_gpu_to_model_uids[dev].remove(model_info)
+                if not self._user_specified_gpu_to_model_uids[dev]:
+                    del self._user_specified_gpu_to_model_uids[dev]
+
+            # Keep strategy bookkeeping in sync for spread逻辑
+            if strategy is not None:
+                strategy.release(model_uid, devices)
+                if strategy.is_idle():
+                    self._launch_strategies.pop(base_model_uid, None)
 
     async def _create_subpool(
         self,