Skip to content

Adding model based round robin policy for DP to CP flow #1266

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
CrowleyRajapakse merged 1 commit into from
Feb 25, 2025
Merged

Adding model based round robin policy for DP to CP flow #1266

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions apim-apk-agent/internal/constants/constants.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,4 +48,7 @@ const (
// Version constants
V1 = "v1"
V2 = "v2"

// Policy Types
CommonType = "common"
)
291 changes: 191 additions & 100 deletions apim-apk-agent/pkg/managementserver/rest_server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,104 @@ func createAPIYaml(apiCPEvent *APICPEvent) (string, string, string) {
provider = config.ControlPlane.Provider
}
context := removeVersionSuffix(apiCPEvent.API.BasePath, apiCPEvent.API.APIVersion)
operations, scopes, operationsErr := extractOperations(*apiCPEvent)

multiEndpoints := apiCPEvent.API.MultiEndpoints
apimEndpints := []APIMEndpoint{}
prodCount := 0
sandCount := 0
primaryProductionEndpointID := ""
primarySandboxEndpointID := ""
primaryProdcutionURL := ""
primarySandboxURL := ""
for _, endpoint := range multiEndpoints.ProdEndpoints {
prodCount++
var endpointName string
if prodCount == 1 {
endpointName = "Default Production Endpoint"
} else {
endpointName = fmt.Sprintf("%d Production Endpoint", prodCount)
}
prodEndpoint := ""
if endpoint.URL != "" {
prodEndpoint = fmt.Sprintf("%s://%s", multiEndpoints.Protocol, endpoint.URL)
}
endpointUUID := uuid.New().String() + "--PRODUCTION"
if prodCount == 1 {
primaryProductionEndpointID = endpointUUID
primaryProdcutionURL = prodEndpoint
}
apimEndpints = append(apimEndpints, APIMEndpoint{
DeploymentStage: "PRODUCTION",
EndpointUUID: endpointUUID,
EndpointName: endpointName,
EndpointConfig: APIMEndpointConfig{
EndpointType: multiEndpoints.Protocol,
ProductionEndpoints: Endpoints{
URL: prodEndpoint,
},
EndpointSecurity: APIMEndpointSecurity{
Production: SecurityConfig{
Enabled: endpoint.SecurityEnabled,
Type: endpoint.SecurityType,
Username: endpoint.BasicUsername,
Password: endpoint.BasicPassword,
APIKeyIdentifier: endpoint.APIKeyName,
APIKeyValue: endpoint.APIKeyValue,
APIKeyIdentifierType: endpoint.APIKeyIn,
ConnectionTimeoutDuration: -1.0,
SocketTimeoutDuration: -1.0,
ConnectionRequestTimeoutDuration: -1.0,
},
},
},
})
}
for _, endpoint := range multiEndpoints.SandEndpoints {
sandCount++
var endpointName string
if sandCount == 1 {
endpointName = "Default Sandbox Endpoint"
} else {
endpointName = fmt.Sprintf("%d Sandbox Endpoint", sandCount)
}

sandEndpoint := ""
if endpoint.URL != "" {
sandEndpoint = fmt.Sprintf("%s://%s", multiEndpoints.Protocol, endpoint.URL)
}
endpointUUID := uuid.New().String() + "--SANDBOX"
if sandCount == 1 {
primarySandboxEndpointID = endpointUUID
primarySandboxURL = sandEndpoint
}
apimEndpints = append(apimEndpints, APIMEndpoint{
DeploymentStage: "SANDBOX",
EndpointUUID: endpointUUID,
EndpointName: endpointName,
EndpointConfig: APIMEndpointConfig{
EndpointType: multiEndpoints.Protocol,
SandboxEndpoints: Endpoints{
URL: sandEndpoint,
},
EndpointSecurity: APIMEndpointSecurity{
Sandbox: SecurityConfig{
Enabled: endpoint.SecurityEnabled,
Type: endpoint.SecurityType,
Username: endpoint.BasicUsername,
Password: endpoint.BasicPassword,
APIKeyIdentifier: endpoint.APIKeyName,
APIKeyValue: endpoint.APIKeyValue,
APIKeyIdentifierType: endpoint.APIKeyIn,
ConnectionTimeoutDuration: -1.0,
SocketTimeoutDuration: -1.0,
ConnectionRequestTimeoutDuration: -1.0,
},
},
},
})
}

operations, scopes, operationsErr := extractOperations(*apiCPEvent, apimEndpints)
if operationsErr != nil {
logger.LoggerMgtServer.Errorf("Error occured while extracting operations from open API: %s, \nError: %+v", apiCPEvent.API.Definition, operationsErr)
operations = []APIOperation{}
Expand Down Expand Up @@ -283,6 +380,7 @@ func createAPIYaml(apiCPEvent *APICPEvent) (string, string, string) {
"accessControlExposeHeaders": apiCPEvent.API.CORSPolicy.AccessControlExposeHeaders,
}
}

maxTps := make(map[string]interface{})

