wehub-resource-sync
409 lines
11 KiB
Python
409 lines
11 KiB
Python
import pytest
|
|
|
|
from mlflow.genai.evaluation.harness import (
|
|
AUTO_INITIAL_RPS,
|
|
_make_rate_limiter,
|
|
_parse_rate_limit,
|
|
)
|
|
from mlflow.genai.evaluation.rate_limiter import (
|
|
NoOpRateLimiter,
|
|
RPSRateLimiter,
|
|
call_with_retry,
|
|
eval_retry_context,
|
|
is_rate_limit_error,
|
|
)
|
|
from mlflow.genai.judges.adapters.litellm_adapter import (
|
|
_get_litellm_retry_policy,
|
|
is_litellm_rate_limit_retries_disabled,
|
|
)
|
|
from mlflow.utils.rest_utils import is_429_retry_disabled
|
|
|
|
|
|
class FakeClock:
|
|
"""Deterministic clock for testing. sleep() advances the clock by the requested amount.
|
|
|
|
Thread safety is not needed here because RPSRateLimiter's internal lock serializes
|
|
all calls to clock() and sleep() — they are never called concurrently for a given limiter.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._now = 0.0
|
|
self.sleep_calls: list[float] = []
|
|
|
|
def monotonic(self) -> float:
|
|
return self._now
|
|
|
|
def sleep(self, seconds: float) -> None:
|
|
self.sleep_calls.append(seconds)
|
|
self._now += seconds
|
|
|
|
|
|
# ── Token bucket tests ──
|
|
|
|
|
|
def test_invalid_rate_raises():
|
|
with pytest.raises(ValueError, match="must be positive"):
|
|
RPSRateLimiter(0)
|
|
with pytest.raises(ValueError, match="must be positive"):
|
|
RPSRateLimiter(-1)
|
|
|
|
|
|
def test_sub_one_rps_can_acquire():
|
|
# rps < 1.0 was broken: _max_tokens was set to rps, so the bucket could never
|
|
# accumulate a full token and acquire() would loop forever.
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(0.5, clock=clock.monotonic, sleep=clock.sleep) # 1 req / 2s
|
|
|
|
# First acquire: initial tokens=0.5, sleeps 1s to reach 1.0, then succeeds.
|
|
# Second acquire: tokens=0, sleeps 2s to reach 1.0, then succeeds.
|
|
limiter.acquire()
|
|
limiter.acquire()
|
|
|
|
total_sleep = sum(clock.sleep_calls)
|
|
assert total_sleep == pytest.approx(3.0, abs=0.1)
|
|
|
|
|
|
def test_burst_tokens_consumed_without_sleeping():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(5, clock=clock.monotonic, sleep=clock.sleep)
|
|
|
|
for _ in range(5):
|
|
limiter.acquire()
|
|
|
|
assert clock.sleep_calls == []
|
|
|
|
|
|
def test_sleep_called_when_tokens_exhausted():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(5, clock=clock.monotonic, sleep=clock.sleep)
|
|
|
|
for _ in range(5):
|
|
limiter.acquire()
|
|
|
|
limiter.acquire()
|
|
assert len(clock.sleep_calls) == 1
|
|
assert clock.sleep_calls[0] == pytest.approx(0.2, abs=0.01)
|
|
|
|
|
|
def test_total_sleep_for_sustained_rate():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(10, clock=clock.monotonic, sleep=clock.sleep)
|
|
|
|
for _ in range(20):
|
|
limiter.acquire()
|
|
|
|
total_sleep = sum(clock.sleep_calls)
|
|
assert total_sleep == pytest.approx(1.0, abs=0.01)
|
|
|
|
|
|
def test_tokens_refill_after_idle():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(10, clock=clock.monotonic, sleep=clock.sleep)
|
|
|
|
for _ in range(10):
|
|
limiter.acquire()
|
|
|
|
clock._now += 1.0
|
|
|
|
sleep_before = len(clock.sleep_calls)
|
|
for _ in range(10):
|
|
limiter.acquire()
|
|
|
|
assert clock.sleep_calls[sleep_before:] == []
|
|
|
|
|
|
def test_partial_refill():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(10, clock=clock.monotonic, sleep=clock.sleep)
|
|
|
|
for _ in range(10):
|
|
limiter.acquire()
|
|
|
|
clock._now += 0.5
|
|
|
|
sleep_before = len(clock.sleep_calls)
|
|
for _ in range(5):
|
|
limiter.acquire()
|
|
|
|
assert clock.sleep_calls[sleep_before:] == []
|
|
|
|
limiter.acquire()
|
|
assert len(clock.sleep_calls) == sleep_before + 1
|
|
|
|
|
|
def test_noop_acquire_does_nothing():
|
|
limiter = NoOpRateLimiter()
|
|
for _ in range(1000):
|
|
limiter.acquire()
|
|
|
|
|
|
# ── _make_rate_limiter / _parse_rate_limit tests ──
|
|
|
|
|
|
def test_make_rate_limiter_positive_rate():
|
|
assert isinstance(_make_rate_limiter(10.0), RPSRateLimiter)
|
|
|
|
|
|
def test_make_rate_limiter_zero_returns_noop():
|
|
assert isinstance(_make_rate_limiter(0.0), NoOpRateLimiter)
|
|
|
|
|
|
def test_make_rate_limiter_none_returns_noop():
|
|
assert isinstance(_make_rate_limiter(None), NoOpRateLimiter)
|
|
|
|
|
|
def test_make_rate_limiter_adaptive():
|
|
limiter = _make_rate_limiter(10.0, adaptive=True)
|
|
assert isinstance(limiter, RPSRateLimiter)
|
|
assert limiter._adaptive is True
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("raw", "expected_rps", "expected_adaptive"),
|
|
[
|
|
("auto", AUTO_INITIAL_RPS, True),
|
|
("AUTO", AUTO_INITIAL_RPS, True),
|
|
(" Auto ", AUTO_INITIAL_RPS, True),
|
|
("25", 25.0, False),
|
|
("0", None, False),
|
|
(None, None, False),
|
|
],
|
|
)
|
|
def test_parse_rate_limit(raw, expected_rps, expected_adaptive):
|
|
rps, adaptive = _parse_rate_limit(raw)
|
|
assert rps == expected_rps
|
|
assert adaptive == expected_adaptive
|
|
|
|
|
|
# ── is_rate_limit_error tests ──
|
|
|
|
|
|
class _FakeRateLimitError(Exception):
|
|
pass
|
|
|
|
|
|
_FakeRateLimitError.__name__ = "RateLimitError"
|
|
|
|
|
|
class _FakeStatusCodeError(Exception):
|
|
def __init__(self, status_code):
|
|
self.status_code = status_code
|
|
super().__init__(f"HTTP {status_code}")
|
|
|
|
|
|
class _FakeResponseError(Exception):
|
|
def __init__(self, status_code):
|
|
self.response = type("R", (), {"status_code": status_code})()
|
|
super().__init__(f"HTTP {status_code}")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("exc", "expected"),
|
|
[
|
|
(_FakeRateLimitError("rate limit"), True),
|
|
(_FakeStatusCodeError(429), True),
|
|
(_FakeResponseError(429), True),
|
|
(Exception("Error 429: too many requests"), True),
|
|
(Exception("rate limit exceeded"), True),
|
|
(_FakeStatusCodeError(500), False),
|
|
(_FakeResponseError(500), False),
|
|
(Exception("something else entirely"), False),
|
|
(ValueError("bad value"), False),
|
|
],
|
|
)
|
|
def test_is_rate_limit_error(exc, expected):
|
|
assert is_rate_limit_error(exc) == expected
|
|
|
|
|
|
# ── AIMD tests ──
|
|
|
|
|
|
def test_throttle_halves_rate():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(10.0, adaptive=True, clock=clock.monotonic, sleep=clock.sleep)
|
|
|
|
limiter.report_throttle()
|
|
assert limiter._rps == pytest.approx(5.0)
|
|
|
|
|
|
def test_throttle_respects_floor():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(2.0, adaptive=True, clock=clock.monotonic, sleep=clock.sleep)
|
|
|
|
# First throttle: 2.0 * 0.5 = 1.0
|
|
limiter.report_throttle()
|
|
assert limiter._rps == pytest.approx(1.0)
|
|
|
|
# Second throttle (after cooldown): should stay at floor 1.0
|
|
clock._now += 10.0
|
|
limiter.report_throttle()
|
|
assert limiter._rps == pytest.approx(1.0)
|
|
|
|
|
|
def test_throttle_cooldown_coalesces_rapid_signals():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(10.0, adaptive=True, clock=clock.monotonic, sleep=clock.sleep)
|
|
|
|
limiter.report_throttle()
|
|
assert limiter._rps == pytest.approx(5.0)
|
|
|
|
# Within cooldown window — should be ignored
|
|
clock._now += 1.0
|
|
limiter.report_throttle()
|
|
assert limiter._rps == pytest.approx(5.0)
|
|
|
|
# After cooldown — should take effect
|
|
clock._now += 10.0
|
|
limiter.report_throttle()
|
|
assert limiter._rps == pytest.approx(2.5)
|
|
|
|
|
|
def test_success_restores_rate():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(10.0, adaptive=True, clock=clock.monotonic, sleep=clock.sleep)
|
|
|
|
limiter.report_throttle()
|
|
assert limiter._rps == pytest.approx(5.0)
|
|
|
|
# Repeatedly report success — rate should climb back past initial
|
|
for _ in range(100):
|
|
limiter.report_success()
|
|
|
|
assert limiter._rps > 10.0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("multiplier", "expected_ceiling"),
|
|
[(5.0, 50.0), (3.0, 30.0)],
|
|
)
|
|
def test_success_climbs_to_multiplier_ceiling(multiplier, expected_ceiling):
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(
|
|
10.0,
|
|
adaptive=True,
|
|
max_rps_multiplier=multiplier,
|
|
clock=clock.monotonic,
|
|
sleep=clock.sleep,
|
|
)
|
|
for _ in range(10000):
|
|
limiter.report_success()
|
|
assert limiter._rps == pytest.approx(expected_ceiling)
|
|
|
|
|
|
def test_adaptive_false_ignores_throttle_and_success():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(10.0, adaptive=False, clock=clock.monotonic, sleep=clock.sleep)
|
|
|
|
limiter.report_throttle()
|
|
assert limiter._rps == pytest.approx(10.0)
|
|
|
|
limiter.report_success()
|
|
assert limiter._rps == pytest.approx(10.0)
|
|
|
|
|
|
# ── call_with_retry tests ──
|
|
|
|
|
|
def test_call_with_retry_success():
|
|
sleep_calls = []
|
|
limiter = NoOpRateLimiter()
|
|
result = call_with_retry(lambda: 42, limiter, max_retries=3, sleep=sleep_calls.append)
|
|
assert result == 42
|
|
assert sleep_calls == []
|
|
|
|
|
|
def test_call_with_retry_retries_on_429_then_succeeds():
|
|
sleep_calls = []
|
|
limiter = NoOpRateLimiter()
|
|
attempts = []
|
|
|
|
def flaky_fn():
|
|
attempts.append(1)
|
|
if len(attempts) < 3:
|
|
raise _FakeRateLimitError("rate limited")
|
|
return "ok"
|
|
|
|
result = call_with_retry(flaky_fn, limiter, max_retries=3, sleep=sleep_calls.append)
|
|
assert result == "ok"
|
|
assert len(attempts) == 3
|
|
# Two retries with exponential backoff: 2^0=1, 2^1=2
|
|
assert sleep_calls == [1, 2]
|
|
|
|
|
|
def test_call_with_retry_non_429_propagates_immediately():
|
|
sleep_calls = []
|
|
limiter = NoOpRateLimiter()
|
|
|
|
def always_raises():
|
|
raise ValueError("bad input")
|
|
|
|
with pytest.raises(ValueError, match="bad input"):
|
|
call_with_retry(always_raises, limiter, max_retries=3, sleep=sleep_calls.append)
|
|
assert sleep_calls == []
|
|
|
|
|
|
def test_call_with_retry_exhausted_retries():
|
|
sleep_calls = []
|
|
limiter = NoOpRateLimiter()
|
|
|
|
def always_rate_limited():
|
|
raise _FakeRateLimitError("rate limited")
|
|
|
|
with pytest.raises(_FakeRateLimitError, match="rate limited"):
|
|
call_with_retry(always_rate_limited, limiter, max_retries=2, sleep=sleep_calls.append)
|
|
# 3 attempts total (initial + 2 retries), 2 sleeps
|
|
assert len(sleep_calls) == 2
|
|
|
|
|
|
def test_call_with_retry_reports_throttle_and_success():
|
|
clock = FakeClock()
|
|
limiter = RPSRateLimiter(10.0, adaptive=True, clock=clock.monotonic, sleep=clock.sleep)
|
|
attempts = []
|
|
|
|
def flaky_fn():
|
|
attempts.append(1)
|
|
if len(attempts) == 1:
|
|
raise _FakeRateLimitError("rate limited")
|
|
return "ok"
|
|
|
|
result = call_with_retry(flaky_fn, limiter, max_retries=3, sleep=clock.sleep)
|
|
assert result == "ok"
|
|
# After throttle: 10.0 * 0.5 = 5.0, then success bumps it back up slightly
|
|
assert limiter._rps < 10.0
|
|
|
|
|
|
# ── eval_retry_context tests ──
|
|
|
|
|
|
def _retry_flags_active():
|
|
"""Check that both downstream retry-suppression flags are set."""
|
|
return is_litellm_rate_limit_retries_disabled() and is_429_retry_disabled()
|
|
|
|
|
|
def test_eval_retry_context_sets_and_resets():
|
|
assert not _retry_flags_active()
|
|
|
|
with eval_retry_context():
|
|
assert _retry_flags_active()
|
|
|
|
assert not _retry_flags_active()
|
|
|
|
|
|
def test_eval_retry_context_nests():
|
|
assert not _retry_flags_active()
|
|
|
|
with eval_retry_context():
|
|
assert _retry_flags_active()
|
|
with eval_retry_context():
|
|
assert _retry_flags_active()
|
|
assert _retry_flags_active()
|
|
|
|
assert not _retry_flags_active()
|
|
|
|
|
|
def test_litellm_retry_policy_disables_rate_limit_retries_when_flag_set():
|
|
with eval_retry_context():
|
|
policy = _get_litellm_retry_policy(3)
|
|
assert policy.RateLimitErrorRetries == 0
|
|
assert policy.TimeoutErrorRetries == 3
|