99 ctrl "sigs.k8s.io/controller-runtime"
1010 "sigs.k8s.io/controller-runtime/pkg/client"
1111
12- llmdv1alpha1 "github.com/llm-d-incubation/workload-variant-autoscaler/api/v1alpha1"
12+ "github.com/llm-d-incubation/workload-variant-autoscaler/internal/indexers"
13+ "github.com/llm-d-incubation/workload-variant-autoscaler/internal/logging"
1314)
1415
1516// PodVAMapper maps pod names to their corresponding VariantAutoscaling objects.
@@ -24,26 +25,39 @@ func NewPodVAMapper(k8sClient client.Client) *PodVAMapper {
2425 }
2526}
2627
27- // FindVAForPod finds the VariantAutoscaling object for a pod by first finding
28- // its deployment and then finding the VA that targets that deployment.
29- //
30- // Returns the VA name if found, empty string otherwise.
28+ // FindVAForPod finds the VariantAutoscaling object for a Pod by:
29+ // 1. finding the Deployment owning the Pod
30+ // 2. finding the VariantAutoscaling that targets that Deployment, using indexed lookups.
31+ // Returns the VariantAutoscaling name if found, empty string otherwise.
3132func (m * PodVAMapper ) FindVAForPod (
3233 ctx context.Context ,
3334 podName string ,
3435 namespace string ,
3536 deployments map [string ]* appsv1.Deployment ,
36- variantAutoscalings map [string ]* llmdv1alpha1.VariantAutoscaling ,
3737) string {
38+ logger := ctrl .LoggerFrom (ctx )
39+
3840 deploymentName := m .findDeploymentForPod (ctx , podName , namespace , deployments )
3941 if deploymentName == "" {
4042 return ""
4143 }
42- return m .findVAForDeployment (deploymentName , namespace , variantAutoscalings )
44+
45+ // Use indexed lookup for VariantAutoscaling targeting this Deployment
46+ va , err := indexers .FindVAForDeployment (ctx , m .k8sClient , deploymentName , namespace )
47+ if err != nil {
48+ logger .V (logging .DEBUG ).Error (err , "failed to find VariantAutoscaling for deployment" , "deployment" , deploymentName , "namespace" , namespace )
49+ return ""
50+ }
51+
52+ if va == nil {
53+ logger .V (logging .DEBUG ).Info ("no VariantAutoscaling matched for deployment" , "deployment" , deploymentName , "namespace" , namespace )
54+ return ""
55+ }
56+
57+ return va .Name
4358}
4459
45- // findDeploymentForPod finds which deployment owns a pod using Kubernetes API.
46- // Uses the deployment's label selector to find matching pods.
60+ // findDeploymentForPod finds which Deployment owns a Pod by traversing owner references.
4761func (m * PodVAMapper ) findDeploymentForPod (
4862 ctx context.Context ,
4963 podName string ,
@@ -52,48 +66,34 @@ func (m *PodVAMapper) findDeploymentForPod(
5266) string {
5367 logger := ctrl .LoggerFrom (ctx )
5468
55- // TODO: optimize
56- for deploymentName , deployment := range deployments {
57- selector , err := metav1 .LabelSelectorAsSelector (deployment .Spec .Selector )
58- if err != nil {
59- logger .Info ("Invalid label selector for deployment" , "deployment" , deploymentName , "error" , err )
60- continue
61- }
62-
63- podList := & corev1.PodList {}
64- listOpts := & client.ListOptions {
65- Namespace : namespace ,
66- LabelSelector : selector ,
67- }
68-
69- if err := m .k8sClient .List (ctx , podList , listOpts ); err != nil {
70- logger .Info ("Failed to list pods for deployment" , "deployment" , deploymentName , "error" , err )
71- continue
72- }
73-
74- for _ , pod := range podList .Items {
75- if pod .Name == podName {
76- return deploymentName
77- }
78- }
69+ pod := & corev1.Pod {}
70+ if err := m .k8sClient .Get (ctx , client.ObjectKey {Namespace : namespace , Name : podName }, pod ); err != nil {
71+ logger .V (logging .DEBUG ).Error (err , "failed to get pod" , "pod" , podName , "namespace" , namespace )
72+ return ""
7973 }
80- return ""
81- }
8274
83- // findVAForDeployment finds the VariantAutoscaling object that targets a deployment.
84- func (m * PodVAMapper ) findVAForDeployment (
85- deploymentName string ,
86- namespace string ,
87- variantAutoscalings map [string ]* llmdv1alpha1.VariantAutoscaling ,
88- ) string {
89- for vaName , va := range variantAutoscalings {
90- if va == nil {
91- continue
92- }
93- if va .Spec .ScaleTargetRef .Name == deploymentName &&
94- va .Namespace == namespace {
95- return vaName
96- }
75+ owner := metav1 .GetControllerOf (pod )
76+ if owner == nil || owner .Kind != "ReplicaSet" {
77+ logger .V (logging .DEBUG ).Info ("Pod has no ReplicaSet owner" , "pod" , podName , "namespace" , namespace )
78+ return ""
79+ }
80+
81+ rs := & appsv1.ReplicaSet {}
82+ if err := m .k8sClient .Get (ctx , client.ObjectKey {Namespace : namespace , Name : owner .Name }, rs ); err != nil {
83+ logger .V (logging .DEBUG ).Error (err , "failed to get ReplicaSet" , "replicaset" , owner .Name , "namespace" , namespace )
84+ return ""
85+ }
86+
87+ rsOwner := metav1 .GetControllerOf (rs )
88+ if rsOwner == nil || rsOwner .Kind != "Deployment" {
89+ logger .V (logging .DEBUG ).Info ("ReplicaSet has no Deployment owner" , "replicaset" , owner .Name , "namespace" , namespace )
90+ return ""
91+ }
92+
93+ // Verify the Deployment is in our map of tracked Deployments
94+ deploymentKey := namespace + "/" + rsOwner .Name
95+ if deploy , ok := deployments [deploymentKey ]; ok && deploy != nil && deploy .Namespace == namespace {
96+ return rsOwner .Name
9797 }
9898 return ""
9999}
0 commit comments