update filter logic

backend/src/cache/deployer/cache-configmap.yaml.template

@@ -10,6 +10,9 @@ webhooks:
         namespace: ${NAMESPACE}
         path: "/mutate"
       caBundle: ${CA_BUNDLE}
+    objectSelector:
+      matchLabels:
+        pipelines.kubeflow.org/cache_enabled: true
     rules:
       - operations: [ "CREATE" ]
         apiGroups: [""]

backend/src/cache/server/mutation.go

@@ -76,8 +76,8 @@ func MutatePodIfCached(req *v1beta1.AdmissionRequest, clientMgr ClientManagerInt
 	}
 
 	// Pod filtering to only cache KFP argo pods except TFX pods
-	if !isValidPod(&pod) {
-		log.Printf("This pod %s is not a valid pod.", pod.ObjectMeta.Name)
+	if isTFXPod(&pod) {
+		log.Printf("This pod %s is created by tfx pipelines.", pod.ObjectMeta.Name)
 		return nil, nil
 	}
 
@@ -195,41 +195,14 @@ func getValueFromSerializedMap(serializedMap string, key string) string {
 	return value
 }
 
-func isValidPod(pod *corev1.Pod) bool {
-	annotations := pod.ObjectMeta.Annotations
-	if annotations == nil || len(annotations) == 0 {
-		log.Printf("The annotation of this pod %s is empty.", pod.ObjectMeta.Name)
-		return false
-	}
-	if !isKFPArgoPod(&annotations, pod.ObjectMeta.Name) {
-		log.Printf("This pod %s is not created by KFP.", pod.ObjectMeta.Name)
-		return false
-	}
+func isTFXPod(pod *corev1.Pod) bool {
 	containers := pod.Spec.Containers
-	if containers != nil && len(containers) != 0 && isTFXPod(&containers) {
-		log.Printf("This pod %s is created by TFX pipelines.", pod.ObjectMeta.Name)
-		return false
-	}
-	return true
-}
-
-func isKFPArgoPod(annotations *map[string]string, podName string) bool {
-	// is argo pod or not
-	if _, exists := (*annotations)[ArgoWorkflowNodeName]; !exists {
-		log.Printf("This pod %s is not created by Argo.", podName)
-		return false
-	}
-	// is KFP pod or not
-	if v, exists := (*annotations)[KFPAnnotationKeyCacheEnabled]; exists && v == KFPAnnotationValueCacheEnabled {
+	if containers == nil || len(containers) == 0 {
+		log.Printf("This pod container does not exist.")
 		return true
 	}
-
-	return false
-}
-
-func isTFXPod(containers *[]corev1.Container) bool {
 	var mainContainers []corev1.Container
-	for _, c := range *containers {
+	for _, c := range containers {
 		if c.Name != "" && c.Name == "main" {
 			mainContainers = append(mainContainers, c)
 		}

backend/src/cache/server/mutation_test.go

@@ -36,12 +36,12 @@ var (
 		},
 		ObjectMeta: metav1.ObjectMeta{
 			Annotations: map[string]string{
-				ArgoWorkflowNodeName:         "test_node",
-				ArgoWorkflowTemplate:         `{"name": "test_template"}`,
-				KFPAnnotationKeyCacheEnabled: KFPAnnotationValueCacheEnabled,
+				ArgoWorkflowNodeName: "test_node",
+				ArgoWorkflowTemplate: `{"name": "test_template"}`,
 			},
 			Labels: map[string]string{
-				ArgoCompleteLabelKey: "true",
+				ArgoCompleteLabelKey:         "true",
+				KFPAnnotationKeyCacheEnabled: KFPAnnotationValueCacheEnabled,
 			},
 		},
 		Spec: corev1.PodSpec{
@@ -107,22 +107,6 @@ func TestMutatePodIfCachedWithDecodeError(t *testing.T) {
 	assert.Contains(t, err.Error(), "could not deserialize pod object")
 }
 
-func TestMutatePodIfCachedWithNonKFPPod(t *testing.T) {
-	nonKFPPod := *fakePod
-	delete(nonKFPPod.Annotations, KFPAnnotationKeyCacheEnabled)
-	patchOperation, err := MutatePodIfCached(GetFakeRequestFromPod(&nonKFPPod), fakeClientManager)
-	assert.Nil(t, patchOperation)
-	assert.Nil(t, err)
-}
-
-func TestMutatePodIfCachedWithNonArgoPod(t *testing.T) {
-	nonArgoPod := *fakePod
-	delete(nonArgoPod.Annotations, ArgoWorkflowNodeName)
-	patchOperation, err := MutatePodIfCached(GetFakeRequestFromPod(&nonArgoPod), fakeClientManager)
-	assert.Nil(t, patchOperation)
-	assert.Nil(t, err)
-}
-
 func TestMutatePodIfCachedWithTFXPod(t *testing.T) {
 	tfxPod := *fakePod
 	mainContainerCommand := append(tfxPod.Spec.Containers[0].Command, "/tfx-src/"+TFXPodSuffix)