Back to working library

unity-sds · Jun 12, 2024 · 46625c8 · 46625c8
1 parent 9754694
commit 46625c8
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 27 deletions.
diff --git a/airflow/dags/karpenter_test.py b/airflow/dags/karpenter_test.py
@@ -44,7 +44,10 @@
     labels={"app": POD_LABEL},
     annotations={"karpenter.sh/do-not-disrupt": "true"},
     affinity=get_affinity(
-        instance_type=["c6i.large", "c5.large"], capacity_type=["spot"], anti_affinity_label=POD_LABEL
+        capacity_type=["spot"],
+        instance_family=["c6i", "c5"],
+        instance_cpu=["2", "4"],
+        anti_affinity_label=POD_LABEL,
     ),
 )
 
@@ -55,7 +58,7 @@
     name="memory-task",
     image="busybox",
     cmds=["sleep"],
-    arguments=["10"],
+    arguments=["60"],
     retries=3,
     in_cluster=True,
     get_logs=True,
@@ -65,7 +68,10 @@
     labels={"app": POD_LABEL},
     annotations={"karpenter.sh/do-not-disrupt": "true"},
     affinity=get_affinity(
-        instance_type=["r6i.large", "r5.large"], capacity_type=["spot"], anti_affinity_label=POD_LABEL
+        capacity_type=["spot"],
+        instance_family=["r6i", "r5"],
+        instance_cpu=["2", "4"],
+        anti_affinity_label=POD_LABEL,
     ),
 )
 
@@ -87,7 +93,10 @@
     labels={"app": POD_LABEL},
     annotations={"karpenter.sh/do-not-disrupt": "true"},
     affinity=get_affinity(
-        instance_type=["m6i.large", "m5.large"], capacity_type=["spot"], anti_affinity_label=POD_LABEL
+        capacity_type=["spot"],
+        instance_family=["m6i", "m5"],
+        instance_cpu=["2", "4"],
+        anti_affinity_label=POD_LABEL,
     ),
 )
 

diff --git a/airflow/plugins/unity_sps_utils.py b/airflow/plugins/unity_sps_utils.py
@@ -3,19 +3,9 @@
 """
 
 
-def get_affinity(instance_type: list[str], capacity_type: list[str], anti_affinity_label: str = None) -> dict:
-    """
-    Function to create an affinity dictionary to specify which EC2 type the Kubernetes Pod should be running on.
-
-    Args:
-        instance_type: the possible EC2 types (e.g.: ["r7i.xlarge"])
-        capacity_type: "spot" and/or "on-demand"
-        anti_affinity_label: optional label to prevent 2 Pods with that label to be provisioned on the same Kubernetes node.
-
-    Returns:
-        the Kubernetes affinity dictionary to be added to the pod specification.
-
-    """
+def get_affinity(
+    capacity_type: list[str], instance_family: list[str], instance_cpu: list[str], anti_affinity_label: str
+):
 
     affinity = {
         "nodeAffinity": {
@@ -38,20 +28,21 @@ def get_affinity(instance_type: list[str], capacity_type: list[str], anti_affini
                     {
                         "matchExpressions": [
                             {
-                                "key": "karpenter.k8s.aws/instance-type",
+                                "key": "karpenter.k8s.aws/instance-family",
                                 "operator": "In",
-                                "values": instance_type,
-                            }
+                                "values": instance_family,
+                            },
+                            {
+                                "key": "karpenter.k8s.aws/instance-cpu",
+                                "operator": "In",
+                                "values": instance_cpu,
+                            },
                         ]
                     }
                 ]
             },
         },
-    }
-
-    # optionally add an anti_affinity_label to constraint each pod to be instantiated on a different node
-    if anti_affinity_label:
-        affinity["podAntiAffinity"] = {
+        "podAntiAffinity": {
             "requiredDuringSchedulingIgnoredDuringExecution": [
                 {
                     "labelSelector": {
@@ -66,6 +57,6 @@ def get_affinity(instance_type: list[str], capacity_type: list[str], anti_affini
                     "topologyKey": "kubernetes.io/hostname",
                 }
             ]
-        }
-
+        },
+    }
     return affinity