wehub-resource-sync
329 lines
11 KiB
Python
329 lines
11 KiB
Python
"""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)
|