# UDF Rate Limiter Design

## Problem

Geneva's built-in Gemini UDFs (`gemini_udf`, `gemini_embedding_udf`) have no client-side rate limiting. Each applier actor fires API requests as fast as possible. When `concurrency × intra_applier_concurrency` exceeds the API provider's rate limit, workers saturate the endpoint with requests, receive 429 (rate limit) responses, and burn their retry budget on backoff — wasting time and compute.

### Current Behavior

```
concurrency=8 appliers, each doing ~30 RPM
  → ~240 RPM aggregate
  → Tier 1 gemini-2.5-flash allows 1,000 RPM → OK

concurrency=40 appliers, each doing ~30 RPM
  → ~1,200 RPM aggregate
  → Tier 1 gemini-2.5-flash allows 1,000 RPM → 429s, retries, wasted wall time
```

The retry policy (`Retry(ClientError, ServerError, max_attempts=7, backoff="exponential")`) uses `wait_random_exponential(min=1, max=60)` — a single failing request can spend up to ~123 seconds in backoff before exhausting retries. This masks the root cause (rate limiting) and wastes worker time.

### When Does the Status Quo Break Down?

To quantify when retries become insufficient, we modeled the steady-state equilibrium between N concurrent workers, the API rate limit R, and the exponential backoff. The model accounts for the feedback loop: as rejection rate rises, workers spend more time in backoff, which reduces aggregate load, which lowers rejection rate.

**Assumptions**: T_success = 2s per Gemini API call, T_fail = 0.1s for a 429 response, `wait_random_exponential(min=1, max=60)` backoff, 7 max attempts per row.

#### Row failure probability vs. concurrency (gemini-2.5-flash, Tier 1, 1000 RPM)

| Workers | p (reject rate) | P(row fails all 7) | Effective RPM/worker | Aggregate RPM |
|--------:|----------------:|--------------------:|---------------------:|--------------:|
| 8 | 0.00 | ~0 | 30.0 | 240 |
| 16 | 0.00 | ~0 | 30.0 | 480 |
| 32 | 0.00 | ~0 | 30.0 | 960 |
| 50 | 0.30 | 0.024% | 28.8 | 1,438 |
| 64 | 0.40 | 0.16% | 25.9 | 1,660 |
| 80 | 0.46 | 0.45% | 23.2 | 1,859 |
| **100** | **0.52** | **0.98%** | **20.7** | **2,069** |
| 120 | 0.56 | 1.66% | 18.8 | 2,255 |
| 150 | 0.60 | 2.86% | 16.7 | 2,511 |
| 200 | 0.66 | 5.19% | 14.5 | 2,901 |

At **~100 workers** the 1% threshold is crossed — 1 in every ~100 rows will exhaust all 7 retries and fail permanently. Note that the backoff self-regulation keeps the aggregate RPM well above the 1000 RPM limit (workers slow themselves down), but the rejection rate still grows because more workers compete for the same quota.

#### Workers needed for >1% row failure (all models/tiers)

| Model | Tier | RPM limit | Workers for >1% failure |
|-------|------|----------:|------------------------:|
| gemini-2.5-flash | Free | 10 | **2** |
| gemini-2.5-flash | Tier 1 | 1,000 | **101** |
| gemini-2.5-flash | Tier 2 | 2,000 | **202** |
| gemini-2.5-flash | Tier 3 | 4,000 | **403** |
| gemini-2.5-pro | Free | 5 | **1** |
| gemini-2.5-pro | Tier 1 | 150 | **16** |
| gemini-2.5-pro | Tier 2 | 1,000 | **101** |
| gemini-2.5-flash-lite | Free | 15 | **2** |
| gemini-2.5-flash-lite | Tier 1 | 300 | **31** |
| gemini-embedding-001 | Free | 100 | **11** |

The relationship is roughly **N ≈ R / 10** — you can run about 1 worker per 10 RPM of quota before retries start failing.

#### Free tier is broken without rate limiting

Free tier limits are so low that even `concurrency=1` is problematic:

