diff --git a/docs/inference.md b/docs/inference.md index b0242c27..0c8b9edb 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -46,6 +46,53 @@ pipeline = mdr.read_jsonl("input.jsonl").map_async( Set `OPENAI_API_KEY` in the worker environment before execution. For cloud jobs, pass it through `secrets={"OPENAI_API_KEY": None}`. +### Request limits + +Refiner uses adaptive request limiting by default. It starts below the configured +maximum, grows after clean success windows, and backs off when an endpoint +returns HTTP `429`. + +```python +import refiner as mdr + +endpoint = mdr.inference.OpenAIEndpointProvider( + base_url="https://api.openai.com", + model="gpt-5-mini", +) + +pipeline = mdr.read_jsonl("input.jsonl").map_async( + mdr.inference.generate( + fn=summarize, + provider=endpoint, + rate_limit=mdr.inference.AdaptiveRateLimit( + max_concurrency=256, + initial_concurrency=16, + ), + ), + max_in_flight=256, +) +``` + +When `rate_limit` is omitted or set to `None`, Refiner uses +`AdaptiveRateLimit()` with its defaults. Growth uses a shrinking multiplier: the +default starts at `2.0`, subtracts `0.1` after each successful growth window, +and never drops below `1.05`. A rate-limit response halves the current request +concurrency by default and subtracts `0.2` from the growth multiplier. If the +endpoint sends `Retry-After`, Refiner waits before starting more requests. + +Use `StaticRateLimit` when you want the old fixed-concurrency behavior: + +```python +pipeline = mdr.read_jsonl("input.jsonl").map_async( + mdr.inference.generate( + fn=summarize, + provider=endpoint, + rate_limit=mdr.inference.StaticRateLimit(max_concurrency=64), + ), + max_in_flight=64, +) +``` + ### Refiner managed VLLM runtime Use `VLLMProvider` when launching on Refiner Cloud and you want the platform to start and manage a VLLM server for you. @@ -96,3 +143,9 @@ Only the following models are currently supported. If you are missing one, pleas - Direct endpoint example: [`examples/inference_endpoint.py`](../examples/inference_endpoint.py) - VLLM-backed example: [`examples/inference_vllm.py`](../examples/inference_vllm.py) + +## Internal Notes + +Adaptive request limiting is process-local. In multi-worker jobs, each worker +adapts independently to the same provider. This reduces per-worker bursts but is +not a global quota coordinator across the full job. diff --git a/examples/lerobot/sarm_annotation.py b/examples/lerobot/sarm_annotation.py index 8da35c6d..631a8c3b 100644 --- a/examples/lerobot/sarm_annotation.py +++ b/examples/lerobot/sarm_annotation.py @@ -187,7 +187,9 @@ async def annotate_dense_subtasks(row, generate): fn=annotate_dense_subtasks, provider=PROVIDER, default_generation_params={"temperature": 0.1}, - max_concurrent_requests=MAX_IN_FLIGHT, + rate_limit=mdr.inference.AdaptiveRateLimit( + max_concurrency=MAX_IN_FLIGHT + ), ), max_in_flight=MAX_IN_FLIGHT, ) diff --git a/src/refiner/inference/__init__.py b/src/refiner/inference/__init__.py index f524752d..2231f442 100644 --- a/src/refiner/inference/__init__.py +++ b/src/refiner/inference/__init__.py @@ -1,10 +1,14 @@ from refiner.inference.generate import generate -from refiner.inference.client import InferenceResponse +from refiner.inference.client import GenerationRateLimitError, InferenceResponse from refiner.inference.providers import OpenAIEndpointProvider, VLLMProvider +from refiner.inference.rate_limit import AdaptiveRateLimit, StaticRateLimit __all__ = [ "generate", + "AdaptiveRateLimit", + "GenerationRateLimitError", "InferenceResponse", "OpenAIEndpointProvider", + "StaticRateLimit", "VLLMProvider", ] diff --git a/src/refiner/inference/_runtime.py b/src/refiner/inference/_runtime.py index ee3e4f19..821f8143 100644 --- a/src/refiner/inference/_runtime.py +++ b/src/refiner/inference/_runtime.py @@ -2,11 +2,19 @@ import asyncio import inspect +import os from collections.abc import Awaitable, Callable, Mapping +from dataclasses import asdict from typing import Any, TypeAlias, cast -from refiner.inference.client import _OpenAIEndpointClient +from refiner.inference.client import GenerationRateLimitError, _OpenAIEndpointClient from refiner.inference.providers import OpenAIEndpointProvider, VLLMProvider +from refiner.inference.rate_limit import ( + AdaptiveRateLimit, + AdaptiveRateLimiter, + RateLimit, + StaticRateLimit, +) from refiner.pipeline.data.row import Row from refiner.pipeline.steps import MapResult from refiner.services import VLLMRuntimeServiceBinding @@ -31,14 +39,28 @@ def inference_map( defaults: Mapping[str, Any] | None, defaults_key: str | None = None, max_concurrent_requests: int = 256, + rate_limit: RateLimit | None = None, + rate_limit_key: str | None = None, call: ClientCall, record: Callable[[Row, Any], None] | None = None, ) -> Callable[[Row], Awaitable[MapResult]]: if max_concurrent_requests <= 0: raise ValueError("max_concurrent_requests must be > 0") + resolved_rate_limit = rate_limit or StaticRateLimit( + max_concurrency=max_concurrent_requests + ) client: _OpenAIEndpointClient | None = None client_lock = asyncio.Lock() - semaphore = asyncio.Semaphore(max_concurrent_requests) + adaptive_limiter: AdaptiveRateLimiter | None = ( + AdaptiveRateLimiter(resolved_rate_limit) + if isinstance(resolved_rate_limit, AdaptiveRateLimit) + else None + ) + semaphore = ( + None + if adaptive_limiter is not None + else asyncio.Semaphore(resolved_rate_limit.max_concurrency) + ) gauges_registered = False waiting_requests = 0 running_requests = 0 @@ -49,6 +71,12 @@ def _register_metrics() -> None: return register_gauge("waiting_requests", lambda: waiting_requests, unit="requests") register_gauge("running_requests", lambda: running_requests, unit="requests") + if adaptive_limiter is not None: + register_gauge( + "adaptive_concurrency", + lambda: adaptive_limiter.limit, + unit="requests", + ) gauges_registered = True async def _client() -> _OpenAIEndpointClient: @@ -59,7 +87,10 @@ async def _client() -> _OpenAIEndpointClient: if client is not None: return client if isinstance(provider, OpenAIEndpointProvider): - client = _OpenAIEndpointClient(base_url=provider.base_url) + client = _OpenAIEndpointClient( + base_url=provider.base_url, + api_key=os.environ.get(provider.api_key_env_var), + ) else: service_name = provider.service_definition().name service_manager = get_active_service_manager() @@ -88,17 +119,37 @@ async def _request(row: Row, payload: Mapping[str, Any]) -> Any: } resolved_client = await _client() waiting_requests += 1 - await semaphore.acquire() - waiting_requests -= 1 + acquired_adaptive = False + try: + if adaptive_limiter is not None: + await adaptive_limiter.acquire() + acquired_adaptive = True + else: + assert semaphore is not None + await semaphore.acquire() + finally: + waiting_requests -= 1 running_requests += 1 try: response = await call(resolved_client, request_payload) + except GenerationRateLimitError as err: + if adaptive_limiter is not None: + await adaptive_limiter.record_rate_limit(err.retry_after_seconds) + row.log_throughput("rate_limited_requests", 1, unit="requests") + row.log_throughput("failed_requests", 1, unit="requests") + raise except Exception: row.log_throughput("failed_requests", 1, unit="requests") raise finally: running_requests -= 1 - semaphore.release() + if acquired_adaptive: + assert adaptive_limiter is not None + await adaptive_limiter.release() + if semaphore is not None: + semaphore.release() + if adaptive_limiter is not None: + await adaptive_limiter.record_success() row.log_throughput("successful_requests", 1, unit="requests") if record is not None: record(row, response) @@ -114,8 +165,16 @@ async def _wrapped(row: Row) -> MapResult: args: dict[str, Any] = { "fn": fn, "provider": provider.to_builtin_args(), - "max_concurrent_requests": max_concurrent_requests, } + if rate_limit_key is None: + args["max_concurrent_requests"] = resolved_rate_limit.max_concurrency + else: + args[rate_limit_key] = { + "type": "adaptive" + if isinstance(resolved_rate_limit, AdaptiveRateLimit) + else "static", + **asdict(resolved_rate_limit), + } if defaults_key is not None: args[defaults_key] = dict(defaults or {}) setattr( diff --git a/src/refiner/inference/client.py b/src/refiner/inference/client.py index 0fcab4e9..2fa7cab3 100644 --- a/src/refiner/inference/client.py +++ b/src/refiner/inference/client.py @@ -1,9 +1,11 @@ from __future__ import annotations import asyncio +import email.utils import logging import os import random +import time from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import Any @@ -17,6 +19,12 @@ logger = logging.getLogger(__name__) +class GenerationRateLimitError(RuntimeError): + def __init__(self, message: str, *, retry_after_seconds: float | None = None): + super().__init__(message) + self.retry_after_seconds = retry_after_seconds + + @dataclass(frozen=True, slots=True) class InferenceResponse: text: str @@ -124,6 +132,11 @@ async def _post_json( message = f"{operation} request failed with HTTP {err.response.status_code}" if detail: message = f"{message}: {detail}" + if err.response.status_code in {429, 503}: + raise GenerationRateLimitError( + message, + retry_after_seconds=_retry_after_seconds(err.response), + ) from err raise RuntimeError(message) from err return response.json() @@ -198,4 +211,23 @@ def _retry_delay_seconds(attempt: int) -> float: return base_delay * (1.0 + jitter) -__all__ = ["InferenceResponse"] +def _retry_after_seconds(response: httpx.Response) -> float | None: + raw = response.headers.get("Retry-After") + if raw is None: + return None + value = raw.strip() + if not value: + return None + try: + return max(0.0, float(value)) + except ValueError: + try: + parsed = email.utils.parsedate_to_datetime(value) + except (TypeError, ValueError): + return None + if parsed is None: + return None + return max(0.0, parsed.timestamp() - time.time()) + + +__all__ = ["GenerationRateLimitError", "InferenceResponse"] diff --git a/src/refiner/inference/generate.py b/src/refiner/inference/generate.py index d6c41c83..ddfcfb1d 100644 --- a/src/refiner/inference/generate.py +++ b/src/refiner/inference/generate.py @@ -6,6 +6,7 @@ from refiner.inference._runtime import inference_map from refiner.inference.client import InferenceResponse, _OpenAIEndpointClient from refiner.inference.providers import OpenAIEndpointProvider, VLLMProvider +from refiner.inference.rate_limit import AdaptiveRateLimit, RateLimit from refiner.pipeline.data.row import Row from refiner.pipeline.steps import MapResult @@ -19,7 +20,7 @@ def generate( fn: InferenceFn, provider: OpenAIEndpointProvider | VLLMProvider, default_generation_params: Mapping[str, Any] | None = None, - max_concurrent_requests: int = 256, + rate_limit: RateLimit | None = None, ) -> Callable[[Row], Awaitable[MapResult]]: return inference_map( name="inference.generate", @@ -27,7 +28,8 @@ def generate( provider=provider, defaults=default_generation_params, defaults_key="default_generation_params", - max_concurrent_requests=max_concurrent_requests, + rate_limit=rate_limit or AdaptiveRateLimit(), + rate_limit_key="rate_limit", call=_generate, record=_record_usage, ) diff --git a/src/refiner/inference/providers.py b/src/refiner/inference/providers.py index 1d9537da..d613da52 100644 --- a/src/refiner/inference/providers.py +++ b/src/refiner/inference/providers.py @@ -8,12 +8,15 @@ class OpenAIEndpointProvider: base_url: str model: str + api_key_env_var: str = "OPENAI_API_KEY" def __post_init__(self) -> None: if not self.base_url.strip(): raise ValueError("base_url must be non-empty") if not self.model.strip(): raise ValueError("model must be non-empty") + if not self.api_key_env_var.strip(): + raise ValueError("api_key_env_var must be non-empty") def service_definition(self) -> None: return None @@ -23,6 +26,7 @@ def to_builtin_args(self) -> dict[str, object]: "type": "openai_endpoint", "base_url": self.base_url, "model": self.model, + "api_key_env_var": self.api_key_env_var, } return payload diff --git a/src/refiner/inference/rate_limit.py b/src/refiner/inference/rate_limit.py new file mode 100644 index 00000000..90188160 --- /dev/null +++ b/src/refiner/inference/rate_limit.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import asyncio +import math +import time +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class StaticRateLimit: + """Fixed request concurrency settings for inference providers.""" + + max_concurrency: int = 256 + + def __post_init__(self) -> None: + if self.max_concurrency <= 0: + raise ValueError("max_concurrency must be > 0") + + +@dataclass(frozen=True, slots=True) +class AdaptiveRateLimit: + """Adaptive request concurrency settings for inference providers.""" + + max_concurrency: int = 256 + min_concurrency: int = 1 + initial_concurrency: int | None = None + initial_growth_multiplier: float = 2.0 + min_growth_multiplier: float = 1.05 + growth_multiplier_step: float = 0.1 + rate_limit_decrease_factor: float = 0.5 + success_window_requests: int = 50 + default_cooldown_seconds: float = 1.0 + + def __post_init__(self) -> None: + if ( + self.max_concurrency <= 0 + or self.min_concurrency <= 0 + or self.success_window_requests <= 0 + ): + raise ValueError("concurrency and success window must be > 0") + if self.max_concurrency < self.min_concurrency: + raise ValueError("max_concurrency must be >= min_concurrency") + if self.initial_concurrency is not None: + if self.initial_concurrency < self.min_concurrency: + raise ValueError("initial_concurrency must be >= min_concurrency") + if not 1 <= self.min_growth_multiplier <= self.initial_growth_multiplier: + raise ValueError("growth multipliers must satisfy 1 <= min <= initial") + if self.growth_multiplier_step < 0 or self.default_cooldown_seconds < 0: + raise ValueError("growth step and cooldown must be >= 0") + if not 0 < self.rate_limit_decrease_factor < 1: + raise ValueError("rate_limit_decrease_factor must be > 0 and < 1") + + +class AdaptiveRateLimiter: + """Async limiter that probes provider capacity and backs off on rate limits.""" + + def __init__(self, config: AdaptiveRateLimit) -> None: + self._config = config + initial = config.initial_concurrency + if initial is None: + initial = min(16, config.max_concurrency) + self._limit = min(max(initial, config.min_concurrency), config.max_concurrency) + self._growth_multiplier = config.initial_growth_multiplier + self._running = 0 + self._successes = 0 + self._cooldown_until = 0.0 + self._condition = asyncio.Condition() + + @property + def limit(self) -> int: + return self._limit + + async def acquire(self) -> None: + async with self._condition: + while True: + cooldown_remaining = self._cooldown_until - time.monotonic() + if cooldown_remaining > 0: + try: + await asyncio.wait_for( + self._condition.wait(), + timeout=cooldown_remaining, + ) + except TimeoutError: + pass + continue + if self._running < self._limit: + self._running += 1 + return + await self._condition.wait() + + async def release(self) -> None: + async with self._condition: + self._running -= 1 + self._condition.notify_all() + + async def record_success(self) -> None: + async with self._condition: + self._successes += 1 + if self._successes < self._config.success_window_requests: + return + self._successes = 0 + next_limit = min( + self._config.max_concurrency, + max(self._limit + 1, math.ceil(self._limit * self._growth_multiplier)), + ) + if next_limit <= self._limit: + return + self._limit = next_limit + self._growth_multiplier = self._reduced_multiplier( + self._config.growth_multiplier_step + ) + self._condition.notify_all() + + async def record_rate_limit(self, retry_after_seconds: float | None = None) -> None: + async with self._condition: + self._successes = 0 + self._limit = max( + self._config.min_concurrency, + math.floor(self._limit * self._config.rate_limit_decrease_factor), + ) + self._growth_multiplier = self._reduced_multiplier( + self._config.growth_multiplier_step * 2 + ) + cooldown = ( + retry_after_seconds + if retry_after_seconds is not None + else self._config.default_cooldown_seconds + ) + if cooldown > 0: + self._cooldown_until = max( + self._cooldown_until, + time.monotonic() + cooldown, + ) + self._condition.notify_all() + + def _reduced_multiplier(self, step: float) -> float: + return max( + self._config.min_growth_multiplier, + self._growth_multiplier - step, + ) + + +RateLimit = StaticRateLimit | AdaptiveRateLimit + + +__all__ = ["AdaptiveRateLimit", "AdaptiveRateLimiter", "RateLimit", "StaticRateLimit"] diff --git a/tests/test_inference.py b/tests/test_inference.py index 925a562b..70b6ebf0 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -10,6 +10,7 @@ import refiner as mdr from refiner.inference import ( + GenerationRateLimitError, InferenceResponse, OpenAIEndpointProvider, VLLMProvider, @@ -66,6 +67,15 @@ def test_openai_endpoint_requires_non_empty_model() -> None: OpenAIEndpointProvider(base_url="https://api.example.com", model=" ") +def test_openai_endpoint_requires_non_empty_api_key_env_var() -> None: + with pytest.raises(ValueError, match="api_key_env_var must be non-empty"): + OpenAIEndpointProvider( + base_url="https://api.example.com", + model="gpt-test", + api_key_env_var=" ", + ) + + def test_inference_generate_invokes_user_fn_and_merges_default_params( monkeypatch, ) -> None: @@ -163,6 +173,56 @@ async def _invoke() -> object: assert seen["headers"] == {"Authorization": "Bearer secret"} +def test_openai_endpoint_accepts_custom_api_key_env_var(monkeypatch) -> None: + seen: dict[str, object] = {} + + class _FakeResponse: + def raise_for_status(self) -> None: + return None + + def json(self) -> Mapping[str, object]: + return { + "choices": [ + { + "text": "ok", + "finish_reason": "stop", + } + ], + "usage": {}, + } + + class _FakeAsyncClient: + def __init__(self, *, base_url, headers, timeout): + seen["headers"] = dict(headers) + + async def post(self, path, *, json): + return _FakeResponse() + + monkeypatch.setattr(openai_module.httpx, "AsyncClient", _FakeAsyncClient) + monkeypatch.setenv("GOOGLE_API_KEY", "google-secret") + + async def _inference_fn(row, generate): + response = await generate({"prompt": row["prompt"]}) + return {"output": response.text} + + infer = mdr.inference.generate( + fn=_inference_fn, + provider=OpenAIEndpointProvider( + base_url="https://api.example.com", + model="gpt-test", + api_key_env_var="GOOGLE_API_KEY", + ), + ) + + async def _invoke() -> object: + return await infer(DictRow({"prompt": "hi"})) + + result = asyncio.run(_invoke()) + + assert result == {"output": "ok"} + assert seen["headers"] == {"Authorization": "Bearer google-secret"} + + def test_openai_endpoint_preserves_base_url_path_prefix(monkeypatch) -> None: seen: dict[str, object] = {} @@ -350,13 +410,14 @@ def test_openai_endpoint_warns_on_null_chat_content(caplog) -> None: ) -def test_openai_endpoint_does_not_retry_on_http_503(monkeypatch) -> None: +def test_openai_endpoint_treats_http_503_as_rate_limit(monkeypatch) -> None: seen: dict[str, int] = {"calls": 0, "sleeps": 0} request = httpx.Request("POST", "https://api.example.com/v1/chat/completions") response = httpx.Response( 503, request=request, + headers={"Retry-After": "2.5"}, json={"error": {"message": "Service unavailable"}}, ) @@ -376,7 +437,7 @@ async def _fake_sleep(delay: float) -> None: monkeypatch.setattr(openai_module.httpx, "AsyncClient", _FakeAsyncClient) monkeypatch.setattr(openai_module.asyncio, "sleep", _fake_sleep) - with pytest.raises(RuntimeError, match="generation request failed with HTTP 503"): + with pytest.raises(GenerationRateLimitError) as error: asyncio.run( openai_module._OpenAIEndpointClient( base_url="https://api.example.com", @@ -388,9 +449,44 @@ async def _fake_sleep(delay: float) -> None: ) ) + assert error.value.retry_after_seconds == 2.5 assert seen == {"calls": 1, "sleeps": 0} +def test_openai_endpoint_raises_rate_limit_with_retry_after(monkeypatch) -> None: + request = httpx.Request("POST", "https://api.example.com/v1/chat/completions") + response = httpx.Response( + 429, + request=request, + headers={"Retry-After": "2.5"}, + json={"error": {"message": "Rate limited"}}, + ) + + class _FakeAsyncClient: + def __init__(self, *, base_url, headers, timeout): + del base_url, headers, timeout + + async def post(self, path, *, json): + del path, json + return response + + monkeypatch.setattr(openai_module.httpx, "AsyncClient", _FakeAsyncClient) + + with pytest.raises(GenerationRateLimitError) as error: + asyncio.run( + openai_module._OpenAIEndpointClient( + base_url="https://api.example.com", + ).generate( + { + "model": "gpt-test", + "messages": [{"role": "user", "content": "hello"}], + } + ) + ) + + assert error.value.retry_after_seconds == 2.5 + + def test_openai_endpoint_provider_builtin_args_do_not_include_api_key() -> None: provider = OpenAIEndpointProvider( base_url="https://api.example.com", @@ -401,6 +497,7 @@ def test_openai_endpoint_provider_builtin_args_do_not_include_api_key() -> None: "type": "openai_endpoint", "base_url": "https://api.example.com", "model": "gpt-test", + "api_key_env_var": "OPENAI_API_KEY", } diff --git a/tests/test_inference_rate_limit.py b/tests/test_inference_rate_limit.py new file mode 100644 index 00000000..0fe81f10 --- /dev/null +++ b/tests/test_inference_rate_limit.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import asyncio + +from refiner.inference.rate_limit import AdaptiveRateLimit, AdaptiveRateLimiter + + +def test_adaptive_rate_limiter_grows_with_shrinking_multiplier() -> None: + limiter = AdaptiveRateLimiter( + AdaptiveRateLimit( + max_concurrency=100, + initial_concurrency=10, + success_window_requests=1, + ) + ) + + asyncio.run(limiter.record_success()) + assert limiter.limit == 20 + + asyncio.run(limiter.record_success()) + assert limiter.limit == 38 + + +def test_adaptive_rate_limiter_halves_on_rate_limit() -> None: + limiter = AdaptiveRateLimiter( + AdaptiveRateLimit( + max_concurrency=100, + initial_concurrency=80, + success_window_requests=1, + default_cooldown_seconds=0, + ) + ) + + asyncio.run(limiter.record_rate_limit()) + + assert limiter.limit == 40 + + asyncio.run(limiter.record_success()) + assert limiter.limit == 72 + + +def test_adaptive_rate_limiter_never_drops_below_minimum() -> None: + limiter = AdaptiveRateLimiter( + AdaptiveRateLimit( + max_concurrency=100, + min_concurrency=3, + initial_concurrency=4, + default_cooldown_seconds=0, + ) + ) + + asyncio.run(limiter.record_rate_limit()) + asyncio.run(limiter.record_rate_limit()) + + assert limiter.limit == 3 + + +def test_adaptive_rate_limiter_respects_maximum() -> None: + limiter = AdaptiveRateLimiter( + AdaptiveRateLimit( + max_concurrency=100, + initial_concurrency=90, + success_window_requests=1, + ) + ) + + asyncio.run(limiter.record_success()) + + assert limiter.limit == 100