Fix

Signed-off-by: Tyler Gu <[email protected]>

performance_measurement/measure_runner.py

@@ -24,6 +24,96 @@ class MeasurementResult:
     condition_2_ts: float
 
 
+def check_tolerations(desired_toleration: dict, sts: k8s_models.V1StatefulSet):
+    if "pod" in input["spec"] and "tolerations" in input["spec"]["pod"] and input["spec"][
+            "pod"]["tolerations"]:
+        toleration_matched = True
+        for toleration in input["spec"]["pod"]["tolerations"]:
+            for toleration_in_sts in sts["spec"]["template"]["spec"]["tolerations"]:
+                if (("effect" not in toleration or toleration["effect"] == toleration_in_sts["effect"])
+                        and ("key" not in toleration or toleration["key"] == toleration_in_sts["key"])
+                        and ("operator" not in toleration or toleration["operator"] == toleration_in_sts["operator"])
+                        and ("value" not in toleration or toleration["value"] == toleration_in_sts["value"])):
+                    break
+            else:
+                toleration_matched = False
+                break
+
+        if not toleration_matched:
+            logging.info("toleration not matched")
+            return False
+
+
+def check_annotations(desired_annotations: dict, sts_object: k8s_models.V1StatefulSet) -> bool:
+    # check if annotations match
+    annotation_matched = True
+    if desired_annotations is not None:
+        # check if input annotations are in sts
+        for pod_annotation in desired_annotations:
+            if pod_annotation not in sts_object["spec"]["template"]["metadata"]["annotations"]:
+                # annotation is not in sts, but in input
+                annotation_matched = False
+                break
+            if (sts_object["spec"]["template"]["metadata"]["annotations"][pod_annotation] 
+                    != desired_annotations[pod_annotation]):
+                # annotation value not match
+                annotation_matched = False
+                break
+        if not annotation_matched:
+            return False
+
+    for pod_annotation in sts_object["spec"]["template"]["metadata"]["annotations"]:
+        if pod_annotation.startswith("anvil.dev/"):
+            # allow anvil annotations
+            continue
+        if desired_annotations is None or pod_annotation not in desired_annotations:
+            # annotation is not in input, but in sts
+            annotation_matched = False
+            break
+        if (sts_object["spec"]["template"]["metadata"]["annotations"][pod_annotation] 
+                != desired_annotations[pod_annotation]):
+            # annotation value not match
+            annotation_matched = False
+            break
+
+    return annotation_matched
+
+
+def check_labels(desired_labels: dict, sts_object: k8s_models.V1StatefulSet) -> bool:
+    # check if labels match
+    label_matched = True
+    if desired_labels is not None:
+        # check if input labels are in sts
+        for pod_label in desired_labels:
+            if pod_label not in sts_object["spec"]["template"]["metadata"]["labels"]:
+                # annotation is not in sts, but in input
+                label_matched = False
+                break
+            if (sts_object["spec"]["template"]["metadata"]["labels"][pod_label] 
+                    != desired_labels[pod_label]):
+                # annotation value not match
+                label_matched = False
+                break
+        if not label_matched:
+            return False
+
+    for pod_label in sts_object["spec"]["template"]["metadata"]["labels"]:
+        if pod_label.startswith("anvil.dev/"):
+            # allow anvil labels
+            continue
+        if desired_labels is None or pod_label not in desired_labels:
+            # annotation is not in input, but in sts
+            label_matched = False
+            break
+        if (sts_object["spec"]["template"]["metadata"]["labels"][pod_label] 
+                != desired_labels[pod_label]):
+            # annotation value not match
+            label_matched = False
+            break
+
+    return label_matched
+
+
 class MeasurementRunner(Runner):
 
     cr_config_to_zk_config = {
@@ -377,53 +467,13 @@ def wait_for_zk_spec(
                     continue
 
             # check if annotations match
-            annotation_matched = True
-            for pod_annotation in input["spec"]["pod"]["annotations"]:
-                if pod_annotation not in sts_object["spec"]["template"]["metadata"]["annotations"]:
-                    annotation_matched = False
-                    break
-                if sts_object["spec"]["template"]["metadata"]["annotations"][pod_annotation] != input["spec"]["pod"][
-                        "annotations"][pod_annotation]:
-                    annotation_matched = False
-                    break
-            if not annotation_matched:
-                continue
-
-            for pod_annotation in sts_object["spec"]["template"]["metadata"]["annotations"]:
-                if pod_annotation.startswith("anvil.dev/"):
-                    continue
-                if pod_annotation not in input["spec"]["pod"]["annotations"]:
-                    annotation_matched = False
-                    break
-                if sts_object["spec"]["template"]["metadata"]["annotations"][pod_annotation] != input["spec"]["pod"][
-                        "annotations"][pod_annotation]:
-                    annotation_matched = False
-                    break
-            if not annotation_matched:
+            input_annotations = input["spec"]["annotations"] if "annotations" in input["spec"] else None
+            if not check_annotations(input_annotations, sts_object):
                 continue
 
             # check if labels match
-            label_matched = True
-            for pod_label in input["spec"]["pod"]["labels"]:
-                if pod_label not in sts_object["spec"]["template"]["metadata"]["labels"]:
-                    label_matched = False
-                    break
-                if sts_object["spec"]["template"]["metadata"]["labels"][pod_label] != input["spec"]["pod"]["labels"][
-                        pod_label]:
-                    label_matched = False
-                    break
-            if not label_matched:
-                continue
-
-            for pod_label in sts_object["spec"]["template"]["metadata"]["labels"]:
-                if pod_label not in input["spec"]["pod"]["labels"]:
-                    label_matched = False
-                    break
-                if sts_object["spec"]["template"]["metadata"]["labels"][pod_label] != input["spec"]["pod"]["labels"][
-                        pod_label]:
-                    label_matched = False
-                    break
-            if not label_matched:
+            input_labels = input["spec"]["labels"] if "labels" in input["spec"] else None
+            if not check_labels(input_labels, sts_object):
                 continue
 
             # make sure the generation is up to date