diff --git a/python/ray/serve/_private/deployment_state.py b/python/ray/serve/_private/deployment_state.py index c2b59091fd3c..98761b32727a 100644 --- a/python/ray/serve/_private/deployment_state.py +++ b/python/ray/serve/_private/deployment_state.py @@ -2145,8 +2145,15 @@ 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=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..2430e6dfa7f3 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( @@ -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() @@ -196,47 +199,57 @@ 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) - 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 + self._model_load_tasks[model_id] += 1 + 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() + + # 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 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: """Unload the least recently used model.""" @@ -244,6 +257,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/schema.py b/python/ray/serve/schema.py index cb9eb9105f62..44cd54f09c30 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 ( @@ -1423,6 +1423,15 @@ class ReplicaDetails(ServeActorDetails): "state from the running replica actor." ) ) + 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. After a controller " + "restart it stays empty until the replica's model cache next changes." + ), + ) version: Optional[str] = Field( default=None, description=( diff --git a/python/ray/serve/tests/test_multiplex.py b/python/ray/serve/tests/test_multiplex.py index 788c884beff5..7ec165162de0 100644 --- a/python/ray/serve/tests/test_multiplex.py +++ b/python/ray/serve/tests/test_multiplex.py @@ -1,21 +1,30 @@ import asyncio import os -from typing import List +from collections import Counter +from typing import List, Set import httpx import pytest 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 -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: @@ -71,7 +80,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", @@ -221,6 +230,183 @@ 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"] + + 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) + # 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", 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 + + 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): @@ -322,7 +508,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 +521,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 bde842429223..533c14a3ce10 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,52 @@ 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") + # 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",) + + # 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()