Admission and pinning
For each request sent to Iris' proxy we decide which "provider" to forward it to. Either our own, referred to here as Iris, or a third party such as TensorX. There may be several of ours, across different deployments of the same model, with and without Dynamo.
This decision can be made depending on current load, health status etc. of Iris. We prefer using Iris and avoid offloading to a third party if we don't have to.
For every request we can heuristically determine what "session" it belongs to where a session is usually consecutive requests for an ongoing chat session. For requests belonging to the same session we strongly prefer sending them to the same provider so they can benefit from prefix caching. Thus we also need a mechanism that "pins" sessions to providers (both Iris and third party) with renewable leases.
All of this has to work atomically (admission decision + pinning) and distributed across several instances of the proxy.
Below is a pseudo-code implementation of the algorithm we have devised:
# Providers with health checks and capacity limits are preferred in this order
# (Iris 1, Iris 2, ...)
MANAGED_PROVIDERS_IN_PRIORITY_ORDER = [PROVIDER_A, PROVIDER_B]
# These providers have neither health checks nor capacity limits (TensorX, ...)
UNLIMITED_PROVIDERS_IN_PRIORITY_ORDER = [PROVIDER_C, PROVIDER_D]
# Example values, up to discussion.
ROUTING_STATE.limits[PROVIDER_X, MODEL_ID_X] = ModelLimits(
max_pinned_sessions=1_000,
max_global_in_flight_request_slots=100,
global_in_flight_request_slots_reserved_for_pinned_sessions=20,
max_per_session_in_flight_request_slots=2,
)
# How soon a pin expires after no more requests are active for a given session.
SESSION.pin_inactivity_timeout_seconds = 300
# How soon a request is considered to be done after it has stopped being
# active. This is just a fallback in case we do not explicitly release a
# request lease (because the program crashes or similar).
REQUEST.lease_timeout_seconds = 30
MAX_ROUTING_ATTEMPTS = 10
RETRY_ROUTING = object()
def execute_request(request) -> status_code:
session_id = determine_session_id(request)
if session_id is None:
return execute_request_without_pinning(request)
for _ in range(MAX_ROUTING_ATTEMPTS):
pin = ROUTING_STATE.get_pin(request.model_id, session_id)
if pin is None:
result = execute_not_yet_pinned_request(request, session_id)
else:
result = execute_pinned_request(request, session_id, pin)
if result is not RETRY_ROUTING:
return result
return reject_request_as_unavailable(request)
def execute_request_without_pinning(request) -> status_code:
return execute_temporary_fallback(request, excluded_provider=None)
def execute_not_yet_pinned_request(request, session_id) -> status_code:
for provider in all_providers_in_priority_order():
if not provider.supports(request.model_id, request.feature):
continue
if provider_is_degraded(provider, request.model_id):
continue
admission = ROUTING_STATE.try_admit_unpinned_session_request(
provider,
request.model_id,
session_id,
)
if admission.result == ALREADY_PINNED:
return RETRY_ROUTING
if admission.result == ADMITTED:
return execute_on_pinned_provider(request, admission)
# A managed provider may reject due to pin or request capacity. Try the
# next provider. Unlimited providers cannot reach this.
return reject_request_as_unsupported_or_unavailable(request)
def execute_pinned_request(request, session_id, pin) -> status_code:
provider = pin.provider
if (
not provider.supports(request.model_id, request.feature)
or provider_is_degraded(provider, request.model_id)
):
return try_repin_to_other_provider_and_retry_execute(
request,
session_id,
pin,
)
# Every pinned request receives a lease. Providers without configured
# limits admit unconditionally, but their leases still prevent pin expiry.
admission = ROUTING_STATE.try_admit_pinned_session_request(
request.model_id,
session_id,
)
if admission.result == SESSION_NOT_PINNED:
# The pin disappeared between lookup and admission. Restart routing so
# the request can establish a new pin.
return RETRY_ROUTING
# The pin may have changed between GetPin and admission.
if admission.pin.id != pin.id:
if admission.result == ADMITTED:
ROUTING_STATE.release_request_lease(
admission.lease,
refresh_pin=False,
)
return RETRY_ROUTING
if admission.result in (
PER_SESSION_IN_FLIGHT_REQUEST_LIMIT_REACHED,
GLOBAL_IN_FLIGHT_REQUEST_LIMIT_REACHED,
):
# A temporary capacity spike must not destroy the existing pin so we
# offload instead.
return execute_temporary_fallback(
request,
excluded_provider=provider,
)
return execute_on_pinned_provider(request, admission)
def try_repin_to_other_provider_and_retry_execute(
request,
session_id,
current_pin,
) -> status_code:
for provider in all_providers_in_priority_order():
if provider == current_pin.provider:
continue
if not provider.supports(request.model_id, request.feature):
continue
if provider_is_degraded(provider, request.model_id):
continue
result = ROUTING_STATE.try_compare_and_swap_pin_and_admit_request(
expected_pin=current_pin,
model_id=request.model_id,
session_id=session_id,
provider=provider,
)
if result.pin_match != MATCHED:
return RETRY_ROUTING
if result.admission.result == ADMITTED:
return execute_on_pinned_provider(request, result.admission)
if result.admission.result == ALREADY_PINNED:
return RETRY_ROUTING
# The target provider lacks pin or request capacity. The original pin
# is unchanged, so try the next provider with the same expected pin.
if result.admission.result in (
PIN_LIMIT_REACHED,
PER_SESSION_IN_FLIGHT_REQUEST_LIMIT_REACHED,
GLOBAL_IN_FLIGHT_REQUEST_LIMIT_REACHED,
UNRESERVED_IN_FLIGHT_REQUEST_LIMIT_REACHED,
):
continue
return reject_request_as_unsupported_or_unavailable(request)
def execute_temporary_fallback(
request,
excluded_provider,
) -> status_code:
for provider in all_providers_in_priority_order():
if provider == excluded_provider:
continue
if not provider.supports(request.model_id, request.feature):
continue
if provider_is_degraded(provider, request.model_id):
continue
if provider_has_configured_limits(provider, request.model_id):
admission = ROUTING_STATE.try_admit_request_without_pin(
provider,
request.model_id,
)
if admission.result != ADMITTED:
continue
return execute_with_request_lease(
request,
provider,
admission.lease,
refresh_pin_on_success=False,
)
# Unlimited providers do not need capacity leases for requests that are
# not associated with one of their pins.
return run(request, provider).status_code
return reject_request_as_unsupported_or_unavailable(request)
def execute_on_pinned_provider(request, admission) -> status_code:
return execute_with_request_lease(
request,
admission.pin.provider,
admission.lease,
refresh_pin_on_success=True,
)
def execute_with_request_lease(
request,
provider,
lease,
refresh_pin_on_success,
) -> status_code:
result = None
renewal = periodically_renew_request_lease(lease)
try:
result = run(request, provider)
return result.status_code
finally:
renewal.stop()
ROUTING_STATE.release_request_lease(
lease,
refresh_pin=(
refresh_pin_on_success
and result is not None
and result.succeeded
),
)
def periodically_renew_request_lease(lease):
# If renewal reports that the lease was lost, cancel the associated request.
return start_background_renewal(
lease,
interval=REQUEST.lease_timeout_seconds / 3,
)
def all_providers_in_priority_order():
return (
MANAGED_PROVIDERS_IN_PRIORITY_ORDER
+ UNLIMITED_PROVIDERS_IN_PRIORITY_ORDER
)
def provider_is_degraded(provider, model_id) -> bool:
if provider not in MANAGED_PROVIDERS_IN_PRIORITY_ORDER:
return False
return provider.degraded(model_id)
def provider_has_configured_limits(provider, model_id) -> bool:
return ROUTING_STATE.limits.contains(provider, model_id)
def determine_session_id(request) -> string | None:
# Some requests, e.g. for embeddings, have no sensible prefix that we can
# use to build a session ID from so these never form sessions.
if not request.has_stable_prompt_prefix():
return None
return hash(
request.user_id,
request.model_id,
request.feature_class,
request.stable_prompt_prefix,
)
# Capacity and pin limits are maintained independently for each configured
# provider/model pair. Existing pinned sessions may consume capacity reserved
# for pinned requests. New, unpinned, and temporary requests may not.
#
# Provider/model pairs without configured limits have unlimited pin and request
# capacity. Their pinned requests still receive leases so pins cannot expire
# while requests are running and can be refreshed safely upon success.
#
# All pin, expiry, lease, and capacity transitions are atomic.