|
1 | 1 | package backend
|
2 | 2 |
|
3 | 3 | import (
|
| 4 | + "context" |
4 | 5 | "errors"
|
5 | 6 | "math/rand"
|
| 7 | + "strconv" |
6 | 8 | "sync"
|
7 | 9 |
|
8 | 10 | corev1 "k8s.io/api/core/v1"
|
| 11 | + "k8s.io/apimachinery/pkg/labels" |
9 | 12 | "k8s.io/klog/v2"
|
| 13 | + "sigs.k8s.io/controller-runtime/pkg/client" |
10 | 14 | "sigs.k8s.io/gateway-api-inference-extension/api/v1alpha1"
|
11 | 15 | logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/ext-proc/util/logging"
|
12 | 16 | )
|
@@ -111,3 +115,41 @@ func IsCritical(model *v1alpha1.InferenceModel) bool {
|
111 | 115 | }
|
112 | 116 | return false
|
113 | 117 | }
|
| 118 | + |
| 119 | +func (ds *K8sDatastore) LabelsMatch(podLabels map[string]string) bool { |
| 120 | + poolSelector := selectorFromInferencePoolSelector(ds.inferencePool.Spec.Selector) |
| 121 | + podSet := labels.Set(podLabels) |
| 122 | + return poolSelector.Matches(podSet) |
| 123 | +} |
| 124 | + |
| 125 | +func (ds *K8sDatastore) flushPodsAndRefetch(ctx context.Context, ctrlClient client.Client, newServerPool *v1alpha1.InferencePool) { |
| 126 | + podList := &corev1.PodList{} |
| 127 | + if err := ctrlClient.List(ctx, podList, &client.ListOptions{ |
| 128 | + LabelSelector: selectorFromInferencePoolSelector(newServerPool.Spec.Selector), |
| 129 | + Namespace: newServerPool.Namespace, |
| 130 | + }); err != nil { |
| 131 | + klog.Error(err, "error listing clients") |
| 132 | + } |
| 133 | + ds.pods.Clear() |
| 134 | + |
| 135 | + for _, k8sPod := range podList.Items { |
| 136 | + pod := Pod{ |
| 137 | + Name: k8sPod.Name, |
| 138 | + Address: k8sPod.Status.PodIP + ":" + strconv.Itoa(int(newServerPool.Spec.TargetPortNumber)), |
| 139 | + } |
| 140 | + ds.pods.Store(pod, true) |
| 141 | + } |
| 142 | + |
| 143 | +} |
| 144 | + |
| 145 | +func selectorFromInferencePoolSelector(selector map[v1alpha1.LabelKey]v1alpha1.LabelValue) labels.Selector { |
| 146 | + return labels.SelectorFromSet(stripLabelKeyAliasFromLabelMap(selector)) |
| 147 | +} |
| 148 | + |
| 149 | +func stripLabelKeyAliasFromLabelMap(labels map[v1alpha1.LabelKey]v1alpha1.LabelValue) map[string]string { |
| 150 | + outMap := make(map[string]string) |
| 151 | + for k, v := range labels { |
| 152 | + outMap[string(k)] = string(v) |
| 153 | + } |
| 154 | + return outMap |
| 155 | +} |
0 commit comments