From 65d38f5b8434f50c24a890069f8e1b2948304fc1 Mon Sep 17 00:00:00 2001 From: Safiya <147792763+safiya2610@users.noreply.github.com> Date: Mon, 29 Jun 2026 15:35:26 +0000 Subject: [PATCH] Use sandbox pod annotation for pod lookup Signed-off-by: Safiya <147792763+safiya2610@users.noreply.github.com> --- go.sum | 2 - pkg/workloadmanager/k8s_client.go | 53 ++++++++++---------- pkg/workloadmanager/k8s_client_test.go | 67 ++++++++++++++++++++++++-- 3 files changed, 90 insertions(+), 32 deletions(-) diff --git a/go.sum b/go.sum index e0f89f392..9b1cf25ae 100644 --- a/go.sum +++ b/go.sum @@ -246,8 +246,6 @@ golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY= golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU= golang.org/x/oauth2 v0.32.0 h1:jsCblLleRMDrxMN29H3z/k1KliIvpLgCkE6R8FXXNgY= golang.org/x/oauth2 v0.32.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA= -golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= -golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= diff --git a/pkg/workloadmanager/k8s_client.go b/pkg/workloadmanager/k8s_client.go index 2c44d8e3c..ac807dd4c 100644 --- a/pkg/workloadmanager/k8s_client.go +++ b/pkg/workloadmanager/k8s_client.go @@ -21,8 +21,6 @@ import ( "fmt" "time" - "k8s.io/klog/v2" - cubeversioned "github.com/volcano-sh/agentcube/client-go/clientset/versioned" cubeinformers "github.com/volcano-sh/agentcube/client-go/informers/externalversions" runtimev1alpha1 "github.com/volcano-sh/agentcube/pkg/apis/runtime/v1alpha1" @@ -30,7 +28,6 @@ import ( apierrors "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" - "k8s.io/apimachinery/pkg/labels" "k8s.io/apimachinery/pkg/runtime" "k8s.io/client-go/dynamic" "k8s.io/client-go/dynamic/dynamicinformer" @@ -40,7 +37,9 @@ import ( "k8s.io/client-go/rest" "k8s.io/client-go/tools/cache" "k8s.io/client-go/tools/clientcmd" + "k8s.io/klog/v2" sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1" + "sigs.k8s.io/agent-sandbox/controllers" extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1" ) @@ -69,7 +68,7 @@ type K8sClient struct { dynamicClient dynamic.Interface scheme *runtime.Scheme baseConfig *rest.Config // Store base config for creating user clients - clientCache *ClientCache // LRU cache for user clients + clientCache *ClientCache // Thread-safe cache for user clients dynamicInformer dynamicinformer.DynamicSharedInformerFactory informerFactory informers.SharedInformerFactory cubeInformerFactory cubeinformers.SharedInformerFactory @@ -311,28 +310,37 @@ func (u *UserK8sClient) DeleteSandboxClaim(ctx context.Context, namespace, sandb } // GetSandboxPodIP gets the IP address of the pod corresponding to the Sandbox -func (c *K8sClient) GetSandboxPodIP(_ context.Context, namespace, sandboxName, podName string) (string, error) { +// It prefers an explicit podName (from annotation or caller) and avoids +// label/ownerReference fallback now that agent-sandbox provides the +// backing pod name via annotation. +func (c *K8sClient) GetSandboxPodIP(ctx context.Context, namespace, sandboxName, podName string) (string, error) { // If podName is provided, try to get it directly from cache first if podName != "" { pod, err := c.podLister.Pods(namespace).Get(podName) if err == nil && pod != nil { return validateAndGetPodIP(pod) } - klog.Infof("failed to get sandbox pod %s/%s: %v, try get pod by sandbox-name label", namespace, podName, err) + klog.Infof("failed to get sandbox pod %s/%s: %v", namespace, podName, err) } - // Find pod through label selector (sandbox-name label we set) - pods, err := c.podLister.Pods(namespace).List(labels.SelectorFromSet(map[string]string{SandboxNameLabelKey: sandboxName})) - if err != nil { - return "", fmt.Errorf("failed to list pods from cache: %w", err) - } - // Find the pod that belongs to this sandbox by checking ownerReferences - for _, pod := range pods { - for _, ownerRef := range pod.OwnerReferences { - if ownerRef.Kind == "Sandbox" && ownerRef.Name == sandboxName { - if ownerRef.Controller == nil || *ownerRef.Controller { - return validateAndGetPodIP(pod) - } + + // If podName is not provided, try to read the annotation on the Sandbox + // resource which is the authoritative source after agent-sandbox v0.3.10. + if podName == "" { + if c.dynamicClient == nil { + return "", fmt.Errorf("no pod found for sandbox %s: dynamic client not configured", sandboxName) + } + sb, err := c.dynamicClient.Resource(SandboxGVR).Namespace(namespace).Get(ctx, sandboxName, metav1.GetOptions{}) + if err != nil { + klog.Infof("failed to get sandbox %s/%s: %v", namespace, sandboxName, err) + return "", fmt.Errorf("no pod found for sandbox %s: %w", sandboxName, err) + } + if ann := sb.GetAnnotations()[controllers.SandboxPodNameAnnotation]; ann != "" { + podName = ann + pod, err := c.podLister.Pods(namespace).Get(podName) + if err == nil && pod != nil { + return validateAndGetPodIP(pod) } + klog.Infof("failed to get sandbox pod %s/%s from annotation: %v", namespace, podName, err) } } @@ -377,13 +385,8 @@ func (c *K8sClient) WaitForSandboxReady(ctx context.Context, namespace, sandboxN continue } - // Check status.phase or status.ready - status, found, err := unstructured.NestedMap(sandbox.Object, "status") - if err != nil || !found { - continue - } - - phase, found, err := unstructured.NestedString(status, "phase") + // OPTIMIZATION: Flatten nested loop parameters using unstructured fields path syntax directly + phase, found, err := unstructured.NestedString(sandbox.Object, "status", "phase") if err != nil || !found { continue } diff --git a/pkg/workloadmanager/k8s_client_test.go b/pkg/workloadmanager/k8s_client_test.go index b14526eb0..ae5c012a3 100644 --- a/pkg/workloadmanager/k8s_client_test.go +++ b/pkg/workloadmanager/k8s_client_test.go @@ -23,8 +23,13 @@ import ( "github.com/stretchr/testify/assert" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/runtime" + dynamicfake "k8s.io/client-go/dynamic/fake" listersv1 "k8s.io/client-go/listers/core/v1" + sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1" + "sigs.k8s.io/agent-sandbox/controllers" ) // Helper function to create a pod with owner reference @@ -126,8 +131,8 @@ func TestGetSandboxPodIP_Success(t *testing.T) { podLister: mockPodLister, } - // Execute - ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "") + // Execute (pass explicit pod name; label/OwnerReference-based lookup removed) + ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "test-pod") // Verify assert.NoError(t, err, "Expected no error for valid pod") @@ -157,7 +162,7 @@ func TestGetSandboxPodIP_PodNotFound(t *testing.T) { podLister: mockPodLister, } - // Execute + // Execute without podName should return not found (no dynamic client in test) ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "") // Verify @@ -200,8 +205,8 @@ func TestGetSandboxPodIP_InvalidPodStatus(t *testing.T) { podLister: mockPodLister, } - // Execute - ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "") + // Execute (pass explicit pod name) + ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "test-pod") // Verify assert.Error(t, err, "Expected error for invalid pod status") @@ -210,3 +215,55 @@ func TestGetSandboxPodIP_InvalidPodStatus(t *testing.T) { }) } } + +// TestGetSandboxPodIP_AnnotationLookup_Success verifies that when the Sandbox +// resource carries the pod name annotation, GetSandboxPodIP resolves it and +// returns the pod IP from the pod lister. +func TestGetSandboxPodIP_AnnotationLookup_Success(t *testing.T) { + pod := createPodWithOwner("pod-anno", "ns-anno", "test-sandbox", corev1.PodRunning, "10.0.0.5") + mockPodLister := newMockPodLister() + mockPodLister.addPod(pod) + + sb := &sandboxv1alpha1.Sandbox{ + TypeMeta: metav1.TypeMeta{ + APIVersion: "agents.x-k8s.io/v1alpha1", + Kind: "Sandbox", + }, + ObjectMeta: metav1.ObjectMeta{ + Name: "test-sandbox", + Namespace: "ns-anno", + Annotations: map[string]string{ + controllers.SandboxPodNameAnnotation: "pod-anno", + }, + }, + } + + scheme := runtime.NewScheme() + if err := sandboxv1alpha1.AddToScheme(scheme); err != nil { + t.Fatalf("failed to add sandbox scheme: %v", err) + } + + dyn := dynamicfake.NewSimpleDynamicClient(scheme) + + unstructuredObj, err := runtime.DefaultUnstructuredConverter.ToUnstructured(sb) + if err != nil { + t.Fatalf("failed to convert sandbox to unstructured: %v", err) + } + _, err = dyn.Resource(SandboxGVR).Namespace("ns-anno").Create( + context.Background(), + &unstructured.Unstructured{Object: unstructuredObj}, + metav1.CreateOptions{}, + ) + if err != nil { + t.Fatalf("failed to create sandbox in fake dynamic client: %v", err) + } + + client := &K8sClient{ + podLister: mockPodLister, + dynamicClient: dyn, + } + + ip, err := client.GetSandboxPodIP(context.Background(), "ns-anno", "test-sandbox", "") + assert.NoError(t, err) + assert.Equal(t, "10.0.0.5", ip) +}