// Handle Production fields
Expand Down Expand Up @@ -395,101 +493,6 @@ func createAPIYaml(apiCPEvent *APICPEvent) (string, string, string) {
}
}
}
multiEndpoints := apiCPEvent.API.MultiEndpoints
apimEndpints := []APIMEndpoint{}
prodCount := 0
sandCount := 0
primaryProductionEndpointID := ""
primarySandboxEndpointID := ""
primaryProdcutionURL := ""
primarySandboxURL := ""
for _, endpoint := range multiEndpoints.ProdEndpoints {
prodCount++
var endpointName string
if prodCount == 1 {
endpointName = "Default Production Endpoint"
} else {
endpointName = fmt.Sprintf("%d Production Endpoint", prodCount)
}
prodEndpoint := ""
if endpoint.URL != "" {
prodEndpoint = fmt.Sprintf("%s://%s", multiEndpoints.Protocol, endpoint.URL)
}
endpointUUID := uuid.New().String() + "--PRODUCTION"
if prodCount == 1 {
primaryProductionEndpointID = endpointUUID
primaryProdcutionURL = prodEndpoint
}
apimEndpints = append(apimEndpints, APIMEndpoint{
DeploymentStage: "PRODUCTION",
EndpointUUID: endpointUUID,
EndpointName: endpointName,
EndpointConfig: APIMEndpointConfig{
EndpointType: multiEndpoints.Protocol,
ProductionEndpoints: Endpoints{
URL: prodEndpoint,
},
EndpointSecurity: APIMEndpointSecurity{
Production: SecurityConfig{
Enabled: endpoint.SecurityEnabled,
Type: endpoint.SecurityType,
Username: endpoint.BasicUsername,
Password: endpoint.BasicPassword,
APIKeyIdentifier: endpoint.APIKeyName,
APIKeyValue: endpoint.APIKeyValue,
APIKeyIdentifierType: endpoint.APIKeyIn,
ConnectionTimeoutDuration: -1.0,
SocketTimeoutDuration: -1.0,
ConnectionRequestTimeoutDuration: -1.0,
},
},
},
})
}
for _, endpoint := range multiEndpoints.SandEndpoints {
sandCount++
var endpointName string
if sandCount == 1 {
endpointName = "Default Sandbox Endpoint"
} else {
endpointName = fmt.Sprintf("%d Sandbox Endpoint", sandCount)
}

sandEndpoint := ""
if endpoint.URL != "" {
sandEndpoint = fmt.Sprintf("%s://%s", multiEndpoints.Protocol, endpoint.URL)
}
endpointUUID := uuid.New().String() + "--SANDBOX"
if sandCount == 1 {
primarySandboxEndpointID = endpointUUID
primarySandboxURL = sandEndpoint
}
apimEndpints = append(apimEndpints, APIMEndpoint{
DeploymentStage: "SANDBOX",
EndpointUUID: endpointUUID,
EndpointName: endpointName,
EndpointConfig: APIMEndpointConfig{
EndpointType: multiEndpoints.Protocol,
SandboxEndpoints: Endpoints{
URL: sandEndpoint,
},
EndpointSecurity: APIMEndpointSecurity{
Sandbox: SecurityConfig{
Enabled: endpoint.SecurityEnabled,
Type: endpoint.SecurityType,
Username: endpoint.BasicUsername,
Password: endpoint.BasicPassword,
APIKeyIdentifier: endpoint.APIKeyName,
APIKeyValue: endpoint.APIKeyValue,
APIKeyIdentifierType: endpoint.APIKeyIn,
ConnectionTimeoutDuration: -1.0,
SocketTimeoutDuration: -1.0,
ConnectionRequestTimeoutDuration: -1.0,
},
},
},
})
}

dataArr := make([]map[string]interface{}, 0, len(apimEndpints))

Expand Down Expand Up @@ -636,9 +639,38 @@ func createAPIYaml(apiCPEvent *APICPEvent) (string, string, string) {
}
}

logger.LoggerMgtServer.Debugf("API Yaml: %+v", data)
var requestOperationPolicies []OperationPolicy
if apiCPEvent.API.AIModelBasedRoundRobin != nil {
aiModelBasedRoundRobin := apiCPEvent.API.AIModelBasedRoundRobin
logger.LoggerMgtServer.Infof("AIModelBasedRoundRobin : %+v", aiModelBasedRoundRobin)
wrr := ModelBasedRoundRobinConfig{
Production: convertAIModelWeightsToModelConfigs(aiModelBasedRoundRobin.ProductionModels, apimEndpints, true),
Sandbox: convertAIModelWeightsToModelConfigs(aiModelBasedRoundRobin.SandboxModels, apimEndpints, false),
SuspendDuration: fmt.Sprintf("%d", aiModelBasedRoundRobin.OnQuotaExceedSuspendDuration),
}
jsonBytes, err := json.Marshal(wrr)
if err != nil {
logger.LoggerMgtServer.Errorf("Error marshaling WeightedRoundRobinConfigs to JSON: %+v", err)
}
jsonStr := string(jsonBytes)
singleQuoted := strings.ReplaceAll(jsonStr, `"`, `'`)
apiPolicy := OperationPolicy{
PolicyName: constants.ModelWeightedRoundRobin,
PolicyVersion: constants.V1,
PolicyType: constants.CommonType,
Parameters: WeightedRoundRobinConfigs{
WeightedRoundRobinConfigs: singleQuoted,
},
}
requestOperationPolicies = append(requestOperationPolicies, apiPolicy)
}
data["data"].(map[string]interface{})["apiPolicies"] = OperationPolicies{
Request: requestOperationPolicies,
}

