|
| 1 | +package backend |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "strconv" |
| 6 | + "time" |
| 7 | + |
| 8 | + discoveryv1 "k8s.io/api/discovery/v1" |
| 9 | + "k8s.io/apimachinery/pkg/runtime" |
| 10 | + "k8s.io/client-go/tools/record" |
| 11 | + klog "k8s.io/klog/v2" |
| 12 | + ctrl "sigs.k8s.io/controller-runtime" |
| 13 | + "sigs.k8s.io/controller-runtime/pkg/builder" |
| 14 | + "sigs.k8s.io/controller-runtime/pkg/client" |
| 15 | + "sigs.k8s.io/controller-runtime/pkg/predicate" |
| 16 | + "sigs.k8s.io/gateway-api-inference-extension/api/v1alpha1" |
| 17 | + logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/ext-proc/util/logging" |
| 18 | +) |
| 19 | + |
| 20 | +var serviceOwnerLabel = "kubernetes.io/service-name" |
| 21 | + |
| 22 | +type EndpointSliceReconciler struct { |
| 23 | + client.Client |
| 24 | + Scheme *runtime.Scheme |
| 25 | + Record record.EventRecorder |
| 26 | + ServiceName string |
| 27 | + Zone string |
| 28 | + Datastore *K8sDatastore |
| 29 | +} |
| 30 | + |
| 31 | +func (c *EndpointSliceReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { |
| 32 | + klogV := klog.V(logutil.DEFAULT) |
| 33 | + |
| 34 | + inferencePool, err := c.Datastore.getInferencePool() |
| 35 | + if err != nil { |
| 36 | + klogV.InfoS("Skipping reconciling EndpointSlice because the InferencePool is not available yet", "error", err) |
| 37 | + return ctrl.Result{Requeue: true, RequeueAfter: time.Second}, nil |
| 38 | + } |
| 39 | + |
| 40 | + klogV.InfoS("Reconciling EndpointSlice", "name", req.NamespacedName) |
| 41 | + |
| 42 | + endpointSlice := &discoveryv1.EndpointSlice{} |
| 43 | + if err := c.Get(ctx, req.NamespacedName, endpointSlice); err != nil { |
| 44 | + klogV.ErrorS(err, "Unable to get EndpointSlice", "name", req.NamespacedName) |
| 45 | + return ctrl.Result{}, err |
| 46 | + } |
| 47 | + c.updateDatastore(endpointSlice, inferencePool) |
| 48 | + |
| 49 | + return ctrl.Result{}, nil |
| 50 | +} |
| 51 | + |
| 52 | +// TODO: Support multiple endpointslices for a single service |
| 53 | +func (c *EndpointSliceReconciler) updateDatastore( |
| 54 | + slice *discoveryv1.EndpointSlice, |
| 55 | + inferencePool *v1alpha1.InferencePool, |
| 56 | +) { |
| 57 | + klogV := klog.V(logutil.DEFAULT) |
| 58 | + podMap := make(map[Pod]bool) |
| 59 | + |
| 60 | + for _, endpoint := range slice.Endpoints { |
| 61 | + klogV.InfoS("Processing endpoint", "zone", c.Zone, "endpoint", endpoint) |
| 62 | + if c.validPod(endpoint) { |
| 63 | + pod := Pod{ |
| 64 | + Name: endpoint.TargetRef.Name, |
| 65 | + Address: endpoint.Addresses[0] + ":" + strconv.Itoa(int(inferencePool.Spec.TargetPortNumber)), |
| 66 | + } |
| 67 | + podMap[pod] = true |
| 68 | + klogV.InfoS("Storing pod", "pod", pod) |
| 69 | + c.Datastore.pods.Store(pod, true) |
| 70 | + } |
| 71 | + } |
| 72 | + |
| 73 | + removeOldPods := func(k, v any) bool { |
| 74 | + pod, ok := k.(Pod) |
| 75 | + if !ok { |
| 76 | + klogV.ErrorS(nil, "Unable to cast key to Pod", "key", k) |
| 77 | + return false |
| 78 | + } |
| 79 | + if _, ok := podMap[pod]; !ok { |
| 80 | + klogV.InfoS("Removing pod", "pod", pod) |
| 81 | + c.Datastore.pods.Delete(pod) |
| 82 | + } |
| 83 | + return true |
| 84 | + } |
| 85 | + c.Datastore.pods.Range(removeOldPods) |
| 86 | +} |
| 87 | + |
| 88 | +func (c *EndpointSliceReconciler) SetupWithManager(mgr ctrl.Manager) error { |
| 89 | + ownsEndPointSlice := func(object client.Object) bool { |
| 90 | + // Check if the object is an EndpointSlice |
| 91 | + endpointSlice, ok := object.(*discoveryv1.EndpointSlice) |
| 92 | + if !ok { |
| 93 | + return false |
| 94 | + } |
| 95 | + |
| 96 | + gotLabel := endpointSlice.ObjectMeta.Labels[serviceOwnerLabel] |
| 97 | + wantLabel := c.ServiceName |
| 98 | + return gotLabel == wantLabel |
| 99 | + } |
| 100 | + |
| 101 | + return ctrl.NewControllerManagedBy(mgr). |
| 102 | + For(&discoveryv1.EndpointSlice{}, |
| 103 | + builder.WithPredicates(predicate.NewPredicateFuncs(ownsEndPointSlice))). |
| 104 | + Complete(c) |
| 105 | +} |
| 106 | + |
| 107 | +func (c *EndpointSliceReconciler) validPod(endpoint discoveryv1.Endpoint) bool { |
| 108 | + validZone := c.Zone == "" || c.Zone != "" && *endpoint.Zone == c.Zone |
| 109 | + return validZone && *endpoint.Conditions.Ready |
| 110 | +} |
0 commit comments