wehub-resource-sync
114 lines
3.6 KiB
Python
114 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Factory for creating request telemetry reporters based on configuration.
|
|
"""
|
|
|
|
# Standard
|
|
from collections.abc import Callable
|
|
from typing import TYPE_CHECKING, Any, Optional
|
|
import importlib
|
|
|
|
if TYPE_CHECKING:
|
|
# First Party
|
|
from lmcache.integration.request_telemetry.base import RequestTelemetry
|
|
|
|
|
|
class RequestTelemetryFactory:
|
|
"""Factory for creating request telemetry reporters."""
|
|
|
|
_registry: dict[str, Callable[[], type["RequestTelemetry"]]] = {}
|
|
_instances: dict[str, "RequestTelemetry"] = {}
|
|
|
|
@classmethod
|
|
def register(cls, name: str, module_path: str, class_name: str) -> None:
|
|
"""Register a telemetry reporter with lazy-loading module and class name.
|
|
|
|
Args:
|
|
name: The name to register the telemetry reporter under.
|
|
module_path: The module path to import the class from.
|
|
class_name: The class name to import from the module.
|
|
|
|
Raises:
|
|
ValueError: If a reporter with the same name is already registered.
|
|
"""
|
|
if name in cls._registry:
|
|
raise ValueError(f"Telemetry reporter '{name}' is already registered.")
|
|
|
|
def loader() -> type["RequestTelemetry"]:
|
|
module = importlib.import_module(module_path)
|
|
return getattr(module, class_name)
|
|
|
|
cls._registry[name] = loader
|
|
|
|
@classmethod
|
|
def create(
|
|
cls,
|
|
telemetry_type: Optional[str] = None,
|
|
config: Optional[dict[str, Any]] = None,
|
|
use_singleton: bool = True,
|
|
) -> "RequestTelemetry":
|
|
"""Create a request telemetry reporter based on the specified type.
|
|
|
|
Args:
|
|
telemetry_type: The type of telemetry reporter to create.
|
|
If None, defaults to "noop".
|
|
Supported values: "noop", "fastapi".
|
|
config: Optional dict of configuration options passed to the reporter.
|
|
Each reporter defines its own expected keys.
|
|
use_singleton: If True, returns a cached singleton instance.
|
|
If False, creates a new instance each time.
|
|
Note: Singleton is keyed by telemetry_type only.
|
|
|
|
Returns:
|
|
A RequestTelemetry instance.
|
|
|
|
Raises:
|
|
ValueError: If an unsupported telemetry type is specified.
|
|
"""
|
|
if telemetry_type is None:
|
|
telemetry_type = "noop"
|
|
|
|
telemetry_type = telemetry_type.lower()
|
|
|
|
if telemetry_type not in cls._registry:
|
|
raise ValueError(
|
|
f"Unsupported request telemetry type: {telemetry_type}. "
|
|
f"Supported types: {list(cls._registry.keys())}"
|
|
)
|
|
|
|
if use_singleton and telemetry_type in cls._instances:
|
|
return cls._instances[telemetry_type]
|
|
|
|
if config is None:
|
|
config = {}
|
|
|
|
telemetry_cls = cls._registry[telemetry_type]()
|
|
instance = telemetry_cls(config)
|
|
|
|
if use_singleton:
|
|
cls._instances[telemetry_type] = instance
|
|
|
|
return instance
|
|
|
|
@classmethod
|
|
def get_registered_types(cls) -> list[str]:
|
|
"""Get a list of all registered telemetry types."""
|
|
return list(cls._registry.keys())
|
|
|
|
|
|
# Register telemetry types here.
|
|
# The registration should not be done in each individual file, as we want to
|
|
# only load the files corresponding to the current telemetry type.
|
|
|
|
RequestTelemetryFactory.register(
|
|
"noop",
|
|
"lmcache.integration.request_telemetry.noop",
|
|
"NoOpRequestTelemetry",
|
|
)
|
|
|
|
RequestTelemetryFactory.register(
|
|
"fastapi",
|
|
"lmcache.integration.request_telemetry.fastapi",
|
|
"FastAPIRequestTelemetry",
|
|
)
|