chore: initial public snapshot for github upload
This commit is contained in:
Your Name
@@ -0,0 +1,214 @@
|
||||
"""Focus export logger orchestrating DB pull/transform/upload."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
from .destinations import FocusTimeWindow
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from .export_engine import FocusExportEngine
|
||||
else:
|
||||
AsyncIOScheduler = Any
|
||||
|
||||
FOCUS_USAGE_DATA_JOB_NAME = "focus_export_usage_data"
|
||||
DEFAULT_DRY_RUN_LIMIT = 500
|
||||
|
||||
|
||||
class FocusLogger(CustomLogger):
|
||||
"""Coordinates Focus export jobs across transformer/serializer/destination layers."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
provider: Optional[str] = None,
|
||||
export_format: Optional[str] = None,
|
||||
frequency: Optional[str] = None,
|
||||
cron_offset_minute: Optional[int] = None,
|
||||
interval_seconds: Optional[int] = None,
|
||||
prefix: Optional[str] = None,
|
||||
destination_config: Optional[dict[str, Any]] = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.provider = (provider or os.getenv("FOCUS_PROVIDER") or "s3").lower()
|
||||
self.export_format = (
|
||||
export_format or os.getenv("FOCUS_FORMAT") or "parquet"
|
||||
).lower()
|
||||
self.frequency = (frequency or os.getenv("FOCUS_FREQUENCY") or "hourly").lower()
|
||||
self.cron_offset_minute = (
|
||||
cron_offset_minute
|
||||
if cron_offset_minute is not None
|
||||
else int(os.getenv("FOCUS_CRON_OFFSET", "5"))
|
||||
)
|
||||
raw_interval = (
|
||||
interval_seconds
|
||||
if interval_seconds is not None
|
||||
else os.getenv("FOCUS_INTERVAL_SECONDS")
|
||||
)
|
||||
self.interval_seconds = int(raw_interval) if raw_interval is not None else None
|
||||
env_prefix = os.getenv("FOCUS_PREFIX")
|
||||
self.prefix: str = (
|
||||
prefix
|
||||
if prefix is not None
|
||||
else (env_prefix if env_prefix else "focus_exports")
|
||||
)
|
||||
|
||||
self._destination_config = destination_config
|
||||
self._engine: Optional["FocusExportEngine"] = None
|
||||
|
||||
def _ensure_engine(self) -> "FocusExportEngine":
|
||||
"""Instantiate the heavy export engine lazily."""
|
||||
if self._engine is None:
|
||||
from .export_engine import FocusExportEngine
|
||||
|
||||
self._engine = FocusExportEngine(
|
||||
provider=self.provider,
|
||||
export_format=self.export_format,
|
||||
prefix=self.prefix,
|
||||
destination_config=self._destination_config,
|
||||
)
|
||||
return self._engine
|
||||
|
||||
async def export_usage_data(
|
||||
self,
|
||||
*,
|
||||
limit: Optional[int] = None,
|
||||
start_time_utc: Optional[datetime] = None,
|
||||
end_time_utc: Optional[datetime] = None,
|
||||
) -> None:
|
||||
"""Public hook to trigger export immediately."""
|
||||
if bool(start_time_utc) ^ bool(end_time_utc):
|
||||
raise ValueError(
|
||||
"start_time_utc and end_time_utc must be provided together"
|
||||
)
|
||||
|
||||
if start_time_utc and end_time_utc:
|
||||
window = FocusTimeWindow(
|
||||
start_time=start_time_utc,
|
||||
end_time=end_time_utc,
|
||||
frequency=self.frequency,
|
||||
)
|
||||
else:
|
||||
window = self._compute_time_window(datetime.now(timezone.utc))
|
||||
await self._export_window(window=window, limit=limit)
|
||||
|
||||
async def dry_run_export_usage_data(
|
||||
self, limit: Optional[int] = DEFAULT_DRY_RUN_LIMIT
|
||||
) -> dict[str, Any]:
|
||||
"""Return transformed data without uploading."""
|
||||
engine = self._ensure_engine()
|
||||
return await engine.dry_run_export_usage_data(limit=limit)
|
||||
|
||||
async def initialize_focus_export_job(self) -> None:
|
||||
"""Entry point for scheduler jobs to run export cycle with locking."""
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
pod_lock_manager = None
|
||||
if proxy_logging_obj is not None:
|
||||
writer = getattr(proxy_logging_obj, "db_spend_update_writer", None)
|
||||
if writer is not None:
|
||||
pod_lock_manager = getattr(writer, "pod_lock_manager", None)
|
||||
|
||||
if pod_lock_manager and pod_lock_manager.redis_cache:
|
||||
acquired = await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=FOCUS_USAGE_DATA_JOB_NAME
|
||||
)
|
||||
if not acquired:
|
||||
verbose_logger.debug("Focus export: unable to acquire pod lock")
|
||||
return
|
||||
try:
|
||||
await self._run_scheduled_export()
|
||||
finally:
|
||||
await pod_lock_manager.release_lock(
|
||||
cronjob_id=FOCUS_USAGE_DATA_JOB_NAME
|
||||
)
|
||||
else:
|
||||
await self._run_scheduled_export()
|
||||
|
||||
@staticmethod
|
||||
async def init_focus_export_background_job(
|
||||
scheduler: AsyncIOScheduler,
|
||||
) -> None:
|
||||
"""Register the export cron/interval job with the provided scheduler."""
|
||||
|
||||
focus_loggers: List[
|
||||
CustomLogger
|
||||
] = litellm.logging_callback_manager.get_custom_loggers_for_type(
|
||||
callback_type=FocusLogger
|
||||
)
|
||||
if not focus_loggers:
|
||||
verbose_logger.debug(
|
||||
"No Focus export logger registered; skipping scheduler"
|
||||
)
|
||||
return
|
||||
|
||||
focus_logger = cast(FocusLogger, focus_loggers[0])
|
||||
trigger_kwargs = focus_logger._build_scheduler_trigger()
|
||||
scheduler.add_job(
|
||||
focus_logger.initialize_focus_export_job,
|
||||
**trigger_kwargs,
|
||||
)
|
||||
|
||||
def _build_scheduler_trigger(self) -> Dict[str, Any]:
|
||||
"""Return scheduler configuration for the selected frequency."""
|
||||
if self.frequency == "interval":
|
||||
seconds = self.interval_seconds or 60
|
||||
return {"trigger": "interval", "seconds": seconds}
|
||||
|
||||
if self.frequency == "hourly":
|
||||
minute = max(0, min(59, self.cron_offset_minute))
|
||||
return {"trigger": "cron", "minute": minute, "second": 0}
|
||||
|
||||
if self.frequency == "daily":
|
||||
total_minutes = max(0, self.cron_offset_minute)
|
||||
hour = min(23, total_minutes // 60)
|
||||
minute = min(59, total_minutes % 60)
|
||||
return {"trigger": "cron", "hour": hour, "minute": minute, "second": 0}
|
||||
|
||||
raise ValueError(f"Unsupported frequency: {self.frequency}")
|
||||
|
||||
async def _run_scheduled_export(self) -> None:
|
||||
"""Execute the scheduled export for the configured window."""
|
||||
window = self._compute_time_window(datetime.now(timezone.utc))
|
||||
await self._export_window(window=window, limit=None)
|
||||
|
||||
async def _export_window(
|
||||
self,
|
||||
*,
|
||||
window: FocusTimeWindow,
|
||||
limit: Optional[int],
|
||||
) -> None:
|
||||
engine = self._ensure_engine()
|
||||
await engine.export_window(window=window, limit=limit)
|
||||
|
||||
def _compute_time_window(self, now: datetime) -> FocusTimeWindow:
|
||||
"""Derive the time window to export based on configured frequency."""
|
||||
now_utc = now.astimezone(timezone.utc)
|
||||
if self.frequency == "hourly":
|
||||
end_time = now_utc.replace(minute=0, second=0, microsecond=0)
|
||||
start_time = end_time - timedelta(hours=1)
|
||||
elif self.frequency == "daily":
|
||||
end_time = now_utc.replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
start_time = end_time - timedelta(days=1)
|
||||
elif self.frequency == "interval":
|
||||
interval = timedelta(seconds=self.interval_seconds or 60)
|
||||
end_time = now_utc
|
||||
start_time = end_time - interval
|
||||
else:
|
||||
raise ValueError(f"Unsupported frequency: {self.frequency}")
|
||||
return FocusTimeWindow(
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
frequency=self.frequency,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["FocusLogger"]
|
||||
Reference in New Issue
Block a user