| Model | RPM | Workers | P(row fails) | ~1 in |
|-------|----:|--------:|-------------:|------:|
| gemini-2.5-flash | 10 | 1 | 0.98% | 101 |
| gemini-2.5-flash | 10 | 2 | 5.2% | 19 |
| gemini-2.5-pro | 5 | 1 | 5.2% | 19 |
| gemini-2.5-pro | 5 | 2 | 15.7% | 6 |

A single worker running `gemini_udf` with gemini-2.5-pro on free tier will lose **1 in 19 rows**. With 2 workers, it's 1 in 6. These are not edge cases — they happen on every free-tier backfill of any meaningful size.

#### Tighter target: >0.1% row failure (current config)

For workloads where losing even 1 in 1,000 rows is unacceptable:

| Model | Tier | RPM | N for >1% | N for >0.1% | N for >0.01% |
|-------|------|----:|----------:|------------:|-------------:|
| gemini-2.5-flash | Free | 10 | 2 | 1 | 1 |
| gemini-2.5-flash | Tier 1 | 1,000 | 101 | 60 | 47 |
| gemini-2.5-flash | Tier 2 | 2,000 | 202 | 119 | 93 |
| gemini-2.5-flash | Tier 3 | 4,000 | 403 | 238 | 186 |
| gemini-2.5-pro | Tier 1 | 150 | 16 | 9 | 7 |
| gemini-2.5-pro | Tier 2 | 1,000 | 101 | 60 | 47 |
| gemini-2.5-flash-lite | Tier 1 | 300 | 31 | 18 | 14 |
| gemini-embedding-001 | Free | 100 | 11 | 6 | 5 |

The safe zone shrinks fast: for 0.1% failure with Tier 1 flash, you drop from 101 to **60 workers**. For 0.01% (1 in 10,000), it's only **47 workers**. The ratio is roughly **N ≈ R/17** for 0.1% and **N ≈ R/21** for 0.01%.

#### Can more retries fix this?

We tested extending the retry policy with more attempts and higher backoff caps:

| Configuration | Max wall time | flash Free N=1 | flash T1 >1% | flash T1 >0.1% |
|--------------|-------------:|---------------:|--------------:|---------------:|
| Current: 7 attempts, max 60s | ~64s | 0.98% | 101 workers | 60 workers |
| 10 attempts, max 60s | ~123s | 0.10% | 213 workers | 101 workers |
| 10 attempts, max 300s | ~5 min | 0.06% | 311 workers | 116 workers |
| 15 attempts, max 300s | ~19 min | 0.001% | 1,119 workers | 370 workers |

More retries help significantly:
- **10 attempts with max 60s backoff** roughly doubles the safe concurrency. A free-tier single worker drops from 0.98% to 0.10% failure — marginal but better.
- **15 attempts with max 300s backoff** allows **1,119 workers** before hitting 1% on Tier 1 flash (11x improvement). But a failing row now takes up to **19 minutes** of wall time, tying up the worker.
- Increasing max backoff from 60s to 300s without more attempts (**7 attempts, max 300s**) barely helps — the bottleneck is the number of chances, not the wait between them.

**Free tier with extended retries:**

| Configuration | flash (10 RPM) N=1 | pro (5 RPM) N=1 | pro (5 RPM) N=2 |
|--------------|-------------------:|-----------------:|----------------:|
| 7 attempts, max 60s | 0.98% | 5.19% | 15.72% |
| 10 attempts, max 60s | 0.10% | 0.86% | 3.84% |
| 15 attempts, max 300s | 0.001% | 0.019% | 0.12% |

15 attempts with 5-minute max backoff finally makes free tier usable with a single worker (~0.001% failure for flash, ~0.02% for pro). But the cost is that a persistently rate-limited row blocks a worker for up to 19 minutes.

**Takeaway**: More retries buy headroom but don't eliminate the fundamental problem — they just push the cliff further out while increasing worst-case latency per row. Client-side rate limiting avoids hitting the cliff entirely.

### Why This Is Hard in a Distributed Setting

Geneva runs UDFs on distributed Ray applier actors. There is no shared state between actors — each runs in its own process, potentially on different machines. A global rate limiter would require external coordination (e.g., Redis, a central token-dispensing actor), adding operational complexity and a single point of failure.

