diff --git a/charts/model-engine/templates/_istio-attribute-match-conditions.tpl b/charts/model-engine/templates/_istio-attribute-match-conditions.tpl index 6e9feeb16..c34a90da0 100644 --- a/charts/model-engine/templates/_istio-attribute-match-conditions.tpl +++ b/charts/model-engine/templates/_istio-attribute-match-conditions.tpl @@ -1,117 +1,135 @@ {{- /* Generated from the OpenAPI schema with model-engine-internal/scripts/generate_istio_metric_tags.py */}} {{- define "modelEngine.istioAttributeMatchConditions" -}} -- condition: request.method == 'GET' && request.url_path == '/healthcheck' - value: get_/healthcheck -- condition: request.method == 'GET' && request.url_path == '/healthz' - value: get_/healthz -- condition: request.method == 'GET' && request.url_path == '/readyz' - value: get_/readyz -- condition: request.method == 'POST' && request.url_path == '/v1/async-tasks' - value: post_/v1/async-tasks -- condition: request.method == 'GET' && request.url_path.matches('^/v1/async-tasks/[[:alnum:]-_]*$') - value: get_/v1/async-tasks/_task_id - condition: request.method == 'POST' && request.url_path == '/v1/batch-jobs' value: post_/v1/batch-jobs - condition: request.method == 'GET' && request.url_path.matches('^/v1/batch-jobs/[[:alnum:]-_]*$') value: get_/v1/batch-jobs/_batch_job_id - condition: request.method == 'PUT' && request.url_path.matches('^/v1/batch-jobs/[[:alnum:]-_]*$') value: put_/v1/batch-jobs/_batch_job_id -- condition: request.method == 'GET' && request.url_path == '/v1/docker-image-batch-job-bundles' - value: get_/v1/docker-image-batch-job-bundles -- condition: request.method == 'POST' && request.url_path == '/v1/docker-image-batch-job-bundles' - value: post_/v1/docker-image-batch-job-bundles -- condition: request.method == 'GET' && request.url_path == '/v1/docker-image-batch-job-bundles/latest' - value: get_/v1/docker-image-batch-job-bundles/latest -- condition: request.method == 'GET' && request.url_path.matches('^/v1/docker-image-batch-job-bundles/[[:alnum:]-_]*$') - value: get_/v1/docker-image-batch-job-bundles/_docker_image_batch_job_bundle_id -- condition: request.method == 'GET' && request.url_path == '/v1/docker-image-batch-jobs' - value: get_/v1/docker-image-batch-jobs - condition: request.method == 'POST' && request.url_path == '/v1/docker-image-batch-jobs' value: post_/v1/docker-image-batch-jobs +- condition: request.method == 'GET' && request.url_path == '/v1/docker-image-batch-jobs' + value: get_/v1/docker-image-batch-jobs - condition: request.method == 'GET' && request.url_path.matches('^/v1/docker-image-batch-jobs/[[:alnum:]-_]*$') value: get_/v1/docker-image-batch-jobs/_batch_job_id - condition: request.method == 'PUT' && request.url_path.matches('^/v1/docker-image-batch-jobs/[[:alnum:]-_]*$') value: put_/v1/docker-image-batch-jobs/_batch_job_id -- condition: request.method == 'GET' && request.url_path == '/v1/files' - value: get_/v1/files -- condition: request.method == 'POST' && request.url_path == '/v1/files' - value: post_/v1/files -- condition: request.method == 'DELETE' && request.url_path.matches('^/v1/files/[[:alnum:]-_]*$') - value: delete_/v1/files/_file_id -- condition: request.method == 'GET' && request.url_path.matches('^/v1/files/[[:alnum:]-_]*$') - value: get_/v1/files/_file_id -- condition: request.method == 'GET' && request.url_path.matches('^/v1/files/[[:alnum:]-_]*/content$') - value: get_/v1/files/_file_id/content -- condition: request.method == 'POST' && request.url_path == '/v1/llm/completions-stream' - value: post_/v1/llm/completions-stream -- condition: request.method == 'POST' && request.url_path == '/v1/llm/completions-sync' - value: post_/v1/llm/completions-sync -- condition: request.method == 'GET' && request.url_path == '/v1/llm/fine-tunes' - value: get_/v1/llm/fine-tunes -- condition: request.method == 'POST' && request.url_path == '/v1/llm/fine-tunes' - value: post_/v1/llm/fine-tunes -- condition: request.method == 'GET' && request.url_path.matches('^/v1/llm/fine-tunes/[[:alnum:]-_]*$') - value: get_/v1/llm/fine-tunes/_fine_tune_id -- condition: request.method == 'PUT' && request.url_path.matches('^/v1/llm/fine-tunes/[[:alnum:]-_]*/cancel$') - value: put_/v1/llm/fine-tunes/_fine_tune_id/cancel -- condition: request.method == 'GET' && request.url_path.matches('^/v1/llm/fine-tunes/[[:alnum:]-_]*/events$') - value: get_/v1/llm/fine-tunes/_fine_tune_id/events -- condition: request.method == 'GET' && request.url_path == '/v1/llm/model-endpoints' - value: get_/v1/llm/model-endpoints -- condition: request.method == 'POST' && request.url_path == '/v1/llm/model-endpoints' - value: post_/v1/llm/model-endpoints -- condition: request.method == 'POST' && request.url_path == '/v1/llm/model-endpoints/download' - value: post_/v1/llm/model-endpoints/download -- condition: request.method == 'DELETE' && request.url_path.matches('^/v1/llm/model-endpoints/[[:alnum:]-_]*$') - value: delete_/v1/llm/model-endpoints/_model_endpoint_name -- condition: request.method == 'GET' && request.url_path.matches('^/v1/llm/model-endpoints/[[:alnum:]-_]*$') - value: get_/v1/llm/model-endpoints/_model_endpoint_name -- condition: request.method == 'GET' && request.url_path == '/v1/model-bundles' - value: get_/v1/model-bundles +- condition: request.method == 'POST' && request.url_path == '/v1/async-tasks' + value: post_/v1/async-tasks +- condition: request.method == 'GET' && request.url_path.matches('^/v1/async-tasks/[[:alnum:]-_]*$') + value: get_/v1/async-tasks/_task_id +- condition: request.method == 'POST' && request.url_path == '/v1/sync-tasks' + value: post_/v1/sync-tasks +- condition: request.method == 'POST' && request.url_path == '/v1/streaming-tasks' + value: post_/v1/streaming-tasks - condition: request.method == 'POST' && request.url_path == '/v1/model-bundles' value: post_/v1/model-bundles +- condition: request.method == 'GET' && request.url_path == '/v1/model-bundles' + value: get_/v1/model-bundles - condition: request.method == 'POST' && request.url_path == '/v1/model-bundles/clone-with-changes' value: post_/v1/model-bundles/clone-with-changes - condition: request.method == 'GET' && request.url_path == '/v1/model-bundles/latest' value: get_/v1/model-bundles/latest - condition: request.method == 'GET' && request.url_path.matches('^/v1/model-bundles/[[:alnum:]-_]*$') value: get_/v1/model-bundles/_model_bundle_id -- condition: request.method == 'GET' && request.url_path == '/v1/model-endpoints' - value: get_/v1/model-endpoints +- condition: request.method == 'POST' && request.url_path == '/v2/model-bundles' + value: post_/v2/model-bundles +- condition: request.method == 'GET' && request.url_path == '/v2/model-bundles' + value: get_/v2/model-bundles +- condition: request.method == 'POST' && request.url_path == '/v2/model-bundles/clone-with-changes' + value: post_/v2/model-bundles/clone-with-changes +- condition: request.method == 'GET' && request.url_path == '/v2/model-bundles/latest' + value: get_/v2/model-bundles/latest +- condition: request.method == 'GET' && request.url_path.matches('^/v2/model-bundles/[[:alnum:]-_]*$') + value: get_/v2/model-bundles/_model_bundle_id - condition: request.method == 'POST' && request.url_path == '/v1/model-endpoints' value: post_/v1/model-endpoints -- condition: request.method == 'GET' && request.url_path == '/v1/model-endpoints-api' - value: get_/v1/model-endpoints-api -- condition: request.method == 'GET' && request.url_path == '/v1/model-endpoints-schema.json' - value: get_/v1/model-endpoints-schema.json -- condition: request.method == 'DELETE' && request.url_path.matches('^/v1/model-endpoints/[[:alnum:]-_]*$') - value: delete_/v1/model-endpoints/_model_endpoint_id +- condition: request.method == 'GET' && request.url_path == '/v1/model-endpoints' + value: get_/v1/model-endpoints - condition: request.method == 'GET' && request.url_path.matches('^/v1/model-endpoints/[[:alnum:]-_]*$') value: get_/v1/model-endpoints/_model_endpoint_id - condition: request.method == 'PUT' && request.url_path.matches('^/v1/model-endpoints/[[:alnum:]-_]*$') value: put_/v1/model-endpoints/_model_endpoint_id -- condition: request.method == 'POST' && request.url_path == '/v1/streaming-tasks' - value: post_/v1/streaming-tasks -- condition: request.method == 'POST' && request.url_path == '/v1/sync-tasks' - value: post_/v1/sync-tasks +- condition: request.method == 'DELETE' && request.url_path.matches('^/v1/model-endpoints/[[:alnum:]-_]*$') + value: delete_/v1/model-endpoints/_model_endpoint_id +- condition: request.method == 'POST' && request.url_path.matches('^/v1/model-endpoints/[[:alnum:]-_]*/restart$') + value: post_/v1/model-endpoints/_model_endpoint_id/restart +- condition: request.method == 'GET' && request.url_path == '/v1/model-endpoints-schema.json' + value: get_/v1/model-endpoints-schema.json +- condition: request.method == 'GET' && request.url_path == '/v1/model-endpoints-api' + value: get_/v1/model-endpoints-api +- condition: request.method == 'POST' && request.url_path == '/v1/docker-image-batch-job-bundles' + value: post_/v1/docker-image-batch-job-bundles +- condition: request.method == 'GET' && request.url_path == '/v1/docker-image-batch-job-bundles' + value: get_/v1/docker-image-batch-job-bundles +- condition: request.method == 'GET' && request.url_path == '/v1/docker-image-batch-job-bundles/latest' + value: get_/v1/docker-image-batch-job-bundles/latest +- condition: request.method == 'GET' && request.url_path.matches('^/v1/docker-image-batch-job-bundles/[[:alnum:]-_]*$') + value: get_/v1/docker-image-batch-job-bundles/_docker_image_batch_job_bundle_id +- condition: request.method == 'POST' && request.url_path == '/v1/llm/model-endpoints' + value: post_/v1/llm/model-endpoints +- condition: request.method == 'GET' && request.url_path == '/v1/llm/model-endpoints' + value: get_/v1/llm/model-endpoints +- condition: request.method == 'GET' && request.url_path.matches('^/v1/llm/model-endpoints/[[:alnum:]-_]*$') + value: get_/v1/llm/model-endpoints/_model_endpoint_name +- condition: request.method == 'PUT' && request.url_path.matches('^/v1/llm/model-endpoints/[[:alnum:]-_]*$') + value: put_/v1/llm/model-endpoints/_model_endpoint_name +- condition: request.method == 'DELETE' && request.url_path.matches('^/v1/llm/model-endpoints/[[:alnum:]-_]*$') + value: delete_/v1/llm/model-endpoints/_model_endpoint_name +- condition: request.method == 'POST' && request.url_path == '/v1/llm/completions-sync' + value: post_/v1/llm/completions-sync +- condition: request.method == 'POST' && request.url_path == '/v1/llm/completions-stream' + value: post_/v1/llm/completions-stream +- condition: request.method == 'GET' && request.url_path == '/v1/llm/fine-tunes' + value: get_/v1/llm/fine-tunes +- condition: request.method == 'POST' && request.url_path == '/v1/llm/fine-tunes' + value: post_/v1/llm/fine-tunes +- condition: request.method == 'GET' && request.url_path.matches('^/v1/llm/fine-tunes/[[:alnum:]-_]*$') + value: get_/v1/llm/fine-tunes/_fine_tune_id +- condition: request.method == 'PUT' && request.url_path.matches('^/v1/llm/fine-tunes/[[:alnum:]-_]*/cancel$') + value: put_/v1/llm/fine-tunes/_fine_tune_id/cancel +- condition: request.method == 'GET' && request.url_path.matches('^/v1/llm/fine-tunes/[[:alnum:]-_]*/events$') + value: get_/v1/llm/fine-tunes/_fine_tune_id/events +- condition: request.method == 'POST' && request.url_path == '/v1/llm/model-endpoints/download' + value: post_/v1/llm/model-endpoints/download +- condition: request.method == 'POST' && request.url_path == '/v1/llm/batch-completions' + value: post_/v1/llm/batch-completions +- condition: request.method == 'GET' && request.url_path == '/v1/files' + value: get_/v1/files +- condition: request.method == 'POST' && request.url_path == '/v1/files' + value: post_/v1/files +- condition: request.method == 'GET' && request.url_path.matches('^/v1/files/[[:alnum:]-_]*$') + value: get_/v1/files/_file_id +- condition: request.method == 'DELETE' && request.url_path.matches('^/v1/files/[[:alnum:]-_]*$') + value: delete_/v1/files/_file_id +- condition: request.method == 'GET' && request.url_path.matches('^/v1/files/[[:alnum:]-_]*/content$') + value: get_/v1/files/_file_id/content - condition: request.method == 'GET' && request.url_path == '/v1/triggers' value: get_/v1/triggers - condition: request.method == 'POST' && request.url_path == '/v1/triggers' value: post_/v1/triggers -- condition: request.method == 'DELETE' && request.url_path.matches('^/v1/triggers/[[:alnum:]-_]*$') - value: delete_/v1/triggers/_trigger_id - condition: request.method == 'GET' && request.url_path.matches('^/v1/triggers/[[:alnum:]-_]*$') value: get_/v1/triggers/_trigger_id - condition: request.method == 'PUT' && request.url_path.matches('^/v1/triggers/[[:alnum:]-_]*$') value: put_/v1/triggers/_trigger_id -- condition: request.method == 'GET' && request.url_path == '/v2/model-bundles' - value: get_/v2/model-bundles -- condition: request.method == 'POST' && request.url_path == '/v2/model-bundles' - value: post_/v2/model-bundles -- condition: request.method == 'POST' && request.url_path == '/v2/model-bundles/clone-with-changes' - value: post_/v2/model-bundles/clone-with-changes -- condition: request.method == 'GET' && request.url_path == '/v2/model-bundles/latest' - value: get_/v2/model-bundles/latest -- condition: request.method == 'GET' && request.url_path.matches('^/v2/model-bundles/[[:alnum:]-_]*$') - value: get_/v2/model-bundles/_model_bundle_id +- condition: request.method == 'DELETE' && request.url_path.matches('^/v1/triggers/[[:alnum:]-_]*$') + value: delete_/v1/triggers/_trigger_id +- condition: request.method == 'POST' && request.url_path == '/v2/batch-completions' + value: post_/v2/batch-completions +- condition: request.method == 'GET' && request.url_path.matches('^/v2/batch-completions/[[:alnum:]-_]*$') + value: get_/v2/batch-completions/_batch_completion_id +- condition: request.method == 'POST' && request.url_path.matches('^/v2/batch-completions/[[:alnum:]-_]*$') + value: post_/v2/batch-completions/_batch_completion_id +- condition: request.method == 'POST' && request.url_path.matches('^/v2/batch-completions/[[:alnum:]-_]*/actions/cancel$') + value: post_/v2/batch-completions/_batch_completion_id/actions/cancel +- condition: request.method == 'POST' && request.url_path == '/v2/chat/completions' + value: post_/v2/chat/completions +- condition: request.method == 'POST' && request.url_path == '/v2/completions' + value: post_/v2/completions +- condition: request.method == 'GET' && request.url_path == '/healthcheck' + value: get_/healthcheck +- condition: request.method == 'GET' && request.url_path == '/healthz' + value: get_/healthz +- condition: request.method == 'GET' && request.url_path == '/readyz' + value: get_/readyz {{- end -}} diff --git a/charts/model-engine/templates/gateway_deployment.yaml b/charts/model-engine/templates/gateway_deployment.yaml index 726d6e85a..fd2c86272 100644 --- a/charts/model-engine/templates/gateway_deployment.yaml +++ b/charts/model-engine/templates/gateway_deployment.yaml @@ -28,6 +28,10 @@ spec: }] sidecar.istio.io/proxyMemoryLimit: "5Gi" sidecar.istio.io/proxyMemory: "1Gi" + {{- /* Without a CPU request the sidecar is starved on nodes packed to their CPU + request capacity, its postStart hook hangs, and new gateway pods never + become Ready. No CPU limit: throttling the proxy adds tail latency. */}} + sidecar.istio.io/proxyCPU: {{ .Values.gateway.sidecarCPURequest | quote }} labels: {{- include "modelEngine.selectorLabels.gateway" . | nindent 8 }} {{- include "modelEngine.labels" . | nindent 8 }} @@ -60,6 +64,7 @@ spec: path: /readyz port: 5000 periodSeconds: 2 + timeoutSeconds: {{ .Values.gateway.readinessProbeTimeoutSeconds }} failureThreshold: 30 command: - dumb-init diff --git a/charts/model-engine/templates/istio-ratelimit-envoyfilter.yaml b/charts/model-engine/templates/istio-ratelimit-envoyfilter.yaml new file mode 100644 index 000000000..0eefd519c --- /dev/null +++ b/charts/model-engine/templates/istio-ratelimit-envoyfilter.yaml @@ -0,0 +1,142 @@ +{{- if (.Values.rateLimits).enabled -}} +{{- $fullName := include "modelEngine.fullname" . -}} +{{- $routes := .Values.rateLimits.routes -}} +{{- $tenants := .Values.rateLimits.throttledTenants | default list -}} +{{- range $tenants }} +{{- /* The Basic clamp matches on the unpadded base64 of the raw user id, which is + only a prefix of base64("userId:password") when len(userId) is divisible by 3 + (a 24-char id). Anything else would half-work (Bearer only), silently. */}} +{{- if ne (mod (len .userId) 3) 0 }} +{{- fail (printf "rateLimits.throttledTenants userId %q length must be divisible by 3 for the Basic-auth prefix match" .userId) }} +{{- end }} +{{- end }} +# Envoy local_ratelimit on the gateway sidecars (inbound), giving per-pod token +# buckets. Buckets are per Envoy process, so the fleet-wide ceiling is +# (limit x Ready pods) and scales with the HPA. Overflow is rejected at the proxy +# as 429 and never consumes a gateway worker. +apiVersion: networking.istio.io/v1alpha3 +kind: EnvoyFilter +metadata: + name: {{ $fullName }}-ratelimit + namespace: {{ .Release.Namespace }} + labels: + {{- include "modelEngine.labels" . | nindent 4 }} +spec: + workloadSelector: + labels: + {{- include "modelEngine.selectorLabels.gateway" . | nindent 6 }} + configPatches: + - applyTo: HTTP_FILTER + match: + context: SIDECAR_INBOUND + listener: + filterChain: + filter: + name: envoy.filters.network.http_connection_manager + subFilter: + name: envoy.filters.http.router + patch: + operation: INSERT_BEFORE + value: + name: envoy.filters.http.local_ratelimit + typed_config: + "@type": type.googleapis.com/envoy.extensions.filters.http.local_ratelimit.v3.LocalRateLimit + stat_prefix: model_engine_ratelimit + status: + code: TooManyRequests + # Envoy 1.23 applies stage-0 virtual-host rate limits automatically. It has + # neither vh_rate_limits nor always_consume_default_token_bucket, and an + # unknown field makes istiod skip the whole patch silently. The default + # bucket is therefore effectively unlimited because every request consumes it. + token_bucket: + max_tokens: 100000 + tokens_per_fill: 100000 + fill_interval: 1s + filter_enabled: + default_value: + numerator: 100 + denominator: HUNDRED + filter_enforced: + default_value: + numerator: 100 + denominator: HUNDRED + response_headers_to_add: + - append_action: OVERWRITE_IF_EXISTS_OR_ADD + header: + key: x-local-rate-limited + value: "true" + descriptors: + {{- range $routes }} + - entries: + - key: route + value: {{ .name }} + token_bucket: + max_tokens: {{ .burst }} + tokens_per_fill: {{ .requestsPerSecondPerPod }} + fill_interval: 1s + {{- end }} + {{- range $tenants }} + - entries: + - key: tenant + value: {{ printf "clamp_%s" .userId | quote }} + token_bucket: + max_tokens: {{ .requestsPerSecondPerPod }} + tokens_per_fill: {{ .requestsPerSecondPerPod }} + fill_interval: 1s + {{- end }} + - applyTo: HTTP_ROUTE + match: + context: SIDECAR_INBOUND + routeConfiguration: + vhost: + # Istio names inbound virtual hosts from the Service port. Verified empirically + # on both Istio 1.15 (Envoy config_dump on ml-training-new: inbound|http|80 + # exists, inbound|http|5000 does not) and Istio 1.30; a wrong name here makes + # the patch skip silently and disables all route/tenant buckets. + name: {{ printf "inbound|http|%v" .Values.service.port | quote }} + patch: + operation: MERGE + # Rate-limit actions must live on the ROUTE, not the virtual host: Envoy 1.23's + # local_ratelimit only reads route-level policies (vh_rate_limits arrived in + # 1.24+), so a VIRTUAL_HOST patch is silently inert on prod sidecars. + value: + route: + rate_limits: + {{- range $routes }} + - actions: + - header_value_match: + descriptor_key: route + descriptor_value: {{ .name }} + headers: + - name: ":method" + string_match: + exact: {{ .method }} + - name: ":path" + string_match: + prefix: {{ .pathPrefix }} + {{- end }} + {{- range $tenants }} + {{- $tenant := . }} + {{- /* Callers authenticate with HTTP Basic (base64 of "userId:password") + or Bearer (raw userId). A 24-char user id encodes to exactly 32 + base64 chars, so that prefix matches any password. The scheme is + matched case-insensitively (FastAPI accepts "bearer"/"basic" too, so + an exact-case match is trivially bypassed), the credential exactly. + The :path match scopes the clamp to the async-task routes only. */}} + {{- range list (printf "(?i:Basic) %s.*" (b64enc $tenant.userId)) (printf "(?i:Bearer) %s" $tenant.userId) }} + - actions: + - header_value_match: + descriptor_key: tenant + descriptor_value: {{ printf "clamp_%s" $tenant.userId | quote }} + headers: + - name: authorization + string_match: + safe_regex: + google_re2: {} + regex: {{ . | quote }} + - name: ":path" + string_match: + prefix: /v1/async-tasks + {{- end }} + {{- end }} +{{- end }} diff --git a/charts/model-engine/templates/istio-virtualservice.yaml b/charts/model-engine/templates/istio-virtualservice.yaml index 1bd26e144..7c08a6060 100644 --- a/charts/model-engine/templates/istio-virtualservice.yaml +++ b/charts/model-engine/templates/istio-virtualservice.yaml @@ -26,6 +26,12 @@ spec: port: number: 80 retries: - attempts: 3 - retryOn: connect-failure,unavailable,gateway-error + attempts: {{ .Values.gateway.retries.attempts }} + {{- /* 503 is deliberately not retried: it is the overload signal, and retrying + it multiplies offered load exactly when the fleet is saturated. Numeric + codes in retryOn map to Envoy retriable-status-codes. */}} + retryOn: {{ .Values.gateway.retries.retryOn | quote }} + {{- with .Values.gateway.retries.perTryTimeout }} + perTryTimeout: {{ . }} + {{- end }} {{- end }} diff --git a/charts/model-engine/values.yaml b/charts/model-engine/values.yaml index 15ba6bf97..66b125fd6 100644 --- a/charts/model-engine/values.yaml +++ b/charts/model-engine/values.yaml @@ -33,6 +33,56 @@ celery_broker_type_redis: null # drop: # - ALL +# gateway [optional] tuning for the gateway deployment. +gateway: + # readinessProbeTimeoutSeconds sets the /readyz probe timeout. The k8s default of 1s + # ejects saturated-but-healthy pods and concentrates load on the survivors. + readinessProbeTimeoutSeconds: 5 + # sidecarCPURequest sets the istio-proxy CPU request (no limit). Without it the + # sidecar is starved on nodes at full CPU request capacity and pods never go Ready. + sidecarCPURequest: 250m + # retries is the Envoy retry policy on the gateway VirtualService. perTryTimeout + # stays unset by default: the route carries streaming and long-lived requests, + # which a per-try timeout would abort mid-flight. + retries: + attempts: 3 + retryOn: connect-failure,unavailable,502,504 + perTryTimeout: null + +# rateLimits [optional] per-pod Envoy local_ratelimit token buckets on the gateway +# sidecars (inbound). Overflow returns 429 at the proxy without reaching a gateway +# worker. Buckets are per pod: the fleet ceiling is limit x Ready pods. +rateLimits: + enabled: false + # Route matching is on Envoy :path, which includes the query string, so prefixes + # (not exact paths) are required. The trailing slash on the poll route keeps + # POST /v1/async-tasks out of the GET bucket. + routes: + # Task-status polls. A runaway-storm ceiling AND the aggregate byte bound: polls + # return task results inline (up to multi-MB), and the 2026-08 incident's binding + # resource was node NIC egress (~500MB/s per pod). 50 rps/pod caps worst-case + # egress at ~265MB/s/pod while still admitting ~15x baseline traffic (1,500 rps + # at a 30-pod fleet, ~= the incident's peak). Revisit upward once inline results + # are replaced by presigned URLs (MLI-8311). + - name: get_async_task + method: GET + pathPrefix: /v1/async-tasks/ + requestsPerSecondPerPod: 50 + burst: 100 + # Task submissions. Proportionally tighter: producing above endpoint drain is pure + # harm, and each submit costs a Postgres reader lookup + SQS send. + - name: post_async_tasks + method: POST + pathPrefix: /v1/async-tasks + requestsPerSecondPerPod: 20 + burst: 40 + # throttledTenants clamps named callers on the async-task routes without Helm + # drift. userId is the caller's user id as logged by the gateway; both accepted + # Authorization forms (Basic and Bearer) are matched. + throttledTenants: [] + # - userId: "000000000000000000000000" + # requestsPerSecondPerPod: 1 + redis: auth: authSecretName: "" diff --git a/docs/guides/rate_limits.md b/docs/guides/rate_limits.md index 2aa59dd47..4ae4c4552 100644 --- a/docs/guides/rate_limits.md +++ b/docs/guides/rate_limits.md @@ -14,6 +14,36 @@ The LLM Engine API is currently in a preview mode, and therefore we currently do As the API moves towards a production release, we will update this section with specific rate limits. For now, the API will return HTTP 429 on an as-needed basis. +# Polling async tasks + +When you submit async inference tasks, poll each task's status with a delay that grows +between attempts instead of a fixed tight interval. With a fixed interval, total poll +traffic scales with the number of outstanding tasks: a large batch that outpaces the +endpoint's drain rate generates an ever-growing poll load on the API while the tasks it +is asking about cannot complete any faster. + +The `tenacity` library (used in the backoff examples below) expresses this directly; +`wait_random_exponential` adds the jitter that keeps polls from many tasks from +synchronizing: + +=== "Polling with exponential backoff in python" + +```python +from tenacity import retry, retry_if_result, stop_after_delay, wait_random_exponential + +@retry( + retry=retry_if_result(lambda response: response.status not in ("SUCCESS", "FAILURE")), + wait=wait_random_exponential(min=1, max=60), + stop=stop_after_delay(3600), +) +def wait_for_task(client, task_id): + return client.get_async_task(task_id) +``` + +For large batches, also bound how many tasks you have outstanding (submit in windows +sized to the endpoint's throughput) rather than submitting everything up front and +polling the whole set. + # Error mitigation ## Retrying with exponential backoff diff --git a/model-engine/model_engine_server/api/app.py b/model-engine/model_engine_server/api/app.py index 534aebb54..4c77e1def 100644 --- a/model-engine/model_engine_server/api/app.py +++ b/model-engine/model_engine_server/api/app.py @@ -313,8 +313,13 @@ def load_redis(): get_or_create_aioredis_pool() -def healthcheck() -> Response: - """Returns 200 if the app is healthy.""" +async def healthcheck() -> Response: + """Returns 200 if the app is healthy. + + Must be async: a sync handler runs in the shared threadpool, so under load the + probe queues behind blocked requests and misses its timeout, ejecting pods that + are saturated but healthy. + """ return Response(status_code=200) diff --git a/model-engine/model_engine_server/api/rate_limits.py b/model-engine/model_engine_server/api/rate_limits.py new file mode 100644 index 000000000..a0c07a0b9 --- /dev/null +++ b/model-engine/model_engine_server/api/rate_limits.py @@ -0,0 +1,154 @@ +"""Per-user request rate limiting, enforced after authentication resolves identity. + +Limits are fixed 1-second windows in Redis, keyed on (route class, user_id). +The limiter fails open: if Redis is unavailable or slow, requests are allowed. +""" + +import asyncio +import time +from typing import Optional + +import redis.asyncio as aioredis +from datadog import statsd +from fastapi import Depends, HTTPException, status +from model_engine_server.api.dependencies import get_or_create_aioredis_pool, verify_authentication +from model_engine_server.common.config import hmi_config +from model_engine_server.common.env_vars import DD_ENV +from model_engine_server.core.auth.authentication_repository import User +from model_engine_server.core.loggers import logger_name, make_logger + +logger = make_logger(logger_name()) + +# The limiter must stay cheap when the pod is saturated; a slow Redis answer is +# treated as an outage and fails open. +_REDIS_TIMEOUT_SECONDS = 0.1 +_LOG_INTERVAL_SECONDS = 60.0 +# Circuit breaker: each timed-out call abandons its pooled connection, so during a +# Redis brownout per-request checks become a reconnect storm against the slow Redis +# (and churn the shared cache pool). After enough consecutive failures, fail open +# without touching Redis for a cooldown period. +_BREAKER_FAILURE_THRESHOLD = 5 +_BREAKER_COOLDOWN_SECONDS = 10.0 + +# Atomic so a partial failure cannot leave a counter key without a TTL. +_INCR_WITH_TTL_LUA = """ +local count = redis.call('INCR', KEYS[1]) +if count == 1 then + redis.call('EXPIRE', KEYS[1], 2) +end +return count +""" + +_client: Optional[aioredis.Redis] = None +_client_pool: Optional[aioredis.ConnectionPool] = None +_last_log_times: dict = {} +_consecutive_failures: int = 0 +_breaker_open_until: float = 0.0 + + +def _emit_decision(outcome: str, route_class: str, user_id: str) -> None: + # Fire-and-forget UDP to the local Datadog agent. This doubles as the per-tenant + # volume signal on the rate-limited routes (outcome:allowed) and as enforcement + # telemetry (throttled/would_throttle/fail_open/breaker_open), replacing the + # postmortem's proposed raw request-volume monitor. + statsd.increment( + "model_engine.user_rate_limit.decision", + tags=[ + f"env:{DD_ENV}", + f"route_class:{route_class}", + f"outcome:{outcome}", + f"user_id:{user_id}", + ], + ) + + +def _should_log(log_key: str) -> bool: + # Both limiter logs fire per request when things go wrong (Redis outage, or a + # noisy tenant in log-only mode), so they are sampled per key. + now = time.monotonic() + if now - _last_log_times.get(log_key, 0.0) < _LOG_INTERVAL_SECONDS: + return False + if len(_last_log_times) > 1000: + _last_log_times.clear() + _last_log_times[log_key] = now + return True + + +def _get_client() -> aioredis.Redis: + # Cached per pool: the pool is rebuilt when its credentials expire, and + # constructing a Redis client per request copies its full callback tables. + global _client, _client_pool + pool = get_or_create_aioredis_pool() + if _client is None or _client_pool is not pool: + _client = aioredis.Redis(connection_pool=pool) + _client_pool = pool + return _client + + +async def _count_request(key: str) -> int: + return int(await _get_client().eval(_INCR_WITH_TTL_LUA, 1, key)) + + +async def enforce_user_rate_limit(route_class: str, user: User) -> None: + """Raises 429 with Retry-After if the user is over their per-route limit. + + No-op unless `user_rate_limits` is configured; counts but does not reject + unless `user_rate_limits.enforce` is true (log-only rollout mode). + """ + config = hmi_config.user_rate_limits or {} + limit = (config.get("routes") or {}).get(route_class) + if not limit: + return + limit = int(limit) + global _consecutive_failures, _breaker_open_until + if time.monotonic() < _breaker_open_until: + _emit_decision("breaker_open", route_class, user.user_id) + return + key = f"user-rate-limit:{route_class}:{user.user_id}:{int(time.time())}" + try: + count = await asyncio.wait_for(_count_request(key), timeout=_REDIS_TIMEOUT_SECONDS) + _consecutive_failures = 0 + except Exception: + _consecutive_failures += 1 + if _consecutive_failures >= _BREAKER_FAILURE_THRESHOLD: + _breaker_open_until = time.monotonic() + _BREAKER_COOLDOWN_SECONDS + _consecutive_failures = 0 + if _should_log("fail-open"): + logger.warning( + f"Rate limiter failing open for route_class={route_class}", exc_info=True + ) + _emit_decision("fail_open", route_class, user.user_id) + return + if count <= limit: + _emit_decision("allowed", route_class, user.user_id) + return + if not config.get("enforce"): + if _should_log(f"over-limit:{user.user_id}:{route_class}"): + logger.warning( + f"Rate limit exceeded (log-only): user_id={user.user_id} " + f"route_class={route_class} count={count} limit={limit}" + ) + _emit_decision("would_throttle", route_class, user.user_id) + return + _emit_decision("throttled", route_class, user.user_id) + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail=( + f"Rate limit exceeded for {route_class}: {limit} requests per second per user. " + "Retry after the indicated delay." + ), + headers={"Retry-After": "1"}, + ) + + +def user_rate_limit(route_class: str): + """FastAPI dependency limiting the authenticated user on this route. + + Composes with verify_authentication (FastAPI caches it per request, so + authentication still runs once). + """ + + async def dependency(auth: User = Depends(verify_authentication)) -> None: + await enforce_user_rate_limit(route_class, auth) + + return dependency diff --git a/model-engine/model_engine_server/api/tasks_v1.py b/model-engine/model_engine_server/api/tasks_v1.py index c256778a0..fbd5bb77e 100644 --- a/model-engine/model_engine_server/api/tasks_v1.py +++ b/model-engine/model_engine_server/api/tasks_v1.py @@ -1,11 +1,14 @@ import asyncio +import functools +import anyio from fastapi import APIRouter, Depends, HTTPException from model_engine_server.api.dependencies import ( ExternalInterfaces, get_external_interfaces_read_only, verify_authentication, ) +from model_engine_server.api.rate_limits import user_rate_limit from model_engine_server.common.dtos.tasks import ( CreateAsyncTaskV1Response, EndpointPredictV1Request, @@ -40,12 +43,22 @@ logger = make_logger(logger_name()) +# Task-status polls do a blocking result-backend read (S3 on AWS). They must not share +# the default anyio threadpool: poll volume scales with the number of outstanding tasks, +# and when polls fill the shared pool every other sync route queues behind them. +# Created lazily because anyio.CapacityLimiter requires a running event loop. +@functools.cache +def _get_task_limiter() -> anyio.CapacityLimiter: + return anyio.CapacityLimiter(40) + + @inference_task_router_v1.post("/async-tasks", response_model=CreateAsyncTaskV1Response) async def create_async_inference_task( model_endpoint_id: str, request: EndpointPredictV1Request, auth: User = Depends(verify_authentication), external_interfaces: ExternalInterfaces = Depends(get_external_interfaces_read_only), + _rate_limit: None = Depends(user_rate_limit("post_async_tasks")), ) -> CreateAsyncTaskV1Response: """ Runs an async inference prediction. @@ -81,10 +94,11 @@ async def create_async_inference_task( @inference_task_router_v1.get("/async-tasks/{task_id}", response_model=GetAsyncTaskV1Response) -def get_async_inference_task( +async def get_async_inference_task( task_id: str, auth: User = Depends(verify_authentication), external_interfaces: ExternalInterfaces = Depends(get_external_interfaces_read_only), + _rate_limit: None = Depends(user_rate_limit("get_async_task")), ) -> GetAsyncTaskV1Response: """ Gets the status of an async inference task. @@ -94,7 +108,10 @@ def get_async_inference_task( use_case = GetAsyncInferenceTaskV1UseCase( model_endpoint_service=external_interfaces.model_endpoint_service, ) - return use_case.execute(user=auth, task_id=task_id) + return await anyio.to_thread.run_sync( + functools.partial(use_case.execute, user=auth, task_id=task_id), + limiter=_get_task_limiter(), + ) except (ObjectNotFoundException, ObjectNotAuthorizedException) as exc: raise HTTPException( status_code=404, diff --git a/model-engine/model_engine_server/common/config.py b/model-engine/model_engine_server/common/config.py index 9e6e39b81..d85573f19 100644 --- a/model-engine/model_engine_server/common/config.py +++ b/model-engine/model_engine_server/common/config.py @@ -81,6 +81,14 @@ class HostedModelInferenceServiceConfig: None # For GCP Memorystore (e.g., redis://MEMORYSTORE_HOST:6379/0) ) sglang_repository: Optional[str] = None + # Per-user request rate limits; route classes are the names passed to + # api.rate_limits.user_rate_limit on each route. E.g. + # user_rate_limits: + # enforce: false # false = log-only mode + # routes: + # get_async_task: 200 # requests/second per user; unset = unlimited + # post_async_tasks: 25 + user_rate_limits: Optional[dict] = None @classmethod def from_json(cls, json): diff --git a/model-engine/model_engine_server/inference/forwarding/celery_forwarder.py b/model-engine/model_engine_server/inference/forwarding/celery_forwarder.py index cb35aa875..cf26a3f2b 100644 --- a/model-engine/model_engine_server/inference/forwarding/celery_forwarder.py +++ b/model-engine/model_engine_server/inference/forwarding/celery_forwarder.py @@ -101,6 +101,9 @@ def create_celery_service( ) monitoring_metrics_gateway = DatadogInferenceMonitoringMetricsGateway() + # requests treats its timeout as socket inactivity. The Celery limit supplies + # the wall-clock bound when an upstream response keeps trickling bytes. + task_time_limit = forwarder.timeout_seconds class ErrorHandlingTask(Task): """Sets a 'custom' field with error in the Task response for FAILURE. @@ -154,6 +157,7 @@ def after_return( name=LIRA_CELERY_TASK_NAME, track_started=True, autoretry_for=(ConnectionError,), + time_limit=task_time_limit, ) def exec_func( payload, @@ -193,6 +197,7 @@ def exec_func( base=ErrorHandlingTask, name=DEFAULT_CELERY_TASK_NAME, track_started=True, + time_limit=task_time_limit, ) def exec_func_pre_lira(payload, arrival_timestamp, *ignored_args, **ignored_kwargs): return exec_func(payload, arrival_timestamp, *ignored_args, **ignored_kwargs) diff --git a/model-engine/model_engine_server/inference/forwarding/forwarding.py b/model-engine/model_engine_server/inference/forwarding/forwarding.py index 5183955b9..64dec2071 100644 --- a/model-engine/model_engine_server/inference/forwarding/forwarding.py +++ b/model-engine/model_engine_server/inference/forwarding/forwarding.py @@ -1,5 +1,6 @@ import ast import json +import math import os import time from dataclasses import dataclass @@ -42,6 +43,8 @@ DEFAULT_PORT: int = 5005 +DEFAULT_SYNC_TIMEOUT_SECONDS: float = 3600 + class ModelEngineSerializationMixin: """Mixin class for optionally wrapping Model Engine requests.""" @@ -169,6 +172,10 @@ class Forwarder(ModelEngineSerializationMixin): # We do this to avoid having to put this data in any sync response and only do it for async responses forward_http_status_in_body: bool post_inference_hooks_handler: Optional[PostInferenceHooksHandler] = None + # Cap on the full round-trip to the user-defined service. Must be explicit: without + # one, aiohttp applies its default total=300s and long-running non-streaming + # generations get cut off with a 500 while the model server keeps computing. + timeout_seconds: float = DEFAULT_SYNC_TIMEOUT_SECONDS async def forward(self, json_payload: Any, trace_config: Optional[str] = None) -> Any: json_payload, using_serialize_results_as_string = self.unwrap_json_payload(json_payload) @@ -185,6 +192,7 @@ async def forward(self, json_payload: Any, trace_config: Optional[str] = None) - self.predict_endpoint, json=json_payload, headers=headers, + timeout=aiohttp.ClientTimeout(total=self.timeout_seconds), ) response = await response_raw.json( content_type=None @@ -239,6 +247,7 @@ def __call__(self, json_payload: Any, trace_config: Optional[str] = None) -> Any self.predict_endpoint, json=json_payload, headers=headers, + timeout=self.timeout_seconds, ) response = response_raw.json() except Exception: @@ -297,8 +306,17 @@ class LoadForwarder: wrap_response: bool = True forward_http_status: bool = False forward_http_status_in_body: bool = False + timeout_seconds: float = DEFAULT_SYNC_TIMEOUT_SECONDS def load(self, resources: Optional[Path], cache: Any) -> Forwarder: + if ( + not isinstance(self.timeout_seconds, (int, float)) + or isinstance(self.timeout_seconds, bool) + or not math.isfinite(self.timeout_seconds) + or self.timeout_seconds <= 0 + ): + raise ValueError(f"timeout_seconds must be a positive number: {self.timeout_seconds=}") + if self.use_grpc: raise NotImplementedError( "User-defined service **MUST** use HTTP at the moment. " @@ -405,6 +423,7 @@ def endpoint(route: str) -> str: wrap_response=self.wrap_response, forward_http_status=self.forward_http_status, forward_http_status_in_body=self.forward_http_status_in_body, + timeout_seconds=self.timeout_seconds, ) diff --git a/model-engine/tests/unit/api/test_app.py b/model-engine/tests/unit/api/test_app.py index 4fe760151..dfdca32f5 100644 --- a/model-engine/tests/unit/api/test_app.py +++ b/model-engine/tests/unit/api/test_app.py @@ -1,10 +1,17 @@ +import asyncio +import threading + +import anyio.to_thread +import pytest from fastapi import FastAPI from fastapi.testclient import TestClient +from httpx import AsyncClient from model_engine_server.api.app import ( OPENAPI_SCHEMA_RENAME_PATTERNS, CustomMiddleware, _convert_openapi_31_to_30, _rename_openapi_schemas, + app, get_openapi_schema, ) from starlette.middleware import Middleware @@ -21,6 +28,30 @@ def test_healthcheck(simple_client: TestClient): assert response.status_code == 200 +@pytest.mark.asyncio +async def test_readyz_answers_while_default_threadpool_blocked(): + release = threading.Event() + limiter = anyio.to_thread.current_default_thread_limiter() + blockers = [ + asyncio.create_task(anyio.to_thread.run_sync(release.wait, 10)) + for _ in range(int(limiter.total_tokens)) + ] + try: + deadline = asyncio.get_event_loop().time() + 5 + while limiter.borrowed_tokens < limiter.total_tokens and ( + asyncio.get_event_loop().time() < deadline + ): + await asyncio.sleep(0.01) + assert limiter.borrowed_tokens == limiter.total_tokens + # The probe must not ride the (saturated) threadpool. + async with AsyncClient(app=app, base_url="http://test") as client: + response = await asyncio.wait_for(client.get("/readyz"), 1) + assert response.status_code == 200 + finally: + release.set() + await asyncio.gather(*blockers) + + def _get_unhandled_exception_response(error_details: str): test_app = FastAPI(middleware=[Middleware(CustomMiddleware)]) diff --git a/model-engine/tests/unit/api/test_rate_limits.py b/model-engine/tests/unit/api/test_rate_limits.py new file mode 100644 index 000000000..06485b5e4 --- /dev/null +++ b/model-engine/tests/unit/api/test_rate_limits.py @@ -0,0 +1,330 @@ +import asyncio +import dataclasses +import logging +import time + +import pytest +import redis.asyncio as aioredis +from fastapi import Depends, HTTPException +from model_engine_server.api import rate_limits +from model_engine_server.api.app import app +from model_engine_server.api.dependencies import basic_auth, oauth2_scheme, verify_authentication +from model_engine_server.common.config import HostedModelInferenceServiceConfig +from model_engine_server.core.auth.authentication_repository import User +from tests.unit.api.conftest import fake_verify_authentication, get_test_auth_repository + + +class FakeRateLimitRedis: + """Counter-only fake for the limiter's eval call.""" + + def __init__(self, count=1, error=None, delay=0.0, sequence=None): + self.count = count + self.error = error + self.delay = delay + self.sequence = list(sequence) if sequence is not None else None + self.calls = 0 + self.keys = [] + + async def eval(self, script, numkeys, *keys): + self.calls += 1 + self.keys.extend(keys) + if self.delay: + await asyncio.sleep(self.delay) + if self.sequence is not None: + item = self.sequence.pop(0) + if isinstance(item, Exception): + raise item + return item + if self.error: + raise self.error + return self.count + + +class FakeClock: + """Stands in for the `time` module inside rate_limits (monotonic + wall clock).""" + + def __init__(self): + self.now = time.monotonic() + self.wall = time.time() + + def monotonic(self): + return self.now + + def time(self): + return self.wall + + +def _install_fake_clock(monkeypatch) -> FakeClock: + clock = FakeClock() + monkeypatch.setattr(rate_limits, "time", clock) + return clock + + +@pytest.fixture(autouse=True) +def _reset_limiter_state(monkeypatch): + monkeypatch.setattr(rate_limits, "_consecutive_failures", 0) + monkeypatch.setattr(rate_limits, "_breaker_open_until", 0.0) + monkeypatch.setattr(rate_limits, "_last_log_times", {}) + monkeypatch.setattr(rate_limits, "_client", None) + monkeypatch.setattr(rate_limits, "_client_pool", None) + + +USER = User(user_id="test-user", team_id="test-team", is_privileged_user=False) +LIMITS = {"enforce": True, "routes": {"get_async_task": 5}} +LIMITS_LOG_ONLY = {"enforce": False, "routes": {"get_async_task": 5}} +LIMITS_STRING_VALUE = {"enforce": True, "routes": {"get_async_task": "5"}} + + +@pytest.mark.parametrize( + "config,route_class,count,expect_429", + [ + pytest.param(None, "get_async_task", 100, False, id="disabled-no-config"), + pytest.param(LIMITS, "get_async_task", 5, False, id="at-limit-allowed"), + pytest.param(LIMITS, "get_async_task", 6, True, id="over-limit-rejected"), + pytest.param(LIMITS, "post_async_tasks", 100, False, id="unconfigured-route-unlimited"), + pytest.param(LIMITS_LOG_ONLY, "get_async_task", 100, False, id="log-only-allows"), + pytest.param(LIMITS_STRING_VALUE, "get_async_task", 6, True, id="string-limit-coerced"), + ], +) +@pytest.mark.asyncio +async def test_enforce_user_rate_limit(monkeypatch, config, route_class, count, expect_429): + monkeypatch.setattr(rate_limits.hmi_config, "user_rate_limits", config, raising=False) + monkeypatch.setattr(rate_limits, "_get_client", lambda: FakeRateLimitRedis(count=count)) + if expect_429: + with pytest.raises(HTTPException) as exc_info: + await rate_limits.enforce_user_rate_limit(route_class, USER) + assert exc_info.value.status_code == 429 + assert exc_info.value.headers is not None + assert "Retry-After" in exc_info.value.headers + else: + await rate_limits.enforce_user_rate_limit(route_class, USER) + + +@pytest.mark.parametrize( + "redis", + [ + pytest.param(FakeRateLimitRedis(error=ConnectionError("redis down")), id="redis-error"), + pytest.param(FakeRateLimitRedis(count=100, delay=1.0), id="redis-slow"), + ], +) +@pytest.mark.asyncio +async def test_fail_open(monkeypatch, redis): + monkeypatch.setattr(rate_limits.hmi_config, "user_rate_limits", LIMITS, raising=False) + monkeypatch.setattr(rate_limits, "_get_client", lambda: redis) + # Over-limit counts must still be allowed when Redis errors or times out. + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + + +@pytest.mark.asyncio +async def test_circuit_breaker_stops_touching_redis(monkeypatch): + monkeypatch.setattr(rate_limits.hmi_config, "user_rate_limits", LIMITS, raising=False) + redis = FakeRateLimitRedis(error=ConnectionError("redis down")) + monkeypatch.setattr(rate_limits, "_get_client", lambda: redis) + for _ in range(rate_limits._BREAKER_FAILURE_THRESHOLD + 5): + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + # After the threshold trips, the cooldown window must not touch Redis at all. + assert redis.calls == rate_limits._BREAKER_FAILURE_THRESHOLD + + +@pytest.mark.asyncio +async def test_breaker_cooldown_expiry_resumes_checks_and_enforcement(monkeypatch): + monkeypatch.setattr(rate_limits.hmi_config, "user_rate_limits", LIMITS, raising=False) + clock = _install_fake_clock(monkeypatch) + failing = FakeRateLimitRedis(error=ConnectionError("redis down")) + monkeypatch.setattr(rate_limits, "_get_client", lambda: failing) + for _ in range(rate_limits._BREAKER_FAILURE_THRESHOLD): + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + + healthy = FakeRateLimitRedis(count=6) + monkeypatch.setattr(rate_limits, "_get_client", lambda: healthy) + # Inside the cooldown window Redis stays untouched even though it recovered. + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + assert healthy.calls == 0 + + clock.now += rate_limits._BREAKER_COOLDOWN_SECONDS + 0.1 + with pytest.raises(HTTPException) as exc_info: + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + assert exc_info.value.status_code == 429 + assert healthy.calls == 1 + + +@pytest.mark.asyncio +async def test_success_resets_consecutive_failure_counter(monkeypatch): + monkeypatch.setattr(rate_limits.hmi_config, "user_rate_limits", LIMITS, raising=False) + error = ConnectionError("redis down") + # 4 failures + 1 success + 4 failures: one short of the threshold on each side. + redis = FakeRateLimitRedis(sequence=[error] * 4 + [1] + [error] * 4) + monkeypatch.setattr(rate_limits, "_get_client", lambda: redis) + for _ in range(9): + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + assert rate_limits._breaker_open_until == 0.0 + assert redis.calls == 9 + + +@pytest.mark.asyncio +async def test_log_sampling_suppresses_per_key_and_reemits(monkeypatch, caplog): + monkeypatch.setattr(rate_limits.hmi_config, "user_rate_limits", LIMITS_LOG_ONLY, raising=False) + # The limiter's logger does not propagate to root by default, which hides it from caplog. + monkeypatch.setattr(rate_limits.logger, "propagate", True) + caplog.set_level(logging.WARNING) + clock = _install_fake_clock(monkeypatch) + error = ConnectionError("redis down") + + def warning_count(): + return sum( + 1 + for record in caplog.records + if record.name == rate_limits.logger.name and record.levelno == logging.WARNING + ) + + # (redis behavior, expected cumulative warning count); None advances the clock + # past the sampling window. The fail-open and over-limit log keys are independent. + steps = [ + (error, 1), # first fail-open logged + (error, 1), # suppressed within 60s + (100, 2), # over-limit key logs despite fail-open suppression window + (100, 2), # suppressed within 60s + (None, 2), + (error, 3), # re-emitted after window + (100, 4), # re-emitted after window + ] + for behavior, expected in steps: + if behavior is None: + clock.now += rate_limits._LOG_INTERVAL_SECONDS + 0.1 + else: + fake = FakeRateLimitRedis( + error=behavior if isinstance(behavior, Exception) else None, + count=behavior if not isinstance(behavior, Exception) else 1, + ) + monkeypatch.setattr(rate_limits, "_get_client", lambda fake=fake: fake) + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + assert warning_count() == expected + + +def test_last_log_times_cleared_past_1000_keys(monkeypatch): + now = time.monotonic() + monkeypatch.setattr(rate_limits, "_last_log_times", {f"key-{i}": now for i in range(1001)}) + assert rate_limits._should_log("new-key") + assert set(rate_limits._last_log_times) == {"new-key"} + + +def test_get_client_cached_per_pool(monkeypatch): + pool_a = aioredis.ConnectionPool.from_url("redis://localhost:6379/0") + pool_b = aioredis.ConnectionPool.from_url("redis://localhost:6379/0") + pools = iter([pool_a, pool_a, pool_b]) + monkeypatch.setattr(rate_limits, "get_or_create_aioredis_pool", lambda: next(pools)) + + client_a = rate_limits._get_client() + assert rate_limits._get_client() is client_a + # A rebuilt pool (credential rotation) must produce a new client bound to it. + client_b = rate_limits._get_client() + assert client_b is not client_a + assert client_b.connection_pool is pool_b + + +@pytest.mark.asyncio +async def test_window_rollover_uses_distinct_key_per_second(monkeypatch): + monkeypatch.setattr(rate_limits.hmi_config, "user_rate_limits", LIMITS, raising=False) + clock = _install_fake_clock(monkeypatch) + redis = FakeRateLimitRedis(count=1) + monkeypatch.setattr(rate_limits, "_get_client", lambda: redis) + + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + clock.wall += 1 + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + + expected_seconds = [int(clock.wall) - 1, int(clock.wall)] + assert redis.keys == [ + f"user-rate-limit:get_async_task:{USER.user_id}:{second}" for second in expected_seconds + ] + + +def test_rate_limit_dependency_composition_over_http(monkeypatch, simple_client, test_api_key): + monkeypatch.setattr(rate_limits.hmi_config, "user_rate_limits", LIMITS, raising=False) + monkeypatch.setattr(rate_limits, "_get_client", lambda: FakeRateLimitRedis(count=100)) + auth_calls = [] + + def counting_auth( + credentials=Depends(basic_auth), + tokens=Depends(oauth2_scheme), + auth_repo=Depends(get_test_auth_repository), + ): + auth_calls.append(1) + return fake_verify_authentication(credentials, tokens, auth_repo) + + app.dependency_overrides[verify_authentication] = counting_auth + + response = simple_client.get("/v1/async-tasks/test_task_id", auth=(test_api_key, "")) + assert response.status_code == 429 + assert response.headers.get("Retry-After") == "1" + # FastAPI caches verify_authentication per request across the auth + limiter deps. + assert len(auth_calls) == 1 + # Healthchecks carry no rate-limit dependency and must be unaffected. + assert simple_client.get("/readyz").status_code == 200 + + +def _minimal_service_config_json(): + return { + field.name: "x" + for field in dataclasses.fields(HostedModelInferenceServiceConfig) + if field.default is dataclasses.MISSING and field.default_factory is dataclasses.MISSING + } + + +@pytest.mark.parametrize( + "user_rate_limits", + [ + pytest.param(None, id="absent-defaults-none"), + pytest.param({"enforce": False, "routes": {"get_async_task": 200}}, id="round-trips"), + ], +) +def test_service_config_round_trips_user_rate_limits(user_rate_limits): + raw = _minimal_service_config_json() + if user_rate_limits is not None: + raw["user_rate_limits"] = user_rate_limits + config = HostedModelInferenceServiceConfig.from_json(raw) + assert config.user_rate_limits == user_rate_limits + + +@pytest.mark.parametrize( + "config,count,error,expected_outcome", + [ + pytest.param(LIMITS, 3, None, "allowed", id="allowed"), + pytest.param(LIMITS, 6, None, "throttled", id="throttled"), + pytest.param(LIMITS_LOG_ONLY, 6, None, "would_throttle", id="would-throttle"), + pytest.param(LIMITS, 1, ConnectionError("down"), "fail_open", id="fail-open"), + ], +) +@pytest.mark.asyncio +async def test_decision_metric_emitted(monkeypatch, config, count, error, expected_outcome): + monkeypatch.setattr(rate_limits.hmi_config, "user_rate_limits", config, raising=False) + monkeypatch.setattr( + rate_limits, "_get_client", lambda: FakeRateLimitRedis(count=count, error=error) + ) + emitted = [] + monkeypatch.setattr( + rate_limits.statsd, "increment", lambda name, tags: emitted.append((name, tags)) + ) + try: + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + except HTTPException: + pass + assert len(emitted) == 1 + name, tags = emitted[0] + assert name == "model_engine.user_rate_limit.decision" + assert f"outcome:{expected_outcome}" in tags + assert "user_id:test-user" in tags + assert "route_class:get_async_task" in tags + + +@pytest.mark.asyncio +async def test_decision_metric_breaker_open(monkeypatch): + monkeypatch.setattr(rate_limits.hmi_config, "user_rate_limits", LIMITS, raising=False) + monkeypatch.setattr( + rate_limits, "_get_client", lambda: FakeRateLimitRedis(error=ConnectionError("down")) + ) + emitted = [] + monkeypatch.setattr(rate_limits.statsd, "increment", lambda name, tags: emitted.append(tags)) + for _ in range(rate_limits._BREAKER_FAILURE_THRESHOLD + 1): + await rate_limits.enforce_user_rate_limit("get_async_task", USER) + assert any("outcome:breaker_open" in tags for tags in emitted[-1:]) diff --git a/model-engine/tests/unit/api/test_tasks.py b/model-engine/tests/unit/api/test_tasks.py index ebd732a8a..30f25af4a 100644 --- a/model-engine/tests/unit/api/test_tasks.py +++ b/model-engine/tests/unit/api/test_tasks.py @@ -1,8 +1,16 @@ +import asyncio +import threading from typing import Any, Dict, Tuple from unittest.mock import AsyncMock, MagicMock, patch +import anyio import pytest -from model_engine_server.common.dtos.tasks import EndpointPredictV1Request +from model_engine_server.api import tasks_v1 +from model_engine_server.common.dtos.tasks import ( + EndpointPredictV1Request, + GetAsyncTaskV1Response, + TaskStatus, +) from model_engine_server.domain.entities import ModelBundle, ModelEndpoint from model_engine_server.domain.exceptions import ( InvalidRequestException, @@ -12,6 +20,15 @@ ) +@pytest.fixture +def fresh_task_limiter_cache(): + # The memoized limiter binds to the event loop that first creates it; clear it so + # this test (and later tests on other loops) each get a limiter on their own loop. + tasks_v1._get_task_limiter.cache_clear() + yield + tasks_v1._get_task_limiter.cache_clear() + + def test_create_async_task_success( model_bundle_1_v1: Tuple[ModelBundle, Any], model_endpoint_1: Tuple[ModelEndpoint, Any], @@ -242,6 +259,72 @@ def test_get_async_task_raises_404_object_not_authorized( assert response.status_code == 404 +@pytest.mark.asyncio +async def test_task_limiter_lazy_and_memoized(fresh_task_limiter_cache): + limiter = tasks_v1._get_task_limiter() + assert tasks_v1._get_task_limiter() is limiter + assert limiter.total_tokens == 40 + + +@pytest.mark.asyncio +async def test_task_polls_saturate_only_dedicated_limiter( + fresh_task_limiter_cache, + model_bundle_1_v1: Tuple[ModelBundle, Any], + model_endpoint_1: Tuple[ModelEndpoint, Any], + test_api_key: str, + get_async_test_client_wrapper, +): + assert model_endpoint_1[0].infra_state is not None + release = threading.Event() + + def blocking_execute(user, task_id): + release.wait(timeout=10) + return GetAsyncTaskV1Response(task_id=task_id, status=TaskStatus.PENDING) + + mock_use_case = MagicMock() + mock_use_case.return_value.execute = blocking_execute + client = get_async_test_client_wrapper( + fake_docker_repository_image_always_exists=True, + fake_model_bundle_repository_contents={ + model_bundle_1_v1[0].id: model_bundle_1_v1[0], + }, + fake_model_endpoint_record_repository_contents={ + model_endpoint_1[0].record.id: model_endpoint_1[0].record, + }, + fake_model_endpoint_infra_gateway_contents={ + model_endpoint_1[0].infra_state.deployment_name: model_endpoint_1[0].infra_state, + }, + fake_batch_job_record_repository_contents={}, + fake_batch_job_progress_gateway_contents={}, + fake_docker_image_batch_job_bundle_repository_contents={}, + ) + try: + with patch( + "model_engine_server.api.tasks_v1.GetAsyncInferenceTaskV1UseCase", + mock_use_case, + ): + polls = [ + asyncio.create_task( + client.get("/v1/async-tasks/test_task_id", auth=(test_api_key, "")) + ) + for _ in range(45) + ] + limiter = tasks_v1._get_task_limiter() + deadline = asyncio.get_event_loop().time() + 10 + while limiter.borrowed_tokens < 40 and asyncio.get_event_loop().time() < deadline: + await asyncio.sleep(0.01) + # The dedicated limiter is saturated at exactly its capacity; the 5 extra polls queue. + assert limiter.borrowed_tokens == 40 + # The default anyio threadpool must remain serviceable while polls are blocked. + assert await asyncio.wait_for(anyio.to_thread.run_sync(lambda: "ok"), 2) == "ok" + release.set() + responses = await asyncio.gather(*polls) + assert [response.status_code for response in responses] == [200] * 45 + finally: + release.set() + await client.aclose() + + def test_create_sync_task_success( model_bundle_1_v1: Tuple[ModelBundle, Any], model_endpoint_2: Tuple[ModelEndpoint, Any], diff --git a/model-engine/tests/unit/inference/test_celery_forwarder_worker.py b/model-engine/tests/unit/inference/test_celery_forwarder_worker.py index 5c15e00fd..4b1ad61bf 100644 --- a/model-engine/tests/unit/inference/test_celery_forwarder_worker.py +++ b/model-engine/tests/unit/inference/test_celery_forwarder_worker.py @@ -6,11 +6,38 @@ (which has no per-child recycling). app.Worker is mocked so .start() does not run a real worker. """ +from types import SimpleNamespace from unittest.mock import MagicMock +import pytest +from celery import Celery +from model_engine_server.common.constants import DEFAULT_CELERY_TASK_NAME, LIRA_CELERY_TASK_NAME from model_engine_server.inference.forwarding import celery_forwarder +@pytest.mark.parametrize("task_name", [LIRA_CELERY_TASK_NAME, DEFAULT_CELERY_TASK_NAME]) +def test_forwarder_tasks_use_configured_wall_clock_limit(monkeypatch, task_name): + app = Celery("test-forwarder", broker="memory://", backend="cache+memory://") + monkeypatch.setattr(celery_forwarder, "celery_app", lambda **_kwargs: app) + monkeypatch.setattr( + celery_forwarder, + "infra_config", + lambda: SimpleNamespace(s3_bucket="test", profile_ml_inference_worker=None), + ) + monkeypatch.setattr(celery_forwarder, "DatadogInferenceMonitoringMetricsGateway", MagicMock) + forwarder = MagicMock(timeout_seconds=123.5, post_inference_hooks_handler=None) + + celery_forwarder.create_celery_service( + forwarder=forwarder, + task_visibility=celery_forwarder.TaskVisibility.VISIBILITY_24H, + broker_type="redis", + backend_protocol="redis", + queue_name="test-queue", + ) + + assert app.tasks[task_name].time_limit == 123.5 + + def _worker_kwargs(monkeypatch, pool, env_value): monkeypatch.setattr(celery_forwarder, "CELERY_WORKER_POOL", pool) if env_value is None: diff --git a/model-engine/tests/unit/inference/test_forwarding.py b/model-engine/tests/unit/inference/test_forwarding.py index 3e0141c83..3938814d9 100644 --- a/model-engine/tests/unit/inference/test_forwarding.py +++ b/model-engine/tests/unit/inference/test_forwarding.py @@ -4,12 +4,14 @@ from unittest import mock from unittest.mock import AsyncMock, MagicMock +import aiohttp import pytest from fastapi import HTTPException from fastapi.responses import JSONResponse from model_engine_server.core.utils.env import environment from model_engine_server.domain.entities import ModelEndpointConfig from model_engine_server.inference.forwarding.forwarding import ( + DEFAULT_SYNC_TIMEOUT_SECONDS, ENV_SERIALIZE_RESULTS_AS_STRING, KEY_SERIALIZE_RESULTS_AS_STRING, Forwarder, @@ -261,6 +263,74 @@ def test_forwarders(post_inference_hooks_handler): _check(json_response) +@mock.patch("requests.post") +def test_sync_forwarder_uses_configured_timeout(mock_post, post_inference_hooks_handler): + mock_post.return_value.status_code = 200 + mock_post.return_value.json.return_value = PAYLOAD + fwd = Forwarder( + "http://user-service/predict", + model_engine_unwrap=True, + serialize_results_as_string=False, + post_inference_hooks_handler=post_inference_hooks_handler, + wrap_response=False, + forward_http_status=False, + forward_http_status_in_body=False, + timeout_seconds=123, + ) + + fwd({"ignore": "me"}) + + mock_post.assert_called_once_with( + "http://user-service/predict", + json={"ignore": "me"}, + headers={"Content-Type": "application/json"}, + timeout=123, + ) + + +@pytest.mark.asyncio +async def test_async_forwarder_uses_configured_timeout(post_inference_hooks_handler): + post_calls = [] + + class FakeAiohttpResponse: + status = 200 + + async def json(self, content_type=None): + return PAYLOAD + + class FakeAiohttpSession: + def __init__(self, *args, **kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def post(self, url, **kwargs): + post_calls.append((url, kwargs)) + return FakeAiohttpResponse() + + fwd = Forwarder( + "http://user-service/predict", + model_engine_unwrap=True, + serialize_results_as_string=False, + post_inference_hooks_handler=post_inference_hooks_handler, + wrap_response=False, + forward_http_status=False, + forward_http_status_in_body=False, + timeout_seconds=123, + ) + with mock.patch("aiohttp.ClientSession", FakeAiohttpSession): + response = await fwd.forward({"ignore": "me"}) + + assert response == PAYLOAD + ((url, kwargs),) = post_calls + assert url == "http://user-service/predict" + assert kwargs["timeout"] == aiohttp.ClientTimeout(total=123) + + def _check(json_response) -> None: json_response = ( json.loads(json_response.body.decode("utf-8")) @@ -480,6 +550,40 @@ def test_forwarder_loader(): _check_responses_not_wrapped(json_response) +@mock.patch("requests.post", mocked_post) +@mock.patch("requests.get", mocked_get) +@mock.patch( + "model_engine_server.inference.forwarding.forwarding.get_endpoint_config", + mocked_get_endpoint_config, +) +@pytest.mark.parametrize( + "loader_kwargs, expected_timeout", + [ + pytest.param({}, DEFAULT_SYNC_TIMEOUT_SECONDS, id="default"), + pytest.param({"timeout_seconds": 123.0}, 123.0, id="override"), + ], +) +def test_forwarder_loader_timeout(loader_kwargs, expected_timeout): + fwd = LoadForwarder(**loader_kwargs).load(None, None) # type: ignore + assert fwd.timeout_seconds == expected_timeout + + +@pytest.mark.parametrize( + "invalid_timeout", + [ + pytest.param(0, id="zero"), + pytest.param(-1, id="negative"), + pytest.param(None, id="null"), + pytest.param(float("inf"), id="non-finite"), + pytest.param("60", id="string"), + pytest.param(True, id="bool"), + ], +) +def test_forwarder_loader_invalid_timeout(invalid_timeout): + with pytest.raises(ValueError, match="timeout_seconds"): + LoadForwarder(timeout_seconds=invalid_timeout).load(None, None) # type: ignore + + @mock.patch("requests.post", mocked_post) @mock.patch("requests.get", mocked_get) @mock.patch( diff --git a/model-engine/tests/unit/inference/test_http_forwarder.py b/model-engine/tests/unit/inference/test_http_forwarder.py index bb255fea2..9e046b1a9 100644 --- a/model-engine/tests/unit/inference/test_http_forwarder.py +++ b/model-engine/tests/unit/inference/test_http_forwarder.py @@ -12,7 +12,10 @@ from fastapi.testclient import TestClient from model_engine_server.common.dtos.tasks import EndpointPredictV1Request from model_engine_server.domain.entities.model_endpoint_entity import ModelEndpointConfig -from model_engine_server.inference.forwarding.forwarding import Forwarder +from model_engine_server.inference.forwarding.forwarding import ( + DEFAULT_SYNC_TIMEOUT_SECONDS, + Forwarder, +) from model_engine_server.inference.forwarding.http_forwarder import ( MultiprocessingConcurrencyLimiter, get_concurrency_limiter, @@ -157,6 +160,26 @@ def test_get_forwarder_loader(): assert loader.predict_route == "/v1/chat/completions" +@pytest.mark.parametrize( + "sync_config_overrides, expected_timeout", + [ + pytest.param({}, DEFAULT_SYNC_TIMEOUT_SECONDS, id="default"), + pytest.param({"timeout_seconds": 123.0}, 123.0, id="override"), + ], +) +def test_get_forwarder_loader_plumbs_timeout_from_config(sync_config_overrides, expected_timeout): + def get_config_with_overrides(): + config = mocked_get_config() + config["sync"].update(sync_config_overrides) + return config + + with mock.patch( + "model_engine_server.inference.forwarding.http_forwarder.get_config", + get_config_with_overrides, + ): + assert get_forwarder_loader().timeout_seconds == expected_timeout + + @mock.patch( "model_engine_server.inference.forwarding.http_forwarder.get_config", mocked_get_config,