diff --git a/deploy/helm/vanity-gateway/README.md b/deploy/helm/vanity-gateway/README.md index a6d406c06..e04dd3823 100644 --- a/deploy/helm/vanity-gateway/README.md +++ b/deploy/helm/vanity-gateway/README.md @@ -107,12 +107,42 @@ or the pod is killed mid-drain. - `openai`: per-endpoint model routes, keyed by endpoint (`chatCompletions`, `completions`, `embeddings`, `responses`, and the image endpoints). Each route - requires `modelName` and `functionID`, and supports shadow-traffic fields such - as `shadowModelName`, `shadowPercentage`, and - `shadowCancelOnClientDisconnect`. + requires `modelName` and `functionID`. - `vanity`: host-based routes, each requiring a `host` and a `paths` map. Each path requires `path` and `functionID`. +Use `shadows` to set policy for each shadow target: + +```yaml +primary: + modelName: example/primary + functionID: primary-function-id + shadows: + - modelName: private/example/shadow-a + percentage: 10 + samplingMethod: perBearerKey + cancelOnClientDisconnect: true + - modelName: private/example/shadow-b + percentage: 50 + samplingMethod: random +shadow-a: + modelName: private/example/shadow-a + functionID: shadow-a-function-id +shadow-b: + modelName: private/example/shadow-b + functionID: shadow-b-function-id +``` + +Each shadow model must be another model route in the same OpenAI endpoint. +`percentage` defaults to `100`, `samplingMethod` defaults to `random`, and +`cancelOnClientDisconnect` defaults to `false`. +Shadowing is not supported for `imageEdits` or `imageVariations` routes. + +Legacy `shadowModelName`, `shadowModelNames`, `shadowPercentage`, +`shadowSamplingMethod`, and `shadowCancelOnClientDisconnect` fields remain +supported. Their policy applies to every legacy shadow target. Do not combine +the `shadows` field with legacy shadow fields on the same route. + Both sections are empty by default. `vanityGateway.config.shadowMaxConcurrent` bounds concurrent shadow requests across all routes. diff --git a/deploy/helm/vanity-gateway/helm-nvcf-vanity-gateway/values.schema.json b/deploy/helm/vanity-gateway/helm-nvcf-vanity-gateway/values.schema.json index 90340b28c..4b5d1395e 100644 --- a/deploy/helm/vanity-gateway/helm-nvcf-vanity-gateway/values.schema.json +++ b/deploy/helm/vanity-gateway/helm-nvcf-vanity-gateway/values.schema.json @@ -219,10 +219,10 @@ "$ref": "#/definitions/openAIEndpointRoutes" }, "imageEdits": { - "$ref": "#/definitions/openAIEndpointRoutes" + "$ref": "#/definitions/openAIMultipartEndpointRoutes" }, "imageVariations": { - "$ref": "#/definitions/openAIEndpointRoutes" + "$ref": "#/definitions/openAIMultipartEndpointRoutes" } }, "additionalProperties": { @@ -235,6 +235,12 @@ "$ref": "#/definitions/openAIModelRoute" } }, + "openAIMultipartEndpointRoutes": { + "type": "object", + "additionalProperties": { + "$ref": "#/definitions/openAIMultipartModelRoute" + } + }, "openAIModelRoute": { "type": "object", "properties": { @@ -271,11 +277,20 @@ "$ref": "#/definitions/nonEmptyString" } }, + "shadows": { + "type": "array", + "items": { + "$ref": "#/definitions/shadowTarget" + } + }, "shadowPercentage": { "type": "integer", "minimum": 1, "maximum": 100 }, + "shadowSamplingMethod": { + "$ref": "#/definitions/legacyShadowSamplingMethod" + }, "shadowCancelOnClientDisconnect": { "type": "boolean" } @@ -284,8 +299,129 @@ "modelName", "functionID" ], + "allOf": [ + { + "not": { + "allOf": [ + { + "required": [ + "shadows" + ] + }, + { + "anyOf": [ + { + "required": [ + "shadowModelName" + ] + }, + { + "required": [ + "shadowModelNames" + ] + }, + { + "required": [ + "shadowPercentage" + ] + }, + { + "required": [ + "shadowSamplingMethod" + ] + }, + { + "required": [ + "shadowCancelOnClientDisconnect" + ] + } + ] + } + ] + } + } + ], "additionalProperties": true }, + "openAIMultipartModelRoute": { + "allOf": [ + { + "$ref": "#/definitions/openAIModelRoute" + }, + { + "type": "object", + "properties": { + "shadows": { + "maxItems": 0 + }, + "shadowModelName": { + "maxLength": 0 + }, + "shadowModelNames": { + "maxItems": 0 + }, + "shadowSamplingMethod": { + "enum": [ + "", + null + ] + }, + "shadowCancelOnClientDisconnect": { + "enum": [ + false + ] + } + }, + "not": { + "required": [ + "shadowPercentage" + ] + } + } + ] + }, + "shadowSamplingMethod": { + "type": "string", + "enum": [ + "random", + "perBearerKey" + ] + }, + "legacyShadowSamplingMethod": { + "type": [ + "string", + "null" + ], + "enum": [ + "", + "random", + "perBearerKey", + null + ] + }, + "shadowTarget": { + "type": "object", + "properties": { + "modelName": { + "$ref": "#/definitions/nonEmptyString" + }, + "percentage": { + "type": "integer", + "minimum": 1, + "maximum": 100 + }, + "samplingMethod": { + "$ref": "#/definitions/shadowSamplingMethod" + }, + "cancelOnClientDisconnect": { + "type": "boolean" + } + }, + "required": [ + "modelName" + ], + "additionalProperties": false + }, "vanityRoutes": { "type": "object", "additionalProperties": { diff --git a/src/invocation-plane-services/vanity-gateway/README.md b/src/invocation-plane-services/vanity-gateway/README.md index 3b9e53b41..9e655a248 100644 --- a/src/invocation-plane-services/vanity-gateway/README.md +++ b/src/invocation-plane-services/vanity-gateway/README.md @@ -142,10 +142,11 @@ v2config: X-Provider-Feature: enabled X-Request-Source: vanity-gateway tooManyRequestsMessage: "Try again later or use a partner endpoint." - shadowModelNames: - - private/meta/llama-3.1-8b-shadow - shadowPercentage: 10 - shadowSamplingMethod: perBearerKey + shadows: + - modelName: private/meta/llama-3.1-8b-shadow + percentage: 10 + samplingMethod: perBearerKey + cancelOnClientDisconnect: false private_meta_llama-3_1-8b-shadow: modelName: private/meta/llama-3.1-8b-shadow functionID: 00000000-0000-0000-0000-000000000002 @@ -188,11 +189,16 @@ v2config: | `eol` | No | RFC3339 timestamp. Future dates add a `Deprecation` header; past dates return `410 Gone` and hide the model from `/v1/models`. | | `offlineMessage` | No | Non-empty value returns `503 Service Unavailable` with this message. | | `tooManyRequestsMessage` | No | Message appended to upstream `429 Too Many Requests` responses for this model. | -| `shadowModelName` | No | Legacy single shadow target. Prefer `shadowModelNames` for new config. | -| `shadowModelNames` | No | Additional model names in the same OpenAI section that receive shadow traffic. Not supported for multipart image edit or variation endpoints. | -| `shadowPercentage` | No | Percentage of primary requests to shadow, from `1` to `100`. Defaults to `100` when shadow targets exist. | -| `shadowSamplingMethod` | No | Shadow admission method. Allowed values are `random` and `perBearerKey`. Missing or empty defaults to `random`. Requires at least one shadow target. Not supported for multipart image edit or variation endpoints. | -| `shadowCancelOnClientDisconnect` | No | When `true`, cancels shadow work if the primary request context is canceled. Requires at least one shadow target. | +| `shadows` | No | List of per-target shadow configs. It cannot be combined with any legacy top-level shadow field. Not supported for multipart image edit or variation endpoints. | +| `shadows[].modelName` | Yes | Target model name in the same OpenAI section. It cannot match the primary model or another shadow target. | +| `shadows[].percentage` | No | Percentage of primary requests sent to this target, from `1` to `100`. Defaults to `100`. | +| `shadows[].samplingMethod` | No | Admission method for this target. Allowed values are `random` and `perBearerKey`. Defaults to `random`. | +| `shadows[].cancelOnClientDisconnect` | No | When `true`, cancels this target if the client disconnects or cancels the request, or if the primary proxy encounters a transport or response-write failure or panics before normal completion. Defaults to `false`. | +| `shadowModelName` | No | Legacy single shadow target. Prefer `shadows` for new config. | +| `shadowModelNames` | No | Legacy list of additional shadow targets. Targets must be in the same OpenAI section. | +| `shadowPercentage` | No | Legacy percentage applied to every legacy target, from `1` to `100`. Defaults to `100`. | +| `shadowSamplingMethod` | No | Legacy admission method applied to every legacy target. Allowed values are `random` and `perBearerKey`. Missing, empty, or `null` defaults to `random`. | +| `shadowCancelOnClientDisconnect` | No | Legacy cancellation policy applied to every legacy target. It uses the same disconnect, client cancellation, transport failure, response-write failure, and panic behavior. Defaults to `false`. | ### Shadow Traffic Support @@ -203,9 +209,21 @@ section as the primary model. The gateway rewrites the request `model` field to the shadow target and marks the replay with `NVCF-Shadow: true` so shadow requests do not recursively shadow. -`shadowSamplingMethod` controls how `shadowPercentage` admits requests. The -default method, `random`, draws a request-local bucket from `0` to `99` and -admits the request when `bucket < shadowPercentage`. +Each `shadows` entry has its own `percentage`, `samplingMethod`, and +`cancelOnClientDisconnect` policy. The gateway evaluates each target against +its policy. This allows one primary model to shadow different request +percentages to different targets. + +Legacy top-level shadow fields remain supported. The gateway combines +`shadowModelName` and `shadowModelNames` into one target list and applies the +top-level `shadowPercentage`, `shadowSamplingMethod`, and +`shadowCancelOnClientDisconnect` policy to every target. A model entry cannot +combine `shadows` with any legacy top-level shadow field. Config validation +rejects the mixed form. + +The default sampling method, `random`, draws one request-local bucket from `0` +to `99` for each primary request. Every `random` shadow on that request uses the +same bucket and admits the request when `bucket < percentage`. The `perBearerKey` method makes admission sticky by bearer credential. It requires exactly one `Authorization` header whose value starts with the @@ -214,22 +232,28 @@ only the `Bearer` scheme and following separator whitespace, then hashes the complete remaining credential as opaque UTF-8 bytes. It computes SHA-256, reads the first 8 digest bytes as a big-endian `uint64`, sets `bucket = value % 100`, and admits the request when -`bucket < shadowPercentage`. Bearer credential prefixes such as `nvapi`, +`bucket < percentage`. Every `perBearerKey` shadow on the request uses the same +credential bucket. Bearer credential prefixes such as `nvapi`, `nvapi-stg`, and `nvapi-nvcf` remain part of the credential and are not stripped. Missing, malformed, duplicate, or non-Bearer authorization skips -shadow dispatch when `shadowPercentage` is below `100`. +each `perBearerKey` target whose `percentage` is below `100`. It does not affect +targets that use `random`. `perBearerKey` is key-level sampling. It is not true user or session sampling and can skew shadow volume when a few bearer keys dominate traffic. Shadow traffic is not supported for multipart image endpoints: `imageEdits` and -`imageVariations`. Config validation rejects shadow fields in those sections. +`imageVariations`. Config validation rejects effective shadow settings in those +sections. Legacy zero-value settings and an empty `shadows` list remain accepted +as no-ops. Admitted shadow requests are bounded by `SHADOW_MAX_CONCURRENT` and the gateway shadow timeout. Normal primary response completion does not cancel shadow work. -When `shadowCancelOnClientDisconnect` is `true`, the gateway cancels shadow work -only if the client disconnects or cancels the primary request before the primary -response completes. +When a target's `cancelOnClientDisconnect` is `true`, the gateway cancels only +that target if the client disconnects or cancels the request, or if the primary +proxy encounters a transport or response-write failure or panics before normal +completion. Other targets continue under their own policies. An HTTP error +status alone does not trigger cancellation. ### Vanity Mapping Fields diff --git a/src/invocation-plane-services/vanity-gateway/config.yaml b/src/invocation-plane-services/vanity-gateway/config.yaml index bfbf09851..c681009ca 100644 --- a/src/invocation-plane-services/vanity-gateway/config.yaml +++ b/src/invocation-plane-services/vanity-gateway/config.yaml @@ -12,8 +12,9 @@ v2config: functionID: eb69df99-3fc9-4272-a1ce-1be218deef3a functionVersionID: eb69df99-3fc9-4272-a1ce-1be218deef3a tooManyRequestsMessage: "Check out this model at a partner!" - shadowModelName: private/facebook/opt-125m-shadow - shadowPercentage: 100 + shadows: + - modelName: private/facebook/opt-125m-shadow + percentage: 100 private/facebook/opt-125m-shadow: modelName: private/facebook/opt-125m-shadow functionID: shadow-func-0001 @@ -22,9 +23,10 @@ v2config: modelName: private/test/cancel-on-disconnect functionID: cancel-primary-func functionVersionID: cancel-primary-ver - shadowModelName: private/test/cancel-on-disconnect-shadow - shadowPercentage: 100 - shadowCancelOnClientDisconnect: true + shadows: + - modelName: private/test/cancel-on-disconnect-shadow + percentage: 100 + cancelOnClientDisconnect: true private/test/cancel-on-disconnect-shadow: modelName: private/test/cancel-on-disconnect-shadow functionID: cancel-shadow-func diff --git a/src/invocation-plane-services/vanity-gateway/gateway/openai_director.go b/src/invocation-plane-services/vanity-gateway/gateway/openai_director.go index 9d0106aed..000e22847 100644 --- a/src/invocation-plane-services/vanity-gateway/gateway/openai_director.go +++ b/src/invocation-plane-services/vanity-gateway/gateway/openai_director.go @@ -90,19 +90,23 @@ type OpenAIError struct { } type FunctionInfo struct { - functionId string - functionVersionId string - pathOverride *string - usePexec bool - sessionTimeout config.SessionTimeoutSeconds - customHeaders config.CustomHeaders - eol time.Time - offlineMessage string - tooManyRequestsMessage string - shadowModelNames []string - shadowPercentage int - shadowSamplingMethod config.ShadowSamplingMethod - shadowCancelOnClientDisconnect bool + functionId string + functionVersionId string + pathOverride *string + usePexec bool + sessionTimeout config.SessionTimeoutSeconds + customHeaders config.CustomHeaders + eol time.Time + offlineMessage string + tooManyRequestsMessage string + shadows []shadowConfig +} + +type shadowConfig struct { + modelName string + percentage int + samplingMethod config.ShadowSamplingMethod + cancelOnClientDisconnect bool } type ModelMapping struct { @@ -123,19 +127,16 @@ type ModelListResponse struct { } type ModelNameToFunctionIdVersionId struct { - FunctionId string - FunctionVersionId string - OutgoingPathOverride string - UsePexec bool - SessionTimeout config.SessionTimeoutSeconds - CustomHeaders config.CustomHeaders - EOL time.Time - OfflineMessage string - TooManyRequestsMessage string - ShadowModelNames []string - ShadowPercentage *int - ShadowSamplingMethod config.ShadowSamplingMethod - ShadowCancelOnClientDisconnect bool + FunctionId string + FunctionVersionId string + OutgoingPathOverride string + UsePexec bool + SessionTimeout config.SessionTimeoutSeconds + CustomHeaders config.CustomHeaders + EOL time.Time + OfflineMessage string + TooManyRequestsMessage string + Shadows []shadowConfig } type openAIRequestBody struct { @@ -236,21 +237,18 @@ func buildModelMapping( if entry.OutgoingPathOverride != "" { pathOverride = &entry.OutgoingPathOverride } - initializeShadowDropMetrics(entry.ShadowModelNames) + initializeShadowDropMetrics(shadowModelNames(entry.Shadows)) modelNameToNVCFUrl[modelName] = FunctionInfo{ - functionId: entry.FunctionId, - functionVersionId: entry.FunctionVersionId, - pathOverride: pathOverride, - usePexec: entry.UsePexec, - sessionTimeout: entry.SessionTimeout, - customHeaders: entry.CustomHeaders, - eol: entry.EOL, - offlineMessage: entry.OfflineMessage, - tooManyRequestsMessage: entry.TooManyRequestsMessage, - shadowModelNames: entry.ShadowModelNames, - shadowPercentage: defaultShadowPercentage(entry.ShadowPercentage), - shadowSamplingMethod: defaultShadowSamplingMethod(entry.ShadowSamplingMethod), - shadowCancelOnClientDisconnect: entry.ShadowCancelOnClientDisconnect, + functionId: entry.FunctionId, + functionVersionId: entry.FunctionVersionId, + pathOverride: pathOverride, + usePexec: entry.UsePexec, + sessionTimeout: entry.SessionTimeout, + customHeaders: entry.CustomHeaders, + eol: entry.EOL, + offlineMessage: entry.OfflineMessage, + tooManyRequestsMessage: entry.TooManyRequestsMessage, + shadows: entry.Shadows, } // build the modelInfo list and modelName to modelInfo map @@ -400,19 +398,16 @@ func convertIntoModelNameToFunctionIdAndVersionIdMappingV2(mapping map[string]co modelNameToFunctionIdVersionId := make(map[string]ModelNameToFunctionIdVersionId) for _, entry := range mapping { modelNameToFunctionIdVersionId[entry.ModelName] = ModelNameToFunctionIdVersionId{ - FunctionId: entry.FunctionID, - FunctionVersionId: entry.FunctionVersionID, - OutgoingPathOverride: entry.OutgoingPathOverride, - UsePexec: entry.UsePexec, - SessionTimeout: entry.SessionTimeout, - CustomHeaders: entry.CustomHeaders, - EOL: entry.EOL, - OfflineMessage: entry.OfflineMessage, - TooManyRequestsMessage: entry.TooManyRequestsMessage, - ShadowModelNames: effectiveShadowModelNames(entry.ShadowModelName, entry.ShadowModelNames), - ShadowPercentage: entry.ShadowPercentage, - ShadowSamplingMethod: entry.ShadowSamplingMethod, - ShadowCancelOnClientDisconnect: entry.ShadowCancelOnClientDisconnect, + FunctionId: entry.FunctionID, + FunctionVersionId: entry.FunctionVersionID, + OutgoingPathOverride: entry.OutgoingPathOverride, + UsePexec: entry.UsePexec, + SessionTimeout: entry.SessionTimeout, + CustomHeaders: entry.CustomHeaders, + EOL: entry.EOL, + OfflineMessage: entry.OfflineMessage, + TooManyRequestsMessage: entry.TooManyRequestsMessage, + Shadows: normalizeShadowConfigs(entry.EffectiveShadows()), } } @@ -433,15 +428,31 @@ func defaultShadowSamplingMethod(shadowSamplingMethod config.ShadowSamplingMetho return shadowSamplingMethod } -func effectiveShadowModelNames(legacyModelName string, modelNames []string) []string { - if legacyModelName == "" && len(modelNames) == 0 { +func normalizeShadowConfigs(shadows []config.ShadowConfig) []shadowConfig { + if len(shadows) == 0 { + return nil + } + result := make([]shadowConfig, 0, len(shadows)) + for _, shadow := range shadows { + result = append(result, shadowConfig{ + modelName: shadow.ModelName, + percentage: defaultShadowPercentage(shadow.Percentage), + samplingMethod: defaultShadowSamplingMethod(shadow.SamplingMethod), + cancelOnClientDisconnect: shadow.CancelOnClientDisconnect, + }) + } + return result +} + +func shadowModelNames(shadows []shadowConfig) []string { + if len(shadows) == 0 { return nil } - shadowModelNames := make([]string, 0, len(modelNames)+1) - if legacyModelName != "" { - shadowModelNames = append(shadowModelNames, legacyModelName) + modelNames := make([]string, 0, len(shadows)) + for _, shadow := range shadows { + modelNames = append(modelNames, shadow.modelName) } - return append(shadowModelNames, modelNames...) + return modelNames } func (d *OpenAIDirector) ServeCompletions(writer http.ResponseWriter, request *http.Request) { @@ -660,12 +671,12 @@ func (d *OpenAIDirector) dispatchShadowIfNeeded(resolved resolvedOpenAIRequest, if d.shadower == nil { return func(error) {} } - if !shouldDispatchShadow(resolved.request, resolved.functionInfo, d.randomShadowBucket) { + shadows := admittedShadows(resolved.request, resolved.functionInfo.shadows, d.randomShadowBucket) + if len(shadows) == 0 { return func(error) {} } - shadowModelNames := resolved.functionInfo.shadowModelNames - shadowCtx, finishShadowPrimary := shadowContext(resolved.request, resolved.functionInfo.shadowCancelOnClientDisconnect) + targetModelNames := shadowModelNames(shadows) // Clone body only for shadowed requests — avoids allocation on the hot path. rawBody, err := io.ReadAll(resolved.request.Body) @@ -673,13 +684,13 @@ func (d *OpenAIDirector) dispatchShadowIfNeeded(resolved resolvedOpenAIRequest, if err != nil { recordShadowDispatchSummary( resolved.request.Context(), - shadowModelNames, + targetModelNames, 0, - len(shadowModelNames), - repeatedStrings(shadowDroppedReasonBodyReadError, len(shadowModelNames)), - shadowModelNames, + len(targetModelNames), + repeatedStrings(shadowDroppedReasonBodyReadError, len(targetModelNames)), + targetModelNames, ) - return finishShadowPrimary + return func(error) {} } // Reset primary request body (pool buffer released above). resolved.request.Body = io.NopCloser(bytes.NewReader(rawBody)) @@ -696,24 +707,21 @@ func (d *OpenAIDirector) dispatchShadowIfNeeded(resolved resolvedOpenAIRequest, droppedCount := 0 var droppedReasons []string var droppedTargetModels []string - for _, shadowModelName := range shadowModelNames { - // Rewrite model field in the shadow body. - shadowBody, err := rewriteShadowRequestModel(rawBody, shadowModelName) + finishers := make([]func(error), 0, len(shadows)) + for _, shadow := range shadows { + shadowBody, err := rewriteShadowRequestModel(rawBody, shadow.modelName) if err != nil { - recordShadowDispatchSummary( - resolved.request.Context(), - shadowModelNames, - 0, - len(shadowModelNames), - repeatedStrings(shadowDroppedReasonBodyRewriteError, len(shadowModelNames)), - shadowModelNames, - ) - return finishShadowPrimary + droppedCount++ + droppedReasons = append(droppedReasons, shadowDroppedReasonBodyRewriteError) + droppedTargetModels = append(droppedTargetModels, shadow.modelName) + continue } // Build shadow request with the rewritten body and recursion guard. + shadowCtx, finishShadow := shadowContext(resolved.request, shadow.cancelOnClientDisconnect) + finishers = append(finishers, finishShadow) shadowReq := newShadowRequest(resolved.request, shadowBody, shadowCtx) - handler := newShadowReplayHandler(shadowModelName, replayHandler) + handler := newShadowReplayHandler(shadow.modelName, replayHandler) // Shadow admission errors are logged by TrafficShadower and summarized below. // They must not affect the primary request. @@ -722,14 +730,14 @@ func (d *OpenAIDirector) dispatchShadowIfNeeded(resolved resolvedOpenAIRequest, droppedCount++ if reason := shadowDroppedReason(err); reason != "" { droppedReasons = append(droppedReasons, reason) - droppedTargetModels = append(droppedTargetModels, shadowModelName) + droppedTargetModels = append(droppedTargetModels, shadow.modelName) } continue } dispatchedCount++ } - recordShadowDispatchSummary(resolved.request.Context(), shadowModelNames, dispatchedCount, droppedCount, droppedReasons, droppedTargetModels) - return finishShadowPrimary + recordShadowDispatchSummary(resolved.request.Context(), targetModelNames, dispatchedCount, droppedCount, droppedReasons, droppedTargetModels) + return finishShadows(finishers) } func (d *OpenAIDirector) randomShadowBucket() int { @@ -739,18 +747,50 @@ func (d *OpenAIDirector) randomShadowBucket() int { return rand.IntN(100) } -func shouldDispatchShadow(req *http.Request, info FunctionInfo, randomBucket func() int) bool { - if len(info.shadowModelNames) == 0 { - return false - } - if info.shadowPercentage >= 100 { - return true +func admittedShadows(req *http.Request, shadows []shadowConfig, randomBucket func() int) []shadowConfig { + var admitted []shadowConfig + var requestRandomBucket int + randomBucketSet := false + var credentialBucket int + credentialBucketSet := false + credentialValid := false + + for _, shadow := range shadows { + if shadow.percentage >= 100 { + admitted = append(admitted, shadow) + continue + } + if shadow.samplingMethod == config.ShadowSamplingMethodPerBearerKey { + if !credentialBucketSet { + credential, ok := bearerCredential(req) + credentialValid = ok + if ok { + credentialBucket = shadowBucketForBearerCredential(credential) + } + credentialBucketSet = true + } + if credentialValid && credentialBucket < shadow.percentage { + admitted = append(admitted, shadow) + } + continue + } + if !randomBucketSet { + requestRandomBucket = randomBucket() + randomBucketSet = true + } + if requestRandomBucket < shadow.percentage { + admitted = append(admitted, shadow) + } } - if info.shadowSamplingMethod == config.ShadowSamplingMethodPerBearerKey { - credential, ok := bearerCredential(req) - return ok && shadowBucketForBearerCredential(credential) < info.shadowPercentage + return admitted +} + +func finishShadows(finishers []func(error)) func(error) { + return func(proxyErr error) { + for _, finish := range finishers { + finish(proxyErr) + } } - return randomBucket() < info.shadowPercentage } func bearerCredential(req *http.Request) ([]byte, bool) { diff --git a/src/invocation-plane-services/vanity-gateway/gateway/openai_director_test.go b/src/invocation-plane-services/vanity-gateway/gateway/openai_director_test.go index 69976d997..54ddae589 100644 --- a/src/invocation-plane-services/vanity-gateway/gateway/openai_director_test.go +++ b/src/invocation-plane-services/vanity-gateway/gateway/openai_director_test.go @@ -45,6 +45,23 @@ func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } +func testShadowConfigs(modelNames []string, percentage int) []shadowConfig { + return testShadowConfigsWithPolicy(modelNames, percentage, config.ShadowSamplingMethodRandom, false) +} + +func testShadowConfigsWithPolicy(modelNames []string, percentage int, samplingMethod config.ShadowSamplingMethod, cancelOnClientDisconnect bool) []shadowConfig { + shadows := make([]shadowConfig, 0, len(modelNames)) + for _, modelName := range modelNames { + shadows = append(shadows, shadowConfig{ + modelName: modelName, + percentage: percentage, + samplingMethod: samplingMethod, + cancelOnClientDisconnect: cancelOnClientDisconnect, + }) + } + return shadows +} + type failingResponseWriter struct { header http.Header } @@ -302,8 +319,7 @@ func TestBuildModelMapping(t *testing.T) { OutgoingPathOverride: "/custom/path", UsePexec: true, TooManyRequestsMessage: "Try a partner API!", - ShadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - ShadowPercentage: &shadowPct, + Shadows: testShadowConfigs([]string{"private/facebook/opt-125m-shadow"}, shadowPct), }, "private/facebook/opt-125m-shadow": { FunctionId: "shadow-func", @@ -319,8 +335,7 @@ func TestBuildModelMapping(t *testing.T) { assert.Equal(t, "/custom/path", *functionInfo.pathOverride) assert.True(t, functionInfo.usePexec) assert.Equal(t, "Try a partner API!", functionInfo.tooManyRequestsMessage) - assert.Equal(t, []string{"private/facebook/opt-125m-shadow"}, functionInfo.shadowModelNames) - assert.Equal(t, shadowPct, functionInfo.shadowPercentage) + assert.Equal(t, testShadowConfigs([]string{"private/facebook/opt-125m-shadow"}, shadowPct), functionInfo.shadows) _, shadowIsPublic := mapping.modelNameToModelInfo["private/facebook/opt-125m-shadow"] assert.False(t, shadowIsPublic) @@ -331,13 +346,17 @@ func TestBuildModelMappingPreservesAndDefaultsShadowSamplingMethod(t *testing.T) mapping, err := buildModelMapping(map[string]ModelNameToFunctionIdVersionId{ "facebook/opt-125m": { - FunctionId: "func-123", - ShadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - ShadowSamplingMethod: config.ShadowSamplingMethodPerBearerKey, + FunctionId: "func-123", + Shadows: testShadowConfigsWithPolicy( + []string{"private/facebook/opt-125m-shadow"}, + 100, + config.ShadowSamplingMethodPerBearerKey, + false, + ), }, "meta/llama-3.1-8b": { - FunctionId: "func-456", - ShadowModelNames: []string{"private/meta/llama-3.1-8b-shadow"}, + FunctionId: "func-456", + Shadows: testShadowConfigs([]string{"private/meta/llama-3.1-8b-shadow"}, 100), }, "private/facebook/opt-125m-shadow": { FunctionId: "shadow-func", @@ -348,8 +367,8 @@ func TestBuildModelMappingPreservesAndDefaultsShadowSamplingMethod(t *testing.T) }, privateModelMatcher) require.NoError(t, err) - assert.Equal(t, config.ShadowSamplingMethodPerBearerKey, mapping.modelNameToNVCFUrl["facebook/opt-125m"].shadowSamplingMethod) - assert.Equal(t, config.ShadowSamplingMethodRandom, mapping.modelNameToNVCFUrl["meta/llama-3.1-8b"].shadowSamplingMethod) + assert.Equal(t, config.ShadowSamplingMethodPerBearerKey, mapping.modelNameToNVCFUrl["facebook/opt-125m"].shadows[0].samplingMethod) + assert.Equal(t, config.ShadowSamplingMethodRandom, mapping.modelNameToNVCFUrl["meta/llama-3.1-8b"].shadows[0].samplingMethod) } func TestResolveModelMappedRequestAddsMetricAttributes(t *testing.T) { @@ -395,11 +414,11 @@ func TestBuildModelMappingPreservesMultipleShadowTargets(t *testing.T) { mapping, err := buildModelMapping(map[string]ModelNameToFunctionIdVersionId{ "facebook/opt-125m": { FunctionId: "func-123", - ShadowModelNames: []string{ + Shadows: testShadowConfigs([]string{ "private/facebook/opt-125m-shadow-a", "private/facebook/opt-125m-shadow-b", "private/facebook/opt-125m-shadow-c", - }, + }, 100), }, "private/facebook/opt-125m-shadow-a": {FunctionId: "shadow-a-func"}, "private/facebook/opt-125m-shadow-b": {FunctionId: "shadow-b-func"}, @@ -409,11 +428,11 @@ func TestBuildModelMappingPreservesMultipleShadowTargets(t *testing.T) { functionInfo, ok := mapping.modelNameToNVCFUrl["facebook/opt-125m"] require.True(t, ok) - assert.Equal(t, []string{ + assert.Equal(t, testShadowConfigs([]string{ "private/facebook/opt-125m-shadow-a", "private/facebook/opt-125m-shadow-b", "private/facebook/opt-125m-shadow-c", - }, functionInfo.shadowModelNames) + }, 100), functionInfo.shadows) } func TestConvertIntoModelNameToFunctionIdAndVersionIdMappingV2(t *testing.T) { @@ -447,13 +466,50 @@ func TestConvertIntoModelNameToFunctionIdAndVersionIdMappingV2(t *testing.T) { assert.Equal(t, config.SessionTimeoutSeconds(900), expected.SessionTimeout) assert.Equal(t, eolDate, expected.EOL) assert.Equal(t, "Try a partner API!", expected.TooManyRequestsMessage) - assert.Equal(t, []string{ - "private/facebook/opt-125m-shadow", - "private/facebook/opt-125m-shadow-b", - }, expected.ShadowModelNames) - assert.Equal(t, &shadowPct, expected.ShadowPercentage) - assert.Equal(t, config.ShadowSamplingMethodPerBearerKey, expected.ShadowSamplingMethod) - assert.True(t, expected.ShadowCancelOnClientDisconnect) + assert.Equal(t, testShadowConfigsWithPolicy( + []string{ + "private/facebook/opt-125m-shadow", + "private/facebook/opt-125m-shadow-b", + }, + shadowPct, + config.ShadowSamplingMethodPerBearerKey, + true, + ), expected.Shadows) +} + +func TestConvertIntoModelNameToFunctionIdAndVersionIdMappingV2PreservesPerShadowPolicies(t *testing.T) { + percentage := 25 + result := convertIntoModelNameToFunctionIdAndVersionIdMappingV2(map[string]config.ModelFunctionDetails{ + "model-key": { + ModelName: "facebook/opt-125m", + FunctionID: "func-123", + Shadows: []config.ShadowConfig{ + { + ModelName: "private/facebook/opt-125m-shadow-a", + Percentage: &percentage, + SamplingMethod: config.ShadowSamplingMethodPerBearerKey, + CancelOnClientDisconnect: true, + }, + {ModelName: "private/facebook/opt-125m-shadow-b"}, + }, + }, + }) + + expected, ok := result["facebook/opt-125m"] + require.True(t, ok) + assert.Equal(t, []shadowConfig{ + { + modelName: "private/facebook/opt-125m-shadow-a", + percentage: percentage, + samplingMethod: config.ShadowSamplingMethodPerBearerKey, + cancelOnClientDisconnect: true, + }, + { + modelName: "private/facebook/opt-125m-shadow-b", + percentage: 100, + samplingMethod: config.ShadowSamplingMethodRandom, + }, + }, expected.Shadows) } func TestDefaultShadowPercentage(t *testing.T) { @@ -470,9 +526,8 @@ func TestResolveModelMappedRequestPreservesShadowConfig(t *testing.T) { modelMapping := map[string]FunctionInfo{ "facebook/opt-125m": { - functionId: "primary-func", - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 100, + functionId: "primary-func", + shadows: testShadowConfigs([]string{"private/facebook/opt-125m-shadow"}, 100), }, "private/facebook/opt-125m-shadow": { functionId: "shadow-func", @@ -481,8 +536,7 @@ func TestResolveModelMappedRequestPreservesShadowConfig(t *testing.T) { resolved, handled := director.resolveModelMappedRequest(writer, req, modelMapping) require.False(t, handled) - assert.Equal(t, []string{"private/facebook/opt-125m-shadow"}, resolved.functionInfo.shadowModelNames) - assert.Equal(t, 100, resolved.functionInfo.shadowPercentage) + assert.Equal(t, testShadowConfigs([]string{"private/facebook/opt-125m-shadow"}, 100), resolved.functionInfo.shadows) body, err := io.ReadAll(resolved.request.Body) require.NoError(t, err) @@ -506,9 +560,8 @@ func TestDispatchShadowIfNeededReplaysHandlerAndRewritesBody(t *testing.T) { modelMapping := map[string]FunctionInfo{ "facebook/opt-125m": { - functionId: "primary-func", - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 100, + functionId: "primary-func", + shadows: testShadowConfigs([]string{"private/facebook/opt-125m-shadow"}, 100), }, "private/facebook/opt-125m-shadow": { functionId: "shadow-func", @@ -526,8 +579,7 @@ func TestDispatchShadowIfNeededReplaysHandlerAndRewritesBody(t *testing.T) { resolved := resolvedOpenAIRequest{ request: req, functionInfo: FunctionInfo{ - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 100, + shadows: testShadowConfigs([]string{"private/facebook/opt-125m-shadow"}, 100), }, } director.dispatchShadowIfNeeded(resolved, modelMapping) @@ -557,10 +609,13 @@ func TestDispatchShadowIfNeededPerBearerKeyUsesBearerBucket(t *testing.T) { modelMapping := map[string]FunctionInfo{ "facebook/opt-125m": { - functionId: "primary-func", - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 47, - shadowSamplingMethod: config.ShadowSamplingMethodPerBearerKey, + functionId: "primary-func", + shadows: testShadowConfigsWithPolicy( + []string{"private/facebook/opt-125m-shadow"}, + 47, + config.ShadowSamplingMethodPerBearerKey, + false, + ), }, "private/facebook/opt-125m-shadow": { functionId: "shadow-func", @@ -589,6 +644,117 @@ func TestDispatchShadowIfNeededPerBearerKeyUsesBearerBucket(t *testing.T) { assert.Equal(t, true, shadowBody["stream"]) } +func TestDispatchShadowIfNeededAppliesCancellationPerShadow(t *testing.T) { + attachedStarted := make(chan struct{}) + attachedCanceled := make(chan struct{}) + detachedStarted := make(chan struct{}) + detachedCanceled := make(chan struct{}) + detachedDone := make(chan struct{}) + releaseDetached := make(chan struct{}, 1) + release := func() { + select { + case releaseDetached <- struct{}{}: + default: + } + } + t.Cleanup(release) + + transport := roundTripFunc(func(req *http.Request) (*http.Response, error) { + switch req.Header.Get("function-id") { + case "attached-func": + close(attachedStarted) + <-req.Context().Done() + close(attachedCanceled) + return nil, req.Context().Err() + case "detached-func": + close(detachedStarted) + select { + case <-req.Context().Done(): + close(detachedCanceled) + return nil, req.Context().Err() + case <-releaseDetached: + close(detachedDone) + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Header: make(http.Header), + Body: io.NopCloser(bytes.NewReader(nil)), + Request: req, + }, nil + } + default: + t.Fatalf("unexpected function-id %q", req.Header.Get("function-id")) + return nil, nil + } + }) + + vanity, err := NewVanityDirector("https://nvcf.example.test", transport) + require.NoError(t, err) + director := &OpenAIDirector{ + shadower: NewTrafficShadower(2, 30*time.Second), + vanityDirector: vanity, + } + modelMapping := map[string]FunctionInfo{ + "attached-shadow": {functionId: "attached-func"}, + "detached-shadow": {functionId: "detached-func"}, + } + ctx, cancel := context.WithCancel(context.Background()) + request := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + bytes.NewBufferString(`{"model":"primary"}`), + ).WithContext(ctx) + request.Header.Set("Content-Type", "application/json") + + finishPrimary := director.dispatchShadowIfNeeded(resolvedOpenAIRequest{ + request: request, + functionInfo: FunctionInfo{shadows: []shadowConfig{ + { + modelName: "attached-shadow", + percentage: 100, + samplingMethod: config.ShadowSamplingMethodRandom, + cancelOnClientDisconnect: true, + }, + { + modelName: "detached-shadow", + percentage: 100, + samplingMethod: config.ShadowSamplingMethodRandom, + }, + }}, + }, modelMapping) + + for name, started := range map[string]<-chan struct{}{ + "attached": attachedStarted, + "detached": detachedStarted, + } { + select { + case <-started: + case <-time.After(time.Second): + t.Fatalf("%s shadow did not start", name) + } + } + + cancel() + finishPrimary(context.Canceled) + select { + case <-attachedCanceled: + case <-time.After(time.Second): + t.Fatal("attached shadow was not canceled") + } + select { + case <-detachedCanceled: + t.Fatal("detached shadow was canceled") + case <-time.After(100 * time.Millisecond): + } + + release() + select { + case <-detachedDone: + case <-time.After(time.Second): + t.Fatal("detached shadow did not finish") + } +} + func TestShadowCancelledWhenPrimaryProxyErrorsBeforeRequestContextCancels(t *testing.T) { shadowStarted := make(chan struct{}) shadowCanceled := make(chan struct{}) @@ -618,10 +784,13 @@ func TestShadowCancelledWhenPrimaryProxyErrorsBeforeRequestContextCancels(t *tes } modelMapping := map[string]FunctionInfo{ "facebook/opt-125m": { - functionId: "primary-func", - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 100, - shadowCancelOnClientDisconnect: true, + functionId: "primary-func", + shadows: testShadowConfigsWithPolicy( + []string{"private/facebook/opt-125m-shadow"}, + 100, + config.ShadowSamplingMethodRandom, + true, + ), }, "private/facebook/opt-125m-shadow": { functionId: "shadow-func", @@ -703,10 +872,13 @@ func TestShadowCancelledWhenPrimaryResponseWriteFails(t *testing.T) { } modelMapping := map[string]FunctionInfo{ "facebook/opt-125m": { - functionId: "primary-func", - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 100, - shadowCancelOnClientDisconnect: true, + functionId: "primary-func", + shadows: testShadowConfigsWithPolicy( + []string{"private/facebook/opt-125m-shadow"}, + 100, + config.ShadowSamplingMethodRandom, + true, + ), }, "private/facebook/opt-125m-shadow": { functionId: "shadow-func", @@ -761,10 +933,13 @@ func TestShadowCancelledWhenPrimaryProxyPanics(t *testing.T) { } modelMapping := map[string]FunctionInfo{ "facebook/opt-125m": { - functionId: "primary-func", - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 100, - shadowCancelOnClientDisconnect: true, + functionId: "primary-func", + shadows: testShadowConfigsWithPolicy( + []string{"private/facebook/opt-125m-shadow"}, + 100, + config.ShadowSamplingMethodRandom, + true, + ), }, "private/facebook/opt-125m-shadow": { functionId: "shadow-func", @@ -806,8 +981,7 @@ func TestDispatchShadowIfNeededSkipsShadowRequests(t *testing.T) { director.dispatchShadowIfNeeded(resolvedOpenAIRequest{ request: req, functionInfo: FunctionInfo{ - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 100, + shadows: testShadowConfigs([]string{"private/facebook/opt-125m-shadow"}, 100), }, }, map[string]FunctionInfo{ "private/facebook/opt-125m-shadow": { diff --git a/src/invocation-plane-services/vanity-gateway/gateway/shadow_test.go b/src/invocation-plane-services/vanity-gateway/gateway/shadow_test.go index 91ca35b5a..f73ad23ee 100644 --- a/src/invocation-plane-services/vanity-gateway/gateway/shadow_test.go +++ b/src/invocation-plane-services/vanity-gateway/gateway/shadow_test.go @@ -337,22 +337,22 @@ func TestShadowPercentage100(t *testing.T) { "expected all %d requests to be shadowed, got %d", iterations, requestCount.Load()) } -func TestShouldDispatchShadowPerBearerKeyUsesBearerBucket(t *testing.T) { +func TestAdmittedShadowsPerBearerKeyUsesBearerBucket(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{}`)) req.Header.Set("Authorization", "Bearer test-key") - info := FunctionInfo{ - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 47, - shadowSamplingMethod: config.ShadowSamplingMethodPerBearerKey, + shadow := shadowConfig{ + modelName: "private/facebook/opt-125m-shadow", + percentage: 47, + samplingMethod: config.ShadowSamplingMethodPerBearerKey, } - assert.True(t, shouldDispatchShadow(req, info, fixedRandomBucket(99))) + assert.Len(t, admittedShadows(req, []shadowConfig{shadow}, fixedRandomBucket(99)), 1) - info.shadowPercentage = 46 - assert.False(t, shouldDispatchShadow(req, info, fixedRandomBucket(0))) + shadow.percentage = 46 + assert.Empty(t, admittedShadows(req, []shadowConfig{shadow}, fixedRandomBucket(0))) } -func TestShouldDispatchShadowPerBearerKeySkipsMalformedBearerBelow100(t *testing.T) { +func TestAdmittedShadowsPerBearerKeySkipsMalformedBearerBelow100(t *testing.T) { tests := []struct { name string authorization []string @@ -371,32 +371,32 @@ func TestShouldDispatchShadowPerBearerKeySkipsMalformedBearerBelow100(t *testing for _, value := range tc.authorization { req.Header.Add("Authorization", value) } - info := FunctionInfo{ - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 99, - shadowSamplingMethod: config.ShadowSamplingMethodPerBearerKey, + shadow := shadowConfig{ + modelName: "private/facebook/opt-125m-shadow", + percentage: 99, + samplingMethod: config.ShadowSamplingMethodPerBearerKey, } - assert.False(t, shouldDispatchShadow(req, info, fixedRandomBucket(0))) + assert.Empty(t, admittedShadows(req, []shadowConfig{shadow}, fixedRandomBucket(0))) }) } } -func TestShouldDispatchShadowRandomUsesInjectedBucket(t *testing.T) { +func TestAdmittedShadowsRandomUsesInjectedBucket(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{}`)) req.Header.Set("Authorization", "Bearer test-key") - info := FunctionInfo{ - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 40, - shadowSamplingMethod: config.ShadowSamplingMethodRandom, + shadow := shadowConfig{ + modelName: "private/facebook/opt-125m-shadow", + percentage: 40, + samplingMethod: config.ShadowSamplingMethodRandom, } - assert.True(t, shouldDispatchShadow(req, info, fixedRandomBucket(39))) - assert.False(t, shouldDispatchShadow(req, info, fixedRandomBucket(40))) + assert.Len(t, admittedShadows(req, []shadowConfig{shadow}, fixedRandomBucket(39)), 1) + assert.Empty(t, admittedShadows(req, []shadowConfig{shadow}, fixedRandomBucket(40))) } -func TestShouldDispatchShadowPerBearerKeyNormalizesBearerScheme(t *testing.T) { +func TestAdmittedShadowsPerBearerKeyNormalizesBearerScheme(t *testing.T) { tests := []string{ "Bearer test-key", "bearer test-key", @@ -408,17 +408,36 @@ func TestShouldDispatchShadowPerBearerKeyNormalizesBearerScheme(t *testing.T) { t.Run(authorization, func(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{}`)) req.Header.Set("Authorization", authorization) - info := FunctionInfo{ - shadowModelNames: []string{"private/facebook/opt-125m-shadow"}, - shadowPercentage: 47, - shadowSamplingMethod: config.ShadowSamplingMethodPerBearerKey, + shadow := shadowConfig{ + modelName: "private/facebook/opt-125m-shadow", + percentage: 47, + samplingMethod: config.ShadowSamplingMethodPerBearerKey, } - assert.True(t, shouldDispatchShadow(req, info, fixedRandomBucket(99))) + assert.Len(t, admittedShadows(req, []shadowConfig{shadow}, fixedRandomBucket(99)), 1) }) } } +func TestAdmittedShadowsAppliesEachPolicyWithSharedRequestBuckets(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{}`)) + req.Header.Set("Authorization", "Bearer test-key") + + randomCalls := 0 + admitted := admittedShadows(req, []shadowConfig{ + {modelName: "key-in", percentage: 47, samplingMethod: config.ShadowSamplingMethodPerBearerKey}, + {modelName: "key-out", percentage: 46, samplingMethod: config.ShadowSamplingMethodPerBearerKey}, + {modelName: "random-in", percentage: 40, samplingMethod: config.ShadowSamplingMethodRandom}, + {modelName: "random-out", percentage: 39, samplingMethod: config.ShadowSamplingMethodRandom}, + }, func() int { + randomCalls++ + return 39 + }) + + assert.Equal(t, []string{"key-in", "random-in"}, shadowModelNames(admitted)) + assert.Equal(t, 1, randomCalls) +} + func TestShadowBucketForBearerCredentialFixedVector(t *testing.T) { assert.Equal(t, 46, shadowBucketForBearerCredential([]byte("test-key"))) } diff --git a/src/invocation-plane-services/vanity-gateway/gateway/shadow_tracing_test.go b/src/invocation-plane-services/vanity-gateway/gateway/shadow_tracing_test.go index 86cff0cef..5bd357615 100644 --- a/src/invocation-plane-services/vanity-gateway/gateway/shadow_tracing_test.go +++ b/src/invocation-plane-services/vanity-gateway/gateway/shadow_tracing_test.go @@ -18,6 +18,7 @@ limitations under the License. package gateway import ( + config "ai-api-gateway-service/gateway_config" "bytes" "context" "net/http" @@ -30,6 +31,7 @@ import ( "github.com/stretchr/testify/require" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/sdk/trace/tracetest" ) @@ -83,9 +85,8 @@ func TestShadowSpanIsolation(t *testing.T) { modelMapping := map[string]FunctionInfo{ "primary-model": { - functionId: "func-primary", - shadowModelNames: []string{"shadow-model"}, - shadowPercentage: 100, + functionId: "func-primary", + shadows: testShadowConfigs([]string{"shadow-model"}, 100), }, "shadow-model": { functionId: "func-shadow", @@ -107,9 +108,12 @@ func TestShadowSpanIsolation(t *testing.T) { resolved := resolvedOpenAIRequest{ request: req, functionInfo: FunctionInfo{ - shadowModelNames: []string{"shadow-model"}, - shadowPercentage: 100, - shadowCancelOnClientDisconnect: true, // attached mode + shadows: testShadowConfigsWithPolicy( + []string{"shadow-model"}, + 100, + config.ShadowSamplingMethodRandom, + true, // attached mode + ), }, } director.dispatchShadowIfNeeded(resolved, modelMapping) @@ -173,10 +177,13 @@ func TestDetachedShadowKeepsParentSpanContext(t *testing.T) { modelMapping := map[string]FunctionInfo{ "primary-model": { - functionId: "func-primary", - shadowModelNames: []string{"shadow-model"}, - shadowPercentage: 100, - shadowCancelOnClientDisconnect: false, // detached + functionId: "func-primary", + shadows: testShadowConfigsWithPolicy( + []string{"shadow-model"}, + 100, + config.ShadowSamplingMethodRandom, + false, // detached + ), }, "shadow-model": { functionId: "func-shadow", @@ -197,9 +204,7 @@ func TestDetachedShadowKeepsParentSpanContext(t *testing.T) { resolved := resolvedOpenAIRequest{ request: req, functionInfo: FunctionInfo{ - shadowModelNames: []string{"shadow-model"}, - shadowPercentage: 100, - shadowCancelOnClientDisconnect: false, + shadows: testShadowConfigs([]string{"shadow-model"}, 100), }, } director.dispatchShadowIfNeeded(resolved, modelMapping) @@ -250,9 +255,8 @@ func TestShadowSpanRecordsTimeoutError(t *testing.T) { modelMapping := map[string]FunctionInfo{ "primary-model": { - functionId: "func-primary", - shadowModelNames: []string{"shadow-model"}, - shadowPercentage: 100, + functionId: "func-primary", + shadows: testShadowConfigs([]string{"shadow-model"}, 100), }, "shadow-model": { functionId: "func-shadow", @@ -273,8 +277,7 @@ func TestShadowSpanRecordsTimeoutError(t *testing.T) { resolved := resolvedOpenAIRequest{ request: req, functionInfo: FunctionInfo{ - shadowModelNames: []string{"shadow-model"}, - shadowPercentage: 100, + shadows: testShadowConfigs([]string{"shadow-model"}, 100), }, } director.dispatchShadowIfNeeded(resolved, modelMapping) @@ -330,9 +333,8 @@ func TestShadowSpanRecordsHTTPError(t *testing.T) { modelMapping := map[string]FunctionInfo{ "primary-model": { - functionId: "func-primary", - shadowModelNames: []string{"shadow-model"}, - shadowPercentage: 100, + functionId: "func-primary", + shadows: testShadowConfigs([]string{"shadow-model"}, 100), }, "shadow-model": { functionId: "func-shadow", @@ -353,8 +355,7 @@ func TestShadowSpanRecordsHTTPError(t *testing.T) { resolved := resolvedOpenAIRequest{ request: req, functionInfo: FunctionInfo{ - shadowModelNames: []string{"shadow-model"}, - shadowPercentage: 100, + shadows: testShadowConfigs([]string{"shadow-model"}, 100), }, } director.dispatchShadowIfNeeded(resolved, modelMapping) @@ -397,9 +398,8 @@ func TestShadowAdmissionAttributes(t *testing.T) { modelMapping := map[string]FunctionInfo{ "primary-model": { - functionId: "func-primary", - shadowModelNames: []string{"shadow-model"}, - shadowPercentage: 100, + functionId: "func-primary", + shadows: testShadowConfigs([]string{"shadow-model"}, 100), }, "shadow-model": { functionId: "func-shadow", @@ -418,7 +418,7 @@ func TestShadowAdmissionAttributes(t *testing.T) { director.dispatchShadowIfNeeded(resolvedOpenAIRequest{ request: req, - functionInfo: FunctionInfo{shadowModelNames: []string{"shadow-model"}, shadowPercentage: 100}, + functionInfo: FunctionInfo{shadows: testShadowConfigs([]string{"shadow-model"}, 100)}, }, modelMapping) assert.Eventually(t, func() bool { @@ -470,7 +470,7 @@ func TestShadowAdmissionAttributes(t *testing.T) { director.dispatchShadowIfNeeded(resolvedOpenAIRequest{ request: req, - functionInfo: FunctionInfo{shadowModelNames: []string{"shadow-model"}, shadowPercentage: 100}, + functionInfo: FunctionInfo{shadows: testShadowConfigs([]string{"shadow-model"}, 100)}, }, modelMapping) primarySpan.End() @@ -497,6 +497,119 @@ func TestShadowAdmissionAttributes(t *testing.T) { }) } +func TestPerTargetSamplingRecordsOnlyAdmittedShadow(t *testing.T) { + tp, exp := newTestTracerProvider() + defer tp.Shutdown(context.Background()) + metricReader := sdkmetric.NewManualReader() + metricProvider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(metricReader)) + t.Cleanup(func() { + require.NoError(t, metricProvider.Shutdown(context.Background())) + }) + counter, err := newShadowDroppedCounter(metricProvider.Meter(shadowMetricsScope)) + require.NoError(t, err) + previousCounter := shadowDroppedCounter + shadowDroppedCounter = counter + t.Cleanup(func() { + shadowDroppedCounter = previousCounter + }) + initializeShadowDropMetrics([]string{"shadow-admitted", "shadow-sampled-out"}) + + tracer := tp.Tracer("test") + replayedFunctions := make(chan string, 2) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + replayedFunctions <- r.Header.Get("function-id") + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + vanity, err := NewVanityDirector(upstream.URL, http.DefaultTransport) + require.NoError(t, err) + + randomCalls := 0 + director := &OpenAIDirector{ + shadower: NewTrafficShadower(2, 30*time.Second), + vanityDirector: vanity, + shadowRandomBucket: func() int { + randomCalls++ + return 39 + }, + } + modelMapping := map[string]FunctionInfo{ + "shadow-admitted": {functionId: "func-admitted"}, + "shadow-sampled-out": {functionId: "func-sampled-out"}, + } + + ctx, primarySpan := tracer.Start(context.Background(), "primary_partial_shadow") + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + bytes.NewBufferString(`{"model":"primary-model"}`)) + req = req.WithContext(ctx) + req.Header.Set("Content-Type", "application/json") + + finishPrimary := director.dispatchShadowIfNeeded(resolvedOpenAIRequest{ + request: req, + functionInfo: FunctionInfo{shadows: []shadowConfig{ + { + modelName: "shadow-admitted", + percentage: 40, + samplingMethod: config.ShadowSamplingMethodRandom, + }, + { + modelName: "shadow-sampled-out", + percentage: 39, + samplingMethod: config.ShadowSamplingMethodRandom, + }, + }}, + }, modelMapping) + assert.Equal(t, 1, randomCalls) + + select { + case functionID := <-replayedFunctions: + assert.Equal(t, "func-admitted", functionID) + case <-time.After(time.Second): + t.Fatal("admitted shadow did not replay") + } + select { + case functionID := <-replayedFunctions: + t.Fatalf("sampled-out shadow replayed with function ID %q", functionID) + case <-time.After(100 * time.Millisecond): + } + + assert.Eventually(t, func() bool { + tp.ForceFlush(context.Background()) + return len(spansByName(exp.GetSpans(), shadowReplaySpanName)) == 1 + }, 5*time.Second, 50*time.Millisecond) + + finishPrimary(nil) + primarySpan.End() + tp.ForceFlush(context.Background()) + + primarySpans := spansByName(exp.GetSpans(), "primary_partial_shadow") + require.Len(t, primarySpans, 1) + + dispatchedCount, ok := spanAttr(primarySpans[0], traceAttrShadowDispatchedCount) + require.True(t, ok) + assert.Equal(t, int64(1), dispatchedCount.AsInt64()) + + droppedCount, ok := spanAttr(primarySpans[0], traceAttrShadowDroppedCount) + require.True(t, ok) + assert.Equal(t, int64(0), droppedCount.AsInt64()) + + targetModels, ok := spanAttr(primarySpans[0], traceAttrShadowTargetModels) + require.True(t, ok) + assert.Equal(t, []string{"shadow-admitted"}, targetModels.AsStringSlice()) + + _, hasDroppedReasons := spanAttr(primarySpans[0], traceAttrShadowDroppedReasons) + assert.False(t, hasDroppedReasons) + _, hasDroppedTargetModels := spanAttr(primarySpans[0], traceAttrShadowDroppedTargetModels) + assert.False(t, hasDroppedTargetModels) + + dropCounts := collectShadowDroppedCounts(t, metricReader) + require.Len(t, dropCounts, len(shadowDroppedReasons)*2) + for labels, count := range dropCounts { + assert.Zero(t, count, "unexpected drop count for model %s and reason %s", labels.model, labels.reason) + } +} + func TestMultiShadowDispatchRecordsAggregatePrimarySpanAttributes(t *testing.T) { tp, exp := newTestTracerProvider() defer tp.Shutdown(context.Background()) @@ -515,9 +628,8 @@ func TestMultiShadowDispatchRecordsAggregatePrimarySpanAttributes(t *testing.T) modelMapping := map[string]FunctionInfo{ "primary-model": { - functionId: "func-primary", - shadowModelNames: []string{"shadow-a", "shadow-b"}, - shadowPercentage: 100, + functionId: "func-primary", + shadows: testShadowConfigs([]string{"shadow-a", "shadow-b"}, 100), }, "shadow-a": {functionId: "func-shadow-a"}, "shadow-b": {functionId: "func-shadow-b"}, @@ -533,7 +645,7 @@ func TestMultiShadowDispatchRecordsAggregatePrimarySpanAttributes(t *testing.T) director.dispatchShadowIfNeeded(resolvedOpenAIRequest{ request: req, - functionInfo: FunctionInfo{shadowModelNames: []string{"shadow-a", "shadow-b"}, shadowPercentage: 100}, + functionInfo: FunctionInfo{shadows: testShadowConfigs([]string{"shadow-a", "shadow-b"}, 100)}, }, modelMapping) primarySpan.End() @@ -590,9 +702,8 @@ func TestMultiShadowDispatchRecordsAllDroppedReasonsAndTargets(t *testing.T) { modelMapping := map[string]FunctionInfo{ "primary-model": { - functionId: "func-primary", - shadowModelNames: []string{"shadow-a", "shadow-b"}, - shadowPercentage: 100, + functionId: "func-primary", + shadows: testShadowConfigs([]string{"shadow-a", "shadow-b"}, 100), }, "shadow-a": {functionId: "func-shadow-a"}, "shadow-b": {functionId: "func-shadow-b"}, @@ -620,7 +731,7 @@ func TestMultiShadowDispatchRecordsAllDroppedReasonsAndTargets(t *testing.T) { director.dispatchShadowIfNeeded(resolvedOpenAIRequest{ request: req, - functionInfo: FunctionInfo{shadowModelNames: []string{"shadow-a", "shadow-b"}, shadowPercentage: 100}, + functionInfo: FunctionInfo{shadows: testShadowConfigs([]string{"shadow-a", "shadow-b"}, 100)}, }, modelMapping) primarySpan.End() diff --git a/src/invocation-plane-services/vanity-gateway/gateway_config/gateway_config.go b/src/invocation-plane-services/vanity-gateway/gateway_config/gateway_config.go index 282b904e5..f98bd99a4 100644 --- a/src/invocation-plane-services/vanity-gateway/gateway_config/gateway_config.go +++ b/src/invocation-plane-services/vanity-gateway/gateway_config/gateway_config.go @@ -36,6 +36,35 @@ const ( ShadowSamplingMethodPerBearerKey ShadowSamplingMethod = "perBearerKey" ) +type ShadowConfig struct { + ModelName string `json:"modelName"` + Percentage *int `json:"percentage,omitempty"` + SamplingMethod ShadowSamplingMethod `json:"samplingMethod,omitempty"` + CancelOnClientDisconnect bool `json:"cancelOnClientDisconnect,omitempty"` +} + +func (s *ShadowConfig) UnmarshalJSON(data []byte) error { + fields := map[string]json.RawMessage{} + if err := json.Unmarshal(data, &fields); err != nil { + return err + } + for field := range fields { + switch field { + case "modelName", "percentage", "samplingMethod", "cancelOnClientDisconnect": + default: + return fmt.Errorf("unknown shadow config field %q", field) + } + } + + type shadowConfigAlias ShadowConfig + var alias shadowConfigAlias + if err := json.Unmarshal(data, &alias); err != nil { + return err + } + *s = ShadowConfig(alias) + return nil +} + type CustomHeaders map[string]string func (h *CustomHeaders) UnmarshalJSON(data []byte) error { @@ -72,25 +101,108 @@ type ModelFunctionDetails struct { EOL time.Time `json:"eol,omitempty"` // RFC3339 timestamp (full ISO 8601) OfflineMessage string `json:"offlineMessage,omitempty"` // non-empty = endpoint is offline TooManyRequestsMessage string `json:"tooManyRequestsMessage"` + Shadows []ShadowConfig `json:"shadows,omitempty"` ShadowModelName string `json:"shadowModelName,omitempty"` ShadowModelNames []string `json:"shadowModelNames,omitempty"` ShadowPercentage *int `json:"shadowPercentage,omitempty"` // 1-100 when set; omitted defaults to 100 ShadowSamplingMethod ShadowSamplingMethod `json:"shadowSamplingMethod,omitempty"` - ShadowCancelOnClientDisconnect bool `json:"shadowCancelOnClientDisconnect,omitempty"` // cancel shadow when primary completes; default false + ShadowCancelOnClientDisconnect bool `json:"shadowCancelOnClientDisconnect,omitempty"` // cancel shadows on client cancellation or primary proxy failure; default false + shadowsPresent bool + legacyShadowFieldsPresent bool } func (m *ModelFunctionDetails) UnmarshalJSON(data []byte) error { type modelFunctionDetailsAlias ModelFunctionDetails + fields := map[string]json.RawMessage{} + if err := json.Unmarshal(data, &fields); err != nil { + return err + } + var alias modelFunctionDetailsAlias if err := json.Unmarshal(data, &alias); err != nil { return err } *m = ModelFunctionDetails(alias) + for field := range fields { + switch { + case strings.EqualFold(field, "shadows"): + m.shadowsPresent = true + case strings.EqualFold(field, "shadowModelName"), + strings.EqualFold(field, "shadowModelNames"), + strings.EqualFold(field, "shadowPercentage"), + strings.EqualFold(field, "shadowSamplingMethod"), + strings.EqualFold(field, "shadowCancelOnClientDisconnect"): + m.legacyShadowFieldsPresent = true + } + } return nil } +func (m ModelFunctionDetails) EffectiveShadows() []ShadowConfig { + if len(m.Shadows) > 0 { + shadows := make([]ShadowConfig, len(m.Shadows)) + for i, shadow := range m.Shadows { + shadows[i] = cloneShadowConfig(shadow) + } + return shadows + } + + modelNames := make([]string, 0, len(m.ShadowModelNames)+1) + if m.ShadowModelName != "" { + modelNames = append(modelNames, m.ShadowModelName) + } + modelNames = append(modelNames, m.ShadowModelNames...) + if len(modelNames) == 0 { + return nil + } + + shadows := make([]ShadowConfig, 0, len(modelNames)) + for _, modelName := range modelNames { + shadows = append(shadows, ShadowConfig{ + ModelName: modelName, + Percentage: cloneInt(m.ShadowPercentage), + SamplingMethod: m.ShadowSamplingMethod, + CancelOnClientDisconnect: m.ShadowCancelOnClientDisconnect, + }) + } + return shadows +} + +func cloneShadowConfig(shadow ShadowConfig) ShadowConfig { + shadow.Percentage = cloneInt(shadow.Percentage) + return shadow +} + +func cloneInt(value *int) *int { + if value == nil { + return nil + } + cloned := *value + return &cloned +} + +func (m ModelFunctionDetails) hasLegacyShadowConfig() bool { + return m.ShadowModelName != "" || + len(m.ShadowModelNames) > 0 || + m.ShadowPercentage != nil || + m.ShadowSamplingMethod != "" || + m.ShadowCancelOnClientDisconnect +} + +func (m ModelFunctionDetails) hasShadowsField() bool { + return m.shadowsPresent || len(m.Shadows) > 0 +} + +func (m ModelFunctionDetails) hasLegacyShadowFields() bool { + return m.hasLegacyShadowConfig() || m.legacyShadowFieldsPresent +} + +func (m ModelFunctionDetails) hasMixedShadowFields() bool { + return m.hasShadowsField() && m.hasLegacyShadowFields() +} + type PathFunctionDetails struct { Path string `json:"path"` // incoming path OutgoingPathOverride *string `json:"outgoingPathOverride"` @@ -192,7 +304,17 @@ func uniqueShadowModelNames(legacyModelName string, modelNames []string) ([]stri return result, nil } -func validateOpenAIShadowConfig(location string, entry ModelFunctionDetails) ([]string, error) { +func validateOpenAIShadowConfig(location string, entry ModelFunctionDetails) ([]ShadowConfig, error) { + if entry.hasMixedShadowFields() { + return nil, fmt.Errorf("%s: shadows cannot be combined with legacy shadow fields", location) + } + if entry.hasShadowsField() { + if err := validatePerTargetShadowConfigs(location, entry.Shadows); err != nil { + return nil, err + } + return entry.EffectiveShadows(), nil + } + shadowTargets, err := uniqueShadowModelNames(entry.ShadowModelName, entry.ShadowModelNames) if err != nil { return nil, fmt.Errorf("%s: %w", location, err) @@ -221,15 +343,44 @@ func validateOpenAIShadowConfig(location string, entry ModelFunctionDetails) ([] } } - return shadowTargets, nil + return entry.EffectiveShadows(), nil +} + +func validatePerTargetShadowConfigs(location string, shadows []ShadowConfig) error { + seen := make(map[string]struct{}, len(shadows)) + for i, shadow := range shadows { + shadowLocation := fmt.Sprintf("%s.shadows[%d]", location, i) + if shadow.ModelName == "" { + return fmt.Errorf("%s: modelName is required", shadowLocation) + } + if _, ok := seen[shadow.ModelName]; ok { + return fmt.Errorf("%s: duplicate shadow target %q", shadowLocation, shadow.ModelName) + } + seen[shadow.ModelName] = struct{}{} + + if shadow.Percentage != nil { + percentage := *shadow.Percentage + if percentage < 1 || percentage > 100 { + return fmt.Errorf("%s: percentage must be between 1 and 100", shadowLocation) + } + } + if err := validateSamplingMethod(shadowLocation, "samplingMethod", shadow.SamplingMethod); err != nil { + return err + } + } + return nil } func validateShadowSamplingMethod(location string, method ShadowSamplingMethod) error { + return validateSamplingMethod(location, "shadowSamplingMethod", method) +} + +func validateSamplingMethod(location string, fieldName string, method ShadowSamplingMethod) error { switch method { case "", ShadowSamplingMethodRandom, ShadowSamplingMethodPerBearerKey: return nil default: - return fmt.Errorf("%s: shadowSamplingMethod must be %q or %q", location, ShadowSamplingMethodRandom, ShadowSamplingMethodPerBearerKey) + return fmt.Errorf("%s: %s must be %q or %q", location, fieldName, ShadowSamplingMethodRandom, ShadowSamplingMethodPerBearerKey) } } @@ -295,7 +446,10 @@ func isMultipartOpenAISection(sectionName string) bool { func validateMultipartOpenAISection(sectionName string, entries map[string]ModelFunctionDetails) error { for modelKey, entry := range entries { - if entry.ShadowModelName != "" || len(entry.ShadowModelNames) > 0 || entry.ShadowPercentage != nil || entry.ShadowSamplingMethod != "" || entry.ShadowCancelOnClientDisconnect { + if entry.hasMixedShadowFields() { + return fmt.Errorf("openai.%s.%s: shadows cannot be combined with legacy shadow fields", sectionName, modelKey) + } + if len(entry.Shadows) > 0 || entry.hasLegacyShadowConfig() { return fmt.Errorf("openai.%s.%s: shadow config is unsupported for multipart image endpoints", sectionName, modelKey) } } @@ -309,20 +463,24 @@ func validateOpenAIShadowTargets(sectionName string, entries map[string]ModelFun if err != nil { return err } - if err := validateShadowTargetNames(location, sectionName, entry.ModelName, shadowTargets, modelNames); err != nil { + if err := validateShadowTargetNames(location, sectionName, entry.ModelName, shadowTargets, modelNames, len(entry.Shadows) > 0); err != nil { return err } } return nil } -func validateShadowTargetNames(location string, sectionName string, modelName string, shadowTargets []string, modelNames map[string]struct{}) error { - for _, shadowTarget := range shadowTargets { - if shadowTarget == modelName { - return fmt.Errorf("%s: shadow target cannot reference the same model", location) +func validateShadowTargetNames(location string, sectionName string, modelName string, shadows []ShadowConfig, modelNames map[string]struct{}, perTarget bool) error { + for i, shadow := range shadows { + shadowLocation := location + if perTarget { + shadowLocation = fmt.Sprintf("%s.shadows[%d]", location, i) + } + if shadow.ModelName == modelName { + return fmt.Errorf("%s: shadow target cannot reference the same model", shadowLocation) } - if _, ok := modelNames[shadowTarget]; !ok { - return fmt.Errorf("%s: shadow target must reference another model in openai.%s", location, sectionName) + if _, ok := modelNames[shadow.ModelName]; !ok { + return fmt.Errorf("%s: shadow target must reference another model in openai.%s", shadowLocation, sectionName) } } return nil diff --git a/src/invocation-plane-services/vanity-gateway/gateway_config/gateway_config_test.go b/src/invocation-plane-services/vanity-gateway/gateway_config/gateway_config_test.go index 4a59f6bc0..1c689054e 100644 --- a/src/invocation-plane-services/vanity-gateway/gateway_config/gateway_config_test.go +++ b/src/invocation-plane-services/vanity-gateway/gateway_config/gateway_config_test.go @@ -18,6 +18,7 @@ limitations under the License. package config import ( + "fmt" "os" "path/filepath" "testing" @@ -106,6 +107,126 @@ func TestGatewayConfigValidateAcceptsMultipleOpenAIShadows(t *testing.T) { require.NoError(t, cfg.Validate()) } +func TestGatewayConfigValidateAcceptsPerTargetShadows(t *testing.T) { + percentage := 10 + cfg := &GatewayConfig{} + cfg.OpenAI.ChatCompletions = map[string]ModelFunctionDetails{ + "primary": { + ModelName: "facebook/opt-125m", + FunctionID: "func-id", + Shadows: []ShadowConfig{ + { + ModelName: "private/facebook/opt-125m-shadow-a", + Percentage: &percentage, + SamplingMethod: ShadowSamplingMethodPerBearerKey, + CancelOnClientDisconnect: true, + }, + { + ModelName: "private/facebook/opt-125m-shadow-b", + SamplingMethod: ShadowSamplingMethodRandom, + }, + }, + }, + "shadow-a": { + ModelName: "private/facebook/opt-125m-shadow-a", + FunctionID: "shadow-a-func-id", + }, + "shadow-b": { + ModelName: "private/facebook/opt-125m-shadow-b", + FunctionID: "shadow-b-func-id", + }, + } + + require.NoError(t, cfg.Validate()) +} + +func TestGatewayConfigValidateTreatsProgrammaticEmptyShadowsAsAbsent(t *testing.T) { + cfg := &GatewayConfig{} + cfg.OpenAI.ChatCompletions = map[string]ModelFunctionDetails{ + "primary": { + ModelName: "facebook/opt-125m", + FunctionID: "func-id", + Shadows: []ShadowConfig{}, + ShadowModelName: "private/facebook/opt-125m-shadow", + }, + "shadow": { + ModelName: "private/facebook/opt-125m-shadow", + FunctionID: "shadow-func-id", + }, + } + + require.NoError(t, cfg.Validate()) + assert.Equal(t, []ShadowConfig{{ModelName: "private/facebook/opt-125m-shadow"}}, + cfg.OpenAI.ChatCompletions["primary"].EffectiveShadows()) +} + +func TestModelFunctionDetailsEffectiveShadowsNormalizesLegacyFields(t *testing.T) { + percentage := 25 + entry := ModelFunctionDetails{ + ShadowModelName: "shadow-a", + ShadowModelNames: []string{"shadow-b"}, + ShadowPercentage: &percentage, + ShadowSamplingMethod: ShadowSamplingMethodPerBearerKey, + ShadowCancelOnClientDisconnect: true, + } + + shadows := entry.EffectiveShadows() + require.Len(t, shadows, 2) + assert.Equal(t, "shadow-a", shadows[0].ModelName) + assert.Equal(t, "shadow-b", shadows[1].ModelName) + for _, shadow := range shadows { + require.NotNil(t, shadow.Percentage) + assert.Equal(t, 25, *shadow.Percentage) + assert.Equal(t, ShadowSamplingMethodPerBearerKey, shadow.SamplingMethod) + assert.True(t, shadow.CancelOnClientDisconnect) + } + + *shadows[0].Percentage = 50 + assert.Equal(t, 25, percentage) + assert.Equal(t, 25, *shadows[1].Percentage) +} + +func TestModelFunctionDetailsEffectiveShadowsReturnsCopy(t *testing.T) { + percentage := 10 + entry := ModelFunctionDetails{ + Shadows: []ShadowConfig{{ModelName: "shadow", Percentage: &percentage}}, + } + + shadows := entry.EffectiveShadows() + shadows[0].ModelName = "changed" + *shadows[0].Percentage = 20 + + assert.Equal(t, "shadow", entry.Shadows[0].ModelName) + assert.Equal(t, 10, *entry.Shadows[0].Percentage) +} + +func TestModelFunctionDetailsEffectiveShadowsLegacyAndPerTargetEquivalent(t *testing.T) { + percentage := 25 + legacy := ModelFunctionDetails{ + ShadowModelName: "shadow-a", + ShadowModelNames: []string{"shadow-b"}, + ShadowPercentage: &percentage, + ShadowSamplingMethod: ShadowSamplingMethodPerBearerKey, + ShadowCancelOnClientDisconnect: true, + } + perTarget := ModelFunctionDetails{Shadows: []ShadowConfig{ + { + ModelName: "shadow-a", + Percentage: intPtr(25), + SamplingMethod: ShadowSamplingMethodPerBearerKey, + CancelOnClientDisconnect: true, + }, + { + ModelName: "shadow-b", + Percentage: intPtr(25), + SamplingMethod: ShadowSamplingMethodPerBearerKey, + CancelOnClientDisconnect: true, + }, + }} + + assert.Equal(t, perTarget.EffectiveShadows(), legacy.EffectiveShadows()) +} + func TestGatewayConfigLoadAcceptsLegacyAndPluralShadowModelNames(t *testing.T) { configPath := filepath.Join(t.TempDir(), "config.yaml") err := os.WriteFile(configPath, []byte(` @@ -136,6 +257,189 @@ v2config: assert.Equal(t, []string{"private/facebook/opt-125m-shadow-b"}, primary.ShadowModelNames) } +func TestGatewayConfigLoadAcceptsStandaloneLegacyEmptyValues(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "config.yaml") + err := os.WriteFile(configPath, []byte(` +v2config: + openai: + chatCompletions: + primary: + modelName: facebook/opt-125m + functionID: func-id + shadowModelName: "" + shadowModelNames: [] + shadowPercentage: null + shadowSamplingMethod: "" + shadowCancelOnClientDisconnect: false + null-sampling: + modelName: facebook/opt-125m-null-sampling + functionID: null-sampling-func-id + shadowSamplingMethod: null +`), 0600) + require.NoError(t, err) + + reloadable, err := SetupConfigWithConfigPath(configPath) + require.NoError(t, err) + + primary := reloadable.Get().OpenAI.ChatCompletions["primary"] + assert.Empty(t, primary.EffectiveShadows()) + assert.Empty(t, reloadable.Get().OpenAI.ChatCompletions["null-sampling"].EffectiveShadows()) +} + +func TestGatewayConfigLoadAcceptsPerTargetShadows(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "config.yaml") + err := os.WriteFile(configPath, []byte(` +v2config: + openai: + chatCompletions: + primary: + modelName: facebook/opt-125m + functionID: func-id + shadows: + - modelName: private/facebook/opt-125m-shadow-a + percentage: 10 + samplingMethod: perBearerKey + cancelOnClientDisconnect: true + - modelName: private/facebook/opt-125m-shadow-b + percentage: 50 + shadow-a: + modelName: private/facebook/opt-125m-shadow-a + functionID: shadow-a-func-id + shadow-b: + modelName: private/facebook/opt-125m-shadow-b + functionID: shadow-b-func-id +`), 0600) + require.NoError(t, err) + + reloadable, err := SetupConfigWithConfigPath(configPath) + require.NoError(t, err) + + shadows := reloadable.Get().OpenAI.ChatCompletions["primary"].Shadows + require.Len(t, shadows, 2) + assert.Equal(t, "private/facebook/opt-125m-shadow-a", shadows[0].ModelName) + assert.Equal(t, 10, *shadows[0].Percentage) + assert.Equal(t, ShadowSamplingMethodPerBearerKey, shadows[0].SamplingMethod) + assert.True(t, shadows[0].CancelOnClientDisconnect) + assert.Equal(t, "private/facebook/opt-125m-shadow-b", shadows[1].ModelName) + assert.Equal(t, 50, *shadows[1].Percentage) +} + +func TestGatewayConfigLoadRejectsUnknownPerTargetShadowFields(t *testing.T) { + tests := []struct { + name string + field string + }{ + {name: "typo", field: "precentage: 10"}, + {name: "nested legacy field", field: "shadowPercentage: 10"}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "config.yaml") + contents := fmt.Sprintf(` +v2config: + openai: + chatCompletions: + primary: + modelName: facebook/opt-125m + functionID: func-id + shadows: + - modelName: private/facebook/opt-125m-shadow + %s +`, tc.field) + err := os.WriteFile(configPath, []byte(contents), 0600) + require.NoError(t, err) + + _, err = SetupConfigWithConfigPath(configPath) + require.Error(t, err) + assert.ErrorContains(t, err, "unknown shadow config field") + }) + } +} + +func TestGatewayConfigLoadRejectsMixedShadowFieldPresence(t *testing.T) { + tests := []struct { + name string + section string + fields string + }{ + { + name: "empty legacy single target", + fields: ` shadows: + - modelName: private/facebook/opt-125m-shadow + shadowModelName: ""`, + }, + { + name: "empty legacy target list", + fields: ` shadows: + - modelName: private/facebook/opt-125m-shadow + shadowModelNames: []`, + }, + { + name: "null legacy percentage", + fields: ` shadows: + - modelName: private/facebook/opt-125m-shadow + shadowPercentage: null`, + }, + { + name: "empty legacy sampling method", + fields: ` shadows: + - modelName: private/facebook/opt-125m-shadow + shadowSamplingMethod: ""`, + }, + { + name: "false legacy cancellation policy", + fields: ` shadows: + - modelName: private/facebook/opt-125m-shadow + shadowCancelOnClientDisconnect: false`, + }, + { + name: "empty per-target list", + fields: ` shadows: [] + shadowModelName: private/facebook/opt-125m-shadow`, + }, + { + name: "null per-target list", + fields: ` shadows: null + shadowModelName: private/facebook/opt-125m-shadow`, + }, + { + name: "multipart empty fields", + section: "imageEdits", + fields: ` shadows: [] + shadowCancelOnClientDisconnect: false`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + section := tc.section + if section == "" { + section = "chatCompletions" + } + configPath := filepath.Join(t.TempDir(), "config.yaml") + contents := fmt.Sprintf(` +v2config: + openai: + %s: + primary: + modelName: facebook/opt-125m + functionID: func-id +%s + shadow: + modelName: private/facebook/opt-125m-shadow + functionID: shadow-func-id +`, section, tc.fields) + err := os.WriteFile(configPath, []byte(contents), 0600) + require.NoError(t, err) + + _, err = SetupConfigWithConfigPath(configPath) + require.Error(t, err) + assert.ErrorContains(t, err, "shadows cannot be combined with legacy shadow fields") + }) + } +} + func TestGatewayConfigLoadAcceptsShadowSamplingMethod(t *testing.T) { configPath := filepath.Join(t.TempDir(), "config.yaml") err := os.WriteFile(configPath, []byte(` @@ -471,6 +775,156 @@ func TestGatewayConfigValidateRejectsDuplicateOpenAIShadows(t *testing.T) { } } +func TestGatewayConfigValidateRejectsMixedShadowForms(t *testing.T) { + percentage := 50 + tests := []struct { + name string + apply func(*ModelFunctionDetails) + }{ + { + name: "legacy single target", + apply: func(entry *ModelFunctionDetails) { + entry.ShadowModelName = "private/facebook/opt-125m-shadow-b" + }, + }, + { + name: "legacy target list", + apply: func(entry *ModelFunctionDetails) { + entry.ShadowModelNames = []string{"private/facebook/opt-125m-shadow-b"} + }, + }, + { + name: "legacy percentage", + apply: func(entry *ModelFunctionDetails) { + entry.ShadowPercentage = &percentage + }, + }, + { + name: "legacy sampling method", + apply: func(entry *ModelFunctionDetails) { + entry.ShadowSamplingMethod = ShadowSamplingMethodRandom + }, + }, + { + name: "legacy cancellation policy", + apply: func(entry *ModelFunctionDetails) { + entry.ShadowCancelOnClientDisconnect = true + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + primary := ModelFunctionDetails{ + ModelName: "facebook/opt-125m", + FunctionID: "func-id", + Shadows: []ShadowConfig{{ + ModelName: "private/facebook/opt-125m-shadow-a", + }}, + } + tc.apply(&primary) + + cfg := &GatewayConfig{} + cfg.OpenAI.ChatCompletions = map[string]ModelFunctionDetails{ + "primary": primary, + "shadow-a": { + ModelName: "private/facebook/opt-125m-shadow-a", + FunctionID: "shadow-a-func-id", + }, + "shadow-b": { + ModelName: "private/facebook/opt-125m-shadow-b", + FunctionID: "shadow-b-func-id", + }, + } + + err := cfg.Validate() + require.Error(t, err) + assert.ErrorContains(t, err, "shadows cannot be combined with legacy shadow fields") + }) + } +} + +func TestGatewayConfigValidateRejectsInvalidPerTargetShadows(t *testing.T) { + tests := []struct { + name string + shadows []ShadowConfig + expected string + }{ + { + name: "empty model name", + shadows: []ShadowConfig{{}}, + expected: "shadows[0]: modelName is required", + }, + { + name: "duplicate target", + shadows: []ShadowConfig{ + {ModelName: "private/facebook/opt-125m-shadow-a"}, + {ModelName: "private/facebook/opt-125m-shadow-a"}, + }, + expected: "duplicate shadow target", + }, + { + name: "percentage below range", + shadows: []ShadowConfig{{ + ModelName: "private/facebook/opt-125m-shadow-a", + Percentage: intPtr(0), + }}, + expected: "percentage must be between 1 and 100", + }, + { + name: "percentage above range", + shadows: []ShadowConfig{{ + ModelName: "private/facebook/opt-125m-shadow-a", + Percentage: intPtr(101), + }}, + expected: "percentage must be between 1 and 100", + }, + { + name: "invalid sampling method", + shadows: []ShadowConfig{{ + ModelName: "private/facebook/opt-125m-shadow-a", + SamplingMethod: ShadowSamplingMethod("weighted"), + }}, + expected: "samplingMethod must be", + }, + { + name: "self reference", + shadows: []ShadowConfig{{ + ModelName: "facebook/opt-125m", + }}, + expected: "shadow target cannot reference the same model", + }, + { + name: "missing target", + shadows: []ShadowConfig{{ + ModelName: "private/facebook/missing-shadow", + }}, + expected: "shadow target must reference another model in openai.chatCompletions", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + cfg := &GatewayConfig{} + cfg.OpenAI.ChatCompletions = map[string]ModelFunctionDetails{ + "primary": { + ModelName: "facebook/opt-125m", + FunctionID: "func-id", + Shadows: tc.shadows, + }, + "shadow-a": { + ModelName: "private/facebook/opt-125m-shadow-a", + FunctionID: "shadow-a-func-id", + }, + } + + err := cfg.Validate() + require.Error(t, err) + assert.ErrorContains(t, err, tc.expected) + }) + } +} + func TestGatewayConfigValidateRejectsInvalidOpenAIShadowPercentage(t *testing.T) { tests := []struct { name string @@ -749,6 +1203,20 @@ func TestGatewayConfigValidateRejectsShadowOnMultipartImageSections(t *testing.T } }, }, + { + name: "imageEdits per-target shadows", + applyTo: func(cfg *GatewayConfig) { + cfg.OpenAI.ImageEdits = map[string]ModelFunctionDetails{ + "edit": { + ModelName: "qwen/qwen-image-edit-2511", + FunctionID: "edit-id", + Shadows: []ShadowConfig{{ + ModelName: "qwen/qwen-image-edit-shadow", + }}, + }, + } + }, + }, { name: "imageVariations shadowPercentage", applyTo: func(cfg *GatewayConfig) {