## Proposed Solution: Per-Worker Token Bucket

Add an optional `max_rpm` parameter to `gemini_udf` and `gemini_embedding_udf`. The rate is divided evenly across workers at UDF construction time, and each worker enforces its share locally using a simple token bucket.

### User-Facing API

```python
# Gemini generative UDF
udf = gemini_udf(
    column="image",
    prompt="Describe this scene",
    mime_type="image/jpeg",
    max_rpm=1000,       # total RPM budget across all workers
    concurrency=8,      # hint: how many workers will run this UDF
)
table.add_columns({"caption": udf})
table.backfill("caption", concurrency=8)

# Gemini embedding UDF
udf = gemini_embedding_udf(
    column="text",
    max_rpm=1000,
    concurrency=8,
)
```

### Per-Worker Rate Calculation

```
per_worker_rpm = max_rpm / concurrency
```

Example: `max_rpm=1000, concurrency=8` → each worker throttles to 125 RPM (~2.08 requests/sec).

For `gemini_embedding_udf`, each `embed()` call covers one batch (one API call), so the per-worker rate applies to batch-level API calls, not individual rows.

### Token Bucket Implementation

Each worker maintains a local token bucket that refills at `per_worker_rpm / 60` tokens per second. Before making an API call, the worker acquires a token — blocking (sleeping) if none are available.

```python
import time
import threading

class TokenBucket:
    """Simple thread-safe token bucket rate limiter."""

    def __init__(self, rate_per_second: float, burst: int = 1) -> None:
        self._rate = rate_per_second
        self._burst = burst
        self._tokens = float(burst)
        self._last_refill = time.monotonic()
        self._lock = threading.Lock()

    def acquire(self) -> None:
        """Block until a token is available."""
        while True:
            with self._lock:
                now = time.monotonic()
                elapsed = now - self._last_refill
                self._tokens = min(self._burst, self._tokens + elapsed * self._rate)
                self._last_refill = now
                if self._tokens >= 1.0:
                    self._tokens -= 1.0
                    return
            # Sleep briefly before retrying
            time.sleep(0.05)
```

### Where It Lives

The `TokenBucket` class would live in a new module `src/geneva/udfs/_ratelimit.py` (internal, not public API). It is attached to the model object (`_GeminiModel`, `_GeminiEmbeddingModel`) and called before each API request.

```python
@attrs.define
class _GeminiModel:
    # ... existing fields ...
    _rate_limiter: TokenBucket | None = attrs.field(default=None)

    def generate(self, value: Any) -> str | None:
        if value is None:
            return None
        if self._rate_limiter is not None:
            self._rate_limiter.acquire()
        response = self.client.models.generate_content(...)
        return response.text
```

### Serialization

`TokenBucket` is transient state — it must not be serialized. On deserialization (worker startup), it is reconstructed from the stored rate parameters:

```python
@attrs.define
class _GeminiModel:
    _max_rpm_per_worker: float | None = attrs.field(default=None)

    # Not serialized — recreated on the worker
    _rate_limiter: TokenBucket | None = attrs.field(init=False, default=None)

    def __getstate__(self) -> dict[str, Any]:
        return attrs.asdict(self, filter=lambda attr, _: attr.init)

    def __setstate__(self, state: dict[str, Any]) -> None:
        self.__init__(**state)
        # Reconstruct rate limiter on worker
        if self._max_rpm_per_worker is not None:
            self._rate_limiter = TokenBucket(
                rate_per_second=self._max_rpm_per_worker / 60.0,
            )
```

## Alternatives Considered

### 1. Central Rate-Limiting via JobTracker

Geneva already has a centralized per-job actor: **JobTracker**. It lives for the duration of every backfill/refresh job, is configured with `max_restarts=-1` (auto-recovery), and workers already communicate with it to push metrics. Instead of creating a new actor, we could add a token-bucket to the existing JobTracker and have applier actors request tokens before each API call.

#### How it would work

```
ApplierActor A ──┐
ApplierActor B ──┤──▶ JobTracker.acquire_token.remote() ──▶ sleep_seconds
ApplierActor C ──┘
```

