Skip to content

Commit de24c11

Browse files
committed
Add test coverage to validate the functionality of the get_cluster method
1 parent 97cd6ea commit de24c11

4 files changed

+56
-13
lines changed

tests/e2e/mnist_raycluster_sdk_aw_kind_test.py

+2-3
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from time import sleep
44

5-
from codeflare_sdk import Cluster, ClusterConfiguration, TokenAuthentication
5+
from codeflare_sdk import Cluster, ClusterConfiguration
66
from codeflare_sdk.ray.client import RayJobClient
77

88
import pytest
@@ -68,6 +68,7 @@ def run_mnist_raycluster_sdk_kind(
6868
cluster.details()
6969

7070
self.assert_jobsubmit_withoutlogin_kind(cluster, accelerator, number_of_gpus)
71+
assert_get_cluster_and_jobsubmit(self, "mnist", accelerator, number_of_gpus)
7172

7273
# Assertions
7374

@@ -106,8 +107,6 @@ def assert_jobsubmit_withoutlogin_kind(self, cluster, accelerator, number_of_gpu
106107

107108
client.delete_job(submission_id)
108109

109-
cluster.down()
110-
111110
def assert_job_completion(self, status):
112111
if status == "SUCCEEDED":
113112
print(f"Job has completed: '{status}'")

tests/e2e/mnist_raycluster_sdk_kind_test.py

+3-5
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from time import sleep
44

5-
from codeflare_sdk import Cluster, ClusterConfiguration, TokenAuthentication
5+
from codeflare_sdk import Cluster, ClusterConfiguration
66
from codeflare_sdk.ray.client import RayJobClient
77

88
import pytest
@@ -44,8 +44,6 @@ def run_mnist_raycluster_sdk_kind(
4444
num_workers=1,
4545
head_cpu_requests="500m",
4646
head_cpu_limits="500m",
47-
head_memory_requests=2,
48-
head_memory_limits=2,
4947
worker_cpu_requests="500m",
5048
worker_cpu_limits=1,
5149
worker_memory_requests=1,
@@ -68,6 +66,8 @@ def run_mnist_raycluster_sdk_kind(
6866

6967
self.assert_jobsubmit_withoutlogin_kind(cluster, accelerator, number_of_gpus)
7068

69+
assert_get_cluster_and_jobsubmit(self, "mnist", accelerator, number_of_gpus)
70+
7171
# Assertions
7272

7373
def assert_jobsubmit_withoutlogin_kind(self, cluster, accelerator, number_of_gpus):
@@ -105,8 +105,6 @@ def assert_jobsubmit_withoutlogin_kind(self, cluster, accelerator, number_of_gpu
105105

106106
client.delete_job(submission_id)
107107

108-
cluster.down()
109-
110108
def assert_job_completion(self, status):
111109
if status == "SUCCEEDED":
112110
print(f"Job has completed: '{status}'")

tests/e2e/mnist_raycluster_sdk_oauth_test.py

+6-5
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@
22

33
from time import sleep
44

5-
from codeflare_sdk import Cluster, ClusterConfiguration, TokenAuthentication
5+
from codeflare_sdk import (
6+
Cluster,
7+
ClusterConfiguration,
8+
TokenAuthentication,
9+
)
610
from codeflare_sdk.ray.client import RayJobClient
711

812
import pytest
@@ -44,8 +48,6 @@ def run_mnist_raycluster_sdk_oauth(self):
4448
num_workers=1,
4549
head_cpu_requests="500m",
4650
head_cpu_limits="500m",
47-
head_memory_requests=4,
48-
head_memory_limits=4,
4951
worker_cpu_requests=1,
5052
worker_cpu_limits=1,
5153
worker_memory_requests=1,
@@ -68,6 +70,7 @@ def run_mnist_raycluster_sdk_oauth(self):
6870

6971
self.assert_jobsubmit_withoutLogin(cluster)
7072
self.assert_jobsubmit_withlogin(cluster)
73+
assert_get_cluster_and_jobsubmit(self, "mnist")
7174

7275
# Assertions
7376

@@ -132,8 +135,6 @@ def assert_jobsubmit_withlogin(self, cluster):
132135

133136
client.delete_job(submission_id)
134137

135-
cluster.down()
136-
137138
def assert_job_completion(self, status):
138139
if status == "SUCCEEDED":
139140
print(f"Job has completed: '{status}'")

tests/e2e/support.py

+45
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import random
44
import string
55
import subprocess
6+
from codeflare_sdk import get_cluster
67
from kubernetes import client, config
78
import kubernetes.client
89
from codeflare_sdk.common.kubernetes_cluster.kube_api_helpers import (
@@ -348,3 +349,47 @@ def get_nodes_by_label(self, node_labels):
348349
label_selector = ",".join(f"{k}={v}" for k, v in node_labels.items())
349350
nodes = self.api_instance.list_node(label_selector=label_selector)
350351
return [node.metadata.name for node in nodes.items]
352+
353+
354+
def assert_get_cluster_and_jobsubmit(
355+
self, cluster_name, accelerator=None, number_of_gpus=None
356+
):
357+
# Retrieve the cluster
358+
cluster = get_cluster(cluster_name, self.namespace)
359+
360+
cluster.details()
361+
362+
cluster.config.verify_tls = False
363+
364+
# Initialize the job client
365+
client = cluster.job_client
366+
367+
# Submit a job and get the submission ID
368+
env_vars = (
369+
get_setup_env_variables(ACCELERATOR=accelerator)
370+
if accelerator
371+
else get_setup_env_variables()
372+
)
373+
submission_id = client.submit_job(
374+
entrypoint="python mnist.py",
375+
runtime_env={
376+
"working_dir": "./tests/e2e/",
377+
"pip": "./tests/e2e/mnist_pip_requirements.txt",
378+
"env_vars": env_vars,
379+
},
380+
entrypoint_num_cpus=1 if number_of_gpus is None else None,
381+
entrypoint_num_gpus=number_of_gpus,
382+
)
383+
print(f"Submitted job with ID: {submission_id}")
384+
385+
# Fetch the list of jobs and validate
386+
job_list = client.list_jobs()
387+
print(f"List of Jobs: {job_list}")
388+
389+
# Validate the number of jobs in the list
390+
assert len(job_list) == 1
391+
392+
# Validate the submission ID matches
393+
assert job_list[0].submission_id == submission_id
394+
395+
cluster.down()

0 commit comments

Comments
 (0)