wehub-resource-sync
263 lines
11 KiB
Python
263 lines
11 KiB
Python
import asyncio
|
|
import inspect
|
|
import logging
|
|
import time
|
|
from collections import OrderedDict
|
|
from typing import Any, Callable, List, Set
|
|
|
|
from ray.serve import metrics
|
|
from ray.serve._private.common import ReplicaID, RequestRoutingInfo
|
|
from ray.serve._private.constants import (
|
|
MODEL_LOAD_LATENCY_BUCKETS_MS,
|
|
PUSH_MULTIPLEXED_MODEL_IDS_INTERVAL_S,
|
|
SERVE_LOGGER_NAME,
|
|
)
|
|
from ray.serve._private.metrics_utils import MetricsPusher
|
|
from ray.serve._private.usage import ServeUsageTag
|
|
from ray.serve.context import _get_global_client, _get_internal_replica_context
|
|
|
|
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
|
|
|
|
|
class _ModelMultiplexWrapper:
|
|
"""A wrapper class that wraps the model load function and
|
|
provides the LRU caching functionality.
|
|
|
|
The model multiplexer is a wrapper class that wraps the model load function
|
|
and provides the LRU caching functionality, and the model load function should
|
|
be a coroutine function that takes the model ID as the first argument and
|
|
returns the user-constructed model object.
|
|
The model multiplexer will also ensure that the number of models on the current
|
|
replica does not exceed the specified limit.
|
|
The model will be unloaded in the LRU order, the model multiplexer will call the
|
|
model's __del__ attribute if it exists to clean up the model resources eagerly.
|
|
|
|
"""
|
|
|
|
_PUSH_MULTIPLEXED_MODEL_IDS_TASK_NAME = "push_multiplexed_model_ids"
|
|
|
|
def __init__(
|
|
self,
|
|
model_load_func: Callable[[str], Any],
|
|
self_arg: Any,
|
|
max_num_models_per_replica: int,
|
|
):
|
|
"""Initialize the model multiplexer.
|
|
Args:
|
|
model_load_func: the model load async function.
|
|
self_arg: self argument when model_load_func is class method.
|
|
max_num_models_per_replica: the maximum number of models to be loaded on the
|
|
current replica. If it is -1, there is no limit for the number of models
|
|
per replica.
|
|
"""
|
|
|
|
ServeUsageTag.MULTIPLEXED_API_USED.record("1")
|
|
|
|
self.models = OrderedDict()
|
|
self._func: Callable = model_load_func
|
|
self.self_arg: Any = self_arg
|
|
self.max_num_models_per_replica: int = max_num_models_per_replica
|
|
|
|
# log MODEL_LOAD_LATENCY_BUCKET_MS
|
|
logger.debug(f"MODEL_LOAD_LATENCY_BUCKET_MS: {MODEL_LOAD_LATENCY_BUCKETS_MS}")
|
|
self.model_load_latency_ms = metrics.Histogram(
|
|
"serve_multiplexed_model_load_latency_ms",
|
|
description="The time it takes to load a model.",
|
|
boundaries=MODEL_LOAD_LATENCY_BUCKETS_MS,
|
|
)
|
|
self.model_unload_latency_ms = metrics.Histogram(
|
|
"serve_multiplexed_model_unload_latency_ms",
|
|
description="The time it takes to unload a model.",
|
|
boundaries=MODEL_LOAD_LATENCY_BUCKETS_MS,
|
|
)
|
|
self.num_models_gauge = metrics.Gauge(
|
|
"serve_num_multiplexed_models",
|
|
description="The number of models loaded on the current replica.",
|
|
)
|
|
|
|
self.registered_model_gauge = metrics.Gauge(
|
|
"serve_registered_multiplexed_model_id",
|
|
description="The model id registered on the current replica.",
|
|
tag_keys=("model_id",),
|
|
)
|
|
self.get_model_requests_counter = metrics.Counter(
|
|
"serve_multiplexed_get_model_requests_counter",
|
|
description="The counter for get model requests on the current replica.",
|
|
)
|
|
self.models_unload_counter = metrics.Counter(
|
|
"serve_multiplexed_models_unload_counter",
|
|
description="The counter for unloaded models on the current replica.",
|
|
)
|
|
self.models_load_counter = metrics.Counter(
|
|
"serve_multiplexed_models_load_counter",
|
|
description="The counter for loaded models on the current replica.",
|
|
)
|
|
|
|
context = _get_internal_replica_context()
|
|
if context is None:
|
|
raise RuntimeError(
|
|
"`@serve.multiplex` can only be used within a deployment "
|
|
"(failed to retrieve Serve replica context)."
|
|
)
|
|
|
|
self._app_name: str = context.app_name
|
|
self._deployment_name: str = context.deployment
|
|
self._replica_id: ReplicaID = context.replica_id
|
|
|
|
# Whether to push the multiplexed replica info to the controller.
|
|
self._push_multiplexed_replica_info: bool = False
|
|
|
|
# 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()
|
|
|
|
self.metrics_pusher = MetricsPusher()
|
|
self.metrics_pusher.register_or_update_task(
|
|
self._PUSH_MULTIPLEXED_MODEL_IDS_TASK_NAME,
|
|
self._push_model_ids_info,
|
|
PUSH_MULTIPLEXED_MODEL_IDS_INTERVAL_S,
|
|
)
|
|
self.metrics_pusher.start()
|
|
|
|
def _get_loading_and_loaded_model_ids(self) -> List[str]:
|
|
"""Get the model IDs of the loaded models & loading models in the replica.
|
|
This is to push the model id information early to the controller, so that
|
|
requests can be routed to the replica.
|
|
"""
|
|
models_list = set(self.models.keys())
|
|
models_list.update(self._model_load_tasks)
|
|
return list(models_list)
|
|
|
|
def _push_model_ids_info(self):
|
|
"""Push the multiplexed replica info to the controller."""
|
|
try:
|
|
self.num_models_gauge.set(len(self.models))
|
|
|
|
for model_id in self.models:
|
|
self.registered_model_gauge.set(1, tags={"model_id": model_id})
|
|
|
|
if self._push_multiplexed_replica_info:
|
|
_get_global_client().record_request_routing_info(
|
|
RequestRoutingInfo(
|
|
replica_id=self._replica_id,
|
|
multiplexed_model_ids=self._get_loading_and_loaded_model_ids(),
|
|
)
|
|
)
|
|
self._push_multiplexed_replica_info = False
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Failed to push the multiplexed replica info "
|
|
f"to the controller. Error: {e}"
|
|
)
|
|
|
|
async def shutdown(self):
|
|
"""Unload all the models when the model multiplexer is deleted."""
|
|
while len(self.models) > 0:
|
|
try:
|
|
await self.unload_model_lru()
|
|
except Exception as e:
|
|
logger.exception(
|
|
f"Failed to unload model. Error: {e}",
|
|
)
|
|
|
|
async def load_model(self, model_id: str) -> Any:
|
|
"""Load the model if it is not loaded yet, and return
|
|
the user-constructed model object.
|
|
|
|
Args:
|
|
model_id: the model ID.
|
|
|
|
Returns:
|
|
The user-constructed model object.
|
|
"""
|
|
|
|
if not isinstance(model_id, str):
|
|
raise TypeError("The model ID must be a string.")
|
|
|
|
if not model_id:
|
|
raise ValueError("The model ID cannot be empty.")
|
|
|
|
self.get_model_requests_counter.inc()
|
|
|
|
if model_id in self.models:
|
|
# Move the model to the end of the OrderedDict to mark it as
|
|
# most-recently-used. Using move_to_end() instead of pop()+reinsert
|
|
# avoids a race condition where concurrent coroutines could see the
|
|
# key as missing during the brief window between pop and reinsert.
|
|
self.models.move_to_end(model_id)
|
|
return self.models[model_id]
|
|
else:
|
|
# Set the flag to push the multiplexed replica info to the controller
|
|
# before loading the model. This is to make sure we can push the model
|
|
# 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
|
|
)
|
|
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
|
|
|
|
async def unload_model_lru(self) -> None:
|
|
"""Unload the least recently used model."""
|
|
|
|
self.models_unload_counter.inc()
|
|
unload_start_time = time.time()
|
|
model_id, model = self.models.popitem(last=False)
|
|
logger.debug(f"Unloading model '{model_id}'.")
|
|
|
|
# If the model has __del__ attribute, call it.
|
|
# This is to clean up the model resources eagerly.
|
|
if hasattr(model, "__del__"):
|
|
if not inspect.iscoroutinefunction(model.__del__):
|
|
await asyncio.get_running_loop().run_in_executor(None, model.__del__)
|
|
else:
|
|
await model.__del__()
|
|
model.__del__ = lambda _: None
|
|
unload_latency_ms = (time.time() - unload_start_time) * 1000.0
|
|
self.model_unload_latency_ms.observe(unload_latency_ms)
|
|
logger.debug(
|
|
f"Successfully unloaded model '{model_id}' in {unload_latency_ms:.1f}ms."
|
|
)
|
|
self.registered_model_gauge.set(0, tags={"model_id": model_id})
|