From a8599a285598c0f8f42a7f7b3949a34f83fe3949 Mon Sep 17 00:00:00 2001 From: w1nstep <3539832429@qq.com> Date: Wed, 2 Sep 2026 10:15:03 +0800 Subject: [PATCH 1/2] chore(redis): add legacy rate-limit key cleanup --- scripts/cleanup_legacy_rate_limit_keys.py | 254 ++++++++++++++++++ .../test_cleanup_legacy_rate_limit_keys.py | 111 ++++++++ 2 files changed, 365 insertions(+) create mode 100644 scripts/cleanup_legacy_rate_limit_keys.py create mode 100644 tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py 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) From e95c77e59316d067da414c697a826c856f2446fc Mon Sep 17 00:00:00 2001 From: w1nstep <3539832429@qq.com> Date: Wed, 2 Sep 2026 10:55:49 +0800 Subject: [PATCH 2/2] fix(redis): make legacy key cleanup safer --- scripts/cleanup_legacy_rate_limit_keys.py | 33 +++++++---- .../test_cleanup_legacy_rate_limit_keys.py | 59 +++++++++++++++++-- 2 files changed, 76 insertions(+), 16 deletions(-) diff --git a/scripts/cleanup_legacy_rate_limit_keys.py b/scripts/cleanup_legacy_rate_limit_keys.py index 4399e5d08d7..c232ffc6d17 100644 --- a/scripts/cleanup_legacy_rate_limit_keys.py +++ b/scripts/cleanup_legacy_rate_limit_keys.py @@ -10,15 +10,18 @@ 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. +namespace can be supplied to narrow the SCAN. Apply mode requires a namespace +and removes only keys with a permanent TTL (``TTL=-1``), so a dry-run is the +only unscoped mode and active finite-TTL counters are retained. 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 os import re import sys from collections import Counter @@ -128,6 +131,8 @@ def cleanup_legacy_keys( normalized_namespace = _normalize_namespace(namespace) if count <= 0: raise ValueError("count must be positive") + if apply and normalized_namespace is None: + raise ValueError("namespace is required when apply=True") scan_iter = getattr(client, "scan_iter", None) ttl = getattr(client, "ttl", None) @@ -153,14 +158,16 @@ def cleanup_legacy_keys( report.ttl_errors += 1 continue + is_permanent = False if ttl_value == -1: report.permanent += 1 + is_permanent = True elif isinstance(ttl_value, int) and ttl_value >= 0: report.finite += 1 else: report.unknown_ttl += 1 - if apply: + if apply and is_permanent: try: _delete_one(client, key) except Exception: @@ -172,8 +179,12 @@ def cleanup_legacy_keys( def _connection_kwargs(args: argparse.Namespace) -> dict[str, object]: + if args.host is None and (args.port is not None or args.db is not None): + raise ValueError("--host is required when using --port or --db") + if args.host is not None and args.port is None: + raise ValueError("--port is required when using --host") + values = { - "url": args.url, "host": args.host, "port": args.port, "db": args.db, @@ -187,7 +198,10 @@ def _connect(args: argparse.Namespace) -> object: # and Cluster/Sentinel connection handling for the real command. from litellm._redis import get_redis_client - return get_redis_client(**_connection_kwargs(args)) + connection_kwargs = _connection_kwargs(args) + if connection_kwargs and (os.getenv("REDIS_CLUSTER_NODES") or os.getenv("REDIS_SENTINEL_NODES")): + raise ValueError("explicit host overrides cannot be combined with Redis Cluster or Sentinel configuration") + return get_redis_client(**connection_kwargs) def _write_report(report: CleanupReport, *, apply: bool, json_output: bool) -> None: @@ -216,10 +230,9 @@ 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("--host", help="Redis host override; use with --port") + parser.add_argument("--port", type=int, help="Redis port override; use with --host") + parser.add_argument("--db", type=int, help="Redis database override; use with --host") parser.add_argument("--json", action="store_true", dest="json_output", help="emit machine-readable JSON") parser.add_argument( "--apply", diff --git a/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py b/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py index f7291d47879..0b5425d26f4 100644 --- a/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py +++ b/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py @@ -2,6 +2,7 @@ import importlib.util import sys +from argparse import Namespace from pathlib import Path import pytest @@ -84,28 +85,74 @@ def test_dry_run_scans_cluster_aware_iterator_without_deleting() -> None: assert redis.unlink_calls == [] -def test_apply_deletes_only_unique_legacy_keys() -> None: +def test_apply_requires_namespace() -> 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) + with pytest.raises(ValueError, match="namespace"): + cleanup.cleanup_legacy_keys(redis, apply=True) - assert report.candidates == 1 + assert redis.unlink_calls == [] + + +def test_apply_deletes_only_unique_permanent_legacy_keys() -> None: + permanent_key = "tenant-a:global_router:id:model:rpm:13-07" + finite_key = "tenant-a:global_router:id:model:tpm:13-07" + redis = FakeRedis( + [permanent_key, permanent_key, finite_key, "tenant-a:global_router:id:model:rpm:v2:29300000"], + {permanent_key: -1, finite_key: 30}, + ) + + report = cleanup.cleanup_legacy_keys(redis, namespace="tenant-a", apply=True) + + assert report.candidates == 2 + assert report.permanent == 1 + assert report.finite == 1 assert report.deleted == 1 assert report.failed == 0 - assert redis.unlink_calls == [old_key] + assert redis.unlink_calls == [permanent_key] def test_apply_falls_back_to_single_key_delete() -> None: - old_key = "13-07:model" + old_key = "tenant-a:13-07:model" redis = DeleteOnlyRedis([old_key], {old_key: 9}) - report = cleanup.cleanup_legacy_keys(redis, apply=True) + report = cleanup.cleanup_legacy_keys(redis, namespace="tenant-a", apply=True) + + assert report.deleted == 0 + assert redis.delete_calls == [] + + +def test_apply_falls_back_to_single_key_delete_for_permanent_key() -> None: + old_key = "tenant-a:13-07:model" + redis = DeleteOnlyRedis([old_key], {old_key: -1}) + + report = cleanup.cleanup_legacy_keys(redis, namespace="tenant-a", apply=True) assert report.deleted == 1 assert redis.delete_calls == [old_key] +def test_connection_overrides_require_an_explicit_host_and_port() -> None: + with pytest.raises(ValueError, match="--host"): + cleanup._connection_kwargs(Namespace(host=None, port=6380, db=None)) + with pytest.raises(ValueError, match="--port"): + cleanup._connection_kwargs(Namespace(host="redis.example", port=None, db=None)) + + +def test_connection_overrides_do_not_accept_a_url() -> None: + kwargs = cleanup._connection_kwargs(Namespace(host="redis.example", port=6380, db=2)) + + assert kwargs == {"host": "redis.example", "port": 6380, "db": 2} + + +def test_connection_overrides_do_not_mix_with_cluster_configuration(monkeypatch) -> None: + monkeypatch.setenv("REDIS_CLUSTER_NODES", '[{"host": "cluster.example", "port": 6379}]') + + with pytest.raises(ValueError, match="Cluster or Sentinel"): + cleanup._connect(Namespace(host="redis.example", port=6380, db=None)) + + def test_count_must_be_positive() -> None: with pytest.raises(ValueError, match="positive"): cleanup.cleanup_legacy_keys(FakeRedis([], {}), count=0)