The JobTracker would maintain a single `TokenBucket` with the full `max_rpm` budget. Each applier calls `acquire_token.remote()` before making an API call. The JobTracker returns the number of seconds the worker should sleep before proceeding (0 if a token is immediately available). The sleep happens on the worker side so the JobTracker actor is not blocked.

```python
# In JobTracker (Ray actor)
class JobTracker:
    def __init__(self, ...):
        ...
        self._rate_limiter: TokenBucket | None = None

    def set_rate_limit(self, max_rpm: float) -> None:
        """Called once by the pipeline driver after job setup."""
        self._rate_limiter = TokenBucket(rate_per_second=max_rpm / 60.0)

    def acquire_token(self) -> float:
        """Return seconds to sleep. 0 means proceed immediately."""
        if self._rate_limiter is None:
            return 0.0
        return self._rate_limiter.wait_time()

# In ApplierActor / UDF
sleep_secs = ray.get(job_tracker.acquire_token.remote())
if sleep_secs > 0:
    time.sleep(sleep_secs)
response = client.models.generate_content(...)
```

#### Advantages over a new dedicated actor

- **No new actor lifecycle to manage** — JobTracker is already created, monitored, and cleaned up by the pipeline.
- **Auto-recovery** — `max_restarts=-1` means if it crashes and restarts, the token bucket resets (briefly allowing a burst, but self-correcting).
- **Knows actual concurrency** — the pipeline thread that creates the JobTracker also knows the real `concurrency` value, so `max_rpm` doesn't need to be divided at UDF construction time. This eliminates the `concurrency` parameter mismatch problem entirely.
- **Globally accurate** — one bucket enforces one rate across all workers, no over/under-provisioning from uneven worker load.

#### Challenges

- **Network round-trip per API call**: Each `acquire_token.remote()` is a Ray object ref that requires a round-trip to the JobTracker actor. For Gemini API calls that take 1–3 seconds, the ~0.1–1ms Ray RPC overhead is negligible. But for fast UDFs or high-throughput embedding batches, this could become a bottleneck.

- **Serialization boundary**: The UDF currently has no reference to the JobTracker actor handle. The handle would need to be injected into the UDF or the applier actor at job setup time. This means the rate limiter can't live inside `_GeminiModel` — it would need to be in the applier/pipeline layer, which changes where the throttling logic lives.

  Possible injection points:
  - **ApplierActor.__init__**: Pass the JobTracker handle when constructing applier actors. The applier calls `acquire_token` before invoking the UDF. This keeps the UDF unaware of rate limiting but requires applier-level awareness of which UDFs need throttling.
  - **UDF.set_job_context()**: A lifecycle hook called after deserialization on the worker, passing the JobTracker handle. This keeps throttling in the UDF layer but adds a new UDF protocol method.

- **Head-of-line blocking**: If many workers call `acquire_token` simultaneously and the rate is low, Ray will queue the calls to the single-threaded actor. Workers will block on `ray.get()` waiting for their turn. This is correct behavior (they should wait) but means workers are idle while holding Ray resources. With `asyncio`-based JobTracker methods this could be mitigated, but adds complexity.

- **Not usable outside Ray**: The per-worker `TokenBucket` approach works for local/non-Ray execution too (e.g., `udf(batch)` in a notebook). The JobTracker approach only works during Ray-managed backfills.

#### Hybrid: JobTracker for coordination, local bucket for execution

A middle ground: the JobTracker tracks the *global observed rate* (it already receives timing metrics) and periodically broadcasts an adjusted per-worker rate to appliers. Each worker still has a local token bucket, but its rate is dynamically tuned by the JobTracker based on actual aggregate throughput.

```
Pipeline setup:
  JobTracker.set_rate_limit(max_rpm=1000, num_workers=8)
  → initial per-worker rate = 125 RPM

During execution:
  JobTracker receives timing metrics from all workers
  → observes actual aggregate = 900 RPM (some workers slower)
  → increases faster workers' rate to compensate

  Workers periodically call:
    new_rate = ray.get(job_tracker.get_my_rate.remote(worker_id))
    local_bucket.update_rate(new_rate)
```

