new strategy

Author: OliverBryant

xinference/core/launch_strategy.py

@@ -16,7 +16,6 @@
 from dataclasses import dataclass
 from typing import Dict, List, Mapping, Optional, Set, Tuple, Union
 
-from ..device_utils import update_gpu_memory_info
 from .utils import parse_replica_model_uid
 
 logger = logging.getLogger(__name__)
@@ -57,10 +56,7 @@ def is_idle(self) -> bool:
 
 
 class IdleFirstLaunchStrategy(LaunchStrategy):
-    """
-    Prefer the GPU running Xinference, otherwise keep allocating onto the emptiest
-    remaining GPU.
-    """
+    """Always place replicas onto the currently emptiest GPU."""
 
     _DEFAULT_BOOKED_MB = 1024  # logical reservation per replica
 
@@ -104,17 +100,29 @@ def _select_emptiest_gpu(
 
         scored: List[Tuple[int, Union[int, float]]] = []
         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], item[0]))
+            score = available - penalty
+            scored.append((dev, score))
+
+        # If scores are infinite (heartbeat missing => infinite available),
+        # fall back to smallest reserved/penalty; tie-break by GPU index.
+        if any(val[1] == float("inf") for val in scored):
+            scored.sort(
+                key=lambda item: (
+                    self._reserved_memory_mb.get(item[0], 0.0)
+                    + pending_gpu_counts.get(item[0], 0)
+                    + len(allocated_gpus.get(item[0], set())),
+                    item[0],
+                )
+            )
+        else:
+            scored.sort(key=lambda item: (-item[1], item[0]))
         return scored[0][0] if scored else None
 
     def allocate(
@@ -133,36 +141,18 @@ def allocate(
             base_model_uid, _ = parse_replica_model_uid(model_uid)
         except Exception:
             base_model_uid = model_uid
-        used_in_spread = self._model_spread_used_gpus.setdefault(base_model_uid, set())
         n_gpu = spec.n_gpu
 
         pending_gpu_counts: Dict[int, int] = {}
         selected: List[int] = []
 
         while len(selected) < n_gpu:
-            # Prefer truly idle GPUs first: those without existing allocations
-            unoccupied_gpus = [
+            # Always pick the emptiest eligible GPU (excludes user-specified ones)
+            candidate_pool = [
                 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 unoccupied_gpus if dev not in used_in_spread
-                ]
-                if not candidate_pool:
-                    candidate_pool = [dev for dev in unoccupied_gpus]
-            else:
-                candidate_pool = [
-                    dev
-                    for dev in available_total
-                    if dev not in user_specified_allocated_devices
-                ]
             emptiest_gpu = self._select_emptiest_gpu(
                 candidate_pool, pending_gpu_counts, allocated_gpus
             )
@@ -173,10 +163,9 @@ def allocate(
             pending_gpu_counts[emptiest_gpu] = (
                 pending_gpu_counts.get(emptiest_gpu, 0) + 1
             )
-            used_in_spread.add(emptiest_gpu)
 
-        # Persist spread history for this base model
-        self._model_spread_used_gpus[base_model_uid] = used_in_spread
+        # Persist spread history for compatibility with release bookkeeping
+        self._model_spread_used_gpus.setdefault(base_model_uid, set()).update(selected)
         self._active_model_counts[base_model_uid] = (
             self._active_model_counts.get(base_model_uid, 0) + 1
         )

xinference/core/supervisor.py

@@ -30,6 +30,7 @@
     List,
     Literal,
     Optional,
+    Set,
     Tuple,
     Type,
     Union,
@@ -48,6 +49,7 @@
 from ..core.status_guard import InstanceInfo, LaunchStatus
 from ..model.utils import get_engine_params_by_name
 from ..types import PeftModelConfig
+from .launch_strategy import create_launch_strategy
 from .metrics import record_metrics
 from .resource import GPUStatus, ResourceStatus
 from .utils import (
@@ -899,6 +901,44 @@ def _get_worker_refs_by_ip(self, ip: str) -> List[xo.ActorRefType["WorkerActor"]
         )
         return refs
 
+    def _build_gpu_memory_info(
+        self, worker_ref
+    ) -> Optional[Dict[int, Dict[str, float]]]:
+        """Use latest heartbeat data for GPU memory snapshot."""
+        worker_status = self._worker_status.get(worker_ref.address)
+        if worker_status is None:
+            return None
+        gpu_info: Dict[int, Dict[str, float]] = {}
+        for dev, status in worker_status.status.items():
+            if isinstance(status, GPUStatus) and str(dev).startswith("gpu-"):
+                try:
+                    idx = int(str(dev).split("-", 1)[1])
+                except Exception:
+                    continue
+                gpu_info[idx] = {
+                    "total": status.mem_total // (1024**2),
+                    "used": status.mem_used // (1024**2),
+                    "available": status.mem_free // (1024**2),
+                }
+        return gpu_info or None
+
+    async def _install_strategy_on_worker(self, model_uid: str, worker_ref) -> None:
+        ctx = await worker_ref.get_launch_strategy_context()
+        gpu_memory_info = self._build_gpu_memory_info(worker_ref)
+        if gpu_memory_info is None:
+            # Heartbeat disabled or missing: assume all visible GPUs are available with "infinite" mem
+            gpu_memory_info = {
+                dev: {"total": float("inf"), "used": 0.0, "available": float("inf")}
+                for dev in ctx["total_gpu_devices"]
+            }
+        strategy = create_launch_strategy(
+            strategy_name=ctx["launch_strategy_name"],
+            total_gpu_devices=ctx["total_gpu_devices"],
+            allowed_devices=ctx["allowed_devices"],
+            gpu_memory_info=gpu_memory_info,
+        )
+        await worker_ref.install_launch_strategy(model_uid, strategy)
+
     @log_async(logger=logger)
     async def launch_builtin_model(
         self,
@@ -1096,9 +1136,6 @@ 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,
@@ -1140,6 +1177,7 @@ async def _launch_model():
             try:
                 # Pre-fetch worker loads for balanced scheduling
                 worker_candidates = []
+                prepared_workers: Set[str] = set()
 
                 if target_worker_refs:
                     workers = target_worker_refs
@@ -1188,6 +1226,11 @@ async def _launch_model():
                         _idx
                     ].append(worker_ref)
 
+                    # Prepare launch strategy per worker once before launching replicas
+                    if worker_ref.address not in prepared_workers:
+                        await self._install_strategy_on_worker(model_uid, worker_ref)
+                        prepared_workers.add(worker_ref.address)
+
                     if enable_xavier and _idx == 0:
                         """
                         Start the rank 0 model actor on the worker that holds the rank 1 replica,
@@ -1359,6 +1402,7 @@ async def _launch_model():
                     "n_worker cannot be larger than the number of available workers."
                 )
             try:
+                prepared_workers: Set[str] = set()
                 for _idx, rep_model_uid in enumerate(
                     iter_replica_model_uid(model_uid, replica)
                 ):
@@ -1375,6 +1419,11 @@ async def _launch_model():
                         ].replica_to_worker_refs[_idx].append(worker_ref)
                         nonlocal model_type
                         model_type = model_type or "LLM"
+                        if worker_ref.address not in prepared_workers:
+                            await self._install_strategy_on_worker(
+                                model_uid, worker_ref
+                            )
+                            prepared_workers.add(worker_ref.address)
                         if i_worker > 1:
                             assert (
                                 driver_info is not None