diff --git a/scripts/cleanup_legacy_rate_limit_keys.py b/scripts/cleanup_legacy_rate_limit_keys.py new file mode 100644 index 00000000000..4399e5d08d7 --- /dev/null +++ b/scripts/cleanup_legacy_rate_limit_keys.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 +"""Find and optionally remove legacy Redis rate-limit keys. + +The rate-limit key format changed from a local-clock ``HH-MM`` suffix/prefix to +versioned UTC epoch-minute keys. This maintenance command finds only the old +formats. It is a dry-run by default; deleting keys requires both ``--apply`` +and the explicit confirmation token. + +Usage: + python scripts/cleanup_legacy_rate_limit_keys.py [options] + +The Redis connection is read from LiteLLM's normal REDIS_* configuration. A +namespace can be supplied to narrow the SCAN. redis-py RedisCluster's +``scan_iter`` is used as-is, so cluster scans remain node-aware. Deletions are +sent one key at a time to avoid a cross-slot multi-key command. +""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from collections import Counter +from dataclasses import dataclass, field +from typing import Any + +CONFIRMATION_TOKEN = "DELETE_LEGACY_RATE_LIMIT_KEYS" +_COUNTER_KINDS = frozenset({"tpm", "rpm", "itpm", "otpm"}) +_WINDOW_RE = re.compile(r"^(?:[01][0-9]|2[0-3])-[0-5][0-9]$") + + +@dataclass +class CleanupReport: + """Counters returned by a legacy-key scan.""" + + scanned: int = 0 + candidates: int = 0 + permanent: int = 0 + finite: int = 0 + unknown_ttl: int = 0 + ttl_errors: int = 0 + deleted: int = 0 + failed: int = 0 + by_kind: Counter[str] = field(default_factory=Counter) + + def as_dict(self) -> dict[str, Any]: + return { + "scanned": self.scanned, + "candidates": self.candidates, + "permanent": self.permanent, + "finite": self.finite, + "unknown_ttl": self.unknown_ttl, + "ttl_errors": self.ttl_errors, + "deleted": self.deleted, + "failed": self.failed, + "by_kind": dict(sorted(self.by_kind.items())), + } + + +def _normalize_namespace(namespace: str | None) -> str | None: + if namespace is None: + return None + if not namespace or namespace != namespace.strip() or any(char.isspace() for char in namespace): + raise ValueError("namespace must be non-empty and must not contain whitespace") + normalized = namespace.rstrip(":") + if not normalized: + raise ValueError("namespace must contain at least one non-colon character") + return normalized + + +def _as_text(value: object) -> str: + if isinstance(value, bytes): + return value.decode("utf-8", errors="replace") + return str(value) + + +def _is_legacy_window(value: str) -> bool: + return _WINDOW_RE.fullmatch(value) is not None + + +def legacy_key_kind(key: str | bytes, namespace: str | None = None) -> str | None: + """Return the legacy key kind, or ``None`` for a v2/unrelated key. + + Counter keys are recognized by their final ``::HH-MM`` segments, + which covers the historical global-router and model-group variants. The + dynamic limiter used ``HH-MM:`` and is recognized separately. + """ + + normalized_namespace = _normalize_namespace(namespace) + key_text = _as_text(key) + if normalized_namespace is not None: + namespace_prefix = f"{normalized_namespace}:" + if not key_text.startswith(namespace_prefix): + return None + key_text = key_text[len(namespace_prefix) :] + + parts = key_text.split(":") + if len(parts) >= 3 and parts[-2] in _COUNTER_KINDS and _is_legacy_window(parts[-1]): + return parts[-2] + if len(parts) >= 2 and _is_legacy_window(parts[0]) and ":".join(parts[1:]): + return "dynamic" + return None + + +def _delete_one(client: object, key: str) -> None: + """Delete one key so a Cluster client never receives a cross-slot batch.""" + + unlink = getattr(client, "unlink", None) + if callable(unlink): + unlink(key) + return + delete = getattr(client, "delete", None) + if not callable(delete): + raise AttributeError("Redis client has neither unlink nor delete") + delete(key) + + +def cleanup_legacy_keys( + client: object, + *, + namespace: str | None = None, + count: int = 1000, + apply: bool = False, +) -> CleanupReport: + """Scan a Redis client and optionally remove unique legacy rate-limit keys.""" + + normalized_namespace = _normalize_namespace(namespace) + if count <= 0: + raise ValueError("count must be positive") + + scan_iter = getattr(client, "scan_iter", None) + ttl = getattr(client, "ttl", None) + if not callable(scan_iter) or not callable(ttl): + raise AttributeError("Redis client must provide scan_iter and ttl") + + match = f"{normalized_namespace}:*" if normalized_namespace is not None else "*" + report = CleanupReport() + seen_candidates: set[str] = set() + for raw_key in scan_iter(match=match, count=count): + report.scanned += 1 + key = _as_text(raw_key) + kind = legacy_key_kind(key, normalized_namespace) + if kind is None or key in seen_candidates: + continue + + seen_candidates.add(key) + report.candidates += 1 + report.by_kind[kind] += 1 + try: + ttl_value = ttl(key) + except Exception: + report.ttl_errors += 1 + continue + + if ttl_value == -1: + report.permanent += 1 + elif isinstance(ttl_value, int) and ttl_value >= 0: + report.finite += 1 + else: + report.unknown_ttl += 1 + + if apply: + try: + _delete_one(client, key) + except Exception: + report.failed += 1 + else: + report.deleted += 1 + + return report + + +def _connection_kwargs(args: argparse.Namespace) -> dict[str, object]: + values = { + "url": args.url, + "host": args.host, + "port": args.port, + "db": args.db, + } + return {name: value for name, value in values.items() if value is not None} + + +def _connect(args: argparse.Namespace) -> object: + # Importing at runtime keeps the unit-testable key classifier independent of + # LiteLLM's optional provider imports, while reusing its masked REDIS_* auth + # and Cluster/Sentinel connection handling for the real command. + from litellm._redis import get_redis_client + + return get_redis_client(**_connection_kwargs(args)) + + +def _write_report(report: CleanupReport, *, apply: bool, json_output: bool) -> None: + if json_output: + payload = report.as_dict() + payload["mode"] = "apply" if apply else "dry-run" + sys.stdout.write(json.dumps(payload, sort_keys=True) + "\n") + return + + mode = "apply" if apply else "dry-run" + summary = report.as_dict() + sys.stdout.write( + f"mode={mode} scanned={summary['scanned']} candidates={summary['candidates']} " + f"permanent={summary['permanent']} finite={summary['finite']} " + f"unknown_ttl={summary['unknown_ttl']} ttl_errors={summary['ttl_errors']} " + f"deleted={summary['deleted']} failed={summary['failed']}\n" + ) + kinds = ", ".join(f"{kind}={amount}" for kind, amount in sorted(report.by_kind.items())) + if kinds: + sys.stdout.write(f"candidate_kinds={kinds}\n") + if not apply: + sys.stdout.write("dry-run: no keys changed\n") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--namespace", help="only inspect keys beginning with NAMESPACE:") + parser.add_argument("--count", type=int, default=1000, help="Redis SCAN count hint (default: 1000)") + parser.add_argument("--url", help="Redis URL; otherwise use the normal REDIS_* configuration") + parser.add_argument("--host", help="Redis host override") + parser.add_argument("--port", type=int, help="Redis port override") + parser.add_argument("--db", type=int, help="Redis database override") + parser.add_argument("--json", action="store_true", dest="json_output", help="emit machine-readable JSON") + parser.add_argument( + "--apply", + action="store_true", + help="delete matching keys; requires --confirm " + CONFIRMATION_TOKEN, + ) + parser.add_argument("--confirm", help=argparse.SUPPRESS) + args = parser.parse_args() + + if args.apply and args.confirm != CONFIRMATION_TOKEN: + parser.error("--apply requires --confirm " + CONFIRMATION_TOKEN) + + try: + report = cleanup_legacy_keys( + _connect(args), + namespace=args.namespace, + count=args.count, + apply=args.apply, + ) + except ValueError as exc: + parser.error(str(exc)) + except Exception as exc: + # Do not echo connection details: a configured URL may contain credentials. + sys.stderr.write(f"ERROR: Redis legacy-key cleanup failed ({type(exc).__name__})\n") + return 1 + + _write_report(report, apply=args.apply, json_output=args.json_output) + return 1 if report.failed or report.ttl_errors else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py b/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py new file mode 100644 index 00000000000..f7291d47879 --- /dev/null +++ b/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py @@ -0,0 +1,111 @@ +"""Tests for the legacy Redis rate-limit key maintenance command.""" + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "cleanup_legacy_rate_limit_keys.py" +_spec = importlib.util.spec_from_file_location("cleanup_legacy_rate_limit_keys", _MODULE_PATH) +assert _spec is not None and _spec.loader is not None +cleanup = importlib.util.module_from_spec(_spec) +sys.modules[_spec.name] = cleanup +_spec.loader.exec_module(cleanup) + + +class FakeRedis: + def __init__(self, keys: list[str], ttls: dict[str, int]): + self.keys = keys + self.ttls = ttls + self.scan_calls: list[tuple[str, int]] = [] + self.unlink_calls: list[str] = [] + + def scan_iter(self, *, match: str, count: int): + self.scan_calls.append((match, count)) + yield from self.keys + + def ttl(self, key: str) -> int: + return self.ttls.get(key, -2) + + def unlink(self, key: str) -> int: + self.unlink_calls.append(key) + return 1 + + +class DeleteOnlyRedis(FakeRedis): + unlink = None + + def __init__(self, keys: list[str], ttls: dict[str, int]): + super().__init__(keys, ttls) + self.delete_calls: list[str] = [] + + def delete(self, key: str) -> int: + self.delete_calls.append(key) + return 1 + + +def test_legacy_key_kind_distinguishes_old_formats_from_v2() -> None: + assert cleanup.legacy_key_kind("global_router:id:model:tpm:13-07") == "tpm" + assert cleanup.legacy_key_kind("model:rpm:13-07") == "rpm" + assert cleanup.legacy_key_kind("13-07:model") == "dynamic" + assert cleanup.legacy_key_kind("global_router:id:model:tpm:v2:29300000") is None + assert cleanup.legacy_key_kind("v2:29300000:model") is None + assert cleanup.legacy_key_kind("global_router:id:model:rpm:24-00") is None + + +def test_namespace_is_required_for_matching_when_supplied() -> None: + key = "tenant-a:global_router:id:model:rpm:13-07" + assert cleanup.legacy_key_kind(key, "tenant-a") == "rpm" + assert cleanup.legacy_key_kind(key, "tenant-b") is None + with pytest.raises(ValueError, match="whitespace"): + cleanup.legacy_key_kind(key, "tenant a") + + +def test_dry_run_scans_cluster_aware_iterator_without_deleting() -> None: + keys = [ + "tenant-a:global_router:id:model:tpm:13-07", + "tenant-a:model:rpm:13-07", + "tenant-a:13-07:model", + "tenant-a:global_router:id:model:tpm:v2:29300000", + "tenant-b:global_router:id:model:rpm:13-07", + ] + redis = FakeRedis(keys, {keys[0]: -1, keys[1]: 30, keys[2]: 4}) + + report = cleanup.cleanup_legacy_keys(redis, namespace="tenant-a", count=17) + + assert redis.scan_calls == [("tenant-a:*", 17)] + assert report.scanned == 5 + assert report.candidates == 3 + assert report.permanent == 1 + assert report.finite == 2 + assert report.by_kind == {"dynamic": 1, "rpm": 1, "tpm": 1} + assert report.deleted == 0 + assert redis.unlink_calls == [] + + +def test_apply_deletes_only_unique_legacy_keys() -> None: + old_key = "global_router:id:model:rpm:13-07" + redis = FakeRedis([old_key, old_key, "global_router:id:model:rpm:v2:29300000"], {old_key: -1}) + + report = cleanup.cleanup_legacy_keys(redis, apply=True) + + assert report.candidates == 1 + assert report.deleted == 1 + assert report.failed == 0 + assert redis.unlink_calls == [old_key] + + +def test_apply_falls_back_to_single_key_delete() -> None: + old_key = "13-07:model" + redis = DeleteOnlyRedis([old_key], {old_key: 9}) + + report = cleanup.cleanup_legacy_keys(redis, apply=True) + + assert report.deleted == 1 + assert redis.delete_calls == [old_key] + + +def test_count_must_be_positive() -> None: + with pytest.raises(ValueError, match="positive"): + cleanup.cleanup_legacy_keys(FakeRedis([], {}), count=0)