This combines global visibility with local execution — no per-request RPC, but adaptive rebalancing. The downside is significantly more complexity: the JobTracker needs to track per-worker throughput, implement a rebalancing algorithm, and workers need a periodic sync loop.

#### Verdict

Using the JobTracker is architecturally appealing — it eliminates the `concurrency` mismatch problem and gives globally accurate enforcement without a new actor. The main cost is coupling the rate limiter to the Ray pipeline layer (injection of the actor handle into UDFs/appliers) and the per-request RPC overhead.

**Recommendation**: Start with the simpler per-worker `TokenBucket` (proposed solution) for v1. If users hit problems with uneven rate distribution — e.g., some workers idle while others are throttled — the JobTracker-based approach is a natural v2 upgrade. The `max_rpm` user-facing API is the same either way, so the migration is internal.

### 2. No Client-Side Limiting (Status Quo)

Rely entirely on server-side 429s and retry backoff.

**Pros**: Zero additional code. Self-regulating.
**Cons**:
- Wastes retry budget on predictable 429s.
- Each rate-limited request costs 1–60 seconds of backoff sleep.
- With `max_attempts=7`, a persistently rate-limited row takes ~2 minutes before failing.
- Noisy — retries consume log budget and make real errors harder to spot.

**Verdict**: Works for small jobs but wasteful at scale. The per-worker approach avoids most 429s while keeping the retry policy as a safety net.

### 3. Adaptive Rate Limiting

Workers start fast and slow down in response to 429s, using an AIMD (Additive Increase Multiplicative Decrease) strategy.

**Pros**: No user configuration needed. Adapts to unknown rate limits.
**Cons**:
- Complex to implement correctly (especially across distributed workers).
- Oscillation: workers constantly probe the limit, causing periodic 429 bursts.
- Still wastes some retry budget during the "probe" phase.
- Hard to reason about — users can't predict throughput.

**Verdict**: Interesting future optimization but over-engineered for the initial implementation. The explicit `max_rpm` parameter is simpler and gives users control.

## Design Decisions

### Why `concurrency` is a parameter on the UDF factory, not discovered at runtime

The UDF object is constructed (and serialized) before the backfill job starts. At construction time, we don't know how many workers will actually run — that depends on `table.backfill(concurrency=N)`, which happens later. So we require the user to provide `concurrency` as a hint at UDF construction time.

If `concurrency` at UDF construction doesn't match `concurrency` at backfill time, the rate won't be perfectly divided. This is acceptable — the per-worker bucket is a best-effort throttle, and the retry policy handles the overflow.

An alternative is to resolve concurrency at backfill time and inject it into the UDF before serialization. This would require plumbing `concurrency` through the pipeline into the UDF. The JobTracker-based approach (Alternative 1) solves this naturally since the pipeline knows the real concurrency when it sets up the JobTracker. Worth exploring as a follow-up but not blocking for v1.

### Why `burst=1` by default

A burst size of 1 means no bursting — requests are spread evenly. This is the safest default for API rate limiting where sustained throughput matters more than latency. Users who need bursting (e.g., batch APIs that benefit from grouped requests) could be given a `burst` parameter in the future.

### Thread safety

The token bucket uses `threading.Lock` because `intra_applier_concurrency > 1` means multiple threads within a single applier actor may call the UDF concurrently. Each applier actor gets its own `_GeminiModel` instance (via deserialization), so the lock is per-actor, not global.

## Gemini Rate Limits Reference

### Tier Structure and Upgrade Requirements

Google gates rate limits behind billing tiers. Tiers are per-project (not per API key) and based on cumulative Google Cloud spend (promotional credits do not count).

| Tier | Requirement | Upgrade time |
|------|------------|--------------|
| **Free** | None | Instant |
| **Tier 1** | Enable Cloud Billing with valid payment method | Instant |
| **Tier 2** | $250 cumulative spend + 30 days since first payment | 24–48 hours |
| **Tier 3** | $1,000 cumulative spend + 30 days, or Google Cloud sales engagement | 24–48 hours, or 2–4 weeks via sales |

