e2e test for heterogenous cluster

.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

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)
+
+ 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])
+ print(queues)
+ 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.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()

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": {"ingress-ready" if default else "worker-1": "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,47 @@ 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
- try:
- self.custom_api.delete_cluster_custom_object(
- group="kueue.x-k8s.io",
- plural="resourceflavors",
- version="v1beta1",
- name=self.resource_flavor,
- )
- print(f"'{self.resource_flavor}' resource-flavor deleted")
- except Exception as e:
- print(f"\nError deleting resource-flavor '{self.resource_flavor}' : {e}")
+ 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(namespace, name):
+ label = f"ray.io/cluster={name}"
+ command = [
+ "get",
+ "pods",
+ "-n",
+ namespace,
+ "-l",
+ label,
+ "-o",
+ "jsonpath={.items[0].spec.nodeName}",
+ ]
+ node_name = run_oc_command(command)
+ if node_name:
+ return node_name.strip()
+ raise ValueError(
+ f"Could not retrieve node for pod with label 'ray.io/cluster={name}' in namespace '{namespace}'"
+ )