From f7870ee2bb500dc29639a6930fa5e87e1a86af29 Mon Sep 17 00:00:00 2001 From: Pardeep Singh Date: Thu, 22 Oct 2026 15:57:59 +0530 Subject: [PATCH 1/5] [Serve] Expose multiplexed_model_ids on ReplicaDetails Record the multiplexed model IDs each replica reports in its actor details, so they appear in the Serve instance details (REST API and dashboard). Closes #65347 Signed-off-by: Pardeep Singh --- python/ray/serve/_private/deployment_state.py | 7 ++- python/ray/serve/schema.py | 8 ++++ python/ray/serve/tests/test_multiplex.py | 26 +++++++++-- .../serve/tests/unit/test_deployment_state.py | 45 +++++++++++++++++++ 4 files changed, 81 insertions(+), 5 deletions(-) diff --git a/python/ray/serve/_private/deployment_state.py b/python/ray/serve/_private/deployment_state.py index 1479534887df..7c3efbb49b3f 100644 --- a/python/ray/serve/_private/deployment_state.py +++ b/python/ray/serve/_private/deployment_state.py @@ -2146,8 +2146,13 @@ def get_running_replica_info( ) def record_multiplexed_model_ids(self, multiplexed_model_ids: List[str]): - """Record the multiplexed model ids for this replica.""" + """Record the multiplexed model ids for this replica. + + The IDs are also exposed in the replica's ``ReplicaDetails`` (sorted, + since the replica reports them in no particular order). + """ self._multiplexed_model_ids = multiplexed_model_ids + self.update_actor_details(multiplexed_model_ids=sorted(multiplexed_model_ids)) def record_routing_stats(self, routing_stats: Optional[Dict[str, Any]]): """Record the routing stats for this replica. diff --git a/python/ray/serve/schema.py b/python/ray/serve/schema.py index d4d6f491d7ff..83d0c2b56fed 100644 --- a/python/ray/serve/schema.py +++ b/python/ray/serve/schema.py @@ -1396,6 +1396,14 @@ class ReplicaDetails(ServeActorDetails): "state from the running replica actor." ) ) + multiplexed_model_ids: List[str] = Field( + default_factory=list, + description=( + "IDs of the multiplexed models that are loaded or loading on this " + "replica, as last reported to the controller. Empty for " + "deployments that don't use model multiplexing." + ), + ) @PublicAPI(stability="alpha") diff --git a/python/ray/serve/tests/test_multiplex.py b/python/ray/serve/tests/test_multiplex.py index 788c884beff5..b7e16b8a54d5 100644 --- a/python/ray/serve/tests/test_multiplex.py +++ b/python/ray/serve/tests/test_multiplex.py @@ -1,6 +1,6 @@ import asyncio import os -from typing import List +from typing import List, Set import httpx import pytest @@ -11,11 +11,15 @@ from ray._common.utils import get_or_create_event_loop from ray.serve._private.common import DeploymentID, ReplicaID from ray.serve._private.config import DeploymentConfig -from ray.serve._private.constants import SERVE_MULTIPLEXED_MODEL_ID +from ray.serve._private.constants import ( + SERVE_DEFAULT_APP_NAME, + SERVE_MULTIPLEXED_MODEL_ID, +) from ray.serve._private.request_router import RequestRouter from ray.serve.context import _get_internal_replica_context from ray.serve.handle import DeploymentHandle from ray.serve.multiplex import _ModelMultiplexWrapper +from ray.serve.schema import ServeInstanceDetails def _get_request_router(handle: DeploymentHandle) -> RequestRouter: @@ -322,7 +326,7 @@ async def __call__(self, model_id: str): replica_id = handle.remote("model1").result() def check_replica_information( - model_ids: List[str], + model_ids: Set[str], ): if not handle.is_initialized: handle._init() @@ -335,7 +339,21 @@ def check_replica_information( ): return False - return True + # The controller also exposes the model IDs in the replica details + # returned by the REST API / `get_serve_details()`. + serve_details = serve_instance.get_serve_details() + app = serve_details["applications"][SERVE_DEFAULT_APP_NAME] + replicas = app["deployments"]["MyModel"]["replicas"] + if len(replicas) != 1 or replicas[0]["replica_id"] != replica_id.unique_id: + return False + if set(replicas[0].get("multiplexed_model_ids", [])) != model_ids: + return False + + details = ServeInstanceDetails(**serve_details) + (replica_details,) = ( + details.applications[SERVE_DEFAULT_APP_NAME].deployments["MyModel"].replicas + ) + return set(replica_details.multiplexed_model_ids) == model_ids wait_for_condition( check_replica_information, diff --git a/python/ray/serve/tests/unit/test_deployment_state.py b/python/ray/serve/tests/unit/test_deployment_state.py index cfa2d7f4d4b6..9135f807baf6 100644 --- a/python/ray/serve/tests/unit/test_deployment_state.py +++ b/python/ray/serve/tests/unit/test_deployment_state.py @@ -26,6 +26,7 @@ ReplicaID, ReplicaMetricReport, ReplicaState, + RequestRoutingInfo, TargetCapacityDirection, TimeStampedValue, ) @@ -1074,6 +1075,50 @@ def test_recent_dead_replicas_bounded(mock_deployment_state_manager): assert len(ds._recent_dead_replicas) == 2 +def test_record_multiplexed_model_ids_updates_replica_details( + mock_deployment_state_manager, +): + """Model IDs reported by a replica are exposed in its ReplicaDetails.""" + create_dsm, _, _, _ = mock_deployment_state_manager + dsm: DeploymentStateManager = create_dsm() + dsm.deploy(TEST_DEPLOYMENT_ID, deployment_info()[0]) + ds = dsm._get_deployment_state_for_testing(TEST_DEPLOYMENT_ID) + + dsm.update() + ds._replicas.get()[0]._actor.set_ready() + dsm.update() + check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) + + replica_id = ds._replicas.get()[0].replica_id + (details,) = ds.list_replica_details() + assert details.multiplexed_model_ids == [] + # Not reported yet, so it is omitted from the user-facing details. + assert "multiplexed_model_ids" not in details.model_fields_set + + ds.record_request_routing_info( + RequestRoutingInfo( + replica_id=replica_id, multiplexed_model_ids=["model-b", "model-a"] + ) + ) + (details,) = ds.list_replica_details() + assert details.replica_id == replica_id.unique_id + assert details.multiplexed_model_ids == ["model-a", "model-b"] + + # An evicted model disappears from the details on the next report. + ds.record_request_routing_info( + RequestRoutingInfo(replica_id=replica_id, multiplexed_model_ids=["model-c"]) + ) + (details,) = ds.list_replica_details() + assert details.multiplexed_model_ids == ["model-c"] + + # Routing stats alone (no model IDs) leave the recorded IDs unchanged. + ds.record_request_routing_info( + RequestRoutingInfo(replica_id=replica_id, routing_stats={"key": 1}) + ) + (details,) = ds.list_replica_details() + assert details.multiplexed_model_ids == ["model-c"] + + def test_force_kill(mock_deployment_state_manager): create_dsm, timer, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() From f340d072c05820d5139645a0373bd40edb97fd23 Mon Sep 17 00:00:00 2001 From: Pardeep Singh Date: Sun, 25 Oct 2026 15:10:43 +0530 Subject: [PATCH 2/5] Drop failed and cancelled model loads from the reported IDs; use a tuple - load_model discards the model ID in a finally block and re-pushes when the load failed or was cancelled (also while waiting for the cache lock), so the controller doesn't keep a stale ID. - ReplicaDetails.multiplexed_model_ids is a Tuple[str, ...] so the frozen model stays hashable, and its description notes that it is empty after a controller restart until the replica's cache changes. Signed-off-by: Pardeep Singh --- python/ray/serve/_private/deployment_state.py | 4 +- python/ray/serve/multiplex.py | 83 ++++++++++--------- python/ray/serve/schema.py | 9 +- python/ray/serve/tests/test_multiplex.py | 64 ++++++++++++++ .../serve/tests/unit/test_deployment_state.py | 10 ++- 5 files changed, 122 insertions(+), 48 deletions(-) diff --git a/python/ray/serve/_private/deployment_state.py b/python/ray/serve/_private/deployment_state.py index 7c3efbb49b3f..b778ba82cd59 100644 --- a/python/ray/serve/_private/deployment_state.py +++ b/python/ray/serve/_private/deployment_state.py @@ -2152,7 +2152,9 @@ def record_multiplexed_model_ids(self, multiplexed_model_ids: List[str]): since the replica reports them in no particular order). """ self._multiplexed_model_ids = multiplexed_model_ids - self.update_actor_details(multiplexed_model_ids=sorted(multiplexed_model_ids)) + self.update_actor_details( + multiplexed_model_ids=tuple(sorted(multiplexed_model_ids)) + ) def record_routing_stats(self, routing_stats: Optional[Dict[str, Any]]): """Record the routing stats for this replica. diff --git a/python/ray/serve/multiplex.py b/python/ray/serve/multiplex.py index e59d2e379805..b42ca9d35eb6 100644 --- a/python/ray/serve/multiplex.py +++ b/python/ray/serve/multiplex.py @@ -197,46 +197,51 @@ async def load_model(self, model_id: str) -> Any: # the replica. self._push_multiplexed_replica_info = True self._model_load_tasks.add(model_id) - async with self._model_cache_lock: - # Check if the model has been loaded by another request. - if model_id in self.models: - return self.models[model_id] - try: - # If the number of models per replica is specified, check - # if the number of models on the current replica has - # reached the limit. - if ( - self.max_num_models_per_replica > 0 - and len(self.models) >= self.max_num_models_per_replica - ): - # Unload the least recently used model. - await self.unload_model_lru() - self._push_multiplexed_replica_info = True - - # Load the model. - logger.debug(f"Loading model '{model_id}'.") - self.models_load_counter.inc() - load_start_time = time.time() - if self.self_arg is None: - self.models[model_id] = await self._func(model_id) - else: - self.models[model_id] = await self._func( - self.self_arg, model_id + try: + async with self._model_cache_lock: + # Check if the model has been loaded by another request. + if model_id in self.models: + return self.models[model_id] + try: + # If the number of models per replica is specified, check + # if the number of models on the current replica has + # reached the limit. + if ( + self.max_num_models_per_replica > 0 + and len(self.models) >= self.max_num_models_per_replica + ): + # Unload the least recently used model. + await self.unload_model_lru() + self._push_multiplexed_replica_info = True + + # Load the model. + logger.debug(f"Loading model '{model_id}'.") + self.models_load_counter.inc() + load_start_time = time.time() + if self.self_arg is None: + self.models[model_id] = await self._func(model_id) + else: + self.models[model_id] = await self._func( + self.self_arg, model_id + ) + load_latency_ms = (time.time() - load_start_time) * 1000.0 + logger.debug( + f"Successfully loaded model '{model_id}' in " + f"{load_latency_ms:.1f}ms." ) - load_latency_ms = (time.time() - load_start_time) * 1000.0 - logger.debug( - f"Successfully loaded model '{model_id}' in " - f"{load_latency_ms:.1f}ms." - ) - self._model_load_tasks.discard(model_id) - self.model_load_latency_ms.observe(load_latency_ms) - return self.models[model_id] - except Exception as e: - logger.error( - f"Failed to load model '{model_id}'. Error: {e}", - ) - self._model_load_tasks.discard(model_id) - raise e + self.model_load_latency_ms.observe(load_latency_ms) + return self.models[model_id] + except Exception as e: + logger.error( + f"Failed to load model '{model_id}'. Error: {e}", + ) + raise e + finally: + # The ID was pushed while loading. If the load failed or was + # cancelled, push again so the controller drops it. + self._model_load_tasks.discard(model_id) + if model_id not in self.models: + self._push_multiplexed_replica_info = True async def unload_model_lru(self) -> None: """Unload the least recently used model.""" diff --git a/python/ray/serve/schema.py b/python/ray/serve/schema.py index 83d0c2b56fed..168a975ab41b 100644 --- a/python/ray/serve/schema.py +++ b/python/ray/serve/schema.py @@ -4,7 +4,7 @@ from collections import Counter from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, Dict, List, Optional, Set, Union +from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union from zlib import crc32 from pydantic import ( @@ -1396,12 +1396,13 @@ class ReplicaDetails(ServeActorDetails): "state from the running replica actor." ) ) - multiplexed_model_ids: List[str] = Field( - default_factory=list, + multiplexed_model_ids: Tuple[str, ...] = Field( + default=(), description=( "IDs of the multiplexed models that are loaded or loading on this " "replica, as last reported to the controller. Empty for " - "deployments that don't use model multiplexing." + "deployments that don't use model multiplexing. After a controller " + "restart it stays empty until the replica's model cache next changes." ), ) diff --git a/python/ray/serve/tests/test_multiplex.py b/python/ray/serve/tests/test_multiplex.py index b7e16b8a54d5..d345163a5817 100644 --- a/python/ray/serve/tests/test_multiplex.py +++ b/python/ray/serve/tests/test_multiplex.py @@ -225,6 +225,70 @@ async def model_load_func(model_id: str): assert "3" in multiplexer.models assert len(multiplexer._model_load_tasks) == 0 + async def test_failed_load_drops_model_id(self, start_serve_with_context): + """A model that fails to load is dropped from the pushed model IDs.""" + + release = asyncio.Event() + + async def model_load_func(model_id: str): + await release.wait() + raise RuntimeError("load failed") + + multiplexer = _ModelMultiplexWrapper( + model_load_func, None, max_num_models_per_replica=1 + ) + await multiplexer.metrics_pusher.graceful_shutdown() + + task = get_or_create_event_loop().create_task(multiplexer.load_model("1")) + await asyncio.sleep(0.1) + assert multiplexer._get_loading_and_loaded_model_ids() == ["1"] + # The loading ID has been pushed. + multiplexer._push_multiplexed_replica_info = False + + release.set() + with pytest.raises(RuntimeError, match="load failed"): + await task + assert multiplexer._get_loading_and_loaded_model_ids() == [] + # Push again so the controller drops the ID pushed while loading. + assert multiplexer._push_multiplexed_replica_info + + @pytest.mark.parametrize("cancel_while", ["loading", "waiting_for_lock"]) + async def test_cancelled_load_drops_model_id( + self, start_serve_with_context, cancel_while + ): + """A model whose load is cancelled is dropped from the pushed model IDs.""" + release = asyncio.Event() + + async def model_load_func(model_id: str): + await release.wait() + return model_id + + multiplexer = _ModelMultiplexWrapper( + model_load_func, None, max_num_models_per_replica=2 + ) + await multiplexer.metrics_pusher.graceful_shutdown() + + loop = get_or_create_event_loop() + if cancel_while == "waiting_for_lock": + # Model "1" holds the cache lock while it loads. + holder = loop.create_task(multiplexer.load_model("1")) + await asyncio.sleep(0.1) + task = loop.create_task(multiplexer.load_model("2")) + await asyncio.sleep(0.1) + assert "2" in multiplexer._get_loading_and_loaded_model_ids() + multiplexer._push_multiplexed_replica_info = False + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert "2" not in multiplexer._get_loading_and_loaded_model_ids() + assert multiplexer._push_multiplexed_replica_info + + if cancel_while == "waiting_for_lock": + release.set() + assert await holder == "1" + assert multiplexer._get_loading_and_loaded_model_ids() == ["1"] + class TestBasicAPI: def test_decorator_validation(self): diff --git a/python/ray/serve/tests/unit/test_deployment_state.py b/python/ray/serve/tests/unit/test_deployment_state.py index 9135f807baf6..803bf2d6c7eb 100644 --- a/python/ray/serve/tests/unit/test_deployment_state.py +++ b/python/ray/serve/tests/unit/test_deployment_state.py @@ -1091,7 +1091,7 @@ def test_record_multiplexed_model_ids_updates_replica_details( replica_id = ds._replicas.get()[0].replica_id (details,) = ds.list_replica_details() - assert details.multiplexed_model_ids == [] + assert details.multiplexed_model_ids == () # Not reported yet, so it is omitted from the user-facing details. assert "multiplexed_model_ids" not in details.model_fields_set @@ -1102,21 +1102,23 @@ def test_record_multiplexed_model_ids_updates_replica_details( ) (details,) = ds.list_replica_details() assert details.replica_id == replica_id.unique_id - assert details.multiplexed_model_ids == ["model-a", "model-b"] + assert details.multiplexed_model_ids == ("model-a", "model-b") + # The frozen model stays hashable. + hash(details) # An evicted model disappears from the details on the next report. ds.record_request_routing_info( RequestRoutingInfo(replica_id=replica_id, multiplexed_model_ids=["model-c"]) ) (details,) = ds.list_replica_details() - assert details.multiplexed_model_ids == ["model-c"] + assert details.multiplexed_model_ids == ("model-c",) # Routing stats alone (no model IDs) leave the recorded IDs unchanged. ds.record_request_routing_info( RequestRoutingInfo(replica_id=replica_id, routing_stats={"key": 1}) ) (details,) = ds.list_replica_details() - assert details.multiplexed_model_ids == ["model-c"] + assert details.multiplexed_model_ids == ("model-c",) def test_force_kill(mock_deployment_state_manager): From 9df5804f1c354b6272a9e27a7cdd5a3cc394e099 Mon Sep 17 00:00:00 2001 From: Pardeep Singh Date: Mon, 26 Oct 2026 16:09:09 +0530 Subject: [PATCH 3/5] Count in-flight loads per model ID so a cancelled waiter keeps it reported Concurrent load_model calls for the same model shared one set entry, so a waiter that was cancelled or failed removed the ID while another request was still loading it. Track the number of in-flight loads per ID and only drop it, and push again, once no request for it is left. Signed-off-by: Pardeep Singh --- python/ray/serve/multiplex.py | 28 +++++++++++++-------- python/ray/serve/tests/test_multiplex.py | 32 ++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 11 deletions(-) diff --git a/python/ray/serve/multiplex.py b/python/ray/serve/multiplex.py index b42ca9d35eb6..9a39ca3df538 100644 --- a/python/ray/serve/multiplex.py +++ b/python/ray/serve/multiplex.py @@ -2,8 +2,8 @@ import inspect import logging import time -from collections import OrderedDict -from typing import Any, Callable, List, Set +from collections import Counter, OrderedDict +from typing import Any, Callable, List from ray.serve import metrics from ray.serve._private.common import ReplicaID, RequestRoutingInfo @@ -109,11 +109,11 @@ def __init__( # Model cache lock to ensure that only one model is loading/unloading at a time. self._model_cache_lock = asyncio.Lock() - # The set of model IDs that are being loaded. This is used to early push - # model ids info to the controller. The tasks will be added when there is cache - # miss, and will be removed when the model is loaded successfully or - # failed to load. - self._model_load_tasks: Set[str] = set() + # The number of in-flight load requests per model ID. This is used to early + # push model ids info to the controller. A request is counted when there is + # a cache miss, and uncounted when it finishes, fails or is cancelled; the ID + # is dropped once no request for it is left. + self._model_load_tasks: Counter[str] = Counter() self.metrics_pusher = MetricsPusher() self.metrics_pusher.register_or_update_task( @@ -196,7 +196,7 @@ async def load_model(self, model_id: str) -> Any: # id info to the controller/router early, so that requests can be routed to # the replica. self._push_multiplexed_replica_info = True - self._model_load_tasks.add(model_id) + self._model_load_tasks[model_id] += 1 try: async with self._model_cache_lock: # Check if the model has been loaded by another request. @@ -238,9 +238,15 @@ async def load_model(self, model_id: str) -> Any: raise e finally: # The ID was pushed while loading. If the load failed or was - # cancelled, push again so the controller drops it. - self._model_load_tasks.discard(model_id) - if model_id not in self.models: + # cancelled and no other request is loading it, push again so + # the controller drops it. + self._model_load_tasks[model_id] -= 1 + if self._model_load_tasks[model_id] <= 0: + del self._model_load_tasks[model_id] + if ( + model_id not in self.models + and model_id not in self._model_load_tasks + ): self._push_multiplexed_replica_info = True async def unload_model_lru(self) -> None: diff --git a/python/ray/serve/tests/test_multiplex.py b/python/ray/serve/tests/test_multiplex.py index d345163a5817..ea2ddd086e82 100644 --- a/python/ray/serve/tests/test_multiplex.py +++ b/python/ray/serve/tests/test_multiplex.py @@ -289,6 +289,38 @@ async def model_load_func(model_id: str): assert await holder == "1" assert multiplexer._get_loading_and_loaded_model_ids() == ["1"] + async def test_cancelled_waiter_keeps_in_flight_model_id( + self, start_serve_with_context + ): + """Cancelling one of two loads of the same model keeps its ID reported.""" + release = asyncio.Event() + + async def model_load_func(model_id: str): + await release.wait() + return model_id + + multiplexer = _ModelMultiplexWrapper( + model_load_func, None, max_num_models_per_replica=1 + ) + await multiplexer.metrics_pusher.graceful_shutdown() + + loop = get_or_create_event_loop() + # The first request loads "1" while the second waits for the lock. + loader = loop.create_task(multiplexer.load_model("1")) + await asyncio.sleep(0.1) + waiter = loop.create_task(multiplexer.load_model("1")) + await asyncio.sleep(0.1) + + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + # The first request is still loading "1". + assert multiplexer._get_loading_and_loaded_model_ids() == ["1"] + + release.set() + assert await loader == "1" + assert multiplexer._get_loading_and_loaded_model_ids() == ["1"] + class TestBasicAPI: def test_decorator_validation(self): From b3e9698616ba67854b439b6f07ecae53710791aa Mon Sep 17 00:00:00 2001 From: Pardeep Singh Date: Tue, 27 Oct 2026 11:03:20 +0530 Subject: [PATCH 4/5] Push after an eviction even if the unload fails, and pin the push flag in tests unload_model_lru sets the push flag right after the evicted model leaves the cache, so an unload that fails or is cancelled while another request for the incoming model is queued still re-pushes. The cancelled-waiter test now checks no re-push is flagged, and test_collect_model_ids uses a Counter. Signed-off-by: Pardeep Singh --- python/ray/serve/multiplex.py | 4 +- python/ray/serve/tests/test_multiplex.py | 53 +++++++++++++++++++++++- 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/python/ray/serve/multiplex.py b/python/ray/serve/multiplex.py index 9a39ca3df538..1f6399b98b9a 100644 --- a/python/ray/serve/multiplex.py +++ b/python/ray/serve/multiplex.py @@ -212,7 +212,6 @@ async def load_model(self, model_id: str) -> Any: ): # Unload the least recently used model. await self.unload_model_lru() - self._push_multiplexed_replica_info = True # Load the model. logger.debug(f"Loading model '{model_id}'.") @@ -255,6 +254,9 @@ async def unload_model_lru(self) -> None: self.models_unload_counter.inc() unload_start_time = time.time() model_id, model = self.models.popitem(last=False) + # The model left the cache, so push even if its unload fails or is + # cancelled. + self._push_multiplexed_replica_info = True logger.debug(f"Unloading model '{model_id}'.") # If the model has __del__ attribute, call it. diff --git a/python/ray/serve/tests/test_multiplex.py b/python/ray/serve/tests/test_multiplex.py index ea2ddd086e82..1aa3231eb6fe 100644 --- a/python/ray/serve/tests/test_multiplex.py +++ b/python/ray/serve/tests/test_multiplex.py @@ -1,5 +1,6 @@ import asyncio import os +from collections import Counter from typing import List, Set import httpx @@ -75,7 +76,7 @@ async def test_collect_model_ids(self): multiplexer = _ModelMultiplexWrapper(None, None, max_num_models_per_replica=1) multiplexer.models = {"1": "1", "2": "2"} assert sorted(multiplexer._get_loading_and_loaded_model_ids()) == ["1", "2"] - multiplexer._model_load_tasks = {"3"} + multiplexer._model_load_tasks = Counter({"3": 1}) assert sorted(multiplexer._get_loading_and_loaded_model_ids()) == [ "1", "2", @@ -310,16 +311,64 @@ async def model_load_func(model_id: str): await asyncio.sleep(0.1) waiter = loop.create_task(multiplexer.load_model("1")) await asyncio.sleep(0.1) + # The loading ID has been pushed. + multiplexer._push_multiplexed_replica_info = False waiter.cancel() with pytest.raises(asyncio.CancelledError): await waiter - # The first request is still loading "1". + # The first request is still loading "1", so the pushed IDs are unchanged. assert multiplexer._get_loading_and_loaded_model_ids() == ["1"] + assert not multiplexer._push_multiplexed_replica_info release.set() assert await loader == "1" assert multiplexer._get_loading_and_loaded_model_ids() == ["1"] + assert not multiplexer._push_multiplexed_replica_info + + async def test_failed_eviction_pushes_with_queued_request( + self, start_serve_with_context + ): + """An evicted model is re-pushed even if its unload fails while another + request for the incoming model is queued.""" + + class MyModel: + def __init__(self, model_id): + self.model_id = model_id + self.fail_unload = model_id == "1" + + def __del__(self): + # Fail only the eviction, not the later garbage collection. + if self.fail_unload: + self.fail_unload = False + raise RuntimeError("unload failed") + + async def model_load_func(model_id: str) -> MyModel: + return MyModel(model_id) + + multiplexer = _ModelMultiplexWrapper( + model_load_func, None, max_num_models_per_replica=1 + ) + await multiplexer.metrics_pusher.graceful_shutdown() + await multiplexer.load_model("1") + + # Queue two requests for "2" behind the cache lock. + loop = get_or_create_event_loop() + await multiplexer._model_cache_lock.acquire() + evicting = loop.create_task(multiplexer.load_model("2")) + queued = loop.create_task(multiplexer.load_model("2")) + await asyncio.sleep(0.1) + # The loading ID has been pushed. + multiplexer._push_multiplexed_replica_info = False + multiplexer._model_cache_lock.release() + + # The first request evicts "1", whose unload fails. + with pytest.raises(RuntimeError, match="unload failed"): + await evicting + assert (await queued).model_id == "2" + assert multiplexer._get_loading_and_loaded_model_ids() == ["2"] + # "1" left the cache, so it must be pushed again. + assert multiplexer._push_multiplexed_replica_info class TestBasicAPI: From 308a98311d1f6ec0d6f239ef5ae24d5228e31f34 Mon Sep 17 00:00:00 2001 From: Pardeep Singh Date: Wed, 28 Oct 2026 13:16:58 +0530 Subject: [PATCH 5/5] Stop pushing model IDs before unloading models on shutdown Unloading in shutdown() sets the push flag while the metrics pusher keeps running, so a stopping replica could push its model IDs after the controller removed it. Stop the pusher first. Signed-off-by: Pardeep Singh --- python/ray/serve/multiplex.py | 3 ++ python/ray/serve/tests/test_multiplex.py | 39 +++++++++++++++++++++++- 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/python/ray/serve/multiplex.py b/python/ray/serve/multiplex.py index 1f6399b98b9a..2430e6dfa7f3 100644 --- a/python/ray/serve/multiplex.py +++ b/python/ray/serve/multiplex.py @@ -156,6 +156,9 @@ def _push_model_ids_info(self): async def shutdown(self): """Unload all the models when the model multiplexer is deleted.""" + # Stop pushing first, so unloading doesn't push model IDs for a + # replica the controller has already removed. + await self.metrics_pusher.graceful_shutdown() while len(self.models) > 0: try: await self.unload_model_lru() diff --git a/python/ray/serve/tests/test_multiplex.py b/python/ray/serve/tests/test_multiplex.py index 1aa3231eb6fe..7ec165162de0 100644 --- a/python/ray/serve/tests/test_multiplex.py +++ b/python/ray/serve/tests/test_multiplex.py @@ -8,7 +8,11 @@ import ray from ray import serve -from ray._common.test_utils import SignalActor, wait_for_condition +from ray._common.test_utils import ( + SignalActor, + async_wait_for_condition, + wait_for_condition, +) from ray._common.utils import get_or_create_event_loop from ray.serve._private.common import DeploymentID, ReplicaID from ray.serve._private.config import DeploymentConfig @@ -370,6 +374,39 @@ async def model_load_func(model_id: str) -> MyModel: # "1" left the cache, so it must be pushed again. assert multiplexer._push_multiplexed_replica_info + async def test_shutdown_stops_pushing_model_ids( + self, start_serve_with_context, monkeypatch + ): + """Unloading models in shutdown doesn't push after the replica is gone.""" + pushed = [] + + class FakeClient: + def record_request_routing_info(self, info): + pushed.append(info.multiplexed_model_ids) + + monkeypatch.setattr( + "ray.serve.multiplex._get_global_client", lambda: FakeClient() + ) + monkeypatch.setattr( + "ray.serve.multiplex.PUSH_MULTIPLEXED_MODEL_IDS_INTERVAL_S", 0.01 + ) + + async def model_load_func(model_id: str): + return model_id + + multiplexer = _ModelMultiplexWrapper( + model_load_func, None, max_num_models_per_replica=2 + ) + await multiplexer.load_model("1") + await multiplexer.load_model("2") + await async_wait_for_condition(lambda: ["1", "2"] in map(sorted, pushed)) + + pushed.clear() + await multiplexer.shutdown() + await asyncio.sleep(0.1) + assert not multiplexer.models + assert pushed == [] + class TestBasicAPI: def test_decorator_validation(self):