Source: [Billing](https://ai.google.dev/gemini-api/docs/billing)

### Generative Models (gemini_udf)

Rate limits are measured in RPM (requests per minute), TPM (tokens per minute), and RPD (requests per day). These limits are approximate — Google adjusts them periodically and actual limits may differ from published values.

| Model | Tier | RPM | TPM | RPD |
|-------|------|-----|-----|-----|
| gemini-2.5-flash | Free | 10 | 250,000 | 250 |
| | Tier 1 | 200–1,000 | 1–4M | 1,500–10,000 |
| | Tier 2 | 1,000–2,000 | 2–4M | Unlimited |
| | Tier 3 | 2,000–4,000+ | Custom | Unlimited |
| gemini-2.5-pro | Free | 5 | 250,000 | 100 |
| | Tier 1 | 150 | 1–2M | 1,500–10,000 |
| | Tier 2 | 500–1,000 | 2–4M | Unlimited |
| | Tier 3 | 1,000–4,000+ | Custom | Unlimited |
| gemini-2.5-flash-lite | Free | 15 | 250,000 | 1,000 |
| | Tier 1 | 300 | 1–2M | 1,500 |
| | Tier 2 | 1,500 | 2M | 10,000 |
| | Tier 3 | 4,000+ | Custom | Unlimited |

Note: ranges reflect variance across sources and Google's frequent updates. Check [AI Studio usage dashboard](https://aistudio.google.com/usage) for your project's actual limits.

Since `gemini_udf` is a scalar UDF (one API call per row), each row = 1 request. RPM is the binding constraint for most workloads. TPM may matter for large input documents or images.

Source: [Google AI rate limits](https://ai.google.dev/gemini-api/docs/rate-limits), [Google AI pricing](https://ai.google.dev/gemini-api/docs/pricing)

### Batch API Limits

For large-scale offline processing, Gemini offers a Batch API with token-based (not RPM-based) quotas:

| Tier | Max enqueued tokens |
|------|-------------------|
| Tier 1 | ~10M |
| Tier 2 | 1B |
| Tier 3 | 5B |

The Batch API is a different integration path from Geneva's current UDF model (which uses the real-time API). It could be relevant for future offline-only backfill modes where latency doesn't matter.

### Embedding Models (gemini_embedding_udf)

Gemini embedding models use a **different quota model** than generative models. The primary limit on paid tiers is **TPM (tokens per minute)** rather than RPM, reflecting that embedding requests can batch up to 250 texts per call.

| Model | Tier | RPM | TPM | RPD |
|-------|------|-----|-----|-----|
| gemini-embedding-001 | Free | 100 | — | 1,000 |
| | Paid (Tier 1+) | — | 5,000,000 | — |

Key differences from generative models:

- **Batch-oriented**: Each `embed_content` call can embed up to 250 texts (20,000 tokens per request). Our `gemini_embedding_udf` sends all valid texts in a `RecordBatch` in a single API call, so one backfill batch = one API request.
- **TPM is the primary throttle on paid tiers**: At 5M TPM with ~100 tokens per text, that's ~50,000 texts/min. The RPM limit is less likely to be the bottleneck since each request can carry many texts.
- **Free tier is RPM-gated**: 100 RPM with 1,000 RPD is more restrictive than paid. Since each batch call counts as 1 request, a backfill with `checkpoint_size=100` (100 rows per batch) could embed 10,000 rows/min at 100 RPM — but will hit the 1,000 RPD limit after just 10 batches.

Source: [Vertex AI quotas](https://docs.cloud.google.com/vertex-ai/generative-ai/docs/quotas), [Google AI rate limits](https://ai.google.dev/gemini-api/docs/rate-limits)

### Implications for Rate Limiter Design

The dual quota model (RPM for generative, TPM for embeddings) means a simple `max_rpm` parameter works well for `gemini_udf` but is an incomplete picture for `gemini_embedding_udf`:

| UDF | Bottleneck | Rate limiter approach |
|-----|-----------|----------------------|
| `gemini_udf` | RPM (1 request per row) | `max_rpm` directly limits requests |
| `gemini_embedding_udf` | TPM on paid tiers, RPM on free | `max_rpm` limits batch API calls; TPM limiting is a future concern |

For v1, `max_rpm` is sufficient for both UDFs — it controls request rate, which is the binding constraint in practice. TPM limiting could be added later if users hit token-based throttling (large documents, high batch sizes).

### RPD (Requests Per Day) Consideration

RPD limits are particularly restrictive on free/Tier 1:
- gemini-2.5-pro free tier: 100 RPD = 100 rows/day
- gemini-2.5-flash Tier 1: 1,500–10,000 RPD = 1,500–10,000 rows/day

For datasets with more rows than the RPD limit, backfills will fail partway through and need to be resumed the next day (Geneva's checkpointing handles this). The rate limiter doesn't address RPD — it only smooths out RPM. Users hitting RPD limits need to upgrade their tier.

At **Tier 2+**, RPD is unlimited for generative models, making it the practical minimum tier for production Geneva workloads.

### Example Configurations

```python
# --- gemini_udf (generative, 1 API call per row) ---

# Free tier, flash — very constrained, useful only for dev/testing
gemini_udf(..., max_rpm=10, concurrency=2)
# → 5 RPM per worker, but will hit 250 RPD after 25 min

# Tier 1, flash — moderate throughput
gemini_udf(..., max_rpm=900, concurrency=8)
# → 112.5 RPM per worker (~1.9 req/sec)

# Tier 2, flash — production workloads
gemini_udf(..., max_rpm=1800, concurrency=16)
# → 112.5 RPM per worker (~1.9 req/sec), unlimited RPD

# Tier 3, flash — high-scale production
gemini_udf(..., max_rpm=3500, concurrency=32)
# → ~109 RPM per worker (~1.8 req/sec), unlimited RPD

# Tier 1, pro — lower limits, higher quality
gemini_udf(..., max_rpm=140, concurrency=4)
# → 35 RPM per worker (~0.6 req/sec)

# Tier 3, pro — enterprise scale
gemini_udf(..., max_rpm=3500, concurrency=16)
# → ~219 RPM per worker (~3.6 req/sec)

# --- gemini_embedding_udf (batch API, 1 API call per RecordBatch) ---

# Free tier, embeddings
gemini_embedding_udf(..., max_rpm=90, concurrency=4)
# → 22.5 batch calls/min per worker
# With checkpoint_size=100: ~2,250 rows/min total
# But will hit 1,000 RPD after ~11 min (44 batches)

# Paid tier, embeddings — TPM is the real limit
gemini_embedding_udf(..., max_rpm=500, concurrency=8)
# → 62.5 batch calls/min per worker
# With checkpoint_size=100 (~100 tokens/text): ~50,000 rows/min total
# TPM: 100 texts × 100 tokens × 500 RPM = 5M TPM (at the limit)
```

Note: `max_rpm` should be set slightly below the tier limit to leave headroom for other API consumers sharing the same project/key. Check your project's actual limits in the [AI Studio usage dashboard](https://aistudio.google.com/usage).

## Scope and Milestones

### v1 (this design)

- `TokenBucket` in `src/geneva/udfs/_ratelimit.py`
- `max_rpm` + `concurrency` parameters on `gemini_udf()` and `gemini_embedding_udf()`
- Per-worker static rate division
- Unit tests with mocked time

### Future (out of scope)

- **JobTracker-based global enforcement**: Move the token bucket into the JobTracker actor so all workers share one global rate (see Alternative 1). Eliminates the `concurrency` parameter mismatch problem and handles uneven worker load. Requires injecting the JobTracker handle into UDFs or appliers.
- **Automatic concurrency inference**: Resolve actual concurrency at backfill time and inject it into the UDF before serialization, removing the need for users to pass `concurrency` to the UDF factory.
- **Adaptive rate limiting (AIMD)**: Workers start fast, slow down on 429s. Could be combined with JobTracker-based global feedback — the JobTracker observes aggregate throughput and broadcasts adjusted per-worker rates.
- **Generic `@udf(max_rpm=...)` support**: Extend rate limiting to user-defined UDFs, not just built-in Gemini UDFs.
- **TPM (tokens per minute) limiting**: For LLM cost control, track token consumption in addition to request count.
