feat: add redis service and RPM traffic contract config and tracking

[noahpodgurski]

What's new

tl;dr

• Added Redis service for fixed-window (every 60s) traffic contract enforcement
  • For managing traffic contracts with methods for checking and incrementing feature RPM and TPM.
  • Added separate Lua scripts for checking contract state and incrementing counters via Redis. I separated checking and incrementing since at the start of the request we don't know the token count (that info is in the response). So the order is 1. Check RPM/TPM limit, 2. Make request, 3. Increment RPM/TPM
  • Added testing for redis service and simple mocked contract testing
• Over-contract requests are not rejected The check returns allowed=False with mode="borrowed" or mode="degraded", and logged, and are emitted on metric requests_total via traffic_contract_rpm_mode and traffic_contract_tpm_mode. No soft degradation is applied yet.
  • normal: feature and basket are within contract
  • borrowed: feature is over contract, but basket still has room
  • degraded: basket is over contract

Redis

There's a new redis_service which uses the same Redis instance as LiteLLM, with the mlpa:traffic_contract key prefix. There's a discussion to be had around the current eviction policy and whether we should change it. Although we've had no evictions since inception, but it's something to keep in mind. It uses small Lua scripts to check fixed-window feature/basket state and increment RPM/TPM counters.

Traffic contract design:

Traffic contracts are defined on a per-feature basis, so SW, S2S, etc... RPM and TPM limits are defined for each (ex: SMART_WINDOW_TRAFFIC_CONTRACT_RPM_LIMIT. RPM and TPM limits for all features (all requests/basket) combined is stored in TOTAL_TRAFFIC_CONTRACT_RPM_LIMIT. I also renamed env.user_feature_budget to service_type_config since it contains non budget info (feature info). Each budget/service type is combined into the env.traffic_contract_config which defines the RPM/TPM based off the service type's feature's definition. Not to be confused with the per user rpm/tpm limits defined in the service_type_config (old user_feature_budget)

The traffic_contract_enforcer runs as a handler in the request after the authorize_chat/search_request function has run. This is slightly less optimal (than in the middleware) but it keeps the code cleaner and reduces duplicated (service type, purpose, etc...) validation code and logic that fastapi already provides - plus it will make TPM tracking logic identical. A tradeoff with this is we do not count 401 requests towards the total.

After the upstream response, redis_service.update_contracts increments RPM and, when token usage is available, TPM.

Detailed example flow (expand)

The Redis key is a fixed-window bucket:

{TRAFFIC_CONTRACT_REDIS_KEY_PREFIX}:rpm:{bucket_start_epoch}

For example:

mlpa:traffic_contract:rpm:1789057680

The key is a Redis hash. Each feature gets its own field, and the shared basket
uses __basket__.

Example assumptions:

• TRAFFIC_CONTRACT_RPM_WINDOW_SECONDS=60
• TRAFFIC_CONTRACT_TPM_WINDOW_SECONDS=60
• TRAFFIC_CONTRACT_COUNTER_TTL_SECONDS=120
• smart-window feature RPM limit = 2
• smart-window feature TPM limit = 5000
• total basket RPM limit = 4
• total basket TPM limit = 10000
• ai and memories both map to smart-window
• s2s maps to s2s
• s2s-android maps to s2s-android
• all requests below land in the same minute bucket
• RPM and TPM are checked before the upstream request; RPM is incremented after the response, and TPM is incremented after the response when usage includes total_tokens

Diagram link

This means the whole shared basket is now over contract. The request still
continues for now, but this is the signal we can use later to degrade behavior:
reduce max tokens, skip retries, route differently, or apply a short wait/jitter.

When the next minute starts, Redis writes to a new key:

  mlpa:traffic_contract:rpm:1789057740

The previous bucket stays around briefly because the TTL is longer than the
window, which makes inspection/debugging easier without affecting enforcement.

If not enforcing limits, why not just create alerts based on this RPM/TPM data we already have in Grafana?

Even though we're not enforcing a 429 on requests that go over the traffic contract, the request metric now includes traffic_contract_rpm_mode and traffic_contract_tpm_mode, which provides insight into the percentage of traffic is normal, borrowed, or degraded and tune contract numbers from there.

https://mozilla-hub.atlassian.net/browse/AIPLAT-1060
https://mozilla-hub.atlassian.net/browse/AIPLAT-1064

[subpath]

Suggestion:

maybe wrap it in try except

finally:
        record_completion_latency(
            authorized_chat_request, result, time.perf_counter() - start_time
        )
        record_chat_availability(authorized_chat_request, availability_reason)
        try:
            await redis_service.update_contracts(
                service_type=authorized_chat_request.service_type, usage=usage
            )
        except Exception as exc:
            logger.error(f"Traffic contract update failed: {exc}")
            if not env.TRAFFIC_CONTRACT_FAIL_OPEN_ON_REDIS_ERROR:
                raise

[noahpodgurski]

forgot to add handling here, thank you 🙏

[noahpodgurski]

I added the handling inside update_contracts and also converted the function to be ran as a task

        asyncio.create_task(
            redis_service.update_contracts(
                service_type=authorized_chat_request.service_type,
                usage=usage,
            )
        )

otherwise it blocks the data from being returned, so it's more efficient 👍

[subpath]

Comment: so now redis service will become a hard dependency, without it MLPA will not start

[noahpodgurski]

Only if ENABLE_TRAFFIC_CONTRACT_ENFORCEMENT=true yes?

But LiteLLM won't work without it either

[subpath]

Suspicious:
I think we can have a race condition ...

so I understood it as:

1. requests comes and we run HGET to read current count, predicts count + 1
2. we check if to allow the request or not
3. we don't persists this predicts count + 1 value
4. After LLM responds, we run increment_feature_traffic_contract runs and it does HINCRBY - pushing the value to Redis

but during the duration of the LLM response more calls will come...

Maybe we should persist the value on the before LLM call? 🤔
But for TPM calls it will not work

[noahpodgurski]

That was my first approach, to check and increment in the same script. But like you said for TPM it doesn't work, that's why I split them both out. I figured it's better to have them unified in that way than for just the RPM check and increment be persisted pre-request, even if we have the capability for it

[subpath]

Suggestion: maybe register enforce_traffic_contract into MIDDLEWARE_EXECUTION_ORDER

[noahpodgurski]

I added a bit about this in the description 🙏

The traffic_contract_enforcer runs as a handler in the request after the authorize_chat/search_request function has run. This is slightly less optimal (than in the middleware) but it keeps the code cleaner and reduces duplicated (service type, purpose, etc...) validation code and logic that fastapi already provides - plus it will make TPM tracking logic identical. A tradeoff with this is we do not count 401 requests towards the total.