Fix router panic when all target model weights are zero

[WHOIM1205]

Summary

This PR fixes a critical router panic caused by a divide-by-zero error when selecting a target model from a weighted list where all weights are set to zero
A minimal defensive guard has been added to prevent rand.Intn(0) from panicking, along with focused unit tests to cover the affected edge cases.

---

Issue Description

When a ModelRoute is configured with TargetModel entries whose weights all evaluate to 0, the router crashes at runtime.

This happens because selectFromWeightedSlice computes a total weight of 0 and passes it directly to rand.Intn, which panics immediately.
Such configurations are valid and realistic during canary rollbacks, maintenance windows, or temporary target disablement.

Since this code runs on the request hot path, the panic causes a full router crash and impacts all traffic.

---

What Was Fixed

• Added a defensive guard in selectFromWeightedSlice to safely handle the case where the total weight is zero
• Ensured existing behavior remains unchanged for all valid non-zero weight scenarios
• Added table-driven unit tests covering both edge cases and normal behavior

The fix is intentionally minimal and localized.

---

Impact of the Fix

• Prevents a hard router crash caused by valid but edge-case configurations
• Eliminates a denial-of-service scenario triggered by zero-weight routes
• Improves router stability during deployments, rollbacks, and partial updates
• No behavior change for existing valid configurations

---

Code Changes

Files modified:

• pkg/kthena-router/datastore/store.go
  • Added a guard in selectFromWeightedSlice to handle totalWeight == 0
• pkg/kthena-router/datastore/store_test.go
  • Added table-driven unit tests for weighted selection logic

---

Test Verification

Unit tests were added using Go’s table-driven testing pattern, covering:

• all weights set to zero (no panic)
• single zero weight
• normal weighted selection
• single non-zero weight
• empty weight slice
• all the test cases are passed locally

[volcano-sh-bot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign hzxuzhonghu for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• pkg/OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[gemini-code-assist]

Summary of Changes

Hello @WHOIM1205, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical vulnerability in the router's weighted model selection mechanism. Previously, if all configured target model weights were zero, the selectFromWeightedSlice function would trigger a divide-by-zero panic, leading to a full router crash and denial of service. This PR introduces a minimal, localized fix by adding a guard to handle this specific edge case gracefully, ensuring router stability during various operational scenarios like canary rollbacks or temporary target disablement, without impacting existing valid configurations.

Highlights

• Router Panic Fix: Implemented a defensive guard in selectFromWeightedSlice to prevent a critical router panic caused by a divide-by-zero error when all target model weights are zero.
• Enhanced Test Coverage: Added comprehensive table-driven unit tests for selectFromWeightedSlice to cover edge cases like all zero weights, single zero weight, and empty weight slices, alongside normal weighted selection.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[WHOIM1205]

/assign @hzxuzhonghu

[gemini-code-assist]

Code Review

This pull request effectively addresses a critical panic caused by a divide-by-zero error in selectFromWeightedSlice when all model weights are zero. The fix is minimal and correct, and the new unit tests provide good coverage for the edge cases.

I've identified a related latent issue where handling an empty slice of weights could lead to a panic in the calling function. I've provided two high-severity comments with suggestions to make the implementation more robust by explicitly handling this case and updating the corresponding test.

[FAUST-BENCHOU]

Anyway we should wait for maintainer's advice

[FAUST-BENCHOU]

Check 0 idea is reasonable.But i dont think we should return zero here(It will cover up configuration errors and directly select the first target.)
for example:

rules:
  - name: my-rule
    targetModels:
      - modelServerName: server-a
        weight: 0
      - modelServerName: server-b
        weight: 0

kthena/pkg/kthena-router/datastore/store.go

Lines 941 to 950 in a8c9193

	func (s *store) selectDestination(targets []*aiv1alpha1.TargetModel) (*aiv1alpha1.TargetModel, error) {
	weightedSlice, err := toWeightedSlice(targets)
	if err != nil {
	return nil, err
	}
	index := selectFromWeightedSlice(weightedSlice)
	return targets[index], nil
	}

selectDestination will select targets[0] (server-a)

Maybe we can add check in toWeightedSlice to throw an error

kthena/pkg/kthena-router/datastore/store.go

Lines 952 to 974 in a8c9193

	func toWeightedSlice(targets []*aiv1alpha1.TargetModel) ([]uint32, error) {
	var isWeighted bool
	if targets[0].Weight != nil {
	isWeighted = true
	}
	res := make([]uint32, len(targets))
	for i, target := range targets {
	if (isWeighted && target.Weight == nil) || (!isWeighted && target.Weight != nil) {
	return nil, fmt.Errorf("the weight field in targetModel must be either fully specified or not specified")
	}
	if isWeighted {
	res[i] = *target.Weight
	} else {
	// If weight is not specified, set to 1.
	res[i] = 1
	}
	}
	return res, nil
	}

	if isWeighted && totalWeight == 0 {
		return nil, fmt.Errorf("the sum of weights in targetModels must be greater than 0")
	}

to inform users that their coniguration is not right

[WHOIM1205]

thanks for the suggestion
for this pr i’m intentionally keeping the fix scoped to preventing the runtime panic on the request path returning a deterministic fallback here avoids crashing the router for edgecase or transitional configurations
validating or rejecting misconfiguration earlier (e.g.,, at config or controller level) makes sense but i’d prefer to keep this change defensive and minimal rather than introducing a new panic in the routing path

[YaoZengzeng]

rules:
  - name: my-rule
    targetModels:
      - modelServerName: server-a
        weight: 0
      - modelServerName: server-b
        weight: 0

We should deny this CR in validating webhook actually.

	if isWeighted && totalWeight == 0 {
		return nil, fmt.Errorf("the sum of weights in targetModels must be greater than 0")
	}

Also this defensive validating make sense to me.

[WHOIM1205]

That makes sense
for this pr i kept the change focused on avoiding the runtime panic in the router path adding validation in the webhook to reject all zero weights sounds like the right place for enforcing configuration correctness and i’m happy to follow up with a separate change for that

[hzxuzhonghu]

I agree with @YaoZengzeng YOu should verify it in validation as well

If all the weight is set to 0. Why do you choose targets[0]

[WHOIM1205]

in the all zero case returning targets[0] is just a deterministic fallback to keep the router from crashing on the request path not a correct routing decision
I agree this configuration should ideally be rejected earlier via validation this pr is intentionally scoped to making the runtime path defensive
Happy to align if you’d prefer validation to be included here.

[WHOIM1205]

@YaoZengzeng is there anything i can change in this pr

[WHOIM1205]

@hzxuzhonghu is there anything i can change in this pr

[LiZhenCheng9527]

/cc @YaoZengzeng PTAL

[volcano-sh-bot]

@LiZhenCheng9527: GitHub didn't allow me to request PR reviews from the following users: PTAL.

Note that only volcano-sh members and repo collaborators can review this PR, and authors cannot review their own PRs.

Details

In response to this:

/cc @YaoZengzeng PTAL

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.