@@ -24,39 +24,39 @@ import (
2424 "sigs.k8s.io/gateway-api-inference-extension/apix/v1alpha2"
2525)
2626
27- // ModelRewriteStore encapsulates the logic for storing and retrieving
27+ // modelRewriteStore encapsulates the logic for storing and retrieving
2828// InferenceModelRewrite rules, handling precedence correctly. This struct is not
2929// thread-safe; concurrency must be managed by its consumer.
30- type ModelRewriteStore struct {
30+ type modelRewriteStore struct {
3131 genericRules []* rewriteRuleWithMetadata
3232 rulesByExactModelMatch map [string ][]* rewriteRuleWithMetadata
33- allReWrites map [types. NamespacedName ]* v1alpha2.InferenceModelRewrite
33+ allReWrites map [string ]* v1alpha2.InferenceModelRewrite
3434}
3535
36- func NewModelRewriteStore () * ModelRewriteStore {
37- return & ModelRewriteStore {
36+ func newModelRewriteStore () * modelRewriteStore {
37+ return & modelRewriteStore {
3838 genericRules : []* rewriteRuleWithMetadata {},
39- rulesByExactModelMatch : map [string ][]* rewriteRuleWithMetadata {},
40- allReWrites : map [types. NamespacedName ]* v1alpha2.InferenceModelRewrite {},
39+ rulesByExactModelMatch : map [string ][]* rewriteRuleWithMetadata {}, // Key is the exact model name.
40+ allReWrites : map [string ]* v1alpha2.InferenceModelRewrite {}, // Key is the rewrites name.
4141 }
4242}
4343
44- // Set adds or updates an InferenceModelRewrite in the store. It deconstructs the
44+ // set adds or updates an InferenceModelRewrite in the store. It deconstructs the
4545// object into individual rules and stores them in the appropriate data structures,
4646// ensuring they remain sorted by precedence.
47- func (ms * ModelRewriteStore ) Set (infModelRewrite * v1alpha2.InferenceModelRewrite ) {
48- nn := getNN (infModelRewrite )
49-
47+ func (ms * modelRewriteStore ) set (infModelRewrite * v1alpha2.InferenceModelRewrite ) {
48+ name := infModelRewrite .Name
5049 // If the rewrite object already exists, remove its old rules before adding new ones.
51- if _ , ok := ms .allReWrites [nn ]; ok {
52- ms .deleteInternal (nn )
50+ if _ , ok := ms .allReWrites [name ]; ok {
51+ ms .deleteInternal (infModelRewrite . Name )
5352 }
54- ms .allReWrites [nn ] = infModelRewrite
53+ ms .allReWrites [name ] = infModelRewrite
5554
5655 for i := range infModelRewrite .Spec .Rules {
57- ruleWithMetadata := newRuleWithMetadata (infModelRewrite , i )
58- if ruleWithMetadata == nil {
59- continue
56+ ruleWithMetadata := & rewriteRuleWithMetadata {
57+ rule : infModelRewrite .Spec .Rules [i ],
58+ createTimestamp : infModelRewrite .CreationTimestamp .Time ,
59+ parentRewriteName : name ,
6060 }
6161
6262 if ruleWithMetadata .isGeneric () {
@@ -80,22 +80,22 @@ func (ms *ModelRewriteStore) Set(infModelRewrite *v1alpha2.InferenceModelRewrite
8080 }
8181}
8282
83- // Delete removes an InferenceModelRewrite and all its associated rules from the store.
84- func (ms * ModelRewriteStore ) Delete (nn types.NamespacedName ) {
85- ms .deleteInternal (nn )
83+ // delete removes an InferenceModelRewrite and all its associated rules from the store.
84+ func (ms * modelRewriteStore ) delete (nn types.NamespacedName ) {
85+ ms .deleteInternal (nn . Name )
8686}
8787
8888// deleteInternal is the non-locking implementation for deleting a rewrite.
89- func (ms * ModelRewriteStore ) deleteInternal (nn types. NamespacedName ) {
90- if _ , ok := ms .allReWrites [nn ]; ! ok {
89+ func (ms * modelRewriteStore ) deleteInternal (n string ) {
90+ if _ , ok := ms .allReWrites [n ]; ! ok {
9191 return
9292 }
93- delete (ms .allReWrites , nn )
93+ delete (ms .allReWrites , n )
9494
9595 // Filter out the generic rules associated with the deleted rewrite.
9696 newGenericRules := make ([]* rewriteRuleWithMetadata , 0 , len (ms .genericRules ))
9797 for _ , ruleWithMd := range ms .genericRules {
98- if ruleWithMd .parentNN () != nn {
98+ if ruleWithMd .parentName () != n {
9999 newGenericRules = append (newGenericRules , ruleWithMd )
100100 }
101101 }
@@ -105,7 +105,7 @@ func (ms *ModelRewriteStore) deleteInternal(nn types.NamespacedName) {
105105 for modelName , rulesWithMd := range ms .rulesByExactModelMatch {
106106 newRules := make ([]* rewriteRuleWithMetadata , 0 , len (rulesWithMd ))
107107 for _ , r := range rulesWithMd {
108- if r .parentNN () != nn {
108+ if r .parentName () != n {
109109 newRules = append (newRules , r )
110110 }
111111 }
@@ -118,9 +118,9 @@ func (ms *ModelRewriteStore) deleteInternal(nn types.NamespacedName) {
118118 }
119119}
120120
121- // GetRule returns the single, highest-precedence rule for a given model name.
121+ // getRule returns the single, highest-precedence rule for a given model name.
122122// It prioritizes exact matches over generic ones, and among those, the oldest rule wins.
123- func (ms * ModelRewriteStore ) GetRule (modelName string ) * v1alpha2.InferenceModelRewriteRule {
123+ func (ms * modelRewriteStore ) getRule (modelName string ) * v1alpha2.InferenceModelRewriteRule {
124124 // Exact matches have the highest precedence.
125125 if rulesWithMd , ok := ms .rulesByExactModelMatch [modelName ]; ok && len (rulesWithMd ) > 0 {
126126 return & rulesWithMd [0 ].rule // The list is pre-sorted, so the first element is the oldest.
@@ -133,39 +133,21 @@ func (ms *ModelRewriteStore) GetRule(modelName string) *v1alpha2.InferenceModelR
133133 return nil
134134}
135135
136- // GetAll returns a slice of all InferenceModelRewrite objects currently in the store.
137- func (ms * ModelRewriteStore ) GetAll () []* v1alpha2.InferenceModelRewrite {
136+ // getAll returns a slice of all InferenceModelRewrite objects currently in the store.
137+ func (ms * modelRewriteStore ) getAll () []* v1alpha2.InferenceModelRewrite {
138138 rewrites := make ([]* v1alpha2.InferenceModelRewrite , 0 , len (ms .allReWrites ))
139139 for _ , rewrite := range ms .allReWrites {
140140 rewrites = append (rewrites , rewrite )
141141 }
142142 return rewrites
143143}
144144
145- func getNN (infModelRewrite * v1alpha2.InferenceModelRewrite ) types.NamespacedName {
146- return types.NamespacedName {
147- Namespace : infModelRewrite .Namespace ,
148- Name : infModelRewrite .Name ,
149- }
150- }
151-
152145// rewriteRuleWithMetadata decorates a rule with metadata from its parent object
153146// to be used in precedence sorting.
154147type rewriteRuleWithMetadata struct {
155- rule v1alpha2.InferenceModelRewriteRule
156- createTimestamp time.Time
157- parentRewriteNN types.NamespacedName
158- }
159-
160- func newRuleWithMetadata (infModelRewrite * v1alpha2.InferenceModelRewrite , ruleIdx int ) * rewriteRuleWithMetadata {
161- if ruleIdx >= len (infModelRewrite .Spec .Rules ) {
162- return nil
163- }
164- return & rewriteRuleWithMetadata {
165- rule : infModelRewrite .Spec .Rules [ruleIdx ],
166- createTimestamp : infModelRewrite .CreationTimestamp .Time ,
167- parentRewriteNN : getNN (infModelRewrite ),
168- }
148+ rule v1alpha2.InferenceModelRewriteRule
149+ createTimestamp time.Time
150+ parentRewriteName string
169151}
170152
171153func (rr rewriteRuleWithMetadata ) isGeneric () bool {
@@ -182,6 +164,6 @@ func (rr rewriteRuleWithMetadata) exactModels() map[string]bool {
182164 return modelSet
183165}
184166
185- func (rr rewriteRuleWithMetadata ) parentNN () types. NamespacedName {
186- return rr .parentRewriteNN
167+ func (rr rewriteRuleWithMetadata ) parentName () string {
168+ return rr .parentRewriteName
187169}
0 commit comments