logger.LoggerMgtServer.Infof("API Yaml: %+v", data)
yamlBytes, _ := yaml.Marshal(data)
logger.LoggerMgtServer.Debugf("Endpoint Yaml: %v", endpointsData)
logger.LoggerMgtServer.Infof("Endpoint Yaml: %v", endpointsData)
endpointBytes, _ := yaml.Marshal(endpointsData)
return string(yamlBytes), definition, string(endpointBytes)
}
Expand Down Expand Up @@ -691,6 +723,7 @@ type OperationPolicy struct {
PolicyName string `yaml:"policyName"`
PolicyVersion string `yaml:"policyVersion"`
PolicyID string `yaml:"policyId,omitempty"`
PolicyType string `yaml:"policyType,omitempty"`
Parameters FilterParameters `yaml:"parameters"`
}

Expand All @@ -699,6 +732,29 @@ type FilterParameters interface {
isFilterParameters()
}

func (m WeightedRoundRobinConfigs) isFilterParameters() {}

// WeightedRoundRobinConfigs holds any additional parameter data for a RequestPolicy
type WeightedRoundRobinConfigs struct {
WeightedRoundRobinConfigs string `yaml:"weightedRoundRobinConfigs"`
}

func (m ModelBasedRoundRobinConfig) isFilterParameters() {}

// ModelConfig holds the configuration details of a model
type ModelConfig struct {
Model string `json:"model" yaml:"model"`
EndpointID string `json:"endpointId" yaml:"endpointId"`
Weight int `json:"weight" yaml:"weight"`
}

// ModelBasedRoundRobinConfig holds the configuration details of the transformer
type ModelBasedRoundRobinConfig struct {
Production []ModelConfig `json:"production" yaml:"production"`
Sandbox []ModelConfig `json:"sandbox" yaml:"sandbox"`
SuspendDuration string `json:"suspendDuration" yaml:"suspendDuration"`
}

func (h Header) isFilterParameters() {}

// Header contains the request and response header modifier information
Expand Down Expand Up @@ -757,7 +813,30 @@ type Scope struct {
Bindings []string `yaml:"bindings"`
}

func extractOperations(event APICPEvent) ([]APIOperation, []ScopeWrapper, error) {
func convertAIModelWeightsToModelConfigs(weights []AIModelWeight, apimEndpoints []APIMEndpoint, isProd bool) []ModelConfig {
var configs []ModelConfig
for _, weight := range weights {
var endpointID string
for _, endpoint := range apimEndpoints {
if endpoint.EndpointConfig.ProductionEndpoints.URL == weight.Endpoint {
endpointID = endpoint.EndpointUUID
break
}
if endpoint.EndpointConfig.SandboxEndpoints.URL == weight.Endpoint {
endpointID = endpoint.EndpointUUID
break
}
}
configs = append(configs, ModelConfig{
Model: weight.Model,
EndpointID: endpointID,
Weight: weight.Weight,
})
}
return configs
}

func extractOperations(event APICPEvent, apimEndpoints []APIMEndpoint) ([]APIOperation, []ScopeWrapper, error) {
var apiOperations []APIOperation
var requestOperationPolicies []OperationPolicy
var responseOperationPolicies []OperationPolicy
Expand Down Expand Up @@ -793,11 +872,23 @@ func extractOperations(event APICPEvent) ([]APIOperation, []ScopeWrapper, error)
Name: scope,
DisplayName: scope,
Description: scope,
Bindings: []string{},
},
Shared: false,
}
}
aiModelBasedRoundRobin := operationFromDP.AIModelBasedRoundRobin
if aiModelBasedRoundRobin != nil {
operationPolicy := OperationPolicy{
PolicyName: constants.ModelWeightedRoundRobin,
PolicyVersion: constants.V1,
Parameters: ModelBasedRoundRobinConfig{
Production: convertAIModelWeightsToModelConfigs(aiModelBasedRoundRobin.ProductionModels, apimEndpoints, true),
Sandbox: convertAIModelWeightsToModelConfigs(aiModelBasedRoundRobin.SandboxModels, apimEndpoints, false),
SuspendDuration: fmt.Sprintf("%d", aiModelBasedRoundRobin.OnQuotaExceedSuspendDuration),
},
}
requestOperationPolicies = append(requestOperationPolicies, operationPolicy)
}
// Process filters
for _, operationLevelFilter := range operationFromDP.Filters {
switch filter := operationLevelFilter.(type) {
Expand Down
Loading
Loading