chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,111 @@
|
||||
import asyncio
|
||||
from typing import AsyncGenerator, Generic, Iterable, List, Optional, TypeVar
|
||||
|
||||
from ray.llm._internal.serve.constants import (
|
||||
MODEL_RESPONSE_BATCH_TIMEOUT_MS,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class Batcher(Generic[T]):
|
||||
"""This class batches multiple responses from a generator into a list of
|
||||
single responses, at some time interval.
|
||||
|
||||
Args:
|
||||
generator: the async generator that this class pulls responses
|
||||
from.
|
||||
interval_ms: the interval at which this class yields the current batch.
|
||||
If None, this class will batch all responses from the generator
|
||||
together and yield the entire batch once.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
generator: AsyncGenerator[T, None],
|
||||
interval_ms: Optional[float] = MODEL_RESPONSE_BATCH_TIMEOUT_MS,
|
||||
):
|
||||
self.generator = generator
|
||||
self.queue: asyncio.Queue = asyncio.Queue()
|
||||
|
||||
if interval_ms is None:
|
||||
self.interval_s = None
|
||||
else:
|
||||
self.interval_s = interval_ms / 1000
|
||||
|
||||
if interval_ms == 0:
|
||||
return
|
||||
|
||||
self.done_event: asyncio.Event = asyncio.Event()
|
||||
|
||||
# We are okay with this task getting cancelled (to propagate cancellations)
|
||||
self.read_task = asyncio.create_task(self.read())
|
||||
|
||||
def _merge_results(self, results: List[T]) -> Iterable[T]:
|
||||
return results
|
||||
|
||||
async def stream(self) -> AsyncGenerator[Iterable[T], None]:
|
||||
"""Drain from the queue every interval_ms and yield the merged results"""
|
||||
|
||||
if self.interval_s == 0:
|
||||
async for item in self.generator:
|
||||
yield [item]
|
||||
|
||||
return
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Wait for the interval or until we finish, whichever is faster.
|
||||
# We use an event to avoid asyncio.wait_for cancelling the real task on timeout.
|
||||
try:
|
||||
if self.interval_s is None:
|
||||
await self.done_event.wait()
|
||||
else:
|
||||
await asyncio.wait_for(
|
||||
self.done_event.wait(), timeout=self.interval_s
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
# Get all elements from the queue
|
||||
results, is_done = self.check_done_and_drain()
|
||||
|
||||
# If there are results, merge and yield them
|
||||
if results:
|
||||
output = self._merge_results(results)
|
||||
yield output
|
||||
|
||||
# If the read task is done, exit the stream task
|
||||
if is_done:
|
||||
# Raise exception, if any
|
||||
self.read_task.result()
|
||||
break
|
||||
finally:
|
||||
# If the stream task is done, make sure to exit the read task
|
||||
if not self.read_task.done():
|
||||
self.read_task.cancel()
|
||||
|
||||
def check_done_and_drain(self):
|
||||
results = self.drain_queue()
|
||||
return results, self.read_task.done()
|
||||
|
||||
async def read(self):
|
||||
"""Read from the generator and put into the queue in a tight loop"""
|
||||
try:
|
||||
async for x in self.generator:
|
||||
self.queue.put_nowait(x)
|
||||
finally:
|
||||
self.done_event.set()
|
||||
|
||||
def drain_queue(self):
|
||||
"""Drain all results currently in the queue"""
|
||||
results = []
|
||||
try:
|
||||
while True:
|
||||
results.append(self.queue.get_nowait())
|
||||
except asyncio.QueueEmpty:
|
||||
pass
|
||||
return results
|
||||
@@ -0,0 +1,153 @@
|
||||
import pickle
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
import ray
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
from ray.serve._private.common import RequestMetadata
|
||||
from ray.serve.handle import DeploymentHandle
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Timeout in seconds for waiting for deployment replicas to be populated
|
||||
BROADCAST_REPLICA_POPULATION_TIMEOUT_S = 30
|
||||
|
||||
|
||||
def broadcast(
|
||||
handle: DeploymentHandle,
|
||||
method_name: str,
|
||||
args: Union[Any, Callable[[Any], Any]] = None,
|
||||
kwargs: Union[Dict[str, Any], Callable[[Any], Dict[str, Any]]] = None,
|
||||
combine: Optional[Callable[[List[Any]], Any]] = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Broadcasts a method call to all replicas of the given handle.
|
||||
|
||||
This is useful for broadcasting a control plane message such as kv-cache
|
||||
reset or weight update to all replicas of the given handle.
|
||||
|
||||
NOTE: This API is experimental and may later be promoted to a public API in
|
||||
Ray Serve directly. For now, it is only available in Ray LLM and is
|
||||
intended to enable control plane operations during RL training which is
|
||||
required when orchestrating trianing and inference loops.
|
||||
|
||||
Args:
|
||||
handle: The DeploymentHandle to broadcast to.
|
||||
method_name: The name of the method to call on the deployment.
|
||||
args: The arguments to pass to the method. Can be a list/tuple of args,
|
||||
or a callable that takes the replica object and returns args.
|
||||
kwargs: The keyword arguments to pass to the method. Can be a dict,
|
||||
or a callable that takes the replica object and returns kwargs.
|
||||
combine: An optional callable that takes the list of results from all
|
||||
replicas and returns an aggregated result. If not provided, returns
|
||||
the list of results. The default combine function is to return the
|
||||
list of results.
|
||||
|
||||
Returns:
|
||||
The result of the method call to all replicas. If combine is provided,
|
||||
returns the aggregated result. Otherwise, returns the list of results.
|
||||
"""
|
||||
if args is None:
|
||||
args = ()
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
if not handle.is_initialized:
|
||||
# If the handle is not initialized, we initialize it here.
|
||||
# We enforce running the router in a separate loop to ensure it can
|
||||
# update its replica set asynchronously while we might be blocking or
|
||||
# waiting.
|
||||
handle._init(_run_router_in_separate_loop=True)
|
||||
|
||||
router = handle._router
|
||||
if router is None:
|
||||
raise RuntimeError("DeploymentHandle router is None.")
|
||||
|
||||
# Wait for both the replica set AND the request router to be populated.
|
||||
# `running_replicas_populated()` flips when DEPLOYMENT_TARGETS long-poll
|
||||
# arrives; `request_router` becomes non-None only after DEPLOYMENT_CONFIG
|
||||
# long-poll arrives and sets `_request_router_class`. These are independent
|
||||
# long-polls, so polling only the former races with the latter.
|
||||
#
|
||||
# In normal request flow this is hidden because `assign_request` awaits
|
||||
# `_request_router_initialized` before routing — but `broadcast()` bypasses
|
||||
# `assign_request` and pokes `_replica_id_set` directly, so it has to
|
||||
# synchronize itself.
|
||||
def _get_request_router():
|
||||
if hasattr(router, "_asyncio_router"):
|
||||
return router._asyncio_router.request_router
|
||||
if hasattr(router, "request_router"):
|
||||
return router.request_router
|
||||
return None
|
||||
|
||||
start_time = time.time()
|
||||
while not handle.running_replicas_populated() or _get_request_router() is None:
|
||||
if time.time() - start_time > BROADCAST_REPLICA_POPULATION_TIMEOUT_S:
|
||||
raise TimeoutError(
|
||||
"Timed out waiting for deployment router/replicas to initialize."
|
||||
)
|
||||
time.sleep(0.1)
|
||||
|
||||
request_router = _get_request_router()
|
||||
|
||||
replica_set = request_router._replica_id_set
|
||||
|
||||
# Execute calls
|
||||
futures = []
|
||||
|
||||
# We copy the set to avoid modification during iteration if that happens
|
||||
replicas = list(replica_set)
|
||||
|
||||
for replica in replicas:
|
||||
actor_name = replica.to_full_id_str()
|
||||
try:
|
||||
actor_handle = ray.get_actor(actor_name, namespace="serve")
|
||||
except ValueError:
|
||||
# Actor might be dead or not found
|
||||
continue
|
||||
|
||||
# Prepare args
|
||||
call_args = args
|
||||
call_kwargs = kwargs
|
||||
|
||||
if callable(args):
|
||||
call_args = args(replica)
|
||||
if callable(kwargs):
|
||||
call_kwargs = kwargs(replica)
|
||||
|
||||
if not isinstance(call_args, (list, tuple)):
|
||||
raise ValueError(f"args must be a list or tuple, got {type(call_args)}")
|
||||
|
||||
if not isinstance(call_kwargs, dict):
|
||||
# Fallback if callable returned something else or initial was not dict
|
||||
# But initial default is dict.
|
||||
if call_kwargs is None:
|
||||
call_kwargs = {}
|
||||
else:
|
||||
raise ValueError(f"kwargs must be a dict, got {type(call_kwargs)}")
|
||||
|
||||
# Prepare Metadata
|
||||
request_id = f"broadcast-{uuid.uuid4()}"
|
||||
dummy_rm = RequestMetadata(
|
||||
request_id=request_id,
|
||||
internal_request_id=request_id,
|
||||
call_method=method_name,
|
||||
)
|
||||
|
||||
pickled_rm = pickle.dumps(dummy_rm)
|
||||
|
||||
# Fire remote call
|
||||
# We collect futures to wait for them
|
||||
futures.append(
|
||||
actor_handle.handle_request.remote(pickled_rm, *call_args, **call_kwargs)
|
||||
)
|
||||
|
||||
# Wait for all calls to complete
|
||||
results = []
|
||||
if futures:
|
||||
results = ray.get(futures)
|
||||
|
||||
if combine:
|
||||
return combine(results)
|
||||
return results
|
||||
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
Serve-specific LoRA utilities that use generic abstractions from lora_utils.py.
|
||||
|
||||
This module provides serve-specific functionality while using the generic
|
||||
LoRA abstractions from common/lora_utils.py. This ensures clean separation
|
||||
between generic and serve-specific concerns.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from ray.llm._internal.common.constants import LORA_ADAPTER_CONFIG_NAME
|
||||
from ray.llm._internal.common.models import global_id_manager, make_async
|
||||
from ray.llm._internal.common.utils.cloud_utils import (
|
||||
LoraMirrorConfig,
|
||||
)
|
||||
from ray.llm._internal.common.utils.lora_utils import (
|
||||
CLOUD_OBJECT_MISSING,
|
||||
clean_model_id,
|
||||
clear_directory,
|
||||
get_base_model_id,
|
||||
get_lora_id,
|
||||
get_object_from_cloud,
|
||||
retry_with_exponential_backoff,
|
||||
sync_files_with_lock,
|
||||
)
|
||||
from ray.llm._internal.serve.core.configs.llm_config import (
|
||||
DiskMultiplexConfig,
|
||||
LLMConfig,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def get_lora_finetuned_context_length(bucket_uri: str) -> Optional[int]:
|
||||
"""Gets the sequence length used to tune the LoRA adapter.
|
||||
|
||||
Return: Returns the max sequence length for the adapter, if it exists.
|
||||
|
||||
Raises: HTTPException if the LoRA adapter config file isn't available
|
||||
in the cloud storage repository.
|
||||
"""
|
||||
if bucket_uri.endswith("/"):
|
||||
bucket_uri = bucket_uri.rstrip("/")
|
||||
object_uri = f"{bucket_uri}/{LORA_ADAPTER_CONFIG_NAME}"
|
||||
|
||||
object_str_or_missing_message = await get_object_from_cloud(object_uri)
|
||||
|
||||
if object_str_or_missing_message is CLOUD_OBJECT_MISSING:
|
||||
raise HTTPException(
|
||||
404,
|
||||
f"Unable to find LoRA adapter config file "
|
||||
f'"{LORA_ADAPTER_CONFIG_NAME}" in folder {bucket_uri}. '
|
||||
"Check that the file exists and that you have read permissions.",
|
||||
)
|
||||
else:
|
||||
adapter_config_str = object_str_or_missing_message
|
||||
adapter_config = json.loads(adapter_config_str)
|
||||
return adapter_config.get("max_length")
|
||||
|
||||
|
||||
async def download_multiplex_config_info(
|
||||
model_id: str, base_path: str
|
||||
) -> tuple[str, Optional[int]]:
|
||||
"""Downloads info needed to create a multiplex config.
|
||||
|
||||
Downloads objects using cloud storage provider APIs.
|
||||
|
||||
Returns: 2-tuple containing
|
||||
1. A bucket_uri for the bucket containing LoRA weights and config.
|
||||
2. The maximum LoRA sequence length.
|
||||
|
||||
Raises: HTTPException if the LoRA adapter config file isn't available
|
||||
in the cloud storage repository.
|
||||
"""
|
||||
bucket_uri = f"{base_path}/{model_id}"
|
||||
ft_context_length = await get_lora_finetuned_context_length(bucket_uri)
|
||||
return bucket_uri, ft_context_length
|
||||
|
||||
|
||||
async def get_lora_model_metadata(model_id: str, base_path: str) -> Dict[str, Any]:
|
||||
"""Get the lora model metadata for a given model id and base path.
|
||||
|
||||
Args:
|
||||
model_id: A lora model id of the form ``base_model_id:suffix:id``.
|
||||
base_path: The cloud storage path under which LoRA adapters are stored
|
||||
(typically ``llm_config.lora_config.dynamic_lora_loading_path``).
|
||||
|
||||
Returns:
|
||||
A dict with keys ``model_id``, ``base_model_id``,
|
||||
``max_request_context_length``, and ``bucket_uri``.
|
||||
"""
|
||||
base_model_id = get_base_model_id(model_id)
|
||||
lora_id = get_lora_id(model_id)
|
||||
|
||||
# Examples of the variables:
|
||||
# model_id: "meta-llama/Meta-Llama-3.1-8B-Instruct:my_suffix:aBc1234"
|
||||
# base_path: "s3://ray-llama-weights"
|
||||
# bucket_uri: "s3://ray-llama-weights/my_suffix:aBc1234"
|
||||
(
|
||||
bucket_uri,
|
||||
ft_context_length,
|
||||
) = await download_multiplex_config_info(lora_id, base_path)
|
||||
|
||||
return {
|
||||
"model_id": model_id,
|
||||
"base_model_id": base_model_id,
|
||||
"max_request_context_length": ft_context_length,
|
||||
# Note (genesu): `bucket_uri` affects where the lora weights are downloaded
|
||||
# from remote location.
|
||||
"bucket_uri": bucket_uri,
|
||||
}
|
||||
|
||||
|
||||
async def get_lora_mirror_config(
|
||||
model_id: str,
|
||||
llm_config: LLMConfig,
|
||||
) -> LoraMirrorConfig:
|
||||
"""Get LoRA mirror configuration for serve-specific LLM config."""
|
||||
metadata = await get_lora_model_metadata(
|
||||
model_id, llm_config.lora_config.dynamic_lora_loading_path
|
||||
)
|
||||
|
||||
return LoraMirrorConfig(
|
||||
lora_model_id=model_id,
|
||||
bucket_uri=metadata["bucket_uri"],
|
||||
max_total_tokens=metadata["max_request_context_length"],
|
||||
sync_args=None,
|
||||
)
|
||||
|
||||
|
||||
class LoraModelLoader:
|
||||
"""Download LoRA weights from remote storage and manage disk cache.
|
||||
|
||||
This class is serve-specific as it depends on DiskMultiplexConfig and
|
||||
other serve-specific concepts.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
lora_root: Optional[str] = None,
|
||||
download_timeout_s: Optional[float] = None,
|
||||
max_tries: int = 1,
|
||||
):
|
||||
self.lora_root = lora_root or "/tmp/ray/llm/lora/cache"
|
||||
self.disk_cache: Dict[str, DiskMultiplexConfig] = {}
|
||||
self.active_syncing_tasks: Dict[str, asyncio.Task[DiskMultiplexConfig]] = {}
|
||||
if download_timeout_s is not None and download_timeout_s <= 0:
|
||||
raise ValueError(
|
||||
f"download_timeout_s must be None or >0, got {download_timeout_s}"
|
||||
)
|
||||
self.download_timeout_s = download_timeout_s
|
||||
if max_tries < 1:
|
||||
raise ValueError(f"max_tries must be >=1, got {max_tries}")
|
||||
self.max_tries = max_tries
|
||||
|
||||
async def load_model_from_config(
|
||||
self, lora_model_id: str, llm_config
|
||||
) -> DiskMultiplexConfig:
|
||||
"""Load a LoRA model by first fetching its mirror config from S3."""
|
||||
lora_mirror_config = await get_lora_mirror_config(lora_model_id, llm_config)
|
||||
return await self.load_model(lora_model_id, lora_mirror_config)
|
||||
|
||||
async def load_model(
|
||||
self, lora_model_id: str, lora_mirror_config: LoraMirrorConfig
|
||||
) -> DiskMultiplexConfig:
|
||||
"""Load a LoRA model."""
|
||||
if lora_model_id in self.disk_cache:
|
||||
return self.disk_cache[lora_model_id]
|
||||
|
||||
if lora_model_id not in self.active_syncing_tasks:
|
||||
task = asyncio.create_task(self._load_model_async(lora_mirror_config))
|
||||
task.add_done_callback(
|
||||
lambda result: self.active_syncing_tasks.pop(lora_model_id, None)
|
||||
)
|
||||
self.active_syncing_tasks[lora_model_id] = task
|
||||
else:
|
||||
task = self.active_syncing_tasks[lora_model_id]
|
||||
|
||||
disk_config = await asyncio.shield(task)
|
||||
self.disk_cache[lora_model_id] = disk_config
|
||||
return disk_config
|
||||
|
||||
async def _load_model_async(
|
||||
self, lora_mirror_config: LoraMirrorConfig
|
||||
) -> DiskMultiplexConfig:
|
||||
return await self._load_model(lora_mirror_config)
|
||||
|
||||
@make_async
|
||||
def _load_model(self, lora_mirror_config: LoraMirrorConfig) -> DiskMultiplexConfig:
|
||||
return self._load_model_sync(lora_mirror_config)
|
||||
|
||||
@make_async
|
||||
def clear_cache(self):
|
||||
"""Clear the disk cache."""
|
||||
clear_directory(self.lora_root)
|
||||
|
||||
def _model_dir_path(self, model_id: str) -> str:
|
||||
"""Construct the path for the lora weight."""
|
||||
lora_id = get_lora_id(clean_model_id(model_id))
|
||||
path = os.path.join(self.lora_root, lora_id)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
return path
|
||||
|
||||
def _download_lora(self, lora_mirror_config: LoraMirrorConfig) -> str:
|
||||
"""Download LoRA weights using generic download primitives."""
|
||||
model_local_path = self._model_dir_path(lora_mirror_config.lora_model_id)
|
||||
sync_files_with_lock(
|
||||
lora_mirror_config.bucket_uri,
|
||||
model_local_path,
|
||||
timeout=self.download_timeout_s,
|
||||
)
|
||||
return model_local_path
|
||||
|
||||
def _load_model_sync(
|
||||
self, lora_mirror_config: LoraMirrorConfig
|
||||
) -> DiskMultiplexConfig:
|
||||
"""Load a model from the given mirror configuration."""
|
||||
download_with_retries = retry_with_exponential_backoff(
|
||||
max_tries=self.max_tries,
|
||||
exception_to_check=Exception,
|
||||
)(lambda config: self._download_lora(config))
|
||||
|
||||
local_path = download_with_retries(lora_mirror_config)
|
||||
return DiskMultiplexConfig.model_validate(
|
||||
{
|
||||
"model_id": lora_mirror_config.lora_model_id,
|
||||
"max_total_tokens": lora_mirror_config.max_total_tokens,
|
||||
"local_path": local_path,
|
||||
"lora_assigned_int_id": global_id_manager.next(),
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,75 @@
|
||||
import asyncio
|
||||
from typing import Optional
|
||||
|
||||
import ray
|
||||
from ray.llm._internal.common.utils.download_utils import (
|
||||
download_model_files,
|
||||
)
|
||||
from ray.llm._internal.common.utils.import_utils import try_import
|
||||
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
torch = try_import("torch")
|
||||
transformers = try_import("transformers")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def initialize_remote_node(llm_config: LLMConfig) -> Optional[str]:
|
||||
|
||||
callback = llm_config.get_or_create_callback()
|
||||
engine_config = llm_config.get_engine_config()
|
||||
|
||||
local_path = download_model_files(
|
||||
model_id=engine_config.actual_hf_model_id,
|
||||
mirror_config=engine_config.mirror_config,
|
||||
download_model=callback.ctx.worker_node_download_model,
|
||||
download_extra_files=True,
|
||||
callback=callback,
|
||||
)
|
||||
|
||||
# Validate that the binary exists
|
||||
if local_path and local_path != engine_config.actual_hf_model_id:
|
||||
engine_config.hf_model_id = local_path
|
||||
|
||||
return local_path
|
||||
|
||||
|
||||
async def initialize_node(llm_config: LLMConfig):
|
||||
"""Implements node initialization for LLM engines.
|
||||
|
||||
Downloads model, tokenizer, and extra files as necessary.
|
||||
"""
|
||||
# Get callback instance (if configured) with context information
|
||||
callback = llm_config.get_or_create_callback()
|
||||
ctx = callback.ctx
|
||||
pg_table = ray.util.placement_group_table(ctx.placement_group)
|
||||
|
||||
node_set = set(pg_table["bundles_to_node_id"].values())
|
||||
download_tasks = []
|
||||
for node_id in node_set:
|
||||
node_affinity_strategy = (
|
||||
ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
|
||||
node_id=node_id,
|
||||
soft=False,
|
||||
)
|
||||
)
|
||||
download_tasks.append(
|
||||
ray.remote(initialize_remote_node).options(
|
||||
num_cpus=0,
|
||||
scheduling_strategy=node_affinity_strategy,
|
||||
runtime_env=ctx.runtime_env,
|
||||
)
|
||||
)
|
||||
|
||||
logger.info("Running tasks to download model files on worker nodes")
|
||||
paths = await asyncio.gather(
|
||||
*[download_task.remote(llm_config) for download_task in download_tasks]
|
||||
)
|
||||
|
||||
# assume that all paths are the same
|
||||
assert paths, "No paths returned from download_model_files"
|
||||
assert (
|
||||
len(set(paths)) == 1
|
||||
), "Paths returned from download_model_files are not the same"
|
||||
llm_config.get_engine_config().hf_model_id = paths[0]
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Placement group utilities for Ray LLM Serve."""
|
||||
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from ray.serve._private.utils import get_head_node_id
|
||||
from ray.util.placement_group import PlacementGroup, placement_group_table
|
||||
|
||||
|
||||
def _sort_bundle_indices_by_node(
|
||||
bundles_to_node_id: Dict[int, str],
|
||||
driver_node_id: str,
|
||||
) -> List[int]:
|
||||
"""Sort bundle indices so that same-node bundles are adjacent, driver node first.
|
||||
|
||||
Args:
|
||||
bundles_to_node_id: Mapping from bundle index to node ID.
|
||||
driver_node_id: The node ID of the driver node.
|
||||
|
||||
Returns:
|
||||
List of bundle indices sorted with driver node bundles first,
|
||||
then remaining nodes in deterministic order.
|
||||
"""
|
||||
node_to_bundles: Dict[str, List[int]] = defaultdict(list)
|
||||
# bundle_idx is already in ascending order: created sequentially during
|
||||
# placement group creation and preserved by the GCS protobuf.
|
||||
for bundle_idx, node_id in bundles_to_node_id.items():
|
||||
node_to_bundles[node_id].append(bundle_idx)
|
||||
|
||||
result: List[int] = []
|
||||
if driver_node_id in node_to_bundles:
|
||||
result.extend(node_to_bundles.pop(driver_node_id))
|
||||
for node_id in sorted(node_to_bundles.keys()):
|
||||
result.extend(node_to_bundles[node_id])
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def get_bundle_indices_sorted_by_node(
|
||||
pg: PlacementGroup,
|
||||
driver_node_id: Optional[str] = None,
|
||||
) -> List[int]:
|
||||
"""Return bundle indices sorted such that same-node bundles are adjacent, driver node first.
|
||||
|
||||
When a placement group is provisioned, adjacent bundle indices don't
|
||||
necessarily map to the same physical node. This utility reorders bundle
|
||||
indices so that bundles on the same node are grouped together.
|
||||
|
||||
The driver node's bundles come first so that global rank 0 (which hosts the
|
||||
distributed rendezvous store) is co-located with the driver.
|
||||
|
||||
Args:
|
||||
pg: A ready placement group.
|
||||
driver_node_id: Node ID whose bundles should be ordered first. Callers
|
||||
should pass the node that advertises the distributed master address
|
||||
so that global rank 0 is co-located with it. Defaults to the cluster
|
||||
head node, which may not own any of the placement group's bundles
|
||||
(e.g. a CPU-only head node in a GPU cluster), in which case rank 0
|
||||
falls back to the lexicographically-first node.
|
||||
|
||||
Returns:
|
||||
List of bundle indices sorted such that same-node bundles are adjacent, driver node first.
|
||||
"""
|
||||
table = placement_group_table(pg)
|
||||
bundles_to_node_id = table["bundles_to_node_id"]
|
||||
if driver_node_id is None:
|
||||
driver_node_id = get_head_node_id()
|
||||
return _sort_bundle_indices_by_node(bundles_to_node_id, driver_node_id)
|
||||
@@ -0,0 +1,328 @@
|
||||
"""Generic registry for LLM serving components using Ray's internal KV store.
|
||||
|
||||
This module provides a reusable registry mechanism that enables components to be
|
||||
registered in the driver process and accessed across all Ray processes in the cluster,
|
||||
including Ray Serve child processes.
|
||||
|
||||
Similar to RLlib/Tune's registry but with a fixed global prefix for cross-job access.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any, Callable
|
||||
|
||||
import ray._private.worker as worker
|
||||
import ray.cloudpickle as pickle
|
||||
from ray.experimental.internal_kv import (
|
||||
_internal_kv_del,
|
||||
_internal_kv_exists,
|
||||
_internal_kv_get,
|
||||
_internal_kv_initialized,
|
||||
_internal_kv_put,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# Fixed prefix for cross-job accessibility (Serve deployments run in different jobs)
|
||||
_SERVE_REGISTRY_PREFIX = "serve_global"
|
||||
|
||||
|
||||
def _make_key(category: str, name: str) -> bytes:
|
||||
"""Generate a binary key for the KV store.
|
||||
|
||||
Args:
|
||||
category: The component category (e.g., "kv_connector_backend")
|
||||
name: The component name
|
||||
|
||||
Returns:
|
||||
The key to use for storing the value
|
||||
"""
|
||||
return (
|
||||
b"LLMServeRegistry:"
|
||||
+ _SERVE_REGISTRY_PREFIX.encode("ascii")
|
||||
+ b":"
|
||||
+ category.encode("ascii")
|
||||
+ b"/"
|
||||
+ name.encode("ascii")
|
||||
)
|
||||
|
||||
|
||||
def _create_loader(value: Any) -> Callable[[], Any]:
|
||||
"""Create a loader callable for a value.
|
||||
|
||||
Handles both direct objects/classes and string paths for lazy loading.
|
||||
|
||||
Args:
|
||||
value: Either:
|
||||
- A class, object, or callable (returns lambda: value)
|
||||
- A string in format "module_path:class_name" (creates import loader)
|
||||
|
||||
Returns:
|
||||
A callable that returns the value when called
|
||||
|
||||
Raises:
|
||||
ValueError: If value is a string but doesn't have the correct format
|
||||
"""
|
||||
if isinstance(value, str):
|
||||
if ":" not in value:
|
||||
raise ValueError(
|
||||
f"Invalid format for string value: '{value}'. "
|
||||
f"Expected format: 'module_path:class_name' or a class/object."
|
||||
)
|
||||
module_path, class_name = value.rsplit(":", 1)
|
||||
# Create a loader callable that imports on demand
|
||||
def loader():
|
||||
module = importlib.import_module(module_path)
|
||||
return getattr(module, class_name)
|
||||
|
||||
return loader
|
||||
else:
|
||||
# For direct objects/classes, create a simple loader
|
||||
return lambda: value
|
||||
|
||||
|
||||
class ComponentRegistry:
|
||||
"""Generic registry for LLM serving components using Ray's internal KV store.
|
||||
|
||||
This registry enables components to be registered in the driver process and
|
||||
accessed across all Ray processes in the cluster, including Ray Serve child processes.
|
||||
|
||||
Similar to RLlib/Tune's registry but with a fixed global prefix for cross-job access.
|
||||
|
||||
**Usage Pattern:**
|
||||
This registry is designed for a "register once, read many" pattern:
|
||||
- Components are typically registered in the driver process before deployment
|
||||
- Ray Serve replicas read from the KV store during initialization
|
||||
- Once a component is resolved and cached in a process, subsequent `get()` calls return the cached value without checking the KV store for updates
|
||||
|
||||
Example:
|
||||
# Create a registry for a component category
|
||||
registry = ComponentRegistry("my_component")
|
||||
|
||||
# Register a component
|
||||
registry.register("my_component", MyComponentClass)
|
||||
|
||||
# Get a registered component
|
||||
component = registry.get("my_component")
|
||||
|
||||
# Check if registered
|
||||
if registry.contains("my_component"):
|
||||
...
|
||||
"""
|
||||
|
||||
def __init__(self, category: str):
|
||||
"""Initialize a registry for a specific component category.
|
||||
|
||||
Args:
|
||||
category: The category name (e.g., "kv_connector_backend")
|
||||
"""
|
||||
self.category = category
|
||||
self._loader_cache: dict[str, Callable[[], Any]] = {}
|
||||
self._resolved_cache: dict[str, Any] = {}
|
||||
self._pending: dict[str, bytes] = {}
|
||||
|
||||
def register(self, name: str, value: Any) -> None:
|
||||
"""Register a component.
|
||||
|
||||
Args:
|
||||
name: The name to register under
|
||||
value: The component to register. Can be:
|
||||
- A class, object, or callable (serialized directly)
|
||||
- A string in format "module_path:class_name" (lazy-loaded via import)
|
||||
|
||||
Raises:
|
||||
ValueError: If the component is already registered. Use unregister() first if you need to change the registration.
|
||||
|
||||
Examples:
|
||||
# Register a class directly
|
||||
registry.register("MyClass", MyClass)
|
||||
|
||||
# Register via module path (lazy loading)
|
||||
registry.register("MyClass", "my.module:MyClass")
|
||||
"""
|
||||
# Prevent double registration to avoid cache inconsistencies
|
||||
if self.contains(name):
|
||||
raise ValueError(
|
||||
f"{self.category} '{name}' is already registered. "
|
||||
f"Use unregister() first if you need to change the registration."
|
||||
)
|
||||
|
||||
# Create a loader callable (handles both direct values and string paths)
|
||||
loader = _create_loader(value)
|
||||
|
||||
# Serialize the loader callable
|
||||
serialized = pickle.dumps(loader)
|
||||
|
||||
# Store loader in cache
|
||||
self._loader_cache[name] = loader
|
||||
|
||||
# Store in KV store if Ray is initialized, otherwise queue for later
|
||||
if _internal_kv_initialized():
|
||||
try:
|
||||
key = _make_key(self.category, name)
|
||||
_internal_kv_put(key, serialized, overwrite=True)
|
||||
logger.debug(f"Registered {self.category} '{name}' in KV store")
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to register {self.category} '{name}' in KV store: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
self._pending[name] = serialized
|
||||
else:
|
||||
self._pending[name] = serialized
|
||||
|
||||
def get(self, name: str) -> Any:
|
||||
"""Get a registered component.
|
||||
|
||||
Args:
|
||||
name: The name of the component
|
||||
|
||||
Returns:
|
||||
The registered component. If registered with a string path,
|
||||
returns the imported class/object. If registered directly,
|
||||
returns the original value.
|
||||
|
||||
Raises:
|
||||
ValueError: If the component is not registered
|
||||
"""
|
||||
# Check resolved cache first.
|
||||
if name in self._resolved_cache:
|
||||
return self._resolved_cache[name]
|
||||
|
||||
loader = self._loader_cache.get(name)
|
||||
# If not in local loader cache, try fetching from KV store.
|
||||
if loader is None and _internal_kv_initialized():
|
||||
try:
|
||||
key = _make_key(self.category, name)
|
||||
serialized = _internal_kv_get(key)
|
||||
if serialized is not None:
|
||||
loader = pickle.loads(serialized)
|
||||
# Cache the loader for future gets.
|
||||
self._loader_cache[name] = loader
|
||||
logger.debug(f"Loaded {self.category} '{name}' from KV store")
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to load {self.category} '{name}' from KV store: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
if loader is not None:
|
||||
value = loader()
|
||||
self._resolved_cache[name] = value
|
||||
return value
|
||||
|
||||
# Not found
|
||||
raise ValueError(
|
||||
f"{self.category} '{name}' not found. "
|
||||
f"Registered: {list(self._loader_cache.keys())}"
|
||||
)
|
||||
|
||||
def contains(self, name: str) -> bool:
|
||||
"""Check if a component is registered.
|
||||
|
||||
Args:
|
||||
name: The name to check
|
||||
|
||||
Returns:
|
||||
True if registered, False otherwise
|
||||
"""
|
||||
if name in self._loader_cache:
|
||||
return True
|
||||
|
||||
if _internal_kv_initialized():
|
||||
try:
|
||||
key = _make_key(self.category, name)
|
||||
return _internal_kv_exists(key)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to check if {self.category} '{name}' exists in KV store: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
|
||||
return False
|
||||
|
||||
def unregister(self, name: str) -> None:
|
||||
"""Unregister a component.
|
||||
|
||||
Removes the component from local cache, pending registrations, and KV store.
|
||||
|
||||
Args:
|
||||
name: The name of the component to unregister
|
||||
"""
|
||||
# Remove from local caches
|
||||
if name in self._loader_cache:
|
||||
del self._loader_cache[name]
|
||||
if name in self._resolved_cache:
|
||||
del self._resolved_cache[name]
|
||||
|
||||
# Remove from pending if present
|
||||
if name in self._pending:
|
||||
del self._pending[name]
|
||||
|
||||
# Remove from KV store if Ray is initialized
|
||||
if _internal_kv_initialized():
|
||||
try:
|
||||
key = _make_key(self.category, name)
|
||||
_internal_kv_del(key)
|
||||
logger.debug(f"Unregistered {self.category} '{name}' from KV store")
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to unregister {self.category} '{name}' from KV store: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def flush_pending(self) -> None:
|
||||
"""Flush pending registrations to KV store.
|
||||
|
||||
This is called automatically when Ray initializes via _post_init_hooks.
|
||||
"""
|
||||
if not _internal_kv_initialized() or not self._pending:
|
||||
return
|
||||
|
||||
for name, serialized in self._pending.items():
|
||||
try:
|
||||
key = _make_key(self.category, name)
|
||||
_internal_kv_put(key, serialized, overwrite=True)
|
||||
logger.debug(
|
||||
f"Flushed pending registration for {self.category} '{name}'"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to flush {self.category} '{name}': {e}", exc_info=True
|
||||
)
|
||||
|
||||
self._pending.clear()
|
||||
|
||||
|
||||
# Global registry instances for different component categories
|
||||
_registries: dict[str, ComponentRegistry] = {}
|
||||
|
||||
|
||||
def get_registry(category: str) -> ComponentRegistry:
|
||||
"""Get or create a registry for a component category.
|
||||
|
||||
Args:
|
||||
category: The component category name
|
||||
|
||||
Returns:
|
||||
The ComponentRegistry instance for this category
|
||||
"""
|
||||
if category not in _registries:
|
||||
_registries[category] = ComponentRegistry(category)
|
||||
return _registries[category]
|
||||
|
||||
|
||||
def _flush_all_registries():
|
||||
"""Flush all pending registrations to KV store.
|
||||
|
||||
This is registered as a Ray post-init hook to ensure registrations
|
||||
made before Ray initialization are available across processes.
|
||||
"""
|
||||
for registry in _registries.values():
|
||||
registry.flush_pending()
|
||||
|
||||
|
||||
if _flush_all_registries not in worker._post_init_hooks:
|
||||
worker._post_init_hooks.append(_flush_all_registries)
|
||||
@@ -0,0 +1,218 @@
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from functools import partial
|
||||
from typing import Awaitable, Callable, TypeVar
|
||||
|
||||
from fastapi import HTTPException, status
|
||||
from httpx import HTTPStatusError as HTTPXHTTPStatusError
|
||||
from pydantic import ValidationError as PydanticValidationError
|
||||
|
||||
from ray import serve
|
||||
from ray.llm._internal.common.errors import VLLM_FATAL_ERRORS
|
||||
from ray.llm._internal.serve.constants import DEFAULT_FATAL_ERROR_COOLDOWN_S
|
||||
from ray.llm._internal.serve.core.configs.openai_api_models import (
|
||||
ErrorInfo,
|
||||
ErrorResponse,
|
||||
OpenAIHTTPException,
|
||||
)
|
||||
from ray.llm._internal.serve.observability.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def _is_fatal_engine_error(e: Exception) -> bool:
|
||||
"""Detect fatal engine errors via isinstance check."""
|
||||
if not VLLM_FATAL_ERRORS:
|
||||
return False
|
||||
return isinstance(e, VLLM_FATAL_ERRORS)
|
||||
|
||||
|
||||
class _FatalEngineErrorLogHandler:
|
||||
"""Rate limits logging for fatal engine errors.
|
||||
|
||||
- First fatal error: logged with full traceback.
|
||||
- Subsequent occurences within ``cooldown_s``: suppressed.
|
||||
- Next fatal error after ``cooldown_s``: emits a summary with suppressed errors.
|
||||
- Fatal error after ``2 * cooldown_s`` of quiet: logs full traceback again.
|
||||
- Non-fatal errors: always logged, unaffected by rate limiting.
|
||||
"""
|
||||
|
||||
def __init__(self, cooldown_s: float = DEFAULT_FATAL_ERROR_COOLDOWN_S):
|
||||
self._cooldown_s = cooldown_s
|
||||
self._first_logged = False
|
||||
self._suppressed_count = 0
|
||||
self._last_summary_time = 0.0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def log(
|
||||
self,
|
||||
e: Exception,
|
||||
request_id: str,
|
||||
status_code: int,
|
||||
) -> None:
|
||||
"""Log the error, rate limiting fatal engine errors."""
|
||||
is_fatal = _is_fatal_engine_error(e)
|
||||
|
||||
if not is_fatal:
|
||||
log_fn = logger.error if status_code >= 500 else logger.warning
|
||||
log_fn(
|
||||
f"Encountered failure while handling request {request_id}",
|
||||
exc_info=e,
|
||||
extra={"ray_serve_extra_fields": {"status_code": status_code}},
|
||||
)
|
||||
return
|
||||
|
||||
with self._lock:
|
||||
now = time.monotonic()
|
||||
|
||||
# If enough quiet time has passed, treat this as a new failure
|
||||
# event. The suppressed count is intentionally dropped since the
|
||||
# original fatal error's full traceback was already emitted.
|
||||
if (
|
||||
self._first_logged
|
||||
and (now - self._last_summary_time) >= 2 * self._cooldown_s
|
||||
):
|
||||
self._first_logged = False
|
||||
self._suppressed_count = 0
|
||||
|
||||
if not self._first_logged:
|
||||
self._first_logged = True
|
||||
self._last_summary_time = now
|
||||
logger.error(
|
||||
"Encountered failure while handling request %s",
|
||||
request_id,
|
||||
exc_info=e,
|
||||
extra={"ray_serve_extra_fields": {"status_code": status_code}},
|
||||
)
|
||||
return
|
||||
|
||||
self._suppressed_count += 1
|
||||
elapsed = now - self._last_summary_time
|
||||
if elapsed >= self._cooldown_s:
|
||||
logger.error(
|
||||
"Suppressed %d fatal engine error(s) in the last %.0fs. "
|
||||
"Engine is dead, awaiting replica restart.",
|
||||
self._suppressed_count,
|
||||
elapsed,
|
||||
)
|
||||
self._suppressed_count = 0
|
||||
self._last_summary_time = now
|
||||
|
||||
|
||||
_fatal_error_log_handler = _FatalEngineErrorLogHandler()
|
||||
|
||||
|
||||
def make_async(_func: Callable[..., T]) -> Callable[..., Awaitable[T]]:
|
||||
"""Take a blocking function, and run it on in an executor thread.
|
||||
|
||||
This function prevents the blocking function from blocking the asyncio event loop.
|
||||
The code in this function needs to be thread safe.
|
||||
"""
|
||||
|
||||
def _async_wrapper(*args, **kwargs) -> asyncio.Future:
|
||||
loop = asyncio.get_event_loop()
|
||||
func = partial(_func, *args, **kwargs)
|
||||
return loop.run_in_executor(executor=None, func=func)
|
||||
|
||||
return _async_wrapper
|
||||
|
||||
|
||||
def extract_message_from_exception(e: Exception) -> str:
|
||||
# If the exception is a Ray exception, we need to dig through the text to get just
|
||||
# the exception message without the stack trace
|
||||
# This also works for normal exceptions (we will just return everything from
|
||||
# format_exception_only in that case)
|
||||
message_lines = traceback.format_exception_only(type(e), e)[-1].strip().split("\n")
|
||||
message = ""
|
||||
# The stack trace lines will be prefixed with spaces, so we need to start from the bottom
|
||||
# and stop at the last line before a line with a space
|
||||
found_last_line_before_stack_trace = False
|
||||
for line in reversed(message_lines):
|
||||
if not line.startswith(" "):
|
||||
found_last_line_before_stack_trace = True
|
||||
if found_last_line_before_stack_trace and line.startswith(" "):
|
||||
break
|
||||
message = line + "\n" + message
|
||||
message = message.strip()
|
||||
return message
|
||||
|
||||
|
||||
def _extract_message(e):
|
||||
if isinstance(e, OpenAIHTTPException) and e.internal_message is not None:
|
||||
internal_message = e.internal_message
|
||||
else:
|
||||
internal_message = extract_message_from_exception(e)
|
||||
|
||||
if isinstance(e, HTTPException):
|
||||
message = e.detail
|
||||
elif isinstance(e, OpenAIHTTPException):
|
||||
message = e.message
|
||||
else:
|
||||
message = internal_message
|
||||
|
||||
return internal_message, message
|
||||
|
||||
|
||||
def get_response_for_error(
|
||||
e: Exception,
|
||||
request_id: str,
|
||||
) -> ErrorResponse:
|
||||
if isinstance(e, HTTPException):
|
||||
status_code = e.status_code
|
||||
elif isinstance(e, OpenAIHTTPException):
|
||||
status_code = e.status_code
|
||||
elif isinstance(e, PydanticValidationError):
|
||||
status_code = 400
|
||||
elif isinstance(e, HTTPXHTTPStatusError):
|
||||
status_code = e.response.status_code
|
||||
else:
|
||||
# Try to get the status code attribute from exception,
|
||||
# if not present, fallback to generic 500
|
||||
status_code = int(
|
||||
getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
)
|
||||
|
||||
_fatal_error_log_handler.log(e, request_id, status_code)
|
||||
|
||||
if status_code == status.HTTP_500_INTERNAL_SERVER_ERROR:
|
||||
internal_message = message = "Internal Server Error"
|
||||
exc_type = "InternalServerError"
|
||||
else:
|
||||
internal_message, message = _extract_message(e)
|
||||
exc_type = e.__class__.__name__
|
||||
|
||||
# TODO make this more robust
|
||||
if "(Request ID: " not in message:
|
||||
message += f" (Request ID: {request_id})"
|
||||
|
||||
if "(Request ID: " not in internal_message:
|
||||
internal_message += f" (Request ID: {request_id})"
|
||||
|
||||
error_info = ErrorInfo(
|
||||
message=f"Message: {message}, Internal exception: {internal_message}, original exception: {str(e)}",
|
||||
code=status_code,
|
||||
type=exc_type,
|
||||
)
|
||||
error_response = ErrorResponse(error=error_info)
|
||||
return error_response
|
||||
|
||||
|
||||
def get_serve_request_id() -> str:
|
||||
"""Get request id from serve request context."""
|
||||
context = serve.context._serve_request_context.get()
|
||||
if context is not None:
|
||||
return context.request_id
|
||||
return ""
|
||||
|
||||
|
||||
def get_model_request_id(model: str):
|
||||
return model + "-" + get_serve_request_id()
|
||||
|
||||
|
||||
def replace_prefix(model: str) -> str:
|
||||
"""Replace -- with / in model name to handle slashes within the URL path segment"""
|
||||
return model.replace("--", "/")
|
||||
Reference in New Issue
Block a user