127 lines
4.8 KiB
Python
127 lines
4.8 KiB
Python
"""Rate limiter for README updates.
|
|
|
|
Implements a cooldown mechanism to prevent excessive HuggingFace API calls
|
|
while ensuring all updates eventually reach the repository.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from helpers.logger import logger
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
|
|
class ReadmeRateLimiter:
|
|
"""Rate limits README updates to prevent API throttling.
|
|
|
|
Ensures updates are batched with a minimum interval between API calls,
|
|
while guaranteeing that pending updates are eventually applied.
|
|
"""
|
|
|
|
def __init__(self, cooldown_seconds: float = 30.0) -> None:
|
|
"""Initialise rate limiter with specified cooldown period.
|
|
|
|
Sets up the rate limiter with the specified cooldown interval to
|
|
prevent excessive API calls whilst ensuring pending updates are
|
|
eventually processed through a timer-based batching mechanism.
|
|
"""
|
|
self.cooldown_seconds = cooldown_seconds
|
|
self.last_update_time = 0.0
|
|
self.pending_update = False
|
|
self.update_lock = threading.Lock()
|
|
self.timer: threading.Timer | None = None
|
|
self.update_func: Callable[..., Any] | None = None
|
|
self.update_args: tuple[Any, ...] | None = None
|
|
self.update_kwargs: dict[str, Any] | None = None
|
|
|
|
def request_update(
|
|
self,
|
|
update_func: Callable[..., Any],
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Request a README update, respecting rate limits.
|
|
|
|
Updates are batched during cooldown periods and executed
|
|
when the cooldown expires. Stores the update function and its
|
|
arguments for deferred execution whilst maintaining thread safety.
|
|
"""
|
|
with self.update_lock:
|
|
current_time = time.time()
|
|
time_since_last = current_time - self.last_update_time
|
|
|
|
# Store the latest update request
|
|
self.update_func = update_func
|
|
self.update_args = args
|
|
self.update_kwargs = kwargs
|
|
|
|
if time_since_last >= self.cooldown_seconds:
|
|
# Enough time has passed, update immediately
|
|
logger.debug(f"README update allowed (last update {time_since_last:.1f}s ago)")
|
|
self._execute_update()
|
|
else:
|
|
# Still in cooldown, schedule for later
|
|
remaining = self.cooldown_seconds - time_since_last
|
|
logger.debug(f"README update delayed ({remaining:.1f}s cooldown remaining)")
|
|
|
|
if not self.pending_update:
|
|
# Schedule an update when cooldown expires
|
|
self.pending_update = True
|
|
if self.timer:
|
|
self.timer.cancel()
|
|
self.timer = threading.Timer(remaining, self._delayed_update)
|
|
self.timer.start()
|
|
else:
|
|
# Update already scheduled, just update the args
|
|
logger.debug("README update already scheduled, updating with latest data")
|
|
|
|
def _execute_update(self) -> None:
|
|
"""Execute the actual update (must be called with lock held)."""
|
|
if self.update_func:
|
|
try:
|
|
args = self.update_args or ()
|
|
kwargs = self.update_kwargs or {}
|
|
self.update_func(*args, **kwargs)
|
|
self.last_update_time = time.time()
|
|
logger.debug("README update completed")
|
|
except Exception as e:
|
|
logger.error(f"README update failed: {e}")
|
|
|
|
self.pending_update = False
|
|
self.update_func = None
|
|
self.update_args = None
|
|
self.update_kwargs = None
|
|
|
|
def _delayed_update(self) -> None:
|
|
"""Execute a delayed update after cooldown expires."""
|
|
with self.update_lock:
|
|
if self.pending_update:
|
|
logger.debug("Executing delayed README update")
|
|
self._execute_update()
|
|
|
|
def flush(self) -> None:
|
|
"""Force any pending updates to execute immediately.
|
|
|
|
Called at script end to ensure final state is uploaded.
|
|
"""
|
|
with self.update_lock:
|
|
if self.timer:
|
|
self.timer.cancel()
|
|
self.timer = None
|
|
|
|
if self.pending_update and self.update_func:
|
|
logger.info("Flushing pending README update...")
|
|
# Wait for cooldown if needed
|
|
current_time = time.time()
|
|
time_since_last = current_time - self.last_update_time
|
|
if time_since_last < self.cooldown_seconds:
|
|
wait_time = self.cooldown_seconds - time_since_last
|
|
logger.info(f"Waiting {wait_time:.1f}s for cooldown before final update...")
|
|
time.sleep(wait_time)
|
|
|
|
self._execute_update()
|