e2e test for heterogenous cluster

Author: Ygnas

Diff for: .github/workflows/e2e_tests.yaml

@@ -73,6 +73,8 @@ jobs:
 
       - name: Setup and start KinD cluster
         uses: ./common/github-actions/kind
+        with:
+          worker-nodes: 1
 
       - name: Install NVidia GPU operator for KinD
         uses: ./common/github-actions/nvidia-gpu-operator
@@ -111,6 +113,8 @@ jobs:
           kubectl create clusterrolebinding sdk-user-localqueue-creator --clusterrole=localqueue-creator --user=sdk-user
           kubectl create clusterrole list-secrets --verb=get,list --resource=secrets
           kubectl create clusterrolebinding sdk-user-list-secrets --clusterrole=list-secrets --user=sdk-user
+          kubectl create clusterrole pod-creator --verb=get,list --resource=pods
+          kubectl create clusterrolebinding sdk-user-pod-creator --clusterrole=pod-creator --user=sdk-user
           kubectl config use-context sdk-user
 
       - name: Run e2e tests

Diff for: tests/e2e/heterogeneous_clusters_test.py

@@ -0,0 +1,76 @@
+from time import sleep
+import time
+from codeflare_sdk import (
+    Cluster,
+    ClusterConfiguration,
+    TokenAuthentication,
+    generate_cert,
+)
+
+from codeflare_sdk.common.kueue.kueue import list_local_queues
+
+import pytest
+import ray
+import math
+
+from support import *
+
+
+@pytest.mark.kind
+class TestRayLocalInteractiveOauth:
+    def setup_method(self):
+        initialize_kubernetes_client(self)
+
+    def teardown_method(self):
+        delete_namespace(self)
+        delete_kueue_resources(self)
+
+    @pytest.mark.nvidia_gpu
+    def test_heterogeneous_clusters(self):
+        create_namespace(self)
+        create_kueue_resources(self)
+        self.run_heterogeneous_clusters()
+
+    def run_heterogeneous_clusters(
+        self, gpu_resource_name="nvidia.com/gpu", number_of_gpus=0
+    ):
+        cluster_name = "test-ray-cluster-li"
+
+        used_nodes = []
+
+        for flavor in self.resource_flavors:
+            queues = list_local_queues(namespace=self.namespace, flavors=[flavor])
+            queue_name = queues[0]["name"] if queues else None
+            print(f"Using flavor: {flavor}, Queue: {queue_name}")
+            cluster = Cluster(
+                ClusterConfiguration(
+                    name=f"{cluster_name}-{flavor}",
+                    namespace=self.namespace,
+                    num_workers=1,
+                    head_cpu_requests="500m",
+                    head_cpu_limits="500m",
+                    head_memory_requests=2,
+                    head_memory_limits=2,
+                    worker_cpu_requests="500m",
+                    worker_cpu_limits=1,
+                    worker_memory_requests=1,
+                    worker_memory_limits=4,
+                    worker_extended_resource_requests={
+                        gpu_resource_name: number_of_gpus
+                    },
+                    write_to_file=True,
+                    verify_tls=False,
+                    local_queue=queue_name,
+                )
+            )
+            cluster.up()
+            time.sleep(2)
+            pod_name = f"{cluster_name}-{flavor}"
+            node_name = get_pod_node(self, self.namespace, pod_name)
+            print(f"Cluster {cluster_name}-{flavor} is running on node: {node_name}")
+            time.sleep(2)
+            assert (
+                node_name not in used_nodes
+            ), f"Node {node_name} was already used by another flavor."
+            used_nodes.append(node_name)
+            cluster.down()

Diff for: tests/e2e/support.py

@@ -65,19 +65,30 @@ def create_namespace(self):
         return RuntimeError(e)
 
 
-def create_new_resource_flavor(self):
-    self.resource_flavor = f"test-resource-flavor-{random_choice()}"
-    create_resource_flavor(self, self.resource_flavor)
+def create_new_resource_flavor(self, num_flavors=2):
+    self.resource_flavors = []
+    for i in range(num_flavors):
+        default = i < 1
+        resource_flavor = f"test-resource-flavor-{random_choice()}"
+        create_resource_flavor(self, resource_flavor, default)
+        self.resource_flavors.append(resource_flavor)
 
 
-def create_new_cluster_queue(self):
-    self.cluster_queue = f"test-cluster-queue-{random_choice()}"
-    create_cluster_queue(self, self.cluster_queue, self.resource_flavor)
+def create_new_cluster_queue(self, num_queues=2):
+    self.cluster_queues = []
+    for i in range(num_queues):
+        cluster_queue_name = f"test-cluster-queue-{random_choice()}"
+        create_cluster_queue(self, cluster_queue_name, self.resource_flavors[i])
+        self.cluster_queues.append(cluster_queue_name)
 
 
-def create_new_local_queue(self):
-    self.local_queue = f"test-local-queue-{random_choice()}"
-    create_local_queue(self, self.cluster_queue, self.local_queue)
+def create_new_local_queue(self, num_queues=2):
+    self.local_queues = []
+    for i in range(num_queues):
+        is_default = i == 0
+        local_queue_name = f"test-local-queue-{random_choice()}"
+        create_local_queue(self, self.cluster_queues[i], local_queue_name, is_default)
+        self.local_queues.append(local_queue_name)
 
 
 def create_namespace_with_name(self, namespace_name):
