[Bug] avoid frequent lookup of the routing strategy env (vllm-project#956)

Signed-off-by: Iceber Gu <[email protected]>

Author: Iceber

pkg/plugins/gateway/gateway_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/vllm-project/aibrix/pkg/cache"
 	routing "github.com/vllm-project/aibrix/pkg/plugins/gateway/algorithms"
+	"github.com/vllm-project/aibrix/pkg/utils"
 )
 
 func Test_ValidateRoutingStrategy(t *testing.T) {
@@ -118,6 +119,9 @@ func TestGetRoutingStrategy(t *testing.T) {
 			_ = os.Unsetenv("ROUTING_ALGORITHM")
 		}
 
+		// refresh default values, the process won't modify this environment variable during normal running
+		defaultRoutingStrategy, defaultRoutingStrategyEnabled = utils.LookupEnv(EnvRoutingAlgorithm)
+
 		routingStrategy, enabled := getRoutingStrategy(tt.headers)
 		assert.Equal(t, tt.expectedStrategy, routingStrategy, tt.message)
 		assert.Equal(t, tt.expectedEnabled, enabled, tt.message)

pkg/plugins/gateway/util.go

@@ -50,30 +50,20 @@ func validateStreamOptions(requestID string, user utils.User, jsonMap map[string
 	return nil
 }
 
+var defaultRoutingStrategy, defaultRoutingStrategyEnabled = utils.LookupEnv(EnvRoutingAlgorithm)
+
 // getRoutingStrategy retrieves the routing strategy from the headers or environment variable
 // It returns the routing strategy value and whether custom routing strategy is enabled.
 func getRoutingStrategy(headers []*configPb.HeaderValue) (string, bool) {
-	var routingStrategy string
-	routingStrategyEnabled := false
-
 	// Check headers for routing strategy
 	for _, header := range headers {
 		if strings.ToLower(header.Key) == HeaderRoutingStrategy {
-			routingStrategy = string(header.RawValue)
-			routingStrategyEnabled = true
-			break // Prioritize header value over environment variable
-		}
-	}
-
-	// If header not set, check environment variable
-	if !routingStrategyEnabled {
-		if value, exists := utils.CheckEnvExists(EnvRoutingAlgorithm); exists {
-			routingStrategy = value
-			routingStrategyEnabled = true
+			return string(header.RawValue), true
 		}
 	}
 
-	return routingStrategy, routingStrategyEnabled
+	// If header not set, use default routing strategy from environment variable
+	return defaultRoutingStrategy, defaultRoutingStrategyEnabled
 }
 
 // getRequestMessage returns input request message field which has user prompt

pkg/utils/util.go

@@ -74,10 +74,10 @@ func TrimMessage(message string) string {
 	return message
 }
 
-// CheckEnvExists checks if an environment variable exists.
+// LookupEnv retrieves an environment variable and returns whether it exists.
 // It returns the value and a boolean indicating its existence.
-func CheckEnvExists(envVar string) (string, bool) {
-	value, exists := os.LookupEnv(envVar)
+func LookupEnv(key string) (string, bool) {
+	value, exists := os.LookupEnv(key)
 	return value, exists
 }