Repository: vllm-project/aibrix
File: pkg/plugins/gateway/algorithms/vtc/vtc_basic_test.go (Go)

Code before:

			// For sanity, make sure we're routing to a valid pod
			assert.NotEqual(t, -1, idx, "Should route to a valid pod")
		}
	})

	// Clamp fairness when tokens exceed adaptive bucket*(nPods-1)
	t.Run("ClampedFairnessHighTokens", func(t *testing.T) {
		highTracker := NewInMemorySlidingWindowTokenTracker(trackerConfig)
		highRouter := &BasicVTCRouter{cache: cache, tokenTracker: highTracker, tokenEstimator: tokenEstimator, config: routerConfig}
		podsHigh := createTestPods(3)
		highList := NewSimplePodList(podsHigh)
		// Equal loads for all pods
		for i := 1; i <= 3; i++ {
			podKey := utils.GeneratePodKey("default", fmt.Sprintf("pod%d", i))
			cache.SetPodMetric(podKey, "model1", metrics.NumRequestsRunning, 0)
		}
		// Use tokens > adaptiveBucket*2 to clamp normalizedTokens to last index=2
		_ = highTracker.UpdateTokenCount(ctx, "highUser", 10000, 0)
		addr, err := highRouter.Route(types.NewRoutingContext(ctx, "vtc-basic", "model1", "test", "req-high", "highUser"), highList)
		assert.NoError(t, err)
		// Should select middle pod (index 1) due to clamped fairness index
		assert.Equal(t, "192.168.1.2:8000", addr, "High tokens clamp fairness to middle pod index")
	})

	// Error when no pods to forward request
	t.Run("NoPods", func(t *testing.T) {
		emptyTracker := NewInMemorySlidingWindowTokenTracker(trackerConfig)
		emptyRouter := &BasicVTCRouter{cache: cache, tokenTracker: emptyTracker, tokenEstimator: tokenEstimator, config: routerConfig}
		emptyList := NewSimplePodList([]*v1.Pod{})
		routingCtx := types.NewRoutingContext(ctx, "vtc-basic", "model1", "test", "req-empty", "userX")
		addr, err := emptyRouter.Route(routingCtx, emptyList)
		assert.Error(t, err)
		assert.Empty(t, addr)
		assert.Contains(t, err.Error(), "no pods to forward request")
	})

	// Error when no ready pods available
	t.Run("NoReadyPods", func(t *testing.T) {
		nrTracker := NewInMemorySlidingWindowTokenTracker(trackerConfig)
		nrRouter := &BasicVTCRouter{cache: cache, tokenTracker: nrTracker, tokenEstimator: tokenEstimator, config: routerConfig}
		podsNR := createTestPods(2)
		for _, p := range podsNR {
			p.Status.Conditions = []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionFalse}}
		}
		nrList := NewSimplePodList(podsNR)
		routingCtx := types.NewRoutingContext(ctx, "vtc-basic", "model1", "test", "req-nr", "userY")
		addr, err := nrRouter.Route(routingCtx, nrList)
		assert.Error(t, err)
		assert.Empty(t, addr)
		assert.Contains(t, err.Error(), "no ready pods")

Review comment (suggestion):
We delete the redundant case here, which is not helpful for testing. Our `Route(ctx *RoutingContext, readyPodList PodList) (string, error)` already ensures that there is no noready pod and at least one readypod scenario.

Code after:

			// For sanity, make sure we're routing to a valid pod
			assert.NotEqual(t, -1, idx, "Should route to a valid pod")
		}
	})

	// Clamp fairness when tokens exceed adaptive bucket*(nPods-1)
	t.Run("ClampedFairnessHighTokens", func(t *testing.T) {
		highTracker := NewInMemorySlidingWindowTokenTracker(trackerConfig)
		highRouter := &BasicVTCRouter{cache: cache, tokenTracker: highTracker, tokenEstimator: tokenEstimator, config: routerConfig}
		podsHigh := createTestPods(3)
		highList := NewSimplePodList(podsHigh)
		// Equal loads for all pods
		for i := 1; i <= 3; i++ {
			podKey := utils.GeneratePodKey("default", fmt.Sprintf("pod%d", i))
			cache.SetPodMetric(podKey, "model1", metrics.NumRequestsRunning, 0)
		}
		// Use tokens > adaptiveBucket*2 to clamp normalizedTokens to last index=2
		_ = highTracker.UpdateTokenCount(ctx, "highUser", 10000, 0)
		addr, err := highRouter.Route(types.NewRoutingContext(ctx, "vtc-basic", "model1", "test", "req-high", "highUser"), highList)
		assert.NoError(t, err)
		// Should select middle pod (index 1) due to clamped fairness index
		assert.Equal(t, "192.168.1.2:8000", addr, "High tokens clamp fairness to middle pod index")
	})

	// Random fallback when user is nil
	t.Run("RandomFallbackNilUser", func(t *testing.T) {
		rfTracker := NewInMemorySlidingWindowTokenTracker(trackerConfig)
		rfRouter := &BasicVTCRouter{cache: cache, tokenTracker: rfTracker, tokenEstimator: tokenEstimator, config: routerConfig}
		podsRF := createTestPods(3)
		rfList := NewSimplePodList(podsRF)
		routingCtx := types.NewRoutingContext(ctx, "vtc-basic", "model1", "test", "req-rf", "")
		addr, err := rfRouter.Route(routingCtx, rfList)
		assert.NoError(t, err)
		assert.Contains(t, []string{"192.168.1.1:8000", "192.168.1.2:8000", "192.168.1.3:8000"}, addr)
	})
}

func TestWeightCombinations(t *testing.T) {
	// Use Milliseconds and set weights for testing
	t.Setenv("AIBRIX_ROUTER_VTC_TOKEN_TRACKER_TIME_UNIT", "milliseconds")
	t.Setenv("AIBRIX_ROUTER_VTC_TOKEN_TRACKER_WINDOW_SIZE", "50")
	t.Setenv("AIBRIX_ROUTER_VTC_BASIC_FAIRNESS_WEIGHT", "0.5")
	t.Setenv("AIBRIX_ROUTER_VTC_BASIC_UTILIZATION_WEIGHT", "0.5")

	ctx := context.Background()
	cache := NewSimpleCache()
	tokenEstimator := NewSimpleTokenEstimator()
	podList := NewSimplePodList(createTestPods(3))

	// Set up pod loads - High, Medium, Low