@@ -132,7 +143,7 @@ def create_cluster_queue(self, cluster_queue, flavor):
                                 {"name": "memory", "nominalQuota": "36Gi"},
                                 {"name": "nvidia.com/gpu", "nominalQuota": 1},
                             ],
-                        }
+                        },
                     ],
                 }
             ],
@@ -161,11 +172,21 @@ def create_cluster_queue(self, cluster_queue, flavor):
     self.cluster_queue = cluster_queue
 
 
-def create_resource_flavor(self, flavor):
+def create_resource_flavor(self, flavor, default=True):
     resource_flavor_json = {
         "apiVersion": "kueue.x-k8s.io/v1beta1",
         "kind": "ResourceFlavor",
         "metadata": {"name": flavor},
+        "spec": {
+            "nodeLabels": {"worker-1" if default else "ingress-ready": "true"},
+            "tolerations": [
+                {
+                    "key": "node-role.kubernetes.io/control-plane",
+                    "operator": "Exists",
+                    "effect": "NoSchedule",
+                }
+            ],
+        },
     }
 
     try:
@@ -190,14 +211,14 @@ def create_resource_flavor(self, flavor):
     self.resource_flavor = flavor
 
 
-def create_local_queue(self, cluster_queue, local_queue):
+def create_local_queue(self, cluster_queue, local_queue, is_default=True):
     local_queue_json = {
         "apiVersion": "kueue.x-k8s.io/v1beta1",
         "kind": "LocalQueue",
         "metadata": {
             "namespace": self.namespace,
             "name": local_queue,
-            "annotations": {"kueue.x-k8s.io/default-queue": "true"},
+            "annotations": {"kueue.x-k8s.io/default-queue": str(is_default).lower()},
         },
         "spec": {"clusterQueue": cluster_queue},
     }
@@ -235,25 +256,43 @@ def create_kueue_resources(self):
 
 def delete_kueue_resources(self):
     # Delete if given cluster-queue exists
-    try:
-        self.custom_api.delete_cluster_custom_object(
-            group="kueue.x-k8s.io",
-            plural="clusterqueues",
-            version="v1beta1",
-            name=self.cluster_queue,
-        )
-        print(f"\n'{self.cluster_queue}' cluster-queue deleted")
-    except Exception as e:
-        print(f"\nError deleting cluster-queue '{self.cluster_queue}' : {e}")
+    for cq in self.cluster_queues:
+        try:
+            self.custom_api.delete_cluster_custom_object(
+                group="kueue.x-k8s.io",
+                plural="clusterqueues",
+                version="v1beta1",
+                name=cq,
+            )
+            print(f"\n'{cq}' cluster-queue deleted")
+        except Exception as e:
+            print(f"\nError deleting cluster-queue '{cq}' : {e}")
 
     # Delete if given resource-flavor exists
+    for flavor in self.resource_flavors:
+        try:
+            self.custom_api.delete_cluster_custom_object(
+                group="kueue.x-k8s.io",
+                plural="resourceflavors",
+                version="v1beta1",
+                name=flavor,
+            )
+            print(f"'{flavor}' resource-flavor deleted")
+        except Exception as e:
+            print(f"\nError deleting resource-flavor '{flavor}': {e}")
+
+
+def get_pod_node(self, namespace, name):
+    label_selector = f"ray.io/cluster={name}"
     try:
-        self.custom_api.delete_cluster_custom_object(
-            group="kueue.x-k8s.io",
-            plural="resourceflavors",
-            version="v1beta1",
-            name=self.resource_flavor,
+        pods = self.api_instance.list_namespaced_pod(
+            namespace, label_selector=label_selector
         )
-        print(f"'{self.resource_flavor}' resource-flavor deleted")
+        if not pods.items:
+            raise ValueError(
+                f"No pods found with label 'ray.io/cluster={name}' in namespace '{namespace}'"
+            )
+        pod = pods.items[0]
+        return pod.spec.node_name
     except Exception as e:
-        print(f"\nError deleting resource-flavor '{self.resource_flavor}' : {e}")
+        print(f"\nError retrieving pod node: {e}")